Flow architectures

The change-of-variables note ended on a cost problem. Every layer of a flow contributes a term $\pm\log\lvert\det J\rvert$ to the log-likelihood (with a plus sign when $J$ is the Jacobian of the data-to-latent map). For a general $d\times d$ Jacobian that determinant costs $O(d^3)$. For a triangular Jacobian it is just the product of the diagonal entries, which costs $O(d)$. The question is therefore how to give $J$ a structure whose determinant is cheap by construction: triangular, a linear map stored in factored form, or the identity plus a low-rank term.

The answer comes in three designs: coupling layers (part A) and autoregressive layers (part B), whose Jacobians are triangular, and invertible $1\times1$ convolutions (part C), whose Jacobian is block-diagonal with dense blocks stored as products of triangular factors. Each one keeps the layer invertible and its determinant cheap, while arbitrary neural networks still do the heavy lifting inside it. Three more parts build on them. Part D replaces the affine map inside a coupling layer with a monotone spline, part E adds ActNorm, the normalization layer of Glow, and part F assembles all of them into flows for images. Part G shows an older route to a cheap determinant that does not use triangular matrices at all: planar and radial flows. A first reading can take parts A, B, C and F; parts D and G are extensions.

Notation. One flow layer maps $z\in\mathbb R^d$ to $z'\in\mathbb R^d$. Its Jacobian is $J = \partial z'/\partial z$, with entries $J_{ij} = \partial z'_i/\partial z_j$: row $i$ is an output coordinate and column $j$ an input coordinate. $\odot$ is the element-wise product, and $\exp$ of a vector acts element-wise. It does not matter here whether a layer is used in the data-to-latent or the latent-to-data direction. The inverse layer has the reciprocal determinant, so a determinant that is cheap in one direction is cheap in the other. A subscript indexes the coordinates of one vector ($z_i$, $z_{<i}$), not the layers of a flow as in the previous note. An image tensor has height $h$, width $w$ and $C$ channels.

A. Coupling layers (NICE, RealNVP)#

Split the input into two parts,

$$z = (z_a,\; z_b), \qquad z_a\in\mathbb R^{k},\quad z_b\in\mathbb R^{d-k},$$

and define the transform

$$\begin{aligned} z'_a &= z_a, \\ z'_b &= z_b \odot \exp\bigl(s(z_a)\bigr) + t(z_a), \end{aligned}$$

where $s, t:\mathbb R^{k}\to\mathbb R^{d-k}$ are arbitrary neural networks, called the scale and the translation network. Two facts make this work:

• $z'_a$ does not depend on $z_b$ at all;
• $z'_b$ depends on $z_b$ only affinely, and even element-wise: the $i$-th coordinate $z'_{b,i}$ involves only $z_{b,i}$, multiplied by $\exp(s_i(z_a))$ and shifted by $t_i(z_a)$.

Figure

An affine coupling layer and its inverse

The Jacobian. Those two facts fix the block structure of $J$:

$$J = \frac{\partial z'}{\partial z} = \begin{bmatrix} I_k & 0 \\[4pt] \dfrac{\partial z'_b}{\partial z_a} & \operatorname{diag}\bigl(\exp(s(z_a))\bigr) \end{bmatrix}.$$

• Top left, $\partial z'_a/\partial z_a = I$, because $z'_a = z_a$.
• Top right, $\partial z'_a/\partial z_b = 0$, because $z'_a$ ignores $z_b$.
• Bottom right, $\partial z'_b/\partial z_b = \operatorname{diag}(\exp(s(z_a)))$, because the dependence on $z_b$ is element-wise with slope $\exp(s_i(z_a))$.
• Bottom left, $\partial z'_b/\partial z_a$ is a dense matrix involving the derivatives of $s$ and $t$. It is complicated, but it does not matter.

So $J$ is lower-triangular, and its determinant is the product of its diagonal:

$$\begin{aligned} \det J &= 1^k\cdot\prod_{i=1}^{d-k} \exp\bigl(s_i(z_a)\bigr) \\ &= \exp\Bigl(\sum_{i=1}^{d-k} s_i(z_a)\Bigr), \\[4pt] \log\lvert\det J\rvert &= \sum_{i=1}^{d-k} s_i(z_a). \end{aligned}$$

Here $i$ runs over the $d-k$ coordinates of $z_b$, and $s_i$ is the $i$-th output of the network $s$.

The determinant is always positive, so no absolute value is needed. The layer is invertible for any choice of the networks $s$ and $t$, not because $\det J\ne0$ (that alone only guarantees local invertibility) but because the inverse can be written down explicitly, as shown below. The dense block never enters the determinant: $s$ and $t$ can be as deep as we like without making $\det J$ any more expensive.

Figure

Why the coupling Jacobian is lower-triangular

For example, with $d = 6$ split $3+3$, suppose that at some point $s(z_a) = (0.2,\,-0.5,\,0.1)$. Then $\log\lvert\det J\rvert = 0.2 - 0.5 + 0.1 = -0.2$ and $\det J = e^{-0.2}\approx 0.819$. Near this point the layer shrinks small volumes by about 18%, whatever the dense block contains.

The inverse reverses only the element-wise step. Since $z'_a = z_a$, the networks can be evaluated again on the same input:

