Accelerating Continuous Normalizing Flow with Trajectory Polynomial Regularization

In this paper, we propose an approach to effectively accelerating the computation of continuous normalizing flow (CNF), which has been proven to be a powerful tool for the tasks such as variational inference and density estimation. The training time cost of CNF can be extremely high because the required number of function evaluations (NFE) for solving corresponding ordinary differential equations (ODE) is very large. We think that the high NFE results from large truncation errors of solving ODEs. To address the problem, we propose to add a regularization. The regularization penalizes the difference between the trajectory of the ODE and its fitted polynomial regression. The trajectory of ODE will approximate a polynomial function, and thus the truncation error will be smaller. Furthermore, we provide two proofs and claim that the additional regularization does not harm training quality. Experimental results show that our proposed method can result in 42.3% to 71.3% reduction of NFE on the task of density estimation, and 19.3% to 32.1% reduction of NFE on variational auto-encoder, while the testing losses are not affected at all.


Introduction
Normalizing flows (Rezende and Mohamed 2015) are a kind of invertible neural networks that have an efficient calculation of Jacobian determinant. They can be used as generative models, density estimation or posterior distribution of variational auto-encoders (VAE). However, their two requirements, invertibility and easy Jacobian determinant calculation, put a great restriction on the design of corresponding neural network architecture.
Recently, continuous normalizing flow (CNF) (Chen et al. 2018;Grathwohl et al. 2019) is proposed that can avoid such design restriction. CNF describes the transformation of hidden states with ordinary differential equations (ODE), instead of layer-wise function mappings. In this way, the invertibility becomes trivial and the determinant of Jacobian becomes simply a trace. With the freedom of network architecture design, CNF has been shown to outperform normalizing flow with discrete layers in terms of lower testing loss.
However, running a CNF model, which is equal to solving a series of ODEs, may take a lot of time. A single training Copyright © 2021, Association for the Advancement of Artificial Intelligence (www.aaai.org). All rights reserved. iteration can need up to hundreds of function evaluations, which is equivalent to running a neural network with hundreds of layers. Moreover, the number of function evaluations required per training iteration can gradually grow up throughout the training process. The reason is that CNF uses ODE solvers with adaptive steps. The step sizes of solving ODEs are determined based on the truncation errors. Such error usually grows as the training goes and thus the number of steps becomes larger.
To make training of CNF faster, we propose Trajectory Polynomial Regularization (TPR). TPR is an additional loss function which regularizes the trajectories of the ODE solutions. It penalizes the difference between the trajectory of solution and its fitted polynomial regression. Therefore, TPR enforces the solutions to approximate polynomial curves, which can reduce the truncation error of solving the ODEs of CNF. Despite adding the regularization, we argue that the method does not harm the testing loss much. We prove that the optimal solutions of the modified loss are still optimal solutions for the original task. The detail of the proofs and argument are in Section 3. In experiments, we find that for both density estimation and variational inference tasks, our model can save as much as 71% of number of function evaluations (NFE), which is the number of evaluating the ODEs. Furthermore, our method barely affect the testing errors. Our code is published in our Github repository 1 .
The remainder of the paper is organized as follows. We first introduce the background and related work of CNF in Section 2. In Section 3, we describe the intuition and mathematics of TPR and provide two proofs to argue the power of its quality. In Section 4, we conduct experiments on two tasks, density estimation and VAE. We visualize the effect of our model with simple 2D toy data. And then we evaluate the efficiency and quality of our model on real data.

