DEV Community

Umesh Tharuka Malaviarachchi
Umesh Tharuka Malaviarachchi

Posted on

Normalizing Flows: Why Invertible Transformations Make Exact Likelihoods Possible

Most generative models force you to choose. You can sample easily but not evaluate likelihoods, as with GANs. You can evaluate a bound on the likelihood but not the exact value, as with VAEs. Or you can evaluate exact likelihoods but only for a narrow family of distributions. Normalizing flows refuse that trade-off, at least in principle: they give you exact likelihoods and exact sampling from the same model, by building a complex distribution out of a simple one through a chain of invertible transformations.

The idea is elegant, the implementation is subtle, and the limitations are instructive. Understanding flows is a good way to understand what "exact likelihood" really buys you, and what it costs.

The core idea: change of variables

Start with a simple base distribution that you understand completely, usually a standard Gaussian. Call a sample from it zz , with density pZ(z)p_Z(z) . Now apply an invertible, differentiable transformation ff to get x=f(z)x = f(z) . What is the density of xx ?

The change-of-variables formula answers this:

pX(x)=pZ(f−1(x))⋅∣det⁡Jf−1(x)∣ p_X(x) = p_Z(f^{-1}(x)) \cdot \left| \det J_{f^{-1}}(x) \right|

where Jf−1(x)J_{f^{-1}}(x) is the Jacobian matrix of the inverse transformation at xx . Equivalently, working in the forward direction:

pX(x)=pZ(z)⋅∣det⁡Jf(z)∣−1,z=f−1(x) p_X(x) = p_Z(z) \cdot \left| \det J_f(z) \right|^{-1}, \quad z = f^{-1}(x)

Intuition: the base density tells you how likely zz was, and the Jacobian determinant tells you how much ff stretches or compresses volume around zz . If ff stretches a region, density there must fall, because the same probability mass now spreads over more volume. If it compresses a region, density rises. A flow is just this accounting, done carefully, many times in a row.

Sampling is easy: draw zz from the base distribution and push it forward through ff . Density evaluation is also easy, provided you can invert ff and compute its Jacobian determinant cheaply. Those two provisos are the whole design problem of normalizing flows.

Composing simple maps into a complex distribution

A single transformation is rarely expressive enough. Flows therefore compose many:

f=fK∘⋯∘f2∘f1 f = f_K \circ \cdots \circ f_2 \circ f_1

If z0z_0 is a base sample and zk=fk(zk−1)z_k = f_k(z_{k-1}) , then the log density of the final output x=zKx = z_K is:

log⁡pX(x)=log⁡pZ(z0)−∑k=1Klog⁡∣det⁡Jfk(zk−1)∣ \log p_X(x) = \log p_Z(z_0) - \sum_{k=1}^{K} \log \left| \det J_{f_k}(z_{k-1}) \right|

Each layer adds its own log-determinant term. Training by maximum likelihood is then straightforward in principle: maximise this log density over the training data, by gradient descent through the composed maps. No variational bound, no adversarial game, no Markov chain. Just a likelihood, computed exactly.

The catch is computational. For a general neural network transformation in DD dimensions, computing a Jacobian determinant costs on the order of D3D^3 operations, which is hopeless for images. Every practical flow architecture is therefore a compromise: a family of transformations whose Jacobian is triangular, or otherwise structured, so its determinant is cheap, usually the product of diagonal entries.

Coupling layers: the workhorse design

The most influential such design is the coupling layer. Split the input into two parts, xax_a and xbx_b . Leave the first part unchanged, and transform the second part using a function that depends only on the first part:

ya=xa y_a = x_a
yb=xb⊙exp⁡(s(xa))+t(xa) y_b = x_b \odot \exp(s(x_a)) + t(x_a)

Here ss and tt are arbitrary neural networks, ⊙\odot is elementwise multiplication, and exp⁡\exp ensures the scale is positive. This is an affine coupling layer.

Why is this clever? First, it is trivially invertible. Given yy , recover xa=yax_a = y_a , then:

xb=(yb−t(ya))⊙exp⁡(−s(ya)) x_b = (y_b - t(y_a)) \odot \exp(-s(y_a))

Second, its Jacobian is lower triangular, because yay_a does not depend on xbx_b . The determinant is therefore just the product of the diagonal scale terms:

log⁡∣det⁡J∣=∑jsj(xa) \log \left| \det J \right| = \sum_j s_j(x_a)

No matrix determinant, no cubic cost. And ss and tt themselves never need to be invertible, so they can be deep, expressive networks. The invertibility lives in the structure of the coupling, not in the networks inside it.

One coupling layer only transforms half the variables, so practical flows alternate which half is transformed, and often permute or shuffle dimensions between layers, so that every dimension gets transformed eventually. Stack enough of these layers and the composed map can represent surprisingly complex distributions.

A minimal PyTorch-style sketch makes the structure concrete:

class AffineCoupling(nn.Module):
    def __init__(self, dim_a, dim_b, hidden):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim_a, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, 2 * dim_b),
        )

    def forward(self, x_a, x_b):
        # returns y_a, y_b, log determinant
        s, t = self.net(x_a).chunk(2, dim=-1)
        s = torch.tanh(s)  # keep scales stable in practice
        y_a = x_a
        y_b = x_b * torch.exp(s) + t
        log_det = s.sum(dim=-1)
        return y_a, y_b, log_det

    def inverse(self, y_a, y_b):
        s, t = self.net(y_a).chunk(2, dim=-1)
        s = torch.tanh(s)
        x_a = y_a
        x_b = (y_b - t) * torch.exp(-s)
        return x_a, x_b
Enter fullscreen mode Exit fullscreen mode

Training then sums the base log density and the negative log determinants across layers, and maximises that total. The code is short. The design decisions around it, such as how to split, permute, and stabilise scales, are where most of the practical work lies.

Other flow families, and why structure matters

Coupling layers are not the only option. Autoregressive flows make each output dimension depend only on earlier dimensions, giving a triangular Jacobian by construction, at the cost of sequential sampling or sequential density evaluation depending on direction. Continuous normalizing flows replace discrete layers with an ordinary differential equation, where the log-density change is governed by the trace of a Jacobian rather than its determinant, which opens up free-form architectures but requires an ODE solver. Residual flows use a contraction condition to guarantee invertibility and estimate the log determinant with a series expansion.

Each family trades off the same three axes: expressiveness, cost of sampling, and cost of density evaluation. A flow that is fast in one direction may be slow in the other. Before choosing a family, decide which direction your application actually needs. Anomaly detection needs fast density evaluation. Data generation needs fast sampling. A model that is elegant in the wrong direction is not useful.

What exact likelihood actually gives you

Exact likelihood is genuinely valuable for several tasks:

  • Anomaly detection and out-of-distribution detection. A test point with very low likelihood under the model is a candidate anomaly, with a principled score rather than a heuristic reconstruction error.
  • Model comparison. Likelihoods on held-out data give a comparable, quantitative score across flow architectures, without the sampling variability that plagues some generative metrics.
  • Compression and coding. An exact density can drive an arithmetic coder, turning a good density model into a good compressor.
  • Posterior inference and sampling. Invertible maps can transform a simple distribution into a complex posterior approximation, or serve as a proposal in a larger inference scheme.

It is worth being precise about what exact likelihood does not give you. It does not guarantee good samples, because likelihood is dominated by typical regions and can be high even when samples look poor in ways humans notice. It does not guarantee good representations, because the latent space is constrained to have the same dimension as the data. And it does not protect you from the classic likelihood pathology: assigning high density to a narrow spike on training data while missing broader structure.

Limitations and common pitfalls

Several limitations recur, and they explain why flows have not replaced other generative models despite their elegance:

  • Dimension preservation. A flow maps DD dimensions to DD dimensions. It cannot compress to a lower-dimensional latent space, unlike a VAE, and it cannot easily model data that truly live on a lower-dimensional manifold embedded in a higher-dimensional space. In that setting, the Jacobian determinant becomes ill-behaved, and training can become unstable.
  • Topological constraints. A continuous invertible map cannot change topology. It cannot split one connected component into two disconnected ones, or map a unimodal base distribution onto a sharply multimodal target without stretching regions thin. In practice, flows approximate such targets with narrow bridges of density, which can show up as poor samples between modes.
  • Expressiveness versus cost. Triangular Jacobians are cheap, but they restrict each layer. Getting high expressiveness requires many layers, careful permutations, and long training. A flow that is theoretically universal may be practically weak at the depth you can afford.
  • Numerical stability. Exponentiated scales, deep compositions, and log determinants can overflow or underflow. Scale clipping, actnorm-style data-dependent initialisation, and careful learning-rate choices are not cosmetic details for flows. They are often the difference between training and divergence.
  • Likelihood is not sample quality. Do not judge a flow only by held-out likelihood, and do not judge it only by samples. Report both, plus task-specific metrics. A flow with excellent likelihood and mediocre samples is telling you something about likelihood as a metric, not just about the model.

When to use a flow

Use a normalizing flow when you need both sampling and exact density evaluation from the same model, when the data dimension is moderate or the structure allows efficient coupling, and when anomaly scoring or likelihood-based comparison is central to the task. Flows are a strong choice for tabular density estimation, for scientific data where likelihood matters, and as a component inside larger systems, such as a flexible posterior or a learned prior.

Prefer another family when you need a low-dimensional representation, when the data are high-dimensional images and sample quality dominates, or when the target distribution is strongly multimodal or manifold-structured. In those cases, diffusion models or VAEs often give better samples for the same compute, even though they sacrifice exact likelihood.

Takeaway

Normalizing flows are a lesson in how a single mathematical identity, the change of variables, can become a practical modelling tool once you design around its computational cost. The Jacobian determinant is not a technical footnote. It is the reason flows work, the reason their architectures look the way they do, and the reason they have the limits they do.

If you remember one equation from this article, make it the log-density decomposition: base log density minus the sum of log determinants. Everything else in flow design is a strategy for making that sum cheap, stable, and expressive enough for the distribution you actually need to model.


Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial

Top comments (0)