PDE and ML (6): Continuous Normalizing Flows and Neural ODE

How do you turn an isotropic Gaussian into a photograph of a cat?

Normalizing flows give the most direct answer: stack a sequence of invertible transformations and let them push the simple distribution into the complex one. This article’s continuous version (CNF) takes that idea to the limit — let the step size go to zero and the discrete chain becomes an ODE. Invertibility is automatic, and the change of density is governed by the instantaneous change of variables formula.

When the right-hand side of that ODE is a neural network, the whole thing is called a Neural ODE, and a few of its properties are genuinely elegant: memory cost decoupled from the number of integration steps (the adjoint method), continuous-time interpolation, a direct connection to optimal transport.

This article walks Neural ODEs, the adjoint method, continuous normalizing flows (FFJORD), and finally the post-2023 method that has largely replaced them in large-scale generative modeling — Flow Matching. Seven figures clarify the core mechanics.

What You Will Learn#

Generative modeling reduces to one geometric question: how do you transform a simple distribution (a Gaussian) into a complex one (faces, molecules, motion)? Discrete normalizing flows stack invertible blocks, but each block needs a Jacobian determinant at $O(d^3)$ cost. Neural ODEs replace discrete depth with a continuous ODE; Continuous Normalizing Flows (CNF) then push densities through that ODE using the instantaneous change-of-variables formula, dropping density computation to $O(d)$ . Flow Matching removes the divergence integral altogether and turns training into plain regression on a target velocity field.

Three threads are braided together throughout the chapter:

1. PDE side — the continuity equation $\partial_t\rho+\nabla\!\cdot(\rho v)=0$ governs how a velocity field $v$ transports a density $\rho$ .
2. ODE side — Picard-Lindelof guarantees existence/uniqueness; Liouville’s theorem links volume change to $\nabla\!\cdot v$ ; the adjoint equation makes backprop $O(1)$ in memory.
3. ML side — Neural ODEs, FFJORD and Flow Matching parameterise $v$ with a network and learn it from data.

Prerequisites: ODE basics (existence/uniqueness), probability change of variables, autograd.

ODE Foundations: Existence, Uniqueness, Volume#

Picard-Lindelof: when do ODEs have unique solutions?#

then a unique solution exists on some interval $[0,T]$ .

Why this matters for ML. If $f_\theta$ is a neural network with Lipschitz activations (ReLU, tanh, GELU) and bounded weights, the Lipschitz condition holds locally. So a Neural ODE is well-posed as long as the network is well-behaved — which is almost always true in practice and which is why Neural ODEs are robust enough to backprop through.

Liouville’s theorem: how flows change volume#

Therefore $\nabla\!\cdot f=0$ preserves volume, $\nabla\!\cdot f<0$ contracts, $\nabla\!\cdot f>0$ expands. In normalizing flows we want a non-zero divergence: that is exactly the lever that lets us reshape probability mass.

Mental picture. A divergence-free $f$ behaves like an incompressible fluid (Hamiltonian / symplectic — Part 5 ). A divergence-rich $f$ behaves like a compressible flow that can squeeze probability mass into thin filaments and then re-inflate it elsewhere — which is what generative modelling needs.

Instantaneous change of variables#

Proof sketch. The continuity equation $\partial_t\rho+\nabla\!\cdot(\rho f)=0$ expands to $\partial_t\rho+f\!\cdot\!\nabla\rho=-\rho\,\nabla\!\cdot f$ . The left-hand side is the material derivative $D\rho/Dt$ along $\mathbf{z}(t)$ . Dividing by $\rho$ gives (1).

Why this single equation matters. Discrete normalizing flows pay $O(d^3)$ for $\log|\det\partial\phi/\partial\mathbf{z}|$ . Equation (1) only ever needs the trace of the Jacobian (i.e. the divergence), which costs $O(d)$ with a vector-Jacobian product (Section 3.2 below). This is the central computational reason CNFs exist.

