[Paper Review] How to train your neural ODE: the world of Jacobian and kinetic regularization

Chris Finlay, Jörn-Henrik Jacobsen+2 more|arXiv (Cornell University)|Feb 7, 2020

Model Reduction and Neural Networks29 references76 citations

TL;DR

RNODE introduces two regularizers (kinetic energy and Jacobian Frobenius norm) to neural ODEs, speeding training by encouraging simpler, more regular dynamics without sacrificing performance.

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.

Motivation & Objective

• Motivate the difficulty of training neural ODEs on large datasets due to adaptive solvers and deep discretizations.
• Propose theoretically grounded regularization terms to encourage simpler, well-behaved dynamics.
• Demonstrate that regularization speeds up training for neural ODE-based generative models without performance loss.

Proposed method

• Introduce two regularization terms derived from optimal transport and flow regularity: kinetic energy of the flow and Frobenius norm of the Jacobian.
• Augment the FFJORD continuous normalizing flow with these penalties via additional state variables to track kinetic energy and Jacobian cost.
• Estimate the divergence and Jacobian terms with Monte Carlo trace estimators to enable scalable training.
• Use adjoint sensitivity for backpropagation to compute gradients of the augmented objective.
• Show that RNODE can use fixed-grid solvers during training while preserving performance, enabling faster training.

Experimental results

Research questions

• RQ1 Can kinetic-energy and Jacobian regularizations yield more trainable neural ODEs without compromising likelihood performance?
• RQ2 Do these regularizers lead to faster wall-clock training by producing smoother, more integrable dynamics?
• RQ3 How do RNODE regularizers relate to optimal transport formulations of the learned flow?
• RQ4 What is the empirical impact of RNODE on training stability and compute (e.g., number of function evaluations) across standard generative modeling benchmarks?

Key findings

• RNODE achieves comparable log-likelihood to FFJORD while reducing training time substantially (e.g., reported speedups and maintained performance).
• Regularizers curb Jacobian growth and promote straight, constant-speed trajectories, improving numerical integrability of the learned dynamics.
• The kinetic-energy term and Jacobian regularization independently reduce solver step counts and stabilize training, as shown in ablation studies.
• RNODE's regularizers leverage existing divergence estimates and Jacobian-vector computations with negligible additional cost in FFJORD, enabling practical deployment on large datasets.