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

Introduction

Recent research has bridged dynamical systems, a workhorse of mathematical modeling, with neural networks, the defacto function approximator for high dimensional data. The great promise of this pairing is that the vast mathematical machinery stemming from dynamical systems can be leveraged for modelling high dimensional problems in a dimension-independent fashion.

Connections between neural networks and ordinary differential equations (ODEs) were almost immediately noted after residual networks (He et al., 2016) were first proposed. Indeed, it was observed that there is a striking similarity between ResNets and the numerical solution of ordinary differential equations (E, 2017; Haber & Ruthotto, 2017; Ruthotto & Haber, 2018; Chen et al., 2018, 2019). In these works, deep networks are interepreted as discretizations of an underlying dynamical system, where time indexes the “depth” of the network and the parameters of the discretized dynamics are learned. An alternate viewpoint was taken by neural ODEs (Chen et al., 2018), where the dynamics of the neural network are approximated by an adaptive ODE solver on the fly. This latter approach is quite compelling as it does not require specifying the number of layers of the network beforehand. Furthermore, it allows the learning of homeomorphisms without any structural constraints on the function computed by the residual block.

Neural ODEs have shown great promise in the physical sciences (Köhler et al., 2019), in modeling irregular time series (Rubanova et al., 2019), mean field games (Ruthotto et al., 2019), continuous-time modeling (Yildiz et al., 2019; Kanaa et al., 2019), and for generative modeling through normalizing flows with free-form Jacobians (Grathwohl et al., 2019). Recent work has even adapted neural ODEs to the stochastic setting (Li et al., 2020). Despite these successes, some hurdles still remain. In particular, although neural ODEs are memory efficient, they can take a prohibitively long time to train, which is arguably one of the main stumbling blocks towards their widespread adoption.

In this work we reduce the training time of neural ODEs by regularizing the learned dynamics, complementing other recent approaches to this end such as augmented neural ODEs (Dupont et al., 2019). Without further constraints on their dynamics, high dimensional neural ODEs may learn dynamics which minimize an objective function, but which generate irregular solution trajectories. See for example Figure 1(b), where an unregularized flow exhibits undesirable properties due to unnecessarily fluctuating dynamics. As a solution, we propose two theoretically motivated regularization terms arising from an optimal transport viewpoint of the learned map, which encourage well-behaved dynamics (see 1(a) left). We empirically demonstrate that proper regularization leads to significant speed-up in training time without loss in performance, thus bringing neural ODEs closer to deployment on large-scale datasets. Our methods are validated on the problem of generative modelling and density estimation, as an example of where neural ODEs have shown impressive results, but could easily be applied elsewhere.

In summary, our proposed regularized neural ODE (RNODE) achieves the same performance as the baseline, while reducing the wall-clock training time by many hours or even days.

Neural ODEs & Continuous normalizing flows

Neural ODEs simplify the design of deep neural networks by formulating the forward pass of a deep network as the solution of a ordinary differential equation. Initial work along these lines was motivated by the similarity of the evaluation of one layer of a ResNet and the Euler discretization of an ODE. Suppose the block in the $t$-th layer of a ResNet is given by the function $\mathbf{f}(\mathbf{x},t;\theta)$, where $\theta$ are the block’s parameters. Then the evaluation of this layer of the ResNet is simply $\mathbf{x}^{t+1}=\mathbf{x}^{t}+\mathbf{f}(\mathbf{x}^{t},t;\theta)$. Now, instead consider the following ODE

The Euler discretization of this ODE with step-size $\tau$ is $\mathbf{z}^{t+1}=\mathbf{z}^{t}+\tau\mathbf{f}(\mathbf{z}^{t},t;\theta)$, which is nearly identical to the forward evaluation of the ResNet’s layer (setting step-size $\tau=1$ gives equality). Armed with this insight, Chen et al. (2018) suggested a method for training neural networks based on (ODE) which abstain from a priori fixing step-size. Chen et al.’s method is a continuous-time generalization of residual networks, where the dynamics are generated by an adaptive ODE solver that chooses step-size on-the-fly.

Because of their adaptive nature, neural ODEs can be more flexible than ResNets in certain scenarios, such as when trading between model speed and accuracy. Moreover given a fixed network depth, the memory footprint of neural ODEs is orders of magnitude smaller than a standard ResNet during training. They therefore show great potential on a host of applications, including generative modeling and density estimation. An apparent drawback of neural ODEs is their long training time: although a learned function $\mathbf{f}(\cdot\,;\theta)$ may generate a map that solves a problem particularly well, the computational cost of numerically integrating (ODE) may be so prohibitive that it is not tractable in practice. In this paper we demonstrate this need not be so: with proper regularization, it is possible to learn $\mathbf{f}(\cdot\,;\theta)$ so that (ODE) is easily and quickly solved.

