PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Ordinary Differential Equations

Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, David Duvenaud

arXiv:1806.07366v5cs.LGcs.AIstat.ML

TL;DR

Training continuous-depth neural networks through ODE solvers creates a memory and numerical-differentiation challenge. This paper parameterizes hidden-state derivatives with neural networks, enables scalable solver backpropagation, and demonstrates constant-memory, adaptive computation across several model types.

  • Problem

    Backpropagating through ODE solvers is technically difficult because direct differentiation incurs high memory costs and additional numerical error.

  • Method

    The paper defines hidden-state dynamics with neural networks, solves them using black-box ODE solvers, and derives reverse-mode differentiation without accessing solver internals.

  • Results

    The resulting models use constant memory, adapt evaluation to requested accuracy, and achieve performance around that of standard residual networks.

  • Takeaways & Limitations

    ODE-based neural models support explicit speed–precision control and continuous-time predictions while extending continuous-depth modeling to learning and density estimation.

  • Takeaways & Limitations

    The framework requires users to choose error tolerances for both forward and reverse passes during training.

Abstract

from arXiv · show

We introduce a new family of deep neural network models. Instead of specifying a discrete sequence of hidden layers, we parameterize the derivative of the hidden state using a neural network. The output of the network is computed using a black-box differential equation solver. These continuous-depth models have constant memory cost, adapt their evaluation strategy to each input, and can explicitly trade numerical precision for speed. We demonstrate these properties in continuous-depth residual networks and continuous-time latent variable models. We also construct continuous normalizing flows, a generative model that can train by maximum likelihood, without partitioning or ordering the data dimensions. For training, we show how to scalably backpropagate through any ODE solver, without access to its internal operations. This allows end-to-end training of ODEs within larger models.

1 Introduction

The paper frames neural networks as continuous-depth systems whose hidden state follows an ODE solved by a black-box differential-equation solver. This formulation enables constant-memory training, adaptive computation, scalable invertible flows, and continuous-time models for irregularly timed data.

  • Core formulation: Neural ODEs define hidden-state dynamics with a neural-network parameterized derivative, mapping h(0) to h(T) through an ODE initial-value problem.A black-box differential-equation solver evaluates the dynamics as needed to reach the desired accuracy.
  • Memory efficiency: Constant-memory training computes gradients without backpropagating through solver operations or storing forward-pass intermediates.The resulting memory cost remains constant as model depth increases.
  • Adaptive computation: Adaptive ODE solvers monitor approximation error and adjust evaluations on the fly to meet requested accuracy.This lets evaluation cost scale with problem complexity while trading numerical precision against speed.
  • Continuous time-series models: Continuously defined dynamics naturally incorporate time-series observations and emissions arriving at arbitrary times without discretizing their intervals.The paper constructs and demonstrates a continuous time-series model based on this property.

2 Reverse-mode automatic differentiation of ODE solutions

The section presents reverse-mode differentiation through an ODE solver as the key training challenge and resolves it by treating the solver as a black box. The adjoint sensitivity method computes gradients via a backward augmented ODE with low memory cost, linear scaling, and explicit numerical-error control.

  • Reverse-mode differentiation through the ODE solver is the main technical difficulty, because differentiating forward operations requires high memory and adds numerical error.
  • The adjoint sensitivity method treats the ODE solver as a black box and computes gradients by solving an augmented ODE backward in time.The augmented system contains the original state and the loss sensitivity with respect to that state.
  • The method applies to all ODE solvers, scales linearly with problem size, uses little memory, and explicitly controls numerical error.
  • Gradients with respect to the state are obtained by integrating the adjoint backward, while recomputing the forward state trajectory backward as needed.The adjoint is a(t) = ∂L/∂z(t), and parameter gradients depend on both the state and adjoint.
  • All gradients can be computed in a single ODE-solver call by concatenating the state, adjoint, and parameter partial derivatives into one augmented vector.Automatic differentiation evaluates the required vector-Jacobian products at a time cost similar to evaluating f.

3 Replacing residual networks with ODEs for supervised learning

