ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs

By A. Gholami, Kurt Keutzer, George Biros. 2019

202 citations recorded by the source corpus.

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 arXiv:1806.07366, claimed that this memory overhead can be reduced from LNt, where Nt 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, a neural ODE framework which avoids the numerical instability related problems noted above. ANODE has a memory footprint of O(L) + O(Nt), with the same computational cost as reversing ODE solve. We furthermore, discuss a memory efficient algorithm which can further reduce this footprint with a tradeoff of additional computational cost. We show results on Cifar-10/100 datasets using ResNet and SqueezeNext neural networks.

Subfields: Computational Mechanics, Statistical and Nonlinear Physics

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computational Mechanics

In Statistical and Nonlinear Physics

Citations in the corpus

Equivariant Flows: exact likelihood generative learning for symmetric densities (2020, background): “However, in general, there are no guarantees that this procedure is stable, which therefore can result in very noisy gradients, leading to long training times and inferior final results (Gholami et al., 2019).”
Continuous-in-Depth Neural Networks (2020, background): “In recent literature, so-called Neural ODEs have attracted attention in both the ML and the physical sciences communities [11, 66, 16, 87, 85, 20].”
JFB: Jacobian-Free Backpropagation for Implicit Networks (2021, background): “A drawback is that the adjoint equation must be solved to high-accuracy; otherwise, a descent direction is not necessarily guaranteed (Gholami, Keutzer, and Biros 2019; Onken and Ruthotto 2020; Onken et al. 2021).”
Stochastic Normalizing Flows (2020, background): “Adaptive solvers prove particularly effective, although, as pointed out in Gholami et al. (2019), the backward solve (4) can often run into stability issues, suggesting a Rosenbrock or other implicit approach (Hairer & Wanner, 1996).”
TO-FLOW: Efficient Continuous Normalizing Flows with Temporal Optimization adjoint with Moving Speed (2022, background): “Some models introduce a check-point mechanism to store some of the state nodes during forward propogation to solve the backward integral accurately [38, 39, 55].”
Universal Differential Equations for Scientific Machine Learning (2020, methodology): “Previous research has shown that the discrete adjoint approach is more stable than continuous adjoints in some cases [41, 37, 42, 43, 44, 45] while continuous adjoints have been demonstrated to be more stable in others [46, 43] and can reduce spurious oscillations [47, 48, 49].”
OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport (2020, methodology): “In particular, DTO is efficient for solving neural ODEs (Li et al. 2017; Gholaminejad, Keutzer, and Biros 2019; Onken and Ruthotto 2020).”
Neural Controlled Differential Equations for Irregular Time Series (2020, methodology): “We do remark that the adjoint method should be used with care, as some systems are not stable to evaluate in both the forward and backward directions [29, 30].”
Liquid Time-constant Networks (2020, methodology): “This phenomenon happens because the adjoint method forgets the forward-time computational trajectories (See Figure 2), which was repeatedly denoted by the community [18].”