PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Checkpoint Adjoint Method for Gradient Estimation in Neural ODE

Juntang Zhuang, Nicha Dvornek, Xiaoxiao Li, Sekhar Tatikonda, Xenophon Papademetris, James Duncan

arXiv:2006.02493v1stat.MLcs.LG

TL;DR

NODEs lag discrete-layer models on benchmark classification, while existing gradient estimators have numerical-error or computation-graph drawbacks. ACA checkpoints the forward trajectory and removes redundant backward computation; on image classification it halves error and training time versus adjoint and naive methods, with further gains reported against ResNet and time-series competitors.

  • Problem

    NODE benchmark classification trails discrete-layer models, while the adjoint method has reverse-integration numerical errors and the naive method creates a redundant graph during adaptive step-size search.

    Source: §1 · p. 1

  • Method

    ACA uses checkpointed forward trajectories as reverse-mode trajectories, removes redundant backward-pass components, and supports adaptive ODE solvers.

    Source: §1 · p. 1

  • Results

    Twice the speed and half the error rate versus adjoint and naive methods on image classification; ACA-enabled NODE also outperforms ResNet in accuracy and reliability, and other methods in time-series modeling.

    Source: §1 · p. 2

  • Takeaways & Limitations

    NODE trained with ACA can incorporate physical knowledge for better accuracy, and the authors report that ACA applies to general ODEs.

    Source: §1 · p. 2, §7 · p. 9

    The study investigates only explicit general-purpose one-step ODE solvers, leaving multi-step and implicit solvers outside its scope.

    Source: §5 · p. 9

Abstract

from arXiv · show

