%matplotlib inline
from d2l import torch as d2l
import numpy as np
import torch28.1 Ordinary Differential Equations and Numerical Solvers
Every continuous-time generative model, from Neural ODEs and continuous normalizing flows to the deterministic samplers of diffusion and flow matching, is an ordinary differential equation that you integrate. An ODE is a velocity rule: at every point of space (and time) it tells you which way to move and how fast. A solution is the path you trace by always following the local arrow; a numerical solver is a recipe for following arrows in finite steps; a deep residual network is such a recipe, with its layers as the steps; and a density transported by the flow changes at a rate read off the field’s Jacobian. This section builds the mathematics behind that picture: what a vector field is, when a trajectory exists and is unique, how eigenvalues decide stability, and how numerical solvers trade step size for accuracy. The stochastic counterpart of the Euler method arrives in Section 28.2.3.2, the probability-flow ODE that turns a diffusion model into a deterministic flow in Section 28.3, and the solvers studied here set the step-count/quality tradeoff of every sampler in Section 28.4.
We lean on Section 23.2 (eigenvalues and eigenvectors; that section already defined the matrix exponential, and here we prove its properties), Section 24.2 (Jacobians, Taylor expansion), Section 24.3 (reverse-mode AD as vector–Jacobian products), Section 24.4 (the integrals being approximated), and Section 26.1 (the change-of-variables formula for densities). The numerical demonstrations are plain NumPy, because every solver is a handful of lines, until we train a Neural ODE below.
%matplotlib inline
from d2l import tensorflow as d2l
import numpy as np
import tensorflow as tf%matplotlib inline
from d2l import jax as d2l
import jax
from jax import numpy as jnp
import numpy as np%matplotlib inline
from d2l import mxnet as d2l
from mxnet import autograd, gluon, init, np as mnp, npx
from mxnet.gluon import nn
import numpy as np
npx.set_np()28.1.1 Vector Fields, Trajectories, and Well-Posedness
28.1.1.1 Velocity Fields and Integral Curves
An initial-value problem specifies a velocity rule and a starting point:
\[ \dot{\mathbf{x}}(t) = \mathbf{f}(\mathbf{x}(t), t), \qquad \mathbf{x}(0) = \mathbf{x}_0, \tag{28.1.1}\]
where \(\mathbf{x}(t) \in \mathbb{R}^d\) is the state, the dot is the time derivative, and \(\mathbf{f} : \mathbb{R}^d \times \mathbb{R} \to \mathbb{R}^d\) is a vector field: a function that attaches an arrow to every point. A trajectory (or integral curve) is a differentiable path whose tangent matches the field everywhere: at each instant, the particle moves with exactly the velocity the field prescribes at its current location. When \(\mathbf{f}\) does not depend on \(t\) we call the system autonomous; the arrows are frozen and only the particle moves. Figure 28.1.1 shows the picture to keep in mind: a grid of arrows, and curves threading through them, everywhere tangent.
Integrating both sides of Equation 28.1.1 from \(0\) to \(t\) (Section 24.4) gives the equivalent integral form
\[ \mathbf{x}(t) = \mathbf{x}_0 + \int_0^t \mathbf{f}(\mathbf{x}(s), s)\,ds, \tag{28.1.2}\]
which trades the derivative for an accumulation of velocities. This form does real work twice over: it is the fixed-point equation behind the existence theorem below, and it is what every numerical solver approximates; Euler’s method is the left-endpoint Riemann sum of this very integral.
Two scalar warm-ups anchor the notation. For \(\dot{x} = -x\), the rule “move toward zero at a speed proportional to your distance” is solved by \(x(t) = e^{-t} x_0\): exponential decay, the prototype of a contracting mode. For the rotational field \(\dot{\mathbf{x}} = \left(\begin{smallmatrix}0&-1\\1&0\end{smallmatrix}\right)\mathbf{x}\), the velocity is always perpendicular to the position, so \(\tfrac{d}{dt}\|\mathbf{x}\|^2 = 2\,\mathbf{x}^\top\dot{\mathbf{x}} = 0\) and trajectories are circles traversed at constant angular speed: pure rotation, the prototype of an oscillating mode. The field of Figure 28.1.1 is exactly the sum of these two behaviors, and we will see shortly that its eigenvalues \(-0.5 \pm i\) say so at a glance.
28.1.1.2 The Flow Map
Fix the field and vary the starting point: for an autonomous system, write \(\Phi_t(\mathbf{x}_0) = \mathbf{x}(t)\) for the solution at time \(t\) started at \(\mathbf{x}_0\). This flow map \(\Phi_t : \mathbb{R}^d \to \mathbb{R}^d\) moves every point of space forward \(t\) seconds along its arrows, and it composes the way time does:
\[ \Phi_0 = \mathrm{id}, \qquad \Phi_{t+s} = \Phi_t \circ \Phi_s . \]
Flowing for \(s\) seconds and then \(t\) more is the same as flowing for \(t+s\), whenever all three solutions exist. This holds because the field is frozen in time: the second leg starts from \(\Phi_s(\mathbf{x}_0)\) and follows the same rule. If solutions exist uniquely both forward and backward over the relevant interval, taking \(s=-t\) gives \(\Phi_t \circ \Phi_{-t}=\mathrm{id}\), so \(\Phi_t\) is a bijection with inverse \(\Phi_{-t}\). A field whose solutions exist for every positive and negative time therefore generates a one-parameter group of bijections. This is the seed of an invertible generative flow: push noise forward through a learned field to get data, and integrate backward to recover the noise (and, as we will see in Section 28.1.5, its likelihood). The qualifiers about existence and uniqueness matter, which is where we turn next.
28.1.1.3 Existence and Uniqueness
Before building on this, we need to know that exactly one trajectory passes through each starting point. The right hypothesis is that the field does not change too abruptly from point to point.
Proposition (Picard–Lindelöf). Let \(\mathbf{f}(\mathbf{x}, t)\) be continuous in \(t\) and Lipschitz in \(\mathbf{x}\): there is an \(L\) with
\[ \|\mathbf{f}(\mathbf{x}, t) - \mathbf{f}(\mathbf{y}, t)\| \le L\,\|\mathbf{x} - \mathbf{y}\| \qquad \textrm{for all } \mathbf{x}, \mathbf{y}, t. \]
Then the initial-value problem Equation 28.1.1 has exactly one solution on \([0, T]\), for every \(\mathbf{x}_0\) and every \(T\) (Coddington and Levinson 1955).
Proof sketch (the integral operator is a contraction). A path \(\boldsymbol{\varphi}\) solves Equation 28.1.1 iff it is a fixed point of the Picard operator built from the integral form Equation 28.1.2,
\[ (P\boldsymbol{\varphi})(t) = \mathbf{x}_0 + \int_0^t \mathbf{f}(\boldsymbol{\varphi}(s), s)\,ds . \]
For two candidate paths on a short interval \([0, \delta]\),
\[ \|(P\boldsymbol{\varphi})(t) - (P\boldsymbol{\psi})(t)\| \le \int_0^t \|\mathbf{f}(\boldsymbol{\varphi}(s), s) - \mathbf{f}(\boldsymbol{\psi}(s), s)\|\,ds \le L\,\delta \cdot \max_{s \le \delta}\|\boldsymbol{\varphi}(s) - \boldsymbol{\psi}(s)\| , \]
so once \(\delta < 1/L\) the operator \(P\) shrinks the distance between any two paths by a fixed factor: it is a contraction, and by the Banach fixed-point theorem a contraction on a complete space has exactly one fixed point (iterate \(P\) from anywhere; the iterates form a Cauchy sequence, and two distinct fixed points would have to be strictly closer to each other than they are). Completeness is what turns the Cauchy sequence into a solution: the space \(C([0, \delta])\) of continuous paths under the sup norm is complete, so the iterates converge uniformly to a continuous fixed point. That settles \([0, \delta]\); restart at \(t = \delta\) and repeat to cover all of \([0, T]\).
Iterating \(P\) from the constant path is Picard iteration; here it is once by hand. For \(\dot{x} = x\), \(x_0 = 1\): starting from \(\varphi_0 \equiv 1\),
\[ \varphi_1(t) = 1 + t, \qquad \varphi_2(t) = 1 + t + \tfrac{t^2}{2}, \qquad \varphi_3(t) = 1 + t + \tfrac{t^2}{2} + \tfrac{t^3}{6}, \;\;\ldots \]
These are the partial sums of \(e^t\): the contraction constructs the exponential, one Taylor term per sweep through the integral.
These hypotheses are a convenient sufficient package, not individually necessary conditions. Two standard counterexamples show the failures that they are designed to rule out. Uniqueness can fail without Lipschitz. The field \(\dot{x} = \sqrt{|x|}\) has slope \(\to \infty\) near \(x = 0\) (no finite \(L\) works there), and through \(x_0 = 0\) it threads infinitely many solutions: \(x(t) \equiv 0\) is one, and for every waiting time \(c \ge 0\),
\[ x_c(t) = \begin{cases} 0, & t \le c, \\ (t - c)^2/4, & t > c \end{cases} \]
is another: sit at the origin for \(c\) seconds, then leak away (check: \(\dot{x}_c = (t-c)/2 = \sqrt{x_c}\) for \(t > c\)). The solutions form a fan leaving the origin, one blade per departure time, and the initial-value problem is genuinely ambiguous (Figure 28.1.2).
Global existence fails without a growth bound. The field \(\dot{x} = x^2\) is Lipschitz on every bounded set but not globally, and from \(x_0 > 0\), separation of variables gives
\[ x(t) = \frac{x_0}{1 - x_0 t}, \]
which escapes to infinity at the finite time \(t = 1/x_0\): the faster you grow, the faster you grow, and superlinear feedback compounds to a singularity. The solution exists only locally, on \([0, 1/x_0)\), which is all that local Lipschitz continuity can promise.
For deep learning the theorem is a license. A neural field \(\mathbf{f}_\theta\) built from linear maps and Lipschitz activations (\(\tanh\), ReLU, GELU) is globally Lipschitz (a composition of globally Lipschitz maps is globally Lipschitz, with constant at most the product of the weight norms and the activations’ Lipschitz constants), which is exactly the global-\(L\) hypothesis of the proposition above. So a Neural ODE (Section 28.1.4) is well-posed: through every input there is exactly one trajectory, distinct inputs can never collide (two trajectories meeting at time \(t\) would be two solutions of the same final-value problem run backward), and the learned map \(\Phi_T\) is automatically a bijection. Invertibility comes free, as a theorem.
28.1.2 Linear ODEs and Stability
28.1.2.1 The Matrix Exponential
One family of ODEs can be solved in closed form, and it is the family that explains all the others locally: the linear system
\[ \dot{\mathbf{x}} = A\mathbf{x}, \qquad \mathbf{x}(0) = \mathbf{x}_0, \qquad A \in \mathbb{R}^{d \times d}. \]
In one dimension, \(\dot{x} = ax\) is solved by \(e^{at} x_0\), so we engineer a matrix version of the exponential to make the same formula true. Define, for any square matrix \(B\), the matrix exponential by the same power series that defines \(e^z\):
\[ e^{B} = \sum_{k=0}^{\infty} \frac{B^k}{k!} = I + B + \frac{B^2}{2} + \frac{B^3}{6} + \cdots . \tag{28.1.3}\]
The series converges for every matrix: term by term, \(\|B^k\| \le \|B\|^k\) (the operator norm is submultiplicative), so the series is dominated by the scalar series for \(e^{\|B\|}\), which converges absolutely.
Proposition (the matrix exponential solves linear ODEs). \(\mathbf{x}(t) = e^{At}\mathbf{x}_0\) is the unique solution of \(\dot{\mathbf{x}} = A\mathbf{x}\), \(\mathbf{x}(0) = \mathbf{x}_0\).
Proof. Differentiate the series Equation 28.1.3 with \(B = At\) term by term (legitimate: on any bounded \(t\)-interval the differentiated series still converges uniformly, by the same domination):
\[ \frac{d}{dt} e^{At} = \sum_{k=1}^{\infty} \frac{A^k\, t^{k-1}}{(k-1)!} = A \sum_{j=0}^{\infty} \frac{(At)^j}{j!} = A\, e^{At}. \]
So \(\mathbf{x}(t) = e^{At}\mathbf{x}_0\) satisfies the ODE, and at \(t = 0\) the series collapses to \(I\), so \(\mathbf{x}(0) = \mathbf{x}_0\). Uniqueness is Picard–Lindelöf: the field \(\mathbf{x} \mapsto A\mathbf{x}\) is Lipschitz with \(L = \|A\|\). \(\blacksquare\)
An infinite series of matrix powers sounds unwieldy, but the eigendecomposition of Section 23.2 collapses it to \(d\) scalar exponentials.
Proposition (matrix exponential by eigendecomposition). If \(A = V \Lambda V^{-1}\) with \(\Lambda = \mathrm{diag}(\lambda_1, \ldots, \lambda_d)\), then
\[ e^{At} = V e^{\Lambda t} V^{-1}, \qquad e^{\Lambda t} = \mathrm{diag}\!\left(e^{\lambda_1 t}, \ldots, e^{\lambda_d t}\right). \tag{28.1.4}\]
Proof. Powers telescope: \((At)^k = V (\Lambda t)^k V^{-1}\), because each adjacent \(V^{-1} V\) cancels. Summing the series Equation 28.1.3 and pulling \(V, V^{-1}\) out of the sum gives \(e^{At} = V\left(\sum_k (\Lambda t)^k / k!\right) V^{-1}\). The inner sum is the exponential of a diagonal matrix, which is diagonal with entries \(\sum_k (\lambda_i t)^k / k! = e^{\lambda_i t}\). \(\blacksquare\)
The formula is a decoupling statement. Expand the start point in the eigenbasis, \(\mathbf{x}_0 = \sum_i c_i \mathbf{v}_i\); then
\[ \mathbf{x}(t) = \sum_{i=1}^{d} c_i\, e^{\lambda_i t}\, \mathbf{v}_i : \]
along each eigendirection the system is the scalar equation \(\dot{x} = \lambda_i x\) we already solved, evolving independently of the others. Eigenvectors are the directions in which a linear ODE is one-dimensional; the \(e^{\lambda_i t}\) are its modes.
28.1.2.2 The Stability Dictionary
Each mode \(e^{\lambda_i t}\) with \(\lambda_i = a_i + i b_i\) has magnitude \(|e^{\lambda_i t}| = e^{a_i t}\): the real part sets growth or decay, the imaginary part sets rotation. Reading Equation 28.1.4 mode by mode gives the dictionary every dynamical argument in this chapter uses:
- \(\operatorname{Re}\lambda_i < 0\) for all \(i\): every mode decays, and \(\mathbf{x}(t) \to \mathbf{0}\) from every start: the origin is asymptotically stable, with decay rate set by the slowest mode, \(\max_i \operatorname{Re}\lambda_i\).
- \(\operatorname{Re}\lambda_i > 0\) for some \(i\): that mode grows exponentially, and almost every trajectory escapes: unstable.
- \(\operatorname{Im}\lambda_i \neq 0\): the conjugate pair contributes \(e^{a t}(\cos bt, \sin bt)\) terms: rotation, decaying or growing with the sign of \(a\).
The rotation entry is the eigenvalue–rotation dictionary of Section 23.2.1.3.3 run in continuous time. On the invariant plane of a conjugate pair \(a \pm ib\), the matrix \(A\) acts as the block \(\left(\begin{smallmatrix}a & -b\\ b & \phantom{-}a\end{smallmatrix}\right)\) of Equation 23.2.7. Write it as \(aI + bJ\) with \(J = \left(\begin{smallmatrix}0 & -1\\ 1 & \phantom{-}0\end{smallmatrix}\right)\); since \(J^2 = -I\), the exponential series splits into cosine and sine parts, \(e^{(aI + bJ)t} = e^{at}(\cos bt \, I + \sin bt \, J)\), which is scaling by \(e^{at}\) composed with rotation by \(bt\). Where the iterated matrix of Section 23.2.1.3.3 turned by a fixed angle \(\theta\) per step, the flow turns at angular velocity \(b\): this is the spiral, and it is exactly where the \(e^{at}(\cos bt, \sin bt)\) terms above come from.
In two dimensions the dictionary is a portrait gallery (Figure 28.1.3). The saddle deserves a note: with real eigenvalues of opposite signs, trajectories approach along the stable eigendirection and escape along the unstable one, the generic fate of an “unstable equilibrium.” The stable spiral is the field of Figure 28.1.1 (\(\lambda = -0.5 \pm i\): rotate while decaying); the center is our rotation warm-up.
| eigenvalues of \(A\) (\(2\times 2\)) | portrait | behavior |
|---|---|---|
| real, both negative | stable node | monotone decay |
| real, both positive | unstable node | monotone growth |
| real, opposite signs | saddle | decay in, escape out |
| complex pair, \(\operatorname{Re}\lambda < 0\) | stable spiral | decaying oscillation |
| complex pair, \(\operatorname{Re}\lambda > 0\) | unstable spiral | growing oscillation |
| purely imaginary pair | center | closed orbits |
(When \(A\) is defective, i.e. not diagonalizable (Section 23.2), the modes pick up polynomial factors \(t^k e^{\lambda t}\), but polynomials never beat exponentials, so the stability verdict still depends only on the signs of \(\operatorname{Re}\lambda_i\).)
28.1.2.3 Linearization at Fixed Points
Nonlinear fields inherit all of this locally. A fixed point of an autonomous system is a state \(\mathbf{x}^\star\) with \(\mathbf{f}(\mathbf{x}^\star) = \mathbf{0}\): the particle parked there never moves. For a small displacement \(\boldsymbol{\delta}(t) = \mathbf{x}(t) - \mathbf{x}^\star\), Taylor expansion (Section 24.2) gives
\[ \dot{\boldsymbol{\delta}} = \mathbf{f}(\mathbf{x}^\star + \boldsymbol{\delta}) = J\,\boldsymbol{\delta} + O(\|\boldsymbol{\delta}\|^2), \qquad J = \frac{\partial \mathbf{f}}{\partial \mathbf{x}}(\mathbf{x}^\star), \]
so near the fixed point the dynamics are the linear system \(\dot{\boldsymbol{\delta}} = J \boldsymbol{\delta}\), and the eigenvalues of the Jacobian at the fixed point decide local stability by the same dictionary. (This is rigorous whenever no eigenvalue sits exactly on the imaginary axis, by the Hartman–Grobman theorem (Hartman 1960; Grobman 1959); on the axis, the neglected quadratic terms get a vote.) The damped pendulum \(\ddot{\theta} = -\sin\theta - \gamma\dot{\theta}\), written as a first-order system in \((\theta, \dot{\theta})\), has Jacobian eigenvalues with negative real part at the hanging rest point \((0, 0)\) (a stable spiral, the ring-down of a released pendulum) and a saddle at the inverted balance \((\pi, 0)\), which is why you can stand a pencil on its tip only along a measure-zero set of initial conditions. The same computation, applied to a trained network’s dynamics or to the mean of the Ornstein–Uhlenbeck process (Section 28.2.4, whose deterministic skeleton is exactly the contracting mode \(\dot{x} = -\theta x\)), is the working stability tool of this whole chapter.
Let us compute all of this. The cell builds \(e^{At}\) for the spiral field of Figure 28.1.1 three independent ways (the eigendecomposition formula Equation 28.1.4, the truncated series Equation 28.1.3, and the compounding limit \((I + tA/n)^n\), forward Euler in disguise and a preview of the next section) and checks the stability dictionary’s prediction that \(\|\mathbf{x}(t)\|\) decays exactly like \(e^{-t/2}\) (the rotation part of this particular \(A\) is norm-preserving).
A = np.array([[-0.5, -1.0], [1.0, -0.5]]) # spiral sink: eigenvalues -0.5 +- i
t = 2.0
lam, V = np.linalg.eig(A) # A = V diag(lam) V^{-1}
expAt_eig = (V @ np.diag(np.exp(lam * t)) @ np.linalg.inv(V)).real
term, expAt_series = np.eye(2), np.eye(2) # partial sums of sum_k (At)^k / k!
for k in range(1, 30):
term = term @ (A * t) / k
expAt_series += term
n = 2 ** 24 # Euler compounding (I + tA/n)^n
expAt_euler = np.linalg.matrix_power(np.eye(2) + (t / n) * A, n)
x0 = np.array([2.0, 0.0])
print('eigenvalues of A:', lam.round(4))
print('|series - eigen formula| :', f'{np.abs(expAt_series - expAt_eig).max():.1e}')
print('|(I+tA/n)^n - eigen| :', f'{np.abs(expAt_euler - expAt_eig).max():.1e}')
print('||e^{At} x0|| =', f'{np.linalg.norm(expAt_eig @ x0):.6f}',
' vs e^{-t/2}||x0|| =', f'{np.exp(-0.5 * t) * np.linalg.norm(x0):.6f}')eigenvalues of A: [-0.5+1.j -0.5-1.j]
|series - eigen formula| : 5.6e-17
|(I+tA/n)^n - eigen| : 5.4e-08
||e^{At} x0|| = 0.735759 vs e^{-t/2}||x0|| = 0.735759Twenty-nine series terms already agree with the eigendecomposition formula to machine precision, the compounded Euler product lands within \(10^{-7}\), and the trajectory norm matches \(e^{-t/2}\|\mathbf{x}_0\|\) to six digits: the spectrum (\(-0.5 \pm i\)) told us the decay rate before we integrated anything.
28.1.3 Numerical Solvers: From Euler to Runge–Kutta
28.1.3.1 Forward Euler and Its Global Error
Almost no ODE beyond the linear family has a closed-form solution, so we march: replace the integral form Equation 28.1.2 over one short step by its left-endpoint approximation. With step size \(h\) and \(t_n = nh\), forward Euler is
\[ \mathbf{x}_{n+1} = \mathbf{x}_n + h\, \mathbf{f}(\mathbf{x}_n, t_n) \tag{28.1.5}\]
Take the arrow under your feet, follow it for time \(h\), look again. Equivalently, the update is the Taylor expansion of the true solution truncated after the linear term, and the truncation error of one step is the first term dropped: if \(\|\ddot{\mathbf{x}}(t)\| \le M\) along the solution, then
\[ \mathbf{x}(t_{n+1}) = \mathbf{x}(t_n) + h\,\mathbf{f}(\mathbf{x}(t_n), t_n) + \mathbf{r}_n, \qquad \|\mathbf{r}_n\| \le \tfrac{M h^2}{2}. \]
A per-step error of \(O(h^2)\) does not mean a final error of \(O(h^2)\): to reach a fixed horizon \(T\) you take \(n = T/h\) steps, and the per-step errors accumulate; worse, each step also inherits and possibly amplifies the error already made. One order of \(h\) is lost to the accumulation, and the bookkeeping is short enough to do in full.
Proposition (Euler converges at order 1). Let \(\mathbf{f}\) be \(L\)-Lipschitz in \(\mathbf{x}\), and let the solution of Equation 28.1.1 on \([0, T]\) satisfy \(\|\ddot{\mathbf{x}}(t)\| \le M\). Then the Euler iterates Equation 28.1.5 obey
\[ \max_{n \le T/h} \|\mathbf{x}_n - \mathbf{x}(t_n)\| \;\le\; \frac{M h}{2L}\left(e^{LT} - 1\right) \;=\; O(h). \]
Proof. Let \(e_n = \|\mathbf{x}_n - \mathbf{x}(t_n)\|\), and subtract the Taylor identity above from the update Equation 28.1.5:
\[ e_{n+1} \le \underbrace{\|\mathbf{x}_n - \mathbf{x}(t_n)\| + h\,\|\mathbf{f}(\mathbf{x}_n, t_n) - \mathbf{f}(\mathbf{x}(t_n), t_n)\|}_{\le\, (1 + hL)\, e_n \textrm{ by Lipschitz}} \;+\; \underbrace{\|\mathbf{r}_n\|}_{\le\, Mh^2/2} . \]
So each step multiplies the existing error by at most \((1 + hL)\) and adds at most \(Mh^2/2\). Unrolling the recursion from \(e_0 = 0\) sums a geometric series:
\[ e_n \le \frac{Mh^2}{2} \sum_{j=0}^{n-1} (1 + hL)^j = \frac{Mh^2}{2} \cdot \frac{(1 + hL)^n - 1}{hL} \le \frac{Mh}{2L}\left(e^{LT} - 1\right), \]
using \((1 + hL)^n \le e^{nhL} \le e^{LT}\). \(\blacksquare\)
The structure of the bound recurs for every solver: a local error of order \(h^{p+1}\) per step becomes a global error of order \(h^p\) after \(T/h\) steps, inflated by a stability factor \(e^{LT}\) that prices how strongly the dynamics can amplify old mistakes. We say Euler is a method of order \(p = 1\): halve the step, halve the error.
28.1.3.2 Runge–Kutta Methods
Euler’s weakness is that it commits to the slope at the start of the step, while the trajectory is already curving away. The fix is to spend a few extra field evaluations probing the slope inside the step and average them. The simplest upgrade, the midpoint method, takes half an Euler step only to measure the slope there, then takes the real step with that better slope: \(\mathbf{x}_{n+1} = \mathbf{x}_n + h\,\mathbf{f}\!\left(\mathbf{x}_n + \tfrac{h}{2}\mathbf{f}(\mathbf{x}_n)\right)\) (autonomous notation). One Taylor line shows why it helps: in one dimension,
\[ x + h f\!\left(x + \tfrac{h}{2} f(x)\right) = x + h f + \tfrac{h^2}{2} f' f + O(h^3), \]
which matches the true expansion \(x(t+h) = x + h\dot{x} + \tfrac{h^2}{2}\ddot{x} + O(h^3)\) because \(\ddot{x} = \tfrac{d}{dt} f(x(t)) = f' f\). The \(h^2\) term is now correct instead of absent: local error \(O(h^3)\), global order \(2\), at the price of two field evaluations per step. Its trapezoid twin, Heun’s method, takes a full trial Euler step and averages the slopes at its two ends, \(\mathbf{x}_{n+1} = \mathbf{x}_n + \tfrac{h}{2}\left[\mathbf{f}(\mathbf{x}_n) + \mathbf{f}\!\left(\mathbf{x}_n + h\,\mathbf{f}(\mathbf{x}_n)\right)\right]\) (same order \(2\), same two evaluations), and is the two-stage sampler of choice for the learned dynamics of Section 28.4.
Pushing the same idea to four probe slopes gives the standard general-purpose method of scientific computing, classical Runge–Kutta (RK4):
\[ \begin{aligned} \mathbf{k}_1 &= \mathbf{f}(\mathbf{x}_n,\, t_n), \\ \mathbf{k}_2 &= \mathbf{f}(\mathbf{x}_n + \tfrac{h}{2}\mathbf{k}_1,\, t_n + \tfrac{h}{2}), \\ \mathbf{k}_3 &= \mathbf{f}(\mathbf{x}_n + \tfrac{h}{2}\mathbf{k}_2,\, t_n + \tfrac{h}{2}), \\ \mathbf{k}_4 &= \mathbf{f}(\mathbf{x}_n + h\,\mathbf{k}_3,\, t_n + h), \\ \mathbf{x}_{n+1} &= \mathbf{x}_n + \tfrac{h}{6}\left(\mathbf{k}_1 + 2\mathbf{k}_2 + 2\mathbf{k}_3 + \mathbf{k}_4\right). \end{aligned} \tag{28.1.6}\]
The weights \(\tfrac16(1, 2, 2, 1)\) are Simpson’s rule (the quadrature rule \(\int_a^b g \approx \tfrac{b-a}{6}\bigl(g(a) + 4g(\tfrac{a+b}{2}) + g(b)\bigr)\)) applied to the probe slopes at the beginning, middle (twice), and end of the step, and they are chosen so that the update’s Taylor expansion matches the true solution’s through the \(h^4\) term.
Proposition (RK4 converges at order 4). For \(\mathbf{f}\) four times continuously differentiable (\(C^4\)), the RK4 iterates Equation 28.1.6 have local truncation error \(O(h^5)\) and global error \(O(h^4)\) on a fixed interval \([0, T]\).
The proof is the same accumulation argument as for Euler, fed with a long but mechanical five-term Taylor match; see Hairer et al. (1993), the standard reference, for the general order theory. What matters for practice is the scaling: halving \(h\) cuts RK4’s error by a factor of 16, so a method of order \(p\) delivers error \(\varepsilon\) at cost proportional to \(\varepsilon^{-1/p}\) field evaluations, the difference between \(10^6\) steps and \(30\) steps for the same accuracy. This is why solver order is the headline spec, and why few-step samplers for diffusion models (Section 28.4) rely on higher-order integrators. Production solvers add one more trick: embedded pairs such as Dormand–Prince RK45 compute two orders at once and use the gap between them as a free error estimate (Dormand and Prince 1980). That estimate drives an adaptive step size, shrinking \(h\) where the field bends fast and stretching it where nothing happens.
Both claims, slope \(1\) and slope \(4\), are measurable. We integrate the spiral field to \(T = 2\), where we know the exact answer from #odes-solvers-matrix-exponential, and sweep the step size across five octaves. On log–log axes, \(\textrm{error} \approx C h^p\) is a line of slope \(p\).
def euler(f, x0, T, n):
"""Forward Euler: n steps of size h = T/n."""
x, h = np.asarray(x0, dtype=float), T / n
for k in range(n):
x = x + h * f(x)
return x
def rk4(f, x0, T, n):
"""Classical Runge-Kutta: four probe slopes per step."""
x, h = np.asarray(x0, dtype=float), T / n
for k in range(n):
k1 = f(x)
k2 = f(x + 0.5 * h * k1)
k3 = f(x + 0.5 * h * k2)
k4 = f(x + h * k3)
x = x + (h / 6) * (k1 + 2 * k2 + 2 * k3 + k4)
return x
f = lambda x: A @ x # the spiral field from above
T = 2.0
x_true = expAt_eig @ x0 # exact endpoint via e^{At}
ns = 2 ** np.arange(3, 10) # 8, 16, ..., 512 steps
hs = T / ns
err_eu = np.array([np.linalg.norm(euler(f, x0, T, n) - x_true) for n in ns])
err_rk = np.array([np.linalg.norm(rk4(f, x0, T, n) - x_true) for n in ns])
print('measured slope, Euler:', f'{np.polyfit(np.log(hs), np.log(err_eu), 1)[0]:.3f}')
print('measured slope, RK4 :', f'{np.polyfit(np.log(hs), np.log(err_rk), 1)[0]:.3f}')
d2l.plot(hs, [err_eu, err_rk, err_eu[-1] * (hs / hs[-1]),
err_rk[-1] * (hs / hs[-1]) ** 4],
'step size h', 'global error at T', xscale='log', yscale='log',
legend=['Euler', 'RK4', 'slope 1', 'slope 4'])measured slope, Euler: 1.040
measured slope, RK4 : 4.022
measured slope, Euler: 1.040
measured slope, RK4 : 4.022
measured slope, Euler: 1.040
measured slope, RK4 : 4.022
measured slope, Euler: 1.040
measured slope, RK4 : 4.022
The fitted slopes land within a few percent of the theoretical orders \(1\) and \(4\), and the RK4 curve sits nearly nine decades below Euler at the smallest step (\(3.6 \times 10^{-3}\) versus \(5.0 \times 10^{-12}\)): same trajectory, same horizon, four times the work per step, nearly a billion times the accuracy.
28.1.3.3 Stiffness and Implicit Methods
Order does not decide everything. Accuracy says how well you track the solution; stability asks whether your errors quietly die out or compound into an explosion, and for explicit methods like Euler and RK4, stability puts a hard ceiling on the step size. The phenomenon is already visible in one dimension.
Proposition (stability of Euler on the test equation). On \(\dot{x} = -\lambda x\) with \(\lambda > 0\) (true solution: decay to zero): forward Euler gives \(x_n = (1 - h\lambda)^n x_0\), which decays iff \(|1 - h\lambda| < 1\), i.e.
\[ h < \frac{2}{\lambda} . \]
Backward (implicit) Euler, \(\mathbf{x}_{n+1} = \mathbf{x}_n + h\,\mathbf{f}(\mathbf{x}_{n+1}, t_{n+1})\), gives \(x_n = (1 + h\lambda)^{-n} x_0\), which decays for every \(h > 0\).
Proof. Forward Euler: \(x_{n+1} = x_n - h\lambda x_n = (1 - h\lambda)x_n\); iterate. The factor has magnitude below \(1\) iff \(-1 < 1 - h\lambda < 1\), and the left inequality is the binding one: \(h\lambda < 2\). Backward Euler: \(x_{n+1} = x_n - h\lambda x_{n+1}\), so \(x_{n+1} = x_n / (1 + h\lambda)\), and \(0 < (1 + h\lambda)^{-1} < 1\) for all \(h > 0\). \(\blacksquare\)
Past the threshold, forward Euler amplifies: each step overshoots the origin and lands farther away on the other side, an oscillating divergence with growth factor \(|1 - h\lambda|\). A method that decays on the test equation for every \(h > 0\), as backward Euler does, is called A-stable (Dahlquist 1963). (Our \(\lambda\) is real, but the full picture lives in the complex plane. Write the test equation as \(\dot{y} = \lambda y\) with complex \(\lambda\), so that a decaying, rotating mode has \(\operatorname{Re}\lambda < 0\); a method’s stability region is the set of \(z = h\lambda\) for which its update decays, and A-stability proper demands that this region contain the entire left half-plane. Figure 28.1.4 draws the regions; the computation above, \(z = -h\lambda\) with \(\lambda > 0\), is their slice along the negative real axis.) The price of implicitness is that each step defines \(\mathbf{x}_{n+1}\) only implicitly: you must solve an equation per step, a linear solve when \(\mathbf{f}\) is linear and a few Newton iterations (Section 24.2) otherwise.
Why tolerate that price? Stiffness. A system is stiff when the step size that stability forces on an explicit method is far smaller than the one accuracy would need. The standard linear source is eigenvalues spread over wildly different scales, say eigenvalues \(-50\) and \(-1\), i.e. decay rates \(\lambda_{\textrm{fast}} = 50\) and \(\lambda_{\textrm{slow}} = 1\). The fast mode dies almost immediately, and what is left to track is the slow, smooth mode, for which a large step would be perfectly accurate. But forward Euler’s stability constraint \(h < 2/\lambda_{\textrm{fast}}\) is set by the fastest decay rate, including modes that died long ago and contribute nothing to the answer. The dead mode governs your budget: that is stiffness. An implicit method deletes the constraint and lets the step size follow the physics you actually care about. The sweep below shows the forward-Euler threshold appear exactly at \(h = 2/\lambda = 0.04\), backward Euler decaying at every step size, and the two-scale system blowing up in the fast component that had already decayed to \(10^{-22}\).
lam_f = 50.0 # the test equation dx/dt = -50 x
T = 1.0
print('forward Euler is stable iff h < 2/lambda =', 2 / lam_f)
print(f'{"h":>8} {"|1-h*lam|":>10} {"forward x(T)":>14} {"backward x(T)":>14}')
for n in [10, 20, 26, 40, 100]:
h = T / n
fwd = (1 - h * lam_f) ** n # forward Euler after n steps
bwd = (1 + h * lam_f) ** (-n) # backward Euler after n steps
print(f'{h:8.3f} {abs(1 - h * lam_f):10.2f} {fwd:14.2e} {bwd:14.2e}')
print('exact x(T) = e^{-50} =', f'{np.exp(-lam_f * T):.2e}')
A_stiff = np.diag([-50.0, -1.0]) # one fast mode, one slow mode
x_stiff = euler(lambda x: A_stiff @ x, np.array([1.0, 1.0]), T, 20)
print('stiff system, forward Euler with h=0.05: x(T) =', x_stiff.round(2),
'(the long-dead fast mode explodes)')forward Euler is stable iff h < 2/lambda = 0.04
h |1-h*lam| forward x(T) backward x(T)
0.100 4.00 1.05e+06 1.65e-08
0.050 1.50 3.33e+03 1.31e-11
0.038 0.92 1.25e-01 7.73e-13
0.025 0.25 8.27e-25 8.18e-15
0.010 0.50 7.89e-31 2.46e-18
exact x(T) = e^{-50} = 1.93e-22
stiff system, forward Euler with h=0.05: x(T) = [3.32526e+03 3.60000e-01] (the long-dead fast mode explodes)Note what backward Euler’s unconditional stability does and does not buy: at \(h = 0.1\) it returns \(1.65 \times 10^{-8}\) where the truth is \(1.93 \times 10^{-22}\): stable (it decays, and further steps decay further) but not accurate. Stability prevents the explosion; only a small step or a higher order makes you right. The practical doctrine, which carries to every learned ODE in Section 28.4: use explicit adaptive solvers by default, and reach for implicit methods when the dynamics are stiff. For a nonlinear or learned field, stiff means that the Jacobian’s eigenvalues along the trajectory are spread over widely different scales.
28.1.3.4 Gradient Descent Is a Solver
Everything above already governs the oldest dynamical system in deep learning: training itself. Descending a loss \(L\) by infinitesimal steps is the gradient flow
\[ \dot{\mathbf{x}} = -\nabla L(\mathbf{x}), \]
a vector field whose fixed points are exactly the critical points of \(L\) and along whose trajectories the loss can only fall: \(\tfrac{d}{dt} L(\mathbf{x}(t)) = \nabla L^\top \dot{\mathbf{x}} = -\|\nabla L\|^2 \le 0\). Gradient descent with learning rate \(\eta\) is forward Euler on this flow (\(\mathbf{x}_{k+1} = \mathbf{x}_k - \eta\, \nabla L(\mathbf{x}_k)\) is Equation 28.1.5 with \(h = \eta\)), so the solver theory of this section is optimization theory. Section 25.1 derived the learning-rate ceiling \(\eta < 2/\lambda_{\max}(H)\) from the quadratic analysis and ran essentially the experiment below; what is new here is the reading. Near a minimum with Hessian \(H\) the dynamics linearize, mode by eigenmode, onto the test equation:
\[ \dot{\mathbf{x}} = -\nabla L(\mathbf{x}) \;\;\xrightarrow{\textrm{linearize}}\;\; \dot{\boldsymbol{\delta}} = -H\boldsymbol{\delta} \;\;\xrightarrow{\textrm{Euler}, \, h = \eta}\;\; \boldsymbol{\delta}_{k+1} = (I - \eta H)\,\boldsymbol{\delta}_k \;\;\xrightarrow{\textrm{stability}}\;\; \eta < \frac{2}{\lambda_{\max}(H)} . \]
The solver bound \(h < 2/\lambda\) is the learning-rate ceiling: a diverging learning rate is a solver instability, and an ill-conditioned Hessian is a stiff gradient flow, with the steep, already-converged curvature directions setting the stability limit under which the shallow directions must then crawl. Momentum, Polyak’s heavy-ball update, discretizes the second-order ODE \(m\, \ddot{\mathbf{x}} + \gamma\, \dot{\mathbf{x}} = -\nabla L(\mathbf{x})\) analyzed in Section 25.1. We rerun that section’s experiment, read now as the twin of the stiffness sweep, on a quadratic with \(\lambda_{\max} = 10\): predicted flip at \(\eta = 2/10 = 0.2\).
H = np.diag([10.0, 1.0]) # curvatures: lambda_max = 10
print('GD on this quadratic is stable iff eta < 2/lambda_max =', 2 / H[0, 0])
print(f'{"eta":>8} {"|1-eta*lam_max|":>16} {"||x|| after 100 GD steps":>25}')
for eta in [0.05, 0.15, 0.19, 0.201, 0.21, 0.25]:
x = np.array([1.0, 1.0])
for _ in range(100):
x = x - eta * (H @ x) # gradient descent = forward Euler
print(f'{eta:8.3f} {abs(1 - eta * H[0, 0]):16.2f} {np.linalg.norm(x):25.2e}')GD on this quadratic is stable iff eta < 2/lambda_max = 0.2
eta |1-eta*lam_max| ||x|| after 100 GD steps
0.050 0.50 5.92e-03
0.150 0.50 8.75e-08
0.190 0.90 2.66e-05
0.201 1.01 2.70e+00
0.210 1.10 1.38e+04
0.250 1.50 4.07e+17The verdict flips exactly at the predicted threshold. At \(\eta = 0.19\) the fast coordinate contracts by \(|1 - 1.9| = 0.9\) per step and a hundred steps crush the iterate to \(10^{-5}\); at \(\eta = 0.201\) the factor is \(-1.01\) and the same hundred steps have grown it, oscillating in sign, past its starting point; at \(\eta = 0.25\) it has exploded seventeen orders of magnitude. Meanwhile the slow direction (\(\lambda = 1\)) was converging at every one of these rates: the divergence is manufactured entirely by the loss surface’s stiffest mode, its own “fast dead mode.”
28.1.4 Neural ODEs and the Adjoint Method
28.1.4.1 Residual Networks Are Euler Steps
Now the bridge to deep learning, and it is one line long.
Proposition (a residual block is an Euler step). The residual update (He et al. 2016)
\[ \mathbf{x}_{l+1} = \mathbf{x}_l + \mathbf{f}_\theta(\mathbf{x}_l) \]
is exactly one forward-Euler step Equation 28.1.5 of the ODE \(\dot{\mathbf{x}} = \mathbf{f}_\theta(\mathbf{x})\) with step size \(h = 1\).
Proof. Set \(h = 1\) and \(\mathbf{f} = \mathbf{f}_\theta\) in Equation 28.1.5; the two updates are the same formula. \(\blacksquare\)
A trivial identification (Figure 28.1.5) with non-trivial consequences. Read in one direction: a ResNet with \(N\) blocks is a solver; it integrates a vector field for \(N\) unit steps of time, and “depth” is a discretization of a continuous deformation of the representation (E 2017). Read in the other direction: shrink the step while adding blocks, \(\mathbf{x}_{l+1} = \mathbf{x}_l + \tfrac{1}{N}\mathbf{f}_\theta(\mathbf{x}_l)\), and by the Euler convergence proposition (for smooth activations, so that the field is \(C^1\)) the network’s output approaches the exact flow map \(\Phi_1(\mathbf{x}_0)\) of the field; the error of the “infinitely deep” limit is the \(O(h)\) of Section 28.1.3 with \(h = 1/N\). That limit object is the Neural ODE (Chen et al. 2018):
\[ \dot{\mathbf{x}} = \mathbf{f}_\theta(\mathbf{x}, t), \qquad \textrm{output} = \mathbf{x}(T) = \mathbf{x}_0 + \int_0^T \mathbf{f}_\theta(\mathbf{x}(t), t)\,dt . \]
Layers became integration time; the architecture became a single learned vector field plus a choice of solver; and everything this section proved now applies to the model itself: a Lipschitz \(\mathbf{f}_\theta\) makes the input–output map well-posed and invertible (Section 28.1.1.3), and a fancier solver (RK4, an adaptive method) is a drop-in upgrade that changes compute, not parameters.
28.1.4.2 Training Through the Solver
How do you fit \(\theta\)? The direct route is discretize, then differentiate: unroll a fixed solver, say \(N\) Euler steps, into the computation graph, compute the loss on \(\mathbf{x}_N\), and let ordinary backpropagation flow through the unrolled steps. The “network” this builds is literally a ResNet with \(N\) blocks that share the same weights \(\theta\), so ordinary automatic differentiation already trains it.
Let us do exactly that, on a task small enough to watch: learn a planar field \(\mathbf{f}_\theta\) (one hidden layer, \(32\) tanh units) whose time-\(1\) flow carries the unit circle onto a shifted, squashed ellipse. We use \(64\) paired points, mean-squared loss on the endpoints, \(10\) unrolled Euler steps of size \(h = 0.1\), and Adam (Section 25.2) as the optimizer. The flow deforms all of \(\mathbb{R}^2\) smoothly and invertibly; we supervise it only at the \(64\) points.
theta = np.linspace(0, 2 * np.pi, 64, endpoint=False)
X0_ring = np.stack([np.cos(theta), np.sin(theta)], axis=1)
X1_ring = np.stack([1.5 + 1.2 * np.cos(theta), 0.4 * np.sin(theta)], axis=1)
torch.manual_seed(0)
X, Y = torch.tensor(X0_ring).float(), torch.tensor(X1_ring).float()
net = torch.nn.Sequential(torch.nn.Linear(2, 32), torch.nn.Tanh(),
torch.nn.Linear(32, 2))
N, h = 10, 0.1 # unrolled Euler: 10 steps to T = 1
def flow(X):
for _ in range(N): # one Euler step = one residual block
X = X + h * net(X)
return X
opt = torch.optim.Adam(net.parameters(), lr=0.05)
for it in range(401):
loss = ((flow(X) - Y) ** 2).mean()
opt.zero_grad()
loss.backward()
opt.step()
if it % 100 == 0:
print(f'iter {it:3d} loss {loss.item():.5f}')
print(f'mean endpoint error: {(flow(X) - Y).norm(dim=1).mean().item():.4f}')iter 0 loss 1.16972iter 100 loss 0.00004iter 200 loss 0.00000iter 300 loss 0.00000iter 400 loss 0.00000
mean endpoint error: 0.0011theta = np.linspace(0, 2 * np.pi, 64, endpoint=False)
X0_ring = np.stack([np.cos(theta), np.sin(theta)], axis=1)
X1_ring = np.stack([1.5 + 1.2 * np.cos(theta), 0.4 * np.sin(theta)], axis=1)
tf.random.set_seed(0)
X, Y = tf.constant(X0_ring, tf.float32), tf.constant(X1_ring, tf.float32)
net = tf.keras.Sequential([tf.keras.layers.Dense(32, activation='tanh'),
tf.keras.layers.Dense(2)])
N, h = 10, 0.1 # unrolled Euler: 10 steps to T = 1
def flow(X):
for _ in range(N): # one Euler step = one residual block
X = X + h * net(X)
return X
opt = tf.keras.optimizers.Adam(learning_rate=0.05)
@tf.function
def train_step():
with tf.GradientTape() as tape:
loss = tf.reduce_mean((flow(X) - Y) ** 2)
opt.apply_gradients(zip(tape.gradient(loss, net.trainable_variables),
net.trainable_variables))
return loss
for it in range(401):
loss = train_step()
if it % 100 == 0:
print(f'iter {it:3d} loss {float(loss):.5f}')
err = tf.reduce_mean(tf.norm(flow(X) - Y, axis=1))
print(f'mean endpoint error: {float(err):.4f}')iter 0 loss 1.27956
iter 100 loss 0.00006
iter 200 loss 0.00001iter 300 loss 0.00000
iter 400 loss 0.00000
mean endpoint error: 0.0011theta = np.linspace(0, 2 * np.pi, 64, endpoint=False)
X0_ring = np.stack([np.cos(theta), np.sin(theta)], axis=1)
X1_ring = np.stack([1.5 + 1.2 * np.cos(theta), 0.4 * np.sin(theta)], axis=1)
k1, k2 = jax.random.split(jax.random.PRNGKey(0))
params = {'W1': 0.25 * jax.random.normal(k1, (2, 32)), 'b1': jnp.zeros(32),
'W2': 0.25 * jax.random.normal(k2, (32, 2)), 'b2': jnp.zeros(2)}
X, Y = jnp.array(X0_ring), jnp.array(X1_ring)
N, h = 10, 0.1 # unrolled Euler: 10 steps to T = 1
def flow(params, X):
field = lambda X: (jnp.tanh(X @ params['W1'] + params['b1'])
@ params['W2'] + params['b2'])
return jax.lax.fori_loop(0, N, lambda i, X: X + h * field(X), X)
loss_fn = lambda params: ((flow(params, X) - Y) ** 2).mean()
@jax.jit
def adam_step(params, m, v, it): # hand-rolled Adam, lr = 0.05
loss, g = jax.value_and_grad(loss_fn)(params)
m = jax.tree.map(lambda m, g: 0.9 * m + 0.1 * g, m, g)
v = jax.tree.map(lambda v, g: 0.999 * v + 0.001 * g * g, v, g)
upd = lambda p, m, v: p - 0.05 * (m / (1 - 0.9 ** it)) / (
jnp.sqrt(v / (1 - 0.999 ** it)) + 1e-8)
return jax.tree.map(upd, params, m, v), m, v, loss
m = jax.tree.map(jnp.zeros_like, params)
v = jax.tree.map(jnp.zeros_like, params)
for it in range(1, 402):
params, m, v, loss = adam_step(params, m, v, jnp.float32(it))
if it % 100 == 1:
print(f'iter {it - 1:3d} loss {float(loss):.5f}')
err = jnp.linalg.norm(flow(params, X) - Y, axis=1).mean()
print(f'mean endpoint error: {float(err):.4f}')iter 0 loss 1.18094
iter 100 loss 0.00011
iter 200 loss 0.00002
iter 300 loss 0.00001
iter 400 loss 0.00000mean endpoint error: 0.0021theta = np.linspace(0, 2 * np.pi, 64, endpoint=False)
X0_ring = np.stack([np.cos(theta), np.sin(theta)], axis=1)
X1_ring = np.stack([1.5 + 1.2 * np.cos(theta), 0.4 * np.sin(theta)], axis=1)
npx.random.seed(0)
X, Y = mnp.array(X0_ring), mnp.array(X1_ring)
net = nn.Sequential()
net.add(nn.Dense(32, activation='tanh'), nn.Dense(2))
net.initialize(init.Xavier())
N, h = 10, 0.1 # unrolled Euler: 10 steps to T = 1
def flow(X):
for _ in range(N): # one Euler step = one residual block
X = X + h * net(X)
return X
trainer = gluon.Trainer(net.collect_params(), 'adam', {'learning_rate': 0.05})
for it in range(401):
with autograd.record():
loss = ((flow(X) - Y) ** 2).mean()
loss.backward()
trainer.step(1)
if it % 100 == 0:
print(f'iter {it:3d} loss {float(loss):.5f}')
err = mnp.sqrt(((flow(X) - Y) ** 2).sum(axis=1)).mean()
print(f'mean endpoint error: {float(err):.4f}')iter 0 loss 1.27961[14:00:22] /home/smola/mxnet/src/base.cc:48: GPU context requested, but no GPUs found.iter 100 loss 0.00003iter 200 loss 0.00001iter 300 loss 0.00000iter 400 loss 0.00000
mean endpoint error: 0.0013A few hundred full-batch Adam iterations drive the loss to \(\sim 10^{-5}\) and the mean endpoint error to \(\sim 10^{-3}\): one tiny vector field, integrated by ten shared-weight “residual blocks,” learns to carry a circle onto a displaced ellipse. And because the model is a flow, you get for free what no plain MLP gives you: integrate the same field backward (negate it, or run time from \(1\) to \(0\)) and the ellipse returns to the circle: invertibility by construction, courtesy of Picard–Lindelöf.
28.1.4.3 The Adjoint Method: Backpropagation in Continuous Time
Differentiating through the unrolled solver works, but it stores every intermediate state; for an adaptive solver taking thousands of internal steps, that is a lot of tape. The adjoint method computes the same gradients by integrating a second ODE backward in time, storing essentially nothing. It is the continuous-time limit of backpropagation, and its central object, the adjoint state \(\mathbf{a}(t) = \partial L / \partial \mathbf{x}(t)\) (“how would the loss change if the trajectory were nudged at time \(t\)?”), is the continuous analogue of the backprop delta of Section 24.3.
Proposition (the adjoint equations). Let \(\mathbf{f}\) be continuously differentiable in \(\mathbf{x}\) and \(\theta\), let \(\mathbf{x}(t)\) solve \(\dot{\mathbf{x}} = \mathbf{f}(\mathbf{x}, \theta, t)\) on \([0, T]\) from fixed \(\mathbf{x}(0)\), and let \(L = \ell(\mathbf{x}(T))\) be a differentiable loss on the endpoint. Then the adjoint state satisfies the linear ODE
\[ \dot{\mathbf{a}}(t) = -\left(\frac{\partial \mathbf{f}}{\partial \mathbf{x}}\right)^{\!\top}\! \mathbf{a}(t), \qquad \mathbf{a}(T) = \nabla \ell(\mathbf{x}(T)), \tag{28.1.7}\]
integrated backward from \(T\), and the parameter gradient is the integral
\[ \frac{\partial L}{\partial \theta} = \int_0^T \left(\frac{\partial \mathbf{f}}{\partial \theta}\right)^{\!\top}\! \mathbf{a}(t)\, dt, \tag{28.1.8}\]
with both partial derivatives evaluated along the trajectory \((\mathbf{x}(t), \theta, t)\) (Chen et al. 2018).
Proof (variational). Perturb the parameters by an infinitesimal \(\delta\theta\) and ask how the trajectory responds. Because \(\mathbf{f}\) is \(C^1\), the flow is differentiable in \(\theta\) and the first-order expansion below is rigorous: the perturbation \(\delta\mathbf{x}(t)\) obeys the linearization of the dynamics (differentiate the ODE in \(\theta\) and use the chain rule):
\[ \delta\dot{\mathbf{x}} = \frac{\partial \mathbf{f}}{\partial \mathbf{x}}\,\delta\mathbf{x} + \frac{\partial \mathbf{f}}{\partial \theta}\,\delta\theta, \qquad \delta\mathbf{x}(0) = \mathbf{0}. \]
Now let \(\mathbf{a}(t)\) be defined by the backward ODE Equation 28.1.7 and watch the pairing \(\mathbf{a}^\top \delta\mathbf{x}\) evolve; the derivative telescopes by design:
\[ \frac{d}{dt}\left(\mathbf{a}^\top \delta\mathbf{x}\right) = \dot{\mathbf{a}}^\top \delta\mathbf{x} + \mathbf{a}^\top \delta\dot{\mathbf{x}} = -\mathbf{a}^\top \frac{\partial \mathbf{f}}{\partial \mathbf{x}}\,\delta\mathbf{x} + \mathbf{a}^\top \frac{\partial \mathbf{f}}{\partial \mathbf{x}}\,\delta\mathbf{x} + \mathbf{a}^\top \frac{\partial \mathbf{f}}{\partial \theta}\,\delta\theta = \mathbf{a}^\top \frac{\partial \mathbf{f}}{\partial \theta}\,\delta\theta . \]
The state-coupling terms cancel exactly; that cancellation is why the adjoint ODE has the form it has. Integrate from \(0\) to \(T\), using \(\delta\mathbf{x}(0) = \mathbf{0}\) on the left endpoint:
\[ \mathbf{a}(T)^\top \delta\mathbf{x}(T) = \int_0^T \mathbf{a}(t)^\top \frac{\partial \mathbf{f}}{\partial \theta}\, dt \;\delta\theta . \]
The left side is \(\nabla\ell(\mathbf{x}(T))^\top \delta\mathbf{x}(T) = \delta L\), the first-order change of the loss; the right side is linear in \(\delta\theta\) with coefficient Equation 28.1.8. (The same telescoping, stopped at an intermediate time \(s\), shows \(\delta L = \mathbf{a}(s)^\top \delta\mathbf{x}(s)\) for a nudge injected at time \(s\), so the solution of Equation 28.1.7 really is \(\partial L/\partial \mathbf{x}(s)\), earning its name.) \(\blacksquare\)
One caveat connects back to Section 28.1.1.3: the existence theorem tolerates ReLU fields, but the adjoint equations as stated need a \(C^1\) field. A ReLU field is differentiable only piecewise; in practice autograd differentiates the unrolled solver piece by piece, and smooth activations (tanh, GELU) remove the issue entirely.
Look at what the backward ODE computes per step: \(-\mathbf{a}^\top \partial\mathbf{f}/\partial\mathbf{x}\) is a vector–Jacobian product, the very primitive that reverse-mode AD is made of (Section 24.3). Discretize Equation 28.1.7 with the same Euler scheme as the forward pass and you get literally the backpropagation recursion through the unrolled network: backprop is the discrete adjoint method, a lineage that runs from optimal control (Pontryagin et al. 1962) straight to the backward pass of every deep-learning framework. The continuous formulation adds one practical twist: instead of storing the forward states for the VJPs, you may re-integrate \(\mathbf{x}(t)\) backward alongside \(\mathbf{a}(t)\), making memory \(O(1)\) in the number of solver steps, at the cost of extra compute and of numerical drift when the reversed dynamics are unstable (a strongly contracting forward flow is, run backward, strongly expanding; in that regime checkpointing or plain unrolling is the sturdier choice). The literature names the two routes discretize-then-optimize (differentiate the unrolled solver, exact for the program you actually ran) versus optimize-then-discretize (discretize the continuous adjoint, memory-free but only approximately the gradient of anything) (Kidger 2022); Exercise 7 probes exactly this fault line.
Everything above is checkable on the linear ODE \(\dot{\mathbf{x}} = A\mathbf{x}\), where every ingredient has a closed form: the trajectory is \(e^{At}\mathbf{x}_0\), the adjoint of \(L = \tfrac12\|\mathbf{x}(T)\|^2\) is \(\mathbf{a}(t) = e^{A^\top (T-t)}\,\mathbf{x}(T)\) (the adjoint ODE \(\dot{\mathbf{a}} = -A^\top\mathbf{a}\) run backward), and with \(\theta\) the entries of \(A\) itself, \(\partial\mathbf{f}/\partial A\) turns Equation 28.1.8 into \(\partial L/\partial A = \int_0^T \mathbf{a}(t)\,\mathbf{x}(t)^\top dt\). The cell computes the gradient three ways: hand-written reverse mode through the unrolled Euler solver (the discrete adjoint: exactly what reverse-mode autograd builds, written out in six lines), central finite differences on the same unrolled program, and the continuous adjoint integral by quadrature.
def expm_eig(B, s):
"""e^{Bs} via the eigendecomposition formula."""
mu, W = np.linalg.eig(B)
return (W @ np.diag(np.exp(mu * s)) @ np.linalg.inv(W)).real
def loss_and_grad_unrolled(A, x0, T, n):
"""L = ||x_n||^2 / 2 through n unrolled Euler steps; dL/dA by hand-written
reverse mode -- the discrete adjoint."""
h, xs = T / n, [x0]
for k in range(n):
xs.append(xs[-1] + h * (A @ xs[-1])) # forward sweep, tape the states
a, gA = xs[-1].copy(), np.zeros_like(A) # a_n = dL/dx_n = x_n
for k in range(n - 1, -1, -1): # backward sweep: one VJP per step
gA += h * np.outer(a, xs[k]) # dL/dA += h a_{k+1} x_k^T
a = a + h * (A.T @ a) # a_k = (I + hA)^T a_{k+1}
return 0.5 * xs[-1] @ xs[-1], gA
T, n = 2.0, 1000
L, g_disc = loss_and_grad_unrolled(A, x0, T, n)
g_fd = np.zeros_like(A) # central finite differences
for i in range(2):
for j in range(2):
E = np.zeros_like(A); E[i, j] = 1e-5
g_fd[i, j] = (loss_and_grad_unrolled(A + E, x0, T, n)[0]
- loss_and_grad_unrolled(A - E, x0, T, n)[0]) / 2e-5
print('discrete adjoint vs finite differences:', f'{np.abs(g_disc - g_fd).max():.1e}')
m = 2000 # continuous adjoint, trapezoid rule
ts = np.linspace(0.0, T, m + 1)
xT = expm_eig(A, T) @ x0
F = np.stack([np.outer(expm_eig(A.T, T - s) @ xT, expm_eig(A, s) @ x0)
for s in ts])
g_cont = (T / m) * (F[1:-1].sum(axis=0) + 0.5 * (F[0] + F[-1]))
for n_steps in [100, 1000, 10000]:
gap = np.abs(loss_and_grad_unrolled(A, x0, T, n_steps)[1] - g_cont).max()
print(f'n = {n_steps:6d} Euler steps: |discrete - continuous adjoint| = {gap:.1e}')discrete adjoint vs finite differences: 1.2e-10n = 100 Euler steps: |discrete - continuous adjoint| = 2.9e-02
n = 1000 Euler steps: |discrete - continuous adjoint| = 2.8e-03n = 10000 Euler steps: |discrete - continuous adjoint| = 2.8e-04The discrete adjoint matches finite differences to \(10^{-10}\) (it is the exact gradient of the unrolled program), and as the solver is refined the discrete gradient converges to the continuous adjoint integral at exactly the solver’s order, \(O(h)\): ten times the steps, one-tenth the gap. Backprop through a solver and the adjoint method are one algorithm, seen at two resolutions.
28.1.5 Continuous Normalizing Flows
28.1.5.1 The Instantaneous Change of Variables
Push samples through a Neural ODE and you have a generative model; to train it by maximum likelihood you must know how the density changes along the flow. The discrete answer we know from Section 26.1: for an invertible map \(\mathbf{g}\) with Jacobian \(J_{\mathbf{g}}\),
\[ \log p_{\textrm{out}}(\mathbf{g}(\mathbf{x})) = \log p_{\textrm{in}}(\mathbf{x}) - \log \left|\det J_{\mathbf{g}}(\mathbf{x})\right| , \]
and a normalizing flow built from \(K\) layers pays one log-det-Jacobian per layer; the \(O(d^3)\) determinant is the reason discrete flows constrain their layers to triangular or low-rank Jacobians. In continuous time the formula simplifies: over one infinitesimal step the flow map is near-identity, and the determinant of a near-identity matrix is governed by the trace.
Lemma. For any square matrix \(J\), \(\det(I + hJ) = 1 + h \operatorname{tr} J + O(h^2)\).
Proof. In the Leibniz expansion of the determinant, any permutation other than the identity uses at least two off-diagonal entries, each of size \(O(h)\), so contributes \(O(h^2)\). The identity permutation contributes \(\prod_i (1 + h J_{ii}) = 1 + h \sum_i J_{ii} + O(h^2)\). \(\blacksquare\)
Proposition (instantaneous change of variables). Let \(\dot{\mathbf{x}} = \mathbf{f}(\mathbf{x}, t)\) with \(\mathbf{f}\) continuously differentiable and \(\mathbf{x}(0) \sim p_0\) for a differentiable density \(p_0\), and let \(p_t\) denote the density of \(\mathbf{x}(t)\), which then exists and is differentiable along trajectories. Then along each trajectory (Chen et al. 2018)
\[ \frac{d}{dt} \log p_t(\mathbf{x}(t)) = -\operatorname{tr}\!\left(\frac{\partial \mathbf{f}}{\partial \mathbf{x}}\right)\!(\mathbf{x}(t), t). \tag{28.1.9}\]
Proof. Over a short interval the flow advances by the near-identity map \(\Phi_h(\mathbf{x}) = \mathbf{x} + h\,\mathbf{f}(\mathbf{x}, t) + O(h^2)\), whose Jacobian is \(I + h J + O(h^2)\) with \(J = \partial\mathbf{f}/\partial\mathbf{x}\). Apply the discrete change-of-variables formula to this one map and expand with the Lemma, which tolerates the perturbed argument: replacing \(J\) by \(J + O(h)\) changes \(\det(I + hJ)\) only at order \(h^2\). Thus:
\[ \log p_{t+h}(\mathbf{x}(t+h)) = \log p_t(\mathbf{x}(t)) - \log\det\!\left(I + hJ + O(h^2)\right) = \log p_t(\mathbf{x}(t)) - h \operatorname{tr} J + O(h^2), \]
using \(\log(1 + u) = u + O(u^2)\) (for small \(h\) the determinant is positive, so the absolute value is moot). Subtract, divide by \(h\), and let \(h \to 0\). \(\blacksquare\)
The trace of the Jacobian is the divergence of the field: the local expansion rate of volume. Where the field spreads (\(\operatorname{tr} J > 0\)), the cloud of samples dilutes and the log-density along each trajectory falls; where it compresses, probability concentrates. Integrating Equation 28.1.9 along the trajectory turns the per-layer log-det sum of a discrete flow into a trace integral:
\[ \log p_T(\mathbf{x}(T)) = \log p_0(\mathbf{x}(0)) - \int_0^T \operatorname{tr}\!\left(\frac{\partial \mathbf{f}}{\partial \mathbf{x}}\right) dt . \tag{28.1.10}\]
This is the engine of the continuous normalizing flow (CNF): augment the state with a running log-density, integrate \((\mathbf{x}, \log p)\) together with one ODE solver call, and you have an exact likelihood for an unconstrained architecture, with no triangular Jacobians required.
28.1.5.2 The Hutchinson Trace Estimator
One cost remains: \(\operatorname{tr}(\partial\mathbf{f}/\partial\mathbf{x})\) naively needs all \(d\) diagonal entries of the Jacobian, i.e. \(d\) derivative passes. The fix is a one-line identity from randomized numerical linear algebra (Hutchinson 1989).
Proposition (Hutchinson estimator). For any matrix \(M \in \mathbb{R}^{d \times d}\) and any random \(\boldsymbol{\epsilon} \in \mathbb{R}^d\) with \(\mathbb{E}[\boldsymbol{\epsilon}] = \mathbf{0}\) and \(\mathbb{E}[\boldsymbol{\epsilon}\boldsymbol{\epsilon}^\top] = I\) (Gaussian or Rademacher, i.e. uniform \(\pm 1\) entries),
\[ \operatorname{tr}(M) = \mathbb{E}\!\left[\boldsymbol{\epsilon}^\top M \boldsymbol{\epsilon}\right]. \tag{28.1.11}\]
Proof. \(\mathbb{E}[\boldsymbol{\epsilon}^\top M \boldsymbol{\epsilon}] = \mathbb{E}[\operatorname{tr}(M \boldsymbol{\epsilon}\boldsymbol{\epsilon}^\top)] = \operatorname{tr}(M\, \mathbb{E}[\boldsymbol{\epsilon}\boldsymbol{\epsilon}^\top]) = \operatorname{tr}(M)\), using the cyclic property of the trace and its linearity (which lets it commute with the expectation). \(\blacksquare\)
The point is what the estimator touches: \(\boldsymbol{\epsilon}^\top J\) is a single vector–Jacobian product, one reverse-mode pass through \(\mathbf{f}\) with cost comparable to one evaluation/backward pass, and no full Jacobian is ever materialized (Section 24.3), followed by a dot product. An unbiased stochastic log-likelihood at the price of one extra backward pass is what makes CNFs scale; this is precisely the FFJORD recipe (Grathwohl et al. 2019).
For a linear field everything is checkable by hand: \(\dot{\mathbf{x}} = A\mathbf{x}\) has constant Jacobian \(A\), so Equation 28.1.10 says the log-density along every trajectory falls at the constant rate \(\operatorname{tr}(A)\): \(\log p_t(\mathbf{x}(t)) = \log p_0(\mathbf{x}_0) - t \operatorname{tr}(A)\). The cell flows standard-normal samples through the spiral field, where the time-\(t\) density is the Gaussian \(\mathcal{N}(\mathbf{0},\, e^{At} e^{A^\top t})\), computable in closed form, and compares; then it verifies the Hutchinson estimator on a random \(6 \times 6\) matrix.
rng = np.random.default_rng(0)
t = 1.5
X0s = rng.standard_normal((4, 2)) # four samples from p_0 = N(0, I)
Et = expm_eig(A, t)
Xts = X0s @ Et.T # flow each sample to time t
Sig = Et @ Et.T # x(t) ~ N(0, E_t E_t^T)
logp0 = -0.5 * (X0s ** 2).sum(axis=1) - np.log(2 * np.pi)
logpt = (-0.5 * np.einsum('ki,ij,kj->k', Xts, np.linalg.inv(Sig), Xts)
- 0.5 * np.log(np.linalg.det(Sig)) - np.log(2 * np.pi))
print('log p_t(x(t)), Gaussian formula :', logpt.round(6))
print('log p_0(x_0) - t tr(A), our rule :', (logp0 - t * np.trace(A)).round(6))
J = rng.standard_normal((6, 6)) # Hutchinson on a bigger matrix
eps = rng.choice([-1.0, 1.0], size=(100000, 6))
est = np.einsum('ki,ij,kj->k', eps, J, eps)
print(f'Hutchinson: tr(J) = {np.trace(J):.4f}, estimate = {est.mean():.4f}'
f' +- {est.std() / np.sqrt(len(est)):.4f}')log p_t(x(t)), Gaussian formula : [-0.354507 -0.54845 -0.546723 -1.636566]
log p_0(x_0) - t tr(A), our rule : [-0.354507 -0.54845 -0.546723 -1.636566]
Hutchinson: tr(J) = -1.7535, estimate = -1.7554 +- 0.0173The two log-density computations agree digit for digit (the trace integral is the log-det-Jacobian, evaluated the cheap way), and the Hutchinson estimate brackets the true trace within its standard error. Finally, Equation 28.1.9 is the trajectory-wise form of the continuity equation governing how a whole density field is transported; Section 28.3 starts there and turns a diffusion’s noisy paths into a deterministic probability-flow ODE integrated by exactly the solvers of this section.
28.1.6 Summary
- An ODE \(\dot{\mathbf{x}} = \mathbf{f}(\mathbf{x}, t)\) is a velocity field; a solution is everywhere tangent to it, and the flow map \(\Phi_t\) moves all of space at once, composing as \(\Phi_{t+s} = \Phi_t \circ \Phi_s\).
- Picard–Lindelöf: a globally Lipschitz field has exactly one trajectory through each point for every finite time (the Picard integral operator is a contraction). Existence and uniqueness in both time directions make its flow maps bijections, the well-posedness behind invertible generative flows. Without Lipschitz, uniqueness fails (\(\dot{x} = \sqrt{|x|}\), a fan of solutions); without a growth bound, solutions can blow up in finite time (\(\dot{x} = x^2\)).
- Linear systems are solved by the matrix exponential \(e^{At} = \sum_k (At)^k / k!\, = V e^{\Lambda t} V^{-1}\): independent modes \(e^{\lambda_i t}\) along eigendirections. Stability dictionary: real parts decide decay vs. growth, imaginary parts rotation; nonlinear fixed points inherit the verdict from the Jacobian’s eigenvalues.
- Explicit solvers march with global error \(O(h^p)\): forward Euler has order \(1\) (local \(O(h^2)\) errors, \(T/h\) of them, times an \(e^{LT}\) stability factor), RK4 order \(4\); both slopes are measurable on a log–log plot. Stiffness: forward Euler is stable on \(\dot{x} = -\lambda x\) only for \(h < 2/\lambda\), so a fast dead mode can dictate the step; implicit (backward) Euler is stable for every \(h\), at the price of solving an equation per step.
- Gradient descent is forward Euler on the gradient flow \(\dot{\mathbf{x}} = -\nabla L\): the solver stability bound becomes the learning-rate rule \(\eta < 2/\lambda_{\max}(H)\), ill-conditioning is stiffness, and momentum discretizes the heavy-ball ODE \(m\ddot{\mathbf{x}} + \gamma\dot{\mathbf{x}} = -\nabla L\).
- A residual block is one Euler step, a ResNet is a solver, and the continuous limit is a Neural ODE: a learned vector field whose flow is fit by differentiating through the solver. The adjoint method integrates \(\dot{\mathbf{a}} = -(\partial\mathbf{f}/\partial\mathbf{x})^\top\mathbf{a}\) backward (each step a VJP) and is backpropagation in continuous time, with \(O(1)\)-memory and numerical-drift tradeoffs.
- Densities transported by a flow obey the instantaneous change of variables \(\tfrac{d}{dt}\log p_t = -\operatorname{tr}(\partial\mathbf{f}/\partial\mathbf{x})\): the log-det-Jacobian becomes a trace integral, estimated unbiasedly by Hutchinson’s \(\boldsymbol{\epsilon}^\top J \boldsymbol{\epsilon}\) trick at the cost of one VJP: the mathematics of continuous normalizing flows.
28.1.7 Exercises
- Verify by differentiation that \(\mathbf{x}(t) = e^{-t}\mathbf{x}_0\) solves \(\dot{\mathbf{x}} = -\mathbf{x}\), and that trajectories of the rotational field \(\dot{\mathbf{x}} = \left(\begin{smallmatrix}0&-1\\1&0\end{smallmatrix}\right)\mathbf{x}\) keep \(\|\mathbf{x}(t)\|\) constant. Write the integral form Equation 28.1.2 for \(\dot{x} = t\) and evaluate it.
- Show that a linear field \(\mathbf{f}(\mathbf{x}) = A\mathbf{x}\) is Lipschitz with constant \(\|A\|\) (the operator norm). Then construct a solution of \(\dot{x} = \sqrt{|x|}\), \(x(0) = 0\) that waits at the origin until time \(c\) before leaving, verify it solves the ODE at every \(t\), and identify exactly where the Lipschitz hypothesis fails. Finally, compute the blow-up time of \(\dot{x} = x^2\), \(x(0) = x_0 > 0\), and explain why Picard–Lindelöf only promises a local solution here.
- From the series Equation 28.1.3, show that \(e^{At}\) commutes with \(A\) and that \(e^{A(t+s)} = e^{At}e^{As}\). Then classify the phase portraits of \(\left(\begin{smallmatrix}-2&0\\0&-1\end{smallmatrix}\right)\), \(\left(\begin{smallmatrix}1&0\\0&-1\end{smallmatrix}\right)\), and \(\left(\begin{smallmatrix}0&-2\\2&0\end{smallmatrix}\right)\) using the stability dictionary.
- Derive Euler’s local truncation error \(O(h^2)\) from Taylor’s theorem with remainder, and rehearse the accumulation argument that turns it into a global \(O(h)\). Then show that the midpoint method’s update matches the true Taylor expansion through \(h^2\) in the vector case, where the scalar product \(f'f\) becomes the Jacobian–vector product \((\partial\mathbf{f}/\partial\mathbf{x})\mathbf{f}\).
- Re-derive the forward-Euler stability bound \(h < 2/\lambda\) and backward Euler’s unconditional stability. For the stiff matrix \(A = \mathrm{diag}(-100, -1)\): what step size does forward Euler need for stability, how many steps to reach \(T = 5\), and how do both change if the fast eigenvalue moves to \(-10^4\)? For nonlinear \(\mathbf{f}\), write the equation one backward-Euler step must solve and the Newton iteration you would use.
- A residual block computes \(\mathbf{x}_{l+1} = \mathbf{x}_l + \mathbf{f}_\theta(\mathbf{x}_l)\). Identify the implied step size; explain what halving the solver step while doubling the step count corresponds to architecturally; and use Section 28.1.1.3 to argue that the time-\(T\) flow of a Lipschitz \(\mathbf{f}_\theta\) is invertible; then explain why a plain (non-residual) layer \(\mathbf{x} \mapsto \sigma(W\mathbf{x})\) has no such guarantee.
- Re-derive the adjoint equations Equation 28.1.7 and Equation 28.1.8 from the variational argument without looking. Show that one Euler step of the adjoint ODE is exactly one vector–Jacobian product, and that discretizing the adjoint ODE with Euler reproduces the backpropagation recursion of the
#odes-solvers-adjoint-checkcell. When the forward dynamics contract strongly (all \(\operatorname{Re}\lambda \ll 0\)), what goes wrong with reconstructing \(\mathbf{x}(t)\) by integrating backward, and why does the \(O(1)\)-memory adjoint suffer where unrolled backprop does not? - Derive the instantaneous change of variables Equation 28.1.9 from the discrete log-det formula as \(h \to 0\), and check that for a single linear “layer” \(\Phi_t = e^{At}\) it reproduces \(-\log|\det e^{At}| = -t\operatorname{tr}(A)\) (use \(\det e^{At} = e^{t \operatorname{tr} A}\), e.g. via Equation 28.1.4). Prove the Hutchinson estimator Equation 28.1.11 is unbiased, compute its variance for Rademacher probes (express it in terms of the off-diagonal entries of \(M + M^\top\)), and explain why estimating the trace costs one VJP while the determinant has no comparably cheap unbiased estimator.