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

Existing reversible generative models make exact likelihood training tractable by restricting neural-network architectures, limiting flexibility. FFJORD uses continuous-time dynamics with Hutchinson’s trace estimator to obtain an unbiased O(D) likelihood estimate and unrestricted architectures. Across density estimation, image generation, and variational inference, it reports improved performance among exact-likelihood methods with efficient sampling, while function evaluations and solver stiffness remain practical limitations.

  • Problem

    Exact likelihood training for reversible generative models requires cheap Jacobian determinants, but this generally imposes restrictive neural-network architectures.

  • Method

    FFJORD uses continuous-time dynamics, Hutchinson’s trace estimator, automatic differentiation, and black-box ODE solvers to estimate log-likelihoods with unrestricted architectures.

  • Results

    FFJORD improves density-estimation and variational-inference performance, matches Glow and Real NVP on MNIST with less than 2% as many parameters as Glow, and samples efficiently.

  • Takeaways & Limitations

    Continuous-time reversible modeling provides exact log-likelihoods and efficient sampling while avoiding the hand-engineered architectures used by comparable methods.

  • Takeaways & Limitations

    Function evaluations can grow prohibitively large during training, and reliance on general-purpose ODE solvers restricts FFJORD to efficiently solvable non-stiff dynamics.

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

Reversible generative models enable likelihood-based training and efficient sampling but traditionally require architecturally restricted networks for tractable Jacobian determinants. FFJORD replaces this restriction with an unbiased O(D) likelihood estimator for continuous-time dynamics and supports unrestricted architectures.

  • Reversible models transform a fixed base distribution using cheaply invertible networks and train by maximum likelihood via change of variables.
  • Computing Jacobian determinants generally costs O(D3), motivating restrictions such as dimension partitioning or rank one weight matrices.
  • Continuous-time normalizing flows use ODE-defined mappings and trace operations, reducing the cost to O(D2) while retaining some architectural restrictions.
  • FFJORD introduces an unbiased stochastic likelihood estimator with O(D) time cost, allowing completely unrestricted neural network architectures.The method also uses GPU-based adaptive ODE solvers for training and evaluation.
  • FFJORD maps a simple base distribution at t0 to a target distribution at t1 by integrating learned continuous dynamics.

2 BACKGROUND: GENERATIVE MODELS AND CHANGE OF VARIABLES

Change-of-variables models represent complex normalized distributions by warping a normalized base distribution, but tractable density computation drives architectural or computational trade-offs. Continuous normalizing flows replace discrete warping with ODE dynamics, whose states and likelihoods are solved forward or backward using numerical integration and adjoint-based gradients.

  • Change of variables constructs a complex normalized distribution by applying an invertible function to samples from a normalized base distribution.
  • The Jacobian determinant in the change-of-variables formula generally costs O(D3), motivating restricted architectures and alternative transformation families.
  • Normalizing flows exploit restricted functional forms, while autoregressive transformations enforce a lower-triangular Jacobian through dimension ordering.
  • Reversible generative models use change of variables while maintaining efficient density estimation and efficient single-pass sampling.
  • GANs use unrestricted networks without closed-form likelihoods, whereas autoregressive models specify the joint distribution through explicit conditional distributions.
  • Continuous Normalizing Flows: Continuous normalizing flows sample z0, solve an ODE to obtain observable data z(t1), and track log-density change through a second differential equation.
  • Continuous Normalizing Flows: For a datapoint x, solving the coupled dynamics backward recovers its generating z0 and computes log p(x) by adding the base log-density.
  • Backpropagating Through ODE Solutions: The adjoint method computes gradients by solving another initial value problem backward from the terminal adjoint state.This provides a continuous-time analogue of backpropagation that can be combined with gradient-based optimization.

3 SCALABLE DENSITY EVALUATION WITH UNRESTRICTED ARCHITECTURES

