F LOW ++: I MPROVING FLOW - BASED GENERATIVE MODELS WITH VARIATIONAL DEQUANTIZATION AND ARCHITECTURE DESIGN
Anonymous authors
Paper under double-blind review
A BSTRACT
Flow-based generative models are powerful exact likelihood models with efﬁcient sampling and inference.
Despite their computational efﬁciency, ﬂow-based models generally have much worse density modeling performance compared to stateof-the-art autoregressive models.
In this paper, we investigate and improve upon three limiting design choices employed by ﬂow-based models in prior work: the use of uniform noise for dequantization, the use of inexpressive afﬁne ﬂows, and the use of purely convolutional conditioning networks in coupling layers.
Based on our ﬁndings, we propose Flow++, a new ﬂow-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 signiﬁcant performance gap that has so far existed between autoregressive models and ﬂow-based models.
I NTRODUCTION
Deep generative models – latent variable models in the form of variational autoencoders (Kingma &
Welling, 2013), implicit generative models in the form of GANs (Goodfellow et al., 2014), and exact likelihood models like PixelRNN/CNN (van den Oord et al., 2016b;c), Image Transformer (Parmar et al., 2018), PixelSNAIL (Chen et al., 2017), NICE, RealNVP, and Glow (Dinh et al., 2014; 2016;
Kingma & Dhariwal, 2018) – have recently begun to successfully model high dimensional raw observations from complex real-world datasets, from natural images and videos, to audio signals and natural language (Karras et al., 2017; Kalchbrenner et al., 2016b; van den Oord et al., 2016a;
Kalchbrenner et al., 2016a; Vaswani et al., 2017).
Autoregressive models, a certain subclass of exact likelihood models, achieve state-of-the-art density estimation performance on many challenging real-world datasets, but generally suffer from slow sampling time due to their autoregressive structure (van den Oord et al., 2016b; Salimans et al., 2017; Chen et al., 2017; Parmar et al., 2018).
Inverse autoregressive models can sample quickly and potentially have strong modeling capacity, but they cannot be trained efﬁciently by maximum likelihood (Kingma et al., 2016).
Non-autoregressive ﬂow-based models (which we will refer to as “ﬂow models”), such as NICE, RealNVP, and Glow, are efﬁcient for sampling, but have so far lagged behind autoregressive models in density estimation benchmarks (Dinh et al., 2014; 2016;
Kingma & Dhariwal, 2018).
In the hope of creating an ideal likelihood-based generative model that simultaneously has fast sampling, fast inference, and strong density estimation performance, we seek to close the density estimation performance gap between ﬂow models and autoregressive models.
In subsequent sections, we present our new ﬂow model, Flow++, which is powered by an improved training procedure for continuous likelihood models and a number of architectural extensions of the coupling layer deﬁned by Dinh et al.
F LOW M ODELS
A ﬂow model f is constructed as an invertible transformation that maps observed data x to a standard Gaussian latent variable z = f ( x ) , as in nonlinear independent component analysis (Bell &
Sejnowski, 1995; Hyv¨arinen et al., 2004; Hyv¨arinen & Pajunen, 1999).
The key idea in the design of a ﬂow model is to form f by stacking individual simple invertible transformations (Dinh et al., 2014; 2016; Kingma & Dhariwal, 2018; Rezende & Mohamed, 2015; Kingma et al., 2016;
Louizos & Welling, 2017).
Explicitly, f is constructed by composing a series of invertible ﬂows as f ( x ) =
f 1 ◦· · ·◦ f L ( x ) , with each f i having a tractable inverse and a tractable Jacobian determinant.
This way, sampling is efﬁcient, as it can be performed by computing f − 1 ( z ) =
f − 1 L ◦ · · · ◦ f − 1 1 ( z ) for z ∼ N ( 0 , I ) , and so is training by maximum likelihood, since the model density
[Equation]
i
[Equation] [Equation]
[Equation]
is easy to compute and differentiate with respect to the parameters of the ﬂows f i .
F LOW ++
In this section, we describe three modeling inefﬁciencies in prior work on ﬂow models: (1) uniform noise is a suboptimal dequantization choice that hurts both training loss and generalization; (2) commonly used afﬁne coupling ﬂows are not expressive enough; (3) convolutional layers in the conditioning networks of coupling layers are not powerful enough.
Our proposed model, Flow++, consists of a set of improved design choices: (1) variational ﬂow-based dequantization instead of uniform dequantization; (2) logistic mixture CDF coupling ﬂows; (3) self-attention in the conditioning networks of coupling layers.
3. D EQUANTIZATION VIA VARIATIONAL INFERENCE
Many real-world datasets, such as CIFAR10 and ImageNet, are recordings of continuous signals quantized into discrete representations.
Fitting a continuous density model to discrete data, however, will produce a degenerate solution that places all probability mass on discrete datapoints (Uria et al., 2013).
A common solution to this problem is to ﬁrst convert the discrete data distribution into a continuous distribution via a process called “dequantization,” and then model the resulting continuous distribution using the continuous density model (Uria et al., 2013;
Dinh et al., 2016; Salimans et al., 2017).
3.1. U NIFORM DEQUANTIZATION
Dequantization is usually performed in prior work by adding uniform noise to the discrete data over the width of each discrete bin: if each of the D components of the discrete data x takes on values in { 0 , 1 , 2 , . .
. , 255 } , then the dequantized data is given by y = x + u , where u is drawn uniformly from [0 , 1) D .
Theis et al.
(2015) note that training a continuous density model p model on uniformly dequantized data y can be interpreted as maximizing a lower bound on the log-likelihood for a certain discrete model P model on the original discrete data x :
[Equation]
The argument of Theis et al.
(2015) proceeds as follows.
Letting P data denote the original distribution of discrete data and p data denote the distribution of uniformly dequantized data, Jensen’s inequality implies that
[Equation]
Consequently, maximizing the log-likelihood of the continuous model on uniformly dequantized data cannot lead to the continuous model degenerately collapsing onto the discrete data, because its objective is bounded above by the log-likelihood of a discrete model.
3.1. V ARIATIONAL DEQUANTIZATION
While uniform dequantization successfully prevents the continuous density model p model from collapsing to a degenerate mixture of point masses on discrete data, it asks p model to assign uniform density to unit hypercubes x + [0 , 1) D around the data x .
It is difﬁcult and unnatural for smooth function approximators, such as neural network density models, to excel at such a task.
To sidestep this issue, we now introduce a new dequantization technique based on variational inference.
Again, we are interested in modeling D -dimensional discrete data x ∼ P data using a continuous density model p model , and we will do so by maximizing the log-likelihood of its associated discrete model P model ( x ) :
= (cid:82) [0 , 1) D p model ( x + u ) d u .
Now, however, we introduce a dequantization noise distribution q ( u | x ) , with support over u ∈ [0 , 1) D .
Treating q as an approximate posterior, we have the following variational lower bound, which holds for all q :
[Equation]
We will choose q itself to be a conditional ﬂow-based generative model of the form u = q x ( (cid:15) ) , where (cid:15) ∼ p
(cid:15) )
=
N ( (cid:15) ; 0 , I ) is Gaussian noise.
In this case, q ( u | x ) = p ( q − 1 x ( u )) · (cid:12)(cid:12)
∂q − 1 x /∂ u (cid:12)(cid:12) , and thus we obtain the objective
[Equation]
which we maximize jointly over p model and q .
When p model is also a ﬂow model x = f − 1 ( z ) (as it is throughout this paper), it is straightforward to calculate a stochastic gradient of this objective using the pathwise derivative estimator, as f ( x + q x ( (cid:15) )) is differentiable with respect to the parameters of f and q .
Notice that the lower bound for uniform dequantization – eqs. (3) to (5) – is a special case of our variational lower bound – eqs.
(6) to (8), when the dequantization distribution q is a uniform distribution that ignores dependence on x .
Because the gap between our objective (8) and the true expected log-likelihood E x ∼ P data [log P model ( x )] is exactly E x ∼ P data [ D KL ( q ( u | x ) (cid:107) p model ( u | x ))] , using a uniform q forces p model to unnaturally place uniform density over each hypercube x + [0 , 1) D to compensate for any potential looseness in the variational bound introduced by the inexpressive q .
Using an expressive ﬂow-based q , on the other hand, allows p model to place density in each hypercube x
+ [0 , 1) D according to a much more ﬂexible distribution q ( u | x ) .
This is a more natural task for p model to perform, improving both training and generalization loss.
3. I MPROVED COUPLING LAYERS
Recent progress in the design of ﬂow models has involved carefully constructing ﬂows to increase their expressiveness while preserving tractability of the inverse and Jacobian determinant computations.
One example is the invertible 1 × 1 convolution ﬂow, whose inverse and Jacobian determinant can be calculated and differentiated with standard automatic differentiation libraries (Kingma &
Dhariwal, 2018).
Another example, which we build upon in our work here, is the afﬁne coupling layer (Dinh et al., 2016).
It is a parameterized ﬂow y =
f θ ( x ) that ﬁrst splits the components of x into two parts x 1 , x 2 , and then computes y
= ( y 1 , y 2 ) , given by
[Equation]
Here, a θ and b θ are outputs of a neural network that acts on x 1 in a complex, expressive manner, but the resulting behavior on x 2 always remains an elementwise afﬁne transformation – effectively, a θ and b θ together form a data-parameterized family of invertible afﬁne transformations.
This allows the afﬁne coupling layer to express complex dependencies on the data while keeping inversion and log-likelihood computation tractable.
Using · and exp to respectively denote elementwise multiplication and exponentiation,
[Equation]
The splitting operation x (cid:55)→ ( x 1 , x 2 ) and merging operation ( y 1 , y 2 ) (cid:55)→ y are usually performed over channels or over space in a checkerboard-like pattern (Dinh et al., 2016).
3.2. E XPRESSIVE COUPLING TRANSFORMATIONS WITH CONTINUOUS MIXTURE CDF S
We found in our experiments that density modeling performance of these coupling layers could be improved by augmenting the data-parameterized elementwise afﬁne transformations by more general nonlinear elementwise transformations.
For a given scalar component x of x 2 , we apply the cumulative distribution function (CDF) for a mixture of K logistics – parameterized by mixture probabilities, means, and log scales π , µ , s – followed by an inverse sigmoid and an afﬁne transformation parameterized by a and b :
[Equation]
12) where MixLogCDF( x ; π , µ , s ) :=
K
[Equation]
(( x − µ
[Equation]
· exp(
[Equation]
The transformation parameters π , µ , s , a, b for each component of x 2 are produced by a neural network acting on x 1 .
This neural network must produce these transformation parameters for each component of x 2 , hence it produces vectors a θ ( x 1 ) and b θ ( x 1 ) and tensors π θ ( x 1 ) , µ θ ( x 1 ) , s θ ( x 1 ) (with last axis dimension K ).
The coupling transformation is then given by:
[Equation]
where the formula for computing y 2 operates elementwise.
The inverse sigmoid ensures that the inverse of this coupling transformation always exists: the range of the logistic mixture CDF is (0 , 1) , so the domain of its inverse must stay within this interval.
The
CDF itself can be inverted efﬁciently with bisection, because it is a monotonically increasing function.
Moreover, the Jacobian determinant of this transformation involves calculating the probability density function of the logistic mixtures, which poses no computational difﬁculty.
3.2. E XPRESSIVE CONDITIONING ARCHITECTURES WITH SELF - ATTENTION
In addition to improving the expressiveness of the elementwise transformations on x 2 , we found it crucial to improve the expressiveness of the conditioning on x 1 – that is, the expressiveness of the neural network responsible for producing the elementwise transformation parameters π , µ , s , a , b .
Our best results were obtained by stacking convolutions and multi-head self attention into a gated residual network (Mishra et al., 2018;
Chen et al., 2017), in a manner resembling the Transformer (Vaswani et al., 2017) with pointwise feedforward layers replaced by 3 × 3 convolutional layers.
Our architecture is deﬁned as a stack of blocks.
Each block consists of the following two layers connected in a residual fashion, with layer normalization (Ba et al., 2016) after each residual connection:
[Equation]
where Gate refers to a 1 × 1 convolution that doubles the number of channels, followed by a gated linear unit (Dauphin et al., 2016).
The convolutional layer is identical to the one used by
PixelCNN++ (Salimans et al., 2017), and the multi-head self attention mechanism we use is identical to the one in the Transformer (Vaswani et al., 2017).
(We always use 4 heads in our experiments, since we found it to be effective early on in our experimentation process.)
With these blocks in hand, the network that outputs the elementwise transformation parameters is simply given by stacking blocks on top of each other, and ﬁnishing with a ﬁnal convolution that increases the number of channels to the amount needed to specify the elementwise transformation parameters.
E XPERIMENTS
Here, we show that Flow++ achieves state-of-the-art density modeling performance among nonautoregressive models on CIFAR10 and 32x32 and 64x64 ImageNet.
We also present ablation experiments that quantify the improvements proposed in section 3, and we present example generative samples from Flow++ and compare them against samples from autoregressive models.
Our experiments employed weight normalization and data-dependent initialization (Salimans &
Kingma, 2016).
We used the checkerboard-splitting, channel-splitting, and downsampling ﬂows of Dinh et al.
(2016); we also used before every coupling ﬂow an invertible 1x1 convolution ﬂows of Kingma & Dhariwal (2018), as well as a variant of their “actnorm” ﬂow that normalizes all activations independently (instead of normalizing per channel).
Our CIFAR10 model used 4 coupling layers with checkerboard splits at 32x32 resolution, 2 coupling layers with channel splits at 16x resolution, and 3 coupling layers with checkerboard splits at 16x16 resolution; each coupling layer used 10 convolution-attention blocks, all with 96 ﬁlters.
More details on architectures, as well as details for the other experiments, will be given in a source code release.
4. D ENSITY MODELING RESULTS
In table 1, we show that Flow++ achieves state-of-the-art density modeling results out of all nonautoregressive models, and it is competitive with autoregressive models: its performance is on par with the ﬁrst generation of PixelCNN models (van den Oord et al., 2016b), and it outperforms
Multiscale PixelCNN (Reed et al., 2017).
As of submission, our models have not fully converged due to computational constraint
and we expect further performance gain in future revision of this manuscript.
[Table]
4. A BLATIONS
We ran the following ablations of our model on unconditional CIFAR10 density estimation: variational dequantization vs. uniform dequantization; logistic mixture coupling vs. afﬁne coupling; and stacked self-attention vs. convolutions only.
As each ablation involves removing some component of the network, we increased the number of ﬁlters in all convolutional layers (and attention layers, if present) in order to match the total number of parameters with the full Flow++ model.
In ﬁg. 1 and table
[Table]
we compare the performance of these ablations relative to Flow++ at 400 epochs of training, which was not enough for these models to converge
[Table]
but far enough to see their relative performance differences.
Switching from our variational dequantization to the more standard uniform dequantization costs the most: approximately 0 .
127 bits/dim.
The remaining two ablations both cost approximately 0 .
03 bits/dim: switching from our logistic mixture coupling layers to afﬁne coupling layers, and switching from our hybrid convolution-and-self-attention architecture to a pure convolutional residual architecture.
Note that these performance differences are present despite all networks having approximately the same number of parameters: the improved performance of
Flow++ comes from improved inductive biases, not simply from increased parameter count.
The most interesting result is probably the effect of the dequantization scheme on training and generalization loss.
At 400 epochs of training, the full Flow++ model with variational dequantization has a train-test gap of approximately 0.02 bits/dim, but with uniform dequantization, the train-test gap is approximately 0.06 bits/dim.
This conﬁrms our claim in Section 3.1.2 that training with variational dequantization is a more natural task for the model than training with uniform dequantization.
[Figure]
[Table]
[Figure]
[Table]
[Table]
S AMPLES
We present the samples from our trained density models of Flow++ on CIFAR10, 32x32 ImageNet, 64x64 ImageNet, and 5-bit CelebA in ﬁgs. 2 to 5.
The Flow++ samples match the perceptual quality of PixelCNN samples, showing that Flow++ captures both local and global dependencies as well as PixelCNN and is capable of generating diverse samples on large datasets.
Moreover, sampling is fast: our CIFAR10 model takes approximately 0.32 seconds to generate a batch of 8 samples in parallel on one NVIDIA
Ti GPU, making it more than an order of magnitude faster than
PixelCNN++ with sampling speed optimizations (Ramachandran et al., 2017). More samples are available in the appendix (section 7).
R ELATED W ORK
Likelihood-based models constitute a large family of deep generative models.
One subclass of such methods, based on variational inference, allows for efﬁcient approximate inference and sampling, but does not admit exact log likelihood computation (Kingma & Welling, 2013; Rezende et al., 2014; Kingma et al., 2016).
Another subclass, which we called exact likelihood models in this work, does admit exact log likelihood computation.
These exact likelihood models are typically [Figure] [Figure] speciﬁed as invertible transformations that are parameterized by neural networks (Deco & Brauer, 1995; Larochelle & Murray, 2011; Uria et al., 2013; Dinh et al., 2014; Germain et al., 2015; van den
Oord et al., 2016b; Salimans et al., 2017; Chen et al., 2017).
There is prior work that aims to improve the sampling speed of deep autoregressive models.
The
Multiscale PixelCNN (Reed et al., 2017) modiﬁes the PixelCNN to be non-fully-expressive by introducing conditional independence assumptions among pixels in a way that permits sampling in a logarithmic number of steps, rather than linear.
Such a change in the autoregressive structure allows for faster sampling but also makes some statistical patterns impossible to capture, and hence reduces the capacity of the model for density estimation.
WaveRNN (Kalchbrenner et al., 2018) [Figure] improves sampling speed for autoregressive models for audio via sparsity and other engineering considerations, some of which may apply to ﬂow models as well.
There is also recent work that aims to improve the expressiveness of coupling layers in ﬂow models.
Kingma & Dhariwal (2018) demonstrate improved density estimation using an invertible 1x1 convolution ﬂow, and demonstrate that very large ﬂow models can be trained to produce photorealistic faces.
M¨uller et al.
(2018) introduce piecewise polynomial couplings that are similar in spirit to our mixture of logistics couplings.
They found them to be more expressive than afﬁne couplings, but reported little performance gains in density estimation.
We leave a detailed comparison between our coupling layer and the piecewise polynomial CDFs for future work.
C ONCLUSION
We presented Flow++, a new ﬂow-based generative model that begins to close the performance gap between ﬂow models and autoregressive models.
Our work considers speciﬁc instantiations of design principles for ﬂow models – dequantization, ﬂow design, and conditioning architecture design – and we hope these principles will help guide future research in ﬂow models and likelihoodbased models in general.
A PPENDIX A: S AMPLES