Ordinary Differential Equations and Numerical Solvers

Dive into Deep Learning · §28.1

A velocity field and the curves it generates
ODEs, numerical solvers, and Neural ODEs.

Every continuous model is an ODE

Motivation

Give a velocity at every point and time; the solution is the curve that follows it. This one idea reappears as:

  • a ResNet = an Euler solver,
  • backprop = the adjoint ODE,
  • a normalizing flow’s log-density = a trace integral.
image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/

01

Vector fields and well-posedness

the IVP, the flow map, Picard–Lindelöf

The initial-value problem

The objects

A velocity field \mathbf f and a start \mathbf x_0 define

\dot{\mathbf x}(t) = \mathbf f(\mathbf x(t),t), \qquad \mathbf x(0)=\mathbf x_0, \qquad \mathbf x(t) = \mathbf x_0 + \int_0^t \mathbf f(\mathbf x(s),s)\,ds.

Euler’s method is just the left-endpoint Riemann sum of that integral.

The flow map

The objects

Collect all solutions into the map \Phi_t(\mathbf x_0) = \mathbf x(t):

\Phi_0 = \mathrm{id}, \qquad \Phi_{t+s} = \Phi_t\circ\Phi_s, \qquad \Phi_t^{-1} = \Phi_{-t}.

Invertibility comes free as a theorem: the starting point of continuous normalizing flows.

Picard–Lindelöf: one and only one

Well-posedness

If \mathbf f is continuous in t and L-Lipschitz in \mathbf x, the IVP has a unique solution.

Proof. The Picard operator (P\varphi)(t)=\mathbf x_0+\int_0^t\mathbf f(\varphi(s),s)\,ds is a contraction on a short interval (\delta < 1/L); Banach’s fixed point gives existence + uniqueness, then patch intervals to cover [0,T]. \blacksquare

When uniqueness fails

Counterexamples

Drop the Lipschitz bound and solutions fan out: \dot x=\sqrt{|x|} from x(0)=0 can wait any time c then leave as (t-c)^2/4:

Drop the growth bound and \dot x = x^2 blows up in finite time at t=1/x_0.

02

Linear ODEs and stability

the matrix exponential, eigenvalues, phase portraits

The matrix exponential

Linear systems

For \dot{\mathbf x}=A\mathbf x the solution is \mathbf x(t)=e^{At}\mathbf x_0, where

e^{At} = \sum_{k=0}^{\infty}\frac{(At)^k}{k!} = V e^{\Lambda t}V^{-1}.

The eigenbasis decouples the system into independent scalar modes e^{\lambda_i t}, exactly the eigendecomposition machinery of the linear algebra part at work.

Three ways to the same map

Linear systems

Power series, eigendecomposition, and the Euler limit (I+tA/n)^n all agree, and the norm decays exactly as the theory predicts:

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.735759

The stability dictionary

Linear systems

Each mode’s size is |e^{\lambda_i t}| = e^{(\operatorname{Re}\lambda_i)t}:

\operatorname{Re}\lambda < 0 → decay; \operatorname{Re}\lambda > 0 → blow-up; \operatorname{Im}\lambda \ne 0 → rotation. Eigenvalues are the stability certificate.

Phase portraits

Linear systems

The eigenvalue signature names the picture (node, saddle, spiral, center):

At a nonlinear fixed point the Jacobian \partial\mathbf f/\partial\mathbf x gives the same verdict (Hartman–Grobman).

03

Numerical solvers

Euler, Runge–Kutta, stiffness, and gradient descent

Forward Euler, order one

Solvers

\mathbf x_{n+1} = \mathbf x_n + h\,\mathbf f(\mathbf x_n,t_n).

Local error O(h^2) per step; unrolling N=T/h steps and summing a geometric series gives global error O(h), amplified by the stability factor e^{LT}. The general rule: local O(h^{p+1}) → global O(h^p).

Runge–Kutta: probe inside the step

Solvers

RK4 samples four slopes (k_1 at the start, two at the midpoint, one at the end) and Simpson-weights them \tfrac16(k_1+2k_2+2k_3+k_4): global order 4, so halving h cuts error by 16.

measured slope, Euler: 1.040
measured slope, RK4  : 4.022

The measured slopes are 1.0 (Euler) and 4.0 (RK4), matching theory.

Stiffness: stable or sorry

Solvers

On \dot x=-\lambda x, forward Euler is stable only for h < 2/\lambda; a fast dead mode then forces a tiny step. Backward Euler is stable for every h (A-stable) at the cost of one solve per step:

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)

The stability picture

Solvers

