ML

Neural ODEs & FFJORD

Continuous-depth networks via neural ordinary differential equations: the adjoint method for memory-efficient training, FFJORD's trace estimator for free-form flows, and the connection between depth and continuous dynamics.

Advanced

Prerequisites

Table of Contents

  1. Learning Objectives
  2. Notation
  3. Core Intuition
  4. From ResNets to Neural ODEs
  5. The Adjoint Method
  6. FFJORD: Free-Form Continuous Flows
  7. ODE Solvers for Neural ODEs
  8. Regularization of Neural ODEs
  9. Augmented Neural ODEs
  10. Common Pitfalls
  11. Summary
  12. Exercises

Learning Objectives

  1. Derive neural ODEs as the continuous-depth limit of ResNets.
  2. Explain the adjoint method for memory-efficient backpropagation.
  3. Derive FFJORD's training objective with Hutchinson's trace estimator.
  4. Compare fixed-step vs adaptive ODE solvers for neural ODEs.
  5. Explain why augmented neural ODEs overcome topological limitations.

Notation

  • z(t)\mathbf{z}(t) — hidden state at continuous time tt
  • fθ(z,t)f_\theta(\mathbf{z}, t) — neural network defining dynamics
  • a(t)=∂L/∂z(t)\mathbf{a}(t) = \partial\mathcal{L}/\partial\mathbf{z}(t) — adjoint state

Core Intuition

A ResNet computes zl+1=zl+fθ(zl)\mathbf{z}_{l+1} = \mathbf{z}_l + f_\theta(\mathbf{z}_l) — a discrete Euler step. Taking the step size to zero: dz/dt=fθ(z,t)d\mathbf{z}/dt = f_\theta(\mathbf{z}, t). This is a neural ODE: a network with "infinite depth" and continuous dynamics. Memory cost becomes independent of depth (via adjoint method), and the output is a diffeomorphism — making it a natural fit for normalizing flows.

Neural ODE Vector Field

t = 0.50
Time t
0.50
Vector field f(x,t)Trajectory
Explore: Neural ODEs parameterize dx/dt with a neural network. Points integrate the learned vector field; adjoint method enables memory-efficient backprop through the ODE solver.

From ResNets to Neural ODEs

ResNet (discrete): zl+1=zl+fθ(zl,l)\mathbf{z}_{l+1} = \mathbf{z}_l + f_\theta(\mathbf{z}_l, l), for l=0,1,…,L−1l = 0, 1, \ldots, L-1.

Neural ODE (continuous limit):

dz(t)dt=fθ(z(t),t),z(0)=x.(1)\frac{d\mathbf{z}(t)}{dt} = f_\theta(\mathbf{z}(t), t), \quad \mathbf{z}(0) = \mathbf{x}. \tag{1}

Output: z(1)=z(0)+∫01fθ(z(t),t) dt\mathbf{z}(1) = \mathbf{z}(0) + \int_0^1 f_\theta(\mathbf{z}(t), t)\,dt.

Properties:

  • Continuous depth: no fixed number of layers.
  • Always invertible (ODE solutions with Lipschitz ff are diffeomorphisms).
  • Adaptive computation: harder inputs use more solver steps.

The Adjoint Method

Problem: Backpropagation through LL layers needs O(L)O(L) memory (store all intermediate states). For neural ODEs with many solver steps, this is prohibitive.

Solution: Compute gradients by solving an adjoint ODE BACKWARD in time:

dadt=−aT∂f∂z,a(1)=dLdz(1).(2)\frac{d\mathbf{a}}{dt} = -\mathbf{a}^T\frac{\partial f}{\partial\mathbf{z}}, \quad \mathbf{a}(1) = \frac{d\mathcal{L}}{d\mathbf{z}(1)}. \tag{2}

Parameter gradients:

dLdθ=−∫10a(t)T∂f∂θ dt.(3)\frac{d\mathcal{L}}{d\theta} = -\int_1^0 \mathbf{a}(t)^T\frac{\partial f}{\partial\theta}\,dt. \tag{3}

Memory: O(1)O(1) in the number of solver steps (only store z(1)\mathbf{z}(1) and a(1)\mathbf{a}(1), then integrate backward).

Tradeoff: Saves memory but doubles compute (forward + backward ODE solve). Also, numerical errors in backward solve can accumulate.

FFJORD: Free-Form Continuous Flows

FFJORD (Grathwohl et al., 2019): Combine neural ODE with the instantaneous change of variables:

log⁡p1(z(1))=log⁡p0(z(0))−∫01tr(∂fθ∂z)dt.(4)\log p_1(\mathbf{z}(1)) = \log p_0(\mathbf{z}(0)) - \int_0^1 \text{tr}\left(\frac{\partial f_\theta}{\partial\mathbf{z}}\right)dt. \tag{4}

Problem: Computing tr(∂f/∂z)\text{tr}(\partial f/\partial\mathbf{z}) requires dd backprop passes (one per dimension).

Hutchinson's estimator:

