PaperlayerView exact revision on arXiv

Source-linked AI summary

FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models

Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, David Duvenaud

arXiv:1810.01367v3cs.LGcs.CVstat.ML

TL;DR

Likelihood-based reversible generative models often restrict architectures to make Jacobian determinant computations tractable. FFJORD uses continuous-time dynamics and an unbiased stochastic log-density estimator, and the paper reports improved performance on density estimation and variational inference.

  • Problem

    Likelihood-based training of reversible generative models requires architectural restrictions to make Jacobian determinant computations tractable.

    Source: §1 · p. 1

  • Method

    FFJORD combines continuous-time dynamics with an unbiased stochastic log-likelihood estimator, allowing unrestricted neural network architectures.

    Source: §1 · p. 1, §3.2 · p. 4

  • Results

    The authors report improved performance on density estimation and variational inference.

    Source: §7 · p. 8

  • Takeaways & Limitations

    FFJORD provides a reversible model with exact log-likelihoods and efficient sampling while using an unrestricted neural network.

    Source: §7 · p. 8

    The adaptive solver’s function-evaluation count can grow during training and become prohibitively large, while regularization to reduce it can slightly hurt performance.

    Source: §6 · p. 8

Abstract

from arXiv · show

A promising class of generative models maps points from a simple distribution to a complex distribution through an 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 log-density. 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, achieving the state-of-the-art among exact likelihood methods with efficient sampling.

1 INTRODUCTION

FFJORD uses an unbiased stochastic likelihood estimator with O(D) time cost to enable unrestricted neural architectures in continuous-time generative models. It addresses architectural restrictions in reversible flows that avoid costly Jacobian determinant computation.

1 INTRODUCTION

  • Reversible generative models are easy to sample and train by maximum likelihood, but avoiding O(D3) determinant computation requires awkward architectural restrictions such as dimension partitioning or rank-one weight matrices.Examples include NICE (Dinh et al., 2014), Real NVP (Dinh et al., 2017), and Glow (Kingma & Dhariwal, 2018).Source: §1 · p. 1
  • Chen et al. (2018) introduced continuous-time normalizing flows using ODEs, where likelihood computation uses relatively cheap trace operations, though avoiding O(D2) cost still restricts network architectures.Source: §1 · p. 1
  • FFJORD introduces an unbiased stochastic likelihood estimator with O(D) time cost, enabling completely unrestricted architectures, and uses GPU-based adaptive ODE solvers for training and evaluation.The approach is called Free-form Jacobian of Reversible Dynamics (FFJORD).Source: §1 · p. 1

2 BACKGROUND: GENERATIVE MODELS AND CHANGE OF VARIABLES

Change-of-variables models warp a normalized base distribution through an invertible function, but computing the Jacobian determinant is generally O(D^3), motivating architectures that make it tractable. Continuous normalizing flows replace a single warp with continuous-time dynamics, while adjoint equations enable gradient-based training.

2 BACKGROUND: GENERATIVE MODELS AND CHANGE OF VARIABLES

  • An invertible map f warps a normalized base distribution into an implicitly specified normalized distribution, but its Jacobian log-determinant generally costs O(D^3).Reversible generative models are defined here as supporting efficient density estimation and sampling in a single pass.Source: §2 · p. 2
  • Normalizing flows exploit determinant identities but lack tractable analytic inverses; autoregressive maps have triangular Jacobians but need D evaluations to invert; partitioned affine maps offer cheap determinants and inverses.Normalizing flows can represent approximate posteriors; autoregressive models excel on tabular density estimation, while partitioned transformations support convolutional architectures and excel on image density estimation.Source: §2 · p. 2

2.1 OTHER GENERATIVE MODELS

  • GANs lack closed-form likelihoods, autoregressive models require at least O(D) evaluations to sample, and VAEs efficiently provide only a lower bound on marginal likelihood.GANs use a discriminator to estimate divergences; autoregressive models specify the joint as conditional distributions, while VAEs specify p(x|z).Source: §2.1 · p. 2

