Continuous Normalising Flows and Neural ODEs

A CIFAR-10 image has \(D = 3072\) dimensions. Computing the exact trace of a network's Jacobian at that size means one vector-Jacobian product per diagonal entry, so 3,072 backward passes, and a continuous flow needs that trace at every step of an ODE solver. FFJORD replaced those 3,072 passes with one, using a random-projection identity from 1989, and with unconstrained networks in a multiscale stack reached 3.40 bits/dim on CIFAR-10, beside Glow's 3.35, using by the authors' count under 2% of Glow's parameters (Grathwohl et al., 2019, FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models, ICLR, arXiv:1810.01367). Training still took six GPUs about five days. The elegance and the five days are the whole story of this model family.

The concept on normalising flows explains why a discrete flow's architecture is dictated by its Jacobian determinant. This one is about what happens when the layers become continuous.

From determinant to trace

A neural ODE defines a hidden state by \(\frac{dz(t)}{dt} = f(z(t), t; \theta)\) and computes the output \(z(t_1)\) with a black-box solver (Chen, Rubanova, Bettencourt & Duvenaud, 2018, Neural Ordinary Differential Equations, NeurIPS, arXiv:1806.07366). A residual block \(z + h\,f(z)\) is one Euler step of this ODE with step \(h\).

Take that Euler step and ask how log-density changes. The Jacobian of the step is \(I + hJ\) with \(J = \partial f / \partial z\), and for small \(h\), \(\log\det(I + hJ) = h\,\mathrm{Tr}(J) + O(h^2)\). Letting \(h \to 0\) gives the instantaneous change of variables:

so \(\log p(z(t_1)) = \log p(z(t_0)) - \int_{t_0}^{t_1} \mathrm{Tr}(\partial f / \partial z)\,dt\). The determinant, which cost \(O(D^3)\) and forced coupling layers or autoregressive masks, becomes a trace, and \(f\) no longer needs any special structure. The map is invertible automatically, because an ODE run backwards recovers its start.

Hutchinson's estimator

An exact trace is still \(O(D^2)\): \(D\) vector-Jacobian products. For any noise vector \(\epsilon\) with \(\mathbb{E}[\epsilon] = 0\) and \(\mathrm{Cov}(\epsilon) = I\),

and \(\epsilon^\top J\) is a single reverse-mode pass. FFJORD samples one Gaussian or Rademacher \(\epsilon\) per solve and holds it fixed through the integration, which keeps the dynamics deterministic for the solver while leaving the log-density estimate unbiased. Cost per function evaluation drops to \(O(D)\).

Variance is the price. FFJORD reduces it with a bottleneck trick: if \(f = g \circ h\) passes through a hidden layer of width \(H < D\), the cyclic property of the trace lets the estimator work in \(H\) dimensions instead of \(D\). Estimator variance, not bias, is what the architecture now has to manage.

The adjoint method

Backpropagating through every internal solver step stores every intermediate state. The adjoint method avoids that. Define the adjoint \(a(t) = \partial L / \partial z(t)\). It obeys its own ODE,

which a second solver call integrates backwards from \(t_1\), reconstructing \(z(t)\) alongside \(a(t)\) instead of storing it. Memory becomes constant in the number of solver steps. Chen et al. also reported that the backward pass needed roughly half the function evaluations of the forward pass.

Not everyone accepts the bargain. Gholami, Keutzer and Biros argued that reconstructing \(z(t)\) by solving backwards can be numerically unstable, and that this optimise-then-discretise gradient can disagree with the true gradient of the discretised forward pass enough to make training diverge (Gholami et al., 2019, ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs, arXiv:1902.10298). Their fix checkpoints the forward states, spending memory to get exact gradients. The disagreement is the classic numerical-analysis choice between discretise-then-optimise and optimise-then-discretise, reappearing inside deep learning.

When it breaks

Function evaluations grow as the model learns. Adaptive solvers pick their own step count, and FFJORD's authors found it tends to rise during training until it becomes prohibitive, even with constant memory. Weight decay and spectral normalisation reduce it and cost a little quality. Finlay et al. showed that penalising kinetic energy and the Jacobian's Frobenius norm pushes the dynamics toward straighter, cheaper paths without losing likelihood (Finlay et al., 2020, How to Train Your Neural ODE, ICML, arXiv:2002.02798), the same instinct that later drove rectified flow.

Stiffness. General-purpose solvers handle non-stiff dynamics efficiently. If the learned field becomes stiff, error control forces tiny steps, and stiff solvers need many more evaluations per unit of accuracy.

Trajectories cannot cross. An ODE flow is a homeomorphism, so it cannot represent even the one-dimensional map sending \(-1 \to 1\) and \(1 \to -1\). Dupont and colleagues proved this and fixed it by adding extra zero-initialised dimensions for the flow to move through (Dupont, Doucet & Teh, 2019, Augmented Neural ODEs, NeurIPS, arXiv:1904.01681). For density estimation the constraint is intrinsic: pinching a Gaussian into disconnected modes needs steep, expensive dynamics.

Simulation in the training loop is the real bottleneck. Every likelihood evaluation integrates an ODE and every gradient integrates another. Flow matching, covered separately, kept the continuous-time model and removed exactly this by regressing a velocity field directly, and continuous-time generators went on to scale once training no longer had to simulate them.