Flow models

FlowModel([rng,] dim::Int, layers::Tuple)
FlowModel([rng,] layers::Tuple)

A flow model of dimension dim, a tuple of layers whose sizes are set and whose parameters are not: compile gives it its parameters, and the Flow node takes the CompiledFlowModel that results.

Arguments

• rng: the generator the random permutations are drawn from, those of a PermutationLayer() and of an AdditiveCouplingLayer with permute = true. Default: the task's generator.
• dim: the dimension of the input, and of the output.
• layers: the layers, applied in order: AdditiveCouplingLayers and PermutationLayers. In the second form the tuple starts with an InputLayer(dim), which gives the dimension.

Throws

An AssertionError when the second form's tuple does not start with an InputLayer.

Examples

julia> model = FlowModel(2, (AdditiveCouplingLayer(PlanarFlow()), AdditiveCouplingLayer(PlanarFlow())));

julia> length(getlayers(model))   # each coupling layer is followed by a permutation
4

julia> nr_params(model)
6

CompiledFlowModel

A flow model with its parameters set, which the Flow node takes through FlowApproximation. It is made by compile from a FlowModel, never by hand, so that every layer and coupling flow is sized for the model's dimension.

Its output and inverse are FlowMessagePassingRules.forward and FlowMessagePassingRules.backward, its Jacobians FlowMessagePassingRules.jacobian and FlowMessagePassingRules.inv_jacobian, with both at once from FlowMessagePassingRules.forward_jacobian and FlowMessagePassingRules.backward_inv_jacobian; getlayers gives its layers, nr_params the number of its parameters, and eltype(model) their element type.

compile([rng,] model::FlowModel) -> CompiledFlowModel

Give a flow model random parameters, drawn from rng, by default the task's generator. Each coupling flow draws its own, in the order of the layers: see PlanarFlow and RadialFlow for the distributions. Permutations are the model's, drawn when it was built.

See also compile(::FlowModel, ::Vector).

compile(model::FlowModel, params::Vector) -> CompiledFlowModel

Give a flow model the parameters params, which the layers take in order, each coupling flow its own consecutive slice. The result is deterministic: nothing is drawn.

Throws

An AssertionError when length(params) is not nr_params(model).

Examples

julia> model = compile(FlowModel(2, (AdditiveCouplingLayer(PlanarFlow(); permute = false),)), [1.0, 2.0, 3.0]);

julia> FlowMessagePassingRules.forward(model, [1.0, 2.0]) ≈ [1.0, 2.0 + 1.0 + tanh(2.0 + 3.0)]
true

nr_params(model)

The number of parameters of a flow model, of a layer or of a coupling flow, compiled or not: the length of the vector compile(model, params) takes.

getlayers(model)

The tuple of layers of a FlowModel or of a CompiledFlowModel.

eltype(model::CompiledFlowModel)

The element type of the compiled flow model's parameters, promoted over its layers.

InputLayer(dim::Int)

The first element of a layer tuple given to FlowModel(layers), which gives the model's dimension dim; FlowModel((InputLayer(dim), layers...)) is FlowModel(dim, layers).

Throws

An AssertionError unless dim > 1.

AdditiveCouplingLayer(flow; partition_dim::Int = 1, permute::Bool = true)

A layer of a FlowModel that adds to each coordinate a coupling flow of the one before it. For an input $x \in \mathbb{R}^d$,

\[y_1 = x_1, \qquad y_k = x_k + f_{k-1}(x_{k-1}), \quad k = 2, \dots, d,\]

with a flow $f_{k-1}$ of its own for each $k$, all of the kind of flow. Here $x_k$ is a partition of partition_dim coordinates, one coordinate by default. It is invertible whatever the flows are, $x_k = y_k - f_{k-1}(x_{k-1})$ in order, and its Jacobian is block lower-triangular with identity blocks on the diagonal, so its determinant is one. With two partitions, partition_dim = d / 2, it is the layer of Dinh, Krueger and Bengio, "NICE: Non-linear independent components estimation", arXiv:1410.8516 (2014).

Arguments

• flow: the coupling flow placeholder, PlanarFlow() or RadialFlow(). The layer's dimension is set when it is wrapped in a FlowModel, and the flows' parameters by compile.

Keywords

• partition_dim: the size of a partition, which must divide the model's dimension: a model of another dimension is an ArgumentError. Default 1, every coordinate its own partition.
• permute: whether the model places a random PermutationLayer after the layer, so that the next coupling layer mixes the coordinates in another order. Default true.

Examples

julia> model = compile(FlowModel(3, (AdditiveCouplingLayer(PlanarFlow(); permute = false),)), zeros(6));

julia> FlowMessagePassingRules.forward(model, [1.0, 2.0, 3.0]) ≈ [1.0, 3.0, 5.0]   # u = 0: each fₖ(x) = x
true

Two halves of four coordinates, as in NICE, with one planar flow on two coordinates:

julia> model = compile(FlowModel(4, (AdditiveCouplingLayer(PlanarFlow(); partition_dim = 2, permute = false),)), zeros(5));

julia> FlowMessagePassingRules.forward(model, [1.0, 2.0, 3.0, 4.0]) ≈ [1.0, 2.0, 4.0, 6.0]   # y₂ = x₂ + x₁
true

PermutationLayer(dim::Int, P::PermutationMatrix)
PermutationLayer(P::PermutationMatrix)
PermutationLayer([rng,] dim::Integer)
PermutationLayer()

