Discrete Denoising Flows
Abstract
Discrete ﬂow-based models are a recently proposed class of generative models that learn invertible transformations for discrete random variables.
Since they do not require data dequantization and maximize an exact likelihood objective, they can be used in a straight-forward manner for lossless compression.
In this paper, we introduce a new discrete ﬂow-based model for categorical random variables: Discrete Denoising Flows (DDFs).
In contrast with other discrete ﬂow-based models, our model can be locally trained without introducing gradient bias.
We show that DDFs outperform
Discrete Flows on modeling a toy example, binary MNIST and Cityscapes segmentation maps, measured in log-likelihood.
1. Introduction
Due to their wide range of applications, ﬂow-based generative models have been extensively studied in recent years (Rezende & Mohamed, 2015; Dinh et al., 2016).
Research has mainly focused on modeling continuous data distributions, in which discretely stored data like audio or image data must be dequantized prior to modeling.
However, two recent publications explore ﬂow-based generative modeling of discrete distributions: Discrete Flows (Tran et al., 2019) for categorical random variables and Integer Discrete Flows (Hoogeboom et al., 2019) for ordinal discrete random variables.
Due to their discrete nature and exact likelihood objective, these discrete ﬂow-based models can be used directly for lossless compression.
Unlike other approaches that use generative models for lossless compression, discrete ﬂow-based models are advantageous because they (i) enable efﬁcient inference and (ii) can encode single data samples efﬁciently. Approaches that use
(ICML 2021). the Variational Autoencoder (VAE) (Kingma & Welling, 2013) for lossless compression typically combine the model with bits-back-coding (Hinton & Van Camp, 1993), which is effective for compressing full data sets but inefﬁcient for encoding single samples.
Autoregressive models such as PixelCNN (Oord et al., 2016) can also be used for lossless compression, however, they are generally expensive to decode.
Unfortunately, both Discrete Flows and Integer Discrete
Flows come with the drawback that each of their layers contains a quantization operation.
When optimizing them with the backpropagation algorithm, the gradient of the quantization operation has to be estimated with a biased gradient estimator, which may compromise their performance.
To improve training efﬁciency, reduce gradient bias and improve overall performance, we introduce a new discrete ﬂow-based generative model for categorical random variables, Discrete Denoising Flows (DDFs).
DDFs can be trained without introducing gradient bias.
They further come with the positive side effect that the training is computationally very efﬁcient.
This efﬁciency results from the local training algorithm of the DDFs, which trains only one layer at a time instead of all at once.
We demonstrate that
Discrete Denoising Flows outperform Discrete Flows in terms of log-likelihood.
2. Related Work & Background
This section ﬁrst introduces normalizing ﬂows as well as discrete ﬂows.
It then goes on to describe alternate approaches that use generative models for lossless compression.
Normalizing Flows The fundamental idea of ﬂow-based modeling is to express a complicated probability distribution as a transformation on a simple probability distribution.
Given the two continuous random variables X and Z and the invertible and differentiable transformation T : Z → X ,
X ’s probability distribution p X ( · ) can be written in terms of Z ’s probability distribution p Z ( · ) as
[Equation]
with z = T − 1 (
[Equation]
using the change of variables formula.
The Jacobian determinant acts as normalizing term and ensures that p X ( · ) is a valid probability distribution.
The distribution p Z ( · ) is referred to as the base distribution and the transformation
T as a normalizing ﬂow .
A composition of invertible and differentiable functions can be viewed as a repeated application of formula 1.
Therefore, such compositions are also referred to as normalizing ﬂows throughout literature.
Discrete Flows In the case of two discrete random variables X
and Z , the change of variables formula for continuous random variables given in formula 1 simpliﬁes to
[Equation]
Normalization with the Jacobian determinant is no longer necessary as it corrects for a change in volume.
Discrete distributions, however, have no volume since they only have support on a discrete set of points.
Discrete Flows (DFs) (Tran et al., 2019) are discrete ﬂowbased models that learn transformations on categorical random variables.
The authors deﬁne a bijective transformation
T : Z → X with X =
Z = { 1 , . .
. , K } D in the form of a bipartite coupling layer (Dinh et al., 2016).
The coupling layer input x is partitioned into two sets s.t. x
= [ x a , x b ] and then transformed into an output z
=
[ z a , z b ] with
[Equation]
where ◦ denotes element-wise multiplication.
Note that the transformation is only invertible if each element of the scale is coprime to the number of classes K .
Scale s θ 1 ( · ) and translation t
θ 2 ( · ) are modeled by a neural network with parameters θ 1 , 2 .
To obtain discrete scale and translation values, the authors use the argmax operator combined with a relaxed softmax as gradient estimator (Jang et al., 2016) to enable backpropagation.
This introduces bias to the model parameter gradients, which harms optimization.
Note that the example describes the bipartite version and not the autoregressive version of DFs.
In (Tomczak, 2020) different partitions for coupling layers are explored.
Generative Models for Lossless Compression
The Variational Autoencoder (VAE) (Kingma & Welling, 2013) can be used for lossless compression is by discretizing the continuous latent vector and applying bits-back coding (Hinton & Van Camp, 1993).
Recent methods that work according to this approach include Bits-Back with ANS (Townsend et al., 2019a), Bit-Swap (Kingma et al., 2019) and HiLLoC (Townsend et al., 2019b).
These methods obtain performances close to the negative ELBO for compressing full datasets.
However, when encoding a single data sample they are rather inefﬁcient because the auxiliary bits needed for bits-back coding cannot be amortized across many samples.
The same problem but in a scaled-up version due to multiple latent variables arises when local bits-back coding is used in normalizing ﬂows (Ho et al., 2019).
In this case, encoding a single image would require more bits than the original image.
Mentzer et al.
(Mentzer et al., 2019) use a VAE with deterministic encoder to transform a data sample into a set of discrete multiscale latent vectors.
Although this method does not require bits-back coding, it optimizes only a lowerbound on the likelihood instead of the likelihood directly.
Another generative model that is well suited for lossless compression is the PixelCNN (Oord et al., 2016).
PixelCNN organizes the pixels of an image as a sequence and predicts the distribution of a pixel conditioned on all previous pixels.
Consequently, drawing samples from PixelCNN requires multiple network evaluations and is very costly.
Nevertheless, PixelCNN achieves state-of-the-art performances in lossless compression.
3. Method: Discrete Denoising Flows
In this section, we introduce Discrete Denoising Flows .
Like other ﬂow-based models, DDFs consist of several bipartite coupling layers that are easily invertible, the so-called denoising coupling layers .
3.1. Denoising
Coupling Layer
Complying to the change of variables formula 2, we deﬁne the denoising coupling layer as an invertible transformation
T : Z → X between two categorical variables X
and Z with domains
X = Z = { 1 , . .
. , K } D .
The inverse T − 1 , representing the forward pass during training, is given as
[Equation]
That is, the input x ∈ { 1 , .., K } D is partitioned into two sets such that x = [ x a , x b ] and x a ∈ { 1 , .., K } d .
The ﬁrst part stays the same while the second part is transformed conditioned on the ﬁrst part.
For this transformation, we use a neural network n as well as the conditional permutation operation cond perm( ·|· ) .
The conditional permutation operation is the core component of the denoising coupling layer.
For notation clarity, we deﬁne the variable θ = n ( x a ) with θ ∈ R ( D − d ) × K as the output of the neural network n .
The conditional permutation operation is then deﬁned as
[Equation]
per dimension i ∈ { 1 , . . .
, D − d } , where we used Cauchy’s two-line notation for permutations to deﬁne the permutation perm θ i .
We also introduced the argsort top h ( · ) operation with the additional hyperparameter h ∈ { 1 , .., K } .
This operation acts similar to the regular argsort operation, with the difference that it only sorts the top h largest elements while leaving the remaining elements in their previous positions.
Figure 1 illustrates this functionality.
The intuition behind the operation is that only the most predictable classes are permuted, leaving more of the structure intact than an entire argsort .
Also, observe that in the binary case K =and for h = K , argsort and argsort top h are equivalent.
The conditional permutation operation is easily invertible as
[Equation]
per dimension i ∈ { 1 , . . .
, D − d } .
Using this deﬁnition, we can write the transformation T representing the denoising coupling layer at inference time as
[Equation]
3.2. Training Denoising Discrete Flows
For training a denoising coupling layer, we simply train a neural network n to predict x b from x a .
To this end, we use the mean cross-entropy loss between n ( x a ) and x b as our objective function.
After training, the ﬁxed neural network n can be employed in a denoising coupling layer.
When we apply the conditional permutation operation in [Equation] ) , the more the argmax of n ( x a ) resembles x b , the more likely it is for a value in x b to be switched to one of the smaller class values.
Consequently, given that the argmax of n ( x a ) somewhat resembles x b , the outcome of the conditional permutation operation z b , is more likely to contain smaller class values than x b .
This makes the value of the random variable
Z more predictable than the value of the random variable
X , when looking at those dimensions in isolation.
In other words, we decorrelated the random variable X into the random variable Z . As a direct consequence, modeling the
Algorithm 1 transform(ddf , S, x ) .
Transforms x
(cid:55)→ z
Input: x , ddf // ddf is a list of classiﬁers [ n 1 , n 2 , . . . ]
Let z = x for n, shuﬄe in ddf , S do
Split [ z a , z b ] =
z z b = cond perm( z b | θ = n ( z a ))
Combine z
=
[ z a , z b ] z = shuﬄe( z ) end for return z
Algorithm
Input: n new , z // n new is a pixel-wise classiﬁer
Split [ z a , z b ] =
z
Optim.
log C ( z b | θ = n new ( z a )) //
Equiv.
to cross-entropy
Algorithm
Input: number of layers L ddf =
[ ] // create DFF with 0 layers
Init S =
[shuﬄe 1 , . . .
shuﬄe L ] //
Init L shufﬂing layers while i < L do
Init n new // new classiﬁer while n new not converged do
Sample x ∼ Data z = transform(ddf , S, x ) optimize( n new , z ) end while ddf .
append( n new ) end while return ddf distribution p Z ( · ) with a D -dimensional i.i.d.
categorical distribution will result in a smaller mismatch than it would for the distribution p X ( · ) .
To give some more intuition on this functionality, we include an illustrating example in appendix A and provide additional augment for the case
K = 2 in appendix B.
Shufﬂing and Splitpriors Algorithms 1, 2 and 3 describe the training process of DDFs.
Note that only one denoising coupling layer is trained at a time.
Moreover, we use invertible shufﬂe operations such that the input is partitioned differently in each coupling layer.
Note that instead of propagating the full input vector x through all layers of the DDF, we factor out parts of the input vector at regular intervals and model these parts conditioned on the other parts following the splitprior approach in (Dinh et al., 2016; Kingma &
Dhariwal, 2018).
As a result, the following coupling layers operate on lower-dimensional data.
Not only is this more efﬁcient, but it also allows for some additional dependencies between parts of the output vector z .
[Figure]
4. Experiments
In this section we explore how the compression rate in bits per dimension (BPD) of Discrete Denoising Flows compares to Discrete Flows on three different data sets.
Each experiment was conducted at least three times; we show the average results as well as the standard deviation in Table 1.
To fully capture the difference in modeling capacity between DFs and DDFs, we trained both models with and without splitpriors.
All experimental details and samples from the models trained without splitpriors can be found in appendices C and D. In each experiment, we use equally sized neural networks in the coupling layers of DFs and DDFs to ensure comparability.
[Table]
Eight Gaussians As a ﬁrst experiment, we train DDFs and DFs on a two-dimensional toy data set also used by Tran et al.
This data set is a mixture of Gaussians withmeans uniformly distributed around a circle and discretized into 91 bins (i.e. K = 91 classes).
We model the data with a single coupling layer per model and set h = K for the DDF coupling layer.
For 2D no splitpriors are used, because that would already make the model universal and we cannot see how well the ﬂow itself performs.
As apparent from the qualitative results in Figure 2 as well as the achieved BPD given in Table 1, DDFs outperform DFs.
Binary MNIST
In a second experiment, we train both
DFs and DDFs on the binarized MNIST data set.
Since the data set has K = 2 classes, we have h = 2 for the DDF coupling layers.
The samples given in Figure 3 and the achieved BPDs in table 1 show that DDFs outperform DFs.
Cityscapes To test the performance of DDFs on imagetype data, we use a 8 -class version of the Cityscapes data set [Figure] [Figure] (Cordts et al., 2016)
modiﬁed by Hoogeboom et al.
This data set contains 32 × 64 segmentation maps, samples are given in ﬁgure 4c.
Since we use multiple coupling layers in this experiment and have a data set with a number of classes K > 2 , there is a trade-off between permuting classes and maintaining structure in each coupling layer.
Therefore, after performing a grid search, we set h = 4 for all DDF coupling layers.
From the samples in ﬁgure 4 and the BPD rates in table 1, we can see that DDFs outperform DFs.
5. Conclusion
In this paper, we have introduced a new discrete ﬂow-based generative model for categorical data distributions, Discrete
Denoising Flows .
We showed that our model outperforms
Discrete Flows in terms of log-likelihood.
A. Denoising Coupling Layer: Example
Consider the distribution
P X 1 Xx 1 \ x   0.
If we were to directly model P X 1 X 2 with a factorized
Bernoulli distribution
P model ( x ) =
Bern( x 1 | p ) · Bern( x 2 | q ) , the highest achievable log likelihood would be
E x ∼ P X 1 X 2 [log 2 P model ( x )] =
E x ∼ P X 1 X 2 [log 2 (Bern( x 1 | p ) · Bern( x 2 | q ))]
≈ − 1 .with p = 0 .
4 and q = 0 .
Suppose now we have trained a neural net n to predict x 2 from x 1 , such that argmax n ( x 1 ) = x 2 .
Using n in a denoising coupling layer T − 1 : X → Z as deﬁned in equation 4 with K =and implicitly h {+= 
+}
we obtain the distribution P Z 1 Z 2 .
In this new distribution essentially the events with probability 0 .
3 and 0 . 1 are swapped.
[Equation] When modeling P Z 1 Z 2 with a factorized Bernoulli distribution, the highest achievable log likelihood is
E [Equation]
∼ [Equation]
[log [Equation] P model [Equation]
] = [Equation]
≈ − 1 .with p = 0 .
4 and q = 0 .
We can see that P Z 1 Z 2 is now modeled higher log-likelihood than the factorized Bernoulli distribution on P X 1 X 2 .
B. Denoising Coupling Layer: General Binary Case
In the following, we generalize the example given in appendix A to provide further insight into the functionality of the denoising coupling layer.
Consider again a twodimensional binary random variable X with probability distribution P X 1 X 2 deﬁned as
P X 1 Xx 1 \ x
pp pp
We train a neural network n to predict x 2 from x 1 , such that argmax n ( x 1 ) = x 2 .
Using n in a denoising coupling layer as deﬁned in equation 4 with K = 2 and implicitly h
= 2 , we obtain the distribution P Z 1 Z 2 .
[Equation] Z 1 Z
[Equation]
We now want to show that applying the denoising coupling layer results in a distribution that can be modelled more accurately with a factorized Bernoulli distribution
P model ( x ) =
Bern( x 1 | p ) · Bern( x 2 | q ) , with parameters 0 ≤ p, q ≤ 1 in z-space than in x-space.
To this end, we demonstrate that the model log-likelihood for P Z 1 Z 2 is always higher than or equal to the model loglikelihood for P X 1 X 2 .
The model log-likelihood of P X 1 X 2 is given as
[Equation]
where the last line assumes the optimal choice for p = p 1 + p 2 and q = p 1 + p 3 and H denotes the binary entropy deﬁned as H ( p )
= − p log p − (1 − p ) log(1 − p ) .
Analogously for
P Z 1 Z 2 using Bernoulli parameters ˆ q, ˆ p E z ∼ P
[Equation] Z 2 [log P model ( z )] = log(ˆ p ) · (max( p 1 , p 2 ) + min( p 1 , p 2 ))
+ log(1 − ˆ p ) · (max( p 3 , p 4 ) + min( p 3 , p 4 )) + log(ˆ q ) · (max( p 1 , p 2 ) + max( p 3 , p 4 )) + log(1 − ˆ q ) · (min( p 1 , p 2 )
+ min( p 3 , p 4 )) = −H (ˆ p ) − H (ˆ q ) with optimal choice ˆ p = max( p 1 , p 2 ) + min( p 1 , p 2 ) =
p 1 + p 2 which is the same, and the new ˆ q = max( p 1 , p 2 ) + max( p 3 , p 4 ) .
Since H ( p ) =
H (ˆ p ) we only need to compare the terms containing ˆ q and q .
We use that −H is a monotonic increasing function when only considering the interval [0 . 5 , 1 . 0] and that H ( p ) =
H (1 − p ) .
First we check if p 1 + p 3 ≥ 0 .
5 or otherwise p 2 + p 4 ≥ 0 .
Let this value be a so that a ≥ 0 .
5 , and the other value b with b ≤ 0 .
From the symmetry we have:
[Equation]
Next, observe that ˆ q = max( p 1 , p 2 ) + max( p 3 , p 4 ) ≥ max( p 1 + p 3 , p 2 + p 4 ) =
a and since −H is monotonically increasing on [0 . 5 , 1 . 0] it follows that:
[Equation]
Plugging this back into the previous equation gives us the desired identity:
[Equation]
under optimal choice of the Bernoulli parameters.
C. Experimental Details
We train Discrete Flows and Discrete Denoising Flows on three data sets.
In each experiment, both models use the same architecture to ensure comparability.
Throughout the experiments, we use the Adam optimizer, a learning rate of 0 . 001 , and a batch size of 64 .
The base distribution is always a factorized categorical distribution with K classes.
K varies between the data sets.
Recall that in the coupling layers of both DFs and DDFs, the Ddimensional coupling layer input x is split into two parts such that x =
[ x a , x b ] , at a split index d .
For all of our experiments, we set d = D 2 .
Eight Gaussians
In this experiment, we train a singlelayer Discrete Flow and a single-layer Discrete Denoising
Flow .
For both models, we use an MLP consisting oflinear layers with 256 hidden units and ReLU activations to parameterize the coupling layer.
This small model size is sufﬁcient for modeling our 2D toy data set.
Consisting of only one coupling layer, preserving the structure of the input vector for later coupling layers is not relevant for the DDF model.
Consequently, we set the parameter h in the denoising coupling layer to the number of classes in the data set K = 91 .
Binary MNIST
In this experiment, we work with binary image data.
Consequently, each DF and DDF coupling layer is parameterized by a DenseNet (Huang et al., 2017) consisting of 8 dense building blocks.
For DDFs, modeling binary data implies that h equals the number of classes K , i.e. h = K = 2 .
We embed the coupling layers into a multi-layer architecture of coupling layers, split priors, and squeeze operations (Dinh et al., 2016).
The squeeze changes the vector size from [ channels × H × W ] to [4 · channels × H 2 × W 2 ] .
The overall model architecture consists of 2 blocks that consisting of the following layers (in that order): [Equation]
Note that each coupling layer is preceded by a shufﬂing operation applied to the channels of the input vector.
Further the splitprior factors out the opposite part that the coupling transformed (so if the coupling layer transforms x b (cid:55)→ z b then z a is factored out).
Cityscapes In this experiment, we’re again dealing with image-type data, this time with K = 8 classes.
Like in the previous experiment, we utilize a DenseNet (Huang et al., 2017) in the DF and DDF coupling layers.
However, here it consists of 15 dense building blocks.
We perform a grid search for the DDF parameter h on { 1 , 2 , 4 , 6 , 8 } and ﬁnd that h
= 4 leads to the best performance.
Since the last coupling layer does not have a trade-off between permuting classes and maintaining structure for the next coupling layer to work on, we can set h
= K = 8 in the last layer.
In analogy to the previous experiment, we embed the coupling layers in a multi-layer architecture of coupling layers, splitpriors and squeeze operations.
For this experiment, the model architecture consists of 3 building blocks with the following layers (in that order): {
squeeze - coupling - splitprior - coupling
- splitprior }
Again, the splitprior factors out the opposite part that the coupling transformed (so if the coupling layer transforms x b (cid:55)→ z b then z a is factored out). D. Training without splitpriors
In the experiments without splitpriors, we propagate the full input vector x through all layers of the Discrete Flow and Discrete Denoising Flow models.
We perform this additional set of experiments to further illustrate the difference in modeling capacity between Discrete Flows and
Discrete Denoising Flows .
When both models are trained with splitpriors, the splitpriors allow for additional dependencies between the input and output of the model, thus increasing modeling capacity.
If we omit the splitpriors, the model performance depends solely on the coupling layers, allowing us to better see the difference in performance between the two types of ﬂows.
D.1. Experimental Setup
Similar to the experiments with splitpriors, we use the Adam optimizer, a learning rate of 0 .
001 , a batch size of 64 and set the split index d
= D 2 with D being the dimensionality of the input x for all experiments without splitpriors.
Binary MNIST Equivalent to the MNIST binary experiment described in Appendix C, each DF and DDF coupling layer is parameterized by a DenseNet (Huang et al., 2017) consisting of 8 dense building blocks, and we set h
= K
= 2 for the DDF model.
The overall model architecture for this experiment consists of 2 building blocks with the following layers (in that order):
[Equation]
where each coupling layer is preceded by a shufﬂing operation applied to the channels of the input vector.
Cityscapes Like in the Cityscapes experiment described in Appendix C, each DF and DDF coupling layer is parameterized by a DenseNet (Huang et al., 2017) consisting of 15 dense building blocks, and we set h = 4 for each DDF coupling layer but the last, where we set h
= K = 8 .
The model architecture for this experiment consists ofbuilding block with layers (in that order) { squeeze - coupling } followed by 2 building blocks with layers (in that order) { squeeze - coupling - coupling } where each coupling layer is preceded by a shufﬂing operation applied to the channels of the input vector.
D.2. Evaluation
Figure 5 and ﬁgure 6 show samples from the trained DF and DDF models on the binary MNIST data set and on the Cityscapes data set.
As we can see from the samples as well as from the quantitative results given in table 1, there is a bigger difference in performance between DFs and DDFs when they are trained without splitpriors.
The reason is that splitpriors already increase the model capacity, making the relative difference smaller.
Therefore, this comparison shows even more clearly the improved modeling performance of our DDFs over DFs.
[Figure]
[Figure]