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: The paper presents a new approach for dequantization ie embedding discrete data in a continuous space using variational inference and truncated flows called TRUFL. Unlike previous approaches, TRUFL aloows the dequantization layer to have a learnable truncated support. The authors perform several experiments to demonstrate the advantages of the proposed method to Categorical Normalizing Flows (CatNF) and Argmax Flows.