PaperlayerView exact revision on arXiv

Source-linked AI summary

ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs

Amir Gholami, Kurt Keutzer, George Biros

arXiv:1902.10298v3cs.LG

TL;DR

Neural ODE backpropagation is memory-intensive because gradients require intermediate trajectories, while reverse-time reconstruction can be unstable or inconsistent. The paper proposes ANODE, which uses checkpointing with Discretize-Then-Optimize differentiation to obtain accurate gradients at lower memory cost, while forward-solve stability remains outside its guarantee.

  • Problem

    Neural ODE gradient computation requires storing intermediate activations, increasing memory from O(L) to O(LNt) across L blocks and Nt time steps.

  • Method

    ANODE uses checkpointing with Discretize-Then-Optimize differentiation instead of relying on reverse-time reconstruction of the forward ODE.

  • Results

    ANODE has memory footprint O(L)+O(Nt) and computes correct numerical gradients consistent with the discretization scheme, including for ReLU activations.

  • Takeaways & Limitations

    Checkpointing addresses neural ODE memory costs without the reverse-solve instability and gradient inconsistency analyzed in the paper.

  • Takeaways & Limitations

    ANODE does not guarantee stability of the forward ODE solve, which must be enforced through architecture or regularization.

Abstract

from arXiv · show

Residual neural networks can be viewed as the forward Euler discretization of an Ordinary Differential Equation (ODE) with a unit time step. This has recently motivated researchers to explore other discretization approaches and train ODE based networks. However, an important challenge of neural ODEs is their prohibitive memory cost during gradient backpropogation. Recently a method proposed in [8], claimed that this memory overhead can be reduced from O(LN_t), where N_t is the number of time steps, down to O(L) by solving forward ODE backwards in time, where L is the depth of the network. However, we will show that this approach may lead to several problems: (i) it may be numerically unstable for ReLU/non-ReLU activations and general convolution operators, and (ii) the proposed optimize-then-discretize approach may lead to divergent training due to inconsistent gradients for small time step sizes. We discuss the underlying problems, and to address them we propose ANODE, an Adjoint based Neural ODE framework which avoids the numerical instability related problems noted above, and provides unconditionally accurate gradients. ANODE has a memory footprint of O(L) + O(N_t), with the same computational cost as reversing ODE solve. We furthermore, discuss a memory efficient algorithm which can further reduce this footprint with a trade-off of additional computational cost. We show results on Cifar-10/100 datasets using ResNet and SqueezeNext neural networks.

I. INTRODUCTION

Neural ODEs connect residual blocks to ODE discretizations but create substantial backpropagation memory costs. Reversing the forward solve can reduce storage, yet may introduce instability, incorrect gradients, and divergent training; ANODE addresses these issues with checkpointing.

  • ODE perspective: A residual block is equivalent to an autonomous ODE solved with forward Euler using one time step.The block computes z1 = z0 + f(z0, θ), corresponding to an Euler discretization of dz/dt = f(z(t), θ).
  • Memory challenge: O(LNt) storage is required when gradients retain all intermediate activations across L ODE blocks and Nt time steps.The memory burden arises because adjoint-gradient computation needs the activation trajectory z(t).
  • Reverse solve: Reverse-time reconstruction was proposed to reduce memory from O(LNt) to O(L), but it can produce O(1) gradient errors and divergent training.The approach assumes activations can be recovered by solving the forward ODE backward, which may fail or accumulate discretization errors.
  • Reverse solve: Reversing neural ODEs may be unstable for ReLU and general convolution or activation operators, including when adaptive time stepping is used.The paper illustrates this instability with residual blocks using convolution and ReLU or Leaky ReLU activations.
  • ANODE: ANODE uses checkpointing and Discretize-Then-Optimize differentiation to avoid reverse-solve instability while reducing memory to O(L)+O(Nt).The footprint can be reduced further through logarithmic checkpointing at additional computational cost.

A. Related Work

Related work connects neural networks with ODEs, reversible architectures, adjoint methods, and checkpointing strategies for reducing memory during backpropagation.

  • ODEs: Prior work viewed residual networks as forward Euler discretizations of ODEs and studied stable or reversible network architectures.These architectures include discrete Hamiltonian systems and reversible residual networks for avoiding intermediate activation storage.
  • Checkpointing: Checkpointing was applied to long recurrent networks and memory-efficient neural networks by storing selected activations and recomputing intermediates during backpropagation.The cited approaches include square-root checkpointing and earlier optimal-checkpointing methods.

II. PROBLEM DESCRIPTION

