FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models (1810.01367v3)

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.

• The paper introduces FFJORD, which leverages continuous neural ODEs and Hutchinson’s trace estimator to reduce computational complexity from O(D³) to O(D).
• The method employs flexible, unrestricted architectures to enhance density estimation and variational inference across datasets such as MNIST and CIFAR10.
• Experimental results validate FFJORD’s scalability and efficiency, opening avenues for advanced applications in high-dimensional generative modeling.

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

In the paper "FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models," Grathwohl et al. introduce a novel approach to generative modeling that leverages continuous-time dynamics to create expressive and scalable reversible generative models.

The cornerstone of the method, FFJORD (Free-form Jacobian of Reversible Dynamics), lies in combining neural ODEs with Hutchinson's trace estimator to derive a continuous normalizing flow with unrestricted architectures. This promising approach addresses the computational constraints linked to traditional generative models that rely on Jacobian determinants, enabling a time complexity reduction from $\mathcal{O}(D^3)$ to $\mathcal{O}(D)$, where $D$ is the data dimensionality.

Background: Traditional and Continuous Generative Models

The authors begin by outlining the landscape of existing generative models, including:

• Normalizing Flows like NICE and Glow, which transform simple base distributions to target distributions using invertible neural networks, constrained by the need for efficient Jacobian determinant computations.
• Autoregressive Models, which are efficient for density estimation but have high invertibility costs.
• GANs and VAEs, which, while powerful, suffer from training inefficiencies and approximate log-likelihood calculations.

Building on continuous normalizing flows proposed by Chen et al. (2018), the authors extend this idea by employing a stochastic estimator, specifically Hutchinson's trace estimator. This innovation significantly reduces computational complexity, offering a scalable way to compute log-densities in continuous-time generative models.

Methodology: FFJORD Framework

The core of the methodology is the application of an unbiased stochastic estimator to the trace of the Jacobian, thus bypassing the expensive computation of exact log-densities. The authors employ an ODE defined by a parametric function $f(z(t), t; \theta)$ to map base distribution samples $z_0 \sim p_{z_0}(z_0)$ to data samples $x$.

Key components include:

1. Stochastic Trace Estimation:

$$\log p(z(t_1)) = \log p(z(t_0)) - E_{p(\epsilon)}\left[\int_{t_0}^{t_1} \epsilon^T \frac{\partial f}{\partial z(t)} dt \right]$$

By using fixed noise over the duration of the ODE solve, the trace estimator approximates the Jacobian efficiently, allowing for an $\mathcal{O}(D)$ time complexity.

1. Bottleneck Trick: For architectures containing bottleneck layers (hidden layers with dimensionality $H$ where $H < D$), the trace estimation can be further optimized by reshaping dimensions to reduce estimator variance.

Experimental Results

The authors validate FFJORD's efficacy through extensive experiments in density estimation and variational inference. Highlights include:

• Density Estimation: On both toy and real datasets (e.g., MNIST, CIFAR10), FFJORD consistently achieves superior or comparable log-likelihood performance relative to existing methods, including Glow and Real NVP. On tabular datasets, FFJORD surpasses most reversible models and presents competitive results against autoregressive models.
• Variational Inference: FFJORD enhances VAEs with normalizing flows, outperforming other flow-based methods such as Planar Flows and Inverse Autoregressive Flow (IAF) in terms of lower evidence lower bound (ELBO) across multiple datasets.

Notably, FFJORD's flexibility allows for deeper, more expressive architectures while maintaining computational feasibility and leveraging GPU-based adaptive ODE solvers for training efficiency.

Implications and Future Work

The implications of FFJORD are twofold:

1. Practical Benefits: The use of unrestricted neural architectures implies potential applications in various high-dimensional tasks, making FFJORD versatile for image and signal processing domains.
2. Theoretical Insights: The work pushes the boundaries of continuous normalizing flows, offering a framework that could inspire further refinements in trace estimation and adjoint sensitivity methods.

Future research directions include exploring methods to reduce the number of function evaluations during ODE integration, possibly through regularization techniques or more advanced numerical solvers. These advancements would be crucial for scaling FFJORD to handle even larger and more complex datasets effectively.

In summary, the FFJORD framework represents a significant step toward more expressive and computationally efficient generative models, unlocking new possibilities in the field of deep learning and beyond.