Background
Normalizing flows (Rezende and Mohamed 2015) are models that can be used to represent a wide range of distributions. They consist of an invertible function. Let f : R d → R d be an invertible function, i.e., z = f (x) and x = f −1 (z).
Due to the invertibility, the log probabilities of the input x and the output z = f (x) have the following relation. (1) Once p(x) is set to be a known distribution, the distribution p(z) is represented by p(x) and the function f . The function f can be parametrized by deep neural networks. However, since f should be invertible and the Jacobian det ∂f ∂x in equation (1) should be easily calculated, the possible architecture of f is highly restricted. As a result, the capacity of representing p(x) is also limited.
Different from normalizing flows, continuous normalizing flow (CNF) (Chen et al. 2018) makes the transformation from x to z a continuous evolution. That is, instead of a direct function mapping from x to z, i.e. z = f (x), the transformation can be represented by a continuous function y(t), with x = y(t 0 ) and z = y(t 1 ). The evolution of y(t) in CNF is defined by an ordinary differential equation (ODE), With this continuous transformation, the invertibility of x and z is inherently satisfied. Therefore, there isn't any other restriction on the form of function v(y, t).
The dynamics of the log likelihood log p(y(t)) is derived in the instantaneous change of variables theorem by Chen et al. (2018), which is The computation complexity for obtaining the exact value of Tr ∂v ∂y is O D 2 , where D is the dimension of data x. It can be reduced to O(D) (Grathwohl et al. 2019) with the Hutchinson's trace estimator (Hutchinson 1989).
where E[ ] = 0 and Cov( ) = I. With this method, computing the estimation of Tr ∂v ∂y costs exactly the same as computing v.
Given the initial values of y(t 0 ) and log p(t 0 ), CNF outputs the final values y(t 1 ) and log p(t 1 ), which can be written as log p(y(t 1 )) = log p(y(t 0 )) − These equations are solved by numerical ODE solvers. Usually, adaptive stepsize solvers, such as Dopri5 (Dormand and Prince 1980), are used because they guarantee a desired precision of output. Adaptive solvers dynamically determine the step sizes of solving ODE by estimating the magnitude of truncation errors. They shrink the step sizes when the truncation errors exceed a tolerable value and increase the step sizes when the errors are too small. Unlike the restriction of f 's architecture in normalizing flows, the v(y, t) in CNF can be represented by arbitrary neural networks. As a result, CNF can have greater capacity to represent the distributions p(z) than other normalizing flow models.

