PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Differential Equations that are Easy to Solve

Jacob Kelly, Jesse Bettencourt, Matthew James Johnson, David Duvenaud

arXiv:2007.04504v2cs.LGstat.ML

TL;DR

Learned differential-equation models can have similar predictions but dynamics that differ in how costly they are to solve, and standard training does not penalize that complexity. The paper regularizes higher-order trajectory derivatives as a differentiable surrogate for solver cost; across tasks, tuning this regularization usually yielded at least a 2× speedup without substantially hurting performance.

  • Problem

    Learned models with nearly indistinguishable predictions can have substantially different dynamics, but standard training methods do not penalize dynamics complexity or solver cost.

    Source: §1 · pp. 1–2

  • Method

    The method regularizes norms of higher-order time derivatives of solution trajectories as a differentiable surrogate for adaptive-solver cost, with Taylor-mode automatic differentiation computing the derivatives efficiently.

    Source: §1 · p. 2, §4 · p. 3

  • Results

    Across all tasks, tuning the regularization usually gave at least a 2× speedup without substantially hurting model performance.

    Source: §10 · p. 8

  • Takeaways & Limitations

    The regularization lets a gradient-based optimizer trade off solver cost against predictive performance.

    Source: §1 · p. 2

    In density-estimation experiments with adaptive solvers, speed regularization made overall training about 70% slower, despite faster test-time evaluation.

    Source: §9 · p. 8

Abstract

from arXiv · show

Differential equations parameterized by neural networks become expensive to solve numerically as training progresses. We propose a remedy that encourages learned dynamics to be easier to solve. Specifically, we introduce a differentiable surrogate for the time cost of standard numerical solvers, using higher-order derivatives of solution trajectories. These derivatives are efficient to compute with Taylor-mode automatic differentiation. Optimizing this additional objective trades model performance against the time cost of solving the learned dynamics. We demonstrate our approach by training substantially faster, while nearly as accurate, models in supervised classification, density estimation, and time-series modelling tasks.

1 Introduction

Because ODE solutions often require numerical approximation, neural ODEs can incur substantial solver cost; this paper proposes penalizing trajectory derivatives to make learned dynamics easier to solve. Across most instances, solver speed can be approximately doubled without substantially increasing training loss.

2 Background

  • ODE solutions usually lack closed forms and require numerical approximation.Differential equations specify a system’s instantaneous dynamics.Source: §1 · p. 1

1 Introduction

  • Neural ODEs with nearly indistinguishable predictions can have substantially different dynamics, while standard training methods do not penalize their complexity.The paper proposes finding equivalent models that are easier and faster to solve.Source: §1 · p. 1
  • Adaptive solvers’ main time cost is repeated dynamics-function evaluation, but their integer-valued evaluation count requires a differentiable surrogate for optimization.The goal is to reduce the number of function evaluations needed to reach a given error tolerance.Source: §1 · p. 2
  • For an order-m adaptive Runge-Kutta solver, larger local mth total derivatives imply smaller steps and more evaluations, motivating their penalization during training.The proposed regularizer minimizes the norm of this total derivative to control the time required to solve learned dynamics.Source: §1 · p. 2
  • In most instances, solver speed can be approximately doubled without substantially increasing training loss.The authors also provide Taylor-mode automatic differentiation, which is asymptotically more efficient for computing the required total derivatives than standard nested gradients.Source: §1 · p. 2

2 Background

  • Runge-Kutta methods approximate ODE trajectories in steps, while neural ODEs can be slower than fixed-depth networks because adaptive solvers’ step counts are uncontrolled.For an mth-order method, local error grows approximately with the size of the mth Taylor coefficient of the true solution.Source: §2 · p. 2

3 Regularizing Higher-Order Derivatives for Speed

The paper regularizes the squared norm of a chosen order’s total time derivative along the solution trajectory, targeting solver step size without unnecessarily constraining lower-order dynamics. Taylor-mode automatic differentiation makes the required higher derivatives more efficient to compute.

3 Regularizing Higher-Order Derivatives for Speed

  • The regularizer integrates the squared norm of the Kth total time derivative of the state along the solution trajectory, weighted by λ in the training objective.Its dependence on the dynamics parameters is implicit through the solution.Source: §3 · p. 3
  • Setting the Kth-order dynamics to zero also sets higher orders to zero, while leaving lower-order dynamics unpenalized; for example, a quadratic trajectory has R3 = 0.On a toy one-dimensional neural ODE, regularizing R3 reduces NFE.Source: §3 · p. 3
  • The proposed rule is to match regularization order to solver order; the figure indicates that m-order Runge-Kutta solvers need small steps when derivatives of order K ≥ m are nonzero.The authors conjecture that regularizing lower orders than the solver order restricts the model unnecessarily.Source: §3 · p. 3, Figure 2 · p. 3

