Adjoint methods for ODEs and neural ODEs

Before this lesson: lesson 01 (vector–Jacobian products, the tape, fan-out) and lesson 02 (Euler and RK4 steps, stability, stiffness). Related lessons in other tracks: calculus 09 (the adjoint of an equation $F(x, \theta) = 0$) and deep learning 05 (residual networks). Time: about 100 minutes, plus the notebook.

---

Why this matters #

Fitting a model that is a differential equation needs $dL/dp$ for one scalar loss $L$ and many parameters $p$. This happens everywhere in ML for science:

• Neural ODEs and continuous normalizing flows (Chen et al. 2018) are trained with the adjoint method or with backpropagation through the solver.
• 4D-Var data assimilation (lesson 06) finds the initial state of a weather model that best fits the observations of a time window. The gradient over millions of state variables comes from the adjoint model, run backward over the window.
• Full-waveform inversion in seismology (lesson 04) gets the gradient with respect to millions of velocity values from one forward and one adjoint wave simulation.
• Differentiable and hybrid climate models (lesson 16) are trained through multi-step rollouts of a solver.

The adjoint method costs about as much as one extra solve, whatever the number of parameters. The hard questions are practical: what to store, what to recompute, whether to run the state backward, and which gradient you get.

1. Predict first #

2. The picture #

A model that is an ODE #

A residual network (ResNet) updates a hidden state by adding a correction: $x_{k+1} = x_k + F(x_k, \theta_k)$. Write the correction as $h\,f(x_k, \theta)$ with a step size $h$ and one shared set of weights. Then

is exactly forward Euler for the ODE $u' = f(u, \theta)$ (lesson 02). The layer index $k$ plays the role of time $t = kh$. A neural ODE (Chen et al. 2018) takes the limit: the model is the ODE

where $f_\theta$ is a small neural network, and any ODE solver computes $u(T)$.

What the continuous view gives.

• Continuous time. The state exists at every $t$, so data observed at irregular times fit in directly (latent ODEs for irregular time series, Rubanova et al. 2019).
• Choice of solver. An adaptive solver uses more steps where the dynamics are fast and fewer where they are slow, at a set tolerance.
• Memory. The backsolve adjoint (section 3) computes gradients with memory that does not grow with the number of steps.
• Invertibility. The flow map $u(0) \mapsto u(T)$ is invertible (run the ODE backward), and the log-density changes by $\frac{d}{dt}\log p(u(t)) = -\operatorname{tr}\big(\partial f_\theta/\partial u\big)$. This is the basis of continuous normalizing flows, and of the probability-flow ODE of diffusion models (lesson 13).

What it costs.

• Every forward pass is a full ODE solve. The number of right-hand-side evaluations depends on the solver and the tolerance, and it often grows during training as $f_\theta$ becomes more complex.
• If $f_\theta$ learns fast-decaying modes, the ODE becomes stiff (lesson 02): explicit solvers need tiny steps.
• Trajectories of an ODE cannot cross (section 3), so some maps cannot be represented in the original dimension. Augmented neural ODEs (Dupont et al. 2019) add extra zero dimensions to fix this.

Training needs $dL/d\theta$. The rest of the lesson answers: how do we differentiate through an ODE solve? Everything works for any parameter vector $p$; for a neural ODE, $p = \theta$.

The adjoint as sensitivity, travelling backward #

Simplest case. Take $u' = au$ with a number $a$, $u(0) = u_0$, and a loss $L = \ell(u(T))$. A small push $\delta u$ given to the state at time $t$ grows to $e^{a(T - t)}\delta u$ at time $T$. So the sensitivity of $L$ to the state at time $t$ is

This number is the adjoint (or adjoint state, or costate). Read the formula from right to left: it starts at the final time, $\lambda(T) = \ell'(u(T))$, and it is known at earlier times only after we know what happens later. So it is computed backward in time. Differentiate in $t$: $\lambda' = -a\lambda$. For a system $u' = Au$ the same argument gives $\lambda(t) = e^{A^\top(T - t)}\lambda(T)$: the propagator appears transposed. This is the vector–Jacobian product of the flow map (lesson 01).

Parameters act as a forcing. A small change $\delta p$ of the parameters changes the slope $f$ by $J_p\,\delta p$ at every time, where $J_p = \partial f/\partial p$. The push $J_p\,\delta p\,dt$ given during $[t, t + dt]$ is worth $\lambda(t)^\top J_p\,\delta p\,dt$ in the loss. Adding all the pushes:

The discrete picture. Forward Euler is a ResNet whose layers share the weights $p$. Backprop through it: $\lambda_n = \partial L/\partial u_n$ is the cotangent of layer $n$, and the cotangent of the shared $p$ is a sum over all layers (fan-out, lesson 01). As $h \to 0$, the backward recursion becomes the adjoint ODE and the sum becomes the integral. The adjoint method is reverse mode in time.

Interactive figure: loading…

Three ways to get the gradient, and their cost #

There are three standard ways to get $dL/dp$ through a solver. The next section derives them; the figure compares their cost first.

Interactive figure: loading…

3. Definitions and main results #

Setting #