FFJORD combines continuous-time dynamics with Hutchinson’s trace estimator to obtain unbiased, scalable log-density estimates using unrestricted neural networks. Its implementation uses fixed-noise ODE solves and adaptive GPU-based solvers, with a likelihood cost reduced to O((DH + D)L̂).

  • Motivation: Continuous-time dynamics reduce the normalizing-flow bottleneck from O(D^3) to O(D^2), while FFJORD targets O(D) trace estimation with unrestricted architectures.The continuous-time formulation introduces a numerical ODE solver but permits more expressive dynamics networks.
  • Unbiased trace estimation: Hutchinson’s estimator computes an unbiased trace estimate from a noise vector and a vector-Jacobian product, avoiding explicit Jacobian construction.Reverse-mode automatic differentiation computes the vector-Jacobian product at approximately the cost of evaluating the dynamics function.
  • Unbiased trace estimation: A fixed noise vector is used throughout each ODE solve, preserving deterministic dynamics without introducing bias.Standard Gaussian and Rademacher distributions are typical choices for the noise vector.
  • Variance reduction: The bottleneck trick chooses the smallest hidden dimension and can reduce the matrix norm, potentially lowering Hutchinson-estimator variance.The method exploits the cyclic property of trace when the dynamics network contains narrower hidden layers.
  • FFJORD: FFJORD provides a scalable reversible generative model with an unconstrained Jacobian by combining continuous dynamics with efficient stochastic log-density estimation.The complete method is specified through the dynamics and log-likelihood estimator, with pseudocode given in Algorithm 1.
  • Computational cost: FFJORD reduces likelihood computation to O((DH + D)L̂), compared with O((DH + D^3)L) for transformations using determinant identities.Here D is data dimensionality, H is the largest hidden dimension, L is the number of transformations, and L̂ is the number of ODE-function evaluations.

4 EXPERIMENTS Data

FFJORD is evaluated across toy, tabular, image, and variational-inference settings using GPU-based ODE solvers and trace estimation. The experiments show flexible density modeling, competitive image results, broad tabular performance, and strong variational-inference results, alongside slower computation.

  • Toy 2D data: FFJORD fits multi-modal and discontinuous 2D distributions by warping an isotropic Gaussian.The ODE solver uses roughly 70–100 evaluations across the datasets.
  • Real-data density estimation: FFJORD performs best among reversible models on five tabular datasets, while recent autoregressive models outperform it overall.It outperforms MAF on all but one dataset and surpasses TAN on MINIBOONE.
  • Image density estimation: On MNIST, a single FFJORD flow matches Glow and Real NVP, while multiscale flows improve MNIST and match Glow on CIFAR10.FFJORD achieves this image performance with less than 2% as many parameters as Glow.
  • Efficiency and implementation: The approach is generally slower than competing methods, but the adjoint method enables much larger batch sizes.Batch sizes reached 10,000 on tabular datasets and 900 on image datasets.
  • Variational inference: FFJORD outperforms all other competing normalizing flows on every tested variational-inference dataset.The comparison uses VAEs with no flow, Planar Flows, IAF, and Sylvester normalizing flows.

5 ANALYSIS AND DISCUSSION

The analyses examine training variance, computational scaling, and architecture trade-offs in FFJORD. Bottlenecks can accelerate training, while solver evaluations are approximately dimension-independent but architecture-dependent.

  • 5.1 FASTER TRAINING WITH BOTTLENECK TRICK: Bottleneck architectures reduce the variance of FFJORD’s log-density estimator and can speed up training.On MNIST, the bottleneck trick accelerated convergence with Gaussian noise but not with Rademacher noise.
  • 5.2 NUMBER OF FUNCTION EVALUATIONS VS. DATA DIMENSION: The full computational cost is O(DH bL), where D is data dimensionality, H is hidden-state size, and bL is the adaptive solver’s number of function evaluations.Each model evaluation costs O(DH), with H typically chosen close to D.
  • 5.2 NUMBER OF FUNCTION EVALUATIONS VS. DATA DIMENSION: NFE converges to approximately the same value across increasing latent dimensions, indicating dependence on distribution complexity rather than dimensionality.The Gaussian-to-Gaussian thought experiment has zero optimal dynamics and zero evaluations.
  • 5.3 MULTISCALE ARCHITECTURES: Multiscale FFJORD uses more function evaluations than single-scale FFJORD but achieves better performance on image data.The comparison uses MNIST models with comparable parameter counts and includes forward and backward passes.

6 SCOPE AND LIMITATIONS

FFJORD’s computational and solver behavior imposes practical limits. Function evaluations can become prohibitively large, and general-purpose solvers work efficiently only for non-stiff dynamics.

  • NUMBER OF FUNCTION EVALUATIONS: Function evaluations can grow during training and become prohibitively large even when adjoint memory usage remains constant.Weight decay and spectral normalization can reduce NFE but tend to hurt performance slightly.
  • LIMITATIONS OF GENERAL-PURPOSE ODE SOLVERS: General-purpose ODE solvers restrict FFJORD in practice to non-stiff differential equations that can be solved efficiently.Stiff solvers exist but require many more evaluations for the same error; small weight decay was sufficient to constrain the ODE to be non-stiff.