4 Efficient Higher Order Differentiation with Taylor Mode

  • Taylor mode computes total derivatives exactly at O(K^2) cost, and its asymptotic improvement reduced wall-clock time by an order of magnitude even for K as low as 3.The authors provide an open-source Taylor-mode AD implementation in JAX.Source: §4 · pp. 3–4
  • Naively nesting first-order JVPs computes higher derivatives at O(exp(K)) cost because lower-order expressions are recomputed rather than shared.Taylor mode propagates truncated Taylor series to compute derivatives with cost O(K^2) or O(K log K), depending on the operations involved.Source: §4 · p. 4

5 Experiments

Across supervised learning, time-series modeling, and density estimation, speed regularization reduced solver evaluations with limited changes in model performance. The experiments also present Taylor-coefficient propagation rules for standard functions.

5 Experiments

  • Table 1 gives rules for propagating Taylor polynomial coefficients through standard functions, generalizing standard first-order derivatives.Source: Table 1 · p. 4

5.1 Supervised Learning

  • Speed regularization decreased NFE throughout supervised training without substantially changing training error.The experiment compares MNIST classifiers trained with and without R3 regularization.Source: Figure 3 · p. 4, §5.1 · p. 4

5.2 Continuous Generative Time Series Models

  • Average NFE fell from 281 to 90 in the PhysioNet latent ODE, with an 8% increase in loss.The figure shows a representative two-dimensional slice of the 20-dimensional dynamics.Source: Figure 4 · p. 4

5.3 Density Estimation with Continuous Normalizing Flows

  • NFE decreased by 38% for a 0.6% increase in log-likelihood measured in bits/dim on density estimation.Source: §5.3 · p. 5

6 Analysis and Discussion

The analysis finds a tradeoff between solver speed and training loss, with substantial performance degradation generally appearing only after a 50% reduction in solver speed. It also examines regularization order, training time, solver calibration, and generalization to held-out data.

6.1 Trading off function evaluations for loss

  • Substantial performance degradation generally begins only after a 50% reduction in solver speed when controlling R2 across three neural ODE applications.Figure 5 plots average function evaluations against unregularized training loss at the end of training.Source: §6.1 · p. 5, Figure 5 · p. 5

6.2 Order of regularization vs. order of solver

  • For a 2nd-order solver, K = 2 gives a strictly better speed–performance tradeoff; regularizing above K = 3 gives little benefit for higher-order solvers.The study partially supports matching regularizer order to solver order; Figure 6 reports mixed results for a 5th-order solver.Source: §6.2 · p. 5, Figure 6 · p. 6

6.3 Do we reduce training time?

  • Adaptive-solver training does not improve overall training time because computing the regularizer adds expense, while fixed-grid training is faster but can be unstable without regularization.For FFJORD, this method is 2.4× slower than Finlay et al. (2020); for MNIST classification, it is 1.7× slower but achieves better test-time NFE.Source: §6.3 · p. 6

6.4 Are we making the solver overconfident?

  • On MNIST classification, speed regularization does not make the solver overconfident in its extrapolation.Figure 8a reports that actual solver error is about equally well-calibrated for regularized and random dynamics.Source: §6.4 · p. 6

6.5 Does speed regularization overfit?

  • On MNIST classification, regularized dynamics have nearly identical average solve time on held-out test data and training data.The authors report this as evidence that speed regularization does not overfit solver speed in this setting.Source: §6.5 · p. 6

7 Related Work

Related work addresses regularization, fixed-step solvers, alternative neural differential equation solvers, and other ways to stabilize or speed up learned dynamics. The authors note that their initial speedups do not make neural ODEs competitive with fixed-depth architectures for standard supervised learning, while identifying continuous-time applications as a possible use case.

