Source-linked AI summary
Flow++: Improving Flow-Based Generative Models with Variational Dequantization and Architecture Design
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, Pieter Abbeel
TL;DR
Flow models provide efficient sampling and inference but have generally underperformed autoregressive models in density estimation. Flow++ addresses this gap with variational dequantization, more expressive coupling transformations, and self-attention conditioning, achieving state-of-the-art non-autoregressive density modeling and competitive performance with autoregressive models.
Problem
Flow-based models offer computational efficiency but generally lag autoregressive models in density modeling performance.
Method
Flow++ combines variational flow-based dequantization, logistic mixture CDF coupling flows, and self-attention in coupling-layer conditioning networks.
Results
Flow++ achieves state-of-the-art density modeling among non-autoregressive models and is competitive with autoregressive models.
Takeaways & Limitations
The study begins to close the performance gap between flow models and autoregressive models through improved dequantization, flow design, and conditioning architecture design.
Abstract
from arXiv · showhide
Flow-based generative models are powerful exact likelihood models with efficient sampling and inference. Despite their computational efficiency, flow-based models generally have much worse density modeling performance compared to state-of-the-art autoregressive models. In this paper, we investigate and improve upon three limiting design choices employed by flow-based models in prior work: the use of uniform noise for dequantization, the use of inexpressive affine flows, and the use of purely convolutional conditioning networks in coupling layers. Based on our findings, we propose Flow++, a new flow-based model that is now the state-of-the-art non-autoregressive model for unconditional density estimation on standard image benchmarks. Our work has begun to close the significant performance gap that has so far existed between autoregressive models and flow-based models. Our implementation is available at https://github.com/aravindsrinivas/flowpp
1. Introduction
Deep generative models have been applied to complex, high-dimensional real-world data, while flow models offer efficient sampling but lag autoregressive models in density estimation. Flow++ is introduced to narrow this gap while retaining fast sampling and inference.
- Deep generative models include variational autoencoders, GANs, exact likelihood models, and flow models for complex real-world observations.Applications include natural images, videos, audio signals, and natural language.
- Autoregressive models achieve state-of-the-art density estimation but generally sample slowly because of their autoregressive structure.
- Flow models sample efficiently but have lagged behind autoregressive models on density estimation benchmarks.
- Flow++ targets an ideal likelihood-based model combining fast sampling, fast inference, and strong density estimation performance.The paper presents Flow++ as a new flow model powered by improved training and coupling-layer architectural extensions.
2. Flow Models
A flow model is an invertible transformation that maps observed data to a standard Gaussian latent variable. By composing simple invertible transformations, it supports efficient likelihood-based training and inference.
- A flow model maps observed data x to a standard Gaussian latent variable z through an invertible transformation f.
- The transformation f is constructed by composing a series of individual invertible flows f_i.
- Because the model density is easy to compute and differentiate, flow models can be trained by maximum likelihood.
3. Flow++
Flow++ addresses three inefficiencies in prior flow models—uniform dequantization, affine coupling transformations, and convolutional conditioning—through variational dequantization, logistic mixture CDF couplings, and self-attention. These changes preserve tractable flow computation while enabling more flexible density modeling.
- Variational dequantization: Flow++ replaces uniform dequantization with a conditional flow-based distribution q(u|x), optimized jointly with the continuous density model.The variational objective can be optimized using pathwise derivatives when the model and dequantization process are flows.
- Variational dequantization: Uniform dequantization forces the model toward uniform density within each discrete-data hypercube, whereas expressive q(u|x) permits more flexible density placement.The variational-bound gap is the expected KL divergence between q(u|x) and pmodel(u|x).
- Improved coupling layers: Prior affine coupling layers keep inversion and likelihood computation tractable but transform the conditioned variables only through data-parameterized elementwise affine maps.Flow++ augments this design with more general nonlinear elementwise transformations.
- Improved coupling layers: Flow++ replaces affine coupling transformations with nonlinear elementwise transformations based on logistic mixture CDFs, inverse sigmoid, and affine parameters.The mixture is parameterized by probabilities, means, and log scales, while the conditioning network produces the transformation parameters.
- Improved coupling layers: The logistic mixture CDF coupling remains tractable because its inverse can be computed by bisection and its Jacobian uses the logistic-mixture density.Monotonicity of the CDF ensures invertibility after the inverse sigmoid.
- Expressive conditioning architectures with self-attention: Flow++ strengthens coupling-layer conditioning by combining convolutions and multi-head self-attention in a gated residual network with residual connections and layer normalization.The architecture resembles a Transformer while replacing pointwise feedforward layers with 3 × 3 convolutions; experiments use four attention heads.
4. Experiments
Flow++ is evaluated on unconditional image density estimation, with comparisons against prior models, ablations of its design choices, and sample-quality and sampling-speed assessments.
- Density modeling results: Flow++ achieves state-of-the-art density modeling among non-autoregressive models on CIFAR10 and 32x32 and 64x64 ImageNet.
- Density modeling results: Flow++ is on par with first-generation PixelCNN models and outperforms Multiscale PixelCNN in density modeling.
- Ablations: 0.127 bits/dim is the approximate cost of replacing variational dequantization with uniform dequantization, compared with approximately 0.03 bits/dim for each other ablation.The other ablations replace logistic mixture coupling with affine coupling or hybrid convolution-and-self-attention with pure convolutional residual conditioning.
- Ablations: The ablation differences occur despite approximately equal parameter counts, attributing Flow++'s improved performance to inductive biases rather than parameter count.
- Ablations: 0.02 bits/dim versus 0.06 bits/dim is the train-test gap with variational versus uniform dequantization after 400 epochs.
- Samples: 0.32 seconds is the time for the CIFAR10 model to generate a batch of 8 samples in parallel on one NVIDIA 1080 Ti GPU.This is reported as more than an order of magnitude faster than PixelCNN++ with sampling speed optimizations.
- Samples: Flow++ samples match the perceptual quality of PixelCNN samples and capture local and global dependencies while generating diverse samples.
5. Related Work
Related work spans likelihood-based models, faster autoregressive variants, and more expressive flow coupling layers. Flow++ is evaluated against these approaches through sample quality, diversity, and design comparisons.
- Variational inference models offer efficient approximate inference and sampling but do not provide exact log-likelihood computation.
- Multiscale PixelCNN accelerates sampling by introducing conditional independence assumptions, but this reduces density-estimation capacity.
- Flow++ samples capture local dependencies and reach PixelCNN-level quality while retaining efficient sampling.
- On 32x32 ImageNet, Flow++ sample diversity matches that of an autoregressive model.
- Invertible 1x1 convolutions, neural-network elementwise transformations, and piecewise polynomial couplings are prior efforts to increase coupling-layer expressiveness.
6. Conclusion
Flow++ begins to close the performance gap between flow-based and autoregressive models by applying specific design principles to dequantization, flow construction, and conditioning architecture.
- Flow++ begins to close the performance gap between flow models and autoregressive models.
- The paper focuses on specific instantiations of dequantization, flow design, and conditioning-architecture principles.
Dequantization and Architecture Design – Supplementary Material
Supplementary figures provide additional Flow++ samples across CIFAR10, ImageNet resolutions and bit depths, and CelebA bit depths. The implementation is publicly available.
- Supplementary samples include Flow++ models trained on CIFAR10, 32x32 ImageNet, and 64x64 ImageNet.
- Additional samples cover 5-bit 64x64 ImageNet and 3-bit and 5-bit 64x64 CelebA HQ.
- The Flow++ implementation is available at the cited GitHub repository.