[Paper Review] ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs

Amir Gholami, Kurt Keutzer+1 more|arXiv (Cornell University)|Feb 27, 2019

Model Reduction and Neural Networks26 references77 citations

TL;DR

ANODE introduces a discretize-then-optimize adjoint framework with checkpointing to compute unconditionally accurate gradients for neural ODEs, reducing memory from O(LN_t) to O(L)+O(N_t) and avoiding instability of prior reverse-time approaches.

ABSTRACT

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.

Motivation & Objective

• Motivate and analyze memory challenges in training neural ODEs and identify failures of reverse-time gradient methods.
• Propose ANODE, a checkpointing-based adjoint framework that achieves unconditionally accurate gradients.
• Demonstrate memory efficiency and stability on CIFAR-10/100 with ResNet/SqueezeNext-based ODE blocks.

Proposed method

• Model neural networks as ODEs with residual blocks and forward Euler or other discretizations.
• Show that reversing forward ODEs to compute gradients yields numerical and consistency issues.
• Adopt a discretize-then-optimize gradient computation with checkpointing to store only necessary activations.
• Define a DTO-based adjoint computation that reuses stored forward trajectories to compute gradients.
• Provide a memory-management scheme that reduces storage from O(LN_t) to O(L)+O(N_t) with optional logarithmic checkpointing.
• Experiment with ANODE on CIFAR-10/100 using ResNet and SqueezeNext backbones with Euler and RK2 schemes.

Experimental results

Research questions

• RQ1 Do neural ODEs admit unconditionally accurate gradients under memory-constrained training?
• RQ2 Can a discretize-then-optimize adjoint framework with checkpointing provide stable and accurate gradients for general NN blocks (including ReLU activations and convolutions)?
• RQ3 What is the memory vs. compute trade-off of ANODE compared to prior reverse-time methods for backprop through ODE layers?
• RQ4 How do discretization choices (Euler, RK2) impact training stability and accuracy in neural ODEs?

Key findings

• Reverse-time (backward) ODE solves can be numerically unstable and yield incorrect gradients for general networks.
• Optimize-then-discretize (OTD) gradients can be inconsistent with discretized forward solves, causing training divergence with small time steps.
• Discretize-then-optimize with checkpointing (ANODE) yields unconditionally accurate gradients and stable training.
• ANODE reduces memory from O(LN_t) to O(L)+O(N_t) with the same forward computational cost as the opposing method.
• Using discrete DTO with checkpointing enables efficient gradient computation without constraining weight norms for stability.
• Experiments on CIFAR-10/100 with ResNet/SqueezeNext ODE blocks show ANODE achieves stable convergence and better performance than the reverse-time baseline.