Neural ODEs: From Discrete to Continuous Depth#

Residual networks as forward Euler#

Three immediate wins:

• Parameter efficiency. One network $f_\theta$ replaces a different $f_l$ at every depth.
• Adaptive depth. Solvers like dopri5 (adaptive Runge-Kutta) automatically take small steps where dynamics are stiff and large steps where they are smooth.
• Memory. The adjoint method drops backprop memory from $O(L)$ to $O(1)$ — the only resource that scales with depth is time, not memory.

Implementation: a minimal Neural ODE. The core idea is to replace a ResNet’s forward pass with an ODE solver call. Here is a complete, runnable implementation:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

import torch
import torch.nn as nn
from scipy.integrate import solve_ivp

class ODEFunc(nn.Module):
    # The vector field f_theta(h, t) that defines the dynamics
    def __init__(self, dim=2, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim + 1, hidden), nn.Tanh(),
            nn.Linear(hidden, hidden), nn.Tanh(),
            nn.Linear(hidden, dim)
        )

    def forward(self, t, h):
        t_vec = t * torch.ones(h.shape[0], 1)
        return self.net(torch.cat([h, t_vec], dim=-1))

def neural_ode_forward(func, h0, t_span=(0, 1), steps=100):
    # Simple fixed-step Euler integration (for clarity)
    dt = (t_span[1] - t_span[0]) / steps
    h = h0
    t = t_span[0]
    for _ in range(steps):
        h = h + dt * func(torch.tensor(t), h)
        t += dt
    return h

func = ODEFunc(dim=2, hidden=64)
opt = torch.optim.Adam(func.parameters(), lr=1e-3)

t_true = torch.linspace(0, 1, 50)
theta = 4 * torch.pi * t_true
r = 2 * (1 - t_true)
target = torch.stack([r * torch.cos(theta), r * torch.sin(theta)], dim=-1)

for step in range(2000):
    h0 = target[0:1]  # start point
    h_pred = neural_ode_forward(func, h0.repeat(50, 1), steps=50)
    loss = ((h_pred - target)**2).mean()
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 500 == 0:
        print(f"Step {step}: loss={loss.item():.5f}")

The Neural ODE learns a smooth vector field that pushes any starting point along the spiral. Unlike a ResNet, the integration depth (number of Euler steps) is a hyperparameter, not a fixed architectural choice — and can be made adaptive.

The adjoint sensitivity method#

A standard backprop through the ODE solver stores every intermediate state, which is $O(L)$ in the number of solver steps — and adaptive solvers can take hundreds of them. The adjoint method avoids that completely.

Algorithm.

1. Forward. Solve (2) from $0\to T$ . Store only $\mathbf{h}(T)$ .
2. Initialise. $\mathbf{a}(T)=\partial\mathcal{L}/\partial\mathbf{h}(T)$ .
3. Backward. Solve $\mathbf{h}$ and $\mathbf{a}$ together backwards $T\to 0$ , accumulating (4).

Memory is $O(1)$ , independent of solver steps. The price is one extra ODE solve in the backward pass — roughly 2x compute for $L\to\infty$ memory savings.

Implementation: adjoint sensitivity in practice. The adjoint method avoids storing intermediate states by re-computing them during the backward pass. Here is the key idea implemented:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

import torch
import torch.nn as nn

def adjoint_backward(func, h_T, loss_grad, theta, T=1.0, steps=100):
    # Backward ODE: re-derive h(t) and compute param gradients
    # h_T: final state, loss_grad: dL/dh(T), theta: model params
    dt = T / steps
    h = h_T.detach().clone().requires_grad_(True)
    a = loss_grad.clone()  # adjoint state
    param_grad = [torch.zeros_like(p) for p in theta]

    for step in range(steps):
        t = T - step * dt
        # Recompute f(h, t)
        f = func(torch.tensor(t), h)
        # Adjoint dynamics: da/dt = -a^T df/dh
        vjp = torch.autograd.grad(f, h, -a, retain_graph=True)[0]
        a = a + dt * vjp
        # Accumulate parameter gradients: dL/dtheta += -a^T df/dtheta
        grads = torch.autograd.grad(f, theta, -a * dt, retain_graph=True)
        for pg, g in zip(param_grad, grads):
            pg += g
        # Reverse Euler step on h
        h = (h - dt * f).detach().requires_grad_(True)

    return param_grad