We solve $u' = f(u, p)$, $u(0) = u_0$ for $t \in [0, T]$, with a state $u \in \R^n$ and parameters $p \in \R^{n_p}$. The loss is $L = \ell(u(T))$. Along the trajectory, the Jacobians $J_u = \partial f/\partial u$ ($n \times n$) and $J_p = \partial f/\partial p$ ($n \times n_p$) are evaluated at $(u(t), p)$. A time stepper uses a fixed step $h$, $N$ steps and $T = Nh$.

Three ways to get $dL/dp$ #

• Central finite differences. Perturb each parameter by $\pm\epsilon$ and solve again: $2n_p$ forward solves. The error is $O(\epsilon^2)$ plus rounding noise, plus solver noise (with adaptive steps the step sequence itself changes with $p$).
• Forward sensitivities (forward mode). The matrix $S(t) = \partial u(t)/\partial p \in \R^{n\times n_p}$ solves the forward sensitivity equation (tangent equation)

and $dL/dp = S(T)^\top\nabla\ell(u(T))$. This is $n_p$ extra ODEs of size $n$ (one per column of $S$), solved together with $u$. Nothing is stored over time. Why this equation: differentiate $u' = f(u, p)$ with respect to $p$ and use the chain rule: $\frac{d}{dt}\frac{\partial u}{\partial p} = J_u\frac{\partial u}{\partial p} + J_p$.

• Adjoint (reverse mode). One backward solve of size $n$ plus an integral, for any $n_p$. It needs $u(t)$ during the backward solve.

method	work, in forward solves	memory over time	exact for the code?
central differences	$2n_p$	none	no: $O(\epsilon^2)$ + rounding
forward sensitivities	about $1 + n_p$	none	yes, if derived from the code's steps
adjoint, all states stored	about 2–3, independent of $n_p$	$N + 1$ states	yes (discrete adjoint)
adjoint with checkpoints	about 3–4	about $2\sqrt N$ states (uniform)	yes (discrete adjoint)
backsolve adjoint	about 3–4	$O(1)$	no: $O(h^q)$, can be unstable

Rule of thumb: many parameters and one loss → adjoint; few parameters → forward sensitivities.

Theorem 1 (continuous adjoint) #

Let $f$ and $\ell$ be continuously differentiable, let the solution $u(t)$ exist on $[0, T]$, and let $\lambda$ solve the adjoint ODE backward in time:

Then

Why this is true (intuition). $\lambda(t)^\top\delta u(t)$ is "the change of the loss that a perturbation $\delta u(t)$ at time $t$ will cause". If nothing pushes the trajectory between $t$ and $T$, this number must be the same at every time. The minus sign and the transpose in the adjoint ODE are exactly what keeps it constant.

Proof.

1. Perturb $u_0$ by $\delta u_0$ and $p$ by $\delta p$. To first order the trajectory changes by $\delta u(t)$, which solves the tangent equation $\delta u' = J_u\,\delta u + J_p\,\delta p$ with $\delta u(0) = \delta u_0$.
2. Differentiate the scalar $\lambda^\top\delta u$ in time, with the product rule: $(\lambda^\top\delta u)' = \lambda'^\top\delta u + \lambda^\top\delta u'$. Insert the adjoint ODE ($\lambda'^\top = -\lambda^\top J_u$) and the tangent equation: $(\lambda^\top\delta u)' = -\lambda^\top J_u\,\delta u + \lambda^\top(J_u\,\delta u + J_p\,\delta p) = \lambda^\top J_p\,\delta p$. The two $J_u$ terms cancel. This is why the adjoint ODE has $-J_u^\top$.
3. Integrate from $0$ to $T$: $\lambda(T)^\top\delta u(T) - \lambda(0)^\top\delta u_0 = \big(\int_0^T \lambda^\top J_p\,dt\big)\,\delta p$.
4. By the chain rule and the final condition, $\delta L = \nabla\ell(u(T))^\top\delta u(T) = \lambda(T)^\top\delta u(T)$.
5. So $\delta L = \lambda(0)^\top\delta u_0 + \big(\int_0^T J_p^\top\lambda\,dt\big)^\top\delta p$ for every perturbation. The coefficients of $\delta u_0$ and $\delta p$ are the two gradients. $\blacksquare$

The same result with a Lagrange multiplier

