IDF++: Analyzing and Improving Integer Discrete Flows for Lossless Compression
Authors: Rianne van den Berg
Alexey A. Gritsenko
Mostafa Dehghani
Casper Kaae Sønderby
Tim Salimans
Venue: ICLR 2021 Poster (2021)
Keywords: normalizing flows
lossless source compression
generative modeling

Abstract: In this paper we analyse and improve integer discrete flows for lossless compression. Integer discrete flows 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 flows for discrete random variables are less flexible than their continuous counterparts. We demonstrate with a proof that this claim does not hold for integer discrete flows due to the embedding of data with finite support into the countably infinite integer lattice. Furthermore, we zoom in on the effect of gradient bias due to the straight-through estimator in integer discrete flows, and demonstrate that its influence is highly dependent on architecture choices and less prominent than previously thought. Finally, we show how different architecture modifications improve the performance of this model class for lossless compression, and that they also enable more efficient compression: a model with half the number of flow layers performs on par with or better than the original integer discrete flow model.

Published as a conference paper at ICLR 2021
IDF++: ANALYZING AND IMPROVING INTEGER DISCRETE FLOWS FOR LOSSLESS COMPRESSION
Rianne den Berg, Alexey A. Gritsenko, Mostafa Dehghani,
Casper Kaae Sønderby & Tim Salimans
Google Research
(rianneväberg, agr tsenko, casaane sal imans @ com
casperkaae,
ABSTRACT
In this paper we analyse and improve integer discrete flows for lossless compression. Integer discrete flows are 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 flows for discrete random variables are less flexible than their continuous counterparts. We demonstrate with a proof that this claim does not hold for integer discrete flows due the embedding of data with finite support into the countably infinite integer lattice. Furthermore, we zoom in on the effect of gradient bias due to the straight-through estimator in integer discrete flows, and demonstrate that its influence is highly dependent on architecture choices and less prominent than previously thought. Finally, we show how different architecture modifications improve the performance of this model class for lossless compression, and that they also enable more efficient compression: a model with half the number of flow layers performs on par with or better than the original integer discrete flow model.
INTRODUCTION
Density estimation algorithms that minimize the cross entropy between a data distribution and 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 & Rezende & Mohamed 2015) have seen practical connections to lossless compression
Welling 2013
for some time, normalizing flows were only recently used for lossless compression. Most normalizing flow models are designed for real-valued data, which complicates efficient connection with entropy coders for lossless compression since entropy coders require discretized data. However, normalizing flows for real-valued data were recently connected to bits- coding by Ho et al. (2019b), opening up the possibility for efficient dataset compression with high compression rates. Orthogonal to this, Tran al. (2019) and [Hoogehoom al. (2019)) introduced normalizing flows for discrete random variables. Hoogeboom et al.] (2019a) demonstrated that integer discrete flows can be connected directly entropy coders without the need for bits-back coding.
In this paper we aim to improve integer discrete flows 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 specifically, we start by discussing the claim on the flexibility of normalizing flows for discrete random variables by (Papamakarios (019). and we show that this limitation on flexibility does not apply to integer discrete flows. We then continue by discussing the potential influence of gradient bias on the training of integer discrete flows. We demonstrate that other less-biased gradient estimators do not improve final 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 finite difference estimates of the gradient. We also demonstrate that the previously observed performance degradation as a function of number of flows is highly dependent the architecture of the coupling layers. Motivated by this last finding, we introduce several architecture changes that improve the performance this model class on lossless image compression.
al.]

---

