How to train your neural ODE: the world of Jacobian and kinetic regularization

Abstract

Training neural ODEs on large datasets has not been tractable due to the necessity of allowing the adaptive numerical ODE solver to refine its step size to very small values. In practice this leads to dynamics equivalent to many hundreds or even thousands of layers. In this paper, we overcome this apparent difficulty by introducing a theoretically-grounded combination of both optimal transport and stability regularizations which encourage neural ODEs to prefer simpler dynamics out of all the dynamics that solve a problem well. Simpler dynamics lead to faster convergence and to fewer discretizations of the solver, considerably decreasing wall-clock time without loss in performance. Our approach allows us to train neural ODE-based generative models to the same performance as the unregularized dynamics, with significant reductions in training time. This brings neural ODEs closer to practical relevance in large-scale applications.

arXiv commentAccepted to ICML 2020

References in corpus (13)

• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• Pixel Recurrent Neural Networks
• Density estimation using Real NVP
• Neural Ordinary Differential Equations
• Sensitivity and Generalization in Neural Networks: an Empirical Study
• A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problems
• Invertible Residual Networks
• Scalable Gradients for Stochastic Differential Equations
• ODE${}^2$VAE: Deep generative second order ODEs with Bayesian neural networks
• Augmented Neural ODEs
• Equivariant Flows: sampling configurations for multi-body systems with symmetric energies
• Monge-Ampère Flow for Generative Modeling
• Neural ODEs with stochastic vector field mixtures

Cited by in corpus (36)

• Normalizing Flows: An Introduction and Review of Current Methods
• HiPPO: Recurrent Memory with Optimal Polynomial Projections
• Dissecting Neural ODEs
• Deep Generative Modelling: A Comparative Review of VAEs, GANs, Normalizing Flows, Energy-Based and Autoregressive Models
• Learning Differential Equations that are Easy to Solve
• Parameterized Neural Ordinary Differential Equations: Applications to Computational Physics Problems
• A Neural Network Approach for High-Dimensional Optimal Control Applied to Multi-Agent Path Finding
• Riemannian Continuous Normalizing Flows
• Accelerating Neural ODEs Using Model Order Reduction
• Discretize-Optimize vs. Optimize-Discretize for Time-Series Regression and Continuous Normalizing Flows
• Continuous-in-Depth Neural Networks
• Model-based Reinforcement Learning for Semi-Markov Decision Processes with Neural ODEs
• A Neural Network Approach Applied to Multi-Agent Optimal Control
• Hypersolvers: Toward Fast Continuous-Depth Models
• Understanding and Mitigating Exploding Inverses in Invertible Neural Networks
• Robust learning with implicit residual networks
• Neural Ordinary Differential Equations on Manifolds
• TorchDyn: A Neural Differential Equations Library
• Unification of Symmetries Inside Neural Networks: Transformer, Feedforward and Neural ODE
• E(n) Equivariant Normalizing Flows
• Neural Ordinary Differential Equations for Data-Driven Reduced Order Modeling of Environmental Hydrodynamics
• Predicting Change, Not States: An Alternate Framework for Neural PDE Surrogates
• ACE-NODE: Attentive Co-Evolving Neural Ordinary Differential Equations
• Neural Spatio-Temporal Point Processes
• Learning High Dimensional Wasserstein Geodesics
• A Shooting Formulation of Deep Learning
• Opening the Blackbox: Accelerating Neural Differential Equations by Regularizing Internal Solver Heuristics
• NanoFlow: Scalable Normalizing Flows with Sublinear Parameter Complexity
• Sparse Flows: Pruning Continuous-depth Models
• Piecewise-constant Neural ODEs
• On Tuning Neural ODE for Stability, Consistency and Faster Convergence
• Go with the Flow: Adaptive Control for Neural ODEs
• "Hey, that's not an ODE": Faster ODE Adjoints via Seminorms
• OCT-GAN: Neural ODE-based Conditional Tabular GANs
• Augmented KRnet for density estimation and approximation
• Trust the Critics: Generatorless and Multipurpose WGANs with Initial Convergence Guarantees