3.4 Differentiating through solvers

3.4 Differentiating through solvers#

Parts 1 and 2 used hand-written explicit time steps, chosen small enough that the integration scheme never asserted itself. Real models are not so polite: their physics spans timescales from seconds to millennia, and their solvers — adaptive, implicit, stiff-aware — are sophisticated numerical objects. Differentiating through a solver raises questions the solver’s forward accuracy does not answer: the derivative of a numerical solution is its own numerical problem, with its own stability theory. This chapter demonstrates the two failure modes that matter in practice, and the library (diffrax [Kidger, 2022]) that packages the solutions.

Learning goals

  • Recognize stiffness and why explicit integrators fail on it

  • Use diffrax to solve and differentiate ODEs with explicit and implicit methods

  • Demonstrate that the continuous (‘backsolve’) adjoint is unstable for dissipative systems

  • Choose an adjoint method deliberately, with a working decision rule

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

jax.config.update("jax_enable_x64", True)   # climate work needs double precision
# Running on Colab? First run:  %pip install -q optax equinox diffrax

A stiff climate problem in two boxes#

The minimal model with a real timescale separation: a thin ocean surface layer (heat capacity of ~2 m of water, responding within weeks) exchanging heat with a deep reservoir (~500 m, responding over decades), under a constant 4 W m⁻² forcing — the classic two-box structure used to emulate the transient climate response:

\[C_s\,\dot{T}_s = F - B\,T_s - \kappa\,(T_s - T_d), \qquad C_d\,\dot{T}_d = \kappa\,(T_s - T_d).\]
C_SURF, C_DEEP = 8.2e6, 2.05e9      # J m^-2 K^-1  (~2 m and ~500 m of water)
B_FEEDBACK, KAPPA = 2.0, 1.0        # W m^-2 K^-1
FORCING = 4.0                       # W m^-2
YEAR = 365.25 * 86400.0

def tendency(t, T, kappa):
    Ts, Td = T
    return jnp.array([(FORCING - B_FEEDBACK * Ts - kappa * (Ts - Td)) / C_SURF,
                      kappa * (Ts - Td) / C_DEEP])

tau_fast = C_SURF / (B_FEEDBACK + KAPPA) / 86400
print(f"surface timescale ~ {tau_fast:.0f} days;  deep timescale ~ {C_DEEP/KAPPA/YEAR:.0f} years")
print(f"stiffness ratio   ~ {C_DEEP / KAPPA / (C_SURF / (B_FEEDBACK + KAPPA)):.0f}")

surface timescale ~ 32 days;  deep timescale ~ 65 years
stiffness ratio   ~ 750

Stiffness is exactly this: the fastest process limits the time step, the slowest sets the simulation length, and their ratio (here ~750; in a real ESM, far larger) multiplies the cost. An explicit scheme must resolve the fast process even after it has equilibrated and become boring — exceed its stability limit and the solution does not degrade gracefully, it detonates:

def explicit_euler(kappa, dt_days, years=50):
    dt_s = dt_days * 86400
    def step(T, _):
        return T + dt_s * tendency(0.0, T, kappa), None
    T, _ = jax.lax.scan(step, jnp.zeros(2), None, length=int(years * YEAR / dt_s))
    return T

for dt_days in [30, 100]:
    T = explicit_euler(KAPPA, dt_days)
    print(f"explicit Euler, dt = {dt_days:3d} days: T_surface = {float(T[0]):11.3e} K"
          f"{'   <- stable (limit ~63 days)' if abs(T[0]) < 100 else '   <- unstable'}")

explicit Euler, dt =  30 days: T_surface =   1.600e+00 K   <- stable (limit ~63 days)
explicit Euler, dt = 100 days: T_surface =  -1.208e+61 K   <- unstable

Solvers as library calls: diffrax#

diffrax provides differentiable ODE/SDE solvers for JAX — adaptive step control, explicit and implicit methods, and several adjoint strategies, all composable with grad. The same 50-year problem, solved by an adaptive explicit method (Tsit5, a Runge-Kutta workhorse) and an implicit one built for stiffness (Kvaerno5, which solves a nonlinear system each step and is differentiated through that solve by the implicit function theorem — Chapter 3.2’s trick, applied per step):

import diffrax

def solve(kappa, solver, adjoint=diffrax.RecursiveCheckpointAdjoint(),
          rtol=1e-8, atol=1e-10):
    sol = diffrax.diffeqsolve(
        diffrax.ODETerm(tendency), solver, t0=0.0, t1=50 * YEAR, dt0=86400.0,
        y0=jnp.zeros(2), args=kappa,
        stepsize_controller=diffrax.PIDController(rtol=rtol, atol=atol),
        adjoint=adjoint, max_steps=1_000_000)
    return sol.ys[0], int(sol.stats["num_steps"])

for name, solver in [("Tsit5 (explicit)", diffrax.Tsit5()),
                     ("Kvaerno5 (implicit)", diffrax.Kvaerno5())]:
    (Ts, Td), n = solve(KAPPA, solver)
    print(f"{name:20s}: T = ({float(Ts):.4f}, {float(Td):.4f}) K   in {n:4d} steps")

Tsit5 (explicit)    : T = (1.6001, 0.8014) K   in  255 steps

Kvaerno5 (implicit) : T = (1.6001, 0.8014) K   in   53 steps

Identical answers; the implicit method takes 5× fewer steps on this mildly stiff toy, and the advantage grows with the stiffness ratio. For the boundary-layer, microphysics, and chemistry timescales inside real ESMs, implicit and semi-implicit methods are not an optimization but the entry fee — which is why “differentiate through an implicit solve” (rather than around it) is core infrastructure for differentiable climate models.