The paper formulates neural-network training as empirical-loss optimization in which ODE layers are solved forward and differentiated through during backpropagation. The central implementation difficulty is the memory required by those gradients.

  • Training objective: Training learns parameters θ for a mapping F(θ) using empirical loss over m examples, with optional regularization R(θ).The data are sampled from an unknown distribution, so optimization uses the available training set.
  • Forward and backward passes: Forward evaluation solves each ODE layer and then computes predictions and a specified loss against the ground-truth labels.The resulting prediction may be a vector of class probabilities, and cross entropy is given as an example loss.
  • Forward and backward passes: Backpropagating through ODE integration requires solving an adjoint equation and can make memory prohibitive even for shallow neural networks.The forward operator involves numerical integration, whose intermediate activations must support gradient computation.

A. Adjoint Based Backpropogation

Adjoint backpropagation computes parameter and input gradients by solving an adjoint ODE, but this requires access to the full activation trajectory. Storing those trajectories causes O(LNt) memory scaling, motivating checkpointing alternatives.

  • Adjoint formulation: For an ODE layer, backpropagation computes gradients with respect to parameters θ and the input activation z0 from the output gradient.The adjoint variable is introduced through a Lagrangian formulation of the constrained ODE problem.
  • Adjoint computation: The adjoint system is solved backward from a terminal condition derived from the gradient with respect to the output activation.The resulting adjoint at the initial time gives the gradient with respect to the input, while parameter gradients use α(t) and z(t).
  • Memory cost: O(LNt) storage results because adjoint and parameter-gradient equations require activation values z(t) throughout each ODE trajectory.L denotes network depth and Nt denotes the number of time discretization steps.
  • Reverse-time reconstruction: Reducing storage to O(L) by reconstructing activations through reverse-time ODE solves can yield incorrect or noisy gradients for general neural networks.The method depends on reversibility and numerical accuracy of the reverse solve.

III. CAN WE REVERSE AN ODE?

ODE reversibility can hold locally under smoothness or Lipschitz conditions, but reverse-time integration may become unstable or inaccurate in practical neural networks. The resulting reconstruction errors can corrupt gradients and contribute to divergent or sub-optimal training.

  • Local reversibility: Smooth ODE flows are reversible only over a time horizon that can depend on the initial condition, so one shared horizon may not work for every training example.For smooth vector fields, the flow is locally a diffeomorphism; the paper notes that the relevant horizon depends on z0.
  • Reverse-time instability: Reversing an ODE changes the derivative spectrum's sign, so a forward-stable system can become numerically unstable backward in time.The paper illustrates this with dz/dt = λz when λ < 0: small reverse-solve errors are exponentially amplified.
  • Reverse-time instability: 200,000 time steps are required to resolve both forward and backward flows to 1% for λ = −100 over unit time.For λ = −1e4, the flow is impossible to reverse numerically in double precision; the paper relates such instability to systems including the linear heat equation.
  • Neural-network reversibility: For ReLU ODEs, ode45 produces 1% error with 11 time steps and 0.4% error with 18 time steps, while single-precision accuracy requires 211 steps.These errors are measured by composing forward and reverse flows and comparing the result with the initial state.
  • Neural-network reversibility: ReLU and repeated residual-style maps can fail to reverse reliably, and reverse-flow gradients may therefore contain O(1) errors.The paper notes that ReLU ODEs inherit numerical issues and that simply reversing the sign is insufficient for the analogous discrete map.
  • Empirical consequence: On Cifar-10 SqueezeNext, ANODE trains stably and reaches higher accuracy than the reverse-solve gradient method, whose RK45 training diverges in the first epoch.The comparison uses ODE blocks solved with Euler and RK-2 methods.

IV. OPTIMIZE-THEN-DISCRETIZE VERSUS DISCRETIZE-THEN-OPTIMIZE

The paper identifies two sources of incorrect gradients in adjoint backpropagation: reversing the ODE can fail, and optimize-then-discretize gradients may be inconsistent with the chosen discretization. These issues can produce large gradient errors and divergent or sub-optimal training.

  • Optimize-Then-Discretize: Optimize-then-discretize derives continuous adjoint conditions without accounting for the numerical discretization scheme, so correct gradient information is not guaranteed.The issue arises when finite-difference approximations are applied to continuous integral equations.
  • Empirical evidence: On Cifar-10 and Cifar-100 ResNet-18 experiments, gradients from [8] produce sub-optimal performance compared with ANODE, while RK45 causes divergent training in the first epoch.Both experiments replace nontransition blocks with ODE blocks and use Euler discretization for the main comparison.
  • Gradient inconsistency: The OTD and discretize-then-optimize gradients can differ completely because OTD evaluates the neural operator at the wrong state.Except for rare cases, their error scales as O(dt), becoming large for small time steps.
  • Discretize-Then-Optimize: Discretize-then-optimize derives optimality conditions after discretization, while self-adjoint schemes such as RK2 or implicit methods offer an alternative at higher computational cost.For non-self-adjoint schemes, the DTO-derived adjoint problem is required.