2.2 CONTINUOUS NORMALIZING FLOWS

  • Chen et al. (2018) define continuous normalizing flows by evolving a base sample through a parameterized ODE, with log-density change governed by an instantaneous change-of-variables equation.The total log-density change is obtained by integrating over time.Source: §2.2 · p. 2
  • Solving the combined state and log-density dynamics backward from the observation recovers its base point and log-density, requiring the dynamics and first derivatives to be Lipschitz continuous.Smooth Lipschitz neural-network activations can satisfy this existence-and-uniqueness condition.Source: §2.2 · p. 3

2.2.1 BACKPROPAGATING THROUGH ODE SOLUTIONS WITH THE ADJOINT METHOD

  • The adjoint state yields a second initial-value problem for loss derivatives; Chen et al. (2018) solve the forward dynamics and then the adjoint dynamics to train parameters θ with gradient-based optimization.This provides a continuous-time analogue of backpropagation.Source: §2.2.1 · p. 3

3 SCALABLE DENSITY EVALUATION WITH UNRESTRICTED ARCHITECTURES

FFJORD combines continuous-time dynamics with an unbiased stochastic log-density estimate to enable scalable reversible generative models with unrestricted architectures. The section explains its trace estimator, a bottleneck-based variance reduction, and the resulting likelihood cost.

3 SCALABLE DENSITY EVALUATION WITH UNRESTRICTED ARCHITECTURES

  • Continuous-time dynamics reduce the primary normalizing-flow bottleneck from O(D^3) to O(D^2), at the cost of a numerical ODE solver, enabling unrestricted architectures and an O(D)-cost unbiased log-density estimate.The section contrasts this flexibility with discrete-time flows and planar continuous normalizing flows, whose layers use only a single hidden unit and many hidden units, respectively.Source: §3 · p. 3

3.1 UNBIASED LINEAR-TIME LOG-DENSITY ESTIMATION

  • Hutchinson’s trace estimator uses a zero-mean, identity-covariance noise vector and a vector-Jacobian product to estimate the Jacobian trace without bias.A fixed noise vector can be used throughout each ODE solve without introducing bias; standard Gaussian and Rademacher choices are noted.Source: §3.1 · p. 4

3.1.1 REDUCING VARIANCE WITH BOTTLENECK CAPACITY

  • When the dynamics network has hidden layers narrower than the input, choosing its smallest hidden dimension for the trace estimate can reduce estimator variance.The method uses the cyclic property of trace; the passages suggest fewer dimensions and a reduced matrix norm may help lower variance.Source: §3.1.1 · p. 4

3.2 FFJORD: A CONTINUOUS-TIME REVERSIBLE GENERATIVE MODEL

  • FFJORD reduces likelihood cost to O((DH + D) Lhat), versus O((DH + D^2) Lhat) for CNFs, while enabling an unconstrained Jacobian.Here, H is the largest hidden dimension and Lhat is the number of ODE-solver evaluations of f.Source: §3.2 · p. 4

4 EXPERIMENTS

Experiments evaluate FFJORD on toy and real-data density estimation and variational inference. The reported results include fitting multimodal and discontinuous densities, density-estimation comparisons, and VAE performance.

4.1 DENSITY ESTIMATION ON TOY 2D DATA

  • FFJORD fits multimodal and discontinuous toy densities, while Glow has difficulty modeling low-probability regions between disconnected modes.The ODE solver takes roughly 70–100 evaluations across datasets, compared with Glow's 100 discrete layers.Source: §4.1 · pp. 5–6, Figure 2 · p. 5

4.2 DENSITY ESTIMATION ON REAL DATA

  • On tabular datasets, FFJORD leads reversible models by a wide margin, beats MAF (Papamakarios et al., 2017) on all but one dataset, and beats TAN (Oliva et al., 2018) on MINIBOONE.Recent autoregressive models still outperform FFJORD overall; MAF-DDSF (Huang et al., 2018), the best-performing method, cannot be sampled from analytically, while the other autoregressive models require O(D) sequential sampling computations.Source: §4.2 · p. 6, Table 2 · p. 6
  • On MNIST, a single FFJORD flow matches Glow and Real NVP, while multiple flows improve MNIST performance and yield CIFAR10 performance comparable to Glow with less than 2% of its parameters.Glow uses a learned base distribution, whereas FFJORD and Real NVP use a fixed Gaussian; Table 2 summarizes density-estimation results and Figure 3 shows samples.Source: §4.2 · p. 6, Table 2 · p. 6, Figure 3 · p. 6
  • FFJORD is generally slower than competing methods, but the adjoint method's memory efficiency permits batches up to 10,000 on tabular data and 900 on image data.The experiments used batch sizes up to those values for the respective datasets.Source: §4.2 · p. 6

