The numerical ODE solver is the program, and the adjoint method backprops by solving a second ODE, giving O(1) memory in integration depth.
lane: substrate · arXiv:1806.07366 (Neural Ordinary Differential Equations) — NeurIPS 2018 (Best Paper) — Chen, Rubanova, Bettencourt, Duvenaud (U. Toronto / Vector Institute); arXiv:1912.04232 (JAX-MD) — NeurIPS 2020 — Schoenholz, Cubuk (Google Research)
Gives. Gradients through a whole simulation at O(1) memory; the solver is the program and physics parameters are fitted.
Needs. Non-stiff dynamics and a differentiable right-hand side.
Breaks. Reversing dz/dt = -10^4 z fails in double precision; a ReLU ODE drifts 1% in 11 steps (ANODE).
Every other approach on this hub treats the program as either LLM-generated code, a sparse symbolic equation, or an edited weight matrix. Here the program is the differential equation's right-hand side f(z,t,θ). Running the program means calling a numerical ODE solver. The solver's internal stepping scheme, not a neural network's layers, defines the forward computation. The orthogonal move is deriving gradients through that solver via the continuous adjoint equation, instead of backpropagating through the solver's unrolled steps. This decouples training memory from integration length entirely (arXiv:1806.07366 Sec. 2, Eq. 4).
1. Forward: the hidden state follows dh/dt = f(h(t), t, θ) (Eq. 2). The paper's own implementation calls the implicit Adams method via LSODE/VODE (Sec. 3). JAX-MD's version of forward is a physics simulation: a (space, energy, simulate) triple, where energy is any differentiable function, for example a pairwise potential or a GNN. Forces come from jax.grad(energy_fn), and an integrator like NVE/NVT advances the trajectory (arXiv:1912.04232 Sec. 3.2-3.3). 2. Adjoint instead of backprop-through-steps: define a(t) = ∂L/∂z(t). It obeys the adjoint ODE da/dt = −a(t)ᵗ ∂f/∂z (Eq. 4), the paper's "instantaneous analog of the chain rule." The method computes the parameter gradient dL/dθ = −∫ a(t)ᵗ ∂f/∂θ dt (Eq. 5) by concatenating state, adjoint, and gradient into one augmented vector and calling the solver once, backward in time (Algorithm 1). 3. Constant memory: the method does not store the forward trajectory. It recomputes the trajectory backward together with the adjoint (Sec. 2), giving O(1) memory versus a ResNet's O(L) (Table 1). This is the paper's headline claim. 4. JAX-MD's generalization: instead of backpropagating through one integration, JAX-MD differentiates an entire simulation trajectory with respect to a simulation parameter, for example particle diameter. The trajectory includes an inner optimization loop, a FIRE energy minimization run inside lax.scan. This enables "meta-optimization" over the physics itself, not just over a learned function evaluated once (Sec. 4.2, App. A).
The ODE's right-hand side f, or JAX-MD's energy function, is continuous and differentiable, fitted end-to-end by gradient descent through the adjoint. This is the most purely continuous approach on the hub, with no discrete search step inside the training loop at all. The one place something is fixed and non-differentiable is the solver's own control logic: step-size adaptation and the choice of integration scheme, Euler, RK, or implicit Adams. ANODE's critique adds a second case: reversing time flips the sign of the Jacobian's eigenvalues, which can make the backward solve numerically ill-conditioned regardless of how well f is fit (arXiv:1902.10298 Sec. III). Gradient flow stops wherever the solver's discretization error dominates the signal. That failure mode is unique to this approach because the solver itself is never trained; only f is.
The constant-memory property makes this the only approach on the hub where training cost is largely decoupled from depth, meaning integration length, rather than growing linearly with it. That property is the explicit motivation for continuous normalizing flows and latent-ODE time-series models that need long or irregularly-sampled trajectories (Sec. 4-5). The trade-off is per-step cost. Each forward and backward pass requires an adaptive solver to do work proportional to problem stiffness and desired tolerance. Minibatching multiple examples into one larger ODE can require up to K times more function evaluations to hold joint error tolerance, though the paper reports this did not increase substantially in their experiments (Sec. 6). JAX-MD scales physics simulation to large particle counts via neighbor and cell lists, near O(N log N) instead of O(N²) (Sec. 3.2.3). The paper's own benchmark (Table 1) shows it running roughly 2.4x slower than LAMMPS and HOOMD-Blue on GPU at N=64k, and roughly 10x slower than both on CPU at N=10k, because it is a general differentiable framework, not a hand-tuned CUDA kernel.
On MNIST, an ODE-Net reaches 0.42% test error with 0.22M parameters, versus a ResNet's 0.41% with 0.60M parameters. That is comparable accuracy at roughly a third of the parameters (arXiv:1806.07366 Table 1, Sec. 3). Latent ODEs beat RNNs on irregularly-sampled time-series extrapolation across every tested observation density: RMSE 0.1642 vs. 0.3937 at 30/100 observations, 0.1502 vs. 0.3202 at 50/100, and 0.1346 vs. 0.1813 at full observation (Table 2, Sec. 5.1). JAX-MD's trained graph-neural-network potential on 64-atom silicon DFT data reaches 2.8 meV/atom RMSE on energies, and the team deployed it in a 64,000-atom NVT simulation (Sec. 4.1, Fig. 1-2). Its meta-optimization experiment located the point of maximum frustration in a bidisperse 2D packing, by differentiating straight through 300 steps of an inner FIRE minimization loop (Sec. 4.2, Fig. 3, App. A).
1. Optimize-then-discretize versus discretize-then-optimize mismatch (ANODE Sec. IV): the continuous adjoint is derived first and then discretized, so its gradient need not match the true gradient of the actual discrete computation. The paper reports the discrepancy scales as O(dt) and causes outright divergence. On CIFAR-10 with a SqueezeNext backbone and adaptive RK45, the original Neural-ODE adjoint "diverges in the first epoch" (Fig. 3), with the same pattern on ResNet-18/CIFAR-100 (Fig. 4-5). ANODE's fix, discretize-then-optimize plus optional checkpointing, trades the O(1) memory win partly away, costing O(L)+O(Nt) instead of O(1), to regain stability. 2. Adjoint reversal instability (ANODE, Gholami, Keutzer, Biros 2019, arXiv:1902.10298 Sec. III): reversing an ODE's time direction flips the sign of its Jacobian's eigenvalues. The paper shows this can make the reverse solve numerically unstable, independent of step-size adaptation. For dz/dt = −10⁴z the flow "cannot be reversed in double precision," and a ReLU-based ODE accumulates 1% reconstruction error after just 11 adaptive RK steps. Adaptive time-stepping does not fix this, per the paper's own stated conclusion. 3. Where continuous time buys nothing (Kidger thesis, arXiv:2202.02435 Sec. 2.2.1): for ordinary image classification, the thesis states plainly that "there is usually little to be gained by using a continuous-time model" and that "traditional residual networks... are simply easier to work with." That is a direct statement from inside the field that the headline MNIST parity result does not generalize into a reason to prefer Neural ODEs for standard supervised tasks. 4. JAX-MD's own stated gaps (arXiv:1912.04232 Sec. 3.1): JAX-MD runs roughly 2.4x slower than LAMMPS and HOOMD-Blue on large GPU systems, and roughly 10x slower on CPU. TPUs "perform badly on large systems due to poor neighbor list performance." XLA's requirement for statically-known shapes makes some data structures hard to express. The authors explicitly expect DiffTaichi to outperform JAX-MD on many workloads (Sec. 2). 5. Stiffness remains an open, separately-worked problem (UNVERIFIED depth — this session's WebFetch pagination of the Kidger thesis could not retrieve Chapter 5/6 on solvers and speed): WebSearch, not independently WebFetch-confirmed, surfaces a 2024 paper, "Neural ODEs for Stiff Systems: Implicit Single-Step Methods" (arXiv:2410.05592). It reports that neural ODEs "frequently become stiff during training, even when the training data comes from a non-stiff ODE model." That failure mode compounds with the adjoint-reversal instability above, since stiffness is exactly the regime where reverse-time integration is most ill-conditioned.
Neural ODEs borrow the adjoint sensitivity method from classical optimal control, cited but not independently re-verified this session. The same paper motivates continuous normalizing flows and latent-ODE time-series models as direct applications (Sec. 4-5). ANODE is the explicit stability critique and fix, trading Neural ODE's O(1) memory claim for O(L)+O(Nt) in exchange for a numerically consistent gradient. Kidger's thesis is a later synthesis extending the family to neural CDEs and neural SDEs, while explicitly naming where the continuous-time framing does not pay off. JAX-MD inherits the "differentiate the whole trajectory" idea directly from the Neural-ODE and adjoint lineage, and applies it to molecular dynamics instead of learned dynamics. It differentiates through physics integrators and nested optimization loops the same way Neural ODEs differentiate through a learned vector field.
Verified: WebFetch on https://ar5iv.labs.arxiv.org/html/1806.07366 (full text, Sec. 1-6, Eq. 2-5, Table 1-2, Fig. 2-6), https://ar5iv.labs.arxiv.org/html/1912.04232 (full text, Sec. 2-4, Table 1, Fig. 1-3, App. A), https://ar5iv.labs.arxiv.org/html/1902.10298 (ANODE, Sec. III-VI, Fig. 1/3-7), https://ar5iv.labs.arxiv.org/html/2202.02435 (Kidger thesis, partial — Sec. 1.3, 2.1.2, 2.2.1, 2.2.3.4, 3.5; Chapters 5-6 on solvers/stiffness/speed not retrievable through WebFetch's pagination this session, flagged UNVERIFIED below). · code: https://github.com/rtqichen/torchdiffeq (Neural ODE), https://github.com/jax-md/jax-md (JAX-MD)