Chen 2018 - Scope, Limitations and Statistical Use

Summary

The authors list four practical caveats (minibatching, uniqueness, tolerance choice, backward reconstruction error) and position the work among reversible nets, learned differential equations, GP-ODE links and ODE-sensitivity code in Stan and dolfin. The statistician’s reading: gradients through a solver are cheap, but identifiability, solver error and non-uniqueness of weights are your problem, not the solver’s.

Overview

Sections 6-8 are short but contain the conditions under which the method is valid and where it is positioned relative to existing differentiable-ODE tools (including Stan’s).

Main Content

Uniqueness of solutions (Picard-Lindelof, Sec. 6)

The initial value problem , has a unique solution if is uniformly Lipschitz in and continuous in (Coddington & Levinson 1955). This holds for neural-network vector fields with finite weights and Lipschitz nonlinearities such as or ReLU.

Issue (Sec. 6)StatementPractical consequence
MinibatchingBatch by concatenating states into one ODE of dimension ; controlling error across all elements could require times more evaluations than solving individually, but in practice evaluations did not rise substantially.Batched likelihoods over many subjects (population models) share one step-size sequence.
UniquenessSee theorem above.Trajectories cannot cross in state space; a planar neural ODE cannot represent some maps.
TolerancesUser must pick error tolerance for forward and reverse passes. Sequence modelling used the default ; classification and density estimation worked down to and respectively.Too-loose forward tolerance corrupts the likelihood and gradients that NUTS/HMC rely on.
Reverse reconstructionRe-integrating backward can introduce error if the reconstruction diverges from the forward path; remedy is checkpointing (store intermediate and re-integrate from them). The authors found it not a practical problem and checked that reversing many CNF layers recovered initial states.For chaotic dynamics (positive Lyapunov exponent) this is a real problem; see Chaos, Rough Likelihoods and Data Assimilation (Concept Note).
ThemeCited work (as summarized in the paper)Contrast
Adjoint for neural netsLeCun et al. 1988; Pearlmutter 1995proposed, not demonstrated practically
Reversible / constant-memory netsGomez et al. 2017; Haber & Ruthotto 2017; Chang et al. 2017require restricted architectures that partition hidden units
Adaptive computationGraves 2016; Figurnov et al. 2017; etc.extra networks and parameters; ODE solvers have well-studied, cheap rules
Learning differential equationsRaissi & Karniadakis 2018; GPs with ODE solvers (Schober et al. 2014)see Solving ODEs as Inference and ODE Filters and Smoothers
Differentiating through ODE solversdolfin (backprop through solver ops); Stan (forward sensitivity analysis, quadratic in the number of variables) vs adjoint (linear)the paper provides a generic VJP so a solver composes with any differentiable component

Statistical use (editorial synthesis, not claims of the paper)

Not from the paper

The following links the paper to Bayesian modelling practice and is the author’s reading.

  1. Forward vs adjoint sensitivities. For a handful of parameters and a small state, forward sensitivities (Stan, Torsten, diffrax forward mode) are simple and exact; adjoint wins when parameters are many (a neural-network vector field), as in universal differential equations (Universal Differential Equations and Scientific ML (Concept Note)).
  2. Bayesian inference over a neural vector field has non-identified weights (permutation/sign symmetries), so weight-space is meaningless; diagnostics belong in function (trajectory) space. See Izmailov 2021 - Bayesian Neural Network Posteriors - Overview.
  3. Solver error is part of the likelihood. With loose tolerances HMC sees a noisy or biased potential energy; either tighten tolerances or model solver uncertainty (ODE Filters and Smoothers).
  4. Deterministic latent trajectories. Latent ODEs have no process noise; if the data show stochastic variation in dynamics, SDE models are needed (Jump-Diffusion and Stochastic Differential Equation Models (Concept Note)).

Examples

Choosing a tolerance by sensitivity (recipe). Fit with rtol=1e-3, compute the log-likelihood and its gradient at the posterior mode; repeat at 1e-6, 1e-9. Accept the largest tolerance for which the log-likelihood differs by well under 0.1 and the gradient direction is unchanged. Applies to any ODE likelihood.

Connections

HMC and Stan in Practice, Efficient MCMC, Automatic Differentiation Variational Inference (ADVI).

See Also

Chen 2018 - Neural ODEs - Overview, Chen 2018 - ODE-Nets for Supervised Learning.