In density estimation and generative modeling, we wish to estimate an unknown data distribution $p(\mathbf{x})$ from which we have drawn $N$ samples. Maximum likelihood seeks to approximate $p(\mathbf{x})$ with a parameterized distribution $p_{\theta}(\mathbf{x})$ by minimizing the Kullback-Leibler divergence between the two, or equivalently minimizing

Evaluating the log determinant of the Jacobian is difficult. Grathwohl et al. (2019) exploit the following identity from fluid mechanics (Villani, 2003, p 114)

where $\operatorname{div}\left(\cdot\right)$ is the divergence operatorIn the normalizing flow literature divergence is typically written explicitly as the trace of the Jacobian, however we use $\operatorname{div}\left(\cdot\right)$ which is more common elsewhere., $\operatorname{div}\left(\mathbf{f}\right)(\mathbf{x})=\sum_{i}\partial_{x_{i}}f_{i}(\mathbf{x})$. By the fundamental theorem of calculus, we may then rewrite (2) in integral form

In (Grathwohl et al., 2019), the divergence is estimated using an unbiased Monte-Carlo trace estimate (Hutchinson, 1990; Avron & Toledo, 2011),

By using the substitution (4), the task of maximizing log-likelihood shifts from choosing $p_{\theta}$ to minimize (1), to learning the flow generated by a vector field $\mathbf{f}$. This results in a normalizing flow with a free-form Jacobian and reversible dynamics, and was named FFJORD by Grathwohl et al..

2 The need for regularity

The vector field learned through FFJORD that maximizes the log-likelihood is not unique, and raises troubling problems related to the regularity of the flow. For a simple example, refer to Figure 1, where we plot two normalizing flows, both mapping a toy one-dimensional distribution to the unit Gaussian, and where both maximize the log-likelihood of exactly the same sample of particles. Figure 1(a) presents a “regular” flow, where particles travel in straight lines that travel with constant speed. In contrast, Figure 1(b) shows a flow that still maximizes the log-likelihood, but that has undesirable properties, such as rapidly varying local trajectories and non-constant speed.

From this simple motivating example, the need for regularity of the vector field is apparent. Without placing demands on the vector field $\mathbf{f}$, it is entirely possible that the learned dynamics will be poorly conditioned. This is not just a theoretical exercise: because the dynamics must be solved with a numerical integrator, poorly conditioned dynamics will lead to difficulties during numerical integration of (ODE). Indeed, later we present results demonstrating a clear correlation between the number of time steps an adaptive solver takes to solve (ODE), and the regularity of $\mathbf{f}$.

How can the regularity of the vector field be measured? One motivating approach is to measure the force experienced by a particle $\mathbf{z}(t)$ under the dynamics generated by the vector field $\mathbf{f}$, which is given by the total derivative of $\mathbf{f}$ with respect to time

Well conditioned flows will place constant, or nearly constant, force on particles as they travel. Thus, in this work we propose regularizing the dynamics with two penalty terms, one term regularizing $\mathbf{f}$ and the other $\operatorname{\nabla}\mathbf{f}$. The first penalty, presented in Section 3, is a measure of the distance travelled under the flow $\mathbf{f}$, and can alternately be interpreted as the kinetic energy of the flow. This penalty term is based off of numerical methods in optimal transport, and encourages particles to travel in straight lines with constant speed. The second penalty term, discussed in Section 4, performs regularization on the Jacobian of the vector field. Taken together the two terms ensure that the force experienced by a particle under the flow is constant or nearly so.

These two regularizers will promote dynamics that follow numerically easy-to-integrate paths, thus greatly speeding up training time.

Optimal transport maps & Benamou-Brenier

The objective function (18a) is a measure of the kinetic energy of the flow. The constraint (18b) ensures probability mass is conserved. The latter two constraints guarantee the learned distribution agrees with the source $p$ and target $q$. Note that the kinetic energy (18a) is an upper bound on the transport cost, with equality only at optimality.

The optimal flow $\mathbf{f}$ minimizing (18) has several particularly appealing properties. First, particles induced by the optimal flow $\mathbf{f}$ travel in straight lines. Second, particles travel with constant speed. Moreover, under suitable conditions on the source and target distributions, the optimal solution map is unique (Villani, 2008). Therefore the solution map $\mathbf{z}(\mathbf{x},t)$ is entirely characterized by the initial and final positions: $\mathbf{z}(\mathbf{x},t)=(1-\frac{t}{T})\mathbf{z}(\mathbf{x},0)+\frac{t}{T}\mathbf{z}(\mathbf{x},T)$. Consequently, given an optimal $\mathbf{f}$ it is extraordinarily easy to solve (ODE) numerically with minimal computational effort.