7 Related Work

  • Finlay et al. (2020) regularized continuous-normalizing-flow trajectories and introduced fixed-step solvers for stable, faster training; Onken & Ruthotto (2020) also explored fixed-step solvers.Source: §7 · pp. 6–7
  • Other approaches speed solving or backpropagation with neural-network-parameterized solvers (Poli et al. (2020)) or by exploiting adjoint-equation structure to reduce NFE (Kidger et al. (2020a)).Morrill et al. (2020) proposed a speed method for neural controlled differential equations designed for irregularly sampled time series.Source: §7 · p. 7
  • Prior stabilization approaches include constraining the largest-eigenvalue Jacobian direction in recurrent networks, GRU-like neural ODE parameterizations, and adding state dimensions to ease training and reduce NFE.These examples are attributed to Simard et al. (1991), De Brouwer et al. (2019), and Dupont et al. (2019), respectively.Source: §7 · p. 7
  • Chang et al. (2017) used the ResNet–ODE connection to gradually deepen ResNets during training; increasing NFE in neural ODEs is described as an automatic but uncontrolled version.The authors suggest a speed-regularization schedule that gradually tapers off during training.Source: §7 · pp. 7–8
  • Table 2 reports that the proposed approach is slowest to train with adaptive FFJORD solvers but requires the fewest NFE after training; fixed-step comparisons show lower bits/dim and NFE at equal step counts.Fixed-step solver runs that diverged are marked as NaN bits/dim.Source: Table 2 · p. 7

8 Scope

  • The initial speedups are insufficient to make neural ODEs competitive in speed with standard fixed-depth architectures for standard supervised learning.The authors identify applications already using continuous-time models, including density and time-series models, as potential beneficiaries.Source: §8 · p. 8
  • The paper investigates ODEs; extension to neural stochastic differential equations and other gradient-fit parametric differential equations is presented as a presumed possibility.Source: §8 · p. 8

9 Limitations

The paper reports speed–performance tradeoffs and several choices that constrain their use: regularization, solver settings, and derivative order require tuning. Its conclusions also identify potential future solver designs and describe the work’s broader impact as indirect.

9 Limitations

  • The regularization weight λ must be chosen to balance speed against training loss; neural ODEs avoid choosing the outer layer count required by standard layer stacks.Source: §9 · p. 8
  • Solver order and tolerances must be selected and can substantially affect speed; loosening tolerances was not investigated, while atol and rtol of 1.4e-8 worked in all experiments.Source: §9 · p. 8
  • The derivative order K also requires selection: higher K generally works better but makes each training step slower, and automatic λ tuning is only a possibility when utility is explicit in training loss and NFE.Source: §9 · p. 8

10 Conclusions

  • The work is an initial attempt at controlling integration time through dynamics regularization, and the authors expect better quantities to optimize may exist.Source: §10 · p. 8
  • Across tasks, tuned regularization usually gave at least a 2x speedup without substantially hurting model performance.Source: §10 · p. 8
  • Overall training time with speed regularization is generally 30% to 50% slower with adaptive solvers.Source: §10 · p. 8
  • For standard solvers, regularizing above R2 or R3 provided little additional benefit.Source: §10 · p. 8
  • Future work could adapt solver architectures to reject dynamics that are difficult to solve; the work’s potential impact is through improved general-purpose machine-learning tools, with unintended consequences not further speculated on.Source: §10 · p. 9, Broader Impact · p. 9

Appendix A Taylor-mode Automatic Differentiation

The appendix distinguishes normalized Taylor coefficients from derivative coefficients and explains how higher derivatives of a composition can be computed from lower-order terms. It describes the resulting coefficient dependencies and JAX interface.

A.1 Taylor Polynomials

  • For a Taylor expansion, the derivative coefficient is the corresponding Taylor coefficient multiplied by its factorial normalization.The passage states that i! times the Taylor coefficient gives the derivative coefficient.Source: §A.1 · p. 12
  • The third Taylor coefficient is expressed as a function of the input coefficients through order three.Source: §A.1 · p. 12
  • The appendix makes the relationship between Taylor coefficients and derivatives explicit, unlike the normalized-coefficient presentation it attributes to Griewank & Walther (2008).Source: §A.1 · p. 12
  • Each derivative coefficient y_i is the corresponding higher-order derivative of the composition f(x(t)), and its expression reuses terms computed at lower orders.Source: §A.1 · p. 13
  • The coefficient y_k depends on all input derivatives through order k, and the JAX jet API returns the k-output for the supplied inputs.Source: §A.1 · p. 13

A.2 Relationship with Differential Equations

The appendix converts an initial value problem into an autonomous dynamical system by augmenting it with the independent variable, making the intended total time derivatives unambiguous.

A.2.1 Autonomous Form

  • Augmenting the system with the independent variable’s trivial dynamics makes derivatives with respect to time explicitly total derivatives, avoiding ambiguity with partial derivatives through f’s time argument.Source: §A.2.1 · pp. 13–14

