Adaptive Checkpoint Adjoint Method for Gradient Estimation in Neural ODE

Abstract

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}.

References in corpus (8)

• On the difficulty of training Recurrent Neural Networks
• Trainable Nonlinear Reaction Diffusion: A Flexible Framework for Fast and Effective Image Restoration
• Training Deep Nets with Sublinear Memory Cost
• Latent ODEs for Irregularly-Sampled Time Series
• Beyond Finite Layer Neural Networks: Bridging Deep Architectures and Numerical Differential Equations
• Newton vs the machine: solving the chaotic three-body problem using deep neural networks
• ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs
• Learning Dynamical Systems from Partial Observations

Cited by in corpus (4)

• Parameterized Neural Ordinary Differential Equations: Applications to Computational Physics Problems
• Galaxy Morphology Classification using Neural Ordinary Differential Equations
• How to pose material design problems for flow through porous media applications?: Sensitivity of dissipation rate to medium's permeability holds the key
• Diffeomorphic Transformations for Time Series Analysis: An Efficient Approach to Nonlinear Warping