Now suppose we wish to minimize (18a), with $q(\mathbf{z})$ a unit normal distribution, and $p(\mathbf{x})$ a data distribution, unknown to us, but from which we have drawn $N$ samples, and which we model as a discrete distribution of Dirac masses. Enforcing the initial condition is trivial because we have sampled from $p$ directly. The continuity equation (18b) need not be enforced because we are tracking a finite number of sampled particles. However the final time condition $\rho_{T}=q$ cannot be implemented directly, since we do not have direct control on the form $\rho_{T}(\mathbf{z})$ takes. Instead, introduce a Kullback-Leibler term to (18a) penalizing discrepancy between $\rho_{T}$ and $q$. This penalty term has an elegant simplification when $p(x)$ is modeled as a distribution of a finite number of masses, as is done in generative modeling. Setting $\rho_{0}=p_{\theta}$ a brief derivation yields

For further details on this derivation consult the supplementary materials.

The connection between the Benamou-Brenier formulation of the optimal transport problem on a discrete set of points and continuous normalizing flows is apparent: the optimal transport problem (11) is a regularized form of the continuous normalizing flow optimization problem (1). We therefore expect that adding a kinetic energy regularization term to FFJORD will encourage solution trajectories to prefer straight lines with constant speed.

Unbiased Frobenius norm regularization of the Jacobian

Refering to equation (7), one can see that even if $\mathbf{f}$ is regularized to be small, via a kinetic energy penalty term, if the Jacobian is large then the force experienced by a particle may also still be large. As a result, the error of the numerical integrator can be large, which may lead an adaptive solver to make many function evaluations. This relationship is apparent in Figure 3, where we empirically demonstrate the correlation between the number of function evaluations of $\mathbf{f}$ taken by the adaptive solver, and the size of the Jacobian norm of $\mathbf{f}$. The correlation is remarkably strong: dynamics governed by a poorly conditioned Jacobian matrix require the adaptive solver to take many small time steps.

Moreover, in particle-based methods, the kinetic energy term forces dynamics to travel in straight lines only on data seen during training, and so the regularity of the map is only guaranteed on trajectories taken by training data. The issue here is one of generalization: the map may be irregular on off-distribution or perturbed images, and cannot be remedied by the kinetic energy term during training alone. In the context of generalization, Jacobian regularization is analagous to gradient regularization, which has been shown to improve generalization (Drucker & LeCun, 1992; Novak et al., 2018).

and is the Euclidean norm of the singular values of a matrix. In trace form, the Frobenius norm lends itself to estimation using a Monte-Carlo trace estimator (Hutchinson, 1990; Avron & Toledo, 2011). For real matrix $B$, an unbiased estimate of the trace is given by

where $\mathbf{\epsilon}$ is drawn from a unit normal distribution. Thus the squared Frobenius norm can be easily estimated by setting $B=AA^{\mathsf{T}}$.

Conveniently, in the FFJORD framework the quantity $\mathbf{\epsilon}^{\mathsf{T}}\operatorname{\nabla}\mathbf{f}(\mathbf{z})$ must be computed during the estimate of the probability distribution under the flow, in the Monte-Carlo estimate of the divergence term (5). Thus Jacobian Frobenius norm regularization is available with essentially no extra computational cost.

Algorithm description

All together, we propose modifying the objective function of the FFJORD continuous normalizing flow (Grathwohl et al., 2019) with the two regularization penalties of Sections 3 & 4. The proposed method is called RNODE, short for regularized neural ODE. Pseudo-code of the method is presented in Algorithm 1. The optimization problem to be solved is

where $\mathbf{z}(\mathbf{x},t)$ is determined by numerically solving (ODE). Note that we take the mean over number of samples and input dimension. This is to ensure that the choice of regularization strength $\lambda_{K}$ and $\lambda_{J}$ is independent of dimension size and sample size.

To compute the three integrals and the log-probability under $q$ of $\mathbf{z}(\mathbf{x},T)$ at final time $T$, we augment the dynamics of the ODE with three extra terms, so that the entire system solved by the numerical integrator is

Here $E$, $l$, and $n$ are respectively the kinetic energy, the log determinant of the Jacobian, and the integral of the Frobenius norm of the Jacobian.

