Title: FFJORD: FREE-FORM CONTINUOUS DYNAMICS FOR SCALABLE REVERSIBLE GENERATIVE MODELS

Abstract: Reversible generative models map points from a simple distribution to a complex distribution through an easily invertible neural network. Likelihood-based training of these models requires restricting their architectures to allow cheap computation of Jacobian determinants. Alternatively, the Jacobian trace can be used if the transformation is specified by an ordinary differential equation. In this paper, we use Hutchinson's trace estimator to give a scalable unbiased estimate of the logdensity. The result is a continuous-time invertible generative model with unbiased density estimation and one-pass sampling, while allowing unrestricted neural network architectures. We demonstrate our approach on high-dimensional density estimation, image generation, and variational inference, improving the state-ofthe-art among exact likelihood methods with efficient sampling.

:

Reversible generative models use cheaply invertible neural networks to transform samples from a fixed base distribution. Examples include NICE (Dinh et al., 2014) , Real NVP (Dinh et al., 2017) , and Glow (Kingma & Dhariwal, 2018) . These models are easy to sample from, and can be trained by maximum likelihood using the change of variables formula. However, this requires placing awkward restrictions on their architectures, such as partitioning dimensions or using rank one weight matrices, in order to avoid an O(D 3 ) cost determinant computation.

Recently, Chen et al. ( 2018) introduced continuous normalizing flows (CNF), defining the mapping from latent variables to data using ordinary differential equations (ODE). In their model, the likelihood can be computed using trace operations costing only O(D 2 ). This allows a more flexible, but still restricted, family of network architectures to be used.

Extending this work, we introduce an unbiased stochastic estimator of the likelihood that has O(D) time cost, allowing completely unrestricted architectures. Furthermore, we have implemented GPU-based adaptive ODE solvers to train and evaluate these models on modern hardware. We call our approach Free-form Jacobian of Reversible Dynamics (FFJORD). Figure 1 shows FFJORD smoothly transforming a Gaussian distribution into a multi-modal distribution.

2. BACKGROUND: GENERATIVE MODELS AND CHANGE OF VARIABLES:

In contrast to directly parameterizing a normalized distribution (e.g. Oord et al. (2016); Germain et al. (2015) ), the change of variables formula allows one to specify a complex normalized distribution p x (x) implicitly by warping a normalized base distribution p z (z) through an invertible function f : R D → R D . Given a random variable z ∼ p z (z) the log density of x = f (z) follows log p x (x) = log p z (z) -log det ∂f (z) ∂z

where ∂f (z) /∂z is the Jacobian of f . In general, computing the log determinant has a time cost of O(D 3 ). Much work has gone into developing restricted neural network architectures which make computing the Jacobian's determinant more tractable. These approaches broadly fall into three categories:

Normalizing flows. By restricting the functional form of f , various determinant identities can be exploited (Rezende & Mohamed, 2015; Berg et al., 2018) . These models cannot be trained as generative models from data because they do not have a tractable inverse f -1 . However, they are useful for specifying approximate posteriors for variational inference (Kingma & Welling, 2014) . Autoregressive transformations. By using an autoregressive model and specifying an ordering of the dimensions, the Jacobian of f is enforced to be lower triangular (Kingma et al., 2016; Oliva et al., 2018) . These models excel at density estimation for tabular datasets (Papamakarios et al., 2017) , but require D sequential evaluations of f to invert, which is prohibitive when D is large. Partitioned transformations. Partitioning the dimensions and using affine transformations makes the determinant of the Jacobian cheap to compute, and the inverse f -1 computable with the same cost as f (Dinh et al., 2014; 2017) . This method allows the use of convolutional architectures, excelling at density estimation for image data (Dinh et al., 2017; Kingma & Dhariwal, 2018) .

Throughout this work, we refer to reversible generative models as those which use the change of variables to transform a base distribution to the model distribution while maintaining both efficient density estimation and efficient sampling capabilities using a single pass of the model.

2.1. OTHER GENERATIVE MODELS:

There exist several approaches to generative modeling approaches which do not use the change of variables equation for training. Generative adversarial networks (GANs) (Goodfellow et al., 2014) use large, unrestricted neural networks to transform samples from a fixed base distribution. Lacking a closed-form likelihood, an auxiliary discriminator model must be trained to estimate divergences or density ratios in order to provide a training signal. Autoregressive models (Germain et al., 2015; Oord et al., 2016) directly specify the joint distribution p(x) as a sequence of explicit conditional distributions using the product rule. These models require at least O(D) evaluations to sample from. Variational autoencoders (VAEs) (Kingma & Welling, 2014) use an unrestricted architecture to explicitly specify the conditional likelihood p(x|z), but can only efficiently provide a stochastic lower bound on the marginal likelihood p(x).