Published as a conference paper at ICLR 2021
2 RELATED WORK
Continuous Generative Models: Continuous generative flow-based models (Chen & Copinath
2001 20nh 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 al. (219)) We refer to Papamakarios et al. 2019) for a recent comprehensive review of the field.
By discretizing the continuous latent vectors of variational auto-encoders and flow-based models, efficient lossless compression can be achieved using bits-back coding (Hinton & Van Camp 1993). Recent examples of such approaches
Local Bits-Back Coding with normalizing flows (Ho et al.. 2019b) and variational auto-encoders with bits-back coding such as Bits-Back with ANS (Townsend
al.] 2019)), 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 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 efficiently. 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 al., 22200) by storing 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. (2020) first encode an image using an efficient lossy compression algorithm and store the residual using generative model conditioned on the lossy image encoding.
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.
Hand-designed Lossless Compression Codecs:
3 BACKGROUND: NORMALIZING FLOWS
In this section we briefly review normalizing flows for real-valued and discrete random variables. A normalizing flow consists of sequence of invertible functions applied to a random variable f o fk valued random variable æ E Ed with unknown distribution Pa (z). Let f Rd + E2 be an invertible function that such that y f(x) with E Rd. If we impose a density Py, (y) on the distribution Pa (a: is obtained by marginalizing out y from the joint distribution Pa.y (x, y) Paly (y)p, (t):
(xx), yielding random variables yx < + y² < y = x. First, consider a real-
Pa Pm( = f0000 - (1))) = f
u u) (f (f()) det af(w)
du = Py p=({())) det 0f(x) du
Öx
where we used Pmly (xly) = 5(z f- with 8(x x²) the Dirac delta distribution, and we applied a change of variables. Repeated application of (T) for a sequence of transformations then yields the log- bog-probability:
K
k ay
Oy*
In (x) In x x (y K
+
Σ
In det
By parameterizing the invertible functions with invertible neural networks and by choosing a tractable distribution py, (y*) these models can be used to optimize the likeliildoood of æ. When modeling discrete data with continuous flow models, dequantization noise must be added to the input data to ensure that a lower bound to the discrete log- likelihodood is optimized (Uria et al., 2013 Theis et al. 2015 HH al. 2019a).
2

---

Published as a conference paper at ICLR 2021
3.1 DISCRETE NORMALIZING FLOWS
Next, consider a to be a discrete random variable with domain X' and define the invertible function f: X' + X. In general ʃ can be mapping between two different domains, but for our discussions will be sufficient to consider single domain. The marginal probability mass of x is given by
(x)
Σ yEA
Paly (a) (zly),,(() = Σ 6 Pa (f(x)) yEX
(3)
with the Kronecker delta function d = if and otherwise. Note the absence of a volume correction in the form of Jacobian determinant owing to the fact that probability mass functions only have support on a discrete set of points.
Recently. Tran et al. (2019) and Hoogeboom et al. (2019a) have both considered normalizing flows for discrete random variables. integer discrete flows (IDF) by HHogeboom al.] (2019a) the random variables are assumed to be integers, i.e. X' zd. The main building block of IDF is an
additive bipartite coupling layer (Dinh et al. 2016):
ya Vr
Fa
(4)
I + to (xx))
Here Ya E ETM, y/ nd are obtained by splitting y E 2d into two pieces such that m n d, and
similarly for I, and Zg. The pre-quantized translation (-) is represented by the output of a neural network with learnable parameters 9. and is rounded to integer values with the rounding operator 1-].
The parameters 0 are optimized with straight-through estimator and tradient-based optimizer.
Tran et al. (019) introduce flows for non-ordinal discrete random values with finite number of
possible values: æ E X = {0, I..,, K - 1]d. They introduce an autoregressive and a coupling bijector, with the coupling layer given by yya y/1 = [x:,, (2:a) o IG tx(a))) modK], with denoting element-wise multiplication and 8 E {1, 2..., K 1} and E 10, K - 11 and s and K constrained to be co-prime. Gradients are again computed with a straight-throups estimator.
One of the main differences between discrete flows (Tran et al. 2019) and integer discrete flows (Hoogeboom et al. 2019a) is that the former treats the random variable as having a finite number of classes while the latter considers countably infinite number of classes. In the next section will see that this difference plays an important role in the theoretical flexibility of invertible flows for discrete random variables.
4 CAN INTEGER DISCRETE FLOWS FACTORIZE ANY DISTRIBUTION?
In this section we discuss the theoretical flexibility of normalizing flows for discrete random variables. In particular, we will focus on the ability to map from a distribution with dependencies across all dimensions distribution which is fully factorized (i.e. independent across all dimensions). As generative flow models require a tractable base distribution with efficient sampling, simple base distributions (e.g. Gaussian, logistic, categorical) which are independent across dimensionsssubpixei. are frequently used. Paaamakarios et al.] (2019) state that invertible flows for discrete random variables cannot map all distributions as compared to flows 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 larger set of possible values. Since integer discrete flows embed the data into the integer lattice Zd, this model class naturally does not suffer from a limited factorization capacity.
factorized base distribution, and name this limitation
The starting point for this discussion lies in the observation by Papamakarios et al. (2019) that invertible normalizing flows 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 + X, then there is always exactly one y. such that y f() and Pw (=(x) (y) In contrast, non-volume preserving normalizing flows for real-valued random variables can increase or decrease the density through the Jacobian determinant in the change of variables formula in (I).
Papamakarios et al. 2019 then discuss an educative example to show that this permutation property can lead to limited ability to factorize distributions. Consider the case of doodemnenioonal random
3

---

Published as a conference paper at ICLR 2021
p(21, 02, F3s)
p(u,, U2. Us)
ply1, , U/, ys)
p(y1, y2, Us)
233
y3
ys
I22
y/2
y2
1/2
'1
9/1
Figure I: Left: 3D probability distribution tensor, only nonzero values are indicated with colored
cubes, all empty space is assumed be filled with zero-valued cubes. Middle left: an example of an additive transformation conditioned on 23: // = [4 [d2)]. 1/2 + Lla (2za] Va 2/3- Middle right: an example of an additive transformation conditioned on I and æg: 3/1 21,
9/2 I2 + It (z1, (21,23)] æ3)] l, V/ Tz. Right: distribution tensor that single additive transformation of cannot generate from the cube on the left.
the form
variable a = (x1, 22), with T1, E {0, 1}, and a data distribution given by
21 N22
0
0
0.1
0.3
Pn (2111,20
(5)
0.2
0.4
In order to map this probability distribution to independent base distribution, the corresponding probability matrix must be of rank one. other words, given a sequence of functions f x+ with X' {0, 1) " the probability distribution of the random variable y y f f
(x) should be represented by matrix with linearly dependent columns or rows. This allows for factorization into an outer product of two vectors that represent the independent base distributions for 91 and y/. Since discrete flows can only permute probability mass tensors, the matrix corresponding to Py (11, 3/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, Paapamakarios et al. (2019) conclude that discrete normalizing flows cannot map any distribution a factorized base distribution.
However, one of the key assumptions made above is that the domain of f is restricted to X' = {0, 1)2. By extending to a larger number of classes this example can in fact be factorized by discrete normalizing flow. More concretely, let us extend the domain to X' = {0, 1, 2, 3)3. The probability distribution matrix of æ is shown below, together with permutation of rank 1.
0.1
0.3
0
0.1
0
0
0
70.1
0.2
0.4
0
0.2
0
0.2
0.3 0 (6)
Pa (11, x2):
Py (11, 02)
0
0
0
0
0.3
0
0
0
0
0.4
0.4/
This example illustrates that the number of classes that considered valid for the discrete random variables plays a crucial role in the flexibility of discrete flows. Here, we claim that this holds more generally through the following two lemmas. First show that by embedding the data into space with more possible values than present in the data itself, one can always construct an invertible mapping one-dimensional variable embedded in d dimensions, and that the resulting variable has distribution which is trivially factorized across all d dimensions.
(s1,...,za) with > and x, {O..... k(O) - 1}, distributed according to an arbitrary distribution Pa (x), can be transformed to one-dimensional random variable y t0,..., (IIL,, K() - with a number of classes that scales exponentially with the dimension d through bijective mapping f. Embedding in d dimensions through padding with zeros, iee. y = (y, O, O,..., 0)²,, the distribution over y is trivially factorized across all dimensions: (x( Py (f())x)) Il Sy.o with 8t,5 the Kronecker delta distribution for all tero-padded dimensions.
Lemma A d-dimensional discrete random variable
The closed form of the mapping is == f(x) 21 (1) x, ((0)
+ 1
such that y {O..... ( x()) x()) 13. See Appendix for the full proof. Intuitively, given enough classes, factorized base distribution is obtained by "Tattening" the hypercube that contains
4

---

Published as a conference paper at ICLR 2021
Pemprical
Pameria 0i
P .44
Punta (2)
+
003
-1
02
2 3.3 2.1 2.3 2.
24
2-
2.
0.
Figure 2: Visualization of an IDF that has learned factorize the probability distribution of the toy example in (5). Left: empirical densities of the data, the data transformed by one additive bijector, and the data transformed by two additive bijectors. Right: similar to the left plot, but with data sampled from the model.
all nonzero entries in the data distribution tensor (see left panel Figure T into one dimension. The following lemma demonstrates that additive coupling layers as in of integer discrete flows can model such an invertible mapping.
Lemma 2. For d-dimensional random variable translation operations of the form Ze Ta. Pb = Tb + [(xx)), with a and b denoting the indices of two splits of the variable æ, are sufficient to map in an invertible manner to random variable y = (y1, .,.... which is one-dimensional variable embedded in d dimensions.
The formal proof an induction on d and relies on 1) the fact that integer discrete flows embed the discrete random variables into the integer lattice X' 2d, which always has sufficient number of classes (countably infinite), and 2) the observation that two translations are sufficient to map two dimensions into single dimension with more classes, illustration, Figure depicts the type of operations that single additive bijector can and cannot perform a dinee-dimeniional 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 first three panels correspond the data distribution and the intermediate and factorized final distributions produced by the first and second coupling layer respectively, For architecture details see Appendix D.3
Appendix A For more intuitive
As pointed out by Paapaaakaiios al. (2019), a discrete flow can only map uniform base distribution into a uniform data distribution due to its permutation property. Therefore, invertible discrete flows require ane-dimensional learnable distributions in order to be able to model all data distributions. Note also that discrete flows as introduced by Tran et al. (2019) are designed for a finite number of classes in X' and therefore cannot factorize distributions. Although not implemented in the original work, this model class can in principle also be extended to embed the data into space with larger finite 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 integers (with countably infinite number of classes) does not have limited factorization capacity. We therefore explore other directions for potential improvements.
DOES GRADIENT BIAS HINDER OPTIMIZATION OF IDFs?
In this section we investigate whether gradient bias is problem when training IDF models. Hooge boom et al. (2019a) demonstrated that the performance 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 traiigh-tthrough estimator was suggested to be the cause.
In our experiments we found that neither stochastic rounding nor replacing the identity function the straight- through estimator with soft approximation of the rounding function improved the results. To disentangle the effect of gradient bias of the straight-through estimator and the difficulty of discrete optimization, we have trained IDFs and their continuous additive counterpart (without the rounding operator and without a quantized base distribution) on CIFAR-IO. 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 dy was modeled with a superposition of scaled sigmoid functions and temperature T that allows to interpolate between the identity function (T = 1) and the rounding function = 0, see Appendix CerendixCffor. more details. The left plot in Figure 3 shows that the continuous model with identity straight-through
5

---

Published as a conference paper at ICLR 2021
3.50
cont: 07,
3.8
contidisc Conv
cont: 07.
contidisc ResNet
3.45
cont: mdd. id
3.7
contidisc DenseNet
cont: id. id
consddisc NetseNet
disc: id,
3.6
3.40
disc: mdd, id
3.5
3.35
3.4
3.30
3.3
0.00
0.25
0.50
0.75
1.00
2
T
flows per level
Figure 3: Left: Bits per dimension for models trained for 2000 epochs on CIFAR-IO. Continuous (cont) and discrete (disc) models are trained with different combinations of translations activation functions (-+) and backward (+-) substitute functions for the ttraightthhoough estimators: soft rounding function (OT), a hard rounding function (rnd) and the identity function (id). or interpolates between the identity function at 7 0 and the hard rounding function at T 1. Right: Bits per dimension of IDF models trained for 300K iterations (#960 epochs) with different architectures for the translation networks the coupling layers: convolutional neural networks, ResNets and DenseNets. The models have 3 levels. The discrete ResNet model for 10 flows per level was unstable so the result for this model omitted. The models with the DenseNet architecture correspond the best models used by [Hoogeboom et al. (2019a). The ConvNets and ResNets have approximately equal number of parameters and depth as the DenseNets. Note that the ConvNets used to produce the results in Figure 5 of (Hoogeboom al. 2019b) are much shallower. The DenseNet++ architectures correspond to the proposed IDF++ model of Section
gradient estimator consistently outperforms the continuous model without gradient bias. The discrete model clearly also does not benefit from 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-throuu 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 significance of gradient bias in more detail dimensional toy example of Section 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 finite differences as an approximate gradient vector gdd with elements gf = (L(0, e, /,) L(0,- e, /,)))/2a /2e. Here is the loss function averaged over single batch of data. For additive continuous flow models, gdd will approach the true loss gradient Vgl. as c + 0. For discrete models, ggd can be thought of as a linear approximation of the loss landscape in a trust-region of radius e around the current parameter vector 0.
extension of the two-
We compare continuous flow models that are trained using the unbiased gradient VoL with discrete flow models that are trained using the trraihhtthhrough gradient estimator ga We estimate the agreement between the (approximate) gradient and the finite difference approximation gd for varying trust-region size e 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 gdd ggd//lg/ which can be interpreted as the uncentered correlation between the elements of g and gdd As long reduce as the agreement is consistently positive, performing gradient descent with g is expected the loss according to the trust-region approximation based on gdd. 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 -ifieeeiffeeeenc gradient and the throughh gradient of the discrete model and the agreement with the real gradient for the continuous model, at initial- ization and after 3000 training iterations. The straight-through gradient estimator for 8-bit data is always positively correlated with small but finite difference estimates, corroborating our results on CIFAR-10 10 that the identity straight-through gradient estimator allows for good optimization. For 1-b data the quality of the straight- trriigh-tthroug gradient estimator clearly deteriorates and the correlation becomes zero or even negative. This is not surprising the random variables and translations are in practice modeled on the rescaled grid Z /its leading to more coarse- coarse-grained rounding for lower bits. As lossless compression mostly deals with source data in the higher bit ranges the poor performance

---

Published as a conference paper at ICLR 2021
After 10 iterations
After 3000 iterations
uncentered conclation
1.0
1.0
continuous, bit
discrete, hit
continuous, bits discrete, bits
0.0
-0.5
0.5
10
10
10
10
10
10
10
10
10
finite difference step size
finite difference step size
Figure 4: Top two panels: Agreement between the approximate (true) gradient and finite difference gradient of the discrete (continuous) model at model initialization (left), and after 3000 training iterations (right). The curves were obtained by averaging over 10 batches. The higher variance for the 8-bit case is due to the larger number of different datapoints, leading to more variance in the data batches. Bottom two panels: Visualization of the loss surface (left) and trajectory the optimizer on the loss surface (right) for the discrete model on the 1-bit toy example, following (Li et al. 2018).
for lower bit data is less relevant. The bottom two panels Figure show the loss surface and optimization trajectory of discrete model trained on the 1-bit data, suggesting that the optimization landscape is hard to traverse for low bit-depths. See Appendix Ba 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 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.] (2166) 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 problem than previously suggested. Alternative gradient estimators with less bias do improve the results, and our analysis on the toy examples shows that for data with large bit depths the straight- through gradient estimates have positive correlation with finite 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 lossless image compression.
estimator is less of
IMPROVING INTEGER DISCRETE FLOWS FOR LOSSLESS COMPRESSION
In this section we introduce changes the architecture of IDF that improve its performance. We briefly summarize the IDF architecture, for more detailed information see Appendix D Similar to other generative flow models, the architecture of IDF contains T. levels. Each level / consists of space-to-depth transformation followed by a sequence of K alternating channel permutations and additive coupling layers. For levels 1, ..., L 1 the output random variable is split into two pieces, with the second part serving as an input the next level: [2(1), yy(") f(.K) ... f((.)) (y (II),,, with (0) For the last level simply have z(4) ((L.K) o... ((L.)) (y(L-1)). The combined

---

Published as a conference paper at ICLR 2021
= [=((1) ,..., (1)] denotes the latent representation of x. The distribution of z for 3 levels is which is equivalent to p(=(9)) 1=*" ). The conditional distributions are modeled as discretized logistic distributions and the unconditional p(z(3) is modeled with mixture of discretized logistics with 5 components. Both the conditioning networks and the pre-quantized translations to in (4) are
then factorized p(z) p(z(1) (2), x( =
(p())
(p(2())
modeled with DenseNets (Huang et al., 2017).
The first modification 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 influences the modeling of the conditional distributions y(2(")pp ()).. By inverting the permutations one ensures that the split into y and 2(1) happens along the channel direction of the space-to-depth transformed version of 1) as such y
and (1) retain the spatial correlation structure of the original image presumably making it easier to model the conditional distribution p( ((2)()) ()) (see visualization in Figure 10 of Appendix D).
The second change is an adaptation of the rezero trick by [Bachlechner et al. (2020). The additive bijectors of 4 are replaced by [ya 3/b [.. 26 + [att((x,)]] 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 μ== u log s = 8 log a with [v. log ø = DenseNet and y and 5 learnable scalar parameters initialized to zero
The third modification consists of an alteration in the dense blocks that make up the translation and logistic conditioning DenseNets by introducing group normalization layers (Wu He| (018) and switching from ReLU activations (Nair& & Hinton| 2010) to Swish activations Raaacchandann et al., 2017):
IDF: Convixl + ReLU + Conv3x3 + ReLU
IDF++: Convixi - GroupNorm - Swish - Conv3x3 - GroupNorm - Swish
The trainability of deep neural networks such as flow models is strongly dependent on careful initialization and normalization. Identity mapping initialization of flow 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 modifications 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 modifications (DenseNet++)) on the validation set (consisting of 20% of the training set) on CIFAR- as a function of flows 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 flows per level due to overfitting (see Figure [T]. These results also show that the modifications allow for more efficient use of the flow layers: for 300K iterations the performance of the IDF++ model with DenseNet++ architecture is on par with that of the IDF model of HHoogeboom al. (2019a) with DenseNet coupling layers.
To further corroborate this point, we have trained an IDF++ model with 4 and 8 flows per level until convergence for CIFAR-10, ImageNet-32 and ImageNet-64, and compared it with the baseline IDF model with 8 flows per level and other related work in Table T To make a fair comparison against other methods like local bits-back coding (LBB) by Ho al. (2019b) we train our final models on the entire training set without holding out part of the training set as validation set. Although there no visible difference for the ImageNet-32 and ImageNet-64 datasets, more data does matter for CIFAR- 10: it reduces the model's 's average negative log- ligeliidl in units of bits per dimension (bpd) on the test set from 3.32 (as reported by HHoogbooo et al. (2019a)) to 3.30 for the baseline IDF model. Table [ displays the bpd resulting from compressing the latent representations with a range-basee 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 shown in separate groups. Where available, the bpd as predicted from the models negative log- likelihood is indicated in brackets.
8

---

Published as a conference paper at ICLR 2021
Table Compression results in bits per dimension (bpd) for IDF++, hand-designed codecs and other deep density estimators based on normalizing flows, super resolution and variational auto-encoders. Where available, the bpd according the model's negative log-likelihood is indicated in parenthesis. Results with are taken from Townsend et al. (2019a), and those with are taken from Hoogeboom ct al. (2019a). The IDF CIFAR-I0 10 result indicated with ** obtained by our implementation of IDF with 100% of the training data used to train the model. In (Hoogeboom al. 2019a) only 80% of the training data was used to train the best CIFAR-10 model. All other results are from the original papers.
Compression models
CIFAR-10
IMAGENET-32
IMAGENET-64
PNG (Boutell & Lane
1997>
5.87*
6.39*
5.71*
JPEG-2000 (Rabbani
2002)
5.201
6.481
5.10f
FLIF (Sneyers & Wuille (2166)
4.19*
4.52"
4.19*
BIT-SWAP (Kingma et al.
(011))
3.82 (3.78)
4.50 (4.48)
HILLOC (Townsend
(2019a))
3.56 (3.55)
4.20 (4.18)
3.90 (3.89)
LBB (Ho et al. 2019b)
3.12 (3.12)
3.88 (3.87)
3.70 (3.70)
SREC (Cao et al. (2020))
4.29
IDF (Hoogeboom et al. (019aa)
3.32 (3.30)** **
4.18 (4.15)
3.90 (3.90)
IDF++, SMALL: FLOWS PER LEVEL IDF++
3.31 (3.29)
4.16 (4.14)
3.85 (3.85)
3.26 (3.24)
4.12 (4.10)
3.81 (3.81)
The results show that, contrary to less-biased gradient estimators, the proposed architecture modifica- tions in IDF++ improve the performance over IDF for 8 flows per level for all datasets. Furthermore, the smaller IDF++ model with only flows per level either performs on par or better than the IDF baseline with 8 flows per level. This effectively reduces the number of parameters and time required for forward pass through the model by a factor of 2, and demonstrates that the proposed modifications in IDF++ enable more efficient compression without sacrificing performance. For more detailed ablation on the contributions of the proposed modifications see Appendix E
7 CONCLUSION
In this paper we investigated several potential directions for improvement of integer discrete flows 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 flows for discrete random variables cannot factorize all distributions does not apply to integer discrete flows. Furthermore, we demonstrated that the effect of gradient bias on optimization previously thought. Other gradient estimators with less bias do not improve optimization, and a numerical analysis of the direction the trraighttthoouge gradient estimates showed that they correlate well with finite difference estimates for large-bit data. We also showed that the previously observed performance deterioration with increasing depth of the flow 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 modifications lead to a more efficient flow-based compression model, as evidenced by our results that show that an IDF++ model with half the number of flows compared to the baseline IDF model performs on par or better.
less severe than
Although we found that the simple traight-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 well on the datasets we consider, they (de)compress significantly faster and require significantly 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 make deep density estimators more practical for source compression.

---

Published as a conference paper at ICLR 2021
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fermanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg. Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https: /tww//wwwww... www.:tensorffo.oo org/ Software available from tensorflow.org.
Olah,
Eirikur Agustsson and Lucas Theis. Universally quantized neural compression. arXiv:2006.09952. 2020.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv:1607.064500 2016.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W. Cottrell, and Julian McAuley, ReZero All You Need: Fast Convergence at Large Depth. arXiv:2003.04885 2020.
Thomas Boutell and T Lane. PNG (portable network graphic)) specification version 1.0. Network Working Group, 1997.
Sheng Cao, Chao-Yuan Wu, and Philipp Krähenbühl. Lossless Image Compression through Super- Resolution. 200400200 2020.
Scott Saobing Chen and Ramesh A. Gopinath. Gaussianization. In Advances in Neural Information Processing Systems 13, 2001.
Peter Deutsch. DEFLATE compressed data format specification version 1.3, 1996.
Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: Non-linear independent components estimation. arXiv:1410.8516. 2014.
Laurent Dinh, Jascha arXiv:1605.08803, 2016.
and Samy Bengio. Density estimation using real NVP.
Jarek Duda. Asymmetric numeral systems. arXiv:0902.0271. 2009.
Jarek Duda. Asymmetric numeral systems: entropy coding combining speed of huffman coding with compression rate of arithmetic coding. arXiv: 2013.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, 2016.
Geoffrey E Hinton and Drew Van Camp. Keeping the neural networks simple by minimizing the description length the weights. In Proceedings of the sixth annual conference on Computational learning theory, 1993.
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. arXiv:1902.00275 2019a.
Jonathan Ho, Evan Lohn, and Pieter Abbeel. Compression with Flows via Local Bits-Back Coding. In Advances in Neural Information Processing Systems 32, 2019b.
Emiel Hoogeboom, Jom Peters, Rianne van den Berg, and Max Welling. Integer Discrete Flows and Lossless Compression. In Advances in Neural Information Processing Systems 32, 2019a.
Emiel Hoogeboom, Rianne Van Den Berg, and Max Welling. Emerging Convolutions for Generative Normalizing Flows. In Proceedings of the 36th International Conference on Machine Learning. 2019b.
10

---

Published as a conference paper at ICLR 2021
G. Huang, Z. Liu, L. Van Der Maaten, and K. Q. Weinberger. Densely Connected Convolutional Networks. In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2017.
Sungwon Kim, Sang-gil Lee, Jongyoon Song, Jachyeon Kim, and Sungroh Yoon. FloWaveNet: A generative flow for raw audio. arXiv:1811.022555 2018.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv:1412.699 2014.
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv:1312.6114, 2013.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions, In Advances in Neural Information Processing Systems, 2018.
Friso H Kingma, Picter Abbeel, and Jonathan Ho. Bit-swap: Recursive bits bits-back coding for lossless compression with hierarchical latent variables. 1905vv/905.06600 2019.
Manoj Kumar, Mohammad Babacizadeh,, Dumitru Erhan, Chelsea Finn, Sergey Levine, Laurent Dinh, and Durk Kingma. Videoflow: flow-based generative model for video. arXiv:1903.01433 2019.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. In Advances in Neural Information Processing Systems, 2018.
Fabian Mentzer, Firikur Agustsson, Michael Tschannen, Radu Timofte, and Luc Van Gool. Practical full resolution learned lossless image compression. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2019.
Fabian Mentzer, Lue Van Gool, and Michael Tschannen. Learning Better Lossless Compression Using Lossy Compression, arXiv:2003. /0184, 2020.
Vinod Nair and Geoffrey E. Hinton. Rectified Linear Units Improve Restricted Boltzmann Machines. In Proceedings of the 27th International Conference on International Conference on Machine Learning, 2010.
Aaron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. arXiv:1601.06759, 2016.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Laksh- minarayanan. Normalizing Flows for Probabilistic Modeling and Inference. arXiv:1912.02700 2019.
Majid Rabbani. JPEG2000: Image compression fundamentals, standards and practice. Journal of Electronic Imaging, 2002.
Richard Rabbat. WebP. a new image format for the Web, 2010. URL 11/2000009//bbp/ new image- format for web. .tml
Prajit Ramachandran, Barret Zoph, and Quoc V. Le. arXin:1710.05941, 2017.
Searching for Activation Functions.
Danilo Rezende and Shakir Mohamed. Variational Inference with Normalizing Flows, In Proceedings of the 32nd International Conference on Machine Learning, 2015.
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. PixeICNN++: Improving the PixeICNN with discretized logistic mixture likelihood and other modifications. arX1v:1701.05517, 2017.
Jon Sneyers and Pieter Wuille. FLIF: Free lossless image format based on MANIAC compression.
In 2016 IEEE International Conference on Image Processing (ICIP). JEEE, 2016.
Lucas Theis and Matthias Bethge. Generative image modeling using spatial Istms. In Advances in Neural Information Processing Systems, 2015.
11

---

Published as a conference paper at ICLR 2021
Lucas Theis, Aron van den Oord, and Matthias Bethge. A note the evaluation of generative models. arXiv:1511.1884.. 2015.
James Townsend. A tutorial on the range variant of asymmetric numeral systems. arXiv:2001.09186, 2020.
James Townsend, Thomas Bird, Julius Kunze, and David Barber. HiLLoC: Lossless Image Compres- sion with Hierarchical Latent Variable Models. arXiv: 2019a.
James Townsend, Tom Bird, and David Barber. Practical lossless compression with latent variables using bits back coding. arXiv:1901.04866 2019b.
Dustin Tran, Keyon Vafa, Kumar Agrawal, Laurent Dinh, and Ben Poole. Discrete Flows: Invertible Generative Models of Discrete Data. In Advances in Neural Information Processing Systems 32, 2019.
Benigno Uria, lain Murray, and Hugo Larochelle. RNADE:
The real-valued neural autoregressive density-estimator. In Advances in Neural Information Processing Systems 26. 2013.
Benigno Uria, lain Murray, and Hugo Larochelle. deep and tractable density estimator. In International Conference Machine Learning, 2014.
Yuxin Wu and Kaiming He. Group Normalization. Lecture Notes in Computer Science, 2018
12

---

Published as a conference paper at ICLR 2021
Supplementary Material of "IDF++: Analyzing and Improving Integer Discrete Flows for Lossless Compression"
APPENDIX A FLEXIBILITY OF INTEGER DISCRETE FLOWS
The proof of Lemma [ 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
Proof of Lemma 7 The invertible mapping is given by = f(x:) = 21 + x 12 + K K k(
d-1 x0 with maximum value of (IT K()) - and minimum value zero. Given y. each element 22 can be obtained recursively: 2; = (V- -E)= ki-1x;)) //Ki where // refers to integer division. Through the change variables formula we know that Py(v) = (f (y)). which combined with zero padding embed in dimensions leads the desired factorized distribution in d dimensions.
Alternative proof of Lemma 1 An alternative proof is via induction on
Base case: Let start with d 2, have the random variable (21, x2)² with 21 (O, K(1) - 1} and 22 E {O. (2) -1.. x is distributed according Pa (2). The following function can map this random variable to a one-dimensional random variable with y E {O....,,ku ((2) - -}: y = x1 + (1)22 Given y, # and z2 can be recovered through
x2 = (//K(2) and æ = y - (()))2
General case: Now let us assume that the lemma holds for d = and examine the case of n + x (I1, In, "n+1)² with x; E {O,... k(t) J. Then the following transformation
maps to n-dimensional random variable y (y1.... y,));
I; for i {l,.... ................................ n 1}
and V/- = In + K(m) In+1. E can be recovered from y through and Tn+1 = yn/ K, In yn KIn+1- We have .1...., Un-1 E (O. K()) 1} and yn {0, ..., Kt (- - 1I. As we assumed that we can transform an n-dimensional random variable
y/ for i = (1, ................................ n 1}
to one-dimensional random variable in an invertible manner, this concludes that we can also do so for the d = case.
Proof of Lemma 2 We will prove this via induction on d.
Base case: Let us start again with d = 2. x = (11,TT) with 2, E (o...., KO} can be transformed
into random variable y (31. 0)7 in an invertible manner with two translation operations of
the form 2n = Las 24 = I6 + [t(x,)], where either b) = (1,2) or (a.b) = (2, 1). The
first translation is given by 21 = 01 + ((() r2, 22 = I2, followed by y1 = 21 = 21 + K()22
1/2 = 22 - (x) K()) /2)/ (L) = x2 22 O.
General case: Let us assume translations are sufficient for d = We can then transform d = + dimensional variable x = (æI,..., In, In+1)² )" an invertible manner to a random variable that is effectively n-dimensional:: y (yı,.... Ya. 0) with two translations operations: The first translation is:
= r; for fI...., 1, +
2n = zn + (n) xn+1,
followed by the second translation:
Y/ = for i E {l,.....
Vn+1 -n+1 zn/ (m) In+1 (In, + knn) Nn+ (m))+++)///3
3ntl Fn41 0.
Which results in y/ 2, for E (1,..., n 1} and yn = 2n, + (1))) En+l- As we assumed integer translations were sufficient to map n-dimensional variable like (y1...., T to one-dimensional random variable embedded in n dimensions in an invertible manner, we have now proven that this also holds for d 1.
13

---

Published as a conference paper at ICLR 2021
APPENDIX B LEARNING TOY EXAMPLES
B.1 GRADIENT BIAS ON THE TOY EXAMPLE
The toy example that is discussed in Section 4 has probability distribution with nonzero probabilities for x, E [0. With only two values per dimension with nonzero probability the data is effectively 1-bit quantized. Here higher bit-depth. This can done by modeling the probability masses with log-lineart spaced logits in the interval (0, 1). The resulting distributions for several bit-depths are shown in Figure
extension of this toy example for input data that is quantized to
1 bit
2 bits
4 bits
8 bits
22
0-
52
4-
8-
150
192
224
256
-2
2
-2 -1 0
12 16
32 64
96
125 160
192224 256
Figure 5: probability distribution of the extension of the toy example of Eq.
( to more bits.
B.2 LOSS 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 loss landscape that is affected by quantization operators can 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 0" we choose two direction vectors 0 and 0₂ and plot the value of f(a, 9) L(0 + a0, 302). To choose direction vectors, we use model parameters at different stages of training. Let O be model parameters after iteration i. Given the set of parameters for iterations, apply PCA to the matrix M ( , 0 -- 0%,] and select the two most explanatory directions as 0, and 02.
The results for the 1-bit toy example are shown in the left two panels of Figure for the discrete flow model. Figure 6 shows similar plots for the continuous case, demonstrating a much smoother loss landscape. Figure suggests that discrete models with unlucky initializations local minimum from which it is hard to escape due to sharp cliffs in the loss landscape. Figures 7 and 8 show significantly 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 less of an issue higher bit-depths.
easily end up in
Figure 6: Visualization of the trajectory of the optimization on the loss surface Li et al.
2018 for
the continuous model for the toy example with 1-bit data.
14

---

2018 for Published as a conference paper at ICLR 2021
Figure 7: Visualization of the trajectory of the optimization on the loss surface (Li al. 2018 for the discrete model for the toy example with 8-bit data.
Figure 8: Visualization of the trajectory of the optimization on the loss surface Li et al. the continuous model for the toy example with 8-bit data.
APPENDIX C SOFT ROUNDING FUNCTION
The soft rounding function or that is used in Section 5 is given
by
( [x]] +
(x = [x]
(9)
The soft rounding function is depicted in Figure 9 for different temperatures.
L-1
0r.7=1.0
0,.8=0.8
0..60..6
or 70.4
or. 7.3.3
o7, 7=0.2
or, 0=1.1
-1
Figure 9: Soft rounding function or for different temperatures T. In the limit T the soft rounding function reduces to the hard rounding function [:]. At the soft rounding function is indistinguishable from the identity function.
15

---

Published as a conference paper at ICLR 2021
APPENDIX D ARCHITECTURES
D.I ARCHITECTURE 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 specifically, the architecture of level of IDF is built up as follows:
x
y (1-1) - space-to-depth + [Permute + Additive transform] - [z(0) y (0).
Except for the last level, the output random variable is split into two equal halves (2 y ( where only ( is transformed by the next level. Furthermore, y æ and [x(1), ................................ 12 denotes the latent representation of æ. Inside the additive transformations of (4) the prequantized translations modeled using DenseNets (Huang al., 2017) with dense blocks of the following structure:
Dense block: Convix x - ReLU - Conv3x3 - ReLU.
All DenseNets have 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) = (1z))p(z(1)) p(=(4)))))))))) (1)1
(1)), which is equivalent to (((4))))((((-))) { .....(()) (1). All conditional distributions are modeled with discretized logistic distributions. The unconditional distribution p(z(2) is mod- eled with mixture of discretized logistics with five 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: [v, log oj DenseNett, Note that instead of modeling the random variables as integers (x zd), they are modeled as discrete random variables on a grid with bin-width 1/256 (x 2n /256).
(2
The model is trained with the Adamax optimizer (Kingma Ba 2014 with an exponential learning rate schedule with base learning rate equal to 1 x 10 and a linear warmup phase of 10 epochs. See Table for more details on the learning rate decay, the number of levels, the batch size and the number of epochs used for training.
Table Architecture and training settings for IDF. Table adapted from Hoogeboom et al. 2019a).
Note that we used batch size of 128 for CIFAR-10 instead of 256 used in the original work, while still reproducing the same number as reported by [Hoogbooom et al. (2019a). For ImageNet-32 and ImageNet-64 the batch sizes the same as in the original work (Hoogeboom et al. 2019a).
Dataset
Levels L
Batch size
Learning rate decay
Epochs
CIFAR-10
128
0.999
2000
ImageNet-32
256
0.99
100
ImageNet-64
64
0.99
10
Range-based Asymmetric Numerical Systems (rANS)
(Dua 2009
2013 Townsend 2020) is
used for lossless compression of the latent variables 2
D) by using the probability
distribution p(z) p(z(4)) ((( level structure.
p(= (1) y (1) corresponding to the model's multi-
D.2 ARCHITECTURE AND TRAINING DETAILS OF IDF++
In IDF++, each of the K blocks containing permutation and additive bijector has an additional inverse channel permutation ensure that the output random variable has the same spatial and channel ordering as the input random variable:
x
(1-1) space-to-depth
PPermute
Additive transform
Inverse permute) -220),00)) (1.
16

---

Published as a conference paper at ICLR 2021
The inversion of each permutation inside a level ensures that the splitting at the end of each level done along the channel dimension tional distributions to be able to the spatial correlation between pixels. See Figure [0] for visualization. The dense blocks of the DenseNets for the prequantized translations and the condi- tional discretized logistics have additional group normalization layers (Wu & Hel (018 and Swish
space-to-dephe transformed image, enabling the condi-
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 determined follows: if the number of channels of the group normalization layer is divisible by then 3 groups are used, else it is divisible by then 2 groups are used, and finally if is neither divisible by 2 then a single group is used.
The additive bijectors in (4) are adjusted to include a learnable scalar parameter that ensures initial- ization the identity operator, similar to the rezero trick by Bachlechner al. (2020)
Ya
Za [t=))
(10)
yb
Here is a learnable scalar parameter initialized to zero. The mean and log scale of the conditional discretized logistic distributions are parameterized as l yv, log s log with [v, log o DenseNeta (y)) Y and 8 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 a larger base learning rate of 2 10-3 the exponential decayed learning rate schedule. For CIFAR-10 10, the IDF++
(a)
5
6
8
10
12
13
14
16
(b)
11
10
12
channel dimension
13
14
16
(c)
channel dimension
13
14
15
Figure 10: Schematic visualization of the effect of inverse permutations on the IDF+ architecture. In multi-level IDF+ model the original image (a) undergoes a space-to-depth transformation (b), after which at the end of the first level half of the dimensions (blue tones) are factored out and their conditional distribution based on the remaining dimension (green tones) 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 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 preserved at all levels. Specifically, 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 better results.
17

---

Published as a conference paper at ICLR 2021
model with 8 flows per level trained for 1400 epochs instead of 2000 epochs as overfitting was observed for more epochs. We evaluated all variations of IDF with an exponential moving average of the trainable parameters with decay rate of 0.9999.
D.3 ARCHITECTURE FOR TOY EXAMPIE MODELS
The models for the toy examples in Section and have single level with an unconditional prior given by mixture of logistics with five components. The models have two coupling layers with DenseNets as translation neural networks. The dense blocks have the following structure
Dense block: Conv1x1 >> LayerNorm - ReLU _ Conv1x1 - LayerNorm -+ ReLU.
We found the inclusion of LayerNorm (Ba et al. 2016) to help make the model less sensitive to initialization. Because DenseNets these models are fairly shallow we did not see a performance difference between LayerNorm and GroupNorm. The DenseNets have depth of 4 blocks and 32 channels. The datapoints are treated as images of shape h x x x 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 reversed allow for alternating conditioning. The learning rates for the prior parameters and coupling layers are set to x 10 and 10 respectively and we used a batchsize of 128.
D.4 DATASETS
The training set of CIFAR-10 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 validation set that held out from the respective training sets for model devel- opment, we trained our final models on the entire training set for each dataset. The effect this is discussed in Section E. We follow Hoogebooo al. (2019a) and augment the CIFAR-10 R- dataset with horizontal flipping, reflect-padding and random cropping during training. No augmentation is used for ImageNet-32 and ImageNet-64.
D.5 HARDWARE AND SOFTWARE
We implemented our models in TensorFlow (Abadi et al. 2015). All experiments were run with NVIDIA V100 GPUs.
APPENDIX E ABLATION STUDIES
Table shows the contributions of the modifications proposed in Section IDF++. As Hoogeboom al. 2019a) only use 80% of the training data for CIFAR-I0 10, we also show the influence of using the entire train set for training, as opposed to keeping 20% held out as 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-33 and ImageNet-66 64 was much smaller and we noticed no significant change performance when training on the entire training set for these
Table Ablation study the modifications of IDF for the CIFAR-10 dataset. The second column indicates the percentage of the dataset's original train set that was used to train the model. EMA exponential moving average. The reported numbers are theoretically achievable bits per dimension.
IDF variations
% Train data
CIFAR- 10
IDF
80%
3.322
IDF
100%
3.298
IDF EMA
100%
3.291
IDF + EMA + rezero + Ir X 2
100%
3.262
IDF EMA + rezero Ir x 2 + invert perm
100%
3.255
IDF EMA + rezero Ir x 2 invert perm + groupnorm (IDF++)
100%
3.241
18

---

Published as a conference paper at ICLR 2021
3.50
cont DenscNet++:: fpl
disc DenseNet++: 8
3.45
cont DenseNet++:: 10 fpl
disc DenseNet+:: fpl
3.40
bpd
3.35
3.30
3.25
0k
100k
200k
300k
400k
steps
Figure 11: The performance of IDF++ models and their continuous counterparts the validation for 8 and 10 flows per level (fpl). The models are trained on 80% of the training data.
datasets. All models without group normalization were trained for 2000 epochs. The final IDF++ model, which includes group normalization layers, was trained for 1400 epochs to avoid overfitting.
In Figure the validation bpd shown IDF++ models and their continuous counterparts for and 10 flows per level. The plot shows that both the continuous and the discrete model with 10 flows per level start to overfit.
19


Reviews:
[{"review_id": "aJzsnzOfdK5", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper aims to analyze and improve IDFs for lossless compression. The authors claim is mainly three following points:\n1. IDFs treat the random variable as integers with a countably infinite number of classes, so it does not have a limited factorization capacity. \n2. Hoogeboom et al. (2019) demonstrated that the performance of IDFs can deteriorate when the number of coupling layers is too large, gradient bias induced by the straight-through estimator (STE) was suggested to be the cause. In this paper, authors claim that gradient bias due to the STE is less of a problem, and it highly depends on architectural choice of coupling layer. \n3. Some architectural changes (invert the channel permutations, rezero trick, group normalization) can get better results.\n\nPros: \n- It is interesting that several types of quantizing functions are evaluated (Fig. 3) and the results shows that STE is less of a problem.\n- Architectural changes seems make better results than conventional IDFs (Table 1).\n- Authors shows theoretical background of IDFs and theoretical analysis and characteristics are well written.\n\nCons: \n- Hoogeboom et al. (2019) shows adding more flow layers than 16 coupling layers per level hurts performance as is depicted (Figure 5 at Hoogeboom et al. (2019)). But in this paper, Fig. 3 shows only 8 coupling layers experimental results is shown, so I think it is little weak as a rationale for above claim 2..\n- In contrast to the author's finding (It can be deeper depending on the architecture) in Chapter 5, IDF++ optimization in Chapter 6 was seems to relatively close to detailed tuning, and relevance is a little weak. It seems more convincing if you examine the effect of deepening the coupling layers with IDF++.", "extra_scores": {}, "extra_text": {}}, {"review_id": "YmY5jVvVBDw", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Update after author response**\nI thank the reviewers for their response.  I appreciated the more careful discussion of the discrete claims (as other reviewers also noted). I also appreciated the efforts to give more justification for your architecture changes.  While the actual experimental results still seem incremental in terms of raw performance, I think the other contributions of the paper are solid and worth having in the literature.  Thus, I've updated my score.\n\n**Summary**\nThis paper focuses on lossless compression using integer discrete flows (IDF).  The paper explores the theoretical flexibility of IDF models showing that (with a countably infinite number of possible values for one dimension) any distribution can be factorized.  The paper also explores the bias in the gradient estimates and gives evidence that architecture seems to be more important than unbiased gradient estimators.  The paper proposes a few architecture changes including inverting channels, initializing transform to identity, initializing base distribution to standard logistic distribution, and using GroupNorm and Swish layers.  Finally, the paper shows better compression results than the original IDF model across CIFAR-10, ImageNet-32 and ImageNet-64.\n\n**Strengths:**\n- Investigates the potential gradient bias issue of IDFs and demonstrates that this may not be the core issue. (a very nice observation)\n- Suggests architecture changes and shows improved performance compared to baseline IDF.\n- Nice illustrations of discrete transformations.\n\n**Weaknesses:**\n- Claims about \"refuting\" the claim in [Papamakarios et al. 2018] seem too strong. It seems that Papamakarios et al. [2018] had a restricted case in mind of finite discrete values. And the point still seems interesting for finite discrete values (e.g., binary data).  The first theoretical result in this paper (Lemma 1) seems to be a still interesting but an almost obvious result that if you allow for infinite number of discrete values you can factorize any distribution (i.e., by just putting all possible non-zero configurations along one dimension). Again, while it is interesting and useful to discuss this, a more nuanced discussion is likely needed.  And this should be seen not as refuting the claim but as showing that relaxing one of the assumptions (i.e., finite to countably infinite values) allows for the potential of an arbitrarily flexible distribution.  There should also be a discussion that this transformation could inherently be bad for estimation since it would require the estimation of an exponential number of values.  Thus, the theoretical result may not be practically useful.\n\n- Architecture changes seem to lack clear motivation and insight. While I appreciated the motivation that the limitation may not be gradient bias, I did not understand the reasons for the architecture changes.  Insights into why these architecture changes lead to an improvement in performance are critical.  Why is Swish better than ReLU in this context?  Why is GroupNorm really needed?  Without deeper insights, it is challenging to build off of this work or generalize this work to other contexts.\n\n- The results seem incremental.  The results, while better than IDF, seem relatively incremental (-0.08,-0.06,-0.09).  I appreciated the context of other compression models but it seemed that IDF++ only marginally improves over the baseline IDF.\n\n**Other comments or questions**\n- The gradient-bias experimental section is a bit hard to parse.  A rewrite for clarity and simplicity would help the paper.\n\n- Wouldn't finite difference be very expensive to compute?  This would require $O(|\\theta|)$ number of model evaluations compared to $O(1)$ model evaluations for gradient estimates. Was this experiment only on a very small dataset and model?\n\n- Why does inverting the channels help?  This still seems unclear to me since it seems to be a trivial modification though I'm probably missing something.\n\n- What does \"We furthermore use an exponential moving average for evaluation\" mean? Is this related to GroupNorm or  to calculating the log likelihood/BPD?\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "MkVf6c-ytBW", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#### DESCRIPTION\nThis paper (i) notes that discrete flows are not limited in representation ability when the discrete data is embedded in the integer lattice, (ii) demonstrates that bias due to the straight-through gradient estimator is not as severe as previously thought, and (iii) investigates some architectural tweaks for integer discrete flows which lead to better performance.  \n\n#### DISCUSSION \nI suppose the degree to which the claim in Papamakarios et al 2019 concerning the representational ability of discrete flows can be considered 'not correct' largely depends on whether you decide to initially embed your discrete data in the integer lattice (Z^D) before formulating a model. If you view a discrete flow on e.g. 8-bit images as a bijection on the support of this data (Z_{2^8}^{D} in the case of 8-bit images), then the claim of Papamakarios et al is correct. I'd argue this is the implicit assumption, since the language is centred around a view of discrete flows as permutations, which are by definition bijections. On the other hand, if you begin by embedding discrete data in the integer lattice before formulating a model as a bijection on this integer lattice, then the same limit on representational ability doesn't hold, as you show in the example and in Lemma 1. My point here is that the distinction made by this paper is useful, but not mutually exclusive with the discussion in Papamakarios et al. \n\nThe analysis of gradient bias is good, and works to dispel the result I had taken from Hoogeboom et al 2019 regarding the limitations of deep IDFs. It's maybe not very comforting that performance degradation seems to rely on careful choice of architecture for the nets parameterizing the translations though, and it's not immediately clear to me why a simple CNN deteriorates more quickly -- do you have any explanation for why this might be the case? \n\nEven though it's maybe expected, it's also nice to see that careful consideration of initialization in the nets parameterizing the translations also leads to better performance.   \n\nFinally, the experimental result on image data are fine, but I'm not necessarily convinced discrete flows will see practical use a compression tool just yet, as you mention in your conclusion. \n\n#### EXTRA NOTES\nFigure 1 is neat and is quite a good visual aid. Despite lacking the same production quality, I also found the 2D grid in eq. 6 helpful for conveying the concept, and a generalization of this 2D grid might have provided a nice geometric proof of Lemma 1. \n\nBecause the discretized logistic base distributions are used to parameterize a distribution over the integers rather than the usual 8-bit range, do you find there are any numerical issues with large outputs from the model, or accurately evaluating the CDF for these possible outliers?\n\nFigures 3 & 4: Could both be made larger, with increased label, tick, and legend font sizes. Also (nitpick), \"The right two panels in Figure 4 show the loss surface and optimization trajectory of a discrete model trained on the 1-bit data, showing that the optimization landscape is hard to traverse for low bit-depths\" -- do those panels really show it definitively, or just suggest it?  \n\n#### CONCLUSION\nOverall I think the contributions are worthwhile to have in the literature, and would like to see the paper accepted. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "3INK3e2RH3n", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper investigated the weakness of Integer Discrete Flow (IDF) and improved it for the lossless compression task.\n Specifically, the contributions of this paper are:\n\n1. The authors theoretically analyzed the flexibility of normalizing flows for discrete random variables, showing that normalizing flows for discrete random are capable of mapping from a distribution with dependencies across all dimensions to a distribution which is fully factorized.  By relaxing the assumption in Papamakarios et al. (2019) that the domain of the invertible function is the same as the original random variables, the authors denied the claim in Papamakarios et al. (2019) invertible flows for discrete random variables cannot map all distributions to a factorized base distribution (Lemma 1). The authors further demonstrated that the additive coupling layers in IDF is sufficient to model such mapping (Lemma 2).\n\n2. The authors empirically discussed the potential influence of gradient bias on training IDFs, showing that the straight-through estimator is not the cause of gradient bias. Furthermore, the  influence of gradient bias is highly dependent on architecture choices.\n\n3. Based on the previous analysis, the authors proposed an improved architecture for IDF, named IDF++. Experiments on three benchmarks -- CIFAR-10, ImageNet-32 and ImageNet-64 -- show that IDF++ outperforms IDF on all the three benchmarks.\n\nOverall, this paper is well-written and well-motivated. Experimental details are provided to reproduce the results. Ablation studies are conducted to help analyze the effect of different modifications of IDF++.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers of this paper unanimously agreed that this paper adds an interesting theoretical and practical discussion to discrete flows. The paper has improved from the first version to the final one, in which the comments and suggestions by the reviewers have been followed. 

The paper is still incremental with respect to the previous paper and the reviewers all recommended a poster presentation.
Decision: Accept (Poster)