Related work
The design of normalizing flows is restricted by the requirements of invertibility and efficient Jacobian determinant computation. Rezende and Mohamed (2015) proposed planar flow and radial flow for variational inference. The layers in the proposed model have to be in a simple form. Real NVP (Dinh, Sohl-Dickstein, and Bengio 2017) is a structure of network such that its Jacobian matrix are triangular, so the determinant are simply the product of diagonal values. Kingma and Dhariwal (2018) proposed Glow, which has 1×1 convolution layers that pass partitioned dimensions into invertible affine transformations. Auto-regressive flow models (Kingma et al. 2016;Papamakarios, Pavlakou, and Murray 2017;Oliva et al. 2018;Huang et al. 2018) also have triangular Jacobian matrix in each layer. Although they have high expressiveness, they require D sequential computations in each layer, which can be costly when D is large. Huang, Dinh, and Courville (2020) proposed Augmented Normalizing Flow (ANF), which augments the data by adding more dimensions in transformation. Although reaching high performance, it requires large number of samples to compute accurate importance sampling.
Continuous normalizing flow (CNF) (Chen et al. 2018) is a continuous version of normalizing flows, which replaces the layer-wise transformations of normalizing flows with ODEs. The model does not restrict the structure of its transformation function, but computing Jacobian determinant costs computing gradient D times. Grathwohl et al. (2019) improved the efficiency of computing Jacobian determinant by using Hutchinson's trace estimator. Chen and Duvenaud (2019) proposed a new architecture of neural network which can efficiently compute trace of Jacobian.
As we mentioned in Section 2.1, training CNF takes a lot of time because solving ODEs can be very slow. To enhance the training stability and reduce the training time, some researchers have proposed to adopt regularization methods. Grathwohl et al. (2019) proposed to use weight decay and spectral normalization (Miyato et al. 2018). Salman et al. (2018) proposed geodesic regularization and inverse consistency regularization. Finlay et al. (2020) and Onken et al. (2020) proposed to regularize the dynamics of CNF with the notion of optimal transport. Kelly et al. (2020) proposed to regularize the higher-order derivatives of the velocity field and applied Taylor-mode automatic differentiation to compute the derivatives efficiently. Our work is different from these works in two ways. First, we provide proofs for the existence of infinite optimal solutions with our regularization. Second, our proposed TPR only requires a few samples point in time. On the contrary, those methods need to solve more than one additional ODEs, which increases computational consumption.
The inefficiency of CNF comes from the large number of steps of solving ODEs. With adaptive ODE solvers, the number of steps is dynamic. The solvers decrease the step sizes when the truncation errors of solving OEDs are too large, and vice versa. Therefore, to reduce the number of steps, we propose a method to reduce the truncation errors when solving ODEs. To do so, we focus on reducing the truncation error of solving y(t) but not log p. Although we actually have two ODEs to solve (equations of (2) and (3)), both their right-hand-sides depend only on y.
Theoretically, for a k-th order ODE solver to solve the equation dy dt = v(y, t), the local truncation error of y is O d k+1 y dt k+1 ∆t k+1 , where ∆t is the size of step. To increase the step size ∆t while maintaining the magnitude of truncation error, d k+1 y dt k+1 should be as small as possible. Obviously, the optimal situation is that y(t) is a polynomial function of degree d ≤ k, because d k+1 y dt k+1 = 0. To force the trajectory y(t) to approximate a polynomial function, we add a loss term L p called Trajectory Polynomial Regularization (TPR) into the total loss.
where L 0 is the original loss function of the corresponding task and α is a constant coefficient. L p is derived as follows. First, a number of time steps {τ 0 , . . . , τ n−1 } are randomly sampled. Next, a polynomial function f (t) is fitted to the points {y(τ 0 ), . . . , y(τ n−1 )}. And finally, L p is calculated as the mean squared error (MSE) between f (τ i ) and y(τ i ).
When minimizing L p , y(t) will approach the polynomial function, f (t). Figure 1 illustrates the mechanism and effect of our method. Figure 1(a) demonstrates how TPR regularizes y(t) to approximate a polynomial function. The loss L p is defined as the squared difference between the two curves. So it act as a force to pull the four points from the trajectory to the fitted polynomial regression. Figure 1(b) shows a possible result without TPR. The trajectories of the solutions y(t) can be very winding, so solving y(t) can be more difficult. Because the truncation errors are higher when solving y(t), ODE solvers then have to shrink step sizes in order to guarantee the precision of the solution, which leads to more computation in terms of higher number of function evaluations. On the other hand, Figure 1(c) shows the result using TPR with degree 1. The trajectories of y(t) are more like a straight line. As a result, y(t) is easier to solve, and thus the required number of function evaluations decreases.
To derive L p , the math expression of L p can be written as where · is Frobenius norm. f (t) = c 0 P 0 (t)+· · ·+c d P d (t) is the polynomial function we want to fit. d is the degree of the polynomial function, {c 0 , . . . , c d } is the coefficients to be fitted and {P 0 (t), . . . , P d (t)} is a polynomial basis. The matrices Y , T and C are written as follows, Y = (y(τ 0 ) · · · y(τ n−1 )) , Solving C to minimize Y − T C 2 yields to the equation, C = (T T ) −1 T Y . After substituting it back to equation (8), we can get To let the matrix computation more numerically stable, we use singular value decomposition (SVD) on T , which gives us T = U ΣV . And thus, the final form of calculating the TPR loss is written as follows, where U 1:d+1 denotes the leftmost d + 1 columns of matrix U .
Although we introduce a new loss term into the total loss, we prove that there are always optimal solutions minimizing both L and L 0 . In other words, when minimizing the total loss L to find an optimal solution, the solution is also an optimal solution of L 0 . In the following two theorems, we prove that, even with the hard constraint d 2 y dt 2 = dv dt = 0 (so L p = 0), there still exist infinitely many functions v(y, t) that can transform any distribution to any other distribution. Because L 0 is a function of log p(y(t 1 )), from the theorems, there are infinitely many v(y, t) that can simultaneously minimize L 0 and L p . Therefore, these v can minimize L and L 0 simultaneously. We present the two theorems as follows.
Theorem 1 Assume that y(t) ∈ R D and log p(y(t)) are governed by the differential equations (2) and (3), respectively. Given any distributions p 0 (x) and p 1 (x), where p 0 (x) and p 1 (x) > 0 for all x ∈ R D , there exists a vector field v(y, t) with the constraint dv dt = 0 everywhere, such that if the initial value of log p(y(t)) is log p 0 (y(t 0 )), then its final value is log p 1 (y(t 1 )).
Theorem 2 For D > 1, there are infinitely many such vector fields v(y, t).
The next question is whether we can approach the optimal solutions with a function approximator, which is usually neural networks. We can't prove or guarantee anything about that, but we provide two arguments. First, the number of optimal solutions is infinite. Compared to finite number, the neural networks are more likely to approach one of the infinitely many optimal solutions. Second, in experiments, we will show that our method doesn't affect testing loss much. The model with TPR has approximately the same testing loss as the model without TPR.   (c): Two transformations of 1D distributions by CNF either with or without our TPR. The input is a Gaussian distribution and the output is a mixture of two Gaussian distributions. The white streamlines represent the trajectories of y(t) and the color represents the density p(y(t)).
The details of the proofs are in Appendix A (see supplementary files). The basic idea to prove Theorem 1 is to actually construct a vector field v(y, t) and the solution y(t) that satisfy all constraints. The two main constraints in the theorem should be carefully considered. First, the solution should satisfy the initial and final values. Second, v(y, t) must be a well-defined function at any points y and t. For Theorem 2, the idea is to find an infinitesimal variation of v such that all the conditions are still satisfied. It should be careful that the value of the variation should be bounded in all space.