print("Standard backprop: O(L * d) memory (stores all intermediate h)")
print("Adjoint method:    O(d) memory (re-derives h on the fly)")
print("Cost:              ~2x compute (one extra ODE solve)")

For deep networks ($L > 100$ ) or high-dimensional states ($d > 1000$ ), the adjoint method’s memory savings are decisive. In practice, torchdiffeq implements this automatically.

Universality#

Neural ODEs are dense in the space of homeomorphisms of $\mathbb{R}^d$ (Zhang et al. 2020). They cannot, however, change topology — a single Neural ODE on $\mathbb{R}^d$ cannot un-link two linked rings. This motivates augmented Neural ODEs that lift to $\mathbb{R}^{d+k}$ , where extra coordinates give the flow room to untangle.

Continuous Normalizing Flows (CNF)#

From discrete flows to continuous#

Each $\det$ is $O(d^3)$ unless the architecture is engineered (coupling layers, autoregressive, etc.) — which restricts expressivity.

No invertibility constraint on the architecture — the ODE is invertible by integrating backwards. No determinant — only a trace.

FFJORD: scalable trace via Hutchinson#

This is Hutchinson’s trace estimator and it costs one vector-Jacobian product per sample — independent of $d$ .

Training and sampling#

where $\mathbf{z}_0$ is obtained by solving (5) backwards from $\mathbf{x}$ . Maximise the log-likelihood with the adjoint method. To sample, draw $\mathbf{z}_0\sim p_0$ and integrate forward.

Trade-offs. CNFs give exact-likelihood density estimation, but each forward/backward pass requires solving an ODE — that’s typically tens to hundreds of network evaluations. Training is also somewhat fragile: solver tolerance, regularisation of $f_\theta$ , and Hutchinson variance all interact.

Implementation: FFJORD density estimation. Combining the Neural ODE with Hutchinson trace estimation, we get a complete density estimator:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

import torch
import torch.nn as nn

class FFJORD(nn.Module):
    def __init__(self, dim=2, hidden=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim + 1, hidden), nn.Softplus(),
            nn.Linear(hidden, hidden), nn.Softplus(),
            nn.Linear(hidden, dim)
        )

    def forward(self, t, z):
        t_vec = t * torch.ones(z.shape[0], 1)
        return self.net(torch.cat([z, t_vec], dim=-1))

    def log_prob(self, x, steps=100):
        # Integrate backward: x at t=1 -> z0 at t=0
        dt = 1.0 / steps
        z = x.clone().requires_grad_(True)
        log_det = torch.zeros(x.shape[0])
        for step in range(steps):
            t = 1.0 - step * dt
            f = self.forward(torch.tensor(t), z)
            # Hutchinson trace estimator
            eps = torch.randn_like(z)
            jvp = torch.autograd.grad(f, z, eps, create_graph=True)[0]
            div = (jvp * eps).sum(dim=-1)  # trace estimate
            log_det = log_det - dt * div
            z = z - dt * f
            z = z.detach().requires_grad_(True)
        # Base distribution: standard Gaussian
        log_p0 = -0.5 * (z**2).sum(dim=-1) - z.shape[-1] * 0.5 * torch.log(torch.tensor(2*torch.pi))
        return log_p0 + log_det

model = FFJORD(dim=2)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for step in range(5000):
    log_px = model.log_prob(x_data)
    loss = -log_px.mean()  # negative log-likelihood
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 1000 == 0:
        print(f"Step {step}: NLL={loss.item():.3f}")

