Normalizing Flows & Invertible Networks

1. The Change of Variables Formula

Normalizing flows are a class of generative models that learn invertible transformations between a simple base distribution and a complex target distribution. The foundation rests on the change of variables formula from probability theory.

Let be a random variable with density , and let be an invertible, differentiable mapping. Then has density:

Taking the log-density:

The critical quantity is the log-determinant of the Jacobian (log-det-Jacobian):

1.1 The Push-Forward Perspective

If is a diffeomorphism (smooth bijection with smooth inverse), then the density transforms under the push-forward operation:

For deep flows, we compose invertible transformations:

The log-density becomes:

1.2 The Inverse Problem

For generation, we sample and transform . For density evaluation, we need the inverse and the sum of log-determinants.

The key design principle: construct such that:

1. The forward pass is efficient
2. The inverse is efficient
3. The log-det-Jacobian is cheap to compute

2. Coupling Layers: The Foundation

2.1 RealNVP Architecture

The Real-valued Non-Volume Preserving (RealNVP) transformation splits the input into two halves: and . The coupling layer is:

where are scale and translation networks (typically neural networks).

Jacobian structure:

The log-determinant is:

This is — dramatically cheaper than the required for a general matrix.

2.2 Inverse of Coupling Layers

The inverse is straightforward:

3. Flow Architecture Comparison

<svg viewBox="0 0 800 420" xmlns="http://www.w3.org/2000/svg">
  <rect width="800" height="400" fill="#f8fafc"/>
  
  {/* Title */}
  <text x="400" y="30" text-anchor="middle" fill="#f8fafc" font-family="monospace" font-size="16" font-weight="bold">Normalizing Flow Architecture</text>
  
  {/* Base Distribution */}
  <ellipse cx="100" cy="200" rx="60" ry="80" fill="none" stroke="#3b82f6" stroke-width="2" stroke-dasharray="5,3"/>
  <text x="100" y="205" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="11">z₀ ~ N(0,I)</text>
  
  {/* Flow Layers */}
  <g transform="translate(200, 120)">
    <rect width="100" height="160" rx="8" fill="#f8fafc" stroke="#3b82f6" stroke-width="2"/>
    <text x="50" y="25" text-anchor="middle" fill="#60a5fa" font-family="monospace" font-size="10">Coupling</text>
    <text x="50" y="45" text-anchor="middle" fill="#60a5fa" font-family="monospace" font-size="10">Layer 1</text>
    <text x="50" y="75" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">s₁(x₁)</text>
    <text x="50" y="95" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">t₁(x₁)</text>
    <text x="50" y="130" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="9">log|det| =</text>
    <text x="50" y="150" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="9">Σ s₁</text>
  </g>
  
  <g transform="translate(330, 120)">
    <rect width="100" height="160" rx="8" fill="#f8fafc" stroke="#a855f7" stroke-width="2"/>
    <text x="50" y="25" text-anchor="middle" fill="#c084fc" font-family="monospace" font-size="10">Coupling</text>
    <text x="50" y="45" text-anchor="middle" fill="#c084fc" font-family="monospace" font-size="10">Layer 2</text>
    <text x="50" y="75" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">s₂(x₂)</text>
    <text x="50" y="95" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">t₂(x₂)</text>
    <text x="50" y="130" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="9">log|det| =</text>
    <text x="50" y="150" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="9">Σ s₂</text>
  </g>
  
  <g transform="translate(460, 120)">
    <rect width="100" height="160" rx="8" fill="#f8fafc" stroke="#f59e0b" stroke-width="2"/>
    <text x="50" y="25" text-anchor="middle" fill="#fbbf24" font-family="monospace" font-size="10">Coupling</text>
    <text x="50" y="45" text-anchor="middle" fill="#fbbf24" font-family="monospace" font-size="10">Layer K</text>
    <text x="50" y="75" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">sK(xK)</text>
    <text x="50" y="95" text-anchor="middle" fill="#94a3b8" font-family="monospace" font-size="9">tK(xK)</text>
    <text x="50" y="130" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="9">log|det| =</text>
    <text x="50" y="150" text-anchor="middle" fill="#22c55e" font-family="monospace" font-size="9">Σ sK</text>
  </g>
  
  {/* Data Distribution */}
  <path d="M 650 120 Q 700 180 680 250 Q 660 300 620 320 Q 580 300 560 250 Q 540 180 600 120 Z" fill="none" stroke="#ec4899" stroke-width="2"/>
  <text x="610" y="210" text-anchor="middle" fill="#ec4899" font-family="monospace" font-size="11">p(x)</text>
  
  {/* Arrows */}
  <defs>
    <marker id="3574_arrowhead" markerWidth="10" markerHeight="7" refX="9" refY="3.5" orient="auto">
      <polygon points="0 0, 10 3.5, 0 7" fill="#60a5fa"/>
    </marker>
  </defs>
  <line x1="160" y1="200" x2="198" y2="200" stroke="#60a5fa" stroke-width="2" marker-end="url(#3574_arrowhead)"/>
  <line x1="300" y1="200" x2="328" y2="200" stroke="#60a5fa" stroke-width="2" marker-end="url(#3574_arrowhead)"/>
  <line x1="430" y1="200" x2="458" y2="200" stroke="#60a5fa" stroke-width="2" marker-end="url(#3574_arrowhead)"/>
  <line x1="560" y1="200" x2="600" y2="200" stroke="#60a5fa" stroke-width="2" marker-end="url(#3574_arrowhead)"/>
  
  {/* Labels */}
  <text x="400" y="380" text-anchor="middle" fill="#64748b" font-family="monospace" font-size="10">Forward: z₀ → x = fK∘...∘f1(z₀)</text>
