Normalizing flows: change of variables

A GAN produces samples but cannot tell us how likely a given image is. A VAE can, but only through a lower bound. Normalizing flows take a third route (Rezende and Mohamed, 2015; Papamakarios et al., 2021; Kobyzev, Prince and Brubaker, 2021): start from a simple density, push it through an invertible map, and compute the density of the result exactly with the change-of-variables formula. This note builds that formula, stacks it into a flow, compares flows with the other model families, and ends with the trick that makes flows computable: a triangular Jacobian. How to build layers with that property is the subject of the next note, Flow architectures.

Notation. The map $f$ always goes from the latent space to the data space, $f: \mathcal Z \to \mathcal X$, so $x = f(z)$ generates data and $z = f^{-1}(x)$ encodes it. $p_Z$ is the base (latent) density and $p_\theta(x)$ the model density of the data. For a map $g$, $J_g(u) = \partial g(u)/\partial u$ is its Jacobian matrix at the point $u$.

The model: invertible maps out of a Gaussian#

Take a continuous random variable $z$ with a simple density, a spherical Gaussian:

$$z \sim p_Z(z) = \mathcal N(z;\, 0,\, I).$$

“Spherical” means the covariance is a multiple of the identity (here exactly $I$): the same spread in every direction and spherical contours. The data are produced by a chain of $K$ maps,

$$x = f_\theta(z) = f_K \circ \dots \circ f_2 \circ f_1 (z),$$

and each $f_i$ is invertible (bijective): every point of the output space is the image of exactly one input point. For the density formula below, each $f_i$ must also be differentiable with a differentiable inverse (a diffeomorphism). Then the whole chain is invertible too, with $f^{-1} = f_1^{-1} \circ f_2^{-1} \circ \dots \circ f_K^{-1}$. A diffeomorphism needs $\dim z = \dim x$. The chain rule applied to $f^{-1}\circ f = \mathrm{id}$ and $f\circ f^{-1} = \mathrm{id}$ gives $J_{f^{-1}}J_f = I$ and $J_fJ_{f^{-1}} = I$, and a product of an $m\times n$ and an $n\times m$ matrix has rank at most $\min(m, n)$, so both identities can hold only when $m = n$. Unlike in a VAE, the latent vector is therefore as large as the data.

Figure

Every point of Z has exactly one partner in X

Why not just read off the base density?#

We want $p_\theta(x)$. The first guess is to map $x$ back and read off the base density, $p_\theta(x) \overset{?}{=} p_Z(f^{-1}(x))$. This is wrong, because the answer depends on how $f$ behaves: on how much it stretches or squeezes space around each point.

A one-line counterexample: in one dimension let $z \sim \mathcal N(0, 1)$ and $x = f(z) = 2z$. The naive guess gives $p_Z(x/2)$, which integrates to $2$ over the real line, so it is not even a density. The true distribution is $x \sim \mathcal N(0, 4)$ with density $p_Z(x/2)/2$. At $x = 1$ the naive value is $0.352$ and the correct one is $0.176$: $f$ spreads the same probability over twice the length, so the density must halve.

Change of variables#

Let $f: \mathcal Z \to \mathcal X$ be invertible, differentiable and with a differentiable inverse, and let $p_Z(z)$ be a density defined over $z \in \mathcal Z$.

One dimension. Probability mass is conserved. A small interval of width $|dz|$ around $z$ holds mass $p_Z(z)\,|dz|$, and $f$ sends it to an interval of width $|dx|$ around $x = f(z)$ that must hold the same mass:

$$p_\theta(x)\,|dx| = p_Z(z)\,|dz|,$$

so

