Year: 2021

Abstract:
Dequantisation is a general technique used for transforming data described by a discrete random variable x into a continuous (latent) random variable z, for the purpose of it being modeled by likelihood-based density models. Dequantisation was first introduced in the context of ordinal data, such as image pixel values. However, when the data is categorical, the dequantisation scheme is not obvious. We learn such a dequantisation scheme q(z|x), using variational inference with TRUncated FLows (TRUFL) — a novel flow-based model that allows the dequantiser to have a learnable truncated support. Unlike previous work, the TRUFL dequantiser is (i) capable of embedding the data losslessly in certain cases, since the truncation allows the conditional distributions q(z|x) to have non-overlapping bounded supports, while being (ii) trainable with back-propagation. Addtionally, since the support of the marginal q(z) is bounded and the support of prior p(z) is not, we propose to renormalise the prior distribution over the support of q(z). We derive a lower bound for training, and propose a rejection sampling scheme to account for the invalid samples. Experimentally, we benchmark TRUFL on constrained generation tasks, and find that it outperforms prior approaches. In addition, we find that rejection sampling results in higher validity for the constrained problems.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND: VARIATIONAL DEQUANTISATION
text: 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.

heading: 3 TRUNCATED DEQUANTISER AND PRIOR
text: 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).

heading: 3.1 TRUNCATED LOGISTIC DISTRIBUTION
text: We choose the truncated logistic distribution to parameterise the dequantiser, because the density, the CDF and the inverse CDF of the logistic distribution can all be easily computed:
f(z) = σ(z) · σ(1− z), F (z) = σ(z), F−1(u) = log u 1− u , (4)
for z ∈ R and u ∈ (0, 1), where σ is the logistic sigmoid function. Additionally, we can parameterise the width of Uniform(F (a), F (b)) directly by s, where 0 < s ≤ 1, and the mean of this distribution as m, such that m+ s2 < 1 and m− s 2 > 0. Given a set of statistical parameters ŝ ∈ R and m̂ ∈ R for each category, we can impose these constraints by the following reparameterisation:
m = σ(m̂), s = 2 ·min(m, 1−m) · σ(ŝ). (5) Then, to sample from Uniform(F (a), F (b)),
u = m+ ( u0 − 1
2
) · s, u0 ∼ Uniform(0, 1). (6)
This allows for a simple implementation for a truncated logistic distribution that is differentiable (see Appendix A.1 for the derivation of the gradient). We can write the pdf as
f̃(z;m, s) =
{ 1 sσ(z) · σ(1− z), if 0 < F (z)−m s + 1 2 ≤ 1
0, otherwise. (7)
Like in Lippe & Gavves (2020), the resulting approximate posterior q(zt|xt) requires category specific parameters for the truncation (m(xt) and s(xt)) and any flow applied on top of the truncated logistic (we denote this flow as g). During training, conditioned on a given category xt, we can sample zt with the method described above, and compute the probability q(zt|xt). Since decoding requires computing g−1(zt, x̂t) for all K categories, it limits the choice of flows applicable. In CatNF implementations, a linear flow is used (Kingma & Dhariwal, 2018), while empirically the authors find that a category conditional scale and shift suffices. Algorithm 1 details the steps taken for computing zt, log q(zt|xt), and log p(xt|zt). With this parameterisation of the decoder and the truncated approximate posterior, some q(zt|x̂t) where xt 6= x̂t may not have support over the sample zt. In the extreme case, when q(zt|xt) for all possible assignments of xt had mutually exclusive support, then decoding will always be deterministic, similar to Argmax flows. This is when TRUFL has the flexibility to embed discrete data losslessly in a continuous space. Empirically we find that q(xt|zt) is sparse, but not always deterministic.
Recovering Argmax Flow and CatNF Specific choices of m and s in Eq. 6 can recover CatNF and Argmax Flow support for each category. By settingm = 0.5 and s = 1, we perform no truncation on the posterior distribution, where we recover CatNF. For Argmax Flows, we can achieve orthant support by setting m to 0.25 or 0.75 for each dimension according to the orthant it was assigned based on the binary encoding, and s = 0.25. This will ensure that the support covers only half of the real line in that dimension, thus matching Argmax Flow support for that category.

heading: 3.2 REJECTION SAMPLING
text: When sampling from the model, the standard process is ancestral sampling: z ∼ p(z), and then x ∼ p(x|z). However, in our model, the sample zt may not have support under any of the q(zt|xt), resulting in an undefined optimal p(xt|zt). As mentioned, we reject the samples that have no support for any of the mixture components. This is equivalent to redefining a new prior ptrunc(z) proportional to p(z) such that its support matches the unioned support of q(z|x) and renormalising. The existing ELBO is still a valid lower-bound of such a distribution:
L = Eq(z|x) [ log
p(x|z) p(z) q(z|x)
] ≤ Eq(z|x) [ log
p(x|z) p(z)Z q(z|x)
] ≤ ∫ p(x|z) ptrunc(z) dz, (8)
where Z := ∫ p(z) · 1 [ ∃x̂ : f̃(z;m(x̂), s(x̂)) > 0 ] dz ≤ 1. (9)
Eq. 9 in principle can be estimated by counting the frequency of accepted samples. For one of our experiments, the overall rejection rate at the end of training is about 0.14, suggesting that rejection sampling from a trained prior is not inefficient (see Section 5.2). Using the decoder as an accepting heuristic, we find empirically that this results in more valid samples in constrained tasks that require sampling from the model.

heading: 4 RELATED WORK
text: Flows on Discrete Data Discrete flows (Tran et al., 2019; Hoogeboom et al., 2019) deal with flows in the discrete space, but because of this they resort to the straight-through method (Bengio et al., 2013) for estimating gradient. Using conditional permutations, Lindt & Hoogeboom (2021) develop a way to use discrete flows that do not have the same gradient bias that Discrete flows introduce due to its gradient estimator. Our approach maps the discrete data into a continuous space, allowing the prior to take on most of the modelling complexity.
Learning the Prior Since most of the modelling happens in the prior after dequantisation, the prior has to have the expressibility to model the resulting continuous data. Learnable priors are not new in the VAE literature (Tomczak & Welling, 2018; Huang et al., 2017). One common way for learning sequential data is to use an autoregressive prior (Ziegler & Rush, 2019). Liu et al. (2019) introduces a graph-structured flows which can be used for modelling such data, and there have been flow-based models developed for molecule generation (Madhawa et al., 2019; Shi et al., 2020). Ziegler & Rush (2019) and Huang et al. (2018) also introduce normalising flows that can learn multi-modal target distributions, which Ziegler & Rush (2019) show is important in modelling text data.
Optimal decoder In the case of categorical data with a sufficiently small K, computing the posterior of the encoder is tractable during training. The optimal decoder in Eq. 3 was used in practice in Lippe & Gavves (2020). Argmax flows (Hoogeboom et al., 2021) are a special case in which the encoding of the categorical variable results in a latent variable that has a deterministic decoding, which Nielsen et al. (2020) generalises and calls surjectivity which can be parameterised in either inference or generation. The optimal decoder also has the added effect of alleviating the posterior collapse problem commonly faced when modelling discrete data such as text (Yang et al., 2017; Bowman et al., 2015; Ziegler & Rush, 2019; Chen et al., 2016).
Bounded-support distributions in latent variable models While normalising flows (Papamakarios et al., 2019) allow for learning a more flexible posterior distribution, Jaini et al. (2020) and Verine et al. (2021) discuss the current limitations of flows due to their bi-lipschitz property which result in nearly no transformation at the tails of distributions. Truncated distributions are an extreme case of a light-tailed distribution, achieved by performing a shift and scaling of the base uniform distribution. There have also been efforts on spherical latent variables, modelled by a von Mises-Fisher (vMF) distribution (Xu & Durrett, 2018; Davidson et al., 2018). Polykovskiy & Vetrov (2020) uses bounded-support kernels as the basis of their posterior distributions, and modify the decoder so that the results are deterministic, but introduce gradient estimations in order to train the model. In the latter two cases, rejection sampling was used in order to attain samples for inference.
Consequences of a binary encoding The binary encoding scheme of Argmax flows bears some resemblance to prior work on hierachical softmax schemes to reduce the computational footprint of the large softmax layer. Mnih & Hinton (2008) proposed an analogous scheme for scaling up the softmax layer by representing words as leaves in a tree, and binary encodings as a traversal of such a tree. However, the partitioning of the dlog2Ke-dimensional space into K equal categories will only be exact if K is a power of 2. Otherwise, K − 2dlog2 Ke orthants of the embedding space will not be utilised by the dequantiser. TRUFL has unsupported regions in the decoder as well, but we take the approach of modifying the prior during generation to account for these regions. Furthermore, because the binary vectors are arbitrary, the similarity between words/categories is not reflected in the embedding space as it is in a learned distributed representation.

heading: 5 EXPERIMENTS
text: In this section, we present experiments on TRUFL. See Appendix B for implementation details.

heading: 5.1 TOY EXAMPLE: TWO CATEGORY JOINT DISTRIBUTION
text: In this simple synthetic example, we demonstrate a case in which TRUFL provides an advantage over CatNF and Argmax flows. Our goal is to model a joint probability over 2 categorical random variables, each with 4 categories. This joint distribution is described in Table 1.
Notice that under this distribution, p(x1|x2 = 0) = p(x1|x2 = 2) and p(x2|x1 = 0) = p(x2|x1 = 2). For the purposes of this subsection we will say categories 0 and 2 have equal functionality, while categories 1 and 3 are similar in functionality, but not equal.
To model this distribution, we perform variational dequantisation over the discrete variables x1, x2 using the above methods, and learn a prior over the latent variables. Conditioned on xi, the dequantisation method produces a distribution over a 2 dimensional space — q(zi|xi). It is this distribution we plot in Figure 3.
In both CatNF and TRUFL, the distribution q(zi|xi = 0) and q(zi|xi = 2) are similar. Intuitively, this aligns with the motivation behind the construction of the joint distribution: that the categories are functionally equivalent. However, the dequantisation distribution may not always have this property, as evidenced by the Argmax Flow, due to the way supports are constrained by the orthants. This bimodality in the distributions of these two categories will have to be modelled by the prior, which evidently resulted in a higher log probability in this particular case. In our implementation of Argmax flows for this task, the training eventually resulted in numerical instability. Here we report the best log probability attained before numerical errors occured.
As mentioned in Lippe & Gavves (2020), CatNF can approach deterministic decoding if all other classes are sufficiently far from the target class. However, if several tokens have similar ‘function’ in the dataset, the corresponding components can be forced to be close together in latent space. When this happens, the closeness in the components again result in decoding becoming non-deterministic. TRUFL allows such components to have minimal and sparse overlapping support but yet have the modes of these components be relatively close. If we consider Figure 3b, we can see that the support and distribution for xi = 0 and xi = 2 overlap the most, while the overlap with xi = 1 and xi = 3 is not as great. Since TRUFL minimises the support for the classes that do not overlap in functionality the TRUFL reconstruction loss is lower in comparison to CatNF.

heading: 5.2 GRAPH COLORING
text: We first consider the graph colouring problem (Bondy et al., 1976) introduced in Lippe & Gavves (2020). The task is to colour the nodes in a graph with 3 different colours such that any two nodes
connected by an edge do not have the same colour, which provides a meaningful benchmark for evaluating how well a model learns unknown constraints from the data. We evaluate the Validity of the generated samples conditioned a given graph.
For baselines, we present results from CatNF (Lippe & Gavves, 2020), and implemented the thresholding version of Argmax flows. We implement the truncated flows using the CatNF codebase, and compare results with and without rejection sampling for TRUFL. Additionally, we find a better set of hyperparameters work better for the LARGE setting, and report the results base on those hyperparameters for the baseline as well. All graph-structured flows used in this task are the same as the ones detailed in CatNF, we simply replace the categorical embedding module. For results without rejection sampling, we also compute the probability over the categories using Eq. 9, but disregard the truncation of the densities. We find we perform marginally better than the baselines in the small graph regime. However, with rejection sampling and better hyperparameters, we significantly outperform the flow-based benchmarks, and perform better than the RNN baseline used in Lippe & Gavves (2020).
Rejection Rate of TRUFL Since the prior will be optimised to fit the aggregate posterior during training, we expect the rejection rate of sampling from the truncated prior to drop as training progresses. We measure the rejection rate when sampling from the model and report its average rate of rejection as training progresses. While the model is conditioned on graphs of different sizes during sampling, we look at the average rate of rejection go obtain a measure of the efficiency of our rejection sampling scheme. Figure 4 suggests that the prior learns to fit the support of the aggregate posterior better during training, as the rejection rate reduces from about 0.45 to about 0.14 at the end of training. This
corresponds to roughly losing logZ ≈ 0.15 nats (or 0.22 bits) per molecule. Note that we did not account for this gap while computing the bpd in Table 2 since we would have to estimate logZ conditioned on each graph size, and this would give us slightly lower value in theory.

heading: 5.3 MOLECULE GENERATION
text: Molecule generation is another task commonly used as a generative modelling benchmark (Jin et al., 2018; Shi et al., 2020; Ma et al., 2018; Madhawa et al., 2019). The atoms constituting the molecules
are naturally represented by tokens that are usually interpreted categorically, but the intrinsic structure of the atoms, such as the electron configurations, is often neglected. For this reason, we believe TRUFL and CatNF will learn a more useful representation than Argmax Flow. Moreover, given the results in graph colouring, we believe that rejection sampling can aid in creating more valid samples than prior unconstrained generation methods. Molecule generation also requires the generation of edge categories in addition to the node categories.
Our benchmarks were again implemented with the Lippe & Gavves (2020) codebase, and the Argmax flow benchmark uses the same thresholding scheme. Rejection sampling in this case requires rejecting invalid edges and vertices both, and so resulted in longer sampling times. The results are in Table 3 Again, we find that TRUFL consistently improves upon CatNF in terms of the validity of the generated samples, which can be further increased by truncating the prior via rejection sampling, and taking the largest sub-graph in cases where multiple disconnected graphs are generated (as in done in Lippe & Gavves (2020)).

heading: 5.4 LANGUAGE MODELLING
text: We perform experiments on character-level and word-level language modelling. For word-level language modelling, many tasks require a large number of categories (∼10k). The naive thresholding method will result in a large latent space, no different from using a one-hot representation, which will result in an unwieldy latent size, so we use the previously mentioned binary encoding of each category in our implementation of the Argmax flow.
For the character level experiments on the Penn Treebank (PTB; Marcus et al. 1993), we use the setup in CatNF, replacing the categorical encoding with TRUFL. For text8 (Mikolov et al., 2014), we use the Argmax flow codebase, modifying the Argmax flow module and replacing it with TRUFL. For word level experiments on Wikitext103 (Merity et al., 2016), we use the setting provided by CatNF as well, which evaluates chunks of 256 words, and initialises the word embeddings using GloVe (Pennington et al., 2014). While this is not the standard setting for language modelling, we include these results for comparison with Argmax Flows and CatNF. In this case, we report the negative log-likelihood in the table. For word level experiments on PTB, we use the setting specified by Kim et al. (2019), which also provides an LSTM baseline. Unlike standard PTB results, this benchmark treats each sentence as i.i.d., instead of treating the entire dataset as a continuous string of text. We provide further details of the architecture in Appendix B.2.
The reconstruction loss for TRUFL is higher than that of CatNF, suggesting that CatNF may have optimised the variance of each dequantising distribution to be fairly small in order to approach deterministic decoding. To verify this, we attempt to compute the Within-Group Standard Deviation (WGSD) across all q(zt|xt). Sampling 200 samples from each category, we first normalise the mean and standard deviation across all 10, 000× 200 categories and samples. We then compute the standard deviation for each category across the 200 samples, and average across all 12 dimensions. CatNF has an WGSD of 0.52, while TRUFL has an WGSD of 0.66. A smaller WGSD would indicate peakier disributions in the latent space, given that each quantising distribution has a lower dispersion on average. Since the prior is modelling the correlations between the categories in the sequence, a peakier distribution will require more capacity to model. In this scenario, we find that both models are comparable, with CatNF performing slightly better. However, recall that even with importance sampling, we are not accounting for the unsupported regions of the prior. This means this is an estimate of the upper bound in Eq. 8, and that estimating Z will result in a lower perplexity score. In general, there has been a gap between autoregressive models (e.g. RNNs, Transformers) that model the discrete distribution directly, and latent variable models, and we believe this gap will close as dequantising techniques improve.
We perform a qualitative analysis on the learned stochastic embeddings, by comparing the t-SNE plots of samples of 10 words. This analysis was done by sampling 200 samples from q(zt|xt) for each word, and performing t-SNE over all 200× 10 vectors. As expected, Argmax Flow’s partitioning the latent space into separate orthants results in word embeddings that are arbitrarily chosen. Although N and three are close, this is perhaps due to a chance occurrence of their orthants being dissimilar in only one dimension. On the other hand, TRUFL learning the topology of these embeddings results in clusters that retain some properties of the words they represent. Interestingly, while flew:fly and walked:walk are different tenses of the same word, the learned distribution for each word appears to group them by tense instead. Since this would inform predictions of the tense of future words, this is perhaps a functional representation of these words.

heading: 6 CONCLUSION
text: In this paper, we propose TRUFL, a flow based dequantiser with truncated support for stochastically embedding categorical data in a continuous space. This allows us to employ density-based models, like normalizing flows to fit the continuous representation of the data. To deal with the unsupported regions in the decoder, we further propose truncating the prior distribution to account for the unioned support of the dequantiser, which proves to be more effective at generating samples with constraints, such as for graph coloring and molecule generation. In language modelling, we find that learned stochastic embeddings has better performance than a strict orthant partitioning of the space, and qualitative analysis reveals that the proximity between similar words under such a partitioning is poor. We believe that further work on variational dequantisation will close the gap between auto-regressive models that directly model the categorical space and latent variable models.

heading: A APPENDIX
text: A.1 GRADIENT ESTIMATES OF a AND b
p(z; a, b) =
{ 1
b−a , if a < z < b, 0, otherwise (10)
Let
z = a+ (b− a) · u (11) u ∼ Uniform(0, 1) (12)
u = z − a b− a
(13)
∇aEz[f(z)] = Eu∼Uniform(0,1) [∇af(a+ (b− a) · u)] (14) = Eu∼Uniform(0,1) [(1− u)(∇zf(z))] (15)
∇bEz[f(z)] = Eu∼Uniform(0,1) [u(∇zf(z))] (16)
A.2 SOFT DEQUANTIZATION
Let x denote the data (e.g. a sentence) and L denote the dimensionality of the data; i.e. x ∈ CL and L ∈ Z+. x can be dequantized by sampling zj ∼ q(zj |xj) for j ∈ 1, ..., L. When {q(zj |xj = c) : c ∈ C} has the same support, we call it soft dequantization. We write p(xj |zj) and p(z) to denote the quantizer and the prior over the dequantized data. We consider the following ELBO
L(x, L) = E∏ q(zj |xj)  L∑ j=1 log p(xj |zj) q(zj |xj) + log p(z)  (17) Theorem 1 (Maximizer of ELBO, dequantization). Assume q(x, L) > 0 for all x ∈ CL and L ∈ Z+, and that x1, ..., xL are not conditionally independent given L. Then
Eq(x,L)[log q(x|L)] > Eq(x,L)[L] (18)
as long as the densities q(zj |xj = c) for different categories c ∈ C have the same support. That is, soft dequantization with shared support is suboptimal.
Proof. The optimal quantizer is
q(xj |zj) ∝ q(xj)q(zj |xj) where q(xj) ∝ ∞∑
L=1 L∑ j′=1 q(x′j = xj , L) (19)
Plugging it into the ELBO gives
L(x, L) = E∏ q(zj |xj)  L∑ j=1 log q(xj |zj) q(zj |xj) + log p(z)  = E∏ q(zj |xj)  L∑ j=1 log q(xj) + log p(z)∏ j q(zj)  (20)
where q(zj) = ∑
x′j q(x′j)q(zj |x′j).
On the other hand, we can rewrite the negentropy of the data as
Eq(x,L)[log q(x|L)] = Eq(x,L)  L∑ j=1 log q(xj) + log q(x|L)∏L j=1 q(xj)  (21)
Now to compare the second term with that of the ELBO we let q(z|L) = ∑ q(x|L) ∏L j=1 q(zj |xj).
For simplicity, we denote by T (z|x) = ∏L
j=1 q(zj |xj) the transition kernel applied to either q(x|L) or ∏L
j=1 q(xj). Then
Eq(x,L)
[ log
q(x|L)∏L j=1 q(xj)
] = Eq(x,L)T (z|x) log q(x|L)T (z|x)(∏L j=1 q(xj) ) T (z|x)  (22) = Eq(z,L)q(x|z,L) − log (∏L j=1 q(xj) ) T (z|x)
q(x|L)T (z|x)  (23) ∗ > Eq(z,L) − log Eq(x|z,L)  (∏ j q(xj) ) T (z|x)
q(x|L)T (z|x)  (24) = Eq(z,L) [ − log ∏ j q(zj)
q(z|L)
] = Eq(z,L) [ log
q(z|L)∏ j q(zj)
] (25)
which is greater or equal to the second term of the ELBO by Gibb’s inequality. The inequality marked by ∗ is due to Jensen and the convexity of − log, and is strict since the dependency of x1,...,L implies there exist some x1 and x2 s.t. ∏ j q(x i j) 6= q(xi|L) for i ∈ {1, 2}. Now, since T (z|x) > 0 over the shared support by assumption, for almost all z and L, q(xi|z, L) > 0 for i ∈ {1, 2}. That is, with non-zero probability the argument to the convex function − log has a different value from its mean. This implies the inequality is strict and concludes the proof.
B IMPLEMENTATION DETAILS
We incorporated the codebases from the following sources:
1. Lippe & Gavves (2020): https://github.com/phlippe/CategoricalNF 2. Hoogeboom et al. (2021): https://github.com/didriknielsen/argmax_flows 3. Kim et al. (2019): https://github.com/harvardnlp/compound-pcfg
B.1 HYPERPARAMETERS FOR LARGE GRAPH COLOURING TASK
B.2 ARCHITECTURE FOR THE PRIOR IN LANGUAGE MODELLING TASK
We develop our own architecture for this benchmark, replacing only the categorical encoding scheme for each of the benchmarks. The prior is autoregressive both ‘in time’ and ‘in hidden’, to use nomenclature from Ziegler & Rush (2019). Autoregression in time is governed by an LSTM, and autoregression in hidden is governed by a MADE model (Germain et al., 2015), conditioned on said LSTM. We use a mixture of 4 logistics per-dimension (Ho et al., 2019) in order to model the multi-modality of the distribution over z, which Ziegler & Rush (2019) suggests is prevalent in text data for characters.

Review summaries:
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. 

This paper proposes a dequantization scheme for categorical data, where unlike with ordinal data, element-wise uniform noise cannot be used. The authors propose to encode categorical data into an interval centre and a deviation, and a (pre)dequantized value is obtained by sampling uniformly in the implied real interval. A normalizing flow is then applied to obtain the dequantized variable. The proposed method can thus in principle learn non-overlapping one-dimensional supports for the dequantized variables, enabling lossless dequantization.

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.