tr(∂f∂z)=Ev∼N(0,I)[vT∂f∂zv].(5)\text{tr}\left(\frac{\partial f}{\partial\mathbf{z}}\right) = \mathbb{E}_{\mathbf{v} \sim \mathcal{N}(0,\mathbf{I})}\left[\mathbf{v}^T\frac{\partial f}{\partial\mathbf{z}}\mathbf{v}\right]. \tag{5}

The vector-Jacobian product vT(∂f/∂z)\mathbf{v}^T(\partial f/\partial\mathbf{z}) costs ONE backward pass. Unbiased estimate with ONE random vector.

Training: Solve augmented ODE:

ddt(zlog⁡p)=(fθ(z,t)−vT∂f∂zv).(6)\frac{d}{dt}\begin{pmatrix}\mathbf{z} \\ \log p\end{pmatrix} = \begin{pmatrix}f_\theta(\mathbf{z}, t) \\ -\mathbf{v}^T\frac{\partial f}{\partial\mathbf{z}}\mathbf{v}\end{pmatrix}. \tag{6}

ODE Solvers for Neural ODEs

Fixed-step methods:

  • Euler: zt+h=zt+hf(zt,t)\mathbf{z}_{t+h} = \mathbf{z}_t + hf(\mathbf{z}_t, t). Error: O(h2)O(h^2).
  • RK4: 4 function evaluations per step. Error: O(h5)O(h^5). Standard choice.

Adaptive methods (Dormand-Prince, dopri5):

  • Adjust step size based on local error estimate.
  • Fewer steps where dynamics are smooth; more where they're rapid.
  • Default for neural ODE training.

Training consideration: Adaptive solvers make training non-deterministic (different samples may use different numbers of steps). This can cause issues with batch normalization and learning rate scheduling.

Regularization of Neural ODEs

Problem: Unregularized neural ODEs can learn extremely complex dynamics with many solver steps (expensive at inference).

Kinetic energy regularization:

Lreg=∫01∥fθ(z(t),t)∥2dt.(7)\mathcal{L}_{\text{reg}} = \int_0^1 \|f_\theta(\mathbf{z}(t), t)\|^2 dt. \tag{7}

Penalizes high-velocity dynamics → smoother trajectories → fewer solver steps.

Jacobian regularization: Penalize ∥∂f/∂z∥F\|\partial f/\partial\mathbf{z}\|_F to encourage simple dynamics.

Effect: 2-3x reduction in NFE (number of function evaluations) with minimal quality loss.

Augmented Neural ODEs

Limitation: A neural ODE in Rd\mathbb{R}^d defines a homeomorphism — it cannot change topology. For example, it cannot map a circle to two disconnected points.

Solution (Dupont et al., 2019): Augment the state with extra dimensions:

ddt(za)=fθ((za),t),a(0)=0.(8)\frac{d}{dt}\begin{pmatrix}\mathbf{z} \\ \mathbf{a}\end{pmatrix} = f_\theta\left(\begin{pmatrix}\mathbf{z} \\ \mathbf{a}\end{pmatrix}, t\right), \quad \mathbf{a}(0) = \mathbf{0}. \tag{8}

Intuition: By lifting to higher dimensions, trajectories can "pass around" each other. Project back to original dimensions at the end.

Analogy: Like a highway overpass — cars can cross paths without collision by using the vertical dimension.

Common Pitfalls

Pitfall 1. Using the adjoint method for short-horizon ODEs. For fewer than 20 solver steps, standard backpropagation is faster (adjoint has overhead from backward ODE solve).

Pitfall 2. Not regularizing dynamics complexity. Without regularization, the model may learn dynamics requiring 200+ solver evaluations — impractical for inference.

Pitfall 3. Expecting neural ODEs to scale to high-dimensional image generation. FFJORD works well for tabular/low-dim data but struggles beyond 64x64 images. Use flow matching instead for high-res.

Summary

  • Neural ODE: Continuous-depth limit of ResNet; always invertible.
  • Adjoint method: O(1)O(1) memory backpropagation through ODE (at 2x compute cost).
  • FFJORD: Free-form continuous flow with Hutchinson trace estimator.
  • Adaptive solvers: Automatically adjust compute based on dynamics complexity.
  • Regularization: Penalize velocity to keep solver steps manageable.
  • Augmented: Extra dimensions overcome topological limitations.

Exercises

Exercise 1. For a neural ODE with d=100d=100 dimensions: compare the memory cost of standard backprop (100 solver steps) vs adjoint method.

Exercise 2. Derive the adjoint ODE (equation 2) using the Lagrangian/KKT approach.

Exercise 3. Show that a 2D neural ODE cannot separate two concentric circles (topological argument).

Exercise 4. Compute the variance of Hutchinson's trace estimator for a 10×1010\times 10 matrix with known trace. How many vectors needed for 10% relative error?

Exercise 5. Compare FFJORD vs RealNVP in BPD on 2D density estimation tasks (two moons, pinwheel). Explain which performs better and why.

Next

Autoregressive Flows: MAF & IAF