L EARNING TO D EQUANTISE WITH T RUNCATED F LOWS
A BSTRACT
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 ﬁrst 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 ﬂow-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 ﬁnd that it outperforms prior approaches.
In addition, we ﬁnd that rejection sampling results in higher validity for the constrained problems.
I NTRODUCTION
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 ﬁtting 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 Flow (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 [Figure] 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 ﬂow.
Our experimental results highlight the effectiveness of our approach.
B ACKGROUND : V ARIATIONAL D EQUANTISATION
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 = { z 1 , . .
, z T } denote this continuous representation, and x = { x 1 , . . .
, x T } describe the observed data, where each x t represent, e.g. a node in a graph or a token in a sentence.
Each x t is assumed to be categorical, i.e. x t ∈ { 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 ( z t | x t ) as the dequantiser and p ( x t | z t ) as the quantiser .
Training can be achieved by maximizing a variational lower bound on the marginal likelihood of the data, i.e. :
[Equation]
We are interested in the case where the representation z t can be inferred from x t alone, so we choose the factorisation p ( x | z )
= (cid:81) t p ( x t | z t ) and q ( z | x ) =
(cid:81) t q ( z t | x t ) , following Lippe & Gavves (2020).
In this case, the “optimal” quantiser p ( x t | z t ) can be conveniently computed as:
[Equation]
where q ( x ) denotes the (empirical) data distribution, and ˜ p ( x t ) 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 ( z t | x t ) .
The factorisation we chose for p ( x | z ) and q ( z | x ) is crucial for the arg max 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:
[Equation]
We note that the ﬁrst 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 ( z t | x t ) .
q ( x t | z t ) is deterministic when q ( z t | x t ) does not overlap with other q ( z t | x (cid:48) t ) , in which case q ( z t | x t ) 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 ( z t | x (cid:48) t ) , therefore resulting in a spikier aggregate posterior distribution (see more discussion on this in Section 5.3).
With this general framework that also accounts for lossy quantisation, we brieﬂy 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 ( z t | x t ) =
Uniform( x t , x t + 1) .
The resulting quantisation process is simply (cid:98) z t (cid:99) , and is deterministic.
More generally, q ( z t | x t ) can be any distribution on [ x t , x t +1] .
See Nielsen & Winther (2020); Hoogeboom et al.
for extensions of the uniform dequantisation scheme.
Argmax Flow For categorical data, uniform dequantisation is not applicable, as there is no intrinsic ordering between the categories.
Argmax Flow (Hoogeboom et al., 2021) dequantise categorical data by letting z t ∈ R K be distributed by q ( z t | x t ) with support { z t : arg max k ( z t ) k
= x t } .
When the support over the latent space is disjoint, p ( x t | z t )
=
q ( x t | z t )
{+= 
+}
[ x t = arg max k ( z t ) k ] 1 ; we depict this in Figure 1, Argmax ﬂow (thresh.).
Argmax Flow 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 Flow 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 (cid:100) log 2 K (cid:101) ; See Figure 1, Argmax ﬂow (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 ﬂow, the dequantising distributions q ( z t | x t = k ) for all 0 ≤ k ≤ K −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 ( z t | x t ) .
In this case, the optimal quantiser q ( x t | z t ) is nearly-deterministic if the locations of the dequantisers are far away from each other and they have sufﬁciently small scale.
For this reason, and unlike the ﬁrst 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 ( z t | x t ) 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 ( z t | x t ) 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.
T RUNCATED D EQUANTISER AND P RIOR
The general approach to optimising the variational lower bound proposed in Kingma & Welling (2014) involves sampling from the proposal distribution q ( z t | x t ) 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 x t , but all of the variational distribution is conditioned on the categorical value x t .
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, and we use it to sample from our dequantiser.
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).
3. T RUNCATED L OGISTIC D ISTRIBUTION
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:
[Equation]
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 + s 2 < 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:
[Equation]
Then, to sample from Uniform( F ( a ) , F ( b )) ,
[Equation]
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
[Equation]
Like in Lippe & Gavves (2020), the resulting approximate posterior q ( z t | x t ) requires category speciﬁc parameters for the truncation ( m ( x t ) and s ( x t ) ) and any ﬂow applied on top of the truncated logistic (we denote this ﬂow as g ).
During training, conditioned on a given category x t , we can sample z t with the method described above, and compute the probability q ( z t | x t ) .
Since decoding requires computing g − 1 ( z t , ˆ x t ) for all K categories, it limits the choice of ﬂows applicable.
In CatNF implementations, a linear ﬂow is used (Kingma & Dhariwal, 2018), while empirically the authors ﬁnd that a category conditional scale and shift sufﬁces.
Algorithm 1 details the steps taken for computing z t , log q ( z t | x t ) , and log p ( x t | z t ) .
With this parameterisation of the decoder and the truncated approximate posterior, some q ( z t | ˆ x t ) where x t (cid:54) =
ˆ x t may not have support over the sample z t .
In the extreme case, when q ( z t | x t ) for all possible assignments of x t had mutually exclusive support, then decoding will always be deterministic, similar to Argmax ﬂow.
This is when TRUFL has the ﬂexibility to embed discrete data losslessly in a continuous space.
Empirically we ﬁnd that q ( x t | z t ) is sparse, but not always deterministic.
[-= 
-]
Recovering Argmax Flow and CatNF Speciﬁc choices of m and s in Eq.
6 can recover CatNF and Argmax Flow support for each category.
By setting m = 0 .
5 and s = 1 , we perform no truncation on the posterior distribution, where we recover CatNF.
For Argmax Flow, 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 .
This will ensure that the support covers only half of the real line in that dimension, thus matching Argmax Flow support for that category.
Input: Categorical data x t , Flow g ( · , · ) Output: z t , log q ( z t | x t ) , log p ( x t | z t ) u 0 ∼ Uniform(0 , 1) (cid:46)
Begin encoding u ← m ( x t ) + (cid:0) u 0 − 12 (cid:1) · s ( x t ) z (cid:48)
t ← F − 1 ( u ) z t ← g ( z (cid:48) t , x t ) (cid:46) End Encoding for ˆ x t = 0 to K − 1 do (cid:46) Compute probability of z t given all possible ˆ x t ˆ z (cid:48) t ← g − 1 ( z t , ˆ x t ) log q ( z t | ˆ x t ) ← log ˜ f (ˆ z (cid:48) t ; m (ˆ x t ) , s (ˆ x t ))
+ log (cid:12)(cid:12)(cid:12) d z (cid:48)
t dˆ z t (cid:12)(cid:12)(cid:12) end for log p ( x t | z t ) ← log q ( z t | x t ) ˜
p ( x t ) − log (cid:80)
ˆ x t q ( z t | ˆ x t ) ˜ p (ˆ x t ) (cid:46) log computation of the q posterior
Output: x z ∼ p ( z ) while ∀ ˆ x : ˜ f
( z ; m (ˆ x ) , s (ˆ x ))
{+= 
+}
do z ∼ p ( z ) end while x ∼ p ( x | z ) p ( z ) p trunc ( z )
3. R EJECTION S AMPLING
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 z t may not have support under any of the q ( z t | x t ) , resulting in an undeﬁned optimal p ( x t | z t ) .
As mentioned, we reject the samples that have no support for any of the mixture components.
This is equivalent to redeﬁning a new prior p trunc ( 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:
[Equation]
where
[Equation]
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 inefﬁcient (see Section 5.1).
Using the decoder as an accepting heuristic, we ﬁnd empirically that this results in more valid samples in constrained tasks that require sampling from the model.
R ELATED W ORK
Flows on Discrete Data Discrete ﬂows (Tran et al., 2019; Hoogeboom et al., 2019) deal with ﬂows 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 ﬂows that do not have the same gradient bias that Discrete ﬂows 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.
[Table]
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 ﬂows which can be used for modelling such data, and there have been ﬂow-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 ﬂows 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 sufﬁciently 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 ﬂow (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 ﬂows (Papamakarios et al., 2019) allow for learning a more ﬂexible posterior distribution, Jaini et al.
(2020) and Verine et al.
(2021) discuss the current limitations of ﬂows 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.
E XPERIMENTS
In this section, we present experiments on TRUFL.
See Appendix C for implementation details.
5. G RAPH COLORING
We ﬁrst 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.
[Table]
For baselines
[Table]
we present results from CatNF (Lippe & Gavves, 2020) [Table] and implemented the thresholding version of Argmax ﬂow.
We implement the truncated ﬂows using the CatNF codebase [Table] and compare results with and without rejection sampling for TRUFL.
Additionally, we ﬁnd a better set of hyperparameters work better for the L ARGE setting
[Table]
and report the results base on those hyperparameters for the baseline as well.
All graph-structured ﬂows 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. [Table] but disregard the truncation of the densities.
We ﬁnd we perform marginally better than the baselines in the small graph regime.
However, with rejection sampling and better hyperparameters, we signiﬁcantly outperform the ﬂow-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 ﬁt 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 efﬁciency of our rejection sampling scheme.
Figure 3 suggests that the prior learns to ﬁt the support of the aggregate posterior better during training, as the rejection rate reduces from about 0.
to about 0.14 at the end of training.
This corresponds to roughly losing log Z ≈ 0 .
15 nats (or 0 .
22 bits) per molecule.
Note that we did not account for this gap while computing the bpd in Table 1 since we would have to estimate log Z conditioned on each graph size, and this would give us slightly lower value in theory.
5. M OLECULE GENERATION
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 conﬁgurations, is often neglected.
For this reason, we believe
TRUFL and CatNF will learn a more useful representation than Argmax Flow.
Moreover, given the [Table] [Table] 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 implemented with the Lippe & Gavves (2020) codebase, and the Argmax ﬂow 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
Again, we ﬁnd 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)).
5. L ANGUAGE M ODELLING
One can view these variational dequantisation methods when applied to language modelling as a method of using stochastic embeddings : each token is represented by a distribution over the embedding space, and language models then operate over samples from these distributions.
Probabilistic embeddings have been suggested in prior work (Li et al., 2018; Dasgupta et al., 2020; Chen et al., 2021), but here, distribution parameters are trained with a language modelling objective.
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, so we use the previously mentioned binary encoding of each category in our implementation of the Argmax ﬂow.
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 ﬂow codebase, modifying the Argmax ﬂow 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 Flow and CatNF.
We report the negative log-likelihood in Table 3.
For word level experiments on PTB, we use the setting speciﬁed by Kim et al.
(2019), which also provides an LSTM baseline.
Unlike standard PTB results, this benchmark treats each sentence Since t-SNE reveals the relative proximities of the embeddings, we should note that the visualisation here reveals what clusters of embeddings are close to others relative to the other words in the plot. 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 C.3.
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 compute the Within-Group Standard Deviation (WGSD) across all q ( z t | x t ) .
Sampling 200 samples from each category, we ﬁrst 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 ﬁnd 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, comparing the t-SNE plots of samples of 10 chosen words.
We took 200 samples from q ( z t | x t ) for each word, and performed t-SNE over all 200 × 10 vectors.
As expected, Argmax ﬂow partitions the latent space into separate orthants, resulting in word embeddings that are arbitrarily chosen.
N and three are an exception, but 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.
C ONCLUSION
In this paper, we propose TRUFL, a ﬂow based dequantiser with truncated support for stochastically embedding categorical data in a continuous space.
This allows us to employ density-based models like normalising ﬂows to ﬁt 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 ﬁnd 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.
A A PPENDIX A. G RADIENT ESTIMATES OF a AND b
[Equation]
Let
[Equation]
A. S OFT DEQUANTIZATION
Let x denote the data ( e.g. a sentence) and L denote the dimensionality of the data; i.e. x ∈ C L and
L ∈ Z + .
x can be dequantized by sampling z j ∼ q ( z j | x j ) for j ∈ 1 , ..., L .
When { q ( z j | x j = c ) : c ∈ C} has the same support, we call it soft dequantization.
We write p ( x j | z j ) and p ( z ) to denote the quantizer and the prior over the dequantized data.
We consider the following ELBO
[Equation]
Theorem 1 (Maximizer of ELBO, dequantization) .
Assume q ( x, L ) > 0 for all x ∈ C L and L ∈ Z + , and that x 1 , ..., x L are not conditionally independent given L .
Then
[Equation]
as long as the densities q ( z j | x j = c ) for different categories c ∈ C have the same support.
That is, soft dequantization with shared support is suboptimal.
Proof.
The optimal quantizer is
[Equation]
where
[Equation]
Plugging it into the ELBO gives
[Equation]

[Equation]
[Equation]
where q ( z j ) = (cid:80) x (cid:48)
j q ( x (cid:48) j ) q
( z j | x (cid:48) j ) .
On the other hand, we can rewrite the negentropy of the data as
[Equation] [Equation]
Now to compare the second term with that of the ELBO we let q ( z | L ) =
(cid:80) q
( x | L ) (cid:81)
Lj =1 q ( z j | x j ) .
For simplicity, we denote by T ( z | x ) =
(cid:81)
Lj =1 q ( z j | x j ) the transition kernel applied to either q ( x | L ) or (cid:81)
Lj =1 q ( x j ) .
Then
[Equation]
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 x 1 ,...,L implies there exist some x 1 and x 2 s.t.
(cid:81)
j q ( x ij ) (cid:54) =
q ( x i | L ) for i ∈ { 1 , 2 } .
Now, since T ( z | x ) > 0 over the shared support by assumption, for almost all z and L , q ( x i | 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 F URTHER R ELATED W ORK Consequences of a
binary encoding
The binary encoding scheme of Argmax ﬂow bears some resemblance to prior work on hierachical softmax schemes to reduce the computational footprint of the large softmax layer (Goodman, 2001; Morin & Bengio, 2005; Mnih & Hinton, 2008; Shen et al., 2017).
Speciﬁcally, 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 (cid:100) log 2 K (cid:101) -dimensional space into K equal categories will only be exact if K is a power of 2.
Otherwise, K − 2 (cid:100) log 2 K (cid:101) 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 reﬂected in the embedding space as it is in a learned distributed representation.
C E XPERIMENTS & I MPLEMENTATION D ETAILS
We incorporated the codebases from the following sources:
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
C. T OY E XAMPLE : T WO C ATEGORY J OINT D ISTRIBUTION
In this simple synthetic example, we demonstrate a case in which TRUFL provides an advantage over CatNF and Argmax ﬂow. Our goal is to model a joint probability over 2 categorical random variables, each with 4 categories.
This joint distribution is described in Table 5.
[Figure]
[Table]
Notice that under this distribution, p ( x 1 | x 2 = 0)
=
p ( x 1 | x 2 = 2) and p ( x 2 | x 1 = 0) =
p ( x 2 | x 1 = 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 x 1 , xusing the above methods, and learn a prior over the latent variables.
Conditioned on x i , the dequantisation method produces a distribution over a 2 dimensional space — q
( z i | x i ) .
It is this distribution we plot in Figure 5.
In both CatNF and TRUFL, the distribution q ( z i | x i = 0) and q ( z i | x i = 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 ﬂow 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 sufﬁciently 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 5b, we can see that the support and distribution for x i = 0 and x i = 2 overlap the most, while the overlap with x i = 1 and x i = 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.
C. H YPERPARAMETERS FOR L ARGE GRAPH COLOURING TASK C. A RCHITECTURE FOR THE P RIOR IN L ANGUAGE M ODELLING 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
[Table]