How to train your neural ODE

By Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, Adam M. Oberman. 2020

366 citations recorded by the source corpus.

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 in just over a day on one GPU, whereas unregularized dynamics can take up to 4-6 days of training time on multiple GPUs. This brings neural ODEs significantly closer to practical relevance in large-scale applications.

Subfields: Artificial Intelligence, Biophysics, 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 Artificial Intelligence

In Biophysics

In Statistical and Nonlinear Physics

Citations in the corpus

Flow Matching for Generative Modeling (2022, background): “, 2019), adding regularization terms (Yang & Karniadakis, 2019; Finlay et al., 2020; Onken et al., 2021; Tong et al., 2020; Kelly et al., 2020), or stochastically sampling the integration interval (Du et al.”
HiPPO: Recurrent Memory with Optimal Polynomial Projections (2020, background): “Though neural ODEs are expressive [27, 78], due to their complex parameterization, they often suffer from slow training [29, 53, 58] because of their need for more complicated ODE solvers.”
GRAND: Graph Neural Diffusion (2021, methodology): “Many follow-up works explored augmentation (Dupont et al., 2019) and regularization (Finlay et al., 2020) and provided extensions into new domains such as stochastic (Li et al., 2020a; Tzen & Raginsky, 2019) differential equations.”
An introduction to deep generative modeling (2021, background): “Several approaches that add transport costs to the CNF problem have been proposed recently [48, 47, 14, 30, 34].”
Model-based Reinforcement Learning for Semi-Markov Decision Processes with Neural ODEs (2020, methodology): “We believe that the efficiency of our methods will not only be significantly improved with a theoretically faster numerical ODE solver (e.g., [Finlay et al., 2020, Kelly et al., 2020]), but also with a better implementation of ODE solvers2 (e.g., a faster C++/Cython implementation, using single…”
Implicit Normalizing Flows (2021, background): “, 2020) and exceedingly many ODE solver steps (Finlay et al., 2020), making their large-scale application still an open problem.”
Improving and generalizing flow-based generative models with minibatch optimal transport (2023, positive): “…distributions with neural networks; however, these require constrained architectures (Leygonie et al., 2019; Makkuva et al., 2020; Bunne et al., 2022) or use a regularized CNF, which is challenging to optimize (Tong et al., 2020; Finlay et al., 2020; Onken et al., 2021; Huguet et al., 2022a).”
TO-FLOW: Efficient Continuous Normalizing Flows with Temporal Optimization adjoint with Moving Speed (2022, methodology): “Finlay et al [56] and Onken et al [41] constrain the movement speed f ( z ( t ) , t ; θ ) of particles to reduce the loss of transmission from the perspective of OT theory.”
Efficient-VDVAE: Less is more (2022, methodology): “In table 4, we test the effect of this removal on FFHQ 256× 256 and CelebAHQ 256× 256 which have been consistently quantized to 5-bits in prior work [16, 23, 47, 48, 38, 37, 32, 39, 34].”
JFB: Jacobian-Free Backpropagation for Implicit Networks (2021, background): “Neural ODEs leverage known connections between deep residual models and discretizations of differential equations (Haber and Ruthotto 2017; Weinan 2017; Ruthotto and Haber 2019; Chang et al. 2018; Finlay et al. 2020; Lu et al. 2018), and replace these discretizations by black-box ODE solvers in forw”
Learning Differential Equations that are Easy to Solve (2020, methodology): “This accounts for discrepancy in bits/dim and NFE reported in Finlay et al. (2020).”
Neural Controlled Differential Equations for Irregular Time Series (2020, background): “[55, 58] consider how a Neural ODE model may be regularised (see also the stochastic regularisation discussed above).”
OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport (2020, methodology): “2019) and RNODE (Finlay et al. 2020) (current state-of-the-art) in speed and performance.”
Building Normalizing Flows with Stochastic Interpolants (2022, neutral): “…to exemplary discrete and continuous ﬂows: MADE (Germain et al., 2015), Real NVP (Dinh et al., 2017), Convex Potential Flows (CPF) (Huang et al., 2021), Neural Spline Flows (NSP) Durkan et al. (2019), Free-form continuous ﬂows (FFJORD) (Grathwohl et al., 2019), and OT-Flow (Finlay et al., 2020).”
Stochastic Interpolants: A Unifying Framework for Flows and Diffusions (2023, methodology): “Optimal transport perspectives in flow-based methods have primarily been employed as a regularization penalty [51, 34, 14, 46] or via imposing structure on the parameterization itself [19, 50].”
Deep Generative Modelling: A Comparative Review of VAEs, GANs, Normalizing Flows, Energy-Based and Autoregressive Models (2021, background): “RNODE [53] applies transport costs to FFJORD as well as regularis-”