Normalizing Flows

ABC's four failures from the last post all pointed at the same fix: stop rejecting simulations, start learning a function. Specifically, learn a conditional density $q_{\varphi }(\theta \mid x)$ from simulated $(\theta ,x)$ pairs. After training, querying it at $x_{\mathrm{obs}}$ gives you the posterior in a forward pass — no rejection, no summary statistics, no per-observation rerun.

That's the plan. The question is what kind of object $q_{\varphi }$ can actually be. Most off-the-shelf density estimators are missing one piece or another. Normalizing flows are the one family that has everything you need at once. This post is about why they're the right tool and what they look like.

What we need from the density estimator

Four requirements, all load-bearing.

Tractable density. Given any $\theta$, we can evaluate $q_{\varphi }(\theta \mid x)$ as a number. Without this we can't compute the loss that trains the model — and we can't say "the model thinks ${\theta }_1$ is more likely than ${\theta }_2$" with anything more than samples.

Tractable sampling. Given $x_{\mathrm{obs}}$, we can draw ${\theta }^{\prime }$ from $q_{\varphi }(\theta \mid x_{\mathrm{obs}})$ cheaply. Without this we can't actually produce the posterior samples we wanted in the first place.

Flexibility. The posteriors that come out of real simulators are not Gaussians. They're skewed (asymmetric around the mode), often multimodal (two or more bumps when the model is genuinely ambiguous), often bounded (parameters live on intervals or have hard physical constraints), and almost always correlated across dimensions (knowing one parameter narrows the others). The density estimator has to be able to represent all of those shapes. A Gaussian fails three out of four. A mixture of Gaussians handles multimodality but not bounded support or skewness cleanly.

Trainable. The whole object is parameterized by some $\varphi$ and trained by gradient descent on a loss. Everything has to be differentiable.

Most generative models fail one of these. GANs sample beautifully but give you no density. VAEs and diffusion models give you flexibility and sampling but only a variational lower bound on density. Gaussians give you density and sampling but no flexibility. Flows are the one design that delivers all four exactly — no surrogate losses, no variational gaps.

The idea

A normalizing flow is one trick repeated. Take a simple base distribution you can sample from — a standard Gaussian $u\sim \mathcal{N}(0,I)$, almost always. Apply an invertible, differentiable transformation $\theta =f_{\varphi }(u)$. By the change-of-variables formula, the result is a new distribution whose density you can write down exactly.

Stack many such transformations. Each is invertible. Each contributes a Jacobian-determinant term to the log-density. The composition can fit very complicated shapes while staying tractable in both directions — you can evaluate $q_{\varphi }(\theta )$ and you can sample from it, both from the same object, both cheaply. The rest is engineering: pick invertible transformations whose Jacobian determinants are easy to compute, and make them conditional on $x$.

Change of variables, briefly

If $u\sim p_u$ and $\theta =f_{\varphi }(u)$ with $f_{\varphi }$ invertible, then the density of $\theta$ follows from change of variables:

The intuition: as $f_{\varphi }$ stretches space, it spreads probability mass over a larger volume; the density at any output point is the density at the corresponding input point, divided by the local stretching factor. The log-determinant is that factor in log form. There's a slower derivation with worked 1D examples on the Change of Variables page if you want to see it geometrically. The rest of this post will use the formula above as a given.

For a composition $f_{\varphi }=f_K\circ f_{K-1}\circ \cdots \circ f_1$, the log-determinants add up:

So if each layer has a cheap log-det-Jacobian, the whole stack does. That cheapness is the engineering constraint. A general $d\times d$ Jacobian determinant costs $O(d^3)$ to compute, which is fatal for $d$ bigger than a few. The whole story of practical flow design is "what kind of invertible transformations have $O(d)$ log-determinants?" The answer in every case is the same: make the Jacobian triangular, so its determinant is the product of its diagonal entries.

Layer families

Flows are built from one of a few standard layer types. They all impose triangularity, but they do it differently.

Affine flows

The simplest. Scale and shift each dimension independently:

The Jacobian is diagonal; the log-determinant is just ${\sum }_i{s}_i$. Useful as a building block — most flow stacks have an affine layer at the input or output for normalization — but if you only ever use affine layers you can't represent anything more complicated than a shifted, axis-scaled Gaussian. Real expressiveness has to come from non-axis-aligned moves.

Coupling layers (Real NVP, NICE)

The workhorse of practical flows. Split the variables in half: $\theta =({\theta }_A,{\theta }_B)$. Pass ${\theta }_A$ through unchanged. Transform ${\theta }_B$ by an affine map whose scale and shift are functions of ${\theta }_A$:

The Jacobian is triangular by construction: ${\theta }_A^{\prime }$ depends only on ${\theta }_A$, ${\theta }_B^{\prime }$ on both halves. The log-determinant is ${\sum }_i{s}_{\varphi }({\theta }_A{)}_i$ over the $B$ coordinates only — cheap to compute, easy to differentiate.

The functions $s_{\varphi }$ and $t_{\varphi }$ can be arbitrary neural networks. They don't have to be invertible themselves; the layer's invertibility comes from the structure (given ${\theta }_A^{\prime }={\theta }_A$, you can compute $s_{\varphi }$ and $t_{\varphi }$ directly, then invert the affine map to recover ${\theta }_B$). Stack many coupling layers, alternating which half is the "passthrough", and you get arbitrarily expressive transformations.

Autoregressive flows (MAF, IAF)

A different way to impose triangularity. Transform each coordinate ${\theta }_i$ as a function of ${\theta }_1,\dots ,{\theta }_{i-1}$:

Triangular by construction. The asymmetry is in which direction is cheap: MAF (Masked Autoregressive Flow) parallelizes the density evaluation in one neural-net pass but samples sequentially (each ${\theta }_i$ needs the previous coordinates filled in first). IAF (Inverse Autoregressive Flow) is the opposite — parallel sampling, sequential density. Pick by which direction you need.

Neural spline flows

Replace the affine inner transform of a coupling or autoregressive layer with a monotone piecewise-rational-quadratic spline. Same triangular structure, dramatically more expressive per layer — a spline can capture skew and heavy tails and bounded support that a single affine map can't, in one transformation instead of a stack of ten. Fewer layers needed for the same quality. This is the current standard for serious SBI work; the sbi package's default conditional flow is a stack of neural spline coupling layers.

Conditional flows

Everything above describes an unconditional density $q_{\varphi }(\theta )$. For SBI we want $q_{\varphi }(\theta \mid x)$ — a different density for each value of $x$. The standard trick: make every internal neural network inside the flow take $x$ as an additional input.

Concretely, an embedding network $h_{\psi }(x)$ maps the observation into a fixed-size latent vector. That latent gets concatenated to the inputs of every coupling sub-net $s_{\varphi }$ and $t_{\varphi }$ (or spline-parameter network, or autoregressive sub-net). The flow's per-layer parameters become functions of both the previous coordinates AND the embedding of $x$.

Now the whole pipeline is a conditional density. Given $x_{\mathrm{obs}}$: run it through $h_{\psi }$ once to get an embedding; sample $u\sim \mathcal{N}(0,I)$; transform forward through the flow conditioned on the embedding; out comes ${\theta }^{\prime }\sim q_{\varphi }(\theta \mid x_{\mathrm{obs}})$. Want a density at some specific $\theta$? Run the inverse, accumulate log-Jacobians, add the base log-density. One trained network handles any $x_{\mathrm{obs}}$ in a single forward pass either direction.

The embedding network $h_{\psi }$ is the part that handles structured or high-dimensional $x$. For images it's a CNN; for time series, a temporal CNN or RNN; for graphs, a GNN; for sets, a permutation-invariant network. Whatever embedding produces a useful vector summary of $x$, plug it in. The summary is learned end-to-end with the rest of the flow — no manual feature engineering, no human-picked summary statistics.

This is the explicit answer to ABC's summary-statistic burden from Post 3: the network learns its own summary as part of training, optimized to be informative about $\theta$ rather than chosen by hand to be tractable.

What this object gives us

A conditional normalizing flow is a learnable, flexible, exact-density, exact-sampling model of $q_{\varphi }(\theta \mid x)$. Each of ABC's four failure modes from the previous post is addressed by this single object:

Wasted simulations — every $(\theta ,x)$ pair contributes to training, regardless of whether $x$ is near any specific $x_{\mathrm{obs}}$. There's no rejection step throwing data away.

Curse of dimensionality — the embedding network $h_{\psi }$ learns a useful low-dimensional summary of $x$ as part of training. The flow then matches the posterior in $\theta$-space, whose dimension is usually modest. No ${\epsilon }^d$ volume collapse.

Summary-statistic burden — covered by the embedding network. You're still committing to an architecture for $h_{\psi }$, but the architecture choice is much weaker than committing to a specific summary statistic, and it's the kind of choice ML practice has a lot of precedent for.

No amortization — inference at any new $x_{\mathrm{obs}}$ is one forward pass through $h_{\psi }$ and one through the flow. Per-observation cost is microseconds after training.

The cost is paid up front in training: simulation budget for the $(\theta ,x)$ dataset, and the structural commitment to a flow architecture. The next post is about how the training works — what the loss function is, what the data is, how much of it you need, and what diagnostics to run before trusting the output.