Learning Discrete Distributions by Dequantization
Abstract
Media is generally stored digitally and is therefore discrete .
Many successful deep distribution models in deep learning learn a density, i.e. , the distribution of a continuous random variable.
Na¨ıve optimization on discrete data leads to arbitrarily high likelihoods, and instead, it has become standard practice to add noise to datapoints.
In this paper, we present a general framework for dequantization via latent variable modelling.
In this framework, we are able to recover existing dequantization schemes as special cases, and we are able to derive natural extensions from variational inference literature.
We investigate two unexplored directions for dequantization: More sophisticated inference objectives, based on importance-weighting ( iw ) and R´enyi variational inference.
In addition, we analyze dequantization for diﬀerent types of distributions, and show that autoregressive dequantization achieves 3.06 bits per dimension in negative log-likelihood on CIFAR10.
1. Introduction
[Figure]
Today, virtually all media is handled digitally.
As such, it is stored in bits and is therefore discrete .
Deep learning models
Larochelle and Murray (2011); Kingma and
Welling (2014) aim to learn a distribution p model ( x ) for high-dimensional data.
Many of these models are density models Uria et al.
(2013); van den Oord and Schrauwen (2014);
Dinh et al.
; Papamakarios et al.
meaning they learn a distribution of a continuous random variable.
Applying a continuous density model to discrete data, may place arbitrarily high likelihood on the discrete locations Theis et al.
Since discrete and continuous spaces are topologically diﬀerent, a probability density does not necessarily approximate a probability mass.
The total probability at a single point under a density is always zero.
In this paper, we present a general approach for dequantization via latent variable modelling.
In this framework, we are able to recover existing dequantization schemes as special cases, and we are able to derive natural extensions from variational inference literature.
We investigate two directions: more sophisticated variational inference objectives, importanceweighted ( iw ) and R´enyi dequantization; and autoregressive distributions, since noise does not need to be inverted.
We aim to not only investigate diﬀerent dequantization methods, but also to highlight its importance.
2. Related Work
A large number of distribution models learn a density, a distribution over a continuous variable (Uria et al., 2013; van den Oord and Schrauwen, 2014;
Dinh et al., 2017; Papamakarios et al., 2017; Kingma and Dhariwal, 2018; Huang et al., 2018; De Cao et al., 2019; Grathwohl et al., 2019; Hoogeboom et al., 2019b; Ho et al., 2019; Chen et al., 2019; Song et al., 2019;
Ma et al., 2019).
A standard approach adds uniform noise to discrete values (Theis et al., 2016; Uria et al., 2013; van den Oord and Schrauwen, 2014).
Recently, it was proposed to consider a learnable dequantization treated as a variational posterior over latent continuous variables (Ho et al., 2019; Winkler et al., 2019).
Here, we derive a new framework for dequantization using latent variable modelling and we present two new dequantization objectives based on Burda et al.
; Li and Turner (2016) for VAEs.
3. Methodology
Let x ∈ X denote a vector of D observable discrete random variables and P data ( x ) be its (unknown) distribution.
We assume there is a set of data D = { x n } given, or, equivalently, an empirical distribution ˆ P data ( x ) is provided.
The likelihood-based approach to learning a distribution is about ﬁnding values of parameters of a model P model ( x ) that maximize the log-likelihood function: log P model ( D ) =
E x ∼ ˆ P data ( x ) [log P model ( x n )].
3.1. Dequantization as a latent variable model
Frequently, a discrete distribution models a proxy of a continuous variable in the physical world.
For instance, a digital photograph of an observed scene represents the light that is reﬂected from observed objects, quantized to a certain precision.
In other words, we can consider a latent variable model where continuous latent variables v ∈ R D correspond to a continuous representation of the world and observable discrete variables x are measured quantities.
This suggests the following model: P model ( x ) =
(cid:82) P
ϑ ( x | v ) p θ ( v )d v , where
P ϑ ( x | v ) is an indicator function of v being contained in a volume B ϑ ( x ) ⊆
R D , namely,
P ϑ ( x | v ) = 1 [ v ∈ B ϑ ( x )], and p θ ( v ) is a continuous distribution, which may be modeled using a ﬂexible density model (MacKay and Gibbs, 1999;
Dinh et al., 2017; Rippel and
Adams, 2013).
We refer to P ϑ ( x | v ) as a quantizer .
Note that in principle the volumes B ϑ can be constructed to induce any type of partition of a volume space, where care should be taken that B ϑ for diﬀerent x do not overlap.
When we set B ( x ) = { x · u | u ∈ R D + } for x ∈ {− 1 , 1 } D we recover half-inﬁnite dequantization for binary variables from Winkler et al.
In this paper, since image data is often represented on a square grid we will focus on hypercubes , namely, B ( x ) =
{ x + u : u ∈ [0 , 1) D } .
Notice that Theis et al.
; Ho et al.
(2019) require the deﬁnition P ( x ) =
(cid:82) [0 , 1) D p ( x + u )d u to relate a discrete and continuous model.
In contrast, our method is derived without this deﬁnition, and the quantizer volume B generalizes to any volumetric partition.
Calculating the integral in the latent variable model is troublesome, and thus, learning is infeasible especially in high dimensional cases.
Therefore, in order to alleviate this issue, we introduce a new distribution q φ ( v | x ) with parameters φ , a dequantizing distribution or dequantizer .
In fact, the dequantizer should have the same support as P ϑ ( x | v ), otherwise it would assign probability mass to regions outside the volume B ( x ).
Therefore, we will use u instead of v in the dequantizing distribution to highlight the fact that the support of q φ ( v | x ) equals B ( x ), where we deﬁne v = x + u .
Including the dequantizer yields:
[Equation]
Introducing the dequantizer allows us to connect dequantization to the broad literature on variational inference.
We propose three approaches to approximate the integral: i) variational inference, ii) weighted importance sampling and iii) variational R´enyi approximation.
3.2. Variational Dequantization
We can interpret the dequantizing distribution as a variational distribution and apply
Jensen’s inequality to obtain the lower-bound on the log-likelihood function:
[Equation]
The dequantizing distribution must be restricted to assign probability mass to B ( x ) only, otherwise the lower-bound is undeﬁned for certain samples u ∼ q φ ( u | x ).
For our choice of
B ( x ) being a hypercube, we can apply the sigmoid function to the output of the dequantizer to ensure the lower-bound has appropriate support.
Thus, we can re-write (1) as follows:
[Equation]
which recovers the variational dequantization ( vi dequantization ) from Ho et al.
3.3. Importance-Weighted Dequantization
Alternatively, we can interpret the dequantizing distribution as a proposal distribution and instead of using Jensen’s inequality we sample K times from q φ ( u | x ), which directly approximates the log-likelihood:
[Equation]
where u k ∼ q φ ( u | x ) and v k = x + u k for k = 1 , 2 , . . .
, K .
The last equality follows if we constrain the proposal distribution (the dequantizer) in the same manner as we did in the case of the variational dequantization ( i.e. , the probability mass should be assigned only to B ( x )).
This objective was studied in the context of VAEs in (Burda et al., 2016).
In general, if K → ∞ , then we obtain an equality in (3).
But since we take a ﬁnite sample, the approximate gives a lower-bound to the log-likelihood function ( iw-bound ).
Importantly, the iw-bound is tighter than the variational lower-bound Burda et al.
; Domke and Sheldon (2018).
Hence, the importance-weighting is preferable over the variational inference and in practice it leads to a better log-likelihood performance.
We refer to this dequantization scheme as iw dequantization .
3.4. R´enyi Dequantization
The variational inference and importance-weighting sampling for a latent variable model could be generalized by noticing that both approaches are special cases of the variational
R´enyi bounds.
The log-likelihood function could be lower-bounded by the R´enyi divergence approximated with the sample from q φ ( u | x ) of size K < ∞ (Li and Turner, 2016):
[Equation]
where α ∈ [0 , 1) is a hyperparameter.
Interestingly, for α → 1 we obtain the variational lower-bound and for α =
we get the iw-bound.
Li and Turner (2016) have further shown that it is advantageous to consider α < 0, because it may give tighter bounds than the iw-bound when the sample size K is low.
1 Setting α =
−∞ corresponds to picking the largest importance weight value.
By restricting the domain of q ( u | x ) to B ( x ) can obtain the variational R´enyi max approximation ( VR-max ):
[Equation]
The maximum weight dominates the contributions of all the gradients Li and Turner (2016).
Thus, the VR-max approach could be seen as a fast approximation to the importanceweighting, since it speeds up computations by considering only one example instead of K in calculating gradients.
We refer to this whole dequantization scheme as R´enyi dequantization .
3.5. Dequantizing distributions
The dequantizing distribution plays an important role in the framework and its ﬂexibility allows to obtain better log-likelihood scores.
Importantly, the dequantizing distribution is a conditional distribution and we use it for sampling instead of calculating probabilities.
Uniform The special case in which q φ ( u | x ) is a uniform distribution over B ( x ), is the setting introduced in (Theis et al., 2016; Uria et al., 2013; van den Oord and Schrauwen, 2014), termed uniform dequantization .
Gaussian A more powerful dequantization scheme than the uniform dequantization is a conditional logit-normal distribution Atchison and Shen (1980), namely, q φ ( u | x ) =
sigm (cid:16) N (cid:0) µ φ ( x ) , Σ φ ( x ) (cid:1)(cid:17) , where µ φ ( x ) and Σ φ ( x ) denote the mean and the covariance matrix for given x , respectively, and sigm( · ) is the sigmoid function.
Flow Instead of using a certain family of distribution, we can deﬁne the quantizer by applying the change of variables formula, that is:
[Equation]
where f : R D → R D is a bijective map to a simple base distribution q φ ( ε | x ), and J = ∂ ε ∂ u denotes a Jacobian matrix.
Notice we highlight the need of using the (inverse) sigmoid function on top of the bijective map in order to ensure correct support, i.e. u ∈ [0 , 1) D .
There are two important parts of a ﬂow-based model, namely, a choice of a base distribution q φ ( ε | x ) and the bijective map f φ .
Here we decide to use a diagonal Gaussian base distribution
Dinh et al.
Whereas Ho et al.
(2019) studied bipartite maps for f φ , here we examine autoregressive bijective maps, since dequantization noise does not have to be inverted.
Bipartite ﬂows The bipartite bijective maps ensure invertibility by splitting an input into two parts, and processing only the second part (Dinh et al., 2017).
Autoregressive ﬂows We can model q φ ( u | x ) with an ‘expensive to invert’ bijective map.
We ﬁnd that an autoregressive model as proposed for variational autoencoders (Kingma et al., 2016) is an appealing choice for dequantization.
We utilize the following:
[Equation]
[Equation] [Equation] where ARM φ is an autoregressive model (an autoregressive neural network), h is a context variable that is calculated based on the conditioning x using a neural network, s .
We refer to this dequantization scheme as Autoregressive Dequantization (ARD) .
3.6. Distributions for the density model
The continuous model p θ ( v ) is the crucial component in the presented framework since the better performance depends on the ﬂexibility of this model.
In principle, any continuous density model could be used as p θ ( v ), e.g. , any model mentioned in section 3.5, with the important diﬀerence that the sigmoid function is not used as the support does not need to be conﬁned.
In practice, however, p θ ( v ) has to be evaluated during training and we are interested in sampling v ∼ p θ ( v ).
Hence, utilizing models with autoregressive components would be prohibitively slow.
Therefore, in our experiments, we consider a Gaussian distribution with diagonal covariance, full covariance, and a bipartite ﬂow-based model (a series of coupling layers and a factored-out base distribution) as a continuous distribution.
4. Experiments
To understand and evaluate diﬀerent dequantization schemes, they are tested on diﬀerent data problems: i) a 2-dimensional binary problem, ii) (statically) binarized MNIST (bMNIST) (Larochelle and Murray, 2011) and centered patches of bMNIST, which is derived directly from MNIST and iii) CIFAR
(Krizhevsky et al., 2009) (8 bit and 5 bit).
Performance is evaluated on a held-out test-set using negative log-likelihood.
This method of evaluation is common in deep distribution learning literature because it allows for an information theoretic interpretation: the negative log 2 -likelihood is expressed in bits or bits per dimension (bpd), where the latter is an average over dimensions and gives the lossless compression size.
In the experiments we consider diagonal Gaussians, covariance Gaussians and ﬂows as distribution models, since these models admit exact likelihood evaluation.
A detailed description of experiments (i) and (ii) , and architectures and optimization can be found in the Appendix. [Table]
4.1. Image distribution modelling
In this experiment the model using ARD is compared with other methods in the literature.
Experiments show that our model outperforms other methods in the literature on both variational inference objective and negative likelihood (Table 1).
In general we compare to models that do not require an autoregressive inverse to sample from, where the exception is marked (cid:63) .
In particular, we report vi evaluation, also referred to as Expected Lower
Bound (ELBO)
[Table]
and we report the approximate negative likelihood - log P ( x ) usingimportance weighted samples following Maaløe et al.
Note that Ho et al.
(2019) use 16384 samples, which skews the experiment in their favour for vi and - log P ( x ), but against them for KL( q φ | p θ ).
Architecturally, the density model in ARD is most similar to IDF (Hoogeboom et al., 2019a), where 1 × 1 convolutions from Glow (Kingma and
Dhariwal, 2018) and scale transformations from RealNVP (Dinh et al., 2017) are added.
Flow++ (Ho et al., 2019) has additional attention layers and MintNet (Song et al., 2019) has autoregressive transformations instead of coupling layers in the density model.
Note that even though our model utilizes autoregressive components similar to MintNet (Song et al., 2019), our model is computationally cheap to invert since it does not require the solution to autoregressive inverses.
Residual Flow (Chen et al., 2019) utilizes invertible
ResNets instead of coupling layers.
5. Conclusion
In this paper we presented a framework for dequantization via latent variable modeling.
Using this interpretation, we derive natural extensions from variational inference literature for dequantization: importance-weighted ( iw ) and R´enyi dequantization.
We show that low bit-depth data combined with simple dequantizers beneﬁt from these objectives.
In addition, we analyze dequantization for diﬀerent types of distributions (simple and complicated) and show that autoregressive dequantization achieves 3.06 bits per dimension in negative loglikelihood on CIFAR10.
Appendix A. Experiments A.1. Analysis in 2d
The diﬀerent dequantization methods and objectives are analyzed in two dimensions using the binary checkerboard , which places equal probability over two of the four states in the binary space { 0 , 1 } 2 , such that P data ( x ) = 0 .
5 if x = (1 , 0) or if x = (0 , 1), and zero otherwise.
Although the theoretical likelihood limit of a dataset is typically unknown, for the binary checkerboard this is exactly 1 bit, because there is an equal probability over two events.
Since the problem is two dimensional, the learned distributions can be visualized.
Figure 2 depicts the probability density of the dequantizer q ( v | x ) and the density model p ( v ), for models trained using vi -dequantization.
Since by construction q ( v | x ) only places density on a bin corresponding to x , the distribution q ( v | x ) can be visualized without overlap in the marginal distribution q ( v ) =
E x ∼ P data ( x ) [ q ( v | x )].
When the model p ( v ) is a ﬂow and the dequantizer q ( v | x ) is uniform (Figure 2a), the model p ( v ) is struggling to adequately model the boundaries of the dequantized density regions.
When the model p ( v ) is a simple diagonal Gaussian and the dequantizer q ( v | x ) is a ﬂow (Figure 2b), the ﬂexible dequantizer compensates the limitations of the density model by shaping itself to the limitations of the simple distribution.
An interesting variant we would like to highlight is when the density model p ( v ) is a Gaussian with covariance, and q ( v | x ) is a ﬂow (Figure 2c).
Aided by the dequantizer, the model p ( v ) aims to place density on the diagonal line which improves the performance to 1.08 bits, which is already close to the theoretical limit.
Surprisingly, when the density model is relatively simple, a ﬂexible dequantizer can compensate a lot.
The best performance is achieved when both q ( v | x ) and p ( v ) are ﬂexible (Figure 2d).
For this problem we observe the density contracts somewhat away from boundaries, and the center has relatively high density.
The eﬀects seen in Figure 2 are also conﬁrmed quantitatively with the likelihood performance of these models (Table 2).
Note that the more ﬂexible the distributions, the better the performance.
An interesting observation is that when p ( v ) is a ﬂow distribution, a Gaussian q ( v | x ) and a ﬂow q ( v | x ) have equal performance.
Presumably, the ﬂexibility of p ( v ) does not require a more complicated dequantizer for this relatively simple problem.
In addition the eﬀects iw and R´enyi dequantization are shown in Table 3.
Uniformly dequantized models that are trained using R´enyi or iw dequantization are considerably
[Table]
Lower is better.
[Table]
[Figure]
better than vi in terms of likelihood.
For more complicated dequantizers though, we ﬁnd that improvements are negligible.
Therefore, these sophisticated objectives appear to be particularly useful when the dequantization distribution is simple.
Note that due to the low dimensionality of the problem R´enyi dequantization will typically result in loose bounds for larger values than K = 2.
A.2. Analysis in high-dim: Image modelling
Similar to the 2d example, more sophisticated training objectives are most advantageous when dequantization distributions are simple, which can be seen in Table 6.
For binary
MNIST and patches of binary MNIST, training using iw dequantization improves negative likelihood performance consistently.
However, for more expressive learnable dequantizers the added beneﬁt of these objectives becomes smaller.
For CIFAR10 we train only lastepochs with the sophisticated objectives and the ﬁrst with vi to reduce the computational cost.
In high-bit depth settings such as CIFAR10, we ﬁnd that the performance gains of importance weighting are minimal.
Hence, for simple dequantizers on data with low bit depth, iw -dequantization may considerably improve performance.
R´enyi dequantization achieves similar but slightly worse performance, which is acceptable since it is a faster approximation.
Autoregressive Dequantization
Experiments show that ARD outperforms all other dequantization distributions, when trained using comparable architectures (see Table 4 and 5).
Two results are particularly striking:
Firstly, even when binary MNIST is learned using a simple density model p ( v ) (covariance normal), ARD achieves a negative log-likelihood of 0.183 bpd.
In constrast, the uniformly dequantized covariance model achieves only
[Table] [Table]
with 256 samples
Lower is better
[Table]
0.472 bpd.
Secondly, notice that dequantization seems to matter more when bit-depths are smaller.
To see this, consider the log-likelihood improvement when comparing uniform dequantization and ARD:
For the 8 bit data the improvement is 0.12 bpd, which is about 3.7% relative to the total bpd.
However, for 5 bit data the improvement is already 0.20 bpd which is about 12% relatively.
Hence, log-likelihood modelling of lower bit depth data may especially beneﬁt from more expressive dequantizers.
Recommendations
This section aims to give the reader recommendations on what dequantization methods to use and what gains are to be expected.
When bit-depths are small and dequantization distributions are ﬁxed, we ﬁnd that iw or R´enyi dequantization objectives improve log-likelihood performance.
Note that by design of the objectives, the approximate posterior q ( v | x ) will diverge more from the (unknown) true posterior p ( v | x ).
Therefore, a downside of these objectives is that a single sample iw dequantization (equivalent to vi ) will be a poor approximation to the log-likelihood, and instead multiple samples are required to obtain accurate estimates.
In contrast, when dequantization noise can be learned or bit-depths are higher, it may be better to simply use vi dequantization and better performance is obtained by increasing the complexity of the dequantizer.
Note that at the cost of some small performance losses,
Gaussian dequantization might be a good simple alternative to ﬂow-based dequantization.
Appendix B. Architecture and Optimization details
In the experiments we consider diagonal Gaussian, covariance Gaussian and ﬂows as distribution models, since these models admit exact likelihood evaluation.
The diagonal Gaussian is parametrized straightforwardly using parameters for mean and log scale.
The covariance
Gaussian is parametrized using a Cholesky decomposition, i.e. , the precision Λ = LL T where
L is the learnable parameter.
The diagonal of L is modelled separately using a log diagonal parameter, which ensures positive-deﬁniteness of Λ.
The covariance matrix is deﬁned then as Σ = Λ − 1 .
Further, ﬂows have an architecture as described in Kingma and Dhariwal (2018) using the densenet coupling networks from Hoogeboom et al. (2019a).
For MNIST data we use the given split of 40000 train, 10000 validation and 10000 test images.
In the bMNIST patches experiments, center patches of the relevant size are taken.
For CIFAR we split the 50000 training images into the ﬁrst 40000 for train and the last 10000 for validation, we use the 10000 test images as provided.
Models were all optimized using Kingma and Ba (2015) with a learning rate of 0 .and
standard β parameters.
Furthermore, during initial 10 epochs the learning rate is multiplied by epoch divided by 10, referred to as warmup Kingma and Dhariwal (2018).
All our code was implemented in PyTorch Paszke et al.
Since experiments are computationally expensive, architectures were trained once.
Models were trained using two
Nvidia Tesla V100 GPUs, with Nvidia driver 410.104, CUDA 10.0, and cuDNN v7.5.1.
In this setting, smaller models (Binary MNIST) take approximately three days to complete, and larger models (CIFAR10) require ﬁve to six days to complete.
Results are obtained by running models a single time, due to the large computational power that normalizing ﬂows require.
The basic architecture was built following Kingma and Dhariwal (2018):
The ﬂow is divided in multiple levels with a decreasing number of dimensions.
At the end of every level, half of the representation is modelled using a factor out layer (splitprior)
Dinh et al.
; Kingma and Dhariwal (2018).
Every level consists of subﬂows, i.e. a coupling layer followed by a 1 × 1 convolution Kingma and Dhariwal (2018).
The coupling layers utilize neural networks as described in Hoogeboom et al.
(2019a).
For the autoregressive transformation, we utilize the masking as described in Song et al.
In terms of autoregressive order, this is equivalent to reshaping a C × H × W image to a vector and applying the autoregressive mask.
This is opposed to masking in Kingma et al.
(2016), which is equivalent to a mask on a reshaped H × W × C image.
In practice, the autoregressive transformation is obtained by masking convolutions.
[Table]
Appendix C. Additional results
Visualization of samples v ∼ p ( v ) from a density model, and the quantizer x ∼ P ( x | v ) are depicted in Figure 3.
The quantizer is simply a Kronecker delta peak and amounts to applying a ﬂoor function in the case of hypercube partitioning.
The density model is a ﬂow trained with autoregressive dequantization on standard 8 bit CIFAR10.
Notice that although the method is trained using autoregressive dequantization, the density model p ( v ) uses bipartite transformations and does not require the solution to autoregressive inverses.
[Table]
Appendix D. Visualizations on Binary Checkerboard
In this section a comprehensive overview of the distributions dequantizer and density model pairs is visualized.
The models trained using variational inference are displayed in Table 9.
In general, the dequantizer q ( v | x ) and density model p ( v ) try to compensate for each other where they are lacking ﬂexibility.
This eﬀect can be seen when q ( v | x ) is a ﬂow and p ( v ) is a diagonal Gaussian, a covariance Gaussian and lastly a ﬂow.
When p ( v ) is a ﬂow, it is generally diﬃcult to capture the edges of the squares when dequantization noise is uniform.
However, both Gaussian and ﬂow dequantization perform equally when the model p ( v ) is a ﬂow.
In this simple problem, Gaussian dequantization is suﬃciently ﬂexible when combined with a ﬂow.
The models trained using iw and R´enyi dequantization objectives are depicted in Table 10.
An important diﬀerence with vi -dequantization is that it is much less important for q ( v | x ) and p ( v ) to match completely.
Rather, more emphasis is placed so that p ( v ) places distribution somewhere in the appropriate bin, where the exact location in the bin matters less.
As a result, when q ( v | x ) is uniform the model p ( v ) is not forced to learn the uniform square and retracts somewhat away from the edges.
[Table] [Table]
x )
Uniform
Normal
Flow
Normal diag.
>5.0 -log q(v) -log
p(v) 2.
>5.0 -log q(v) -log
p(v) 2.
>5.0 -log q(v) -log
p(v) 2.
[Table]
Normal
[Table]
>5.0 -log q(v) -log
p(v) 1.
>5.0 -log q(v) -log
p(v) 1.
>5.0 -log q(v) -log
p(v) 1.
Flow  2     0.
>5.0 -log q(v) -log
p(v) 1.
>5.0 -log q(v) -log
p(v) 1.
>5.0 -log q(v) -log
p(v) 1.
[Table]
R´enyi
[Table]
K = 2)
[Table] [Table]
[Table] [Table]
log q(v)
[Table]
log p
[Table]
v)
[Table]
>5.0 -log q(v) -log
p(v) 1.