This section evaluates neural ODEs for supervised learning, comparing ODE-based residual blocks with standard and directly differentiated Runge–Kutta alternatives. It shows that solver tolerance controls an accuracy–cost trade-off, while adjoint-based training reduces memory and backward evaluations.

  • Model Architectures: An ODE-Net replaces six standard residual blocks in a small downsampling residual network, alongside an RK-Net that backpropagates directly through a Runge–Kutta integrator.The architectures are evaluated by test error, parameter count, and memory cost.
  • Error Control in ODE-Nets: ODE solver tolerance controls the network’s output accuracy and creates a trade-off between computational cost and accuracy.Forward-call time is proportional to the number of function evaluations, and lower test-time accuracy can follow high-accuracy training.
  • Error Control in ODE-Nets: The backward pass uses roughly half as many function evaluations as the forward pass, suggesting adjoint sensitivity is more computationally efficient than direct integrator backpropagation.Direct backpropagation must differentiate through every function evaluation in the forward pass.
  • Network Depth: The number of function evaluations increases during training, adapting to the model’s increasing complexity.For ODE solutions, function evaluations provide a solver-dependent notion of network depth and depend on the input or initial state.

4 Continuous Normalizing Flows

Continuous normalizing flows replace discrete flow layers with continuous-time dynamics, turning the log-density change into a trace-based ODE and enabling scalable, reversible density modeling. This reduces computation for wide flows and supports maximum-likelihood training followed by efficient sampling.

  • Computational cost: Continuous flows avoid the cubic Jacobian-determinant bottleneck and can evaluate wide flow models with cost linear in the number of hidden units M, rather than O(M^3).The trace operation is linear, whereas standard normalizing-flow architectures commonly use many layers with a single hidden unit.
  • Instantaneous change of variables: The instantaneous change-of-variables theorem expresses log-probability evolution through an ODE driven by the trace of the dynamics’ Jacobian.Under uniformly Lipschitz dynamics in z and continuity in t, the continuous transformation has a well-defined density evolution.
  • Continuous normalizing flows: Time-dependent dynamics and learned gates define continuous normalizing flows whose hidden-unit dynamics can activate selectively over time.The dynamics may depend on t through a hypernetwork, with gating functions σ_n(t) ∈ (0, 1).
  • Density matching: CNFs generally achieve lower density-matching loss than NFs, using 10,000 Adam iterations versus 500,000 RMSprop iterations in the comparison.The task minimizes KL(q(x)∥p(x)); the NF training setup follows the procedure suggested by Rezende and Mohamed (2015).
  • Maximum likelihood training: Reverse transformation costs are approximately the same as forward transformation, enabling maximum-likelihood density estimation followed by random sampling from the learned density.The model maximizes E_p(x)[log q(x)] using the continuous change-of-variables theorem, then reverses the CNF to generate samples.
  • Empirical behavior: For Two Circles, the CNF smoothly rotates planar flows to spread particles into circles, while NF transformations are unintuitive and struggle with Two Moons.The comparison used 64 hidden units for CNF and 64 stacked one-hidden-unit layers for NF.

5 A generative latent function time-series model

This section presents a continuous-time generative latent ODE model for irregularly sampled time series, representing each series as a trajectory governed by shared dynamics and a local initial state. The model supports arbitrary-time prediction, Poisson-process observation likelihoods, and smooth latent representations that enable reconstruction and extrapolation.

  • Model: The latent ODE model represents each time series with a trajectory determined by a local initial state and global latent dynamics shared across series.An ODE solver produces latent states at the observation times, which generate the observations.
  • Model: A neural network parameterizes the time-invariant dynamics function f, making each latent state determine a unique trajectory that can be extrapolated forwards or backwards.The dynamics satisfy ∂z(t)/∂t = f(z(t), θf).
  • Training and Prediction: Training uses a variational autoencoder with a backward RNN recognition network, while the latent ODE generates predictions at arbitrary time points on a continuous timeline.The recognition network outputs qφ(z0|x1, x2, . . . , xN).
  • Poisson Process likelihoods: A Poisson process likelihood models observation times from latent-state-dependent event rates and can be combined with a data likelihood to model observations and their sampling times jointly.The event-rate function λ(·) can be parameterized by another neural network and evaluated with the latent trajectory in one ODE-solver call.
  • Experiments: On noisy spiral series with irregular timestamps, the latent ODE achieved substantially lower predictive RMSE than the recurrent neural-net baseline and produced consistent reconstructions and extrapolations.Timestamps were sampled with n = {30, 50, 100}, and predictions were evaluated on 100 future time points.
  • Latent space interpolation: Latent trajectories formed separate clusters for clockwise and counter-clockwise spirals, while varying the initial point changed trajectories smoothly between these behaviors.The learned latent representation distinguished the two spiral directions and decoded smoothly into data-space trajectories.