4.3 VARIATIONAL AUTOENCODER

  • For VAE use, FFJORD's encoder outputs a low-rank update to a global weight matrix and an input-dependent bias because directly predicting flow parameters made the ODEs too difficult to integrate.FFJORD layers use data-dependent parameters returned by the encoder.Source: §4.3 · p. 7
  • On every dataset tested, FFJORD outperforms all other competing normalizing flows. Table 3 reports VAE test-data negative ELBO (lower is better), with means and standard deviations over 3 runs, in nats except Frey Faces (bits per dimension).Source: §4.3 · p. 7, Table 3 · p. 7

5 ANALYSIS AND DISCUSSION

The analysis uses ablations to examine bottleneck effects on training, how function-evaluation counts scale with data dimension, and the tradeoff between single- and multiscale architectures.

5.1 FASTER TRAINING WITH BOTTLENECK TRICK

  • The bottleneck trick speeds convergence with Gaussian noise, but not Rademacher noise; the figure attributes the effect to reduced log-density-estimator variance.The comparison uses MNIST with an encoder-decoder architecture.Source: §5.1 · p. 7, Figure 4 · p. 7

5.2 NUMBER OF FUNCTION EVALUATIONS VS. DATA DIMENSION

  • Across latent dimensions, NFE rises during training but converges to the same value; evaluation count depends on distributional complexity, not data dimension.The integration cost is O(DH bL), with each model evaluation costing O(DH).Source: Figure 5 · p. 7, §5.2 · pp. 7–8

5.3 SINGLE-SCALE VS. MULTI-SCALE FFJORD

  • On MNIST, single-scale FFJORD uses about half as many function evaluations as multiscale FFJORD but does not achieve the same performance.The compared models have a comparable number of parameters.Source: §5.3 · p. 8

6 SCOPE AND LIMITATIONS

FFJORD’s computational cost depends on the dynamics and architecture: function evaluations can grow during training, and general-purpose solvers limit practical use to efficiently solvable non-stiff dynamics. Multiscale image architectures require more evaluations but achieve better performance.

6 SCOPE AND LIMITATIONS

  • For image data, multiscale architectures require more ODE function evaluations but achieve better performance.Source: Figure 6 · p. 8
  • NFE depends on the data, architecture, and parameters, tends to grow during training, and can become prohibitive; weight decay or spectral normalization reduce it but slightly hurt performance.The adjoint method keeps memory constant even when NFE grows.Source: §6 · p. 8
  • General-purpose ODE solvers practically restrict models to efficiently solvable non-stiff dynamics; stiff solvers require many more evaluations for the same error, while slight weight decay can constrain dynamics to be non-stiff.Source: §6 · p. 8

7 CONCLUSION

FFJORD uses continuous-time dynamics and unrestricted neural networks to provide exact log-likelihoods and efficient sampling. Its flexibility improves density estimation and variational inference, while reducing ODE-solver evaluations remains a priority for scaling to higher-dimensional datasets.

7 CONCLUSION

  • FFJORD computes exact log-likelihoods and supports efficient sampling through continuous-time dynamics parameterized by an unrestricted neural network.Training and sampling quantities use automatic differentiation, Hutchinson’s trace estimator, and black-box ODE solvers.Source: §7 · p. 8
  • The added flexibility improves density estimation and variational inference, and FFJORD models distributions that Glow and Real NVP cannot.The conclusion contrasts FFJORD with methods that rely on restricted, hand-engineered neural network architectures.Source: §7 · p. 8
  • Reducing ODE-solver function evaluations without hurting predictive performance is a priority for scaling FFJORD to higher-dimensional datasets.The authors identify this as an area for future work.Source: §7 · p. 8

APPENDIX A QUALITATIVE SAMPLES

Figure 7 presents qualitative samples from FFJORD models trained on MNIST and CIFAR10.

APPENDIX A QUALITATIVE SAMPLES

  • Figure 7 shows samples from FFJORD models trained on MNIST and CIFAR10.Source: Appendix A · p. 11