V. ANODE

ANODE uses checkpointing and discretize-then-optimize differentiation to reduce neural ODE memory use while preserving gradients consistent with the forward discretization. It stores block inputs, recomputes trajectories during backpropagation, and can trade additional recomputation for lower memory.

  • Memory-efficient differentiation: ANODE reduces memory from O(LN_t) to O(L)+O(N_t) while retaining the computational cost of the reversing ODE approach.The framework stores the input activation of each ODE block and reuses trajectory memory across blocks.
  • Checkpointing procedure: ANODE recomputes each ODE block’s intermediate trajectory, solves the DTO adjoint backward through automatic differentiation, then releases the memory before processing the next block.This multi-stage procedure preserves only the memory needed for the current block’s trajectory.
  • Gradient accuracy: The DTO procedure gives correct gradients for ReLU activations and avoids numerical instability from solving the forward ODE backward in time.Its numerical gradient follows the discretization scheme used for the forward ODE solve.
  • Additional memory reduction: Logarithmic checkpointing can reduce the O(N_t) activation-storage cost when memory is scarce, at the expense of additional forward recomputation.With only one checkpointed time step, recomputation can grow to O(N_t^2).
  • Experiments: ANODE trains stably and reaches higher accuracy than neural ODE method [8] in the reported SqueezeNext Cifar-10 experiment.The experiment uses Euler discretization and also evaluates the RK2 trapezoidal rule.

VI. CONCLUSIONS

The paper concludes that reversing ODE solves and optimize-then-discretize differentiation can cause instability and inconsistent gradients. ANODE addresses these issues through checkpointing and DTO differentiation, while forward-solve stability remains an unresolved limitation.

  • Conclusions: The analyzed adjoint method may be numerically unstable for general convolution and activation operators and may cause divergence through inconsistent OTD gradients.The paper connects these problems to reverse-time solvability and discretization mismatch.
  • Conclusions: ANODE reduces memory from O(LN_t) to O(L)+O(N_t) with the same computational cost as the neural ODE method proposed by [8].Checkpointing avoids restrictions on the weight-matrix norm that were required for numerical stability in the analyzed approach.
  • Limitations: ANODE does not guarantee stability of the forward ODE solve, which must be enforced through architecture or regularization.The paper gives Hamiltonian systems as an architectural example.

B. Optimize Then Discretize Approach

The OTD derivation yields continuous optimality conditions, but discretizing those equations can produce inconsistent and incorrect gradients. Numerical reversal can also fail because general neural ODEs are not reliably reversible.

  • Optimality conditions: The KKT conditions are obtained by taking variations of the Lagrangian with respect to the adjoint, state, and neural-network parameters.Variations with respect to the adjoint recover the state equation, while state and parameter variations produce the remaining optimality conditions.
  • OTD formulation: The resulting equations are the Optimize-Then-Discretize form of neural ODE backpropagation.They represent continuous optimality conditions before a numerical discretization is selected.
  • OTD limitations: Discretizing OTD equations can create inconsistencies between continuous and discrete equations, leading to incorrect gradients.The paper identifies this as a known issue in scientific computing and motivates deriving Discretize-Then-Optimize conditions instead.
  • Numerical instability: Adaptive RK45 reversal remains numerically unstable across no, ReLU, Leaky ReLU, and Softplus activation settings.The backward-solved output can differ completely from the input image, illustrating instability when solving the forward problem backward.

C. Discretize Then Optimize Approach

DTO derives gradients from the discretized forward solve and stores intermediate activation maps before marching the adjoint backward. This produces the discrete gradient corresponding to the chosen time-stepping scheme.

  • Discrete formulation: DTO computes the correct discrete gradient for an Euler time-stepping scheme with N_t = n + 1 forward steps.The intermediate solutions z_i are defined at time points t = iΔt.
  • Discrete formulation: The discrete Lagrangian uses adjoint variables at each time point and a loss depending directly on the final activation map z_n+1.The spatial-domain integral notation is used as a shorthand, while time is handled by a Riemann sum.
  • Backward gradient computation: Taking variations of the discrete Lagrangian yields adjoint equations for the DTO backward pass.The adjoint is initialized at the final time step and marched backward to compute the gradient with respect to the input activation.
Loading 1902.10298v3…