6 Scope and Limitations

The framework has limitations in minibatching, tolerance selection, and reverse-time trajectory reconstruction, while solution uniqueness depends on regularity assumptions. Checkpointing addresses reconstruction error, and reported experiments found relaxed tolerances without performance degradation.

  • Minibatching: Minibatching concatenates K states into a combined ODE of dimension D × K, but controlling joint error can require K times more evaluations.In practice, the passage reports that the number of evaluations did not increa…
  • Uniqueness: Unique solutions are guaranteed when the dynamics are uniformly Lipschitz continuous in z and continuous in t, satisfied by finite-weight networks with tanh or relu.This follows from Picard’s existence theorem.
  • Setting tolerances: 1.5e-8 was the default sequence-modeling tolerance, while classification used 1e-3 and density estimation used 1e-5 without degrading performance.The framework requires choosing error tolerances for both forward and reverse passes during training.
  • Reconstructing forward trajectories: Reverse integration can accumulate numerical error when reconstructed trajectories diverge from the originals, but checkpointing can recover the exact trajectory.Checkpointing stores intermediate z values and re-integrates from them; the authors did not find reconstruction to be a practical problem.

7 Related Work

Prior work explored adjoint training, reversible residual networks, adaptive computation, and differentiable ODE solvers, while this work integrates black-box ODE solvers into automatic differentiation without architectural restrictions.

  • Adjoint methods and reversible networks: The adjoint method had been proposed for continuous-time neural networks, while residual-network research explored reversibility and approximate computation; this work demonstrates both through direct ODE solving.Prior demonstrations were not practical, whereas the approach applies these properties more generally using an ODE solver.
  • Adaptive computation: ODE solvers adapt computation using well-studied, computationally cheap, and generalizable rules, avoiding secondary networks’ training and testing overhead and extra parameters.Secondary networks choose recurrent or residual evaluations but introduce additional fitting requirements.
  • Learning differential equations: Prior work learned differential equations with feed-forward or recurrent networks and connected Gaussian Processes to ODE solvers, including applications such as fluid simulation.These studies fit or approximate differential equations from data.
  • Differentiating through ODE solvers: The approach trains black-box ODE solvers end-to-end through a generic vector-Jacobian product, unlike methods that backpropagate through solver operations or use quadratic-time forward sensitivity analysis.Adjoint sensitivity analysis is linear in the number of variables, and vector-Jacobian products connect black-box solvers with automatic differentiation.

8 Conclusion

The paper develops models built from black-box ODE solvers for time-series modeling, supervised learning, and density estimation. These models support adaptive evaluation, explicit speed–accuracy control, and continuous-time normalizing flows based on an instantaneous change-of-variables formula.

  • Applications: Black-box ODE solvers serve as model components for time-series modeling, supervised learning, and density estimation.The paper develops models across all three application areas.
  • Adaptive computation: Adaptive evaluation allows explicit control of the tradeoff between computation speed and accuracy.The models adjust evaluation strategy while exposing a speed–accuracy tradeoff.
  • Continuous-time normalizing flows: An instantaneous change-of-variables formula enables continuous-time normalizing flows that scale to large layer sizes.The conclusion identifies both the formula and the resulting flow model as final developments.

Appendix A Proof of the Instantaneous Change of Variables Theorem … Appendix C Full Adjoint sensitivities algorithm