</svg>

4. Glow Architecture

Glow (Kingma & Dhariwal, 2018) extends RealNVP with three innovations:

4.1 Actnorm (Activation Normalization)

Channel-wise affine transformation learned per-sample, initialized via data-dependent initialization:

where are initialized such that and over the first batch.

4.2 Invertible 1×1 Convolution

Replaces the channel permutation in RealNVP. For :

where is a learned weight matrix. The log-determinant:

Using LU decomposition , the log-det becomes:

4.3 Multi-Scale Architecture

Factorizes the latent space into scales:

where are the "detail" latents and is the final coarse representation.

5. Autoregressive Flows

5.1 Neural Autoregressive Flow (NAF)

Autoregressive flows use the chain rule of probability:

Each conditional is modeled as an invertible transformation:

The MAF (Masked Autoregressive Flow) parameterizes the density:

The IAF (Inverse Autoregressive Flow) parameterizes the sample:

5.2 Complexity Analysis

6. Continuous Normalizing Flows (CNF)

6.1 Neural ODE Formulation

Instead of discrete layers, CNFs define a continuous transformation via an ODE:

where , , and .

The density evolution is governed by the continuity equation:

6.2 Instantaneous Change of Variables

This is equivalent to computing the divergence of . Using the Hutchinson trace estimator:

6.3 Adjoint Sensitivity Method

For training, we solve the ODE backward in time with the adjoint state :

This gives memory complexity at the cost of backward passes.

import torch
from torchdiffeq import odeint_adjoint as odeint

class ContinuousNormalizingFlow(torch.nn.Module):
    def __init__(self, f, trace_estimator='hutchinson'):
        super().__init__()
        self.f = f
        self.trace_estimator = trace_estimator
    
    def forward(self, z_0, log_pz0, integration_times=None, reverse=False):
        if integration_times is None:
            integration_times = torch.tensor([0.0, 1.0])
        
        if reverse:
            integration_times = integration_times.flip(0)
            z_0 = z_0.flip(0)
        
        # Solve ODE for state
        z_t = odeint(self.f, z_0, integration_times,
                     method='dopri5', rtol=1e-5, atol=1e-5)
        
        # Solve ODE for log-density
        def div_f(z, t):
            z.requires_grad_(True)
            f = self.f(z, t)
            if self.trace_estimator == 'hutchinson':
                u = torch.randn_like(z)
                tr = (u * torch.autograd.grad(f, z, grad_outputs=u)[0]).sum(-1)
            else:
                tr = torch.autograd.grad(f, z, grad_outputs=torch.ones_like(f))[0].diagonal(0,-1,-2).sum(-1)
            return tr
        
        log_p_t = odeint(div_f, log_pz0, integration_times,
                        method='dopri5', rtol=1e-5, atol=1e-5)
        
        return z_t[-1], log_p_t[-1]

7. Coupling Layer Mathematics

7.1 General Affine Coupling

For a general partition with :

The log-determinant:

7.2 Non-Linear Coupling (NICE)

Setting (pure addition coupling):

The log-determinant is exactly zero — volume preserving.

7.3 Additive Coupling with Neural Spline Flows

The Neural Spline Flow (NSF) uses monotonic rational-quadratic splines:

where the spline is defined on bins with learned parameters .

8. Training Objective

The maximum likelihood objective for normalizing flows:

8.1 KL Divergence Perspective

Minimizing KL divergence is equivalent to maximizing the log-likelihood.

9. Advanced Flow Architectures

9.1 Flow++ (Ho et al., 2019)

Uses variational dequantization and neural spline coupling layers:

where is the dequantization noise.

9.2 Residual Flows

Based on the Banach fixed-point theorem:

with . The log-determinant via the matrix determinant lemma:

approximated by Neumann series.

9.3 Free-Form Flows

The Free-form Normalizing Flow (FFJORD, Grathwohl et al., 2019) allows arbitrary neural network architectures:

with stochastic trace estimation:

where and is computed via reverse-mode autodiff.

10. Summary Table

Normalizing flows remain one of the few generative model families that provide exact density evaluation, making them invaluable for likelihood-based tasks and scientific applications where precise probability computation is essential.