Second-Order Neural ODE Optimizer

Abstract

We propose a novel second-order optimization framework for training the emerging deep continuous-time models, specifically the Neural Ordinary Differential Equations (Neural ODEs). Since their training already involves expensive gradient computation by solving a backward ODE, deriving efficient second-order methods becomes highly nontrivial. Nevertheless, inspired by the recent Optimal Control (OC) interpretation of training deep networks, we show that a specific continuous-time OC methodology, called Differential Programming, can be adopted to derive backward ODEs for higher-order derivatives at the same O(1) memory cost. We further explore a low-rank representation of the second-order derivatives and show that it leads to efficient preconditioned updates with the aid of Kronecker-based factorization. The resulting method -- named SNOpt -- converges much faster than first-order baselines in wall-clock time, and the improvement remains consistent across various applications, e.g. image classification, generative flow, and time-series prediction. Our framework also enables direct architecture optimization, such as the integration time of Neural ODEs, with second-order feedback policies, strengthening the OC perspective as a principled tool of analyzing optimization in deep learning. Our code is available at https://github.com/ghliu/snopt.

arXiv commentAccepted to Advances in Neural Information Processing Systems (NeurIPS) 2021 as Spotlight

References in corpus (23)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation
• New insights and perspectives on the natural gradient method
• How Does Batch Normalization Help Optimization?
• Latent ODEs for Irregularly-Sampled Time Series
• FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models
• The UEA multivariate time series classification archive, 2018
• Dissecting Neural ODEs
• ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs
• Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis
• Natural Neural Networks
• Learning Differential Equations that are Easy to Solve
• Deep Learning Theory Review: An Optimal Control and Dynamical Systems Perspective
• Riemannian Continuous Normalizing Flows
• Neural Manifold Ordinary Differential Equations
• STEER: Simple Temporal Regularization For Neural ODEs
• Dissecting Hessian: Understanding Common Structure of Hessian in Neural Networks
• MALI: A memory efficient and reverse accurate integrator for Neural ODEs
• Mean-field Langevin System, Optimal Control and Deep Neural Networks
• InfoCNF: An Efficient Conditional Continuous Normalizing Flow with Adaptive Solvers
• DDPNOpt: Differential Dynamic Programming Neural Optimizer
• A Trace-restricted Kronecker-Factored Approximation to Natural Gradient
• Model Based Reinforcement Learning with Final Time Horizon Optimization