IDF++: A NALYZING AND I MPROVING I NTEGER D ISCRETE F LOWS FOR L OSSLESS C OMPRESSION
In this paper we analyse and improve integer discrete ﬂows for lossless compression.
Integer discrete ﬂows are a recently proposed class of models that learn invertible transformations for integer-valued random variables.
Their discrete nature makes them particularly suitable for lossless compression with entropy coding schemes.
We start by investigating a recent theoretical claim that states that invertible ﬂows for discrete random variables are less ﬂexible than their continuous counterparts.
We demonstrate with a proof that this claim does not hold for integer discrete ﬂows due to the embedding of data with ﬁnite support into the countably inﬁnite integer lattice.
Furthermore, we zoom in on the effect of gradient bias due to the straight-through estimator in integer discrete ﬂows, and demonstrate that its inﬂuence is highly dependent on architecture choices and less prominent than previously thought.
Finally, we show how different architecture modiﬁcations improve the performance of this model class for lossless compression, and that they also enable more efﬁcient compression: a model with half the number of ﬂow layers performs on par with or better than the original integer discrete ﬂow model.
I NTRODUCTION
Density estimation algorithms that minimize the cross entropy between a data distribution and a model distribution can be interpreted as lossless compression algorithms because the cross-entropy upper bounds the data entropy.
While autoregressive neural networks (Uria et al., 2014; Theis &
Bethge, 2015; Oord et al., 2016; Salimans et al., 2017) and variational auto-encoders (Kingma &
Welling, 2013; Rezende & Mohamed, 2015) have seen practical connections to lossless compression for some time, normalizing ﬂows were only recently used for lossless compression.
Most normalizing ﬂow models are designed for real-valued data, which complicates an efﬁcient connection with entropy coders for lossless compression since entropy coders require discretized data.
However, normalizing ﬂows for real-valued data were recently connected to bits-back coding by Ho et al.
(2019b)
, opening up the possibility for efﬁcient dataset compression with high compression rates.
Orthogonal to this,
Tran et al.
and Hoogeboom et al.
(2019a)
introduced normalizing ﬂows for discrete random variables.
Hoogeboom et al. (2019a) demonstrated that integer discrete ﬂows can be connected directly to entropy coders without the need for bits-back coding.
In this paper we aim to improve integer discrete ﬂows for lossless compression.
Recent literature has proposed several hypotheses on the weaknesses of this model class, which we investigate as potential directions for improving compression performance.
More speciﬁcally, we start by discussing the claim on the ﬂexibility of normalizing ﬂows for discrete random variables by Papamakarios et al.
(2019), and we show that this limitation on ﬂexibility does not apply to integer discrete ﬂows.
We then continue by discussing the potential inﬂuence of gradient bias on the training of integer discrete ﬂows.
We demonstrate that other less-biased gradient estimators do not improve ﬁnal results.
Furthermore, through a numerical analysis on a toy example we show that the straight-through gradient estimates for 8-bit data correlate well with ﬁnite difference estimates of the gradient.
We also demonstrate that the previously observed performance degradation as a function of number of ﬂows is highly dependent on the architecture of the coupling layers.
Motivated by this last ﬁnding, we introduce several architecture changes that improve the performance of this model class on lossless image compression.
R ELATED WORK
Continuous Generative Models: Continuous generative ﬂow-based models (Chen & Gopinath, 2001;
Dinh et al., 2014; 2016) are attractive due to their tractable likelihood computation.
Recently these models have demonstrated promising performance in modeling images (Ho et al., 2019a;
Kingma & Dhariwal, 2018), audio (Kim et al., 2018), and video (Kumar et al., 2019).
We refer to
Papamakarios et al.
for a recent comprehensive review of the ﬁeld.
By discretizing the continuous latent vectors of variational auto-encoders and ﬂow-based models, efﬁcient lossless compression can be achieved using bits-back coding (Hinton & Van Camp, 1993).
Recent examples of such approaches are Local Bits-Back Coding with normalizing ﬂows (Ho et al., 2019b) and variational auto-encoders with bits-back coding such as Bits-Back with ANS (Townsend et al., 2019b),
Bit-Swap (Kingma et al., 2019) and HiLLoC (Townsend et al., 2019a).
These methods achieve good performance when compressing a full dataset, such as the ImageNet test set, since the auxiliary bits needed for bits-back coding can be amortized across many samples.
However, encoding a single image would require more bits than the original image itself (Ho et al., 2019b).
Learned discrete lossless compression: Producing discrete codes allows entropy coders to be directly applied to single data instances.
Mentzer et al.
(2019) encode an image into a set of discrete multiscale latent vectors that can be stored efﬁciently.
Fully autoregressive generative models condition unseen pixels directly on the previously observed pixel values and have achieved the best likelihood values compared to other models (Oord et al., 2016; Salimans et al., 2017).
However, decoding with these models is impractically slow since the conditional distribution for each pixel has to be computed sequentially.
Recently, super-resolution networks were used for lossless compression (Cao et al., 2020) by storing a low resolution image in raw format and by encoding the corrections needed for lossless up-sampling to the full image resolution with a partial autoregressive model.
Finally, Mentzer et al.
ﬁrst encode an image using an efﬁcient lossy compression algorithm and store the residual using a generative model conditioned on the lossy image encoding.
Hand-designed Lossless Compression Codecs: The popular PNG algorithm (Boutell & Lane, 1997) leverages a simple autoregressive model and the DEFLATE algorithm (Deutsch, 1996) for compression.
WebP (Rabbat, 2010) uses larger patches for conditional compression coupled with a custom entropy coder.
In its lossless mode, JPEG 2000 (Rabbani, 2002) transforms an image using wavelet transforms at multiple scales before encoding.
Lastly, FLIF (Sneyers & Wuille, 2016) uses an adaptive entropy coder that selects the local context model using a per-image learned decision tree.
B ACKGROUND : NORMALIZING FLOWS
In this section we brieﬂy review normalizing ﬂows for real-valued and discrete random variables.
A normalizing ﬂow consists of a sequence of invertible functions applied to a random variable x : f K ◦ f K − 1 ◦ ... ◦ f 1 ( x ) , yielding random variables y K ← ...
← y 1 ← y
= x .
First, consider a realvalued random variable x ∈ R d with unknown distribution p x ( x ) .
Let f : R d (cid:55)→ R d be an invertible function that such that y =
f
( x ) with y ∈ R d .
If we impose a density p y
( y ) on y , the distribution p x ( x ) is obtained by marginalizing out y from the joint distribution p x , y ( x , y ) =
p x | y ( x | y )
p y
( y ) :
p x (
[Equation]
)= (cid:90)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
[Table]
u )
[Equation]
u
[Table]
[Equation] [Equation]
where we used p x | y ( x | y )
=
δ ( x − f − 1 ( y )) , with δ ( x − x (cid:48) )
the Dirac delta distribution, and we applied a change of variables.
Repeated application of (1) for a sequence of transformations then yields the log-probability:
[Equation]
By parameterizing the invertible functions with invertible neural networks and by choosing a tractable distribution p y K ( y K ) these models can be used to optimize the log-likelihood of x .
When modeling discrete data with continuous ﬂow models, dequantization noise must be added to the input data to ensure that a lower bound to the discrete log-likelihood is optimized (Uria et al., 2013; Theis et al., 2015; Ho et al., 2019a).
3. D ISCRETE NORMALIZING FLOWS
Next, consider x to be a discrete random variable with domain X , and deﬁne the invertible function f : X (cid:55)→ X .
In general f can be a mapping between two different domains, but for our discussions it will be sufﬁcient to consider a single domain.
The marginal probability mass of x is given by
[Equation]
with the Kronecker delta function δ i,j = 1 if i = j and 0 otherwise.
Note the absence of a volume correction in the form of a Jacobian determinant owing to the fact that probability mass functions only have support on a discrete set of points.
Recently, Tran et al.
and Hoogeboom et al.
(2019a) have both considered normalizing ﬂows for discrete random variables.
In integer discrete ﬂows (IDF) by Hoogeboom et al. (2019a)
the random variables are assumed to be integers, i.e. X =
Z d .
The main building block of IDF is an additive bipartite coupling layer (Dinh et al., 2016):
[Equation]
Here y a ∈ Z m , y b ∈
Z n are obtained by splitting y ∈
Z d into two pieces such that m + n = d , and similarly for x a and x b .
The pre-quantized translation t θ ( · ) is represented by the output of a neural network with learnable parameters θ , and is rounded to integer values with the rounding operator (cid:98)·(cid:101) .
The parameters θ are optimized with a straight-through estimator and a gradient-based optimizer.
Tran et al.
(2019) introduce ﬂows for non-ordinal discrete random values with a ﬁnite number of possible values: x ∈ X = { 0 , 1 , ..., K − 1 } d .
They introduce an autoregressive and a coupling bijector, with the coupling layer given by [ y a , y b ] =
[ x a , ( s θ ( x a ) ◦ x b + t θ ( x a ))
mod K ] , with ◦ denoting element-wise multiplication and s ∈ { 1 , 2 ..., K − 1 } and t ∈ { 0 , 1 , ..., K − 1 } and s and
K constrained to be co-prime.
Gradients are again computed with a straight-through estimator.
One of the main differences between discrete ﬂows (Tran et al., 2019) and integer discrete ﬂows (Hoogeboom et al., 2019a) is that the former treats the random variable as having a ﬁnite number of classes while the latter considers a countably inﬁnite number of classes.
In the next section we will see that this difference plays an important role in the theoretical ﬂexibility of invertible ﬂows for discrete random variables.
C AN INTEGER DISCRETE FLOWS FACTORIZE ANY DISTRIBUTION ?
In this section we discuss the theoretical ﬂexibility of normalizing ﬂows for discrete random variables.
In particular, we will focus on the ability to map from a distribution with dependencies across all dimensions to a distribution which is fully factorized (i.e. independent across all dimensions).
As generative ﬂow models require a tractable base distribution with efﬁcient sampling, simple base distributions (e.g. Gaussian, logistic, categorical) which are independent across dimensions/subpixels are frequently used.
Papamakarios et al.
(2019) state that invertible ﬂows for discrete random variables cannot map all distributions to a factorized base distribution, and name this as a limitation as compared to ﬂows for real-valued random variables.
We will analyze this claim and show that this limitation can be overcome by embedding the discrete data into a space with a larger set of possible values.
Since integer discrete ﬂows embed the data into the integer lattice Z d , this model class naturally does not suffer from a limited factorization capacity.
The starting point for this discussion lies in the observation by Papamakarios et al.
(2019) that invertible normalizing ﬂows for discrete random variables can only permute probability masses in the probability tensor that represents the probability distribution of the random variable.
In other words, if we have an invertible function f : X (cid:55)→ X , then there is always exactly one y , such that
y
=
f
( x ) and p x ( x ) =
p y ( y ) .
In contrast, non-volume preserving normalizing ﬂows for real-valued random variables can increase or decrease the density through the Jacobian determinant in the change of variables formula in (1).
Papamakarios et al.
(2019) then discuss an educative example to show that this permutation property can lead to a limited ability to factorize distributions.
Consider the case of a two-dimensional random variable x = ( x 1 , x 2 ) , with x 1 , x 2 ∈ { 0 , 1 } , and a data distribution given by p
[Equation]
In order to map this probability distribution to an independent base distribution, the corresponding probability matrix must be of rank one.
In other words, given a sequence of functions f k : X (cid:55)→ X with X = { 0 , 1 } 2 , the probability distribution of the random variable y =
y K
=
f K ◦ f K − 1 ◦ ... ◦
f 1 ( x ) should be represented by a matrix with linearly dependent columns or rows.
This allows for a factorization into an outer product of two vectors that represent the independent base distributions for y 1 and y 2 .
Since discrete ﬂows can only permute probability mass tensors, the matrix corresponding to p y
( y 1 , y 2 ) must be a permutation of the matrix in (5).
However, there is no permutation of the elements in (5) that results in a rank-one matrix.
Therefore, Papamakarios et al.
(2019) conclude that discrete normalizing ﬂows cannot map any distribution to a factorized base distribution.
However, one of the key assumptions made above is that the domain of f k is restricted to X = { 0 , 1 } 2 .
By extending X to a larger number of classes this example can in fact be factorized by a discrete normalizing ﬂow.
More concretely, let us extend the domain to X = { 0 , 1 , 2 , 3 } 2 .
The probability distribution matrix of x is shown below, together with a permutation of rank 1.
[Equation]
This example illustrates that the number of classes that are considered valid for the discrete random variables plays a crucial role in the ﬂexibility of discrete ﬂows.
Here, we claim that this holds more generally through the following two lemmas.
First we show that by embedding the data into a space with more possible values than present in the data itself, one can always construct an invertible mapping to one-dimensional variable embedded in d dimensions, and that the resulting variable has a distribution which is trivially factorized across all d dimensions.
Lemma 1.
A d -dimensional discrete random variable x = ( x 1 , . . .
, x d ) T with d > 1 and x i ∈ { 0 , . . .
, K ( i ) − 1 } , distributed according to an arbitrary distribution p x ( x ) , can be transformed to a one-dimensional random variable y ∈ { 0 , . . .
, (cid:16)(cid:81) di =1 K ( i ) (cid:17) − 1 } with a number of classes that scales exponentially with the dimension d through a bijective mapping f .
Embedding y in d dimensions through padding with d − 1 zeros, i.e. y =
( y, 0 , . .
. , 0) T , the distribution over y is trivially factorized across all dimensions: p x ( x ) =
p y ( f ( x )) (cid:81)
d i =
δ y i , 0 with δ i,j the Kronecker delta distribution for all zero-padded dimensions.
The closed form of the mapping is y =
f
( x ) = x 1 + K (1) x
2 + K (1) K (2) x 3 + · · · + (cid:81) d − 1 i
=1 K ( i ) x d , such that y ∈ { 0 , . . .
, (cid:16)(cid:81) di =1 K ( i ) (cid:17) − 1 } .
See Appendix A for the full proof.
Intuitively, given enough classes, a factorized base distribution is obtained by “ﬂattening” the hypercube that contains all nonzero entries in the data distribution tensor (see left panel Figure 1) into one dimension.
The following lemma demonstrates that additive coupling layers as in (4) of integer discrete ﬂows can model such an invertible mapping.
Lemma 2.
For a d -dimensional random variable x , translation operations of the form z a = x a , z b = x b + (cid:98) t ( x a ) (cid:101) , with a and b denoting the indices of two splits of the variable x , are sufﬁcient to map x in an invertible manner to a random variable y
=
( y 1 , 0 , . .
. , 0) T , which is a one-dimensional variable embedded in d dimensions.
The formal proof is an induction on d and relies on 1) the fact that integer discrete ﬂows embed the discrete random variables into the integer lattice X = Z d , which always has a sufﬁcient number of classes (countably inﬁnite), and 2) the observation that two translations are sufﬁcient to map two dimensions into a single dimension with more classes, see Appendix A. For a more intuitive illustration, Figure 1 depicts the type of operations that a single additive bijector can and cannot perform on a three-dimensional probability tensor.
In Figure 2 we show that, although sensitive to initialization, an IDF can be trained to factorize the toy example of (5).
The ﬁrst three panels correspond to the data distribution and the intermediate and factorized ﬁnal distributions produced by the ﬁrst and second coupling layer respectively.
For architecture details see Appendix D.3.
As pointed out by Papamakarios et al.
(2019), a discrete ﬂow can only map a uniform base distribution into a uniform data distribution due to its permutation property.
Therefore, invertible discrete ﬂows require one-dimensional learnable distributions in order to be able to model all data distributions.
Note also that discrete ﬂows as introduced by Tran et al.
(2019) are designed for a ﬁnite number of classes in X and therefore cannot factorize all distributions.
Although not implemented in the original work, this model class can in principle also be extended to embed the data into a space with a larger ﬁnite number of possible values than present in the support of the data in order to help alleviate this issue.
Since IDFs treat the random variable as integers (with a countably inﬁnite number of classes) it does not have a limited factorization capacity.
We therefore explore other directions for potential improvements.
D OES GRADIENT BIAS HINDER OPTIMIZATION OF IDF S ?
In this section we investigate whether gradient bias is a problem when training IDF models.
Hoogeboom et al. (2019a) demonstrated that the performance of IDFs can deteriorate when the number of coupling layers is too large, and that this does not occur in the continuous additive counterpart without rounding operators or straight-through gradient estimators.
The gradient bias induced by the straight-through estimator was suggested to be the cause.
In our experiments we found that neither stochastic rounding nor replacing the identity function in the straight-through estimator with a soft approximation of the rounding function improved the results.
To disentangle the effect of gradient bias of the straight-through estimator and the difﬁculty of discrete optimization, we have trained IDFs and their continuous additive counterpart (without the rounding operator and without a quantized base distribution) on CIFAR-10.
We use different combinations of translation activation functions and backward substitute functions in the straight-through estimator: soft rounding, hard rounding and the identity function.
The soft rounding function σ T was modeled with a superposition of scaled sigmoid functions and a temperature T that allows to interpolate between the identity function ( T = 1 ) and the rounding function T = 0 , see Appendix C for more details.
The left plot in Figure 3 shows that the continuous model with an identity straight-through gradient estimator consistently outperforms the continuous model without gradient bias.
The discrete model clearly also does not beneﬁt from a soft-rounding operator in the straight-through estimator, even though this estimator reduces the bias.
Note that the continuous model trained with a hard rounding function and identity straight-through estimator ( → rnd, ← id ) performs worse than the discrete model because the continuous model needs to maximize a lower bound to the likelihood.
We numerically study the signiﬁcance of gradient bias in more detail on an extension of the twodimensional toy example of Section 4 to 8 bits with a model with two coupling layers.
For details on the 8-bit extension of the toy example see Appendix B and for architecture details see Appendix
D.3.
We study the search directions using ﬁnite differences as an approximate gradient vector g fd with elements g fd i =
( L ( θ i + (cid:15), θ /i )
− L ( θ i − (cid:15), θ /i )) / 2 (cid:15) .
Here L is the loss function averaged over a single batch of data.
For additive continuous ﬂow models, g fd will approach the true loss gradient ∇ θ L as (cid:15) → 0 .
For discrete models, g fd can be thought of as a linear approximation of the loss landscape in a trust-region of radius (cid:15) around the current parameter vector θ .
We compare continuous ﬂow models that are trained using the unbiased gradient ∇ θ L with discrete ﬂow models that are trained using the straight-through gradient estimator g st .
We estimate the agreement between the (approximate) gradient and the ﬁnite difference approximation g fd for varying trust-region size (cid:15) at various stages of training and for varying bit-depth of the input data.
We compute the agreement of the (approximate) gradient directions with the cosine similarity g · g fd / || g
||
|| g fd
|| 2 , which can be interpreted as the uncentered correlation between the elements of g and g fd .
As long as the agreement is consistently positive, performing gradient descent with g is expected to reduce the loss according to the trust-region approximation based on g fd .
If the agreement is consistently zero, or even negative, gradient descent with g is not expected to improve the training loss.
The agreements are estimated over multiple batches.
As the base distribution parameters are not affected by the gradient bias we only consider the gradients for the parameters of the bijectors.
Figure 4 shows the agreement between the ﬁnite-difference gradient and the straight-through gradient of the discrete model and the agreement with the real gradient for the continuous model, at initialization and after 3000 training iterations.
The straight-through gradient estimator for 8-bit data is always positively correlated with small but ﬁnite difference estimates, corroborating our results on
CIFAR-10 that the identity straight-through gradient estimator allows for good optimization.
For 1-bit data the quality of the straight-through gradient estimator clearly deteriorates and the correlation becomes zero or even negative.
This is not surprising as the random variables and translations are in practice modeled on the rescaled grid Z / 2 bits , leading to more coarse-grained rounding for lower bits.
As lossless compression mostly deals with source data in the higher bit ranges the poor performance for lower bit data is less relevant.
The bottom two panels in Figure 4 show the loss surface and optimization trajectory of a discrete model trained on the 1-bit data, suggesting that the optimization landscape is hard to traverse for low bit-depths.
See Appendix B for loss landscapes for 8-bit data and for similar plots for the corresponding continuous models.
Finally, we observed that the performance degradation of IDF as a function of the number of coupling layers is highly dependent on the architecture used for the translation networks in the coupling layers.
The right plot in Figure 3 shows the performance of continuous and discrete models with identity straight-through estimators, under different coupling layer architectures: stacked convolutions,
ResNets (He et al., 2016) or DenseNets (Huang et al., 2017).
All coupling layer architectures are of approximately the same size and depth, and the DenseNet models have the same architecture as the best model of (Hoogeboom et al., 2019b).
The performance of the convolutional model clearly deteriorates quicker than the ResNet and DenseNet models.
The results in this section show that gradient bias due to the straight-through estimator is less of a problem than previously suggested.
Alternative gradient estimators with less bias do not improve the results, and our analysis on the toy examples shows that for data with large bit depths the straightthrough gradient estimates have positive correlation with ﬁnite difference estimates.
We show that the previously observed performance deterioration as a function of the number of coupling layers is highly dependent on the architecture choice of the coupling layers.
This motivates our search for architectural changes that can help improve the performance of IDFs on lossless image compression.
I MPROVING I NTEGER D ISCRETE F LOWS FOR L OSSLESS C OMPRESSION
In this section we introduce changes to the architecture of IDF that improve its performance.
We brieﬂy summarize the IDF architecture, for more detailed information see Appendix D. Similar to other generative ﬂow models, the architecture of IDF contains L levels.
Each level l consists of a space-to-depth transformation followed by a sequence of K alternating channel permutations and additive coupling layers.
For levels l = 1 , ..., L − 1 the output random variable is split into two pieces, with the second part serving as an input to the next level:
[ z ( l ) , y ( l ) ]
= f ( l,K ) ◦ ... ◦
f ( l, 1) ( y ( l − 1) ) , with y (0) = x .
For the last level we simply have z ( L ) =
f ( L,K ) ◦ ... ◦
f ( L, 1) ( y ( L − 1) ) .
The combined z =
[ z (1) , ..., z ( L ) ] denotes the latent representation of x .
The distribution of z for 3 levels is then factorized as p ( z ) =
p ( z (1) , z (2) , z (3) )
=
p ( z (3) ) p ( z (2) | y (2) )
p ( z (1) | y (1) ) , which is equivalent to p ( z (3) ) p ( z (2) | z (3) ) p ( z (1)
| z (2) , z (3) ) .
The conditional distributions are modeled as discretized logistic distributions and the unconditional p ( z (3) ) is modeled with a mixture of discretized logistics with 5 components.
Both the conditioning networks and the pre-quantized translations t θ in (4) are modeled with DenseNets (Huang et al., 2017).
The ﬁrst modiﬁcation is to invert the channel permutations after every coupling layer.
Although inverting the permutation after each translation does not affect the coupling layers of the network, it affects the way the data is split after each level and therefore inﬂuences the modeling of the conditional distributions p ( z ( l ) | y ( l ) ) .
By inverting the permutations one ensures that the split into y ( l ) and z ( l ) happens along the channel direction of the space-to-depth transformed version of y ( l − 1) ; as such y ( l ) and z ( l ) retain the spatial correlation structure of the original image presumably making it easier to model the conditional distribution p ( z ( l ) | y ( l ) )
(see visualization in Figure 10 of Appendix D).
The second change is an adaptation of the rezero trick by Bachlechner et al.
The additive bijectors of (4) are replaced by [ y a , y b ] =
[ x a , x b + (cid:98) α t θ ( x a ) (cid:101) ] , where α is a learnable scalar parameter initialized to zero, such that the bijectors are initialized to the identity function.
The mean and log-scale of the conditional discretized logistic distributions are parameterized as µ = γ ν , log s = δ log σ with [ ν , log σ ]
=
DenseNet φ ( y ) and γ and δ learnable scalar parameters initialized to zero.
The third modiﬁcation consists of an alteration in the dense blocks that make up the translation and logistic conditioning DenseNets by introducing group normalization layers (Wu & He, 2018) and switching from ReLU activations (Nair & Hinton, 2010) to Swish activations (Ramachandran et al., 2017):
[Equation]
The trainability of deep neural networks such as ﬂow models is strongly dependent on careful initialization and normalization.
Identity mapping initialization of ﬂow models, which we achieve by virtue of the rezero trick, is a known technique for improving training stability of these models (Kingma & Dhariwal, 2018).
The use of group normalization layers ensures that the dense blocks within the coupling layers receive properly scaled and centered inputs.
Combined with the Swish activation function this allows DenseNet activations to remain non-zero and the coupling layers to utilize their capacity more effectively.
Empirically, these intuitions are supported by the fact that collectively these modiﬁcations allow for the use of a higher base learning rate during training while achieving better results.
Finally, during training we also keep track of a slow exponential moving average of the models weights, and use these average weights during evaluation.
Figure 3 shows the performance of models with the proposed modiﬁcations (DenseNet++) on the validation set (consisting of 20% of the training set) on CIFAR-10 as a function of ﬂows per level, after 300K iterations.
For this new architecture the performance of both the continuous and the discrete version of the model deteriorates for 10 ﬂows per level due to overﬁtting (see Figure 11).
These results also show that the modiﬁcations allow for a more efﬁcient use of the ﬂow layers: for 300K iterations the performance of the IDF++ model with a DenseNet++ architecture is on par with that of the IDF model of Hoogeboom et al. (2019a) with DenseNet coupling layers.
To further corroborate this point, we have trained an IDF++ model with 4 and 8 ﬂows per level until convergence for CIFAR-10, ImageNet-32 and ImageNet-64, and compared it with the baseline IDF model with 8 ﬂows per level and other related work in Table 1.
To make a fair comparison against other methods like local bits-back coding (LBB) by Ho et al. (2019b)
we train our ﬁnal models on the entire training set without holding out part of the training set as a validation set.
Although there is no visible difference for the ImageNet-32 and ImageNet-64 datasets, more data does matter for
CIFAR-10: it reduces the model’s average negative log-likelihood in units of bits per dimension (bpd) on the test set from 3 .
32 (as reported by Hoogeboom et al. (2019a)) to 3 .
30 for the baseline
IDF model.
Table 1 displays the bpd resulting from compressing the latent representations z with a range-based Asymmetric Numerical Systems (rANS) entropy coder (Duda, 2013).
Results for hand-designed codecs, models with bits-back coding and models without bits-back coding are shown in separate groups.
Where available, the bpd as predicted from the models negative log-likelihood is indicated in brackets.
[Table]
The results show that, contrary to less-biased gradient estimators, the proposed architecture modiﬁcations in IDF++ improve the performance over IDF for 8 ﬂows per level for all datasets.
Furthermore, the smaller IDF++ model with only 4 ﬂows per level either performs on par or better than the
IDF baseline with 8 ﬂows per level.
This effectively reduces the number of parameters and time required for a forward pass through the model by a factor of 2, and demonstrates that the proposed modiﬁcations in IDF++ enable more efﬁcient compression without sacriﬁcing performance.
For a more detailed ablation on the contributions of the proposed modiﬁcations see Appendix E.
C ONCLUSION
In this paper we investigated several potential directions for improvement of integer discrete ﬂows for lossless compression.
We analyzed the hypothesis and claims of recent works on the weaknesses of this model class, and we showed that the claim by Papamakarios et al.
(2019) that ﬂows for discrete random variables cannot factorize all distributions does not apply to integer discrete ﬂows.
Furthermore, we demonstrated that the effect of gradient bias on optimization is less severe than previously thought.
Other gradient estimators with less bias do not improve optimization, and a numerical analysis of the direction of the straight-through gradient estimates showed that they correlate well with ﬁnite difference estimates for large-bit data.
We also showed that the previously observed performance deterioration with increasing depth of the ﬂow model is highly dependent on the architecture of the coupling layers.
Motivated by this, we proposed several architecture changes that improve the performance of IDFs on lossless image compression for an equal computational budget.
The proposed modiﬁcations lead to a more efﬁcient ﬂow-based compression model, as evidenced by our results that show that an IDF++ model with half the number of ﬂows compared to the baseline IDF model performs on par or better.
Although we found that the simple straight-through gradient estimator outperformed all of the other estimators that we considered, future work could consider other estimators inspired by universal quantization and the statistics of quantization noise, similar to the work by Agustsson & Theis (2020).
Another important direction for future work is to further reduce the computational complexity of deep density estimators.
Although hand-designed codecs such as JPEG-2000 do not compress as well on the datasets we consider, they (de)compress signiﬁcantly faster and require signiﬁcantly less memory while not being tuned for each dataset.
More work in directions like optimizing for (de)compression speed (Cao et al., 2020) or generalizing learnable compressors to other datasets (Townsend et al., 2019a) is needed to make deep density estimators more practical for source compression.
A PPENDIX A F LEXIBILITY OF INTEGER DISCRETE FLOWS
The proof of Lemma 1 can be done in two ways, one is by giving the exact mapping, and the second way is through induction.
We include both below as well as the proof for Lemma 2.
Proof of Lemma 1.
The invertible mapping is given by y = f
( x ) = x 1 + K (1) x
2 + K (1) K (2) x 3 + · · · + (cid:81) d − 1 i
=1 K ( i ) x d , with a maximum value of (cid:16)(cid:81) di =1 K ( i ) (cid:17) − 1 and minimum value of zero.
Given y , each element x i can be obtained recursively: x i = (cid:16) y − (cid:16)(cid:80)
dj =
i +1 K j − 1 x j (cid:17)(cid:17) //K
i − 1 , where // refers to integer division.
Through the change of variables formula we know that p y
( y ) =
p x ( f − 1 ( y )) , which combined with zero padding to embed in d dimensions leads to the desired factorized distribution in d dimensions.
Alternative proof of Lemma 1.
An alternative proof is via induction on d .
Base case: Let us start with d = 2 , we have the random variable x
= ( x 1 , x 2 ) T with x 1 ∈ { 0 , ..., K (1) − 1 } and x 2 ∈ { 0 , ..., K (2) − 1 } .
x is distributed according to p x ( x 1 , x 2 ) .
The following function can map this random variable to a one-dimensional random variable y with y ∈ { 0 , . .
, K (1) K (2) − 1 } :
y = x 1 + K (1) x 2 .
Given y , x 1 and x 2 can be recovered through x 2 = y//K (1) and x 1 =
y − K (1) x 2 .
General case: Now let us assume that the lemma holds for d = n , and examine the case of d = n + 1 : x = ( x 1 , ..., x n , x n +1 ) T with x i ∈ { 0 , ..., K ( i ) } .
Then the following transformation maps to an n -dimensional random variable y =
( y 1 , . . .
y n ) T : y i
= x i for i = { 1 , . . .
, n − 1 } and y n = x n + K ( n ) x n +1 .
x can be recovered from y through x i =
y i for i = { 1 , . . .
, n − 1 } and x n +1 = y n //K , x n = y n − Kx n +1 .
We have y 1 , ..., y n − 1 ∈ { 0 , ..., K ( i ) − 1 } and y n ∈ { 0 , ..., K ( n ) K ( n − 1) − 1 } .
As we assumed that we can transform an n -dimensional random variable to a one-dimensional random variable in an invertible manner, this concludes that we can also do so for the d = n + 1 case.
Proof of Lemma 2.
We will prove this via induction on d .
Base case: Let us start again with d = 2 .
x = ( x 1 , x 2 ) T with x i ∈ { 0 , . . .
, K ( i ) } can be transformed into a random variable y =
( y 1 , 0) T in an invertible manner with two translation operations of the form z a = x a , z b = x b + (cid:98) t
( x a ) (cid:101) , where either ( a, b ) =
(1 , 2) or ( a, b ) =
The ﬁrst translation is given by z 1 = x 1 + K (1) x 2 , z 2 = x 2 , followed by y 1 =
z 1 = x 1 + K (1) x 2 , y
= x 2 − ( x 1 + K (1) x 2 ) //K
(1) =
x 2 − x 2 = 0 .
General case: Let us assume translations are sufﬁcient for d = n .
We can then transform a d = n +dimensional variable x = ( x 1 , . . .
, x n , x n +1 ) T in an invertible manner to a random variable that is effectively n -dimensional:
y =
( y 1 , . . .
, y n , 0) with two translations operations:
The ﬁrst translation is:
[Equation]
followed by the second translation:
[Equation]
Which results in y i
= x i for i ∈ { 1 , ..., n − 1 } and y n = x n + K ( n ) x n +1 .
As we assumed integer translations were sufﬁcient to map an n -dimensional variable like ( y 1 , . . .
, y n ) T to a one-dimensional random variable embedded in n dimensions in an invertible manner, we have now proven that this also holds for d = n + 1 .
A PPENDIX B L EARNING TOY EXAMPLES B. G RADIENT BIAS ON THE TOY EXAMPLE
The toy example that is discussed in Section 4 has a probability distribution with nonzero probabilities for x i ∈ { 0 , 1 } .
With only two values per dimension with nonzero probability the data is effectively 1-bit quantized.
Here we use an extension of this toy example for input data that is quantized to a higher bit-depth.
This can be done by modeling the probability masses with log-linearly spaced logits in the interval (0 , 1) .
The resulting distributions for several bit-depths are shown in Figure 5.
[Figure]
B. L OSS LANDSCAPE AND OPTIMIZER TRAJECTORY
In Section 5, we presented the visualization of the loss landscape and the optimization trajectory of the discrete model on a toy example with 1-bit data.
The details on how these plots were obtained will be explained in this section.
Navigating through a loss landscape that is affected by quantization operators can be challenging due to discontinuities.
To illustrate this effect, we visualize the loss landscape and the optimization path for the discrete and continuous models trained on the 1-bit and 8-bit toy examples.
We use the method proposed by Li et al.
(2018), where for given model parameters θ ∗ , we choose two direction vectors θ 1 and θ 2 and plot the value of f ( α, β ) =
L ( θ ∗ + α θ 1 + β θ 2 ) .
To choose the direction vectors, we use model parameters at different stages of training.
Let θ ∗ i be model parameters after iteration i .
Given the set of parameters for n iterations, we apply PCA to the matrix M =
[ θ ∗ 0 − θ ∗ n , . . .
, θ ∗ n − 1 − θ ∗ n ] and select the two most explanatory directions as θ 1 and θ 2 .
The results for the 1-bit toy example are shown in the left two panels of Figure 4 for the discrete ﬂow model.
Figure 6 shows similar plots for the continuous case, demonstrating a much smoother loss landscape.
Figure 4 suggests that discrete models with unlucky initializations can easily end up in a local minimum from which it is hard to escape due to sharp cliffs in the loss landscape.
Figures 7 and 8 show a signiﬁcantly less pronounced difference between the loss landscape of the continuous and discrete model for the 8-bit toy example.
This supports our observation in Section 5 that the gradient bias is less of an issue at higher bit-depths.
A PPENDIX C S OFT ROUNDING FUNCTION
The soft rounding function σ T that is used in Section 5 is given by
[Equation]
The soft rounding function is depicted in Figure 9 for different temperatures.
A PPENDIX D A RCHITECTURES D. A RCHITECTURE AND TRAINING DETAILS OF IDF
The architecture of IDF consists of L levels, each consisting of an alternating sequence of channel permutations and additive bijectors.
More speciﬁcally, the architecture of level l of IDF is built up as follows:
[Equation]
Except for the last level, the output random variable is split into two equal halves [ z ( l ) , y ( l ) ] , where only y ( l ) is transformed by the next level.
Furthermore, y (0) = x and z =
[ z (1) , . . .
, z ( L ) ] denotes the latent representation of x .
Inside the additive transformations of (4) the prequantized translations are modeled using DenseNets (Huang et al., 2017) with dense blocks of the following structure:
[Equation]
All DenseNets have a depth of 12 blocks and 512 channels.
The additive bijector splits the random variable into two parts along the channel dimension with splitting fractions 3 / 4 and 1 / 4 for the untransformed and transformed parts of the random variable.
After translating, the resulting variables are concatenated again along the channel axis.
The distribution of z for L levels is factorized as p ( z ) =
p ( z ( L ) ) p ( z ( L − 1) | y
( L − 1) ) . . .
p ( z (1) | y (1) ) , which is equivalent to p ( z ( L ) ) p ( z ( L − 1)
| z ( L ) ) . . .
p ( z (1) |
z (2) , . . .
z ( L ) ) .
All conditional distributions are modeled with discretized logistic distributions.
The unconditional distribution p ( z ( L ) ) is modeled with a mixture of discretized logistics with ﬁve components.
The log-scale and mean of the conditional logistic distributions are modeled as the outputs of DenseNets with the same structure as the DenseNets of the prequantized translations: [ ν , log σ ] =
DenseNet φ ( y ) .
Note that instead of modeling the random variables as integers ( x ∈ Z d ), they are modeled as discrete random variables on a grid with bin-width 1 / 256 ( x ∈ Z d / 256 ).
The model is trained with the Adamax optimizer (Kingma & Ba, 2014) with an exponential learning rate schedule with base learning rate equal to 1 × 10 − 3 and a linear warmup phase of 10 epochs.
See Table 2 for more details on the learning rate decay, the number of levels, the batch size and the number of epochs used for training.
Note that we used a size original reproducing the as reported by Hoogeboom et al
2019a
ImageNet
ImageNet sizes are
Hoogeboom al
[Table]
Range-based Asymmetric Numerical Systems (rANS) (Duda, 2009; 2013; Townsend, 2020) is used for lossless compression of the latent variables
z =
[ z (1) , . . .
, z ( L ) ] by using the probability distribution p ( z ) =
p ( z ( L ) ) p ( z ( L − 1) | y
( L − 1) ) . . .
p ( z (1) | y (1) ) corresponding to the model’s multilevel structure.
D. A RCHITECTURE AND TRAINING DETAILS OF IDF++
In IDF++, each of the K blocks containing a permutation and additive bijector has an additional inverse channel permutation to ensure that the output random variable has the same spatial and channel ordering as the input random variable: y
[Equation]
→ space-to-depth →
[Equation]
transform
[Equation]
permute
[Equation]
The inversion of each permutation inside a level ensures that the splitting at the end of each level is done along the channel dimension of a space-to-depth transformed image, enabling the conditional distributions to be able to use the spatial correlation between pixels.
See Figure 10 for a visualization.
The dense blocks of the DenseNets for the prequantized translations and the conditional discretized logistics have additional group normalization layers (Wu & He, 2018) and Swish activations (Ramachandran et al., 2017) instead of ReLU activations (Nair & Hinton, 2010):
Dense block: Conv1x1 → GroupNorm → Swish → Conv3x3 → GroupNorm → Swish .
The number of groups for each group normalization layer are determined as follows: if the number of channels of the group normalization layer is divisible by 3 then 3 groups are used, else if it is divisible by 2 then 2 groups are used, and ﬁnally if it is neither divisible by 3 or 2 then a single group is used.
The additive bijectors in (4) are adjusted to include a learnable scalar parameter that ensures initialization to the identity operator, similar to the rezero trick by Bachlechner et al.
[Equation]
Here α is a learnable scalar parameter initialized to zero.
The mean and log-scale of the conditional discretized logistic distributions are parameterized as µ = γ ν , log s = δ log σ with [ ν , log σ ] =
DenseNet θ ( y ) .
γ and δ are learnable scalar parameters initialized to zero, such that the scale is initialized to one and the mean is initialized to zero.
The combination of the rezero trick and the group normalization layers allows us to use a larger base learning rate of 2 × 10 − 3 in the exponential decayed learning rate schedule.
For CIFAR-10, the IDF++
[Equation]
16 c hanne l d i m en s i on               16 c
hanne l d i m en s i on
[Figure]
In a multi-level IDF++ model the original image ( a ) undergoes a space-to-depth transformation ( b ), after which at the end of the ﬁrst level half of the dimensions (blue tones) are factored out and their conditional distribution based on the remaining dimension (green tones) is learned.
This procedure is repeated for the remaining dimensions ( c ) in the next level where again half the dimensions (light green) are factored out and their conditional distribution is modelled based on the remaining dimensions (dark green).
By inverting permutations introduced in the coupling layers of each level we guarantee that the spatial structure of the original image is preserved at all levels.
Speciﬁcally, with this structure, the factored out and remaining dimensions correspond to nearby rows of the original image.
We expect that the spatial correlation present in nearby rows is effectively leveraged by the conditional distributions parameterized with convolutions, leading to better results.
model with 8 ﬂows per level was trained for 1400 epochs instead of 2000 epochs as overﬁtting was observed for more epochs.
We evaluated all variations of IDF with an exponential moving average of the trainable parameters with a decay rate of 0 . 9999 .
D. A RCHITECTURE FOR TOY EXAMPLE MODELS
The models for the toy examples in Section 4 and 5 have a single level with an unconditional prior given by a mixture of logistics with ﬁve components.
The models have two coupling layers with
DenseNets as translation neural networks.
The dense blocks have the following structure
Dense
[Equation]
We found the inclusion of LayerNorm (Ba et al., 2016) to help make the model less sensitive to initialization.
Because the DenseNets in these models are fairly shallow we did not see a performance difference between LayerNorm and GroupNorm.
The DenseNets have a depth of 4 blocks andchannels.
The datapoints are treated as images of shape h × w × c
= 1 × 1 × 2 .
The additive bijectors split the random variable into two parts along the channel dimension.
After translating, the resulting variables are concatenated again along the channel axis.
In between the coupling layers the channels are reversed to allow for alternating conditioning.
The learning rates for the prior parameters and coupling layers are set to 1 × 10 − 3 and 1 × 10 − 4 respectively and we used a batchsize of 128.
D.
D ATASETS
The training set of CIFAR-10 consists of 50000 images and the test set contains 10000 images.
ImageNet-32 and ImageNet-64 contain approximately 1250000 train images and 50000 test images.
While we used a validation set that was held out from the respective training sets for model development, we trained our ﬁnal models on the entire training set for each dataset.
The effect of this is discussed in Section E.
We follow Hoogeboom et al.
(2019a)
and augment the CIFAR-10 dataset with horizontal ﬂipping, reﬂect-padding and random cropping during training.
No augmentation is used for ImageNet-32 and ImageNet-64.
D. H ARDWARE AND SOFTWARE
We implemented our models in TensorFlow (Abadi et al., 2015).
All experiments were run with
NVIDIA V100 GPUs.
A PPENDIX E A BLATION STUDIES
Table 3 shows the contributions of the modiﬁcations proposed in Section 6 for IDF++.
As Hoogeboom et al.
(2019a) only use 80% of the training data for CIFAR-10, we also show the inﬂuence of using the entire train set for training, as opposed to keeping 20% held out as a validation set.
Note that the result for the original IDF model on 80% of the training data reproduces the result of Hoogeboom et al. (2019a).
The held out validation set for ImageNet-32 and ImageNet-64 was much smaller and we noticed no signiﬁcant change in performance when training on the entire training set for these
[Table]
IDF + EMA
[Table]
IDF + EMA + rezero + lr ×
[Table]
IDF + EMA + rezero + lr × 2 + invert perm
[Table]
IDF + EMA + rezero + lr × 2 + invert perm + groupnorm ( IDF++ )
[Table]
datasets.
All models without group normalization were trained for 2000 epochs.
The ﬁnal IDF++ model, which includes group normalization layers, was trained for 1400 epochs to avoid overﬁtting.
In Figure 11 the validation bpd is shown for IDF++ models and their continuous counterparts forand 10 ﬂows per level.
The plot shows that both the continuous and the discrete model with 10 ﬂows per level start to overﬁt.