A layer of a FlowModel that permutes its input, $y = P x$ for a PermutationMatrix $P$. It has no parameters, and its inverse is $x = P^\top y$.

PermutationLayer(P) takes the matrix, which must be square. PermutationLayer([rng,] dim) draws a random dim × dim one from rng, by default the task's generator. PermutationLayer() is a placeholder: its dimension and its random permutation are set when it is wrapped in a FlowModel, from the model's generator.

Examples

julia> layer = PermutationLayer(PermutationMatrix([2, 3, 1]));

julia> FlowMessagePassingRules.forward(layer, [10.0, 20.0, 30.0])
3-element Vector{Float64}:
 20.0
 30.0
 10.0

PlanarFlow(u, w, b)
PlanarFlow([rng,] dim::Int)
PlanarFlow()

The planar flow, a coupling flow for an AdditiveCouplingLayer:

\[f(x) = x + u \tanh(w^\top x + b), \qquad x, u, w \in \mathbb{R}^D, \; b \in \mathbb{R}.\]

It contracts and expands the space along $u$. It has $2D + 1$ parameters, taken by compile in the order $u, w, b$. It was introduced by Rezende and Mohamed, "Variational inference with normalizing flows", ICML 2015.

PlanarFlow(u, w, b) takes the parameters, u and w both scalars or both vectors of the same length. PlanarFlow([rng,] dim) draws them from a standard normal, using rng, by default the task's generator: scalars for dim == 1, vectors of length dim otherwise. PlanarFlow() is the placeholder an AdditiveCouplingLayer takes: its dimension is set by the model, and its parameters by compile, which draws them the same way.

Examples

julia> FlowMessagePassingRules.forward(PlanarFlow(1.0, 2.0, 3.0), 1.0) ≈ 1.0 + tanh(5.0)
true

RadialFlow(z0, α, β)
RadialFlow([rng,] dim::Int)
RadialFlow()

The radial flow, a coupling flow for an AdditiveCouplingLayer:

\[f(x) = x + \frac{\beta (x - z_0)}{\alpha + \lVert x - z_0 \rVert}, \qquad x, z_0 \in \mathbb{R}^D, \; \alpha > 0, \; \beta \in \mathbb{R}.\]

It contracts and expands the space around $z_0$. It has $D + 2$ parameters, taken by compile in the order $z_0, \alpha, \beta$. It was introduced by Rezende and Mohamed, "Variational inference with normalizing flows", ICML 2015.

RadialFlow(z0, α, β) takes the parameters, α > 0. RadialFlow([rng,] dim) draws them using rng, by default the task's generator: z0 and β from a standard normal, α uniformly from $[0, 1)$. RadialFlow() is the placeholder an AdditiveCouplingLayer takes: its dimension is set by the model, and its parameters by compile, which draws them the same way.

Throws

An AssertionError unless α > 0.

Examples

julia> FlowMessagePassingRules.forward(RadialFlow(0.0, 1.0, 2.0), 1.0) ≈ 1.0 + 2.0 / 2.0
true

forward(model::CompiledFlowModel, input::AbstractVector{<:Real})

The output of the compiled flow model at input, passing it through every layer in order. It has methods for the layers and the coupling flows too, and broadcasts over a vector of inputs with the model fixed, forward.(model, inputs).

backward(model::CompiledFlowModel, output::AbstractVector{<:Real})

The inverse of FlowMessagePassingRules.forward: the input of the compiled flow model that gives output, passing it back through every layer in reverse order. It broadcasts over a vector of outputs with the model fixed.

jacobian(model::CompiledFlowModel, input::AbstractVector{<:Real})

The Jacobian matrix of the compiled flow model at input, the product of its layers' Jacobians. It has methods for the layers and the coupling flows too, and broadcasts over a vector of inputs with the model fixed.

inv_jacobian(model::CompiledFlowModel, output::AbstractVector{<:Real})

The Jacobian matrix of the inverse flow, FlowMessagePassingRules.backward, at output: the inverse of jacobian(model, backward(model, output)). It broadcasts over a vector of outputs with the model fixed.

forward_jacobian(model::CompiledFlowModel, input::AbstractVector{<:Real})

The output and the Jacobian matrix of the compiled flow model at input in one pass, the tuple (forward(model, input), jacobian(model, input)).

backward_inv_jacobian(model::CompiledFlowModel, output::AbstractVector{<:Real})

The input and the inverse flow's Jacobian matrix at output in one pass, the tuple (backward(model, output), inv_jacobian(model, output)).

PermutationMatrix(ind::Vector{<:Integer})
PermutationMatrix([rng,] dim::Integer; switch_first::Bool = true)

A permutation matrix, with ones at (k, ind[k]) for k = 1:length(ind) and zeros elsewhere, so that P * v is v[ind]. Its inverse is its adjoint, and products with vectors, square matrices and other permutations reorder entries without multiplying.

Arguments

• ind: the permutation, a vector of the indices 1:n in some order.
• rng: the generator the random permutation is drawn from. Default: the task's generator.
• dim: the size of the random permutation, dim × dim.

Keywords

• switch_first: when true, a draw that keeps the first index in place swaps the first two entries, so the first element always moves. Default true.

Throws

A DimensionMismatch for a product with a matrix that is not square.

Examples

julia> PermutationMatrix([2, 1, 3]) * [10, 20, 30]
3-element Vector{Int64}:
 20
 10
 30