$$\begin{aligned} z_a &= z'_a, \\ z_b &= \bigl(z'_b - t(z'_a)\bigr)\odot\exp\bigl(-s(z'_a)\bigr), \\[4pt] \log\lvert\det J_{\text{inv}}\rvert &= -\sum_{i=1}^{d-k} s_i(z'_a). \end{aligned}$$

Neither $s$ nor $t$ ever has to be inverted. Both directions cost one pass through the two networks.

B. Autoregressive layers (MAF, IAF)#

Ensure that each output dimension depends only on the previous ones:

$$z'_i = f_i(z_1, \dots, z_i), \qquad i = 1,\dots,d.$$

Now $J$ is automatically lower-triangular:

• there is no dependence on dimensions $> i$, so $J_{ij} = \partial z'_i/\partial z_j = 0$ for every $j > i$;
• the diagonal entries are $J_{ii} = \partial f_i/\partial z_i$.

$$J = \begin{bmatrix} \frac{\partial f_1}{\partial z_1} & 0 & \cdots & 0 \\ \ast & \frac{\partial f_2}{\partial z_2} & \cdots & 0 \\ \vdots & & \ddots & \vdots \\ \ast & \ast & \cdots & \frac{\partial f_d}{\partial z_d} \end{bmatrix},$$

so that

$$\det J = \prod_{i=1}^{d} \frac{\partial f_i}{\partial z_i}.$$

Again, the determinant is a product of diagonal entries. A coupling layer is the special case with just two blocks: $z_a$ conditions $z_b$.

The concrete layer. In MAF, the masked autoregressive flow (Papamakarios, Pavlakou and Murray, 2017), each $f_i$ is affine in its own coordinate, with a scale and a shift computed from the coordinates strictly before it:

$$\begin{aligned} z'_i &= z_i\,\exp\bigl(s_i(z_{<i})\bigr) + t_i(z_{<i}), \\[4pt] J_{ii} &= \exp\bigl(s_i(z_{<i})\bigr), \\[4pt] \log\lvert\det J\rvert &= \sum_{i=1}^{d} s_i(z_{<i}). \end{aligned}$$

Here $z_{<i} = (z_1,\dots,z_{i-1})$. Because $s_i$ and $t_i$ do not see $z_i$, the diagonal entry is simply $\exp(s_i) > 0$, and the layer can be inverted one coordinate at a time. The same layer is often written as $f(x_i) = (x_i - \mu(x_{<i}))/\exp(\alpha(x_{<i}))$, which is the form above with $s = -\alpha$ and $t = -\mu\,e^{-\alpha}$.

Masking. We could give every coordinate its own network that reads $z_{<i}$, at the price of $d$ networks and $d$ sequential passes. MAF instead uses one network and masks its weights, so that

$$\frac{\partial z'_i}{\partial z_j} = 0 \qquad \text{for all } j > i .$$

This network is MADE, the masked autoencoder for distribution estimation (Germain et al., 2015). The masks come from degrees, written $m(\cdot)$ here so as not to clash with the dimension $d$:

• input $z_j$ has degree $m(j) = j$;
• each hidden unit $u$ gets a degree $m(u)\in\{1,\dots,d-1\}$ (Germain et al., 2015, §4). The unit is allowed to depend on $z_1,\dots,z_{m(u)}$ and nothing later. A degree of $d$ would be useless: such a unit could only feed outputs that come after $z_d$, and there are none;
• the output unit that produces $(s_i, t_i)$ has degree $m(i) = i$.

In a hidden layer, hidden unit $u$ may read unit $j$ of the layer below only if the sender’s degree is at most the receiver’s:

$$\begin{gathered} M_{uj} = \begin{cases} 1 & \text{if } m(j) \le m(u), \\ 0 & \text{otherwise,} \end{cases} \\[6pt] h = Wx + b \;\;\longrightarrow\;\; h = (W\odot M)\,x + b . \end{gathered}$$

The mask has the same shape as the layer’s weight matrix $W$, and $W\odot M$ simply zeroes the forbidden weights. The output layer uses the strict inequality: its mask is $M^{\text{out}}_{iu} = 1$ if $m(u) < i$ and $0$ otherwise, so output $i$ may read hidden unit $u$ only if $m(u) < i$. This is what makes $(s_i, t_i)$ depend on $z_{<i}$ and not on $z_i$ itself.

A worked example. Take $d = 3$ inputs, one hidden layer of four units with degrees $(1, 2, 1, 2),$ and outputs of degrees $(1, 2, 3)$. In both masks the rows are the receiving units and the columns the sending units. The hidden mask (rule $m(j)\le m(u)$) and the output mask (rule $m(u) < i$) are

$$\begin{gathered} M^{(1)} = \begin{bmatrix} 1&0&0\\ 1&1&0\\ 1&0&0\\ 1&1&0 \end{bmatrix}, \\[6pt] M^{(2)} = \begin{bmatrix} 0&0&0&0\\ 1&0&1&0\\ 1&1&1&1 \end{bmatrix}. \end{gathered}$$

Entry $(i, j)$ of their product counts the paths from input $j$ to output $i$:

$$M^{(2)}M^{(1)} = \begin{bmatrix} 0&0&0\\ 2&0&0\\ 4&2&0 \end{bmatrix}.$$

It is strictly lower-triangular. A nonlinearity such as ReLU cannot create a path that the masks removed, so $s_i, t_i$ depend only on $z_{<i}$. The element-wise step then adds the diagonal $\exp(s_i).$ Output 1 receives no connections at all, so $s_1, t_1$ are learned constants (biases), as they should be: $z_1$ has nothing before it. Input $z_3$ feeds nothing, because no later coordinate needs it.

Figure

Masked weights make the network autoregressive

MAF versus IAF. The masked network gives all the $s_i, t_i$ in one pass when the conditioning variables are known. In MAF they are known in the data-to-latent direction, so the density of a data point is cheap. Sampling has to invert the layer coordinate by coordinate: $z_i = (z'_i - t_i(z_{<i}))\,e^{-s_i(z_{<i})}$ needs $z_{<i}$ first, which takes $d$ sequential passes. IAF, the inverse autoregressive flow (Kingma et al., 2016), swaps the two directions: its $s_i, t_i$ read the latent-side coordinates before $i$. Sampling is one pass, but evaluating the density of an arbitrary given point is sequential. IAF can still score its own samples in one pass, because their latents are known, and that is what variational inference needs. Coupling layers are cheap in both directions because $s$ and $t$ read only $z_a$, which passes through unchanged and is known in both directions. The price is that only half of the coordinates change per layer.

C. Invertible convolutions (Glow)#

Glow layers use $1\times1$ invertible convolutions:

$$z' = W z .$$

For an image-shaped tensor of height $h$, width $w$ and $C$ channels, a $1\times1$ convolution applies the same $C\times C$ matrix $W$ to the vector of $C$ channel values at every pixel. It mixes channels and does nothing across space. Mixing channels is what the coupling layers need between them: it is a learned generalization of the fixed channel permutation (for example, reversing the channel order) that RealNVP uses between couplings (Kingma and Dhariwal, 2018, §3.2). The Jacobian of the whole layer is block-diagonal, with $h\cdot w$ copies of $W$ on the diagonal, so

$$\log\lvert\det J\rvert = h\cdot w\cdot\log\lvert\det W\rvert .$$

$W$ must be invertible, and a general $\det W$ costs $O(C^3)$. So $W$ is learned through its LU decomposition. Not every invertible matrix has an LU factorization without pivoting: the swap $\left[\begin{smallmatrix}0&1\\1&0\end{smallmatrix}\right]$ has none, because $(LU)_{11} = U_{11}$ would have to be $0$ and then $(LU)_{21} = L_{21}U_{11} = 0 \ne 1$. Glow therefore adds a fixed permutation matrix $P$:

$$\begin{gathered} W = PLU, \\[6pt] L = \begin{bmatrix} 1 & & 0 \\ \vdots & \ddots & \\ L_{C1} & \cdots & 1 \end{bmatrix}, \\[6pt] U = \begin{bmatrix} U_{11} & \cdots & U_{1C} \\ & \ddots & \vdots \\ 0 & & U_{CC} \end{bmatrix}. \end{gathered}$$

The entries $L_{ij}$ and $U_{ij}$ are real numbers. $P$ stays fixed; the free entries of $L$ (below its unit diagonal) and of $U$ (on and above its diagonal) are the learnable parameters. Since $\det P = \pm1$, $\det L = 1$ and $\det U = \prod_{i=1}^{C} U_{ii}$,

$$\log\lvert\det W\rvert = \sum_{i=1}^{C} \log\lvert U_{ii}\rvert .$$

The layer is invertible exactly when every $U_{ii}\neq 0$, and its inverse is $W^{-1} = U^{-1}L^{-1}P^{\top}$: a permutation and two triangular solves. Unlike $\exp(s_i)$, a free parameter $U_{ii}$ is not automatically non-zero.

Figure

Glow's invertible 1×1 convolution

For example, take $C = 3$ and $P = I$ with

$$\begin{gathered} L = \begin{bmatrix} 1&0&0\\ 0.5&1&0\\ -1&2&1 \end{bmatrix},\quad U = \begin{bmatrix} 2&1&-1\\ 0&-0.5&3\\ 0&0&1.5 \end{bmatrix} \\[6pt] \Longrightarrow\quad W = LU = \begin{bmatrix} 2&1&-1\\ 1&0&2.5\\ -2&-2&8.5 \end{bmatrix}. \end{gathered}$$

Expanding $\det W$ directly gives $-1.5$, which equals $2\cdot(-0.5)\cdot1.5$. Hence $\log\lvert\det W\rvert = \log 1.5 \approx 0.405$. The negative pivot $U_{22} = -0.5$ shows why the absolute value is needed: $\log(-0.5)$ is undefined.

D. Spline couplings: neural spline flows#

Optional depth: a more flexible element-wise map and the derivation of its inverse. The rest of the note does not depend on it.

Inside a coupling or autoregressive layer, each transformed coordinate goes through a one-dimensional map whose parameters come from a network: $z_{b,i}\mapsto z_{b,i}\,e^{s_i}+t_i$ in part A. An affine map can shift and stretch a conditional distribution, but it cannot change its shape: a unimodal conditional stays unimodal after one layer. Nothing in the triangular argument needs the map to be affine. Any strictly increasing, differentiable function $g_\theta$ with a computable inverse will do. The diagonal entry of $J$ becomes $g_\theta'(z_{b,i})$, and $\theta$ is produced by the same kind of network as before: $\theta = \mathrm{NN}(z_a)$ in a coupling layer, $\theta_i = \mathrm{NN}(z_{<i})$ in an autoregressive one.

The rational-quadratic spline. Neural spline flows (Durkan et al., 2019, §3.1) use a monotone spline built from rational-quadratic pieces, the ratio of two quadratic polynomials, with the construction of Gregory and Delbourgo (1982). In this part $x$ is the scalar input of the map, not a data point. The spline maps $[-B, B]$ onto itself and is the identity outside it. Inside, it is set by $K$ bins with $K + 1$ knots $(x_k, y_k)$, increasing from $(x_0, y_0) = (-B, -B)$ to $(x_K, y_K) = (B, B)$, and by positive derivatives $\delta_k$ at the knots. The two edge derivatives are fixed, $\delta_0 = \delta_K = 1$, so that the slope matches the identity tails. In bin $k$ let

$$s_k = \frac{y_{k+1} - y_k}{x_{k+1} - x_k}, \qquad \xi = \frac{x - x_k}{x_{k+1} - x_k}\in[0, 1]$$

($s_k$ is the slope of the chord across the bin, not the scale network of part A). The spline in bin $k$ is

$$g(x) = y_k + \frac{(y_{k+1} - y_k)\bigl[s_k\,\xi^2 + \delta_k\,\xi(1-\xi)\bigr]}{s_k + (\delta_{k+1} + \delta_k - 2s_k)\,\xi(1-\xi)} .$$

At $\xi = 0$ it gives $y_k$, and at $\xi = 1$ the fraction is $s_k/s_k = 1$, so it gives $y_{k+1}$: every piece joins two consecutive knots. The denominator can be rewritten as $s_k\bigl(\xi^2 + (1-\xi)^2\bigr) + (\delta_k + \delta_{k+1})\,\xi(1-\xi)$, which is positive. Its derivative is

$$g'(x) = \frac{s_k^2\bigl[\delta_{k+1}\xi^2 + 2s_k\,\xi(1-\xi) + \delta_k(1-\xi)^2\bigr]}{\bigl[s_k + (\delta_{k+1} + \delta_k - 2s_k)\,\xi(1-\xi)\bigr]^2} ,$$

with every term of the numerator positive, so $g$ is strictly increasing. At $\xi = 0$ the derivative equals $\delta_k$ and at $\xi = 1$ it equals $\delta_{k+1}$: neighbouring pieces meet with the same slope, and the spline meets the tails with slope 1. The layer’s log-determinant is $\sum_i\log g'(z_{b,i})$ over the transformed coordinates; coordinates in the tails contribute $\log 1 = 0$.

Parameters. For every transformed coordinate the network outputs $3K - 1$ unconstrained numbers (Durkan et al., 2019, §3.1): $K$ for the bin widths and $K$ for the bin heights, each passed through a softmax and multiplied by $2B$, so that they are positive and fill the interval, and $K - 1$ for the inner derivatives, passed through a softplus. Cumulative sums of the widths and heights, starting at $-B$, give the knots.

The inverse is a quadratic. Since $g$ increases, the bin of an output $y$ is found by comparing $y$ with the $y_k$. Setting $g(x) = y$ and multiplying out the fraction gives $a\xi^2 + b\xi + c = 0$ with

$$\begin{aligned} a &= (y_{k+1} - y_k)(s_k - \delta_k) + (y - y_k)(\delta_{k+1} + \delta_k - 2s_k), \\ b &= (y_{k+1} - y_k)\,\delta_k - (y - y_k)(\delta_{k+1} + \delta_k - 2s_k), \\ c &= -s_k\,(y - y_k). \end{aligned}$$

Exactly one root lies in $[0, 1]$ (Durkan et al., 2019, App. A.3):

$$\xi = \frac{2c}{-b - \sqrt{b^2 - 4ac}},$$

the usual formula $(-b + \sqrt{b^2 - 4ac})/(2a)$ rewritten so that it does not lose precision when $4ac$ is small. Then $x = x_k + \xi\,(x_{k+1} - x_k)$. Evaluation and inversion both cost a bin search and a few arithmetic operations, so a spline coupling layer, like an affine one, takes one network pass in each direction.

Figure

A monotone rational-quadratic spline and its derivative

A worked example (toy numbers, the spline of the figure). $K = 4$ bins on $[-3, 3]$ with knots $(-3,-3)$, $(-1.2,-2.4)$, $(0, 1.2)$, $(1.5, 2.2)$, $(3, 3)$, inner derivatives $(\delta_1, \delta_2, \delta_3) = (0.4, 2.5, 0.5)$ and $\delta_0 = \delta_4 = 1$. The input $x = -0.6$ lies in bin $k = 1$, with $x_1 = -1.2$, width $1.2$, $\xi = 0.5$, height $y_2 - y_1 = 3.6$ and $s_1 = 3$. Then $\xi(1-\xi) = 0.25$ and

$$\begin{aligned} g(-0.6) &= -2.4 + \frac{3.6\,(3\cdot0.25 + 0.4\cdot0.25)}{3 + (2.5 + 0.4 - 6)\cdot0.25} \\ &= -2.4 + \frac{3.06}{2.225} \approx -1.025, \\[4pt] g'(-0.6) &= \frac{9\,(2.5\cdot0.25 + 2\cdot3\cdot0.25 + 0.4\cdot0.25)}{2.225^2} \approx 4.04 . \end{aligned}$$

This coordinate adds $\log 4.04 \approx 1.40$ to $\log\lvert\det J\rvert$. Inverting $y = -1.0247$: $y - y_1 = 1.3753$, so $a \approx 5.097$, $b \approx 5.703$, $c \approx -4.126$, $b^2 - 4ac \approx 116.6$, and $\xi = 2c/(-b - \sqrt{b^2 - 4ac}) \approx -8.252/-16.503 = 0.500$, which gives back $x = -1.2 + 0.5\cdot1.2 = -0.6$.

Coupling and autoregressive versions. A spline replaces the affine map in either design. Durkan et al. call the results RQ-NSF (C) and RQ-NSF (AR) (§3.2): the coupling version is Glow with affine couplings replaced by spline couplings, and the autoregressive version is MAF or IAF with affine maps replaced by splines. In their coupling layer the identity half also passes through element-wise splines whose parameters are trained directly rather than predicted. With enough bins a monotone spline approximates any differentiable monotone function on $[-B, B]$; in practice they fix $B = 3$ and $K = 8$ after finding little difference for $B$ between 1 and 5 (§5).

E. ActNorm: per-channel scale and shift with data-dependent initialization#

Deep flows train better with a normalization layer, and in a flow that layer must be invertible with a known determinant. RealNVP uses batch normalization (Dinh, Sohl-Dickstein and Bengio, 2017, §3.7): with the batch statistics $\tilde\mu, \tilde\sigma^2$, the map $x\mapsto(x - \tilde\mu)/\sqrt{\tilde\sigma^2 + \epsilon}$ is a per-dimension rescaling with log-determinant $-\tfrac12\sum_i\log(\tilde\sigma_i^2 + \epsilon)$. That formula treats $\tilde\mu$ and $\tilde\sigma^2$ as constants supplied from outside the example, as RealNVP does with a running average over recent minibatches (App. E). Ordinary train-time batch normalization, differentiated as a function of the current batch, is a different map: each output depends on every example of the batch, and standardizing removes the batch’s own mean and scale, so the outputs alone do not determine the inputs. Batch normalization adds noise through its batch statistics, the noise grows as the batch per device shrinks, and Glow trains on large images with one image per device (§3.1). Its ActNorm layer keeps the affine form and drops the batch statistics (Kingma and Dhariwal, 2018, §3.1, Table 1). With a log-scale $\gamma_c$ and a shift $\tau_c$ per channel $c$,

$$\begin{gathered} z'_{ijc} = e^{\gamma_c}\,z_{ijc} + \tau_c, \\[4pt] \log\lvert\det J\rvert = h\cdot w\cdot\sum_{c=1}^{C}\gamma_c . \end{gathered}$$

The Jacobian is diagonal, with the entry $e^{\gamma_c}$ repeated at each of the $h\cdot w$ positions of channel $c$. Glow writes the scale itself, $s_c$, and the log-determinant $h\cdot w\sum_c\log\lvert s_c\rvert$; the log-scale keeps the scale positive automatically. The inverse is $z_{ijc} = (z'_{ijc} - \tau_c)\,e^{-\gamma_c}$. On flat vectors the same layer has one $(\gamma_i, \tau_i)$ per coordinate and $h = w = 1$; it generalizes the diagonal scaling layer that NICE places at the top of its flow (Dinh, Krueger and Bengio, 2015, §3.3).

Data-dependent initialization. The parameters are set from the first minibatch so that each channel of the layer’s output has zero mean and unit variance on it:

$$\gamma_c = -\log\operatorname{std}_c(z), \qquad \tau_c = -\operatorname{mean}_c\bigl(e^{\gamma_c} z\bigr),$$

which makes $z'_{ijc} = (z_{ijc} - \operatorname{mean}_c)/\operatorname{std}_c$ on that batch. From then on $\gamma$ and $\tau$ are ordinary trainable parameters that do not depend on the data, unlike batch normalization. Every layer of a fresh flow therefore starts with well-scaled activations, whatever the depth.

For example (toy numbers), an ActNorm with log-scales $\gamma = (0.5, -0.2, 0.1)$ on a $32\times32\times3$ tensor contributes $\log\lvert\det J\rvert = 1024\cdot0.4 = 409.6$ nats.

One step of Glow is ActNorm, then a $1\times1$ convolution, then an affine coupling layer (Kingma and Dhariwal, 2018, §3). The last convolution of each coupling network is initialized at zero, so that every coupling layer starts as the identity (§3.3). Both initializations aim at the same thing: a deep stack of invertible layers that starts close to a well-scaled identity map.

F. Flows on images: squeeze, masks and multi-scale#

An RGB image of $64\times64$ pixels has $d = 64\cdot64\cdot3 = 12288$ coordinates. The layers above work on any vector, but an image flow should use convolutional networks for $s$ and $t$, which means that the split into $z_a$ and $z_b$ must respect the pixel grid.

Masks on the pixel grid#

RealNVP writes the coupling with a binary mask $b$ of the same shape as the image tensor (Dinh, Sohl-Dickstein and Bengio, 2017, §3.4):

$$\begin{gathered} z' = b\odot z + (1 - b)\odot\Bigl(z\odot\exp\bigl(s(b\odot z)\bigr) + t(b\odot z)\Bigr), \\[4pt] \log\lvert\det J\rvert = \sum (1 - b)\odot s(b\odot z), \end{gathered}$$

where the sum runs over all entries of the tensor. Entries with $b = 1$ pass through and condition the rest, exactly as $z_a$ does in part A, and $s$, $t$ are convolutional networks that map an $h\times w\times C$ tensor to one of the same shape. Two masks are used:

• checkerboard: $b_{ij} = 1$ where $i + j$ is odd, the same in every channel, so the four direct neighbours of every transformed pixel are among the conditioning pixels;
• channel-wise: $b = 1$ on the first half of the channels and $0$ on the second half.

Consecutive layers alternate $b$ and $1 - b$, so that every entry is transformed.

Squeeze: trading space for channels#

The squeeze operation cuts each channel into $2\times2$ squares and stacks the four values of each square as four channels, turning an $h\times w\times C$ tensor into an $\tfrac h2\times\tfrac w2\times 4C$ tensor (§3.6). With channel $4c + 2p + q$ of the output holding the pixels at offset $(p, q)\in\{0,1\}^2$ of channel $c$, a $4\times4\times1$ input becomes

$$\begin{bmatrix} 1&2&3&4\\ 5&6&7&8\\ 9&10&11&12\\ 13&14&15&16 \end{bmatrix} \;\longrightarrow\; \begin{bmatrix} 1&3\\ 9&11 \end{bmatrix},\; \begin{bmatrix} 2&4\\ 10&12 \end{bmatrix},\; \begin{bmatrix} 5&7\\ 13&15 \end{bmatrix},\; \begin{bmatrix} 6&8\\ 14&16 \end{bmatrix}.$$

Squeezing only reorders the coordinates, so its Jacobian is a permutation matrix and its log-determinant is $0$. It also explains why the two masks complement each other. With 0-based indices the checkerboard of this $4\times4$ input is

$$b = \begin{bmatrix} 0&1&0&1\\ 1&0&1&0\\ 0&1&0&1\\ 1&0&1&0 \end{bmatrix},$$

and after squeezing its ones fill channels 1 and 2 (offsets $(0,1)$ and $(1,0)$) and its zeros channels 0 and 3. The channel-wise mask splits the squeezed channels into $\{0, 1\}$ and $\{2, 3\}$, a different partition, so channels 0 and 3, which the checkerboard put on the same side, end up in different halves, as do channels 1 and 2: couplings after the squeeze condition each group on pixels that the checkerboard couplings kept together.

Factoring out half the variables at each scale#

RealNVP’s multi-scale architecture (§3.6) repeats one block per scale: three coupling layers with alternating checkerboard masks, a squeeze, and three coupling layers with alternating channel-wise masks. After each block, half of the variables are factored out: they leave the flow and become part of the latent $z$, and only the other half continues to the next scale. In the data-to-latent direction, with $\Phi^{(i)}$ the block of scale $i$ and $h^{(0)} = x$,

$$\begin{aligned} \bigl(z^{(i)},\, h^{(i)}\bigr) &= \Phi^{(i)}\bigl(h^{(i-1)}\bigr), \qquad i = 1,\dots,L-1, \\ z^{(L)} &= \Phi^{(L)}\bigl(h^{(L-1)}\bigr), \\ z &= \bigl(z^{(1)}, \dots, z^{(L)}\bigr). \end{aligned}$$

The last scale has four checkerboard couplings and no squeeze, and the recursion stops when its input is $4\times4\times C$ (§4.1). Every factored-out part is scored by the Gaussian base as soon as it leaves, so the loss is spread through the network, and the later scales are cheaper because they process fewer variables. For a $64\times64\times3$ image the recipe gives five scales:

Scale	Input	After squeeze	Factored out	Continues
1	$64\times64\times3$	$32\times32\times12$	6144	$32\times32\times6$
2	$32\times32\times6$	$16\times16\times24$	3072	$16\times16\times12$
3	$16\times16\times12$	$8\times8\times48$	1536	$8\times8\times24$
4	$8\times8\times24$	$4\times4\times96$	768	$4\times4\times48$
5	$4\times4\times48$	none	768

The factored-out sizes add up to $6144 + 3072 + 1536 + 768 + 768 = 12288 = d$: a multi-scale flow still keeps every dimension. Its 28 coupling layers process $6\cdot(1 + \tfrac12 + \tfrac14 + \tfrac18) + 4\cdot\tfrac1{16} = 11.5$ full-image equivalents of variables, about 41% of what 28 couplings on all $12288$ variables would process. Variables factored out early have passed through fewer layers, and RealNVP describes them as finer, more local features (§3.6, App. D). In RealNVP the scale network ends in $\tanh$ times a learned factor and the translation network has an affine output (§4.1).

The Glow step on image tensors#

Glow keeps the multi-scale skeleton with $L$ levels: each level is a squeeze followed by $K$ steps of ActNorm, $1\times1$ convolution and affine coupling, and every level but the last ends with a split that factors out half of the variables (Kingma and Dhariwal, 2018, §3, Fig. 2). Its couplings split only along channels, a simplification over RealNVP (§3.3), and the learned $1\times1$ convolutions replace the fixed channel reversal between couplings. All three layers act on $h\times w\times C$ tensors with per-channel parameters, and their log-determinants carry the factor $h\cdot w$ (part C and part E). A $1\times1$ convolution after one squeeze of an RGB image has $C = 12$ and only $12^2 = 144$ weights. The coupling networks have three convolutions ($3\times3$, $1\times1$, $3\times3$) with 512 hidden channels, and the CIFAR-10 comparisons use $K = 32$ and $L = 3$ (§5). Glow’s best results are 3.35 bits/dim on CIFAR-10 against 3.49 for RealNVP, and 3.81 against 3.98 on ImageNet $64\times64$ (§5, Table 2). For $256\times256$ CelebA-HQ faces it uses $K = 32$ and $L = 6$ on 5-bit images (§6).

Flattening the image instead#

The simplest image flow reshapes the image into a vector of $d = 12288$ numbers and applies the layers of parts A–D with fully connected networks. It works, and it needs no masks or squeezes, but it throws away the grid. An index split of the flattened channel-first tensor puts the whole red channel and the top half of the green channel into $z_a$; the networks have to learn from scratch which coordinates are neighbours; and the first layer of a coupling network with 1024 hidden units already has $6144\cdot1024\approx6.3$ million weights. A learned dense mixing matrix between couplings would have $d^2\approx1.5\cdot10^8$ entries, so a fixed permutation, for example a random one, is the practical choice. On CIFAR-10, NICE, whose couplings are fully connected networks on the flattened image, reaches 4.48 bits/dim against 3.49 for the convolutional multi-scale RealNVP (table). The two models also differ in depth and in additive versus affine couplings, so not all of the gap is due to the grid.

Likelihoods of image flows are reported in bits per dimension after undoing every preprocessing step, which adds the log-determinant of the preprocessing to the model’s own; Images: dequantization, rescaling and bits per dimension works this out for the common rescalings and for the logit transform of RealNVP.

Interpolating and encoding with the exact inverse#

A flow has an exact encoder: any image $x$, including one from outside the training set, has exactly one latent $z = f^{-1}(x)$, and $f(z)$ gives back $x$ up to floating-point error. A VAE instead encodes into the parameters of $q_\phi(z\mid x)$, and its reconstruction is only as faithful as its decoder: the common pixel-wise Gaussian or squared-error decoder returns the mean of $p_\theta(x\mid z)$, which tends to be blurred, while a more expressive decoder likelihood can give sharp samples (VAE; variants in VAE variants). Three uses follow.

Interpolation. Encode two images, move along a path between $z_1$ and $z_2$, and decode. The straight line $z_\lambda = (1-\lambda)z_1 + \lambda z_2$ leaves the region where the base puts its mass. For $z \sim \mathcal N(0, I_d)$ the length $\lVert z\rVert$ concentrates near $\sqrt d$ with a spread of about $0.7$ (the Gaussian shell). If $z_1$ and $z_2$ behave like independent draws, their midpoint has length about $\sqrt{d/2}$. For $d = 12288$ that is $78.4$ instead of $110.9$, about 46 spreads below the typical length. Weights whose squares sum to one keep the length typical. RealNVP’s paths between latents use such $\cos$ and $\sin$ weights (§4.2, eq. 19), and the midpoint $\cos\tfrac\pi4\,z_1 + \sin\tfrac\pi4\,z_2$ has length about $\sqrt d$. Glow interpolates linearly between the latents of real faces and still reports realistic intermediate images (§6).

Temperature. Glow samples from a reduced-temperature model: it multiplies the standard deviation of the base by $T = 0.7$, trading diversity for quality (§6). For a volume-preserving flow, $z\sim\mathcal N(0, T^2 I)$ gives the density $\propto\exp\bigl(-\lVert f^{-1}(x)\rVert^2/(2T^2)\bigr) \propto p_\theta(x)^{1/T^2}$, a sharpened version of the model (the paper writes the exponent as $T^2$).

Encoding an outside image. Preprocess it exactly as the training data, compute $z = f^{-1}(x)$, and compare it with the latents of training images, for example by Euclidean distance. Two cautions. If the latents behave like independent draws from the base, all pairwise distances are close to $\sqrt{2d}$, so nearest neighbours differ only slightly from typical points. And the exact likelihood of the outside image does not say whether it looks like the training data: a Glow trained on CelebA assigns higher likelihood to SVHN house numbers, a dataset it has never seen, than to CelebA faces (Nalisnick et al., 2019, §3, Fig. 2).

G. Planar and radial flows: cheap determinants without triangles#

Optional depth: the older rank-one route to a cheap determinant. The rest of the note does not depend on it.

The first flows of Rezende and Mohamed (2015), §4.1, take another route to a cheap determinant. Their Jacobian is the identity plus a rank-one matrix, and the matrix determinant lemma gives its determinant in $O(d)$ (Papamakarios et al., 2021, §3.3.2): for $a, c\in\mathbb R^d$,

$$\det\bigl(I + a\,c^\top\bigr) = 1 + c^\top a .$$

Every vector orthogonal to $c$ is left unchanged by $I + ac^\top$, so the eigenvalue $1$ has multiplicity at least $d - 1$; the trace $d + c^\top a$ then fixes the last eigenvalue at $1 + c^\top a$, and the determinant is the product of the eigenvalues.

Planar flows. With $u, w\in\mathbb R^d$ and a scalar bias $b$ (not the mask of part F),

$$\begin{gathered} f(z) = z + u\,\tanh(w^\top z + b), \\[4pt] J = I + \tanh'(w^\top z + b)\,u\,w^\top, \\[4pt] \det J = 1 + \tanh'(w^\top z + b)\,w^\top u . \end{gathered}$$

Rezende and Mohamed allow any smooth element-wise nonlinearity in place of $\tanh$; the invertibility condition below is for $\tanh$.

The map moves $z$ only along $u$, by an amount that depends on the signed distance from the hyperplane $w^\top z + b = 0$, so it expands or contracts the density across that hyperplane.

When is it invertible? Split $z$ into a part orthogonal to $w$ and a part $\eta\,w/\lVert w\rVert^2$ along it, so that $w^\top z = \eta$. Taking $w^\top$ of $y = f(z)$ gives the scalar equation

$$w^\top y = \eta + w^\top u\,\tanh(\eta + b) =: g(\eta),$$

and once $\eta$ is known, $z = y - u\tanh(\eta + b)$. So $f$ is invertible exactly when $g$ is. Since $0 < \tanh' \le 1$, with $\tanh' = 1$ only at $\eta = -b$, we get $g'(\eta) = 1 + w^\top u\,\tanh'(\eta + b)\ge 1 + \min(w^\top u, 0)$. If $w^\top u\ge -1$, $g$ is strictly increasing and onto $\mathbb R$, and $f$ is a bijection (Rezende and Mohamed, 2015, App. A.1). If $w^\top u < -1$, then $g'(-b) = 1 + w^\top u < 0$ while $g' \to 1$ far away, so $g$ is not monotone and $f$ folds space onto itself. A diffeomorphism also needs $\det J > 0$ everywhere, which holds for $w^\top u > -1$; Rezende and Mohamed enforce this by replacing $u$ with $u + [m(w^\top u) - w^\top u]\,w/\lVert w\rVert^2$, where $m(x) = -1 + \log(1 + e^x) > -1$. Even for an invertible planar flow, the inverse has no closed form: each layer needs a one-dimensional root search for $\eta$.

For example (toy numbers), take $d = 2$, $w = (1, 0)$, $u = (-0.5, 1)$, $b = 0$, so $w^\top u = -0.5$. At $z = (0.5, 0)$: $\tanh'(0.5) \approx 0.786$, $\det J \approx 1 - 0.5\cdot0.786 \approx 0.607$, and $y = f(z) \approx (0.269, 0.462)$. Inverting needs the root of $\eta - 0.5\tanh\eta = 0.269$, which is $\eta = 0.5$, and then $z = y - u\tanh 0.5 = (0.5, 0)$. With $u = (-2, 1)$ instead, $w^\top u = -2$ and $g(\eta) = \eta - 2\tanh\eta$ vanishes at $\eta = 0$ and $\eta\approx\pm1.915$: every point $y$ with $w^\top y = 0$ has three preimages.

Radial flows. With a centre $z_0\in\mathbb R^d$, $\alpha > 0$, $\beta\in\mathbb R$ and $r = \lVert z - z_0\rVert$,

$$f(z) = z + \frac{\beta}{\alpha + r}\,(z - z_0).$$

The map keeps the direction from $z_0$ and changes only the distance, $r\mapsto\rho(r) = r\,\bigl(1 + \tfrac{\beta}{\alpha + r}\bigr)$. The Jacobian of such a radial map is $(\rho(r)/r)\,I$ plus a rank-one term along $z - z_0$: it stretches the $d - 1$ directions orthogonal to $z - z_0$ by $\rho(r)/r = 1 + \beta/(\alpha + r)$ and the radial direction by $\rho'(r) = 1 + \beta/(\alpha + r) - \beta r/(\alpha + r)^2 = 1 + \alpha\beta/(\alpha + r)^2$. The determinant is the product of these stretches:

$$\det J = \Bigl(1 + \frac{\beta}{\alpha + r}\Bigr)^{d-1}\Bigl(1 + \frac{\alpha\beta}{(\alpha + r)^2}\Bigr).$$

Both factors are positive for every $r\ge 0$ when $\beta > -\alpha$, and then $\rho$ increases from $0$ to $\infty$, so $f$ is a bijection; Rezende and Mohamed state the condition $\beta\ge-\alpha$ and enforce it with $\beta = -\alpha + \log(1 + e^{\hat\beta})$ (App. A.2). Unlike the planar flow, the radial flow can be inverted in closed form: the distance $R = \lVert y - z_0\rVert$ of an output satisfies $R(\alpha + r) = r(\alpha + \beta + r)$, a quadratic in $r$,

$$r^2 + (\alpha + \beta - R)\,r - \alpha R = 0,$$

whose roots have the product $-\alpha R\le 0$, so exactly one is non-negative; then $z = z_0 + (y - z_0)/\bigl(1 + \tfrac{\beta}{\alpha + r}\bigr)$. For example (toy numbers), with $d = 2$, $\alpha = 1$, $\beta = -0.8$, a point at distance $r = 1$ from the centre has $\det J = (1 - 0.4)\cdot(1 - 0.2) = 0.48$ and is mapped to distance $R = 0.6$; the quadratic $r^2 - 0.4r - 0.6 = 0$ has the non-negative root $r = 1$.

Where they are used. Each layer reshapes the density along one direction or around one point, so many layers are needed, and the planar inverse is iterative. Neither matters for variational inference, the setting they were designed for: there one samples $\varepsilon$ from a simple $q_0$, pushes it forward through $K$ layers to $z_K$, and accumulates $\log q_K(z_K) = \log q_0(\varepsilon) - \sum_k\log\lvert\det J_k\rvert$ with forward Jacobians only, as in an IAF. Planar and radial flows have mostly been used this way, as flexible posteriors for VAEs (Papamakarios et al., 2021, §3.3.2), while the layers of parts A–F were designed for exact densities of data.

How the layers warp space and stay invertible#

The Jacobian is how the flow warps space. Every layer above stretches some regions and squeezes others, and $\lvert\det J\rvert$ at a point is the local factor by which it changes small volumes. Pushing a regular grid through a trained flow, one layer at a time, and drawing the bent grid lines makes this visible. The figure below does it for a single coupling layer with simple hand-picked $s$ and $t$.

Figure

The Jacobian measures how a layer warps space

Invertible by construction. Exact maximum likelihood needs the unique latent point $f^{-1}(x)$ of every data point and a Jacobian that never collapses volume; the general argument is in the change-of-variables section of the previous note, and the triangular-Jacobian section shows when a triangular map is invertible. For the layers of this note:

• Affine coupling and autoregressive layers are invertible for any $s$ and $t$: their diagonals $\exp(s_i)$ are positive and their inverses are explicit.
• Spline layers are invertible for any network output: every piece is strictly increasing and the inverse is the root of a quadratic.
• ActNorm is invertible because its scales $e^{\gamma_c}$ are positive; a layer that learns the scale $s_c$ itself, as in the Glow paper, needs $s_c \ne 0$.
• Squeeze and fixed channel permutations only reorder coordinates.
• The $1\times1$ convolution is invertible as long as every pivot $U_{ii}\ne0$, which a parameterization through $\log\lvert U_{ii}\rvert$ with fixed signs guarantees (part C). The Glow paper’s $\mathbf s$ has no such safeguard; there $U_{ii}\ne0$ holds only in practice.
• Planar and radial flows are diffeomorphisms only under a constraint, $w^\top u > -1$ or $\beta > -\alpha$, enforced by a reparameterization.

Limitations#

• No compression. The latent has the dimension of the data: $12288$ numbers for a $64\times64\times3$ image, however simple the image. Multi-scale factoring organizes the latent; it does not shrink it.
• Constrained layers. Every layer must be invertible with a tractable Jacobian, so a flow needs many simple layers, and in each Glow step two of the three layers only rescale or mix channels. Glow uses as many as 320 sub-transformations (Papamakarios et al., 2021, §3.4), and on CIFAR-10 its 3.35 bits/dim still trails the 2.92 of the autoregressive PixelCNN++.
• No holes in the density. A diffeomorphism of $\mathbb R^d$ applied to a Gaussian gives a density that is positive everywhere. Data that lie on separate clusters or near a lower-dimensional surface can only be approximated, with very large Jacobians at the gaps.
• Likelihood is not quality. An exact likelihood measures compression, not how samples look, and it can rank unfamiliar images above familiar ones (part F).

Summary#

Normalizing flows are invertible neural networks that move probability mass in controlled ways. Each layer reports, through its Jacobian determinant, how much it stretched or squeezed space. Autoregressive flows whose element-wise transformation is a sufficiently flexible monotone function (neural or spline transformers) are universal density approximators under mild conditions (Huang et al., 2018, §4, Theorem 1; Papamakarios et al., 2021, §3.1). Affine layers are universal as well, in a weaker sense. Compositions of affine couplings and invertible linear maps (permutations suffice) approximate any $C^2$ diffeomorphism in $L^p$ on compact sets, provided the networks $s$ and $t$ are universal function approximators, and pushing a base density through such compositions reaches any target distribution in the limit, in the sense of convergence in distribution (Teshima et al., 2020, Theorems 2 and 3). The theorem is the right answer to “can affine flows express everything?”, but it leaves the practical questions open. It gives no bound on the number of layers; its construction transforms one coordinate per coupling and covers flat vectors, not directly Glow’s convolutional networks; and its approximating flows can be arbitrarily ill-conditioned. Koehler, Mehta and Risteski (2021) show that shallow affine coupling networks are universal in Wasserstein distance once such ill-conditioning is allowed, and Draxler et al. (2024) prove a universality theorem, in a sense built on the KL divergence, for well-conditioned coupling flows such as RealNVP.

The layers differ in where the cheap determinant comes from:

Layer	Why $\det J$ is cheap	$\log\lvert\det J\rvert$	Cost of the two directions
Coupling (NICE, RealNVP)	$z'_a = z_a$; $z'_b$ element-wise in $z_b$	$\sum_{i=1}^{d-k} s_i(z_a)$ (NICE: $0$)	one pass each way
Autoregressive (MAF)	$z'_i$ depends on $z_1,\dots,z_i$ only (masks)	$\sum_{i=1}^{d} s_i(z_{<i})$	density: one pass; sampling: $d$ passes
Autoregressive (IAF)	same, but $s_i, t_i$ read the latent-side coordinates before $i$	$\sum_{i=1}^{d} s_i$ (of the latent-side coordinates before $i$)	sampling: one pass; density: $d$ passes
$1\times1$ conv (Glow)	$W = PLU$, $P$ fixed, $L, U$ triangular	$h\cdot w\cdot\sum_{i=1}^{C}\log\lvert U_{ii}\rvert$	one matrix product each way
Spline coupling (NSF)	as coupling, with a monotone spline $g$ in place of the affine map	$\sum_i \log g'(z_{b,i})$	one pass each way; the inverse solves a quadratic per coordinate
ActNorm	diagonal	$h\cdot w\cdot\sum_{c=1}^{C}\gamma_c$	one affine map each way
Squeeze, permutation	reordering	$0$	free
Planar	$I$ plus a rank-one matrix (determinant lemma)	$\log\lvert 1 + \tanh'(w^\top z + b)\,w^\top u\rvert$	forward $O(d)$; inverse by a scalar root search
Radial	$I$ scaled plus a rank-one matrix	$(d-1)\log\bigl(1 + \tfrac{\beta}{\alpha + r}\bigr) + \log\bigl(1 + \tfrac{\alpha\beta}{(\alpha + r)^2}\bigr)$	forward $O(d)$; inverse by a quadratic in $r$

• On images, couplings use checkerboard and channel-wise masks, squeezes trade space for channels, and half of the variables are factored out at each scale; Glow’s step is ActNorm, $1\times1$ convolution and affine coupling.
• A flow encodes any image exactly, which allows interpolation and the encoding of outside images; linear paths between latents leave the typical shell of the base.

Notebook#

The normalizing-flows notebook and its library nflib (HSE-LAMBDA Deep Generative Models, branch spring-2025-long) implement the layers of parts A–E on flat vectors. Cell numbers count every cell, markdown included, from 1. The notebook has no saved outputs, so it documents procedures rather than results. In nflib, the forward method of every layer runs data → latent and returns the output with its log-determinant.

• flows.py, coupling. AffineHalfFlow is part A: the forward pass computes x_b * exp(s) + t and log_det = sum(s), the inverse (z_b - t) * exp(-s) and log_det = sum(-s). Its parity flag swaps which half conditions the other, so stacked layers alternate, and scale=False gives NICE’s additive coupling. For odd $d$ the transformed half is the larger one.
• flows.py, autoregressive. SlowMAF gives every coordinate its own MLP (a learned constant, LeafParam, for the first) and loops over coordinates in both directions. MAF uses one ARMLP that returns $(s, t)$ for all coordinates in one pass: MAF.forward is parallel and MAF.inverse loops over dimensions. IAF simply swaps the two methods. Both MAF classes have a parity flag that reverses the coordinate order of the output, but with opposite meanings: MAF flips when parity=True (z.flip), SlowMAF when parity=False.
• made.py. MaskedLinear computes F.linear(input, mask * weight, bias), that is $h = (W\odot M)x + b$, and MADE.update_masks applies the degree rules of part B: self.m[l-1][:, None] <= self.m[l][None, :] for hidden layers and self.m[L-1][:, None] < self.m[-1][None, :] for the output layer. These are built as (in × out) arrays with 0-based degrees and transposed in MaskedLinear.set_mask.
• nets.py. MLP has three hidden layers of width nh with LeakyReLU(0.2). ARMLP is MADE with three hidden layers of width nh, a single mask and the natural ordering. PosEncMLP feeds $\sin$ and $\cos$ features at frequencies 0.5, 1, 2, 4 and 8 into an MLP; nothing in the notebook uses it.
• flows.py, scale layers. AffineConstantFlow is a per-dimension affine layer, z = x * exp(s) + t with log_det = sum(s), a generalization of NICE’s scaling layer; its s and t start as standard normal draws. ActNorm subclasses it and implements part E on flat vectors: on its first forward call it sets s = -log(x.std(0)) and t = -mean(x * exp(s)), then trains them as ordinary parameters. The flag data_dep_init_done that guards this is a plain attribute, not part of the state_dict, so a model rebuilt and loaded from a checkpoint re-initializes on its first forward call and overwrites the loaded s and t. Set the flag to True after loading.
• flows.py, Invertible1x1Conv. It follows the paper’s parameterization of part C and learns $\mathbf s$ directly, without a safeguard against $s_i\to0$. It starts from a random orthogonal matrix, factors it with torch.linalg.lu, keeps $P$ as a fixed buffer, learns $L$, $\mathbf s$ and $U$, applies the matrix to row vectors (z = x @ W, which leaves the determinant unchanged) and returns log_det = sum(log(abs(S))). The inverse uses torch.linalg.inv(W).
• spline_flows.py. unconstrained_RQS and RQS implement part D: the identity outside $[-B, B]$ with zero log-determinant, the edge derivatives set to 1, minimum bin widths, heights and derivatives of $10^{-3}$, a bin search that counts the knots below the input, and the inverse 2c / (-b - sqrt(discriminant)) with -logabsdet. NSF_CL is a spline coupling layer with two MLPs: f1 reads the first half and outputs $3K - 1$ parameters per coordinate of the second half, which it transforms; f2 then reads the new second half and transforms the first. Each layer thus updates both halves, whereas the coupling layer of Durkan et al. (2019), eq. 9–11, gives the conditioning half splines with directly trained parameters. Its reshapes assume an even $d$. NSF_AR is the autoregressive version, with a learned parameter vector for the first coordinate and one MLP per later coordinate; like SlowMAF it loops over coordinates in both directions.
• Two quirks of spline_flows.py. NSF_CL and NSF_AR already pass widths and heights through a softmax and multiply them by $2B$, and RQS applies a softmax again. The final widths are therefore proportional to $\exp(2B\,p_k)$ with $p_k$ a probability vector, so two bins can differ in width by a factor of at most $e^{2B}\approx403$ for $B = 3$. The derivatives also pass through a softplus twice, so an inner knot derivative is never below $\ln 2 + 10^{-3}\approx 0.694$. The spline stays monotone and invertible, but it is less flexible than the paper’s, which applies each transform once.
• Notebook cell 2 lists planar and radial flows, RealNVP, MAF (in the form $f(x_i) = (x_i - \mu(x_{<i}))/\exp(\alpha(x_{<i}))$ of part B), the $1\times1$ convolution, ActNorm (as $f(x) = Wx + b$ with diagonal $W$) and NSF. nflib implements neither planar nor radial flows. Cell 3 imports AffineConstantFlow, ActNorm, AffineHalfFlow, SlowMAF, MAF, IAF, Invertible1x1Conv, NSF_AR and NSF_CL, but cell 7 uses only ActNorm, Invertible1x1Conv and NSF_CL ($K = 8$, $B = 3$, hidden width 16): three blocks of them on 2-D data, with a logistic base distribution. Each block is the Glow step of part E with a spline coupling, an RQ-NSF (C) layer; with $C = 2$ and no spatial grid the “convolution” is a learned invertible $2\times2$ mixing matrix.
• Cell 15 pushes a $20\times20$ grid backward through the trained flow (model.inverse) and plots each layer’s point movements and grid warp, the experiment that the grid-warp figure imitates for a single coupling layer. Cell 17 sets up 500 more training steps with a redraw every 10 steps (point movements for up to nine layers and the grid warp of one layer), but its loop ends with break after the first redraw, so as written it takes a single optimizer step and draws once.

The homework notebook 03_hw_vapnev of the same repository asks for three models of CelebA and a comparison of them. Its only saved outputs are the device (mps), the image shape [3, 64, 64] and a grid of 25 training images; the model cells are left empty.

• Cells 3–6: data. CelebADataset from the repository’s seminars/utils/datasets/celeba.py, with Resize(64), CenterCrop(64), ToTensor() and x * 2 - 1, that is $x' = x/127.5 - 1\in[-1, 1]$ and $d = 12288$. The pipeline adds no dequantization noise. A flow trained on these values sees only 256 distinct values per coordinate and can push its likelihood up without bound by concentrating mass on them. Add uniform noise of width $1/127.5$ before training, and add $\log_2 127.5\approx 6.99$ bits/dim to the reported negative log-likelihood (how this offset arises).
• Cell 10, task 1: a normalizing flow on CelebA, evaluated by FID and negative log-likelihood. Any layers are allowed: convolutional flows (part F) or flows on the flattened image, which the task text expects to be simpler and weaker.
• Cell 12, task 2: a VAE on CelebA, evaluated by FID.
• Cell 14, task 3: VAPNEV (Agrawal and Dukkipati, 2016), a VAE whose likelihood $p(x\mid z)$ is computed exactly by a real NVP instead of a pixel-wise Gaussian or Bernoulli, and a conditional version of it; see VAE variants.
• Cell 16, task 4: latent analysis. Which model interpolates better, encoding an image from outside CelebA and finding its nearest latents, and the most interesting sample. The flow side is in Interpolating and encoding with the exact inverse.
• Cell 18 asks for a written comparison with samples from every model.

Next#

Flows buy an exact likelihood by restricting the architecture: every layer must be invertible, with a Jacobian determinant that is cheap to compute, and the latent must keep the dimension of the data. An energy-based model drops every one of these restrictions. Any network $E_\theta(x)$ that outputs a number defines a density $p_\theta(x) = e^{-E_\theta(x)}/Z_\theta$, as long as $Z_\theta$ is finite, but the normalizing constant $Z_\theta = \int e^{-E_\theta(x)}\,dx$ is unknown, so neither sampling nor the likelihood gradient can be computed directly. Both then need Markov chain Monte Carlo. The next note, sampling unnormalized densities, builds the samplers that need the density only up to its constant: MALA, Hamiltonian Monte Carlo and NUTS.

References#

• L. Dinh, D. Krueger, Y. Bengio. NICE: Non-linear independent components estimation. ICLR Workshop, 2015. §3.3 (scaling layer).
• L. Dinh, J. Sohl-Dickstein, S. Bengio. Density estimation using Real NVP. ICLR, 2017. §3.4 (masks), §3.6 (squeeze and multi-scale architecture), §3.7 (batch normalization), §4.1 (preprocessing and architecture), App. D. Reference code: tensorflow/models, research/real_nvp.
• M. Germain, K. Gregor, I. Murray, H. Larochelle. MADE: Masked autoencoder for distribution estimation. ICML, 2015.
• G. Papamakarios, T. Pavlakou, I. Murray. Masked autoregressive flow for density estimation. NeurIPS, 2017.
• D. P. Kingma, T. Salimans, R. Jozefowicz, X. Chen, I. Sutskever, M. Welling. Improved variational inference with inverse autoregressive flow. NeurIPS, 2016.
• D. P. Kingma, P. Dhariwal. Glow: Generative flow with invertible 1×1 convolutions. NeurIPS, 2018. §3.1 (ActNorm), §3.2 (1×1 convolution), §3.3 (coupling, zero initialization), §5 (Table 2), §6 (temperature, interpolation). Reference code: openai/glow.
• C. Durkan, A. Bekasov, I. Murray, G. Papamakarios. Neural spline flows. NeurIPS, 2019. §3.1–3.2, §5, App. A.
• J. A. Gregory, R. Delbourgo. Piecewise rational quadratic interpolation to monotonic data. IMA Journal of Numerical Analysis, 2(2):123–130, 1982.
• D. J. Rezende, S. Mohamed. Variational inference with normalizing flows. ICML, 2015. §4.1 (planar and radial flows), App. A (invertibility conditions).
• E. Nalisnick, A. Matsukawa, Y. W. Teh, D. Gorur, B. Lakshminarayanan. Do deep generative models know what they don’t know? ICLR, 2019. §3.
• C.-W. Huang, D. Krueger, A. Lacoste, A. Courville. Neural autoregressive flows. ICML, 2018. §4, Theorem 1 (universal approximation).
• T. Teshima, I. Ishikawa, K. Tojo, K. Oono, M. Ikeda, M. Sugiyama. Coupling-based invertible neural networks are universal diffeomorphism approximators. NeurIPS, 2020. Theorems 2 and 3 (universality of affine coupling flows).
• F. Koehler, V. Mehta, A. Risteski. Representational aspects of depth and conditioning in normalizing flows. ICML, 2021.
• F. Draxler, S. Wahl, C. Schnörr, U. Köthe. On the universality of volume-preserving and coupling-based normalizing flows. ICML, 2024. §5.
• G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, B. Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. JMLR, 22(57), 2021. §3.2, §3.3.2, §3.4.
• HSE-LAMBDA, Deep Generative Models course repository: nf.ipynb, 03_hw_vapnev.ipynb.