Learning to Dequantise with Truncated Flows

1 INTRODUCTION . Deep generative models aim to model a distribution of high-dimensional natural data . Many of these methods assume that the data is continuous , despite it being digitally stored in bits and therefore intrinsically discrete . This discrepancy has led to recent interest in dequantising discrete data types to avoid some of the degeneracies of fitting continuous models to discrete data ( Theis et al. , 2015 ) . When data is ordinal ( such as pixel intensities ) a naive dequantisation scheme can be obtained by adding uniform noise to the discrete values ( Theis et al. , 2015 ) . More recently , a generalisation of this approach where dequantisation is seen as inference in a latent variable model has also been proposed ( Ho et al. , 2019 ; Hoogeboom et al. , 2020 ; Nielsen et al. , 2020 ) . However , these methods may not be directly applied in cases where the data is categorical ( Hoogeboom et al. , 2021 ) , because the data is not naturally represented in a vector space . Attempts at devising dequantisation schemes for categorical data by building upon the variational dequantisation scheme have been recently proposed in Hoogeboom et al . ( 2021 ) and Lippe & Gavves ( 2020 ) . These approaches dequantise a categorical input into a latent continuous space . Ideally , a dequantisation scheme for categorical data should be : ( i ) easily learnable by standard optimization techniques and ( ii ) possibly lossless , in the sense that quantisation should recover the input category . Argmax Flows ( Hoogeboom et al. , 2021 ) offer lossless dequantisation but the support of the stochastic embedding is chosen arbitrarily and not optimised , and the dimensionality of the continuous ( dequantised ) variable is required to be at least logarithmic in the number of categories of the input data . Moreover , the method makes minimal assumptions about the topology of the categorical data , disregarding the possible relationships between categories , which can occur for example between word indices in natural language ( Bengio et al. , 2003 ) or the atomic representations of a molecule ’ s constituents . On the other hand , Categorical Normalizing Flows ( CatNF ; Lippe & Gavves 2020 ) can learn a more compact representation of the input category but the dequantisation might be lossy given that the posteriors over the continuous variables have overlapping support . Is there a trade-off between these two schemes ? In this paper , we propose TRUFL , which builds upon the aforementioned variational dequantisation techniques . We achieve that by using truncated posterior distributions over the continuous variables with potentially bounded and disjoint support . In addition , we present a parametrisation of truncated distributions that can be optimised with standard stochastic reparametrisation techniques . Overall , our method inherits strengths of both CatNF and Argmax flows . Our experimental results highlight the effectiveness of our approach . 2 BACKGROUND : VARIATIONAL DEQUANTISATION . Dequantisation refers to the process of embedding discrete-valued data into a continuous space , which allows us to employ density-based models to capture the distribution of the continuous representation . Concretely , let z = { z1 , . . . , zT } denote this continuous representation , and x = { x1 , . . . , xT } describe the observed data , where each xt represent , e.g . a node in a graph or a token in a sentence . Each xt is assumed to be categorical , i.e . xt ∈ { 0 , · · · , K − 1 } for some integer K > 1. z can be interpreted as a latent variable , which follows a prior distribution p ( z ) . We refer to q ( zt|xt ) as the dequantiser and p ( xt|zt ) as the quantiser . Training can be achieved by maximizing a variational lower bound on the marginal likelihood of the data , i.e . : log p ( x ) ≥ Eq ( z|x ) [ log p ( x|z ) p ( z ) q ( z|x ) ] = : L ( x ) ( 1 ) We are interested in the case where the representation zt can be inferred from xt alone , so we choose the factorisation p ( x|z ) = ∏ t p ( xt|zt ) and q ( z|x ) = ∏ t q ( zt|xt ) , following Lippe & Gavves ( 2020 ) . In this case , the “ optimal ” quantiser p ( xt|zt ) can be conveniently computed as : argmax p ( xt|zt ) Eq ( x ) [ L ( x ) ] = q ( zt|xt ) p̃ ( xt ) ∑K−1 x′t=0 q ( zt|x′t ) p̃ ( x′t ) = q ( xt|zt ) = : p ( xt|zt ) ( 2 ) where q ( x ) denotes the ( empirical ) data distribution , and p̃ ( xt ) denotes the estimate of the marginal distribution of each category ( which can be obtained by counting and , in the case of textual data , this corresponds to the unigram distribution over words ) . This equation shows that the optimal quantiser can be obtained implicitly by applying Bayes ’ rule with the parametric dequantiser q ( zt|xt ) . The factorisation we chose for p ( x|z ) and q ( z|x ) is crucial for the argmax above to be represented in this simple form . Without this assumption , the solution will involve a combinatorial sum or an integral , which results in the choice of a parametric quantiser in Ziegler & Rush ( 2019 ) for computational tractability . Plugging the optimal decoder into Eq . 1 yields : L ( x ) = Eq ( z|x ) [ ∑ t log p̃ ( xt ) + log p ( z ) ∑K−1 x′t=0 q ( zt|x′t ) p̃ ( x′t ) ] ( 3 ) We note that the first term is a constant . Therefore , the expression above implies that accurately modelling the dependencies in x boils down to learning an expressive prior p ( z ) and regularising the dequantiser q ( zt|xt ) . q ( xt|zt ) is deterministic when q ( zt|xt ) does not overlap with other q ( zt|x′t ) , in which case q ( zt|xt ) is encouraged to be expanded to maximize the entropy . If there is certain amount of overlapping , the denominator in the second term will push down the density of other q ( zt|x′t ) , therefore resulting in a spikier aggregate posterior distribution ( see more discussion on this in Section 5.4 ) . With this general framework that also accounts for lossy quantisation , we briefly present some of the previously proposed strategies for dequantisation . Ordinal dequantisation In the case where the data is ordinal , such as the case of image pixel values ( e.g . for an 8-bit representation , K = 256 ) , a dequantisation scheme can be obtained by setting q ( zt|xt ) = Uniform ( xt , xt + 1 ) . The resulting quantisation process is simply bztc , and is deterministic . More generally , q ( zt|xt ) can be any distribution on [ xt , xt+1 ] . See Nielsen & Winther ( 2020 ) ; Hoogeboom et al . ( 2019 ) for extensions of the uniform dequantisation scheme . Argmax Flows For categorical data , uniform dequantisation is not applicable , as there is no intrinsic ordering between the categories . Argmax Flows ( Hoogeboom et al. , 2021 ) dequantises categorical data by letting zt ∈ RK be distributed by some q ( zt|xt ) supported on { zt : argmaxk ( zt ) k = xt } . When the support over the latent space is disjoint , p ( xt|zt ) = q ( xt|zt ) = 1 [ xt = argmaxk ( zt ) k ] 1 ; we depict this in Figure 1 , Argmax flow ( thresh. ) . Argmax Flows makes minimal assumptions on the topology of the data : the support of the dequantiser partitions the continuous space evenly and the representations are equally far from each other . As an example , synonyms in text may still have very distinct dequantised representations despite having similar functions and meaning in a language modelling setting . In the naive formulation , Argmax Flows requires the dimensionality to the latent space to be the same as the number of the input categories K. To accomodate for larger categorical spaces , the authors suggest a binary factorisation , reducing the required latent space dimension to dlog2Ke ; See Figure 1 , Argmax flow ( binary ) . Categorical Normalising Flows ( CatNF ) In the previous cases , the quantisation is deterministic , and there is no loss of information . This is because in a cleanly partitioned latent space like in the ordinal setting or argmax flow , the dequantising distributions q ( zt|xt = k ) for all 0 ≤ k ≤ K − 1 have non-overlapping support . CatNF learns a dequantiser that can “ softly ” partition the space . Lippe & Gavves ( 2020 ) propose using a conditional logistic distribution as q ( zt|xt ) . In this case , the optimal quantiser q ( xt|zt ) is nearly-deterministic if the locations of the dequantisers are far away from each other and they have sufficiently small scale . For this reason , and unlike the first two approaches , CatNF is not capable of losslessly dequantising the data ( we provide a formal discussion on this in Appendix A.2 , which is based on an data-processing inequality argument , using the dequantiser as a transitional kernel ) . It can approximate the lossless limit by pushing the bulk of the mass of q ( zt|xt ) away from each other , but that could potentially lead to a highly complex and multi-modal empirical distribution over the representation space for p ( z ) to approximate . Next , we consider the case where q ( zt|xt ) is a truncated distribution , and as such has the ability to encode the data losslessly while learning a meaningful latent topology of the different categories . 3 TRUNCATED DEQUANTISER AND PRIOR . The general approach to optimising the variational lower bound proposed in Kingma & Welling ( 2014 ) involves sampling from the proposal distribution q ( zt|xt ) to estimate the expectation ( Eq . 1 ) . In our case , we want to parameterise this using a TRUncated FLow , which we will refer to as TRUFL . For simplicity , we will drop the dependency on t and xt , but all of the variational distribution is conditioned on the categorical value xt . If we want to bound a scalar distribution between ( a , b ) , and we have a density function f where its cumulative distribution function F ( CDF ) and the inverse of its CDF F−1 are tractable , we can easily sample from f by sampling u from Uniform ( F ( a ) , F ( b ) ) , and then evaluating F−1 ( u ) . Note that this method is differentiable . We use it to sample from our dequantiser . 1We use 1 [ · ] to denote an indicator function . Algorithm 1 Truncated Categorical Encoding for a timestep t Input : Categorical data xt , Flow g ( · , · ) Output : zt , log q ( zt|xt ) , log p ( xt|zt ) u0 ∼ Uniform ( 0 , 1 ) . Begin encoding u← m ( xt ) + ( u0 − 12 ) · s ( xt ) z′t ← F−1 ( u ) zt ← g ( z′t , xt ) . End Encoding for x̂t = 0 to K − 1 do . Compute probability of zt given all possible x̂t ẑ′t ← g−1 ( zt , x̂t ) log q ( zt|x̂t ) ← log f̃ ( ẑ′t ; m ( x̂t ) , s ( x̂t ) ) + log ∣∣∣dz′tdẑt ∣∣∣ end for log p ( xt|zt ) ← log q ( zt|xt ) p̃ ( xt ) − log ∑ x̂t q ( zt|x̂t ) p̃ ( x̂t ) . log computation of the q posterior In general , however , multi-variate distributions may not simply be truncated at the tails , but rather have a support which is a strict subset of its base distribution . One general approach to sampling from such a distribution is via rejection sampling ( Murphy , 2012 ) . This approach has been used in prior work for sampling from bounded-support distributions ( Polykovskiy & Vetrov , 2020 ; Xu & Durrett , 2018 ; Davidson et al. , 2018 ) . Computing gradients for this method is possible via implicit gradients ( Figurnov et al. , 2018 ) , but we do not need gradients in our case , as we use rejection sampling for generating samples from the generative model ( See Section 3.2 ) .

Summary: This paper proposes a Flow model called TRUFL designed for discrete data. The main selling point of the proposed model is that it can handle discrete data better than other dequantization schemes. The authors target two key difficulties of this dequantization problem, which are (i) making the dequantization lossless, and (ii) allowing the dequantizer to be learned easily. Building on the Categorical Normalizing Flows, the authors propose to truncate the dequantized latent space to make the latent space of different input categories less correlated. To compute the probabilities w.r.t. the truncated latent space, the authors use rejection sampling to approximate these probabilities.