The log_prob method integrates the ODE backwards while accumulating the log-density change via Hutchinson’s estimator. The base log-probability $\log p_0(\mathbf{z}_0)$ plus the accumulated divergence integral gives the data log-likelihood — all without computing a single determinant.

Optimal Transport and Flow Matching#

The Benamou-Brenier connection#

The minimiser $v_t^\star$ is exactly the velocity field of a CNF — and one whose trajectories are straight lines (in the Euclidean OT case). This is the cleanest geometric reason to combine CNFs with OT.

Flow Matching#

Flow Matching (Lipman et al. 2022) is the killer-app simplification. Instead of optimising NLL through an ODE solver — and instead of solving an OT problem — it picks a conditional probability path and regresses on the corresponding velocity.

Key theorem (Lipman et al.). The marginal velocity $\mathbb{E}[u_t^\star\mid\mathbf{z}_t]$ satisfies the continuity equation transporting $p_0\to p_1$ . So minimising (8) over a flexible $v_\theta$ recovers a valid CNF — without ever computing a divergence at training time.

What you actually use in 2024#

In 2024 most production-scale continuous-flow systems (image, audio, molecule generation) are some flavour of Flow Matching or Rectified Flow.

Implementation: Flow Matching — the modern standard. Flow Matching replaces the entire FFJORD machinery with simple regression. The training loop is remarkably short:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

import torch
import torch.nn as nn

class VelocityNet(nn.Module):
    # Time-conditioned velocity field v_theta(z, t)
    def __init__(self, dim=2, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim + 1, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, dim)
        )

    def forward(self, z, t):
        t_vec = t.unsqueeze(-1) if t.dim() == 1 else t
        return self.net(torch.cat([z, t_vec], dim=-1))

def flow_matching_loss(model, x1, sigma_min=1e-4):
    # x1: data samples
    batch = x1.shape[0]
    t = torch.rand(batch, 1)  # uniform t in [0, 1]
    x0 = torch.randn_like(x1)  # noise samples
    # Conditional path: z_t = (1 - t) * x0 + t * x1
    zt = (1 - t) * x0 + t * x1
    # Target velocity: dx/dt = x1 - x0
    target = x1 - x0
    # Predict and regress
    pred = model(zt, t)
    return ((pred - target)**2).mean()

model = VelocityNet(dim=2)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for step in range(5000):
    loss = flow_matching_loss(model, x_data)
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 1000 == 0:
        print(f"Step {step}: FM loss={loss.item():.4f}")

@torch.no_grad()
def sample_flow_matching(model, n=2000, steps=100):
    # Generate samples by integrating the learned velocity forward
    z = torch.randn(n, 2)
    dt = 1.0 / steps
    for i in range(steps):
        t = torch.full((n, 1), i * dt)
        z = z + dt * model(z, t)
    return z

samples = sample_flow_matching(model, n=3000)
print(f"Generated {samples.shape[0]} samples, mean={samples.mean(0).numpy().round(3)}")

Compare the simplicity: FFJORD needs an ODE solver in training, Hutchinson trace estimation, careful tolerance tuning. Flow Matching is just regression — sample noise, sample data, interpolate, predict velocity. The ODE solver is only needed at inference time, and even then a simple Euler with 50-100 steps works well.

Continuous Depth in Pictures#

The “continuous depth” idea is what unifies everything in this chapter — a Neural ODE is the continuous limit of a deep network, and CNFs are the continuous limit of a normalizing flow. The picture is the same in both cases.

This is also why one ODE function $f_\theta$ “replaces” hundreds of layers in a deep ResNet: the time variable takes over the role of layer index, and the solver decides discretisation.

Augmented Neural ODEs#

Standard Neural ODEs on $\mathbb{R}^d$ cannot change topology — a connected component stays connected, two separate clusters cannot be merged. This is because the ODE flow is a homeomorphism (continuous bijection). In practice this means that a 2D Neural ODE cannot transform a single Gaussian into two separated clusters — the trajectories would have to cross, violating uniqueness.