Choosing the adjoint: a consequential checkbox#

Differentiating an ODE solution admits two philosophies. Discretize-then- differentiate applies reverse-mode AD to the solver’s actual arithmetic — the exact gradient of the computed solution, with the memory story of Chapter 3.2 (RecursiveCheckpointAdjoint, diffrax’s default, uses the optimal checkpointing schedule). Differentiate-then-discretize (BacksolveAdjoint, the “neural ODE” adjoint) instead derives a continuous adjoint ODE and integrates it backward in time alongside a re-integration of the state, using only \(O(1)\) memory.

The backsolve looks like a free lunch, and for the dissipative systems of climate science it is a trap with a clean mathematical cause: a system that damps perturbations forward in time amplifies them backward. Re-integrating our stable two-box model in reverse is exponentially unstable, and the gradient does not degrade subtly:

def grad_of_deep_T(adjoint, rtol, atol):
    def deep_T(kappa):
        (Ts, Td), _ = solve(kappa, diffrax.Kvaerno5(), adjoint, rtol, atol)
        return Td
    return float(jax.grad(deep_T)(KAPPA))

g_ref = grad_of_deep_T(diffrax.RecursiveCheckpointAdjoint(), 1e-10, 1e-12)
print(f"reference dT_deep/dkappa = {g_ref:.6f}  K per (W m^-2 K^-1)")
print()
print("tolerance   checkpointed adjoint      backsolve adjoint")
for rtol, atol in [(1e-8, 1e-10), (1e-4, 1e-6)]:
    gc = grad_of_deep_T(diffrax.RecursiveCheckpointAdjoint(), rtol, atol)
    gb = grad_of_deep_T(diffrax.BacksolveAdjoint(), rtol, atol)
    print(f"  {rtol:g}     error {abs(gc - g_ref):10.2e}       error {abs(gb - g_ref):10.2e}")

reference dT_deep/dkappa = 0.409117  K per (W m^-2 K^-1)

tolerance   checkpointed adjoint      backsolve adjoint

  1e-08     error   6.91e-09       error  2.74e+232

  0.0001     error   6.24e-07       error  3.62e+238

The checkpointed adjoint is accurate to solver tolerance; the backsolve gradient is wrong by hundreds of orders of magnitude — not noisy, destroyed, at any tolerance. The failure is silent: the forward solution from the same call is perfect. This checkbox has burned enough practitioners that it deserves stating as a rule: for dissipative (climate-like) dynamics, use the checkpointed/discrete adjoint; reserve the backsolve for the non-stiff, non-dissipative problems it was designed for.

One more cell connects this chapter to 3.1. On a chaotic system, both adjoints agree — and both are equally useless at long windows, because the pathology is the question, not the integrator:

K_L96 = 40
def l96(t, u, F):
    return (jnp.roll(u, -1) - jnp.roll(u, 2)) * jnp.roll(u, 1) - u + F

u0 = 8.0 + 0.5 * jax.random.normal(jax.random.key(0), (K_L96,))
term = diffrax.ODETerm(l96)
ctrl = diffrax.PIDController(rtol=1e-8, atol=1e-8)
u_spun = diffrax.diffeqsolve(term, diffrax.Tsit5(), 0.0, 20.0, 0.01, u0, args=8.0,
                             stepsize_controller=ctrl, max_steps=100_000).ys[0]

def mean_state(F, adjoint, t_end):
    sol = diffrax.diffeqsolve(term, diffrax.Tsit5(), 0.0, t_end, 0.01, u_spun, args=F,
                              stepsize_controller=ctrl, adjoint=adjoint, max_steps=100_000,
                              saveat=diffrax.SaveAt(ts=jnp.linspace(0.0, t_end, 200)))
    return jnp.mean(sol.ys)

print("window   checkpointed      backsolve")
for t_end in [2.0, 5.0, 10.0]:
    gc = float(jax.grad(mean_state)(8.0, diffrax.RecursiveCheckpointAdjoint(), t_end))
    gb = float(jax.grad(mean_state)(8.0, diffrax.BacksolveAdjoint(), t_end))
    print(f"  {t_end:4.0f}   {gc:+12.4f}   {gb:+12.4f}")

window   checkpointed      backsolve

     2        +0.0840        +0.0840

     5        +3.0780        +3.0780

    10    -15110.5301    -15110.5304

At window 10 the two methods agree to five digits — on a gradient four orders of magnitude too large to use. Integrator choices cannot repair an ill-posed question; they can only avoid adding numerical damage to a well-posed one.

Decision guide#

Situation

Approach

Non-stiff, short rollout

hand-rolled step + scan, plain grad (Parts 1–2)

Long rollout, memory-bound

checkpointing (jax.checkpoint, or diffrax’s default adjoint)

Stiff dynamics

implicit solver (Kvaerno/ImplicitEuler), gradients via its built-in implicit differentiation

Output is an equilibrium

implicit-function-theorem gradient at the fixed point (Chapter 3.2)

Dissipative dynamics

never BacksolveAdjoint

Chaotic dynamics, long window

not a solver problem — reformulate the question (Chapter 3.1)

Exercises.

  1. Sweep the backsolve gradient’s error as the integration length grows from 1 to 50 years. Estimate the instability’s e-folding time and compare it to \(\tau_{fast}\).

  2. Solve the two-box model with diffrax.ImplicitEuler at fixed 1-year steps. How accurate are the solution and its \(\kappa\)-gradient compared to Kvaerno5?

  3. Add a seasonal cycle to the forcing. Which solver settings change, and does the adjoint conclusion survive?