APPENDIX B EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS · B.1 DENSITY ESTIMATION

Appendix B.1 reports the density-estimation architecture search and implementation details. It describes separate image-model approaches, their dataset-specific scope, and training settings.

B.1 DENSITY ESTIMATION

  • Tabular models were selected by grid-searching flows, hidden layers, dimension multipliers, and tanh or softplus nonlinearities.The search used 1, 2, 5, or 10 flows; 1–4 hidden layers per flow; and hidden dimensions 5, 10, or 20 times the data dimension.Source: §B.1 · p. 11
  • Image experiments compared a single-flow encoder-decoder architecture with a multiscale architecture built from multiple flows.Source: §B.1 · p. 11
  • Encoder-decoder models fit MNIST competitively but could not fit more complex image datasets, including CIFAR10 and Street View House Numbers.Source: §B.1 · p. 11
  • The MNIST encoder-decoder used four convolutional layers with 64, 64, 128, and 128 filters, then four transpose-convolutional layers that mirrored them.Strided convolutions down-sampled by two every other layer, and transpose convolutions up-sampled by two every other layer; all layers used softplus.Source: §B.1 · pp. 11–12
  • Multiscale models used 3 scale blocks for MNIST and 4 for CIFAR10, with N = 2 flows around each squeeze operation.Each flow had three convolutional layers with 64 filters and kernel size 3; all layers used softplus.Source: §B.1 · p. 12
  • Both models used Adam for 500 epochs, with the learning rate reduced from .001 to .0001 after 250 epochs.Training used six GPUs and took approximately five days.Source: §B.1 · p. 12

B.2 VARIATIONAL AUTOENCODER · B.3 STANDARD DEVIATIONS FOR TABULAR DENSITY ESTIMATION

The VAE experiments follow Berg et al. (2018), changing only the normalizing flow in the approximate posterior; FFJORD architectures are selected through a grid search. The tabular-density appendix reports test negative log-likelihood means and standard deviations across three runs.

B.2 VARIATIONAL AUTOENCODER

  • The VAE setup matches Berg et al. (2018) in encoder and decoder depth, learning rate, optimizer, batch size, and early stopping; only the approximate-posterior flow differs.The networks have 7 layers, learning rate .001, Adam, batch size 100, and stop after 100 epochs without validation improvement.Source: §B.2 · p. 12
  • FFJORD dynamics architectures were grid-searched across 1 or 2 hidden layers, dimensions 512, 1024, or 2048, 1, 2, or 5 flow steps, and rank 1, 20, or 64 matrix updates.Source: §B.2 · p. 12
  • Softplus was used for all datasets except Caltech Silhouettes, which used tanh.Source: §B.2 · p. 12
  • Models were trained on a single GPU, with training taking four hours to three days depending on the dataset.Source: §B.2 · p. 12
  • Table 5 presents the best-performing FFJORD model architectures for VAEs.Source: Table 5 · p. 12, §B.2 · p. 12

B.3 STANDARD DEVIATIONS FOR TABULAR DENSITY ESTIMATION

  • Table 6 reports test-data negative log-likelihood for density-estimation models as means and standard deviations over three runs.Source: Table 6 · p. 13

APPENDIX C NUMERICAL ERROR FROM THE ODE SOLVER

This appendix examines how ODE-solver tolerance affects density estimates: numerical error tracks the tolerance, and both overly high and overly low tolerances pose risks.

APPENDIX C NUMERICAL ERROR FROM THE ODE SOLVER

  • Adaptive solvers adjust step size to keep error below a user-set tolerance, but tolerance that is too high can bias training and cause divergence, while tolerance that is too low can trigger machine-precision errors.Because solver outputs determine the reported density estimates, numerical error affects density estimation and comparisons with other methods.Source: Appendix C · p. 13
  • The density integrates to one at sufficiently low tolerance; error tracks tolerance, becomes non-negligible above 10^-5, and most experiments use 10^-5, while tabular experiments use atol=10^-8 and rtol=10^-6.The error is measured by numerically integrating the density with a fine-grid Riemann sum, setting atol and rtol equal for the tolerance sweep.Source: Figure 8 · p. 13, Appendix C · p. 13

Related papers

Loading 1810.01367v3…