2.2. CONTINUOUS NORMALIZING FLOWS:

Chen et al. ( 2018) define a generative model for data x ∈ R D similar to those based on (1), but replace the warping function with an integral of continuous-time dynamics. The generative process first samples from a base distribution z 0 ∼ p z0 (z 0 ). Then, given an ODE whose dynamics are defined by the parametric function ∂z(t) /∂t = f (z(t), t; θ), we solve the initial value problem with z(t 0 ) = z 0 to obtain a data sample x = z(t 1 ). Given a datapoint x, we can compute both the point z 0 which generates x, as well as log p(x) under the model by solving the combined initial value problem:

EQUATION

which integrates the combined dynamics of z(t) and the log-density of the sample backwards in time from t 1 to t 0 . We can then compute log p(x) using the solution of (4) and adding log p z0 (z 0 ). The existence and uniqueness of (4) require that f and its first derivatives be Lipschitz continuous (Khalil, 2002) , which can be satisfied in practice using neural networks with smooth Lipschitz activations, such as softplus or tanh.

2.2.1. BACKPROPAGATING THROUGH ODE SOLUTIONS WITH THE ADJOINT METHOD:

CNFs are trained to maximize (3). This objective involves the solution to an initial value problem with dynamics parameterized by θ. For any scalar loss function which operates on the solution to an initial value problem

EQUATION

then Pontryagin (1962) shows that its derivative takes the form of another initial value problem

EQUATION

The quantity -∂L /∂z(t) is known as the adjoint state of the ODE. Chen et al. ( 2018) use a black-box ODE solver to compute z(t 1 ), and then a separate call to a solver to compute (6) with the initial value ∂L /∂z(t1). This approach is a continuous-time analog to the backpropgation algorithm (Rumelhart et al., 1986; Andersson, 2013) and can be combined with gradient-based optimization to fit the parameters θ by maximum likelihood.

3. SCALABLE DENSITY EVALUATION WITH UNRESTRICTED ARCHITECTURES:

Switching from discrete-time dynamics to continuous-time dynamics reduces the primary computational bottleneck of normalizing flows from O(D 3 ) to O(D 2 ), at the cost of introducing a numerical ODE solver. This allows the use of more expressive architectures. For example, each layer of the original normalizing flows model of Rezende & Mohamed ( 2015) is a one-layer neural network with only a single hidden unit. In contrast, the instantaneous transformation used in planar continuous normalizing flows (Chen et al., 2018) is a one-layer neural network with many hidden units. In this section, we construct an unbiased estimate of the log-density with O(D) cost, allowing completely unrestricted neural network architectures to be used.

3.1. UNBIASED LINEAR-TIME LOG-DENSITY ESTIMATION:

In general, computing Tr ( ∂f /∂z(t)) exactly costs O(D 2 ), or approximately the same cost as D evaluations of f , since each entry of the diagonal of the Jacobian requires computing a separate derivative of f (Griewank & Walther, 2008) . However, there are two tricks that can help. First, vector-Jacobian products v T ∂f ∂z can be computed for approximately the same cost as evaluating f using reverse-mode automatic differentiation. Second, we can get an unbiased estimate of the trace of a matrix by taking a double product of that matrix with a noise vector:

EQUATION

The above equation holds for any D-by-D matrix A and distribution p( ) over D-dimensional vectors such that E[ ] = 0 and Cov( ) = I. The Monte Carlo estimator derived from ( 7) is known as Hutchinson's trace estimator (Hutchinson, 1989; Adams et al., 2018) .

To keep the dynamics deterministic within each call to the ODE solver, we can use a fixed noise vector for the duration of each solve without introducing bias:

EQUATION

Typical choices of p( ) are a standard Gaussian or Rademacher distribution (Hutchinson, 1989) .

3.1.1. REDUCING VARIANCE WITH BOTTLENECK CAPACITY:

Often, there exist bottlenecks in the architecture of the dynamics network, i.e. hidden layers whose width H is smaller than the dimensions of the input D. In such cases, we can reduce the variance of Hutchinson's estimator by using the cyclic property of trace. Since the variance of the estimator for Tr(A) grows asymptotic to ||A|| 2 F (Hutchinson, 1989) , we suspect that having fewer dimensions should help reduce variance. If we view the dynamics as a composition of two functions f = g•h(z) then we observe 

When f has multiple hidden layers, we choose H to be the smallest dimension. This bottleneck trick can reduce the norm of the matrix which may also help reduce the variance of the trace estimator.

As introducing a bottleneck limits our model capacity, we do not use this trick in our experiments. However this trick can reduce variance when a bottleneck is used, as shown in our ablation studies.

3.2. FFJORD: A CONTINUOUS-TIME REVERSIBLE GENERATIVE MODEL:

Our complete method uses the dynamics defined in (2) and the efficient log-likelihood estimator of (8) to produce the first scalable and reversible generative model with an unconstrained Jacobian.

We call this method Free-Form Jacobian of Reversible Dyanamics (FFJORD). Pseudo-code of our method is given in Algorithm 1, and 

[z 0 , ∆ logp ] ← odeint(f aug , [x, 0], t 0 , t 1 )

Solve the ODE We demonstrate FFJORD on a variety of density estimation tasks, and for approximate inference in variational autoencoders (Kingma & Welling, 2014) . Experiments were conducted using a suite of GPU-based ODE-solvers and an implementation of the adjoint method for backpropagation1 . In all experiments the Runge-Kutta 4(5) algorithm with the tableau from Shampine (1986) was used to solve the ODEs. We ensure tolerance is set low enough so numerical error is negligible; see Appendix C.

We used Hutchinson's trace estimator (7) during training and the exact trace when reporting test results. This was done in all experiments except for our density estimation models trained on MNIST and CIFAR10 where computing the exact Jacobian trace was too expensive.

The dynamics of FFJORD are defined by a neural network f which takes as input the current state z(t) ∈ R D and the current time t ∈ R. We experimented with several ways to incorporate t as an input to f , such as hyper-networks, but found that simply concatenating t on to z(t) at the input to every layer worked well and was used in all of our experiments.

4.1. DENSITY ESTIMATION ON TOY 2D DATA:

We first train on 2 dimensional data to visualize the model and the learned dynamics. 2 In Figure 2 , we show that by warping a simple isotropic Gaussian, FFJORD can fit both multi-modal and even discontinuous distributions. The number of evaluations of the ODE solver is roughly 70-100 on all datasets, so we compare against a Glow model with 100 discrete layers.

The learned distributions of both FFJORD and Glow can be seen in Figure 2 . Interestingly, we find that Glow learns to stretch the unimodal base distribution into multiple modes but has trouble modeling the areas of low probability between disconnected regions. In contrast, FFJORD is capable of modeling disconnected modes and can also learn convincing approximations of discontinuous density functions (middle row in Figure 2 ). Since the main benefit of FFJORD is the ability to train with deeper dynamics networks, we also compare against planar CNF (Chen et al be viewed as a single hidden layer network. Without the benefit of a flexible network, planar CNF is unable to model complex distributions.

4.2. DENSITY ESTIMATION ON REAL DATA:

We perform density estimation on five tabular datasets preprocessed as in Papamakarios et al. (2017) and two image datasets; MNIST and CIFAR10. When reproducing Glow, we use the same configurations for Real NVP as Papamakarios et al. (2017) and add invertible fully connected layer between all coupling layers. On the tabular datasets, FFJORD performs the best out of reversible models by a wide margin but is outperformed by recent autoregressive models. Of those, FFJORD outperforms MAF (Papamakarios et al., 2017 ) on all but one dataset and manages to outperform TAN Oliva et al.

(2018) on the MINIBOONE dataset. These models require O(D) sequential computations to sample from while the best performing method, MAF-DDSF (Huang et al., 2018) , cannot be sampled from without resorting to correlated or expensive sampling algorithms such as MCMC.

On MNIST we find that FFJORD can model the data as effectively as Glow and Real NVP using only a single flow defined by a single neural network. This is in contrast to Glow and Real NVP which must compose many flows to achieve similar performance. When we use multiple flows in a multiscale architecture (like those used by Glow and Real NVP) we obtain better performance on MNIST and comparable performance to Glow on CIFAR10. Notably, FFJORD is able to achieve this performance while using less than 2% as many parameters as Glow. We also note that Glow uses a learned base distribution whereas FFJORD and Real NVP use a fixed Gaussian. A summary of our results on density estimation can be found in Table 2 and samples can be seen in Figure 3 . Full details on architectures used, our experimental procedure, and additional samples can be found in Appendix B.1.

In general, our approach is slower than competing methods, but we find the memory-efficiency of the adjoint method allows us to use much larger batch sizes than those methods. On the tabular datasets we used a batch sizes up to 10,000 and on the image datasets we used a batch size of 900.

4.3. VARIATIONAL AUTOENCODER:

We compare FFJORD to other normalizing flows for use in variational inference. We train a VAE (Kingma & Welling, 2014) In VAEs it is common for the encoder network to also output the parameters of the flow as a function of the input x. With FFJORD, we found this led to differential equations which were too difficult to integrate numerically. Instead, the encoder network outputs a low-rank update to a global weight matrix and an input-dependent bias vector. When used in recognition nets, neural network layers defining the dynamics inside FFJORD take the form

EQUATION

)