A.2.2 Taylor Coefficients for ODE Solution with jet

The appendix describes recursively computing local solution Taylor coefficients, motivates regularizing higher-order terms, and outlines gradient computation and supervised-learning experiments. In the toy supervised task, regularizing sixth-order trajectory derivatives reduced solver evaluations.

A.2.2 Taylor Coefficients for ODE Solution with jet

  • Recursive use of jet with x_(k+1) = y_k computes solution-polynomial coefficients; the truncation order determines the local approximation order.Source: §A.2.2 · p. 14, §A.3 · p. 14

A.3 Regularizing Taylor Terms

  • Large higher-order Taylor coefficients make a truncated local approximation diverge more quickly, motivating regularization of those terms.Source: §A.3 · p. 14

Appendix B Experimental Details

  • The experiments used GPU-based ODE solvers and the adjoint method to reconstruct trajectories backward for training gradients; the regularization term was normalized by trajectory dimension.Source: Appendix B · p. 14

B.1 Efficient computation of the gradient of regularization term

  • The adjoint method computes the regularization gradient through an augmented ODE, and the weighted regularizer is added to the original loss.Source: §B.1 · p. 14

B.2 Supervised Learning

  • Regularizing sixth-order trajectory derivatives reduced the toy supervised-learning solve from 92 to 68 NFE with a fifth-order Runge–Kutta solver.Source: Figure 9 · p. 15
  • The MNIST supervised model used a 100-hidden-unit MLP and was trained for 160 epochs on 60,000 training images, with 10,000 test images used for validation and testing.Source: §B.2 · p. 15

B.3 Continuous Generative Modelling of Time-Series

The appendix details time-series and flow experiment settings, then evaluates whether NFE regularization tracks solve-count differences and preserves tradeoffs on losses closer to the tasks’ evaluation objectives.

B.3 Continuous Generative Modelling of Time-Series

  • PhysioNet uses 8,000 trajectories with 41 time-varying traits, hourly averaging over 48 hours into 49 timestamps, and a random 20% validation split.The coarse quantization was chosen for computational reasons related to this model implementation.Source: §B.3 · p. 15

B.4 Continuous Normalizing Flows

  • MNIST used a single-flow model, unlike Finlay et al. (2020)’s multi-flow comparison, accounting for discrepancies in reported bits/dim and NFE.For MINIBOONE, training used 500 epochs and a manually tuned learning-rate schedule based on validation loss, unlike the automatic tuning in Grathwohl et al. (2019).Source: §B.4 · p. 16

B.5 Hardware

  • MNIST supervised learning, PhysioNet time-series, and MNIST FFJORD experiments used NVIDIA Tesla P100 GPUs; all but MNIST FFJORD used double precision.Tabular FFJORD was trained on a Tesla T4 and evaluated on a Tesla P100.Source: §B.5 · p. 16

C.1 Overfitting of NFE

  • Variance in NFE across train and test examples closely tracks their absolute NFE differences, while its relationship with regularization is non-monotonic.The authors suggest average train–test NFE differences may reflect noise in estimating the true averages.Source: §C.1 · p. 16, Figure 10 · p. 16

C.2 Trading off function evaluations with a surrogate loss

  • MNIST classification and time-series modelling retain similar Pareto tradeoffs when evaluated with classification error and MSE rather than cross-entropy and IWAE loss.The time-series curves compare R3 and R2, respectively.Source: §C.2 · pp. 16–17, Figure 11 · p. 17

C.3 Wall-clock Time

The appendix points to wall-clock and fixed-grid solver results in tables, then contrasts R2 with the regularization terms of Finlay et al. (2020), including for autonomous systems.

C.3 Wall-clock Time

  • Wall-clock times and fixed-grid solver results are provided in Tables 3 and 4, whose captions identify MNIST classification and MINIBOONE density estimation.The supplied passages do not report the tables’ numerical outcomes.Source: §C.3 · p. 17, Table 3 · p. 17, Table 4 · p. 18

Appendix D Comparison to How to Train Your Neural ODE

  • Unlike the first-total-derivative and spatial-derivative terms of Finlay et al. (2020), R2 directly regularizes the second total derivative, including explicit time dependence.The distinction remains even for autonomous systems, where ∂f/∂t is zero; Finlay et al.’s terms separately penalize the norms of f and ∇_z f.Source: Appendix D · p. 18

Related papers

Loading 2007.04504v2…