Neural ODE

A residual network updates its state in discrete jumps, $h_{t+1} = h_t + f(h_t, \theta_t)$. Shrink the jump and add a step size and that recurrence is Euler's method; take the step to zero and the network stops being a sequence of layers and becomes a differential equation,

$$\frac{d\mathbf{z}(t)}{dt} = f(\mathbf{z}(t), t, \theta),$$

whose solution at $t_1$ is whatever an ODE solver says it is. Depth is no longer a hyper-parameter but an accuracy request handed to the solver, and the memory the backward pass needs stops growing with it — the adjoint method reconstructs what it needs by integrating backwards rather than storing every intermediate state.

What that buys a generative model is Theorem 1, the instantaneous change of variables. A discrete flow pays for every bijection with the log-determinant of its Jacobian, which is why NICE and RealNVP go to such lengths to keep that Jacobian triangular. A continuous flow pays with a trace instead:

$$\frac{\partial \log p(\mathbf{z}(t))}{\partial t} = -\operatorname{tr}\!\left( \frac{\partial f}{\partial \mathbf{z}(t)}\right).$$

A determinant constrains the architecture; a trace does not. The vector field can be any network at all — no masking, no ordering of the dimensions, no triangular structure — and the density still comes out exactly, by integrating the state and its log-density together as one augmented system.

The paper's vector field is a sum of $M$ functions with learned time-dependent gates, $d\mathbf{z}/dt = \sum_n \sigma_n(t) f_n(\mathbf{z})$ with $\sigma_n(t) \in (0, 1)$ — a hypernetwork, so that the dynamics themselves change as the flow proceeds rather than being one fixed map applied for longer.

Computing the trace exactly costs one vector-Jacobian product per dimension, which is why the paper's experiments are two-dimensional density matching. Replacing it with the Hutchinson estimator,

$$\operatorname{tr}(A) = \mathbb{E}_{p(\boldsymbol{\epsilon})} \left[\boldsymbol{\epsilon}^{\mathsf{T}} A \boldsymbol{\epsilon}\right], \qquad \mathbb{E}[\boldsymbol{\epsilon}] = \mathbf{0},\; \operatorname{Cov}(\boldsymbol{\epsilon}) = I,$$

turns that per-dimension cost into a single product for an unbiased estimate, and is what lets the same flow run on images — the step Grathwohl et al. take in FFJORD ("FFJORD: Free-Form Continuous Dynamics for Scalable Reversible Generative Models", ICLR 2019). Both are available here; trace_method chooses, and the default follows the dimension.