The appendices prove the instantaneous change-of-variables result, relate it to volume-preserving flows and Liouville dynamics, and develop adjoint sensitivities for hidden states, parameters, and integration times. They conclude with a complete reverse-time ODE algorithm for computing these gradients.

  • Appendix A Proof of the Instantaneous Change of Variables Theorem: The instantaneous change-of-variables proof takes the infinitesimal limit of finite-time log-density changes, using Taylor expansion and boundedness assumptions to justify exchanging limits and products.The proof assumes f is Lipschitz in z, continuous in t, and that z(t) is bounded.
  • A.1 Special Cases: The appendix also specifies the parameterization used in all experiments.This implementation statement accompanies the special-case flow constructions.
  • A.1 Special Cases: Hamiltonian continuous normalizing flows have zero Jacobian trace and therefore preserve volume.They split the data into two equal partitions, forming the continuous analog of NICE.
  • A.2 Connection to Fokker-Planck and Liouville PDEs: The instantaneous change-of-variables formulation follows particle trajectories and can be solved as an ODE with state size D + 1, avoiding a finite-difference grid exponential in D.The Liouville PDE requires spatial density derivatives, whereas trajectory integration is more practical numerically.
  • A Modern Proof of the Adjoint Method: The adjoint method is presented through an alternative proof described as short and easy to follow.This establishes the appendix’s proof-oriented treatment of continuous reverse-mode differentiation.
  • B.1 Continuous Backpropagation: The adjoint state is the gradient of the loss with respect to the hidden state at a specified time and obeys a differential equation solved backward from the final-time loss gradient.Intermediate-time loss terms are handled by repeating the backward adjoint step over each interval and summing the resulting gradients.
  • B.2 Gradients wrt. θ and t: Augmenting the ODE with parameters and integration times yields gradients for all inputs to an initial value problem solver.Parameter gradients are obtained by integrating the parameter adjoint over the full interval, while start- and end-time gradients are also recovered.
  • Appendix C Full Adjoint sensitivities algorithm: The full adjoint-sensitivities algorithm initializes an augmented state with the final state and loss gradients, computes vector-Jacobian products, and solves the augmented dynamics backward from t1 to t0.Its inputs include θ, t0, t1, z(t1), and ∂L/∂z(t1).

Appendix D Autograd Implementation

The appendix implements custom Autograd differentiation for SciPy’s ODE integrator by defining an augmented reverse-time system and registering argument-specific vector-Jacobian products. The implementation propagates gradients with respect to states, times, and flattened function arguments through backward ODE solves.

  • Augmented dynamics: The augmented dynamics concatenates the original derivative with vector-Jacobian dynamics for states, time, and function arguments.Autograd’s make_vjp differentiates the flattened function with respect to all three argument groups.
  • Backward pass: Reverse-mode differentiation iterates from the final observation to the first, updating state and time sensitivities while accumulating parameter gradients.At each interval, the implementation accounts for time movement and integrates the augmented system backward to the previous observation.
  • Autograd registration: A generic wrapper selects only requested vector-Jacobian products, and Autograd registers these custom derivatives for odeint.grad_argnums_wrapper filters the full VJP output according to argnums before defvjp_argnums attaches it.

Appendix E Algorithm for training the latent ODE model · Appendix F Extra Figures

Appendix E describes latent ODE training as a VAE procedure combining an RNN variational posterior, ODE-based latent evolution, and ELBO maximization. Appendix F presents spiral reconstructions from noisy observations sampled at 30, 50, and 100 time points.

  • Appendix E Algorithm for training the latent ODE model: The latent initial state zt0 is inferred by traversing the observed sequence with an RNN encoder.The encoder produces parameters for q(zt0|{xti, ti}i, θenc).
  • Appendix E Algorithm for training the latent ODE model: The training algorithm follows a standard VAE with an RNN variational posterior and an ODESolve model.This combines sequence encoding with continuous-time latent dynamics.
  • Appendix E Algorithm for training the latent ODE model: Latent states zt1 through ztM are obtained by solving ODESolve from zt0 across the observation times.The function f defines the gradient dz/dt as a function of z.
  • Appendix E Algorithm for training the latent ODE model: Training maximizes the ELBO, combining reconstruction log-likelihoods with the prior p(zt0) and posterior q(zt0|{xti, ti}i, φ).The stated prior is p(zt0) = N(0, 1).
  • Appendix F Extra Figures: 30, 50, and 100 time points are used as the three observation settings for the spiral reconstructions.The panels are labeled (a) 30 time points, (b) 50 time points, and (c) 100 time points.
  • Appendix F Extra Figures: Spiral reconstructions are evaluated with a latent ODE using a variable number of noisy observations.Figure 10 identifies this experiment as spiral reconstruction under differing observation counts.
Loading 1806.07366v5…