[Paper Review] FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models

Will Grathwohl, Ricky T. Q. Chen+3 more|arXiv (Cornell University)|Oct 2, 2018

Cellular Automata and Applications320 citations

TL;DR

FFJORD introduces a continuous-time, invertible generative model with unbiased log-density estimation using Hutchinson’s trace estimator, enabling unrestricted neural architectures and scalable sampling.

ABSTRACT

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.

Motivation & Objective

• Motivated by the need for exact likelihood in reversible generative models without restrictive architectural constraints.
• Develop a scalable log-density estimator that works with unrestricted neural networks.
• Leverage continuous-time dynamics to replace discrete layers, enabling efficient sampling and density evaluation.
• Demonstrate effectiveness on density estimation, image generation, and variational inference.

Proposed method

• Model data via a base distribution transformed by a continuous-time invertible dynamics z'(t)=f(z(t),t;θ).
• Use instantaneous change of variables: log p(z(t1))=log p(z(t0))−∫ Tr(∂f/∂z) dt.
• Estimate Tr(∂f/∂z) unbiasedly with Hutchinson’s trace estimator: Tr(A)=E[ε^T A ε].
• Compute log-likelihood via augmented ODE solving the state and log-density jointly with an adjoint method for backpropagation.
• Achieve O(D) trace estimation cost and allow unrestricted architectures; solve ODEs with GPU-accelerated solvers.

Experimental results

Research questions

• RQ1 Can continuous-time reversible dynamics provide exact log-likelihoods with unrestricted neural architectures?
• RQ2 Does Hutchinson’s trace estimator enable unbiased log-density estimation at linear time in data dimension?
• RQ3 How does FFJORD perform on high-dimensional density estimation and image generation compared to restricted normalizing flows?
• RQ4 What are the trade-offs in training speed, memory, and function evaluations when using FFJORD versus discrete-flow models?

Key findings

• FFJORD achieves unbiased log-density estimation with O(D) cost, enabling unrestricted architectures.
• On density estimation, FFJORD outperforms other reversible flows and matches or surpasses certain autoregressive methods on tabular data, while using fewer parameters on some image tasks.
• FFJORD can model disconnected and multi-modal densities that some discrete-flow models struggle with.
• Bottleneck architectures can reduce variance of the trace estimator and speed up training under certain trace estimation schemes.
• The number of ODE solver evaluations is not fixed and tends to grow with training, but is largely independent of data dimensionality D.