Experiments
We evaluate our proposed trajectory polynomial regularization (TPR) on two tasks, density estimation and variational autoencoder (VAE) (Kingma and Welling 2018). We compare our results with FFJORD (Grathwohl et al. 2019), the current state of the art of CNF, and the corresponding stateof-the-art normalizing flow models.
Two metrics are evaluated, testing loss and number of function evaluations (NFE). We want to see whether our model leads to lower computational cost compared to FFJORD in training, while keeping comparable training quality. To compare the computational cost, we count the average NFE per training iteration. NFE is defined as the number of evaluating the right-hand-side of the ODEs (2) and (3) when solving them. The lower NFE, the less computational cost.
In all experiments, We use exactly the same architectures and hyper-parameters of neural network as in FFJORD. Also the same as FFJORD, the package of ODE solvers Torchd-iffeq 2 (Chen et al. 2018) is used. The only thing we do is adding our TPR loss to training loss. The hyper-parameters of TPR are described below. We choose the number of sampled point as n = 4, with τ 0 = t 0 and τ n−1 = t 1 as the start and end of time steps. The other τ s are randomly sampled from a uniform distribution. The degree of the polynomial regression is set to be d = 1. We adopted the same ODE solver as FFJORD (Grathwohl et al. 2019), which is Do-pri5 (Dormand and Prince 1980). Dopri5 is a 4th order ODE solver, so we can use any degree with d ≤ 4. In this work, we only choose n = 1 in experiments to prove our concept. Despite the seemingly small degree, we have showed in Section 3 that it is powerful enough because the optimal solutions of L 0 still exist. The coefficient of the polynomial regularization is α = 5. The tolerance coefficient of Dopri5 is set to be atol = rtol = 10 −4 in training, and atol = rtol = 10 −5 in testing to ensure precision of solving ODEs. These hyperparameters above are shared in all experiments.