The fix: augmentation. Lift the state from $\mathbb{R}^d$ to $\mathbb{R}^{d+k}$ by appending $k$ extra coordinates initialised to zero:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

def augmented_neural_ode(func, x, k=1, steps=100):
    # Augment: append k zeros
    aug = torch.zeros(x.shape[0], k)
    h = torch.cat([x, aug], dim=-1)  # (batch, d+k)
    # Integrate in augmented space
    dt = 1.0 / steps
    for step in range(steps):
        t = step * dt
        h = h + dt * func(torch.tensor(t), h)
    # Project back to original d dimensions
    return h[:, :x.shape[-1]]

In $\mathbb{R}^{d+k}$ , the flow has room to “lift” one cluster up, slide it over the other, and bring it back down — like an overpass on a highway. This costs $k$ extra dimensions but removes the topological limitation entirely. Dupont et al. (2019) showed that $k=1$ suffices for most 2D problems; for general $d$ -dimensional data, $k = d$ gives maximal flexibility.

Putting It Together: 2D Density Estimation#

To make the whole pipeline concrete, here is what density estimation actually looks like end-to-end on the canonical two-moons toy.

This dual nature — exact-likelihood density estimation and sampling, through one network $v_\theta$ and the ODE $\dot{\mathbf{z}}=v_\theta(\mathbf{z},t)$ — is what makes continuous flows so attractive theoretically.

Experiments#

Spiral ODE fitting#

A 3-layer MLP (hidden dim 64, tanh) parameterises $f_\theta$ . Trained with the adjoint method on a 2D damped spiral target via dopri5 (rtol $=10^{-5}$ ). After 1000 steps the average trajectory error falls below $10^{-3}$ , with peak GPU memory ~40 MB regardless of the ~80 internal solver steps.

Gaussian -> two moons CNF#

A 4-layer MLP (hidden dim 128, softplus) trained as FFJORD with Hutchinson trace estimation, dopri5 solver, 5000 steps. Generated samples cover both moons and capture their crescent thickness; KDE comparison gives Wasserstein-2 $\approx 0.07$ versus the reference target.

Adjoint vs standard backprop (illustrative; numbers from the original Neural ODE paper, scaled to 1024-dim hidden state)#

Memory drops by ~87%; wall-clock cost increases by ~20-30%.

Flow Matching vs CNF on 2D moons#

Flow Matching converges $\sim 2.7\times$ faster and produces qualitatively cleaner moons. On real image data the gap is even larger (one to two orders of magnitude in both training time and sampling NFE).

Exercises#

Exercise 1. Derive the instantaneous change-of-variables formula (1) directly from the continuity equation.

Solution. Continuity: $\partial_t\rho+\nabla\!\cdot(\rho f)=0$ , i.e. $\partial_t\rho+f\!\cdot\!\nabla\rho+\rho\,\nabla\!\cdot f=0$ . Along $\mathbf{z}(t)$ , $\frac{d}{dt}\rho(\mathbf{z}(t),t)=\partial_t\rho+f\!\cdot\!\nabla\rho=-\rho\,\nabla\!\cdot f$ . Divide by $\rho$ .

Exercise 2. Why is the adjoint method $O(1)$ in memory?

Solution. It only stores the current $\mathbf{h}(t)$ , $\mathbf{a}(t)$ , and the running parameter-gradient accumulator. Whenever the backward solver needs an old $\mathbf{h}(s)$ , it re-derives it by integrating the forward ODE in reverse — so no internal solver states are stored. The $O(1)$ refers to depth-independence; spatial dim still costs $d$ .

Exercise 3. Hutchinson’s estimator: prove (6) is unbiased.