where h is the input to the layer, σ is an element-wise activation function, D in and D out are the input and output dimension of this layer, and Û (x), V (x), b(x) are input-dependent parameters returned from an encoder network. A full description of the model architectures used and our experimental setup can be found in Appendix B.2.

On every dataset tested, FFJORD outperforms all other competing normalizing flows. A summary of our variational inference results can be found in Table 3 . We performed a series of ablation experiments to gain a better understanding of the proposed model.

5. ANALYSIS AND DISCUSSION:



5.1. FASTER TRAINING WITH BOTTLENECK TRICK:

We plotted the training losses on MNIST using an encoder-decoder architecture (see Appendix B.1 for details). Loss during training is plotted in Figure 4 , where we use the trace estimator directly on the D×D Jacobian, or we use the bottleneck trick to reduce the dimension to H × H. Interestingly, we find that while the bottleneck trick (9) can lead to faster convergence when the trace is estimated using a Gaussian-distributed , we did not observe faster convergence when using a Rademacherdistributed .

5.2. NUMBER OF FUNCTION EVALUATIONS VS. DATA DIMENSION:

The full computational cost of integrating the instantaneous change of variables (2) is O(DH L) where D is dimensionality of the data, H is the size of the hidden state, and L is the number of function evaluations (NFE) that the adaptive solver uses to integrate the ODE. In general, each evaluation of the model is O(DH) and in practice, H is typically chosen to be close to D. Since the general form of the discrete change of variables equation (1) requires O(D 3 )-cost, one may wonder whether the number of evaluations L depends on D. We train VAEs using FFJORD flows with increasing latent dimension D. The NFE throughout training is shown in Figure 5 . In all models, we find that the NFE increases throughout training, but converges to the same value, independent of D. We conjecture that the number of evaluations is not dependent on the dimensionality of the data but the complexity of its distribution, or more specifically, how difficult it is to transform its density into the base distribution. Crucial to the scalability of Real NVP and Glow is the multiscale architecture originally proposed in Dinh et al. (2017) . We compare a single-scale encoder-decoder style FFJORD with a multiscale FFJORD on the MNIST dataset where both models have a comparable number of parameters and plot the total NFE-in both forward and backward passes-against the loss achieved in Figure 6 . We find that while the single-scale model uses approximately one half as many function evaluations as the multiscale model, it is not able to achieve the same performance as the multiscale model.

5.3. SINGLE-SCALE:



6. SCOPE AND LIMITATIONS:

Number of function evaluations can be prohibitive. The number of function evaluations required to integrate the dynamics is not fixed ahead of time, and is a function of the data, model architecture, and model parameters. This number tends to grow as the models trains and can become prohibitively large, even when memory stays constant due to the adjoint method. Various forms of regularization such as weight decay and spectral normalization (Miyato et al., 2018) can be used to reduce the this quantity, but their use tends to hurt performance slightly.

Limitations of general-purpose ODE solvers. In theory, our model can approximate any differential equation (given mild assumptions based on existence and uniqueness of the solution), but in practice our reliance on general-purpose ODE solvers restricts us to non-stiff differential equations that can be efficiently solved. ODE solvers for stiff dynamics exist, but they evaluate f many more times to achieve the same error. We find that a small amount of weight decay regularizes the ODE to be sufficiently non-stiff.

7. CONCLUSION:

We have presented FFJORD, a reversible generative model for high-dimensional data which can compute exact log-likelihoods and can be sampled from efficiently. Our model uses continuoustime dynamics to produce a generative model which is parameterized by an unrestricted neural network. All required quantities for training and sampling can be computed using automatic differentiation, Hutchinson's trace estimator, and black-box ODE solvers. Our model stands in contrast to other methods with similar properties which rely on restricted, hand-engineered neural network architectures. We demonstrated that this additional flexibility allows our approach to achieve on-par or improved performance on density estimation and variational inference.

We believe there is much room for further work exploring and improving this method. FFJORD is empirically slower to evaluate than other reversible models like Real NVP or Glow, so we are interested specifically in ways to reduce the number of function evaluations used by the ODE-solver without hurting predictive performance. Advancements like these will be crucial in scaling this method to even higher-dimensional datasets.

filters mirror the first four layers and up-sample by two every other layer. The softplus activation function is used in every layer.

The multiscale architectures were inspired by those presented in Dinh et al. (2017) . We compose multiple flows together interspersed with "squeeze" operations which down-sample the spatial resolution of the images and increase the number of channels. These operations are stacked into a "scale block" which contains N flows, a squeeze, then N flows. For MNIST we use 3 scale blocks and for CIFAR10 we use 4 scale blocks and let N = 2 for both datasets. Each flow is defined by 3 convolutional layers with 64 filters and a kernel size of 3. The softplus nonlinearity is used in all layers.

Both models were trained with the Adam optimizer (Kingma & Ba, 2015) . We trained for 500 epochs with a learning rate of .001 which was decayed to .0001 after 250 epochs. Training took place on six GPUs and completed after approximately five days.

B.2 VARIATIONAL AUTOENCODER:

Our experimental procedure exactly mirrors that of Berg et al. (2018) . We use the same 7-layer encoder and decoder, learning rate (.001), optimizer (Adam Kingma & Ba ( 2015)), batch size (100), and early stopping procedure (stop after 100 epochs of no validaiton improvment). The only difference was in the nomralizing flow used in the approximate posterior.

We performed a grid-search over neural network architectures for the dynamics of FFJORD. We searched over networks with 1 and 2 hidden layers and hidden dimension 512, 1024, and 2048. We used flows with 1, 2, or 5 steps and wight matrix updates of rank 1, 20, and 64. We use the softplus activation function for all datasets except for Caltech Silhouettes where we used tanh. The best performing models can be found in the Huang et al. (2018) . In reproducing Glow, we were able to get comparable results to the reported Real NVP by removing the invertible fully connected layers.

APPENDIX C NUMERICAL ERROR FROM THE ODE SOLVER ODE solvers are numerical integration methods so there is error inherent in their outputs. Adaptive solvers (like those used in all of our experiments) attempt to predict the errors that they accrue and modify their step-size to reduce their error below a user set tolerance. It is important to be aware of this error when we use these solvers for density estimation as the solver outputs the density that we report and compare with other methods. When tolerance is too low, we run into machine precision errors. Similarly when tolerance is too high, errors are large, our training objective becomes biased and we can run into divergent training dynamics.

Since a valid probability density function integrates to one, we take a model trained on Figure 1 and numerically find the area under the curve using Riemann sum and a very fine grid. We do this for a range of tolerance values and show the resulting error in Figure 8 . We set both atol and rtol to the same tolerance. The numerical error follows the same order as the tolerance, as expected. During training, we find that the error becomes non-negligible when using tolerance values higher than 10 -5 . For most of our experiments, we set tolerance to 10 -5 as that gives reasonable performance while requiring few number of evaluations. For the tabular experiments, we use atol=10 -8 and rtol=10 -6 .

Code can be found at https://github.com/rtqichen/ffjord and https://github.com/rtqichen/torchdiffeq.

Videos of the learned dynamics can be found at https://imgur.com/a/Rtr3Mbq.

Figures and Tables:

Figure 1: FFJORD transforms a simple base distribution at t 0 into the target distribution at t 1 by integrating over learned continuous dynamics.

Figure 2: Comparison of trained Glow, planar CNF, and FFJORD models on 2-dimensional distributions, including multi-modal and discontinuous densities.

Figure 5: NFE used by the adaptive ODE solver is approximately independent of data-dimension. Lines are smoothed using a Gaussian filter.

Figure6: For image data, a single FFJORD flow can achieve near performance to multi-scale architecture while using half the number of evaluations.

Figure 8: Numerical integration shows that the density under the model does integrate to one given sufficiently low tolerance. Both log and non-log plots are shown.