Density estimation
Given a data distribution p(x), the task of density estimation aims at approximating this distribution. Let q θ (x) be the approximated distribution. The loss of this task is the negative log-likelihood shown below.
We experiment density estimation on three 2D toy data and real data sets, which are five tabular data sets from (Papamakarios, Pavlakou, and Murray 2017) and two image data sets, MNIST and CIFAR10. The reconstructed distributions and the trajectories of y(t) by FFJORD and our model. To generate the trajectories, we select some grid points to be the initial values y(t 0 ) and solve the ODE in Equation (2) to obtain y(t). Each curve in the figures represents a single trajectory of y(t), and the color represents the value of time t.
Results on 2D toy data In this experiment, we test with three simple 2D toy data. The distributions of the three 2D data used for the experiment are shown in Figure 2(a). The approximated distributions using FFJORD and our model are in the upper halves of Figures 2(b) and 2(c), respectively. Both FFJORD and our model successfully recover the data distributions in decent quality, which implies that our new regularization loss doesn't harm approximation capability, and the model is able to find a good solution.
We visualize the trajectories of y(t) in the lower halves of Figures 2(b) and 2(c). We want to know whether our model actually performs how we expected as in Figure 1(c), which is that the trajectories of y(t) are approximately straight lines. It can be seen that the trajectories of FFJORD are very winding, while our trajectories are less winding and most are straight. Since we are using polynomial regularization of degree 1, the result is what we expected. Because of this, our ODEs are simpler to solve and the required step sizes for solving ODEs can be larger, so the number of steps can be smaller in training. In reality, the average NFEs of FFJORD on the three data are 90.42, 60.48 and 86.7 respectly, while the NFE of our model are just 44.17, 47.39 and 33.66. That is, our TPR leads to respectively 51.15%, 21.64% and 60.90% drop of computational cost.
Results on real data The real data we used here are tabular data, which are five data sets processed in (Papamakarios, Pavlakou, and Murray 2017), and the image data including the MNIST and CIFAR10 data sets. Besides FFJORD, we compare our model with two discrete normalizing flows: Real NVP (Dinh, Sohl-Dickstein, and Bengio 2017) and Glow (Kingma and Dhariwal 2018), and four autoregressive flows: MADE (Kingma et al. 2016), MAF (Papamakarios, Pavlakou, and Murray 2017), TAN (Oliva et al. 2018) , and MAF-DDSF (Huang et al. 2018). But the main purpose of our work is improving continuous normalizing flow, so we mainly compare our result with FFJORD.
For tabular data, we use the same neural network architectures as those in FFJORD, and for image data, the same multiscale (Dinh, Sohl-Dickstein, and Bengio 2017) architectures as in FFJORD are used. For image data, due to our hardware limitation, we used a smaller batch size (200) and smaller number of total training epoch (200). Despite these adjustments, we find the training quality of FFJORD is not affected. For tabular data the setup is the same as those in FFJORD.
The results of average NFE and negative log-likelihood on testing data are shown in Table 1. For average NFE, our model significantly outperforms FFJORD. On all data sets, our model uses fewer NFE in training. The reduction of NFE ranges from 42.3% to 71.3%. The largest reduction, 71.3%, occurs on POWER data set. Note that the actually training time is approximately proportional to NFE. Therefore, our model takes significantly less time in training.
In terms of testing loss, our model is comparable to the original FFJORD. It can be seen that our testing loss is very similar to that of FFJORD*, which is the FFJORD model run by us. Our model even produces lower testing loss on four out of seven data sets than FFJORD*. This result matches our expectation that the additional TPR loss does not affect the performance much, as discussed in Section 3. After all, For image data sets, MNIST and CIFAR10, our model also works great. The two data sets have the highest dimension of data. Therefore, they are the most difficult to train among all these data sets. The fact that our model yields good results on the two data sets indicates that our model can work on difficult data with high dimensions. For the two image data sets, the generated images sampled from our trained model are shown in Figure 3.

Variational autoencoder (VAE)
We perform VAE experiments on four data sets obtained from (Berg et al. 2018). The loss of this task is the negative evidence lower bound (ELBO).
In this experiment, we further compare the state-of-the-art discrete models, including Planar Flow (Rezende and Mohamed 2015), Inverse Autoregressive Flow (IAF) (Kingma et al. 2016) and Sylvester normalizing flow (Berg et al. 2018).
The hyper-parameters of neural network architecture and training of our model are the same as those in FFJORD. In short, the neural network is composed of the low-rank parameter encoding layers, to encode the input data into CNF. The learning rate is set to be 5 × 10 −4 and is divided by 10 when the validation error does not decrease for 35 epochs. Each model is run three times in experiment.
The result of average training NFE and testing negative ELBO is shown in Table 2. For average NFE, our model still significantly outperforms FFJORD. Our model has reduced NFE ranging from 19.3% to 32.1%. The largest reduction is on the Caltech Silhouettes data set. For the negative ELBO, our model performs very similarly to FFJORD. This exactly meets our expectation that simply adding TPR reduces NFE but keeps testing loss the same.

Conclusion
We have improved the computational efficiency of continuous normalizing flow models. High computational cost was the largest limitation of CNF. We proposed to add a loss function based on polynomial regression to regularize the trajectory shape of the ODE solution. Our method is proposed to reduce the truncation errors of solving ODE and can result in fewer NFE. Furthermore, we proved that with this regularization, there are always optimal solutions of the vector field for the original loss function, and we argued that our new regularization doesn't harm testing loss. Empirically, our model reduces a great amount of computation cost, while reaching a comparable testing to FFJORD.