Neural ordinary differential equations (NODEs) have recently attracted increasing attention; however, their empirical performance on benchmark tasks (e.g. image classification) are significantly inferior to discrete-layer models. We demonstrate an explanation for their poorer performance is the inaccuracy of existing gradient estimation methods: the adjoint method has numerical errors in reverse-mode integration; the naive method directly back-propagates through ODE solvers, but suffers from a redundantly deep computation graph when searching for the optimal stepsize. We propose the Adaptive Checkpoint Adjoint (ACA) method: in automatic differentiation, ACA applies a trajectory checkpoint strategy which records the forward-mode trajectory as the reverse-mode trajectory to guarantee accuracy; ACA deletes redundant components for shallow computation graphs; and ACA supports adaptive solvers. On image classification tasks, compared with the adjoint and naive method, ACA achieves half the error rate in half the training time; NODE trained with ACA outperforms ResNet in both accuracy and test-retest reliability. On time-series modeling, ACA outperforms competing methods. Finally, in an example of the three-body problem, we show NODE with ACA can incorporate physical knowledge to achieve better accuracy. We provide the PyTorch implementation of ACA: \url{https://github.com/juntang-zhuang/torch-ACA}.

1. Introduction

NODEs trail discrete-layer models on benchmark classification, and existing gradient estimators have numerical-error or computation-graph drawbacks. The paper proposes ACA and reports stronger performance across image classification, time-series modeling, and a three-body example.

1. Introduction

  • NODE benchmark classification lags discrete-layer models, while existing adjoint and naive gradient methods have numerical-error and redundant-computation drawbacks.The introduction gives CIFAR10 error rates of 19% for NODE and 5% for ResNet.Source: §1 · p. 1
  • ACA checkpoints the forward trajectory for reverse-mode integration, deletes redundant backward-pass components, and supports adaptive ODE solvers.Using the forward trajectory as the reverse trajectory is intended to guarantee numerical accuracy.Source: §1 · p. 1
  • On image classification, ACA achieves twice the speed and half the error rate versus adjoint and naive methods, and enables NODEs to outperform ResNet in accuracy and test-retest reliability.The authors also report that ACA outperforms other methods on time-series modeling and helps NODE incorporate physical knowledge for better accuracy.Source: §1 · p. 2

2. Preliminaries

NODE represents a model as a continuous-time initial-value problem, with its input as the initial condition and the terminal hidden state as output. Gradient calculation requires both the forward state and adjoint across time, while numerical solvers advance the state and adapt step size using error estimates.

2.1. Neural Ordinary Differential Equations

  • NODE models continuous-depth feature evolution with a neural-network-parameterized ODE, taking input x as the initial condition and returning hidden state z(T).The model's feature map evolves smoothly with depth, unlike the integer-depth representation of discrete-layer models.Source: §2.1 · p. 2

2.2. Analytical Form of Gradient for NODE

  • Gradient calculation solves the state forward, sets the terminal adjoint, solves the adjoint backward, then calculates the parameter gradient; both state and adjoint are needed at every time.Because the state and adjoint are solved in opposite directions, the state must be stored or recovered.Source: §2.2 · p. 2

2.3. Numerical Integration

  • Adaptive solvers advance time using error estimates to adjust step size; repeated step-size trials create a redundant computation graph, while global error accumulates propagated local errors.Adaptive step counts vary with input and error tolerance.Source: §2.3 · pp. 2–3, Figure 2 · p. 3

3. Methods

This section compares numerical gradient methods for NODEs: the adjoint method saves memory but can incur reverse-time errors, while the naive method creates a deep graph. ACA checkpoints the forward trajectory and removes redundant graph components.

3.1. Summary of Different Methods

  • The adjoint method forgets the forward trajectory and reconstructs it in reverse, using O(Nf) memory but incurring extra computation and gradient error; the naive method has graph depth O(Nf × Nt × m).The adjoint method’s reverse-time trajectory can differ from the forward trajectory.Source: §3.1 · pp. 3–4
  • ACA stores accepted time points and state values, then reuses saved stepsizes during the backward pass to align reverse- and forward-time trajectories.Its stated memory cost is O(Nf + Nt), and its computation cost is O(Nf × Nt × (m + 1)).Source: §3.1 · p. 4

3.2. Adjoint Method has Numerical Errors

  • The van der Pol equation and a convolution-function ODE serve as numerical examples validating the paper’s argument about adjoint-method error.The figures show forward and reverse trajectories for the van der Pol equation and the input and reverse-time reconstruction for the convolution function.Source: §3.2 · p. 4, Figure 4 · p. 4, Figure 5 · p. 4
  • Reverse accuracy requires e_k = 0 for all k; odd-order solvers cannot satisfy this, while even-order solvers require an identity mapping, so truncation errors cause adjoint inaccuracies.The paper assumes the Picard-Lindelöf conditions hold over the integration interval for its analysis.Source: §3.2 · pp. 4–5

3.3. Naive Method has Deep Computation Graph

  • When stepsize search takes m iterations, naive backpropagation traverses all m, whereas ACA treats the accepted stepsize as constant, yielding graph depths O(Nf × Nt × m) and O(Nf × Nt), respectively.The paper also notes that the naive method’s deep graph takes more memory and might cause vanishing or exploding gradients.Source: §3.3 · p. 5

3.4. ACA Guarantees Reverse-accuracy and has Shallow Computation Graph

  • ACA guarantees reverse-trajectory accuracy, has a shallower graph and medium memory cost, and is reported to have the lowest computation cost among the compared methods.The adjoint method’s reverse-time trajectory can be inaccurate, while the naive method has a deeper graph and higher memory cost.Source: §3.4 · p. 6, Table 1 · p. 6

4. Experiments

Across the experiments, ACA improves NODE gradient estimation and benchmark performance, including image classification and irregularly sampled time series. The three-body example compares models using no, partial, and full physical knowledge.

4.1. Toy Example

  • ACA consistently yields lower toy-example numerical error than the naive and adjoint methods across the tested T values.The methods use Dopri5 with error tolerance 10^-5.Source: §4.1 · p. 6

4.2. Supervised Learning on Image Classification

  • On CIFAR10, for the same NODE model and training epochs, ACA achieves about 5% test error versus about 10% for adjoint and naive methods, substantially outperforming both.For 90 epochs on a single GTX 1080Ti GPU, ACA takes about 9 hours, versus about 18 hours for adjoint and more than 30 hours for naive.Source: §4.2 · p. 6, Figure 7 · p. 7
  • On CIFAR10 and CIFAR100, NODE trained with ACA significantly outperforms ResNet in the reported 90-epoch comparison.The models use the same number of parameters; the figure shows distributions across 10 runs.Source: §4.2 · p. 6, Figure 7 · p. 7
  • Changing ODE solver at inference increases NODE-ACA error by only about 1%, while the paper reports that solver changes generally cause large errors for discrete-layer models.NODE-ACA is tested with different solvers without retraining.Source: §4.2 · p. 7, Table 2 · p. 7
  • NODE-ACA has higher test-retest ICC than ResNet on both datasets, including when ICC is measured only on misclassified examples.The comparison reports ICC1 and ICC1k from 10 independently initialized runs.Source: §4.2 · p. 7, Table 3 · p. 7

4.3. Time-series Model for Irregularly-sampled Data

  • On the Mujoco irregular-time-series task, ACA consistently outperforms the other gradient-estimation methods on interpolation MSE.The comparison uses the setting of Rubanova et al. (2019), changing only the gradient-estimation method.Source: Table 4 · p. 7, §4.3 · p. 8

4.4. Incorporate Physical Knowledge into Models

  • For the chaotic three-body system with unequal masses and arbitrary initial conditions, LSTM fails with no knowledge; with partial knowledge, NODE-ACA outperforms LSTM, and with full knowledge, ODE performs best.Training uses [0,1] year, and mean trajectory MSE is measured on [0,2] years.Source: §4.4 · p. 8

5. Scopes and limitations

The study scope is limited to explicit, general-purpose one-step ODE solvers.

5. Scopes and limitations

  • The project investigates only explicit general-purpose one-step ODE solvers, leaving multi-step and implicit solvers outside its scope.The authors note that acceleration methods could further improve ACA.Source: §5 · p. 9

6. Related works

The related-work discussion contrasts ACA with methods for NODE training, physical modeling, and gradient checkpointing.

Training of NODE

  • Unlike spectral-element NODE training, which requires ground truth for the full trajectory, ACA supports adaptive solvers; unlike ANODE's fixed-step discretization, it retains adaptive error control.The discussion also distinguishes augmented NODE work as addressing performance rather than gradient estimation.Source: §6 · p. 9

Dynamics and Physics

  • The authors distinguish their three-body experiments from Breen et al. (2019), whose methods are limited to equal masses and zero initial velocities.The passage also situates the work among studies of physical priors and dynamical systems.Source: §6 · p. 9

Gradient Checkpointing

  • ACA is distinguished from naive gradient checkpointing: it simplifies graph depth to O(Nf × Nt), whereas naive-GC retains O(Nf × Nt × m) depth.The passage states that naive-GC shares naive's exploding-or-vanishing-gradient issue, while ACA is more numerically stable.Source: §6 · p. 9

7. Conclusion

The conclusion presents ACA as an accurate gradient-estimation method for NODE and summarizes its practical findings.

7. Conclusion

  • The paper concludes that ACA enables accurate NODE gradient estimation, with reported accuracy, speed, and initialization robustness, and that NODE can use physical knowledge for better accuracy.The authors also report implementing ACA as a package.Source: §7 · p. 9

Appendix A, ACA as an AutoDiff Function in PyTorch Style

ACA’s forward pass adaptively solves the ODE and caches accepted states and times. Its backward pass replays those cached steps to compute gradients.

Appendix A, ACA as an AutoDiff Function in PyTorch Style

  • The forward pass initializes the solver, adapts step size, accepts steps within tolerance, and caches the resulting trajectory and time points.Rejected steps reduce the step size; their error-estimate computation graph is deleted.Source: Appendix A · p. 12
  • The backward pass iterates over cached times in reverse, locally recomputes each forward step, and computes gradients while deleting local computation graphs.It initializes the adjoint from the terminal loss gradient and accumulates parameter gradients.Source: Appendix A · p. 12

Appendix B, Proof of Theorem 3.2

The proof analyzes numerical error when an ODE solution is integrated forward and then backward. It concludes that the adjoint method generally cannot reconstruct the initial state exactly.

Appendix B, Proof of Theorem 3.2

  • For a solver of order p, the local truncation error is O(h^(p+1)).The analysis assumes matching forward and reverse grid points with a sufficiently small constant step size.Source: Appendix B · pp. 13–14
  • The local-error terms generally cannot cancel; reverse accuracy for every starting time requires each step’s error contribution to be zero.For odd solver order the terms have the same sign; for even order, the stated condition requires the variational flow to be identity, which is generally not satisfied by a trained network.Source: Appendix B · pp. 14–15
  • The forward-then-reverse solve generally reconstructs an initial state different from the original because of numerical integration errors.The proof denotes the discrepancy as the combined forward and reverse global error.Source: Appendix B · pp. 14–15

Appendix C. Proof of Theorem 2.1

The appendix derives NODE gradients by formulating training as an optimization problem and applying the Lagrangian multiplier method. It outlines the forward, backward, and gradient-computation stages.

Appendix C. Proof of Theorem 2.1

  • The derivation frames NODE training as an optimization problem and uses Lagrangian multipliers and KKT conditions to obtain gradient conditions.The appendix introduces perturbations of the adjoint and hidden state in deriving those conditions.Source: Appendix C · pp. 15–16
  • The derivation treats one ODE block for simplicity and states that the analysis extends to chained blocks when the gradient at the block output is given.For multiple blocks, the output state of one block is the input to the next.Source: Appendix C · p. 15
  • The stated gradient procedure solves the ODE forward, sets a boundary condition, solves backward, and then calculates the gradient; λ corresponds to the negative adjoint variable.The appendix assumes the vector field is differentiable almost everywhere and defines the forward pass using the ODE.Source: Appendix C · pp. 15, 17

1. Experiments with van der Pol Equation

The experiments cover a van der Pol simulation, image classification, irregularly sampled time series, and three-body trajectory modeling. The supplied passages specify some experimental settings and image-classification comparisons, while several figures identify results without stating their outcomes.

1. Experiments with van der Pol Equation

  • The van der Pol experiment uses initial conditions y1(0) = 2 and y2(0) = 0 with MATLAB’s ode45 solver under default settings.Figure 9 is identified as the simulation results.Source: §1 · p. 18, Figure 9 · p. 18

2. Experiments on supervised image classification

  • NODE’s same-method train/test error rate is 4.85% versus ResNet’s 8.47%; solver changes raise error by ∼1% for NODE and ∼7% for ResNet in most tested cases.The comparison tests models with different solvers without retraining; smaller error increases indicate greater robustness.Source: Table 7 · p. 19, Table 6 · p. 19, §2 · p. 19

3. Experiments on time-series modeling with irregularly sampled data

  • The time-series experiments train models for 300 epochs on Mujoco; Figures 10–12 identify evaluations using 10%, 20%, and 50% of training data.The supplied passages do not state the comparative outcomes shown in those figures.Source: §3 · p. 19, Figure 10 · p. 19, Figure 11 · p. 19

4. Experiments on the three-body problem

  • The three-body experiment compares fitted and ground-truth trajectories for three planets across LSTM, NODE, and ODE model variants.The supplied figure captions encode fitted trajectories as orange dashed lines and ground truth as blue solid lines; they do not report which model fits best.Source: Figure 13 · p. 20, Figure 14 · p. 21, Figure 15 · p. 21

More results of ODE-ACA with different initial conditions

Loading 2006.02493v1…