Table 1 summarizes the capabilities of our model compared to other recent generative modeling approaches. Unbiased stochastic log-density estimation using the FFJORD model Require: dynamics f θ , start time t 0 , stop time t 1 , data samples x, data dimension D. ← sample unit variance(x.shape) Sample outside of the integral function f aug ([z t , log p t ], t): Augment f with log-density dynamics. f t ← f θ (z(t), t)

Table content: <table><tr><td/><td>Evaluate neural network</td></tr><tr><td>g ← T ∂f ∂z z(t)</td><td>Compute vector-Jacobian product with automatic differentiation</td></tr><tr><td>Tr = g</td><td>Unbiased estimate of Tr( ∂f ∂z ) with T ∂f</td></tr></table>

Negative log-likehood on test data for density estimation models; lower is better. In nats for tabular data and bits/dim for MNIST and CIFAR10. *Results use multi-scale convolutional architectures.

Table content: <table><tr><td/><td colspan="5">POWER GAS HEPMASS MINIBOONE BSDS300</td><td>MNIST</td><td>CIFAR10</td></tr><tr><td>Real NVP</td><td>-0.17</td><td>-8.33</td><td>18.71</td><td>13.55</td><td>-153.28</td><td>1.06*</td><td>3.49*</td></tr><tr><td>Glow</td><td>-0.17</td><td>-8.15</td><td>18.92</td><td>11.35</td><td>-155.07</td><td>1.05*</td><td>3.35*</td></tr><tr><td>FFJORD</td><td>-0.46</td><td>-8.59</td><td>14.92</td><td>10.43</td><td colspan="2">-157.40 0.99* (1.05  † )</td><td>3.40*</td></tr><tr><td>MADE</td><td>3.08</td><td>-3.56</td><td>20.98</td><td>15.59</td><td>-148.85</td><td>2.04</td><td>5.67</td></tr><tr><td>MAF</td><td colspan="2">-0.24 -10.08</td><td>17.70</td><td>11.75</td><td>-155.69</td><td>1.89</td><td>4.31</td></tr><tr><td>TAN</td><td colspan="2">-0.48 -11.19</td><td>15.12</td><td>11.01</td><td>-157.03</td><td>-</td><td>-</td></tr><tr><td>MAF-DDSF</td><td colspan="2">-0.62 -11.96</td><td>15.09</td><td>8.86</td><td>-157.73</td><td>-</td><td>-</td></tr></table>

Negative ELBO on test data for VAE models; lower is better. In nats for all datasets except Frey Faces which is presented in bits per dimension. Mean/stdev are estimated over 3 runs. et al., 2016), and Sylvester normalizing flows(Berg et al., 2018). To provide a fair comparison, our encoder/decoder architectures and learning setup exactly mirror those ofBerg et al. (2018).

Table content: <table><tr><td/><td>MNIST</td><td>Omniglot</td><td colspan="2">Frey Faces Caltech Silhouettes</td></tr><tr><td>No Flow</td><td colspan="3">86.55 ± .06 104.28 ± .39 4.53 ± .02</td><td>110.80 ± .46</td></tr><tr><td>Planar</td><td colspan="3">86.06 ± .31 102.65 ± .42 4.40 ± .06</td><td>109.66 ± .42</td></tr><tr><td>IAF</td><td colspan="3">84.20 ± .17 102.41 ± .04 4.47 ± .05</td><td>111.58 ± .38</td></tr><tr><td colspan="2">Sylvester 83.32 ± .06</td><td>99.00 ± .04</td><td>4.45 ± .04</td><td>104.62 ± .29</td></tr><tr><td colspan="4">FFJORD 82.82 ± .01 98.33 ± .09 4.39 ± .01</td><td>104.03 ± .43</td></tr></table>

Models were trained on a single GPU and training took between four hours and three days depending on the dataset.

Table content: <table><tr><td>Dataset</td><td colspan="5">nonlinearity # layers hidden dim multiplier # flow steps batchsize</td></tr><tr><td>POWER</td><td>tanh</td><td>3</td><td>10</td><td>5</td><td>10000</td></tr><tr><td>GAS</td><td>tanh</td><td>3</td><td>20</td><td>5</td><td>1000</td></tr><tr><td>HEPMASS</td><td>softplus</td><td>2</td><td>10</td><td>10</td><td>10000</td></tr><tr><td>MINIBOONE</td><td>softplus</td><td>2</td><td>20</td><td>1</td><td>1000</td></tr><tr><td>BSDS300</td><td>softplus</td><td>3</td><td>20</td><td>2</td><td>10000</td></tr></table>