In the complex plane of z = h\lambda:

  • forward Euler decays only inside the disc |1+z| < 1;
  • RK4’s lobes stretch to \approx -2.79 on the real axis;
  • backward Euler is stable everywhere but a disc, in particular on the entire left half-plane, which is A-stability.

Gradient descent is a solver

Training as dynamics

Gradient descent with learning rate \eta is forward Euler on the gradient flow \dot{\mathbf x} = -\nabla L(\mathbf x); near a minimum with Hessian H it linearizes onto the test equation:

\dot{\boldsymbol\delta} = -H\boldsymbol\delta \;\;\xrightarrow{\ \text{Euler},\ h=\eta\ }\;\; \boldsymbol\delta_{k+1} = (I-\eta H)\,\boldsymbol\delta_k \;\;\Longrightarrow\;\; \eta < \frac{2}{\lambda_{\max}(H)}.

A diverging learning rate is a solver instability; an ill-conditioned Hessian is a stiff gradient flow. Momentum is one derivative up: heavy-ball discretizes m\ddot{\mathbf x}+\gamma\dot{\mathbf x}=-\nabla L.

The verdict flips at the predicted threshold

Training as dynamics

On a quadratic with \lambda_{\max}=10 the predicted flip is \eta = 0.2: the sweep diverges already at \eta = 0.201, while the slow mode (\lambda=1) converges at every rate shown:

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+17

The divergence is manufactured entirely by the loss surface’s own fast dead mode.

04

Neural ODEs and continuous flows

ResNet = Euler, the adjoint, the trace trick

A ResNet is an Euler step

Architecture

\mathbf x_{l+1} = \mathbf x_l + \mathbf f_\theta(\mathbf x_l) is one Euler step at h=1; depth becomes integration time:

The Neural ODE limit

Architecture

Shrink the step and the stack becomes \dot{\mathbf x}=\mathbf f_\theta(\mathbf x,t), \mathbf x(T)=\mathbf x_0+\int_0^T\mathbf f_\theta\,dt. A 1-hidden-layer net learns a circle→ellipse flow:

iter   0  loss 1.16972
iter 100  loss 0.00004
iter 200  loss 0.00000
iter 300  loss 0.00000
iter 400  loss 0.00000
mean endpoint error: 0.0011

Lipschitz \mathbf f_\theta → the map is invertible by Picard–Lindelöf: invertibility comes for free.

The adjoint = backprop in continuous time

Adjoint

Run an adjoint \mathbf a(t)=\partial L/\partial\mathbf x(t) backward:

\dot{\mathbf a} = -\Bigl(\tfrac{\partial\mathbf f}{\partial\mathbf x}\Bigr)^{\!\top}\mathbf a, \qquad \frac{\partial L}{\partial\theta} = \int_0^T \Bigl(\tfrac{\partial\mathbf f}{\partial\theta}\Bigr)^{\!\top}\mathbf a\,dt.

discrete adjoint vs finite differences: 1.2e-10
n =    100 Euler steps: |discrete - continuous adjoint| = 2.9e-02
n =   1000 Euler steps: |discrete - continuous adjoint| = 2.8e-03
n =  10000 Euler steps: |discrete - continuous adjoint| = 2.8e-04

Discretizing the adjoint ODE is the backprop recursion, but with O(1) memory instead of storing every step.

Densities along the flow

Continuous flows

A flow transports density by the trace of its Jacobian, replacing a log-determinant:

\frac{d}{dt}\log p_t(\mathbf x(t)) = -\operatorname{tr}\Bigl(\tfrac{\partial\mathbf f}{\partial\mathbf x}\Bigr).

Hutchinson’s estimator \operatorname{tr}(M)=\mathbb E[\boldsymbol\epsilon^\top M\boldsymbol\epsilon] turns that into one vector–Jacobian product (the FFJORD trick):

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.0173

Recap

Wrap-up

  • An IVP is a velocity field; Picard–Lindelöf gives a unique, invertible flow under a Lipschitz bound.
  • Linear systems: e^{At}=Ve^{\Lambda t}V^{-1}; eigenvalues set decay, growth, and rotation.
  • Euler is O(h), RK4 is O(h^4); stiffness forces implicit steps.
  • GD = Euler on the gradient flow: \eta < 2/\lambda_{\max} is a stability bound.
  • ResNet = Euler step; the Neural ODE is its continuous limit, invertible for free.
  • The adjoint ODE is backprop in continuous time, at O(1) memory.
  • Density flows by -\operatorname{tr}(\partial\mathbf f/\partial\mathbf x); Hutchinson makes it one VJP.

Next we add noise to the velocity field: stochastic differential equations.