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

TL;DR

The paper tackles the slow training of neural ODE-based models on large datasets caused by adaptive ODE solvers. It introduces RNODE, which adds two regularizers—one kinetic-energy term inspired by optimal transport and a Frobenius-norm penalty on the Jacobian—to bias the learned dynamics toward simpler, easier-to-solve trajectories. Empirically, RNODE achieves substantial wall-clock speedups (about 2.8×) over FFJORD on MNIST and CIFAR-10 while preserving or improving log-likelihood performance, and reduces solver evaluations by stabilizing dynamics. The method is simple to implement (augmenting the dynamics with three extra state channels) and scales to large generative modeling tasks, advancing the practical use of neural ODEs.

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.

Figures (8)

• Figure 1: Optimal transport map and a generic normalizing flow.
• Figure 2: Log-likelihood (measured in bits/dim) on the validation set as a function of wall-clock time. Rolling average of three hours, with 90% confidence intervals.
• Figure 3: Number of function evaluations vs Jacobian Frobenius norm of flows on CIFAR10 during training with vanilla FFJORD, using an adaptive ODE solver.
• Figure 4: Quality of generated samples samples on 5bit CelebA-HQ64 with RNODE. Here temperature annealing glow with $T=0.7$ was used to generate visually appealing images. For full sized CelebA-HQ256 samples, consult the supplementary materials.
• Figure 5: Ablation study of the effect of the two regularizers, comparing two measures of flow regularity during training with a fixed step-size ODE solver. Figure \ref{['fig:jf']}: mean Jacobian Frobenius norm as a function of training epoch. Figure \ref{['fig:ke']}: mean kinetic energy of the flow as a function of training epoch. Figure \ref{['fig:nfes']}: number of function evaluations.

Theorems & Definitions (1)

• Remark 2.1: Divergence trace estimate