Both the divergence term and the Jacobian Frobenius norm are approximated with Monte-Carlo trace estimates. In our implementation, the Jacobian Frobenius estamate reuses the computatian $\mathbf{\epsilon}^{\mathsf{T}}\operatorname{\nabla}f$ from the divergence estimate for efficiency. We remark that the kinetic energy term only requires the computation of a dot product. Thus just as in FFJORD, our implementation scales linearly with the number of time steps taken by the ODE solver.

Gradients of the objective function with respect to the network parameters are computed using the adjoint sensitivity method (Pontryagin et al., 1962; Chen et al., 2018).

Experimental design

On MNIST and CIFAR10 we train with a batch size of 200 and train for 100 epochs on a single GPUGeForce RTX 2080 Ti, using the Adam optimizer (Kingma & Ba, 2015) with a learning rate of $110-3$. On the two larger datasets, we train with four GPUs, using a per-GPU batch size of respectively 3 and 50 for CelebA-HQ and ImageNet. Data is preprocessed by perturbing with uniform noise followed by the logit transform.

The reference implementation of FFJORD solves the dynamics using a Runge-Kutta 4(5) adaptive solver (Dormand & Prince, 1980) with error tolerances $110-5$ and initial step size $110-2$. We have found that using less accurate solvers on the reference implementation of FFJORD results in numerically unstable training dynamics. In contrast, a simple fixed-grid four stage Runge-Kutta solver suffices for RNODE during training on MNIST and CIFAR10, using a step size of $0.25$. The step size was determined based on a simple heuristic of starting with $0.5$ and decreasing the step size by a factor of two until the discrete dynamics were stable and achieved good performance. The Runge-Kutta 4(5) adaptive solver was used on the two larger datasets. We have also observed that RNODE improves the training time of the adaptive solvers as well, requiring many fewer function evaluations; however in Python we have found that the fixed grid solver is typically quicker at a specified number of function evaluations. At test time RNODE uses the same adaptive solver as FFJORD.

We always initialize RNODE so that $\mathbf{f}(z,t)=0$; thus training begins with an initial identity map. This is done by zero-ing the parameters of the last layer in each piece (block), following Goyal et al. (2017). The identity map is an appropriate choice because it has zero transport cost and zero Frobenius norm. Moreover the identity map is trivially solveable for any numerical solver, thus training begins without any effort required on the solver’s behalf.

On all datasets we set both the kinetic energy regularization coefficient $\lambda_{K}$ and the Jacobian norm coefficient $\lambda_{J}$ to 0.01.

Results

A comparison of RNODE against FFJORD and other flow-based generative models is presented in Table 1. We report both our running of “vanilla” FFJORD and the results as originally reported in (Grathwohl et al., 2019). We highlight that RNODE runs roughly 2.8x faster than FFJORD on both datasets, while achieving or surpassing the performance of FFJORD. This can further be seen in Figure 2 where we plot bits per dimension ( $-\frac{1}{d}\log_{2}p(x)$, a normalized measure of log-likelihood) on the validation set as a function of training epoch, for both datasets. Visual inspection of the sample quality reveals no qualitative difference between regularized and unregularized approaches; refer to Figure 6. Generated images for downsampled ImageNet and CelebA-HQ are deferred to the supplementary materials; we provide smaller generated images for networks trained on CelebA-HQ 64x64 in Figure 4.

Surprisingly, our run of “vanilla” FFJORD achieved slightly better performance than the results reported in (Grathwohl et al., 2019). We suspect the discrepancy in performance and run times between our implementation of FFJORD and that of the original paper is due to batch size: Grathwohl et al. use a batch size of 900 and train on six GPUs, whereas on MNIST and CIFAR10 we use a batch size of 200 and train on a single GPU.

We were not able to train vanilla FFJORD on ImageNet64, due to numerical underflow in the adaptive solver’s time step. This issue cannot be remedied by increasing the solver’s error tolerance, for this would bias the log-likelihood estimates on validation.

In Figure 5, we compare the effect of each regularizer by itself on the training dynamics with the fixed grid ODE solver on the MNIST dataset. Without any regularization at all, training dynamics are numerically unstable and fail after just under 50 epochs. This is precisely when the Jacobian norm grows large; refer to Figure 5(a). Figure 5(a) demonstrates that each regularizer by itself is able to control the Jacobian norm. The Jacobian regularizer is better suited to this task, although it is interesting that the kinetic energy regularizer also improves the Jacobian norm. Unsurprisingly Figure 5(b) demonstrates the addition of the kinetic energy regularizer encourages flows to travel a minimal distance. In addition, we see that the Jacobian norm alone also has a beneficial effect on the distance particles travel. Overall, the results support our theoretical reasoning empirically.

Previous generative flows inspired by optimal transport

