Neural ODEs: Neural Ordinary Differential Equations
Let an ODE solver run the network, and train it with memory that does not grow with depth.
A small neural network gives the derivative of the hidden state. An off-the-shelf ODE solver integrates it to produce the output, and a second ODE, solved backward in time, produces the gradients.
Explaining the paperNeural Ordinary Differential EquationsThe same construction turns the determinant in a normalizing flow into a trace, and fits time series whose observations arrive at irregular times.
A deep network fixes its depth before training: 6 residual blocks, or 50, or 152. This paper replaces the stack with a differential equation. A neural network outputs the rate of change of the hidden state, a numerical solver integrates that rate from the input to the output, and the solver decides how many times to evaluate along the way.
The page follows the paper in order. A residual block is one step of Euler's method, so a stack of blocks is a crude ODE solver; defining the model as the ODE hands the step count to an adaptive solver; the adjoint method computes gradients by solving a second ODE backward, without storing the forward pass; the same continuous view replaces a determinant with a trace in normalizing flows; and a latent ODE models time series sampled at arbitrary times. The last sections cover the experiments and what the continuous model cannot do.
A residual block is one Euler step
A residual block computes a correction with a small network and adds it to its input:
In (1), is the hidden state entering block , is the block's network (a couple of convolutions, say), and are that block's own weights; the network outputs a change to the state, not a new state.
Now take an ordinary differential equation , a rule that gives the velocity of the state at every point and time. Euler's method solves it numerically: from the current point, evaluate the velocity and move a distance along it, then repeat.
With this is equation (1), except that each residual block carries its own weights. The paper states the connection in one line ("these iterative updates can be seen as an Euler discretization of a continuous transformation") and cites Lu et al. and Haber & Ruthotto (both 2017) for it. Haber & Ruthotto write the residual update with an explicit step size and analyse the network's stability with the standard stability condition for forward Euler.
A small example shows what the step size does. Take the 2-D rotation field and start at . The exact solution moves around the unit circle and keeps radius 1. One Euler step with lands at , radius 1.414, which is 41% off. Four steps of reach at radius 1.129 (13% off); a hundred steps of 0.01 reach it at radius 1.005. Each Euler step multiplies the radius by , so the error at a fixed end time shrinks roughly in proportion to .
Defining the network by its derivative
The paper asks what happens as a network adds layers and takes smaller steps. If the number of blocks grows while the step shrinks and the total depth stays fixed, the Euler iteration approaches the solution of an ODE. For a trained ResNet this is a heuristic, since its weights can change arbitrarily from one block to the next, and a continuous limit needs weights that vary smoothly with depth. The paper does not rely on the limit. It defines the model as the ODE:
One network with one set of weights gives the velocity at every depth ; the depth enters as an input. From Section 2 on, the paper writes the state as and the interval as , and this page does the same. The output is the state at , and a loss is computed from it:
ODESolve is any numerical solver, treated as a black box: it calls at whatever points it chooses and returns . Figure 1 draws the relation between the two models. The amber curves are exact solutions of one ODE, with depth running up the page; the teal polylines are Euler approximations with steps, one dot per step, as a residual network with blocks would compute them. At the single step overshoots badly; by the polyline sits on the curve.
In the paper's MNIST classifier, the baseline ResNet downsamples the image twice and then applies 6 residual blocks; the ODE-Net replaces those 6 blocks with one ODESolve module. In the authors' later PyTorch example, is two 3×3 convolutions on 64 channels, each seeing as one extra constant input channel, integrated from to . The parameter counts follow from weight sharing: six residual blocks of two 64-to-64 convolutions each, against one whose convolutions see 65 input channels,
convolution weights, plus biases. With the shared downsampling layers (about 132k weights) the totals come to about 0.58M and 0.21M, close to the paper's 0.60M and 0.22M (Table 1, from its own implementation). In code, the block is a solver call:
# ODE block: one solver call replaces the stack of residual blocks
def f(t, h, theta):
# one small net shared by every t; t enters as an extra input channel
x = concat(full_like(h[:, :1], t), h)
return conv2(relu(conv1(x))) # (norm layers omitted)
def ode_block(h0, theta, tol=1e-3):
# the solver chooses how many times to call f; only h(1) comes back
return odesolve(f, h0, t0=0.0, t1=1.0, theta=theta, rtol=tol, atol=tol)Adaptive solvers set the number of steps
Euler's method uses a fixed step. Modern solvers, descended from Runge (1895) and Kutta (1901), estimate their own local error at each step, typically by comparing two approximations of different order. If the estimate exceeds the requested tolerance, the step is rejected and retried smaller; if it is well below, the next step grows. Steps come out short where the trajectory bends and long where it is nearly straight.
The number of times the solver evaluates during one forward pass is the number of function evaluations, NFE, written . The paper calls it "an implicit number of layers", and forward time is proportional to it. Unlike a layer count it is not fixed: it depends on the input, the tolerance and the current weights. In Figure 2, drag the tolerance and watch the evaluation points move to the bends as NFE rises and the maximum error falls.
Because the tolerance is a solver argument, the same trained weights can run at different accuracies. The paper suggests training with a tight tolerance and loosening it at test time for real-time or low-power use. Its own settings: the supervised experiments used the implicit Adams method from LSODE and VODE, called through scipy.integrate; the default tolerance of 1.5e-8 was used for sequence modeling, and the classification and density-estimation experiments ran at 1e-3 and 1e-5 without losing accuracy.
The NFE also moved during training. In the MNIST runs it increased throughout training, which the paper attributes, with a "presumably", to the model adapting to increasing complexity, so each training step cost more late in training than early, where a ResNet's cost per step stays fixed.
Gradients from a backward ODE
Training needs , and the direct route is to record every operation the solver performed and backpropagate through them, as autodiff would for any program. That route stores every intermediate state, including each internal stage of each step, so memory grows with , a number the solver controls. It also did not suit the paper's solver: implicit Adams solves a nonlinear optimization problem at every step, which makes backpropagating through it difficult. Reversible residual networks already gave constant-memory training by recomputing activations, but they require architectures that partition the hidden units.
The paper instead uses the adjoint sensitivity method (Pontryagin et al., 1962), from optimal control. Define the adjoint as the gradient of the loss with respect to the state at time :
At the end of the interval it is available directly: is the gradient of the loss with respect to the output. To move it earlier, apply the chain rule across one tiny Euler step of size . The step and its Jacobian with respect to are
and backpropagation through that step gives
and letting gives the adjoint ODE, which the paper calls "the instantaneous analog of the chain rule" (Appendix B has this proof):
The minus sign comes from the direction of the chain rule: stepping from back to adds to the adjoint, so per unit of forward time the adjoint changes by the negative of that. Equation (4) is solved backward, from to , and is the gradient with respect to the input.
The weights get the same treatment. Treat as extra state with ; its adjoint starts at zero at and changes at the rate . Integrated from down to , that is equation (5), with the leading minus and the reversed limits both as printed:
Both integrands are vector-Jacobian products, and . Reverse-mode autodiff computes them in one backward sweep through , at a cost similar to one evaluation of , without forming either Jacobian.
A one-dimensional example can be checked by hand. Let , with , , target , and loss . The forward solve, the adjoint from (4), and the weight gradient from (5) are:
Equation (4) here reads , solved backward from . In (5), , and the integrand happens to be constant for this linear . The integral from 1 down to 0 of a constant is , here , and the leading minus of (5) turns it into . Both results agree with differentiating the closed form directly: and both equal here, which is . The gradient is negative because falls short of the target, so a larger lowers the loss.
Figure 3 runs the same computation numerically, with RK4, on the nonlinear , with two weights and ( here). Play the reverse solve: the state is rebuilt from , the adjoint starts at , and the two running integrals for and start at zero at . When the solve reaches they land on the hollow circles, which mark finite-difference gradients computed by perturbing each weight and re-running the forward pass.
Equations (4) and (5) need at every time the backward solver visits, and the forward trajectory was not kept. The paper recomputes it: the original ODE is integrated backward from alongside the adjoint. Algorithm 1 stacks the state, the adjoint and the parameter gradient into one augmented vector and makes a single solver call from to :
# Algorithm 1: gradients of an ODE solve from one backward solver call
def ode_backward(z1, dL_dz1, f, theta, t0, t1):
def aug_dynamics(s, t):
z, a, _ = s
f_val, vjp = vjp_of(f, z, t, theta) # one autodiff pass through f
a_dfdz, a_dfdtheta = vjp(a) # a^T df/dz and a^T df/dtheta
return [f_val, -a_dfdz, -a_dfdtheta] # eq (4) and the eq (5) integrand
s0 = [z1, dL_dz1, zeros_like(theta)] # nothing stored from the forward
z0, dL_dz0, dL_dtheta = odesolve(aug_dynamics, s0, t1, t0) # t1 -> t0
return dL_dz0, dL_dthetaNothing from the forward pass is stored except its output, so memory does not depend on how many steps the solver took. Figure 4 plays the forward and backward passes and meters the two approaches.
With numbers: say has 100 entries and the solver takes 50 steps on some input. Backpropagating through the solver keeps at least one copy of the state per step, 5,000 numbers, and 50,000 if the solver takes 500 steps. The adjoint pass holds the current state (100 numbers), the current adjoint (100), and a gradient buffer the size of , whatever the step count. Table 1 of the paper lists the memory of the ODE-Net as , against for backpropagating through a Runge–Kutta solver and for the ResNet. The is in depth: the parameters, the state and the adjoint buffers still occupy memory in proportion to their sizes.
The adjoint also saved time in the paper's measurements. On the trained MNIST model the backward solve used roughly half as many function evaluations as the forward solve, while backpropagating through the integrator has to pass back through every forward evaluation; the paper concludes the adjoint method is "not only more memory efficient, but also more computationally efficient" than that route. When the loss depends on the state at several times (as in the time-series model below), the backward solve is split at each observation time, and the adjoint jumps by there before continuing.
The idea of training continuous-time networks with the adjoint method is also older than the paper: LeCun et al. (1988) and Pearlmutter (1995) proposed it, and the paper's related-work section notes it "was not demonstrated practically". The paper supplies the general recipe of Algorithm 1 for any black-box solver, with code for autograd (its Appendix D) and later for PyTorch.
When the backward reconstruction drifts
The adjoint equations give the exact gradient of the continuous problem. The computed gradient uses the rebuilt , and the paper names the risk in its limitations section: running the dynamics backward "can introduce extra numerical error if the reconstructed trajectory diverges from the original".
The mechanism shows up in the scalar ODE with . Forward, every trajectory decays toward zero, and two starting values that differ by end up differing by at . Backward, the same equation expands: an absolute error in the stored , from the forward solver's tolerance or from rounding, reaches as , and this holds even for a perfect reverse solver because the reverse-time ODE itself amplifies the error. At , forward contraction is , and an error of in grows to about 4.9 by , larger than the state. The rebuilt then feeds and in (4) and (5) with wrong values.
The ANODE paper (Gholami, Keutzer & Biros, 2019) made this failure precise for Neural ODEs. It shows that reverse-time reconstruction "can lead to mathematically incorrect gradient information" and to divergent training, and gives the scalar case its own numbers: for , reversing to 1% accuracy takes about 200,000 time steps, and for the reversal is impossible in double precision. Its fix is checkpointing: store the state at some points of the forward pass and re-integrate forward from those, at a memory cost of , where is the number of ODE blocks and the number of time steps inside one block, instead of .
Chen et al. proposed the same remedy themselves, storing intermediate values of on the forward pass and re-integrating from them, which is the technique of gradient checkpointing applied to a solver. They did not use it: they report that drift "was not a practical problem" in their experiments, and that they informally checked that reversing many layers of continuous normalizing flows at default tolerances recovered the initial states. Strongly contracting dynamics are the case where that check fails and the memory cost of checkpointing has to be paid.
Continuous normalizing flows
The authors call the next result an "unexpected side-benefit" of continuous transformations. A normalizing flow models a complicated density by drawing from a simple one (a standard Gaussian) and pushing it through an invertible function. The density of the result follows from the change of variables formula, which corrects for how much the function stretches or compresses volume near each point:
The determinant of a general Jacobian costs . Discrete flows keep it cheap by restricting . The planar flow of Rezende & Mohamed (2015) uses one hidden unit per layer:
Here and set the unit's input, is a smooth scalar nonlinearity such as tanh (not the hidden state from Eq. (1)), and is the direction along which the layer moves . With a single unit the Jacobian is the identity plus a rank-one matrix, and its determinant collapses to the scalar . A layer with hidden units has a rank- correction, and the paper puts the cost as cubic in either the dimension of or the number of hidden units. So standard architectures stack many layers of a single hidden unit each.
Now let the transformation be an ODE, . Over a tiny step the map is , with Jacobian where . For a 2×2 example, take with rows and , and :
The two differ by . In general
because every term of the determinant other than the product of the diagonal entries contains at least two off-diagonal entries, each carrying a factor of . Taking the log, the log-density changes by per step, and in the limit the determinant becomes a trace. This is the paper's Theorem 1, the instantaneous change of variables, which assumes is uniformly Lipschitz in and continuous in :
The trace of is the divergence of the velocity field, the net rate at which the flow spreads volume at a point. Where the flow compresses volume (negative trace) the density of the moving sample rises, and where it spreads volume the density falls. The continuous version of the planar flow and its density change are
and because the trace is linear, many hidden units cost no more per unit than one. If the dynamics are a sum , the log-density ODE is a sum of traces:
The paper prints (10) without the minus sign; (8) and (9) both carry it, and the sum of traces inherits it, so this page restores it (see the Provenance panel). For a planar unit the trace is the scalar , so a layer of units costs where a wide discrete planar layer costs . The paper calls these models continuous normalizing flows (CNFs). The flow map is invertible whenever the ODE has a unique solution, so itself does not need to be bijective.
The linear cost relies on that additive structure. For a free-form , reverse-mode autodiff produces the Jacobian one row per backward pass, so the diagonal entries of an exact trace still take about passes. FFJORD (Grathwohl et al., October 2018) states the exact trace as and replaces it with Hutchinson's estimator, which needs one pass: for a random vector with zero mean and identity covariance,
and is a single vector-Jacobian product. In Figure 6, switch between the three ways of paying for the density term and slide . At small the cost of evaluating itself dominates and the three curves nearly coincide.
Figure 7 shows the density bookkeeping. A grid over a Gaussian is carried by a smooth velocity field toward three clusters; the field is hand-made, not a trained CNF. Where grid cells shrink, the density of the points inside rises by the integral of along their paths; where cells stretch, it falls.
The experiments are two-dimensional. For density matching, where the target density can be evaluated, the flows minimize : the CNF trained for 10,000 iterations with Adam, the planar NF for 500,000 with RMSprop, and the CNF generally reached lower loss. A planar CNF with hidden units is at least as expressive as a planar NF with layers, and in the paper's words sometimes much more expressive. For maximum-likelihood training on samples, the CNF uses 64 hidden units against 64 stacked one-unit NF layers. Because a CNF runs backward at about the cost of its forward pass, it trains by maximum likelihood (data to noise, scoring each sample) and then samples by running forward; the paper describes the CNF's transformations as smooth and interpretable and the NF's as very unintuitive. Later generative models kept this machinery: flow matching trains CNFs without simulating the ODE during training, and score-based diffusion computes likelihoods under its probability-flow ODE with equation (8).
Latent ODEs for irregular time series
Medical records, network traffic and neural spike trains arrive at irregular times. A recurrent network advances in fixed steps, so the usual approach bins the observations into fixed intervals, which creates missing data and blurs the timing. A continuous-time model can be evaluated at whatever times the data happens to have.
The paper's latent ODE represents each series by a latent trajectory. A latent initial state is drawn, one ODE shared across all series evolves it, and a decoder emits an observation at each requested time:
Here is time-invariant: it takes the latent state and no time input, so the latent state at one time determines the whole trajectory, forward and backward. All uncertainty about a series sits in its initial state . The model trains as a variational autoencoder: a recognition RNN reads the observations backward in time and outputs an approximate posterior , a sample from it is integrated to every observation time, and the decoder's likelihood at those times, together with the KL term, gives the loss. Gradients come from the adjoint method with a jump at each observation time, as described above. Extrapolation is the same solver call with later times.
The benchmark is a set of 1,000 two-dimensional spirals, half clockwise and half counter-clockwise, each sampled at 100 equally spaced times with Gaussian noise added. From each spiral, 30, 50 or 100 of the points are drawn at random without replacement, and error is measured on 100 time points extending beyond the training points. The model is small: a 4-dimensional latent state, a recognition RNN with 25 hidden units, and one-hidden-layer networks of 20 units for and for the decoder. The baseline is an RNN with 25 hidden units; a second version receives the time gap to the next observation as input. Figure 8 shows the setup and the paper's Table 2 errors.
Predictive RMSE for the latent ODE is 0.1642, 0.1502 and 0.1346 at 30, 50 and 100 observations; for the RNN it is 0.3937, 0.3202 and 0.1813. At 30 points the latent ODE's error is 0.42 times the RNN's, and it stays lower at every setting. The inferred 4-dimensional latent trajectories separate into two clusters, one decoding to clockwise spirals and one to counter-clockwise.
The time of an observation can carry information too: a patient may be more likely to have a test when they are sick. The paper attaches an event rate to the latent state and scores the observation times with an inhomogeneous Poisson process,
which rewards a high rate at the times events occurred and penalizes the total rate over the window. The integral is computed in the same solver call as the latent trajectory, by adding it to the state.
What the experiments show
The MNIST comparison is Table 1 of the paper:
| Model | Test error | Params | Memory | Time |
|---|---|---|---|---|
| 1-layer MLP | 1.60% | 0.24M | - | - |
| ResNet | 0.41% | 0.60M | O(L) | O(L) |
| RK-Net | 0.47% | 0.22M | O(L̃) | O(L̃) |
| ODE-Net | 0.42% | 0.22M | O(1) | O(L̃) |
The ODE-Net's 0.42% error ties the ResNet's 0.41% (it is a hair worse) with 0.22M parameters instead of 0.60M, about 2.7 times fewer, and memory that is constant in depth. The RK-Net isolates the two effects. It has the same architecture as the ODE-Net and differs only in backpropagating directly through a Runge–Kutta integrator. It has the same 0.22M parameters, so the parameter saving comes from sharing one across depth, and its memory is , so the constant memory comes from the adjoint. The paper does not present the ODE-Net as a more accurate classifier, and its acknowledgements thank Yuval Frommer "for pointing out an unsupported claim about parameter efficiency", which the authors withdrew.
Across the three applications: on MNIST the ODE-Net matched the ResNet's accuracy with memory constant in depth and a tolerance that can be loosened after training; the CNFs matched or beat planar flows using wide layers whose density term is a trace; and the latent ODE fit irregularly timed spirals with lower extrapolation error than the RNN at every sampling rate.
Limits of the continuous model
An ODE with a unique solution cannot let two trajectories cross. If two trajectories met at the same point at the same time, the solution from that point onward would have to follow both, and uniqueness allows only one. So the map from input to output is a homeomorphism: it can stretch and bend the space but not tear it or pass one region through another. On the real line a Neural ODE cannot map and , because the two paths would have to cross. A residual network can, since its finite steps jump over each other. Augmented Neural ODEs (Dupont et al., 2019) add extra dimensions to the state so the trajectories have room to go around each other, and report more stable training, better generalization and lower computational cost. Scrub time in Figure 9 to see the forced crossing in one dimension and the detour in two.
The paper lists practical limits of its own. Minibatching stacks examples into one ODE of dimension , and the solver must meet its tolerance on all of them at once; in the worst case that needs up to times as many evaluations as solving each example separately, but in practice the paper found the count did not increase substantially. Uniqueness needs uniformly Lipschitz in and continuous in (Picard's theorem), which holds for networks with finite weights and Lipschitz nonlinearities such as tanh or ReLU. The user has to choose a tolerance for both the forward and the reverse solve.
Questions you might still have
If a residual network is already an Euler discretization, what does this paper add?
The Euler reading is from 2017 (Haber & Ruthotto; Lu et al.; E). The paper adds three things: training through any black-box ODE solver with the adjoint method at memory that does not grow with the number of solver steps, continuous normalizing flows whose density change is a trace instead of a determinant, and the latent ODE for irregularly sampled time series.
How can memory be constant if the backward pass still needs the forward trajectory?
The forward trajectory is not stored. The backward solve starts from the final state z(t1) and integrates the original ODE backward alongside the adjoint, so at any moment it holds the current state, the current adjoint and the running parameter gradient. That is constant in the number of solver steps, not constant in bytes: those buffers scale with the state size and the parameter count.
Is the adjoint gradient exact?
It is the exact gradient of the continuous problem. The computed gradient is only as good as the backward reconstruction of z(t), which can drift from the forward path; for strongly contracting dynamics the drift grows exponentially, as the ANODE paper (2019) shows. The paper proposes checkpointing as the remedy but reports that drift was not a practical problem in its experiments and did not use it.
Is the adjoint method slower than ordinary backpropagation?
Compared with backpropagating through the solver's operations, the paper found it faster as well as lighter: the backward solve used roughly half as many function evaluations as the forward solve. Compared with a ResNet of fixed depth L, both passes cost a number of evaluations the solver picks, and that number grew during training in the paper's MNIST runs.
How does the network f know the time t?
It takes t as an input. In the authors' PyTorch MNIST example each convolution in f sees t as an extra constant input channel. The latent ODE uses a time-invariant f, so its dynamics depend only on the current latent state.
What can a Neural ODE not represent?
Its trajectories cannot cross, so the input-to-output map is a homeomorphism. On a line it cannot swap -1 and +1, which a residual network with discrete jumps can. Augmented Neural ODEs (Dupont et al., 2019) add extra state dimensions so the trajectories can pass each other.
Footnotes & further reading
- The paper: Chen, Rubanova, Bettencourt, Duvenaud, Neural Ordinary Differential Equations (NeurIPS 2018, one of four best paper awards). torchdiffeq is the authors' PyTorch library of differentiable ODE solvers; its MNIST example builds the ODE-Net described above.
- Residual networks as discretized dynamics, the reading this paper builds on: Weinan E, A Proposal on Machine Learning via Dynamical Systems (2017); Lu et al., Beyond Finite Layer Neural Networks; Haber & Ruthotto, Stable Architectures for Deep Neural Networks. For residual blocks themselves, see the ResNet explainer.
- The adjoint (costate) method comes from optimal control: Pontryagin, Boltyanskii, Gamkrelidze, Mishchenko, The Mathematical Theory of Optimal Processes (1962).
- The discrete change of variables and normalizing flows: Rezende & Mohamed, Variational Inference with Normalizing Flows (2015), and Dinh et al., NICE (2014).
- The one-pass trace estimator for free-form CNFs (a follow-up, not in this paper): Grathwohl et al., FFJORD (2018).
- The two follow-ups on the limits: Gholami et al., ANODE (2019), on the accuracy of the backward reconstruction and checkpointing, and Dupont et al., Augmented Neural ODEs (2019), on the no-crossing limitation and its fix.
How could this explainer be improved? Found an error, or something unclear? I read every message.