Best performing model architectures for density estimation on tabular data with FFJORD.

Table content: <table><tr><td>Dataset</td><td colspan="5">nonlinearity # layers hidden dimension # flow steps rank</td></tr><tr><td>MNIST</td><td>softplus</td><td>2</td><td>1024</td><td>2</td><td>64</td></tr><tr><td>Omniglot</td><td>softplus</td><td>2</td><td>512</td><td>5</td><td>20</td></tr><tr><td>Frey Faces</td><td>softplus</td><td>2</td><td>512</td><td>2</td><td>20</td></tr><tr><td>Caltech</td><td>tanh</td><td>1</td><td>2048</td><td>1</td><td>20</td></tr></table>

Best performing model architectures for VAEs with FFJORD. ± 0.01 -8.33 ± 0.14 18.71 ± 0.02 13.55 ± 0.49 -153.28 ± 1.78 Glow -0.17 ± 0.01 -8.15 ± 0.40 18.92 ± 0.08 11.35 ± 0.07 -155.07 ± 0.03 FFJORD -0.46 ± 0.01 -8.59 ± 0.12 14.92 ± 0.08 10.43 ± 0.04 -157.40 ± 0.19 MADE 3.08 ± 0.03 -3.56 ± 0.04 20.98 ± 0.02 15.59 ± 0.50 -148.85 ± 0.28 MAF -0.24 ± 0.01 -10.08 ± 0.02 17.70 ± 0.02 11.75 ± 0.44 -155.69 ± 0.28 TAN -0.48 ± 0.01 -11.19 ± 0.02 15.12 ± 0.02 11.01 ± 0.48 -157.03 ± 0.07 MAF-DDSF -0.62 ± 0.01 -11.96 ± 0.33 15.09 ± 0.40 8.86 ± 0.15 -157.73 ± 0.04

Table content: <table><tr><td colspan="4">B.3 STANDARD DEVIATIONS FOR TABULAR DENSITY ESTIMATION</td></tr><tr><td/><td>POWER</td><td>GAS</td><td>HEPMASS MINIBOONE</td><td>BSDS300</td></tr><tr><td>Real NVP</td><td>-0.17</td><td/><td/></tr></table>

Negative log-likehood on test data for density estimation models. Means/stdev over 3 runs. Real NVP, MADE, MAF, TAN, and MAF-DDSF results on are taken from

Table content: <table/>