1. Add zero to the loss: $\mathcal L = \ell(u(T)) + \int_0^T \lambda^\top\big(f(u, p) - u'\big)\,dt$ for any $\lambda(t)$.
2. Perturb: $\delta\mathcal L = \nabla\ell^\top\delta u(T) + \int_0^T \lambda^\top\big(J_u\delta u + J_p\delta p - \delta u'\big)\,dt$.
3. Integrate $-\int_0^T\lambda^\top\delta u'\,dt$ by parts: it equals $-\lambda(T)^\top\delta u(T) + \lambda(0)^\top\delta u(0) + \int_0^T\lambda'^\top\delta u\,dt$.
4. Collect terms: $\delta\mathcal L = \big(\nabla\ell - \lambda(T)\big)^\top\delta u(T) + \lambda(0)^\top\delta u_0 + \int_0^T\big(\lambda' + J_u^\top\lambda\big)^\top\delta u\,dt + \int_0^T\lambda^\top J_p\,dt\,\delta p$.
5. Choose $\lambda$ so that the terms with the unknown $\delta u(t)$ and $\delta u(T)$ vanish: $\lambda' = -J_u^\top\lambda$ and $\lambda(T) = \nabla\ell(u(T))$. What remains is Theorem 1. $\blacksquare$

Extensions (each follows from the same proof).

• A running cost $L = \int_0^T g(u, p)\,dt$ adds a source: $\lambda' = -J_u^\top\lambda - \nabla_u g$, and $dL/dp$ gets the extra term $\int_0^T\nabla_p g\,dt$.
• Losses at several times $t_1 < \dots < t_K$, as in fitting observed data: in the backward sweep, add $\nabla\ell_k(u(t_k))$ to $\lambda$ each time you pass $t_k$. Between the observations, $\lambda$ follows the adjoint ODE.
• An initial state that depends on $p$ adds $(\partial u_0/\partial p)^\top\lambda(0)$ to $dL/dp$.
• The augmented backward system. Let $\mu' = -J_p^\top\lambda$ with $\mu(T) = 0$; then $\mu(0) = \int_0^T J_p^\top\lambda\,dt = dL/dp$. Solving $[u, \lambda, \mu]$ together from $T$ back to $0$, starting from $[u(T), \nabla\ell(u(T)), 0]$, re-creates $u(t)$ backward instead of storing it. This is the backsolve adjoint of the neural ODE paper (Chen et al. 2018; odeint_adjoint in torchdiffeq).

Stability of the adjoint. In reversed time $s = T - t$ the adjoint reads $d\lambda/ds = J_u^\top\lambda$, with the same eigenvalues as $J_u$. So, for a constant $J_u$, if the forward problem damps a mode, the adjoint run backward damps it too: the adjoint equation is stable backward. The re-created state is different. In reversed time it solves $du/ds = -f$, so every mode that decays forward grows backward.

Theorem 2 (discrete adjoint of a one-step method) #

A one-step method computes $u_{n+1} = \Phi(u_n, p)$ for $n = 0, \dots, N - 1$, and the code's loss is $L_h = \ell(u_N)$. Let $\Phi_u$ and $\Phi_p$ be the Jacobians of one step, and

Then

Proof (the same argument as Theorem 1, in discrete time).

1. Tangent recursion: $\delta u_{n+1} = \Phi_u\,\delta u_n + \Phi_p\,\delta p$.
2. Multiply by $\lambda_{n+1}^\top$ and use $\lambda_n^\top = \lambda_{n+1}^\top\Phi_u$: $\lambda_{n+1}^\top\delta u_{n+1} - \lambda_n^\top\delta u_n = \lambda_{n+1}^\top\Phi_p\,\delta p$.
3. Sum over $n$. The left side telescopes: $\lambda_N^\top\delta u_N - \lambda_0^\top\delta u_0 = \sum_n\lambda_{n+1}^\top\Phi_p\,\delta p$.
4. $\delta L_h = \lambda_N^\top\delta u_N$; read off the two gradients. $\blacksquare$

This is the tape sweep of lesson 01 with one node per time step, and it is what automatic differentiation computes through the solver. It is the exact gradient of the loss that the code computes, up to rounding. It stores the $N + 1$ states (or recomputes them).

Forward Euler. $\Phi(u, p) = u + hf(u, p)$, so $\Phi_u = I + hJ_u(u_n)$ and $\Phi_p = hJ_p(u_n)$:

Compare explicit Euler applied backward to the adjoint ODE: $\lambda_n = \lambda_{n+1} + h\,J_u(u_{n+1})^\top\lambda_{n+1}$. The Jacobian sits at the other end of the step. Both are first-order approximations of the same ODE, but they give different numbers.

Discrete adjoint of an explicit Runge–Kutta step

An explicit RK method with $s$ stages, coefficients $a_{ij}$ ($j < i$) and weights $b_i$ takes one step by

$$\begin{gathered} U_i = u_n + h\sum_{j<i}a_{ij}k_j, \qquad k_i = f(U_i, p), \\ u_{n+1} = u_n + h\sum_{i=1}^s b_ik_i . \end{gathered}$$

Given $\lambda_{n+1}$, go through the stages in reverse order, $i = s, \dots, 1$:

$$\begin{gathered} \bar k_i = h\Big(b_i\lambda_{n+1} + \sum_{j>i}a_{ji}\bar U_j\Big), \\ \bar U_i = J_u(U_i)^\top\bar k_i, \qquad \bar p \mathrel{+}= J_p(U_i)^\top\bar k_i . \end{gathered}$$

Then $\lambda_n = \lambda_{n+1} + \sum_{i=1}^s\bar U_i$, and $\bar p$ is this step's term of $dL_h/dp$.

Proof (reverse mode through the graph of one step).

1. $u_{n+1}$ uses $u_n$ with weight $I$ and each $k_i$ with weight $hb_i$. This sends $\lambda_{n+1}$ to $\bar u_n$ and $hb_i\lambda_{n+1}$ to $\bar k_i$.
2. $k_i$ is also used by every later stage point $U_j$ ($j > i$), with weight $ha_{ji}$: fan-out adds $ha_{ji}\bar U_j$ to $\bar k_i$. Going through the stages in reverse order makes every $\bar U_j$ final before it is needed.
3. $k_i = f(U_i, p)$: its VJP sends $J_u(U_i)^\top\bar k_i$ to $U_i$ and $J_p(U_i)^\top\bar k_i$ to $p$.
4. Every $U_i$ uses $u_n$ with weight $I$, so $\bar u_n$ also collects $\sum_i\bar U_i$. $\blacksquare$

The Jacobians are taken at the stage points $U_i$ (only $U_1 = u_n$), not at $u_{n+1}$. Cost per step: $s$ products with $J_u^\top$ and $s$ with $J_p^\top$.

What this scheme is. Write $\bar k_i = hb_i\Lambda_i$ (when $b_i \ne 0$). The recursion becomes $\Lambda_i = \lambda_{n+1} + h\sum_{j>i}(b_ja_{ji}/b_i)\,J_u(U_j)^\top\Lambda_j$ and $\lambda_n = \lambda_{n+1} + h\sum_i b_iJ_u(U_i)^\top\Lambda_i$. This is a Runge–Kutta step for the adjoint ODE, backward from $t_{n+1}$ to $t_n$, with weights $b_i$ and coefficients $b_ja_{ji}/b_i$ (Hager 2000), but with the Jacobians frozen at the forward stage points $U_i$. For classic RK4 the new coefficients, read in reverse stage order, are again the RK4 coefficients ($b_2a_{21}/b_1 = 1$, $b_3a_{32}/b_2 = \tfrac12$, $b_4a_{43}/b_3 = \tfrac12$). For many other methods they are not. When $f$ is smooth, $L_h = L + C(p)\,h^q + O(h^{q+1})$ for a method of order $q$, and differentiating this expansion shows that $dL_h/dp$ converges to the exact $dL/dp$ at the same order $q$.

When do the continuous and the discrete adjoints differ? #

The two choices have names. Discretize-then-optimize is the discrete adjoint: differentiate the code. Optimize-then-discretize is the continuous adjoint: derive the adjoint ODE, then solve it numerically.

• They agree in one clean case: a linear autonomous system $u' = Au$, an explicit RK method, and the gradient with respect to $u_0$. One step is $u_{n+1} = R(hA)u_n$ with a polynomial $R$ (the stability polynomial of lesson 02), so the discrete adjoint is $\lambda_n = R(hA)^\top\lambda_{n+1}$. The same method on $\lambda' = -A^\top\lambda$ backward with step $h$ gives $R(hA^\top)\lambda_{n+1}$, and $R(hA^\top) = R(hA)^\top$. The parameter gradient still differs, because $J_p$ depends on $u$ and the two methods evaluate it at different points.
• Finite step. In general the two differ by $O(h^q)$. The discrete adjoint is exact for $L_h$. The continuous adjoint approximates $dL/dp$, but it is not the gradient of any loss that the code computes. Near an optimum this mismatch can be as large as the gradient itself, and line searches can fail.
• Adaptive steps. The backward solve chooses its own steps, so the two time grids differ.
• Non-smooth $f$ (events, if branches, clipping). The discrete adjoint differentiates exactly the code that runs.
• Backward re-creation of $u$ (backsolve) can blow up; see the second figure below.
• Chaos is not a difference between the two. For a chaotic system both gradients grow exponentially backward in time (like $e^{0.906\,(T - t)}$ for Lorenz-63, whose largest Lyapunov exponent is about 0.906; lesson 05). Gradients over long windows are huge and not useful for either method.

Interactive figure: loading…

Why backsolve fails for damped modes. Re-creating $u$ backward means solving $du/ds = -f$ in reversed time. A mode that decays like $e^{-kt}$ forward grows like $e^{k(T - t)}$ backward. Every rounding error in that mode at time $T$ (size about $10^{-16}$ relative to the state) is multiplied by $e^{kT}$ when it reaches $t = 0$. With diffusion, damping, seismic attenuation or stiff chemistry, the re-created $u$, and every Jacobian evaluated on it, is wrong. The adjoint $\lambda$ itself stays stable.

Interactive figure: loading…

Checkpointing: memory against recomputation #

The discrete adjoint needs the states $u_n$ in reverse order. For $N$ steps:

strategy	states in memory	extra forward steps
store every state	$N + 1$	0
store only $u_0$; recompute each $u_n$ from it	1	about $N^2/2$
uniform: a checkpoint every $k = N/C$ steps	about $C + N/C$, smallest ($\approx 2\sqrt N$) at $C = \sqrt N$	fewer than $N$ (at most one more forward pass)
recursive (binomial, Revolve)	$O(\log N)$	$O(N\log N)$
backsolve (continuous adjoint)	$O(1)$	one backward re-integration of $u$; not exact

Uniform checkpointing. In the forward pass, store only $u_0, u_k, u_{2k}, \dots$. In the backward pass, take the segments from last to first: recompute the $k - 1$ missing states of the segment from its checkpoint, run the adjoint through them, and free them. The result equals full storage to rounding, because the recomputed states are the same numbers.

Why $\sqrt N$. During the backward pass of one segment we hold at most $C$ checkpoints plus the $k = N/C$ states of the segment. The function $C + N/C$ is smallest where its derivative $1 - N/C^2$ is zero, at $C = \sqrt N$, and there it equals $2\sqrt N$.

Interactive figure: loading…

Binomial checkpointing and reversible solvers

Binomial checkpointing (Griewank 1992; the Revolve algorithm of Griewank and Walther 2000) places checkpoints recursively. With $c$ checkpoints and at most $r$ recomputations of any step, it can reverse $\binom{c + r}{c}$ steps, and this number is the largest possible. So for a fixed memory the number of reversible steps grows very fast with $r$. Diffrax's default RecursiveCheckpointAdjoint is a discrete adjoint with this kind of checkpointing.

Reversible solvers are time steppers whose step can be inverted exactly in floating point (for example the reversible Heun method used by Kidger et al. 2021 for neural SDEs). They give the $O(1)$ memory of backsolve without its instability, because the backward pass recovers exactly the forward states.

Activation checkpointing in deep learning (torch.utils.checkpoint, jax.checkpoint) is the same idea for the layers of a network.

Neural ODEs: trajectories cannot cross #

Let $f_\theta$ be Lipschitz in $u$. Then the solution through any starting point is unique (Picard–Lindelöf). Two consequences:

• Trajectories never cross. If two solutions met at some time, the unique solution through that meeting point would have to be both of them.
• In one dimension the flow map is increasing. If $a < b$, then $u_a(t) < u_b(t)$ for all $t$: the two solutions are continuous, start in this order, and cannot pass through each other. So the map $u(0) \mapsto u(T)$ is a strictly increasing function. In general dimension the flow map is a homeomorphism (continuous, invertible, with a continuous inverse).

This is why a 1D neural ODE cannot learn $u \mapsto -u$, and why augmented neural ODEs add dimensions. A ResNet with finite steps can cross, because each Euler step is a jump, not a continuous path.

Answers to section 1. (1) Central differences: $2n_p = 2 \cdot 10^5$ solves. Forward sensitivities: $n_p = 10^5$ tangent ODEs of size $10^6$, so about $10^5$ solves of work and $10^{11}$ numbers in $S$. Adjoint: one forward and one backward solve (about 2–3 solves of work) for any $n_p$, plus memory for the trajectory or for checkpoints. (2) No. Automatic differentiation gives the discrete adjoint, the exact gradient of the computed loss. The RK4 adjoint ODE gives another $O(h^4)$ approximation of the exact gradient. For RK4 the two schemes have the same coefficients, but they evaluate $J_u$ and $J_p$ at different states. They agree to rounding only in special cases, such as $dL/du_0$ for a linear autonomous system. (3) No. The flow map of a 1D ODE is strictly increasing, and $u \mapsto -u$ is decreasing.

4. Worked example #

Take $u' = -pu$ with $u_0 = 1$, $p = 1$, $T = 1.5$, and the loss $\ell(u) = \tfrac12u^2$. Here $f = -pu$, so $J_u = -p$ and $J_p = -u$. Use forward Euler with $h = 0.5$, so $N = 3$ and each step multiplies the state by $1 - hp = \tfrac12$.

Step 1: forward pass. $u_0 = 1$, $u_1 = \tfrac12$, $u_2 = \tfrac14$, $u_3 = \tfrac18$. The computed loss is $L_h = \tfrac12u_3^2 = \tfrac{1}{128} = 0.0078125$.

Step 2: discrete adjoint (Theorem 2 with $\Phi_u = 1 - hp = \tfrac12$ and $\Phi_p(u_n) = -hu_n$).

• Start: $\lambda_3 = \ell'(u_3) = u_3 = \tfrac18$.
• Backward: $\lambda_2 = \tfrac12\cdot\tfrac18 = \tfrac{1}{16}$, $\lambda_1 = \tfrac{1}{32}$, $\lambda_0 = \tfrac{1}{64}$.
• Parameter gradient: $dL_h/dp = \sum_n(-hu_n)\lambda_{n+1} = -\tfrac12\big(1\cdot\tfrac{1}{32} + \tfrac12\cdot\tfrac{1}{16} + \tfrac14\cdot\tfrac18\big) = -\tfrac12\cdot\tfrac{3}{32} = -\tfrac{3}{64} = -0.046875$. The shared $p$ collects one term per step (fan-out); here each term is $-\tfrac{1}{64}$.

Step 3: check by hand. $u_3 = (1 - p/2)^3u_0$, so $L_h(p, u_0) = \tfrac12(1 - p/2)^6u_0^2$. Then $dL_h/dp = -\tfrac32(1 - p/2)^5u_0^2 = -\tfrac32\cdot\tfrac{1}{32} = -\tfrac{3}{64}$ ✓, and $dL_h/du_0 = (1 - p/2)^6u_0 = \tfrac{1}{64} = \lambda_0$ ✓. Exact for the code.

Step 4: forward sensitivity gives the same number. The discrete tangent recursion is $S_{n+1} = (1 - hp)S_n - hu_n$ with $S_0 = 0$: $S_1 = -\tfrac12$, $S_2 = \tfrac12(-\tfrac12) - \tfrac12\cdot\tfrac12 = -\tfrac12$, $S_3 = \tfrac12(-\tfrac12) - \tfrac12\cdot\tfrac14 = -\tfrac38$. Then $dL_h/dp = S_3\,u_3 = -\tfrac38\cdot\tfrac18 = -\tfrac{3}{64}$ ✓. Forward and reverse mode of the same code give the same gradient.

Step 5: backsolve (the augmented system $[u, \lambda, \mu]$ integrated backward with Euler, step $-h$, using the values at the right end of each step).

• $u$: $\hat u_n = \hat u_{n+1} - h f(\hat u_{n+1}) = (1 + hp)\,\hat u_{n+1}$, so $\hat u = \tfrac18 \to \tfrac{3}{16} \to \tfrac{9}{32} \to \tfrac{27}{64}$. The re-created $\hat u_0 \approx 0.42$, not $u_0 = 1$.
• $\lambda$: $\lambda_n = \lambda_{n+1} - h\,p\,\lambda_{n+1} = (1 - hp)\,\lambda_{n+1}$: the same values as the discrete adjoint (a linear autonomous system), so $dL/du_0 \approx \tfrac{1}{64}$ again.
• $\mu$: $\mu' = -J_p^\top\lambda = u\lambda$, so $\mu_n = \mu_{n+1} - h\,\hat u_{n+1}\lambda_{n+1}$. This gives $\mu_0 = -\tfrac12\big(\tfrac18\cdot\tfrac18 + \tfrac{3}{16}\cdot\tfrac{1}{16} + \tfrac{9}{32}\cdot\tfrac{1}{32}\big) = -\tfrac{37}{2048} \approx -0.01807$.

Step 6: exact gradient of the continuous problem. $u(t) = e^{-pt}$, $\lambda(t) = e^{p(t - T)}u(T) = e^{p(t - 2T)}$, and $dL/dp = \int_0^T(-u)\lambda\,dt = -\int_0^T e^{-2pT}\,dt = -Te^{-2pT} \approx -0.07468$; $dL/du_0 = e^{-2pT} \approx 0.04979$.

	$dL/dp$	$dL/du_0$
discrete adjoint (exact for the code)	$-0.046875$	$0.015625$
backsolve with Euler	$-0.018066$	$0.015625$
exact, continuous problem	$-0.074681$	$0.049787$

The table shows three things. The discrete adjoint is the gradient of what the code computes. The backsolve value is a third number, partly because the re-created $\hat u$ is far from the true $u$. And with $h = 0.5$ neither is close to the exact gradient; both converge as $h \to 0$ (the third figure).

5. In code #

NumPy checks the worked example, the convergence of both gradients, and the backsolve failure. The notebook has these cells and the exercises.

import numpy as np

# Worked example: u' = -p u, l = u(T)^2 / 2, forward Euler with N steps
def euler_loss(p, u0, T=1.5, N=3):
    h, u = T / N, u0
    for _ in range(N):
        u = u + h * (-p * u)
    return 0.5 * u**2

def discrete_adjoint(p, u0, T=1.5, N=3):
    h = T / N
    u = [u0]
    for _ in range(N):                      # forward pass: store every state
        u.append(u[-1] + h * (-p * u[-1]))
    lam, dp = u[N], 0.0                     # lambda_N = dl/du_N = u_N
    for n in range(N - 1, -1, -1):          # backward pass
        dp += (-h * u[n]) * lam             # Phi_p(u_n)^T lambda_{n+1}
        lam = (1 - h * p) * lam             # Phi_u^T lambda_{n+1}
    return dp, lam                          # dL_h/dp, dL_h/du0

print(euler_loss(1.0, 1.0))                 # 0.0078125
print(discrete_adjoint(1.0, 1.0))           # (-0.046875, 0.015625)
eps = 1e-6
fd = (euler_loss(1 + eps, 1.0) - euler_loss(1 - eps, 1.0)) / (2 * eps)
print(f"{fd:.9f}")                          # -0.046875000

# Both gradients converge to the exact one at order 1, but they stay different numbers
T, p = 1.5, 1.0
exact = -T * np.exp(-2 * p * T)
for N in [3, 12, 48, 192]:
    h = T / N
    uN = (1 - h * p) ** N
    disc = -T * (1 - h * p) ** (2 * N - 1)                               # discrete adjoint
    back = -h * uN**2 * sum((1 - (h * p) ** 2) ** j for j in range(N))   # Euler backsolve
    print(N, f"{abs(disc - exact):.1e} {abs(back - exact):.1e} {abs(disc - back):.1e}")
# 3 2.8e-02 5.7e-02 2.9e-02
# 12 5.1e-03 1.9e-02 1.4e-02
# 48 1.2e-03 5.1e-03 3.9e-03
# 192 2.9e-04 1.3e-03 1.0e-03

# Backsolve: re-integrating u backward is unstable for a fast-decaying mode
def rk4(f, y, h):
    k1 = f(y); k2 = f(y + h / 2 * k1); k3 = f(y + h / 2 * k2); k4 = f(y + h * k3)
    return y + h / 6 * (k1 + 2 * k2 + 2 * k3 + k4)

th = 0.3
Q = np.array([[np.cos(th), -np.sin(th)], [np.sin(th), np.cos(th)]])
A = Q @ np.diag([-1.0, -40.0]) @ Q.T        # a slow mode and a fast-decaying mode
u = np.array([1.0, 1.0])
for _ in range(1000):
    u = rk4(lambda y: A @ y, u, 1e-3)       # forward to T = 1: accurate
for _ in range(1000):
    u = rk4(lambda y: A @ y, u, -1e-3)      # backward from T to 0
print(u)                                    # about [ 3.2 -6.2], not [1. 1.]

The two closed forms in the middle part repeat section 4 with $N$ steps: the discrete adjoint gives $dL_h/dp = -T(1 - hp)^{2N - 1}$, and the Euler backsolve gives $\mu_0 = -h\,u_N^2\sum_{j=0}^{N-1}(1 - h^2p^2)^j$, because $\hat u_n\lambda_n = u_N^2(1 - h^2p^2)^{N - n}$. In the last part the forward solve is accurate. The backward solve starts from a state whose fast component is far below rounding level, and the backward flow multiplies its errors by up to $e^{40} \approx 2\cdot10^{17}$. The exact digits of the wrong answer depend on rounding, so they can differ on your machine.

6. Where it shows up #

• Neural ODEs and differentiable solvers. torchdiffeq's odeint_adjoint is the backsolve adjoint. Diffrax's default, RecursiveCheckpointAdjoint, is the discrete adjoint with checkpointing, and its documentation marks BacksolveAdjoint as not recommended. Kidger (2022, §5.1) recommends discretize-then-optimize first. Continuous normalizing flows (FFJORD, Grathwohl et al. 2019) and latent ODEs for irregularly sampled time series (Rubanova et al. 2019) are trained this way.
• Weather and ocean. 4D-Var (lesson 06) minimizes a cost over the initial state of an assimilation window. At ECMWF the gradient of this cost comes from the adjoint of a linearized, lower-resolution version of the forecast model, run backward over the window. Ocean state estimation (ECCO) uses the adjoint of the MITgcm ocean model, generated from the model code by algorithmic differentiation: a discrete adjoint of the code.
• Seismology. Adjoint tomography and full-waveform inversion (lesson 04) get the gradient of a waveform misfit from one forward and one adjoint wave simulation per earthquake, for millions of model parameters. The adjoint run needs the forward wavefield in reverse time: store it, checkpoint it, or re-create it backward. Without attenuation the wave equation is time-reversible, so the forward field can be re-created backward; with attenuation this re-creation is unstable (Komatitsch et al. 2016). This is the backsolve problem of section 3.
• Learned and hybrid weather models. Training on multi-step rollouts, as for NeuralGCM (Kochkov et al. 2024) or the autoregressive fine-tuning of an emulator, is the discrete adjoint through the rollout. Activation memory grows with the number of steps, so each step is checkpointed. Chaos limits the useful rollout length.

7. Common mistakes #

• $J_u$ instead of $J_u^\top$, a wrong sign, or starting $\lambda$ at $t = 0$. The adjoint has a final condition, $\lambda(T) = \nabla\ell(u(T))$, and runs backward with $-J_u^\top$.
• Forgetting that $p$ is shared by all steps. Its cotangent is a sum (or an integral) over all steps. Also forgetting the $(\partial u_0/\partial p)^\top\lambda(0)$ term when the initial state depends on $p$.
• Jacobians at the wrong points in the discrete adjoint. They belong at the stage points $U_i$ of the forward step. Only $U_1 = u_n$; the other stage points are neither $u_n$ nor $u_{n+1}$.
• Checking a continuous adjoint against finite differences of the code and calling the $O(h^q)$ gap a bug. Check the discrete adjoint against finite differences (to about $10^{-8}$ relative) and against automatic differentiation (to about $10^{-12}$). Check the continuous adjoint by its convergence order.
• Backsolve on a dissipative or stiff problem. The re-created state is wrong, so the gradient is wrong, even though the adjoint equation is stable.
• Finite-difference checks through an adaptive solver. The step sequence changes with $p$, which adds noise. Use fixed steps for gradient checks.
• Long chaotic windows. The exact gradient is huge and of little use for optimization; no adjoint method fixes that.

8. Check yourself #

1. For $L = \ell(u(T))$, state the adjoint ODE, its condition, and the formulas for $dL/du_0$ and $dL/dp$. Which step of the proof forces the minus sign and the transpose?
   Answer

   $\lambda' = -J_u^\top\lambda$ with $\lambda(T) = \nabla\ell(u(T))$; $dL/du_0 = \lambda(0)$ and $dL/dp = \int_0^T J_p^\top\lambda\,dt$. In $(\lambda^\top\delta u)' = \lambda'^\top\delta u + \lambda^\top(J_u\delta u + J_p\delta p)$, the $\delta u$ terms cancel only if $\lambda'^\top = -\lambda^\top J_u$, that is, $\lambda' = -J_u^\top\lambda$.
2. $u' = -pu$, $u_0 = 2$, $p = 1$, $\ell = \tfrac12u^2$, one Euler step with $h = 0.5$ ($T = 0.5$). Compute $L_h$, the discrete-adjoint $dL_h/dp$, the Euler backsolve value, and the exact $dL/dp$.
   Answer

   $u_1 = 2(1 - 0.5) = 1$ and $L_h = 0.5$. Discrete: $\lambda_1 = u_1 = 1$, $dL_h/dp = -hu_0\lambda_1 = -1$ (check: $L_h = \tfrac12u_0^2(1 - hp)^2$, derivative $-hu_0^2(1 - hp) = -1$). Backsolve: $\mu_0 = -h\,\hat u_1\lambda_1 = -0.5$. Exact: $-Tu_0^2e^{-2pT} = -2e^{-1} \approx -0.736$.
3. Write the discrete adjoint of forward Euler. Where is $J_u$ evaluated, and how does that differ from explicit Euler applied backward to the adjoint ODE?
   Answer

   $\lambda_n = \lambda_{n+1} + hJ_u(u_n)^\top\lambda_{n+1}$ and $dL_h/dp = h\sum_n J_p(u_n)^\top\lambda_{n+1}$. $J_u$ is at the start point $u_n$ of the forward step. Explicit Euler backward on $\lambda' = -J_u^\top\lambda$ uses $J_u(u_{n+1})$, the start of the backward step. Same order, different numbers.
4. A rollout has $N = 10\,000$ steps. With uniform checkpointing, what number of checkpoints minimizes memory, how many states are in memory, and what is the extra compute?
   Answer

   Memory $\approx C + N/C$, smallest at $C = \sqrt N = 100$: about $200$ states instead of $10\,001$. Each missing state is recomputed once, so the extra work is less than one forward pass ($N - C = 9\,900$ steps).
5. Why can backsolve fail for a discretized heat equation, although the adjoint equation itself is stable backward in time?
   Answer

   Re-creating $u$ backward solves the backward heat equation: a Fourier mode that decays like $e^{-\nu k^2 t}$ forward grows like $e^{\nu k^2 t}$ backward, so rounding and truncation errors in high wavenumbers explode. The adjoint $\lambda' = -J_u^\top\lambda$ with a symmetric negative definite $J_u$ decays when run backward, so it is stable. Only the re-creation of $u$ fails.
6. A model has $n = 10^7$ state variables and 3 parameters, with a loss on the final state. Adjoint or forward sensitivities?
   Answer

   Forward sensitivities: 3 tangent solves run together with the forward solve and store no trajectory. The adjoint needs only one backward solve, but it must store or recompute $N$ states of size $10^7$. With 3 parameters, forward mode is simpler and needs less memory.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. State the adjoint ODE and its condition for the loss $L = \ell(u(T))$.

   Show answer

2. Give $dL/dp$ and $dL/du_0$ from the adjoint $\lambda$.

   Show answer

3. Write the forward sensitivity equation for $S = \partial u/\partial p$.

   Show answer

4. Why is a residual network a discretized ODE?

   Show answer

5. Which gradient does automatic differentiation through a solver compute?

   Show answer

6. Why can the backsolve adjoint fail on a damped or stiff problem?

   Show answer

7. How many checkpoints minimize memory for $N$ steps, and what does it cost?

   Show answer

8. Cost of the gradient for $n_p = 10^6$ parameters: central differences versus adjoint?

   Show answer

9. Why can a one-dimensional neural ODE not learn the map $u \mapsto -u$?

   Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "How does a neural ODE compute its gradients, and what can go wrong?" — The original paper uses the backsolve adjoint with $O(1)$ memory. The gradient is approximate, and the backward re-creation of the state can be unstable. Backpropagation through the solver with checkpointing is usually preferred.
• "Discretize-then-optimize or optimize-then-discretize?" — The first is the exact gradient of the computed loss; the second is an $O(h^q)$ approximation of the true gradient. They agree only in special cases.
• "Backprop through a 1000-step simulation runs out of memory. What are your options?" — Checkpointing ($\sqrt N$ uniform or binomial), reversible solvers, forward sensitivities if there are few parameters, shorter rollouts, lower-precision stored states.
• "Derive the adjoint ODE." — The five steps of Theorem 1: tangent equation, product rule, the $J_u$ terms cancel, integrate, read off the gradients.
• "What is an adjoint model in data assimilation?" — The transpose of the tangent-linear model, run backward over the window. It gives the gradient of the 4D-Var cost with respect to the whole initial state for a small multiple of the cost of one forecast.

10. Practice #

The notebook for this lesson has five exercises, each with a check cell and a folded solution:

1. Forward sensitivity for $u' = -pu$: solve $S' = -pS - u$ together with $u$ and compare with the exact derivative.
2. The discrete adjoint of forward Euler by hand, checked against central finite differences.
3. Lorenz-63 with RK4: the discrete adjoint through the RK4 stages, the continuous (backsolve) adjoint, and the convergence order of the continuous one.
4. Uniform checkpointing: the same gradient with about $2\sqrt N$ stored states.
5. A tiny neural ODE $u' = \theta_1u + \theta_2u^2$, fitted to observations by gradient descent with the adjoint gradient.

11. Watch and read #

Lecture 6 Part 1: Adjoint Differentiation of ODE SolutionsMIT OpenCourseWare · 18.S096 Matrix Calculus (Edelman and Johnson) · 58 min

The forward and the adjoint sensitivity equations derived on the board, with their costs.

Adjoint State Method for an ODE | Adjoint Sensitivity AnalysisMachine Learning & Simulation · 43 min

A slow derivation of the adjoint ODE with the Lagrangian, as in the folded proof of Theorem 1.

Neural ODEs (NODEs) [Physics Informed Machine Learning]Steve Brunton · 25 min

Neural ODEs as models: the ResNet connection and the uses in science.

• S. G. Johnson, Notes on adjoint methods for 18.335: §2–3 (adjoints of linear and nonlinear equations) and §6 (initial-value problems).
• Bright, Edelman and Johnson, Matrix Calculus (for Machine Learning and Beyond), chapter 9: forward and reverse sensitivity analysis, discretize-then-differentiate.
• P. Kidger, On Neural Differential Equations, §5.1: discretize-then-optimize, optimize-then-discretize, reversible solvers.
• Chen, Rubanova, Bettencourt and Duvenaud, Neural Ordinary Differential Equations (2018).
• Dupont, Doucet and Teh, Augmented Neural ODEs (2019): why trajectories cannot cross.
• Blondel and Roulet, The Elements of Differentiable Programming: §12.6 (continuous adjoint against reverse mode on the discretization) and §8.5 (checkpointing).