7 CONCLUSION

The paper concludes that FFJORD combines exact log-likelihoods, efficient sampling, and unrestricted neural networks through continuous-time dynamics. The authors identify reducing ODE-solver evaluations as key future work for scaling further.

  • 7 CONCLUSION: FFJORD computes exact log-likelihoods and supports efficient sampling with continuous-time dynamics parameterized by unrestricted neural networks.Training and sampling use automatic differentiation, Hutchinson’s trace estimator, and black-box ODE solvers.
  • 7 CONCLUSION: Reducing ODE-solver function evaluations without hurting predictive performance is identified as important for scaling FFJORD to higher-dimensional datasets.The authors describe this as a direction for further work.

APPENDIX A QUALITATIVE SAMPLES

Samples from FFJORD models trained on MNIST and CIFAR10 are presented in Figure 7.

  • Figure 7 contains samples from FFJORD models trained on MNIST and CIFAR10.

APPENDIX B EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS

The appendix describes architecture searches for tabular and image datasets, including encoder-decoder and multiscale flows. Training used Adam, with image models trained on six GPUs for approximately five days.

  • Tabular density estimation: Tabular experiments searched over 1–10 flows, 1–4 hidden layers, dimension-scaled hidden widths, and tanh or softplus nonlinearities.The best-performing configurations are reported in Table 4.
  • Image architectures: Image experiments compared a single encoder-decoder flow with a multiscale architecture composed of multiple flows.The encoder-decoder architecture was evaluated separately from the multiscale design.
  • Image architectures: Encoder-decoder models fit MNIST but could not fit CIFAR10 or Street View House Numbers.The MNIST configuration used mirrored convolutional and transpose-convolutional layers with softplus activations.
  • Image architectures: Multiscale models stack scale blocks containing flows and squeeze operations; MNIST used 3 blocks, while CIFAR10 used 4, with N = 2 in both.Squeeze operations down-sample spatial resolution and increase channel count.
  • Training: Image models were trained with Adam for 500 epochs, using a learning rate of .001 decayed to .0001 after 250 epochs.Training used six GPUs and completed after approximately five days.

B.2 VARIATIONAL AUTOENCODER

The variational-autoencoder experiments followed the procedure of Berg et al. while varying the FFJORD approximate posterior architecture. Results and selected architectures are summarized in Tables 4–6.

  • Experimental procedure: The experimental procedure matched Berg et al., including the 7-layer encoder-decoder, .001 learning rate, Adam optimizer, batch size 100, and early stopping.The stated difference was the normalizing flow used in the approximate posterior.
  • Architecture search: FFJORD dynamics were searched over 1–2 hidden layers, widths 512–2048, 1–5 flow steps, and rank-1, 20, or 64 weight-matrix updates.Softplus was used except for Caltech Silhouettes, which used tanh.
  • Reported configurations: Table 4 reports best-performing FFJORD architectures for tabular density estimation, while Table 5 reports architectures for FFJORD VAEs.The supplied table captions identify the respective experimental settings.
  • Datasets: The tabular datasets listed are POWER, GAS, HEPMASS, MINIBOONE, and BSDS300.These names identify the density-estimation benchmarks referenced in the appendix materials.
  • Reported results: Table 6 reports test-data negative log-likelihood means and standard deviations over 3 runs for density-estimation models.The table caption specifies the metric and aggregation across runs.

APPENDIX C NUMERICAL ERROR FROM THE ODE SOLVER

The appendix evaluates numerical integration error in the ODE solver and shows that tolerance controls density-estimation accuracy. It selects tolerances balancing reasonable performance with the number of function evaluations.

  • Solver error: ODE solver outputs contain numerical integration error, and adaptive solvers adjust step sizes to keep predicted error below a user-set tolerance.This error directly affects the density reported and compared with other methods.
  • Evaluation procedure: The density-integral test numerically estimates the area under a model density using a fine-grid Riemann sum across tolerance values.Both absolute and relative tolerances were set to the same value.
  • Evaluation results: Figure 8 shows that the modeled density integrates to one when the numerical tolerance is sufficiently low.The figure includes both log and non-log plots.
  • Tolerance selection: 10^-5 is the usual experimental tolerance because higher values produce non-negligible error, while lower values can encounter machine-precision errors.Tabular experiments instead used atol=10^-8 and rtol=10^-6.
Loading 1810.01367v3…