$$p_\theta(x) = p_Z(z)\left|\frac{dz}{dx}\right| = \frac{p_Z(f^{-1}(x))}{\left|f'(f^{-1}(x))\right|}.$$

Any dimension. Widths become volumes, and we need the factor by which $f$ changes a small volume. Close to a point $z$ the map is approximately linear, $f(z + \delta) \approx f(z) + J_f(z)\,\delta$, with the Jacobian matrix

$$J_f(z) = \begin{bmatrix} \dfrac{\partial f^{(1)}}{\partial z_1} & \cdots & \dfrac{\partial f^{(1)}}{\partial z_d} \\ \vdots & \ddots & \vdots \\ \dfrac{\partial f^{(d)}}{\partial z_1} & \cdots & \dfrac{\partial f^{(d)}}{\partial z_d} \end{bmatrix},$$

where $f^{(1)},\dots,f^{(d)}$ are the output coordinates of $f$: row $k$ of $J_f$ holds the partial derivatives of output $k$, column $j$ those with respect to input $z_j$. A tiny cube with side $\varepsilon$ at $z$, spanned by the vectors $\varepsilon e_1,\dots,\varepsilon e_d$, is therefore sent to a tiny parallelepiped spanned by $\varepsilon J_f(z) e_1, \dots, \varepsilon J_f(z) e_d$, that is, by the columns of $J_f(z)$ scaled by $\varepsilon$. The volume of the parallelepiped spanned by the columns of a matrix $A$ is $\lvert\det A\rvert$, so $f$ multiplies small volumes near $z$ by $\lvert\det J_f(z)\rvert$. In two dimensions a small square becomes a small parallelogram whose area is $\lvert\det J_f(z)\rvert$ times the area of the square.

Figure

The Jacobian is local geometry

Mass is conserved as in one dimension: a cube of volume $\varepsilon^d$ holds mass $p_Z(z)\,\varepsilon^d$, and its image of volume $\lvert\det J_f(z)\rvert\,\varepsilon^d$ must hold the same mass. Hence, with $z = f^{-1}(x)$,

$$\begin{aligned} p_\theta(x) &= p_Z\big(f^{-1}(x)\big)\,\left|\det \frac{\partial f^{-1}(x)}{\partial x}\right| \\ &= p_Z(z)\,\left|\det \frac{\partial z}{\partial x}\right| \\ &= p_Z(z)\,\left|\det \frac{\partial f(z)}{\partial z}\right|^{-1}. \end{aligned}$$

The two forms say the same thing. The Jacobian of the inverse is evaluated at the data point, while the Jacobian of $f$ is evaluated at the latent point $z = f^{-1}(x)$. They are linked by the inverse function theorem, $J_{f^{-1}}(x) = \big[J_f(z)\big]^{-1}$, and $\det(A^{-1}) = 1/\det A$. The factor $|\det J_f|$ measures the expansion or contraction of space around the point (Papamakarios et al., 2021, §2.1). Where $f$ expands space, the same probability is spread thinner and the density drops; where $f$ contracts space, the density rises:

• $|\det J_f| = 2$: volumes are doubled and the density is halved;
• $|\det J_f| = 0.5$: volumes are halved and the density is doubled;
• $|\det J_f| = 1$: volumes are preserved (the shape may still change), and so is the density;
• $\det J_f = 0$: the volume is crushed to zero. The linear map $J_f(z)$ is not invertible, and then $f$ has no differentiable inverse at $f(z)$, since a differentiable inverse would give $J_{f^{-1}}J_f = I$. The density factor $1/|\det J_f|$ is infinite there, and the formula above does not apply. Flow layers are therefore built with $\det J_f \neq 0$ everywhere.

Figure

What |det J| does to a small square

A worked one-dimensional example#

Figure

Change of variables in one dimension

Take $p_Z = \mathcal N(0, 1)$ and the map in the figure,

$$\begin{gathered} f(z) = z + 0.8\tanh(2z), \\ f'(z) = 1 + 1.6\,\operatorname{sech}^2(2z) > 0 . \end{gathered}$$

The derivative is positive everywhere, so $f$ is strictly increasing and therefore invertible (the inverse has no closed form but is easy to find numerically). Then $p_\theta(x) = p_Z(z)/f'(z)$ with $z = f^{-1}(x)$:

$z$	$x = f(z)$	$f'(z)$	$p_Z(z)$	$p_\theta(x) = p_Z(z)/f'(z)$
$0$	$0$	$2.600$	$0.3989$	$0.1534$
$0.5$	$1.109$	$1.672$	$0.3521$	$0.2106$
$1$	$1.771$	$1.113$	$0.2420$	$0.2174$
$2$	$2.799$	$1.002$	$0.0540$	$0.0539$

Near the origin $f$ stretches space by a factor of $2.6$, so the density there falls well below the Gaussian peak; in the tails $f' \approx 1$ and the density is almost unchanged. The single bump becomes two bumps with maxima at $x \approx \pm 1.50$ (that is $z \approx \pm 0.77$), $p_\theta \approx 0.234$. Even one smooth invertible map can turn a unimodal base into a bimodal model. In log form, at $z = 0.5$: $\log p_\theta(x) = \log p_Z(z) - \log f'(z) = -1.044 - 0.514 = -1.558$.

The slab $z \in [0.1, 0.5]$ holds probability $\Phi(0.5) - \Phi(0.1) = 0.152$. Its image is $x \in [0.258, 1.109]$, $2.13$ times wider, so the average density over it drops from $0.379$ to $0.178$, by the same factor: same probability mass, different width.

Example: stretching space by a factor of 2#

The simplest multivariate case is a linear map $x = Jz$, whose Jacobian $\partial x/\partial z = J$ is the same everywhere. With

$$J = \begin{bmatrix} 2 & 0 \\ 0 & 2 \end{bmatrix}$$

every direction is stretched 2 times, and $\det J = 2\cdot 2 - 0\cdot 0 = 4$. A unit square of $\mathcal Z$ becomes a $2\times 2$ square of $\mathcal X$, so

$$p_\theta(x) = p_Z(z)\left|\frac{1}{\det\left(\frac{\partial x}{\partial z}\right)}\right| = \frac{p_Z(z)}{4}.$$

Check with $z \sim \mathcal N(0, I)$ in two dimensions: $x = 2z \sim \mathcal N(0, 4I)$, and indeed $p_\theta(0) = \frac{1}{2\pi\cdot 4} = 0.0398 = \frac14\cdot\frac{1}{2\pi}$.

Orientation does not matter. For

$$J = \begin{bmatrix} 2 & 0 \\ 0 & -2 \end{bmatrix}, \qquad \det J = -4,$$

the square is also mirrored, but its area is still $4$. The sign of the determinant records orientation; the change of volume, and hence of density, uses only $|\det J| = 4$. That is why the formula has an absolute value.

Figure

A Jacobian that stretches each direction by a factor of 2

Stacking maps: a normalizing flow#

Now apply the formula to the whole chain $z \sim p_Z(z)$, $x = f_\theta(z) = f_K\circ\dots\circ f_2\circ f_1(z)$. Name the intermediate points

$$z_0 = z, \quad z_i = f_i(z_{i-1}) \;\; (i = 1,\dots,K),$$

so that $z_K = x$. The inverse runs the chain backwards, $f^{-1} = f_1^{-1}\circ\dots\circ f_K^{-1}$, and by the chain rule its Jacobian is a product of the layer Jacobians, each taken at its own intermediate point: $J_{f^{-1}}(x) = J_{f_1^{-1}}(z_1)\,J_{f_2^{-1}}(z_2)\cdots J_{f_K^{-1}}(z_K)$. Since $\det(AB) = \det A\,\det B$,

$$\begin{aligned} p_\theta(x) &= p_Z(z_0)\,\left|\det\frac{\partial f^{-1}(x)}{\partial x}\right| \\ &= p_Z(z_0)\prod_{i=1}^{K}\left|\det\frac{\partial f_i^{-1}(z_i)}{\partial z_i}\right| \\ &= p_Z(z_0)\prod_{i=1}^{K}\left|\det\frac{\partial f_i(z_{i-1})}{\partial z_{i-1}}\right|^{-1}. \end{aligned}$$

To evaluate it we push $x$ backwards through the layers, $z_K = x \to z_{K-1} \to \dots \to z_0$, and collect one determinant per layer. This construction is a normalizing flow: the density flows through a sequence of invertible maps, and read backwards the chain normalizes the data, turning it (approximately, once the model fits) into a sample of the standard normal base. The name was popularized in machine learning by Rezende and Mohamed (2015), building on earlier work on density estimation by transforming data towards a Gaussian (Tabak and Vanden-Eijnden, 2010). Rezende and Mohamed’s flows were planar and radial maps, whose determinants are cheap for a different reason than the triangular layers below (Flow architectures, part G).

Sampling needs only the forward direction: draw $z \sim \mathcal N(0, I)$ and compute $x = f_\theta(z)$; evaluating a density needs the inverse. Which of the two is cheap depends on the architecture: masked autoregressive flows (MAF; Papamakarios, Pavlakou and Murray, 2017) evaluate densities fast and sample slowly, inverse autoregressive flows (IAF; Kingma et al., 2016) the reverse, and coupling layers make both cheap (see Flow architectures). Sampling turns a simple $z$ into a sample of a complicated distribution; the density formula is what lets us fit $f_\theta$ so that this distribution matches the data.

Maximum likelihood#

Take the logarithm and the product becomes a sum, with $J_{f_i^{-1}}(z_i) = \partial f_i^{-1}(z_i)/\partial z_i$:

$$\begin{aligned} \log p_\theta(x) &= \log p_Z(z_0) \\ &\quad + \sum_{i=1}^{K}\log\big|\det J_{f_i^{-1}}(z_i)\big| . \end{aligned}$$

This is an exact log-likelihood evaluation, with no bound and no sampling. Training is plain maximum likelihood, which minimizes the forward KL divergence from the data distribution to the model (why cross-entropy is the loss we minimize). As in the fitting note, the data density is unknown, so we minimize the negative log-likelihood over the dataset $x^{(1)},\dots,x^{(N)}$, now with a neural model:

$$\begin{aligned} \mathcal L(\theta) &= -\frac1N\sum_{n=1}^{N}\Bigg[\log p_Z\big(z_0^{(n)}\big) \\ &\qquad + \sum_{i=1}^{K}\log\left|\det\frac{\partial f_i^{-1}(z_i^{(n)})}{\partial z_i^{(n)}}\right|\Bigg], \end{aligned}$$

where $z_i^{(n)}$ are the intermediate points of $x^{(n)}$: $z_K^{(n)} = x^{(n)}$ and $z_{i-1}^{(n)} = f_i^{-1}(z_i^{(n)})$, so that $z_0^{(n)} = f_\theta^{-1}(x^{(n)})$.

The two terms pull against each other. The first wants every encoded point near the centre of the Gaussian. The second penalizes the inverse for squeezing space, so the model cannot score well by collapsing all the data to $z = 0$.

Images: dequantization, rescaling and bits per dimension#

Optional depth: the preprocessing that turns a flow’s density on images into a reported number of bits per dimension, with a worked offset check. The rest of the note does not depend on it.

The formula gives a density, so it is exact for continuous data. Images are discrete, for example 8-bit pixel values $x \in \{0,\dots,255\}^d$, and a continuous density could put arbitrarily high values on those points: a flow fitted to the raw values can push its likelihood up without bound by squeezing its mass onto them. Flows are therefore trained on dequantized data $\tilde x = x + u$ with uniform noise $u \sim \mathcal U[0, 1)^d$. The probability the model gives to the discrete value $x$ is $P(x) = \int_{[0,1)^d} p_\theta(x + u)\,du$, and Jensen’s inequality gives $\mathbb E_u\big[\log p_\theta(x + u)\big] \le \log P(x)$: the continuous log-likelihood of the dequantized data is a lower bound on the discrete log-likelihood (Theis, van den Oord and Bethge, 2016).

Bits per dimension. Image likelihoods are reported as $-\log_2 P(x)/d$, the number of bits per coordinate that a lossless code built from the model needs (bits per dimension in the autoregressive note). A model that spreads its probability evenly over the 256 values scores exactly 8. For a flow the reported number is $-\mathbb E_u\big[\log_2 p_\theta(x + u)\big]/d$, which by the inequality above is an upper bound on the discrete value.

Rescaling changes the density. Networks are rarely fed pixel values in $[0, 256)$. A common choice is $x' = \tilde x/127.5 - 1$, which puts the data roughly in $[-1, 1]$. The flow then models the density of $x'$, and the change of variables relates it to the density of $\tilde x$: the map $\tilde x\mapsto x'$ has the constant Jacobian $\tfrac{1}{127.5}I$, so

$$\begin{gathered} p_{\tilde x}(\tilde x) = p_{x'}(x')\cdot 127.5^{-d}, \\[4pt] -\frac{\log_2 p_{\tilde x}(\tilde x)}{d} = -\frac{\log_2 p_{x'}(x')}{d} + \log_2 127.5 . \end{gathered}$$

The bits per dimension in pixel units are those in model units plus $\log_2 127.5 \approx 6.99$. In general, if one quantization level has width $a$ after preprocessing, the offset is $-\log_2 a$ bits per dimension (Kingma and Dhariwal, 2018, §2, eq. 2): $\log_2 255\approx 7.99$ for $x/255$, as produced by common image loaders, and exactly 8 for $x/256$. For example (toy numbers), a flow that reports $-3.50$ bits/dim on $x'$ has $-3.50 + 6.99 = 3.49$ bits/dim in pixel units. A negative value in model units is normal, because a density on a narrow range can exceed 1; only the pixel-unit number can be compared with published results such as the likelihood table of the autoregressive note.

The noise must live in the same units. Dequantization noise has to span exactly one level of whatever the network sees. With $x' = x/127.5 - 1$ that means adding $u/127.5$ with $u\sim\mathcal U[0,1)^d$, which is the same as adding $u$ before rescaling. Noise of width 1 in $x'$ units would smear each value over about 127 levels, and no noise at all brings back the unbounded likelihood. NICE, for instance, dequantizes CIFAR-10 with uniform noise of width $\tfrac{1}{128}$ on data rescaled to $[-1, 1]$, where one level is $\tfrac{2}{256} = \tfrac{1}{128}$ wide (Dinh, Krueger and Bengio, 2015, §5.1).

The logit transform. Pixel values are bounded, while a Gaussian base lives on all of $\mathbb R^d$. To reduce boundary effects, RealNVP models the logit of the rescaled pixels (Dinh, Sohl-Dickstein and Bengio, 2017, §4.1):

$$\begin{gathered} y = \operatorname{logit}(p) = \log\frac{p}{1 - p}, \\[4pt] p = \alpha + (1 - 2\alpha)\,\frac{\tilde x}{256}, \qquad \alpha = 0.05 . \end{gathered}$$

The shift $\alpha$ keeps $p$ inside $[0.05, 0.95)$, so $y$ stays within about $\pm2.94$ instead of running to $\pm\infty$ at black and white pixels. This is the form of the reference implementation; the paper prints $1 - \alpha$ in place of $1 - 2\alpha$. The derivative of the transform is $dy/d\tilde x = (1 - 2\alpha)/\big(256\,p(1-p)\big)$ per coordinate, so

$$\begin{aligned} -\log_2 p_{\tilde x}(\tilde x) &= -\log_2 p_y(y) \\ &\quad + \sum_{i=1}^{d}\log_2\frac{256\,p_i(1 - p_i)}{1 - 2\alpha} . \end{aligned}$$

Unlike a linear rescaling, the correction depends on the pixel: about $6.15$ bits for a mid-grey value ($\tilde x = 128.5$) and $3.80$ bits for the darkest and brightest values ($\tilde x = 0.5$ or $255.5$). RealNVP includes this term in its bits per dimension. The Glow paper states that it follows RealNVP’s preprocessing (§5); its reference code instead rescales to $[-\tfrac12, \tfrac12)$, adds noise of width $\tfrac{1}{256}$ and adds the constant $-d\log 256$ to the log-likelihood, which is the linear case above with $a = \tfrac{1}{256}$. Flows built for images are the subject of Flow architectures, part F.

Recap: GANs, VAEs and flows#

Figure

Recap: three ways to connect data x and latent z

• GAN (Goodfellow et al., 2014; see the GAN note). There is no log-likelihood evaluation: the generator defines its distribution only through samples. There is also no latent-variable inference: nothing maps $x$ back to $z$, and nothing guarantees that the generator covers the full support of the data, so a real $x$ may have no latent code at all.
• VAE (Kingma and Welling, 2014; see the VAE note). It optimizes a lower bound on the log-likelihood (the ELBO) and infers latents through an approximate posterior $q_\phi(z\mid x)$.
• Flow-based generative models, for example NICE (Dinh, Krueger and Bengio, 2015) and RealNVP (Dinh, Sohl-Dickstein and Bengio, 2017). They minimize the exact negative log-likelihood, and posterior inference is exact: $z = f^{-1}(x)$, a deterministic answer.
• Diffusion models complete the list; they come later in the course and, like VAEs, derive their training objective from a bound on the likelihood.

	Log-likelihood	Latent inference
GAN	none	none
VAE	lower bound (ELBO)	approximate, $q_\phi(z\mid x)$
Flow	exact	exact, $z = f^{-1}(x)$

The trick: a triangular Jacobian#

Everything above hinges on $|\det J|$: if it can be computed, the density is exact; if it cannot be computed cheaply, the likelihood is impractical to evaluate and to train on at the scale of images. For a general network it is expensive. Computing the determinant of a general $d\times d$ matrix costs $O(d^3)$ with Gaussian elimination or an LU decomposition (Golub and Van Loan, 2013); naive cofactor expansion is even worse, $O(d!)$. Take $d = 1000$ (a $32\times 32$ grayscale image already has $d = 1024$): that is on the order of $10^9$ operations (about $\tfrac23 d^3 \approx 6.7\times 10^8$ for LU) for every example, every layer and every training step. For a general network even forming $J$ takes $d$ backward passes, one per row.

The trick that makes flows practical is to ensure the Jacobian is triangular. For a (lower-)triangular matrix, expanding the determinant along the first row and repeating on the remaining block leaves only the diagonal:

$$\begin{gathered} A = \begin{bmatrix} a_{11} & 0 & \cdots & 0 \\ a_{21} & a_{22} & \ddots & \vdots \\ \vdots & \vdots & \ddots & 0 \\ a_{n1} & a_{n2} & \cdots & a_{nn} \end{bmatrix} \\ \Longrightarrow\quad \det A = a_{11}\, a_{22} \cdots a_{nn}. \end{gathered}$$

So if $J$ is triangular,

$$\begin{gathered} \det J = \prod_{i=1}^{d} J_{ii}, \\ \log|\det J| = \sum_{i=1}^{d} \log|J_{ii}|, \end{gathered}$$

which costs $O(d)$: about $10^3$ operations instead of about $10^9$, and only the diagonal entries are needed.

A triangular $J$ is also invertible exactly when every diagonal entry is non-zero, which is easy to guarantee by construction. For the map itself the condition does more. If $J_f$ is lower-triangular everywhere, output $i$ depends only on inputs $1,\dots,i$. If each diagonal entry is continuous and non-zero everywhere, it keeps one sign, so each output is strictly monotone in its own input. The map is then injective and can be inverted one coordinate at a time: $z_1$ from $x_1$, then $z_2$ from $x_2$ and the known $z_1$, and so on. It is a bijection of $\mathbb R^d$ when each of these one-dimensional maps is onto $\mathbb R$, as for the affine maps $z_i \mapsto e^{s}z_i + t$ of coupling layers.

Figure

Why a triangular Jacobian is the trick

Notebook#

The normalizing-flows notebook of the HSE-LAMBDA Deep Generative Models repository (pinned commit; cell numbers count every cell, markdown included, from 1) builds on Andrej Karpathy’s pytorch-normalizing-flows library (Karpathy, 2019), included as nflib. The notebook has no saved outputs, so it documents procedures rather than results.

• Cell 2 writes the change of variables from the generative side, $\log p(z_K) = \log p(z_0) - \sum_i \log\lvert\det \partial f_i/\partial z_{i-1}\rvert$.
• Cells 4–5 define three 2-D toy datasets: two moons (the one used), a mixture of four Gaussians with standard deviation $0.5$, and the SIGGRAPH point set, which needs a siggraph.pkl file that is not included in the repository.
• Cell 7 builds the model: a logistic prior, a uniform pushed through an inverse sigmoid, with $\log p_Z(z) = \sum_i\big[-\log(1 + e^{z_i}) - \log(1 + e^{-z_i})\big]$ (the Gaussian prior is commented out). NICE uses the same prior and prefers it for its better-behaved gradients (Dinh, Krueger and Bengio, 2015, §3.4). The change-of-variables formula holds for any base density, so only $\log p_Z$ changes. The flow is three blocks of ActNorm, an invertible $1\times1$ “convolution” and a neural-spline coupling layer (8 spline bins, tail bound 3, hidden width 16), described in Flow architectures.
• Cells 9–11 train with Adam (learning rate $10^{-4}$, weight decay $10^{-5}$) on the loss above: zs, prior_logprob, log_det = model(x_batch), logprob = prior_logprob + log_det, loss = -torch.sum(logprob), for 1000 iterations, each on a fresh batch of 128 samples. The loss is summed over the batch rather than averaged as in $\mathcal L(\theta)$, which only rescales the gradient. nflib names the directions the other way round from this note: its forward runs data → latent (our $f^{-1}$) and accumulates our $\sum_i \log\lvert\det \partial f_i^{-1}/\partial z_i\rvert$, while model.sample draws from the prior and runs inverse (our $f$).
• Cells 13–17 map data to latents and prior samples to data, and show how a grid is warped layer by layer. Cell 13 draws its prior points with model.prior.sample([128, 2]); the prior already has two coordinates, so this returns a $128\times2\times2$ tensor and the scatter plot shows 256 prior points instead of 128, which is harmless (sample((128,)) gives the intended shape). Cell 17 runs one more training step and draws the grid warp once: its 500-step loop breaks at $k = 0$.
• Cell 19 colours a grid by prior_logprob.exp(), that is $p_Z(f^{-1}(x))$ without the log-determinant. This is exactly the naive formula from the start of this note; the model density is (prior_logprob + log_det).exp().

Next#

Flows need layers that are expressive and still have a cheap Jacobian determinant by construction. Flow architectures builds them: coupling layers (NICE, RealNVP), masked autoregressive layers (MAF, IAF) and Glow’s invertible $1\times1$ convolutions, then spline couplings, ActNorm and multi-scale flows for images, and finally the planar and radial flows of Rezende and Mohamed.

References#

• D. J. Rezende, S. Mohamed. Variational inference with normalizing flows. ICML, 2015.
• G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, B. Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. JMLR, 22(57), 2021.
• I. Kobyzev, S. J. D. Prince, M. A. Brubaker. Normalizing flows: An introduction and review of current methods. IEEE TPAMI, 43(11), 2021.
• E. G. Tabak, E. Vanden-Eijnden. Density estimation by dual ascent of the log-likelihood. Communications in Mathematical Sciences, 8(1), 2010.
• L. Dinh, D. Krueger, Y. Bengio. NICE: Non-linear independent components estimation. ICLR Workshop, 2015. §3.4 (prior), §5.1 (dequantization).
• L. Dinh, J. Sohl-Dickstein, S. Bengio. Density estimation using Real NVP. ICLR, 2017. §4.1 (logit preprocessing). Reference code: tensorflow/models, research/real_nvp.
• D. P. Kingma, P. Dhariwal. Glow: Generative flow with invertible 1×1 convolutions. NeurIPS, 2018. §2, eq. 2 (discretization constant), §5. Reference code: openai/glow.
• I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, Y. Bengio. Generative adversarial nets. NeurIPS, 2014.
• D. P. Kingma, M. Welling. Auto-encoding variational Bayes. ICLR, 2014.
• 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.
• L. Theis, A. van den Oord, M. Bethge. A note on the evaluation of generative models. ICLR, 2016.
• G. H. Golub, C. F. Van Loan. Matrix Computations, 4th ed. Johns Hopkins University Press, 2013. §3.2: LU factorization.
• A. Karpathy. pytorch-normalizing-flows. GitHub repository, 2019.
• HSE-LAMBDA, Deep Generative Models course repository: nf.ipynb.