Peer review (iclr2019):
 Summary: This paper discusses an advance in the framework of normalizing flows for generative modeling, named FFJORD. The authors consider normalizing flows in the form of ordinary differential equations, as also discussed in [1]. Their contributions are two-fold: (1) they use an unbiased estimator of the likelihood of the model by approximating the trace of the jacobian with Hutchinson’s trace estimator, (2) they have implemented the required ODE solvers on GPUs. The models are evaluated on a density estimation task on tabular data and two image datasets (MNIST and CIFAR10), as well as on variational inference for auto-encoders, where the datasets MNIST, Omniglot, Freyfaces and Caltech Silhouettes are considered. The authors argue that the trace estimator, in combination with reverse-mode automatic differentiation to compute vector-Jacobian products, leads to a computational cost of O(D), instead of O(D^2) for the exact trace of the jacobian. They compare this to the cost of computing a Jacobian determinant for finite flows, which is O(D^3) in general. They argue that in general all works on finite flows have adjusted their architectures for the flows to avoid the O(D^3) complexity, and that FFJORD has no such restriction. However, I would like the authors to comment on the following train of thought: autoregressive models, such as MAF, as well as IAF (inverse of an autoregressive model) do not require O(D^3) to compute jacobian determinants as the jacobian is of triangular form. Note however, they are still universal approximators if sufficient flows are applied, as any distribution can be factorized in an autoregressive manner. With this in mind, I find the red cross for MAF under free-form Jacobian slightly misleading. Perhaps I misunderstood something, so please clarify. Another topic that I would like the authors to comment on is efficiency and practical use. One of the main points that the authors seem to emphasise, is that contrary to autoregressive models, which require D passes through the model to sample a datapoint of size D, FFJORD is a ‘single-pass’ model, requiring only one pass through the model. They therefore indicate that they can do efficient sampling. However, for FFJORD every forward pass requires a pass through an ODE solver, which as the authors also state, can be very slow. I could imagine that this is still faster than an autoregressive model, but I doubt this is actually of comparable speed to a forward pass of a finite flow such as glow or realNVP. On the other hand, autoregressive models do not require D passes during training, whereas, if I understand correctly, FFJORD relies on two passes through ODE solvers, one for computing the loss, and a second to compute the gradient of the loss with respect to model parameters. So autoregressive models should train considerably faster. The authors do comment on the fact that FFJORD is slower than other models, but they do not give a hint as to how much slower it is. This would be of importance for practical use, and for other people to consider using FFJORD in future work. For the density estimation task, FFJORD does not have the best performance compared other baselines, except for MNIST, for which the overall best model was not evaluated (MAF-DDSF). For variational inference FFJORD is stated to outperform all other flows, but the models are only evaluated on the negative evidence lower bound, and not on the negative log-likehood (NLL). I suspect the NLL to be absent from the paper as it requires more computation, and this takes a long time for FFJORD. Without an evaluation on NLL the improvement over other methods is questionable. Even if the improvement still holds for the NLL, the relative improvement might not weigh heavily enough against increased runtime. FFJORD does require less memory than its competitors. The improved runtime by implementing the ODE solvers on GPU versus the runtime on a CPU would be useful, given that this is listed as one of the main contributions. Besides these questions/comments, I do think the idea of using Hutchinsons trace estimator is a valid contribution, and the experimental validation of continuous normalizing flows is of interest to the research community. Therefore, in my opinion, the community will benefit from the information in this paper, and it should be accepted. However I do wish for the authors to address the above questions as it would give a clearer view of the practical use of the proposed model. See below for comments and questions: Quality The paper has a good setup, and is well structured. The scope and limitations section is very much appreciated. Clarity The paper is clearly written overall. The only section I can comment on is the related work section, which is not the best part of the paper. The division in normalizing flows and partitioned transformations is a bit odd. Partitioned transformations surely are also normalizing flows. Furthermore IAF by Kingma et al. is put in the box of autoregressive models, whereas it is the inverse of an autoregressive model, such that it does not have the D-pass sample problem. For a reader who is not too familiar with normalizing flows literature, I think this section is a little confusing. Furthermore, there is no related work discussed on continuous time flows, such as (but not limited to) [2]. Originality The originality of the paper is not stellar, but sufficient for acceptance. Significance The community can benefit from the experimental analysis of continuous time flows, and the GPU implementation of the ODE solver. Therefore I think this work is significant. Detailed questions/comments: 1. In section 4.2, as an additional downside to MAF-DDSF, the authors argue that sampling cannot be performed analytically. Since FFJORD needs to numerically propagate the ODE, I do not think FFJORD can sample analytically either. Is this correct? 2. The authors argue that they have no restriction on the architecture of the function f, even if they have O(D) estimation of the trace of the jacobian. However, they also say they make use of the bottle-neck trick to reduce the variance that arises due to Hutchinson’s estimate of the trace. This seems like a limitation on the architecture to me. Can the authors comment? 3. In B.1 in the appendix, the street view house numbers dataset is mentioned, but no results appear in the main text, why not? 4. In the results section, it is not clear to me which numbers of the baselines for different datasets are taken from other papers, and which numbers are obtained by the authors of this paper. Please clarify. 5. In the conclusions, when discussing future work, the authors state that they are interested in reducing the number of function evaluations in the ODE solvers. In various disciplines many people have worked on this problem for a long time. Do the authors think major improvements are soon to be made? 6. In section 5.2 the dependence of the number of function evaluations (NFE) on the data dimension D is discussed. As a thought experiment they use the fact that going from an isotropic gaussian distribution (in any D), to an isotropic gaussian distribution has a corresponding differential equation of zero. This should convince the reader that NFE is independent of D. However, this seems to me to be such a singular example, that I gain no insight from it, and it is not very convincing. Do the authors agree that this particular example does not add much? If not, please explain. [1] Chen et al. Neural ordinary differential equations. NIPS 2018 [2] Chen et al. Continuous-time flows for deep generative models. **** EDIT ***** I have read the response of the authors and appreciate their clarifications and the additional information on the runtimes. See my response below for the concern that remains about the absence of the estimate of the log likelihood for the VAE experiments. Besides this issue, the other comments/answers were satisfactory, and I think this paper is of interest to the research community, so I will stick with my score.