Zhang et al. (2018) define a neural ODE flow where the dynamics are given as the gradient of a scalar potential function. This interpretation has deep connections to optimal transport: the optimal transport map is the gradient of a convex potential function. Yang & Karniadakis (2019) continue along these lines, and define an optimal transport again as a scalar potential gradient. Yang & Karniadakis (2019) enforce that the learned map is in fact an optimal transport map by penalizing their objective function with a term measuring violations of the continuity equation. Ruthotto et al. (2019) place generative flows within a broader context of mean field games, and as an example consider a neural ODE gradient potential flow solving the optimal transport problem in up to 100 dimensions. We also note the recent work of Twomey et al. (2019), who proposed regularizing neural ODEs with an Euler-step discretization of the kinetic energy term to enforce ‘straightness’, although connections to optimal transport were not discussed.

When a flow is the gradient of a scalar potential, the change of variables formula (4) simplifies so that the divergence term is replaced by the Laplacian of the scalar potential. Although mathematically parsimonious and theoretically well-motivated, we chose not to implement our flow as the gradient of a scalar potential function due to computational constraints: such an implementation would require ‘triple backprop’ (twice to compute or approximate the Laplacian, and once more for the parameter gradient). Ruthotto et al. (2019) circumvented this problem by utilizing special structural properties of residual networks to efficiently compute the Laplacian.

Discussion

In practice, RNODE is simple to implement, and only requires augmenting the dynamics (ODE) with two extra scalar equations (one for the kinetic energy term, and another for the Jacobian penalty). In the setting of FFJORD, because we may recycle intermediary terms used in the divergence estimate, the computational cost of evaluating these two extra equations is minimal. RNODE introduces two extra hyperparameters related to the strength of the regularizers; we have found these required almost no tuning.

Although the problem of classification was not considered in this work, we believe RNODE may offer similar improvements both in training time and the regularity of the classifier learned. In the classification setting we expect the computional overhead of calculating the two extra terms should be marginal relative to gains made in training time.

Conclusion

We have presented RNODE, a regularized method for neural ODEs. This regularization approach is theoretically well-motivated, and encourages neural ODEs to learn well-behaved dynamics. As a consequence, numerical integration of the learned dynamics is straight forward and relatively easy, which means fewer discretizations are needed to solve the dynamics. In many circumstances, this allows for the replacement of adaptive solvers with fixed grid solvers, which can be more efficient during training. This leads to a substantial speed up in training time, while still maintaining the same empirical performance, opening the use of neural ODEs to large-scale applications.

Acknowledgements

C. F. and A. O. were supported by a grant from the Innovative Ideas Program of the Healthy Brains and Healthy Lives initiative (HBHL) through McGill University.

L. N. was supported by AFOSR MURI FA9550-18-1-0502, AFOSR Grant No. FA9550-18-1-0167, and ONR Grant No. N00014-18-1-2527.

A. O. was supported by the Air Force Office of Scientific Research under award number FA9550-18-1-0167

Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute (www.vectorinstitute.ai/#partners).

References

Appendix A Details of Section 3.1: Benamou-Brenier formulation in Lagrangian coordinates

The Benamou-Brenier formulation of the optimal transportation (OT) problem in Eulerian coordinates is

The connection between continuous normalizing flows (CNF) and OT becomes transparent once we rewrite (18) in Lagrangian coordinates. Indeed, for regular enough velocity fields $\mathbf{f}$ one has that the solution of the continuity equation (18b), (18c) is given by $\rho_{t}=\mathbf{z}(\cdot,t)\sharp p$ where $\mathbf{z}$ is the flow

The relation $\rho_{t}=\mathbf{z}(\cdot,t)\sharp p$ means that for arbitrary test function $\phi$ we have that

Note that $\rho_{t}$ is eliminated in this formulation. The terminal condition (18d) is trivial to implement in Eulerian coordinates (grid-based methods) but not so simple in Lagrangian ones (19d) (grid-free methods). To enforce (19d) we introduce a penalty term in the objective function that measures the deviation of $\mathbf{z}(\cdot,T)\sharp p$ from $q$. Thus, the penalized objective function is

where $\lambda>0$ is the penalization strength. Next, we observe that this objective function can be written as an expectation with respect to $\mathbf{x}\sim p$. Indeed, the Kullback-Leibler divergence is invariant under coordinate transformations, and therefore

Finally, if we assume that $\{\mathbf{x}_{i}\}_{i=1}^{N}$ are iid sampled from $p$, we obtain the empirical objective function

Appendix B Additional results

Here we present additional generated samples on the two larger datasets considered, CelebA-HQ and ImageNet64. In addition bits/dim on clean images are reported in Table 2.