Solution. For any matrix $A$ and $\boldsymbol\epsilon$ with $\mathbb{E}[\boldsymbol\epsilon]=0$ , $\mathrm{Cov}[\boldsymbol\epsilon]=\mathbf{I}$ , $\mathbb{E}[\boldsymbol\epsilon^\top A\,\boldsymbol\epsilon]=\sum_{i,j}A_{ij}\,\mathbb{E}[\epsilon_i\epsilon_j]=\sum_i A_{ii}=\mathrm{tr}\,A$ .

Exercise 4. Compare Flow Matching and DDPM at a high level.

Solution. Both transport noise $\to$ data. DDPM learns a denoiser via score matching on a stochastic forward (SDE) noising process; sampling solves a reverse SDE or its probability-flow ODE. Flow Matching learns a velocity field $v_\theta$ on a deterministic ODE matching a chosen conditional path; sampling integrates that ODE. Flow Matching’s training loss is plain regression with no time-varying noise schedule.

Exercise 5. Show that for the linear conditional path $\mathbf{z}_t=(1-t)\mathbf{z}_0+t\mathbf{z}_1$ , the marginal velocity $\mathbb{E}[\mathbf{z}_1-\mathbf{z}_0\mid\mathbf{z}_t]$ pushes $p_0$ to $p_1$ via the continuity equation.

Solution sketch. Write $\rho_t(\mathbf{z})=\int q(\mathbf{z}_0,\mathbf{z}_1)\,\delta(\mathbf{z}-(1-t)\mathbf{z}_0-t\mathbf{z}_1)\,d\mathbf{z}_0\,d\mathbf{z}_1$ . Differentiate $\rho_t$ in $t$ and use the identity $\partial_t\delta=-\nabla\!\cdot[(\mathbf{z}_1-\mathbf{z}_0)\delta]$ . Marginalising over $\mathbf{z}_0,\mathbf{z}_1$ given $\mathbf{z}_t$ yields $\partial_t\rho_t+\nabla\!\cdot(\rho_t\,\bar v_t)=0$ with $\bar v_t(\mathbf{z})=\mathbb{E}[\mathbf{z}_1-\mathbf{z}_0\mid\mathbf{z}_t=\mathbf{z}]$ .

What’s next#

The handful of core ideas in this chapter (PDE residual as loss, operators on function spaces, Wasserstein geometry, symplectic structure, scores, diffusion) recur throughout the rest of the series. If a section stalls you, jot the question down and keep reading — the next chapter usually re-explains it from a different angle.

The fastest sanity check on your own understanding is to run this chapter’s equation on a minimal example: a 1-D heat equation, a single pendulum, a 2-D Gaussian mixture. The code is short, but it converts “looks right” into “it’s right on my machine.”

References#

[1] Chen, R. T. Q., Rubanova, Y., Bettencourt, J., & Duvenaud, D. (2018). Neural ordinary differential equations. NeurIPS.

[2] Grathwohl, W., Chen, R. T. Q., Bettencourt, J., Sutskever, I., & Duvenaud, D. (2018). FFJORD: Free-form continuous dynamics for scalable reversible generative models. ICLR.

[3] Onken, D., Fung, S. W., Li, X., & Ruthotto, L. (2021). OT-Flow: Fast and accurate continuous normalizing flows via optimal transport. AAAI.

[4] Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., & Le, M. (2022). Flow matching for generative modeling. ICLR.

[5] Liu, X., Gong, C., & Liu, Q. (2022). Flow straight and fast: Learning to generate and transfer data with rectified flow. ICLR.

[6] Tzen, B., & Raginsky, M. (2019). Theoretical guarantees for sampling and inference in generative models with latent diffusions. COLT.

[7] Zhang, H., Gao, X., Unterman, J., & Arodz, T. (2020). Approximation capabilities of neural ODEs and invertible residual networks. ICML.

This is Part 6 of the PDE and Machine Learning series. Next: Part 7 — Diffusion Models . Previous: Part 5 — Symplectic Geometry .