Learning to Dequantise with Truncated Flows
Authors: Shawn Tan
Chin-Wei Huang
Alessandro Sordoni
Aaron Courville
Venue: ICLR 2022 Poster (2022)
Keywords: variational inference
variational bayes
dequantisation
normalizing flows

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 renormalising 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 during generation.
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.

Published as a conference paper at ICLR 2022
LEARNING TO DEQUANTISE WITH TRUNCATED FLOWs
Shawn Tan & Chin-Wei Huang
Mila, University of Montreal
[iing. shan. shawn. tan, chin-wei ehuanngeumottrraaa.
Alessandro Sordoni
Aaron Courville
Microsoft Research
Mila, University of Montreal
Montreal al 1isordonmmeerooott.com
Canada CIFAR AI Chair
courvila@ airo..rr..... umontreal.co
ABSTRACT
Dequantisation is general technique used for transforming data described by a discrete random variable æ into a continuous (latent) random variable z, for the pur- pose of it being modeled by likelihood-base 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 g( æ),, using variational inference with TRUncated FLows (TRUFL) a learnable truncated support. Unlike previous work, the TRUFL dequantiser is capable of embedding the data losslessly in certain cases, since the truncation allows the conditional distributions g(= to have non-overlapping bounded sup- ports, while being (ii) trainable with back-propagation
novel flow-based model that allows the dequantiser to have
Addtionally, since the support of the marginal g(z) is bounded and the support of prior p(z) is not, we propose renormalise the prior distribution over the support of q(z). We derive a lower bound for training. and propose 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.
INTRODUCTION
Deep generative models aim 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 the discrete values (Theis et al., 2015). More recently, generalisation of this approach where dequantisation seen as inference in a latent variable model has also been proposed (Ho al., 2019; Hoogeboom al., 2020; Nielsen 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 naturally represented in 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) the atomic representations

---

Published as a conference paper at ICLR 2022
Do_
p(xelze)
CatNF
Argmax flow (thresh.) Argmax flow (binary)
TRUFL
Figure 1: Decoder for CNFs (left) is the posterior of the encoder and the posteriors have unbounded support resulting probability mass in all categories in p(zz)2z)). Argmax flow (center) have bounded supports that partitions the full latent space into regions where the posterior for each category is deterministic, but (1) the naive thresholding (thresh.) requires the dimensionality of the space be equal to the number of categories, and (2) the binary encoding partitions data into separate octants, sometimes leaving octants unsupported. TRUFL (right) allows bounded support for each category, and for latent dimensions to be different from the category size.
of a molecule's constituents, On the other hand, Categorical Normalizing Flows (CatNF; Lippe & Gavves 2020) can learn more compact representation of the input category but the dequantisation might be lossy given that the posteriors over the continuous variables have overlapping support.
there a trade-off between these two schemes? 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 parametrisation of truncated distributions that can be optimised with standard stochastic reparamctrisation techniques. Overall, our method inherits strengths both CatNF and Argmax flow. Our experimental results highlight the effectiveness of our approach.
2 BACKGROUND: VARIATIONAL DEQUANTISATION
Dequantisation refers the process of embedding discrete-valued data into continuous space, which allows us to employ density- based models to capture the distribution of the continuous representation. Concretely, let {21,... ........ 27 denote this continuous representation, and describe the observed data, where each 22 represent, e.g. a node in a graph or a token in sentence. Each v, is assumed to be categorical, i.e. 3y E {0, K - 1) for some integer K 1, z can be interpreted as a latent variable, which follows prior distribution p(z). We refer to ((ze (el)))))) as the dequantiser and p(zz)=z) the quantiser. Training can be achieved by maximizing a variational lower bound on the marginal likelihood of the data, iee.
{_11...
logp(x) 2 By(z(=) log p(x[z) p(z)
=: L(x)
q(> (+)
We are interested in the case where the representation 2, can be inferred from 2, alone, so we choose the factorisation p(œ|z) II, p(zalze) and q(z[z)) = II (()))xz), following Lippe & Gavves (2020). In this case, the "optimal" quantiser ((0.)22) can be conveniently computed as:
((2,1x,)) p(a,) ((())))))) q( p(a{) = q( =: p(xele)
arg max [[(])[((x)] p(xx(ae)
(2)
where ((x) denotes the (empirical) data distribution, and p(zz) denotes the estimate of the marginal distribution of each category (which can be obtained by counting and, in the case textual data, this words). This equation shows that the optimal quantiser
corresponds to the unigram distribution can obtained implicitly by applying Bayes' rule with the parametric dequantiser q(2 (x))). The factorisation chose for (():)2) and g(2)2z is crucial for the arg max above to be represented in this
2

---

Published as a conference paper at ICLR 2022
simple form. Without this assumption, the solution will involve combinatorial sum 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:
Σ p(z) + log Σ 2x/=0 (22
p(z)
[(x)= = B)(z((x)
We note that the first term is a constant. Therefore, the expression above implies that accurately modelling the dependencies in æ boils down to learning an expressive prior p(z) and regularising the q(<r)=e)) is deterministic when ((2/(1)) does not overlap with other q (r(). dequantiser q (ze).. (Z)(1z)).
in which case (< (2)) 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(2 (x)), 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 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), dequantisation scheme can be obtained by setting q(,, (xx) = Uniform(2,, (Z,, I2 + 1). The resulting quantisation process is simply [2r], and is deterministic. More generally, g( ((zz]e (a)) can be any distribution on [att, + 1]. See Nielsen & Winther (2020); Hoogeboom et al. (2019) for extensions of the uniform dequantisation scheme.
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 2 E RKK be distributed by q(< (are) with support { : arg max (Ft)a = xe). When the support over the latent space is disjoint, p(zal==) g(=x|==) = [ arg (=) we depict this in Figure 1, Argmax flow (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 distinet dequantised representations despite having similar functions and meaning in language modelling setting. In the naive formulation, Argmax Flow requires the dimensionality to the latent space be the same as the number of the input categories K. To accomodate for larger categorical spaces, the authors suggest binary factorisation, reducing the required latent space dimension to flog, K1; See Figure I, Argmax flow (binary).
Argmax Flow
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 argmax flow, the dequantising distributions g [ = 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 conditional logistic distribution as ( (az)). In this case, the optimal quantiser (xx,) 2x) 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 proc-spoo inequality argument, using the dequantiser as transitional kernel). It can approximate the lossless limit by pushing the bulk of the mass of q( = (m:)) away from each other, but that could potentially lead highly complex and multi-moda empirical distribution over the representation space for p() to approximate.
Next, we consider the case where ((=e(a:t) is a truncated distribution, and as such has the ability to encode the data losslessly while learning 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 (2 to estimate the expectation (Eq. 1). In our case, we want to parameterise this using TRUncated FLow, which we will refer to as TRUFL.
We use 1-) to denote indicator function.
3

---

Published as a conference paper at ICLR 2022
For simplicity, we will drop the dependency on t and 22, but all of the variational distribution is conditioned on the categorical value T/..
If we want to bound scalar distribution between (a, b). and we have 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 from Uniform( (F(a), F())), and then evaluating (u). Note that this method is differentiable, and we use it to sample from our dequantiser.
However, vulti-variate distributions may not simply be truncated at the tails, but rather have 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.1 TRUNCATED LOGISTIC DISTRIBUTION
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:
[()) = (z) a1 =),,
F(=)= =
(-) = log u
(4)
for = E and u E (0, 1), where a the logistic sigmoid function. Additionally, we can parameterise
the width of Uniform(F(aa) F(b)) directly by s, where 0 < s S 1, and the mean of this distribution as m, such that m + 3 < 1 and m > 0. Given a set of statistical parameters 3 E R and m e R for each category, we can impose these constraints by the following reparameterisattion:
m (mm,,
s = 2- min(m, m) (8).
(5)
Then, to sample from Uniform (F(a), F(b)).
m +
u0
20 Uniform 1).
(6)
2
This allows for a simple implementation for truncated logistic distribution that is differentiable (see Appendix A.1 for the derivation of the gradient). We can write the pdf as
5 [=)=) a(1 2),, if < F(=)-m
2 S1
f(z:m, 8)
0,
otherwise.
Like in Lippe & Gavves (2020), the resulting approximate posterior g(2e|2z) requires category specific parameters for the truncation (m(x;) and s(x:)) and any flow applied on top of the truncated logistic (we denote this flow g). During training, conditioned sample 2t with the method described above, and compute the probability q( ((2/22)) (z.). Since decoding requires computing 9 (2,, .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 category conditional scale and shift suffices. Algorithm I details the steps taken for computing 2,, log q (z:), and log (z).
a given category 2t, we can
With this parameterisation of the decoder and the truncated approximate posterior, some q(2, (+t) where It f ft may not have support over the sample 2z- In the extreme case, when q(2 (z)) for all possible assignments of x, had mutually exclusive support, then decoding will always be deterministic, similar to Argmax flow. This is when TRUFL has the flexibility to embed discrete data losslessly in a continuous space. Empirically we find that (<xe/2=) is sparse, but not always deterministic.
Recovering Argmax Flow and CatNF Specific choices of m and 8 in Eq. 6 can recover CatNF and Argmax Flow support for each category. By setting m 0.5 and 1, we perform truncation on the posterior distribution, where we recover CatNF. For Argmax Flow, we can achieve orthant support by setting to 0.25 0.75 for each dimension according to the orthant it assigned based on the binary encoding, and = 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.
4

---

Published as a conference paper at ICLR 2022
Algorithm Truncated Categorical Encoding for timestep
Input: Categorical data Flow g(-.)) Output: 2,, log (=2x [24, log p(z zx) uo Uniform(O, 1)
Begin encoding
a - mmx)) + [uu - 2) s(x1)
z/ < (u)
21 < (((
b End Encoding
for 8; = t K - 1--11 do
P Compute probability of 2 given all possible žt
8 g (2)
log ((x (6x) t- log m())) + log
dzf ds
end for
log (2)) + log ((()))))))) p(xx) - log >,, g<ze (i.) p()))
log computation of the posterior
Algorithm Rejection sampling
Output: æ
z ~ p(z) while væ:
- do
2 ~~ >~p(2=) p(z)
end while
x p(z|z)
p(z)
Ptrune (2)
Figure 2: Left: Pseudo-code for rejection sampling. Centre: The untruncated prior p(z). Right: The prior truncated according to the unioned support of q. We denote this Ptrunc (2).
3.2 REJECTION SAMPLING
When sampling from the model, the standard process is ancestral sampling: p(z), and then p(a) ()) However, in our model, the sample 2, may not have support under any the q( (xz). resulting an undefined optimal p(zz|=z). As mentioned, for any of the mixture components. This is equivalent to redefining a new prior Ptrame (z) proportional to p(z) such that its support matches the unioned support of q(> (2) (z) and renormalising. The existing ELBO is still a valid fower-bound of such a distribution:
reject the samples that have support
log p(x|2 2) p(z)
log p(a|2) p(z)
$ / p(x]))) (2) dz,
L = (=1())
< E (x(l))
(8)
9(2(x=)
where
7 = p(z)
> 00 dz 1.
Eq. 9 in principle can be estimated by counting the frequency of accepted samples. For one of our experiments, the overall rejection rate the end of training about 0.14, suggesting that rejection sampling from a trained prior is not inefficient (see Section 5.1). Using the decoder as accepting heuristic, we find empirically that this results in more valid samples in constrained tasks that require sampling from the model.
4 RELATED WORK
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 way to 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 take most of the modelling complexity.
5

---

Published as a conference paper at ICLR 2022
Table 1: Results on the graph coloring problem (Lippe & Gavves, 2020). SMALL are graphs of size 10 I | I 20, and LARGE 25 I [V] I 50. All results are attained using the CatNF codebase, and averaged across random seeds. Results in the rounded box are using different set of hyperparameters than the ones used in CatNF.
SMALL
LARGE
Method
Validity
bpd
Validity
bpd
RaNg-Largest first
93.41% +0.42%
0.68 +0.01
71.32%
+0.77%
0.43 =0.01
66.80% +1.14%
0.45 =0.01
CatNF
94.56% 10.55%
0.67 10.00
68.06%
=0.04%
0.45 =0.00
Argmax (thresh.)
94.81% +0.37%
0.66+0.00
63.65%
+0.24%
0.46 =0.00
No rejection sampling Rejection sampling
95.40% +0.35%
18%10%
+.000%
0.65 +0.01
0.45 +0.00
95.90% +0.29%
74.20%
+0.01%
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; 8: Huang et al., 2017). One common way for learning sequential data is a graph-structured flows which can be used for modelling such data, and there have been flow-based models developed for molecule generation (Madhawa al., 2019; Shi al., 2020). Ziegler & Rush (2019) and Huang et al. (2018) also introduce normalising flows that can learn multi-modie target distributions, which Ziegler& & Rush (2019) show important in modelling text data.
autoregressive prior (Ziegler & Rush, 2019). Liu et al. (2019) introduces
Optimal decoder In the case of categorical data with 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 flow (Hoogeboom al., 2021) are special case in which the encoding of the categorical variable results in a latent variable that has deterministic decoding, which Nielsen 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-suppon distributions in latent variable models While normalising flows (Papamakar- ios 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 light-tailed distribution, achieved by performing a shift and scaling of the base uniform distri- bution. 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- -upport kernels as the basis of their posterior distributions, and modify the decoder so that the results are deterministic, but introduce gradient estimations in order train the model. In the latter two cases, rejection sampling was used in order to attain samples for inference.
5 EXPERIMENTS
In this section, we present experiments on TRUFL. See Appendix C for implementation details.
5.1 GRAPHI COLORING
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 edge do not have the same colour, which provides meaningful benchmark for evaluating how well a model learns unknown constraints from the data. We evaluate the Validity of the generated samples conditioned given graph.

---

Published as a conference paper at ICLR 2022
Table 2:
Performance on molecule generation trained on Zinc250k (Irwin et al., 2012), calculated on 10k samples and averaged over 4 random seeds.
Method
Validity
Uniqueness
Novelty
Constrained generation
JT-VAE (Jin et al., 2018)
100%
100%
100%
GraphAF (Shi et al., 2020)
68%
99.10%
100%
R-VAE (Ma et al., 2018)
34.90%
100%
GraphNVP (Madhawa al., 2019)
42.60%
94.80%
100%
Argmax flow (thresh.)
78.36% +4.93%
100%
99.99%
CNF (Lippe & Gavves, 2020)
83.41% +2.34%
99.99%
100%
+ Sub-graphs
96.35% +2.58%
99.98%
99.98%
No rejection sampling
84.46%++4.24%
100%
100%
4 Sub-graphs
97.51 7.3%%
99.98%
99.995
+ Rejection sampling
85.01 1..0%
100%
100%
Sub- Sub-graphs
97.87 0.006
100%
100%
For baselines, we present results from CatNF (Lippe & Gavves, 2020), and implemented the thresh- olding version of Argmax flow. 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 the aggregate posterior during training, expect the jection rate of sampling from the truncated prior drop as training progresses. We
0.45
Seed 1
Average rejection rete
3.40
Seed
Seed
measure the rejection rate when sampling from the model and report its average rate of rejection as training progresses. While the model conditioned on graphs of dif- ferent sizes during sampling, we look at the average rate of rejection go obtain a measure of the efficiency of our rejection sampling scheme. Figure 3 suggests that the prior learns to fit the support of the ag- gregate posterior better during training, as the rejection rate reduces from about 0.45 to about 0.14 at the end of training. This
0:25
0.30
0.35
2.30
Training iterations
leb
Figure 3: Rejection rate vs training iterations. The re-
jection rate when sampling for evaluation for the graph colouring task reduces during iteration, indicating (I) the rejection rate is not intractably high, and (2) the prior tries to match the support of TRUFL over time.
corresponds roughly losing log 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.2 MOLECULE GENERATION
Molecule generation is another task commonly used as generative modelling benchmark (Jin et al., 2018; Shi et al., 2020; Ma et al., 2018; Madhawa et al., 2019). The atoms constituting the molecules naturally represented by tokens that 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 more useful representation than Argmax Flow. Moreover, given the
7

---

Published as a conference paper at ICLR 2022
Table Results on character-level and word-level language modelling, an average across 3 different random seeds. Results reported with indicates results attained with the Argmax flow codebase. indicates results attained with the CatNF codebase. All other results are from the CatNF paper.
CHARACTER
WORD
Model
PTB (char.)
Text8
Wikitext103
LSTM
1.28 10.01
1.44 10.01
4.81 10.05
Latent NF (Ziegler & Rush, 2019)
1.30 10.01
1.61 1:002
6.39 10.19
Categorical NF (Lippe & Gavves, 2020)
1.27 10.01
1.45 10.01
5,43 10.09
Argmax Flow (Hooggeooom et al., 2021)
11.66 10.01
11.39 10.00
*5.42 10.01
Ours
*1.26 +0.02
$1.40 +0.01
*5.35 +0.01
Table 4: Results on word-level language modelling, each sentence is modelled as a separate data point. CatNF, Argmax Flow, and TRUFL results are averaged across 4 random seeds. Setting and baseline LSTM provided by Kim et al. (2019).
Model
Dim.
PPL
NLL
Recon.
KL
LSTM
86.2
4.46
CatNF
12
139.7
+
4.93
+ 0.02
1.01
4.18
Argmax Flow (binary)
14
242.7
+ 2.7
5.49
+: 0.01
0.00
5.75
Ours
12
143.6
+ 4.9
4.97
+ 0.03
1.47
3.71
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 flow benchmark uses the same thresholding scheme. Rejection sampling in this case requires rejecting invalid edges and vertices both, and resulted in longer sampling times. The results are in Table 2 Again, we find that TRUFL consistently improves upon CatNF in terms of the validity of the generated samples, which can 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.3 LANGUAGE MODELLING
One can view these variational dequantisation methods when applied to language modelling method of using stochastic embeddings: each token is represented by a distribution over the embed- ding 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 per- form experiments on character-level and word-level language modelling. For word-level language modelling, many tasks require large number of categories (~<10)). The naive thresholding method will result in a large latent space, no different from using one-hot representation, 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 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 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 specified by Kim et al. (2019), which also provides an LST baseline. Unlike standard PTB results, this benchmark treats each sentence
8

---

Published as a conference paper at ICLR 2022
throe
singapore
Figure 4: Scatter plot of t-SNE embeddings of samples from Armgax Flows (left) and TRUFL (right). Since t-SNE reveals the relative proximities of the embeddings, we should note that the visualisation here reveals what clusters of embeddings are close others relative to the other words in the plot.
as i..dd, instead of treating the entire dataset as continuous string of text. We provide further details of the architecture in Appendix C.3.
The reconstruction loss for TRUFL 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 (( Sampling 200 samples from each category, we first normalise the mean and standard deviation across all 10, 000 x 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 lower dispersion on average. Since the prior 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 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 ((ze(2)) for each word, and performed t-SNE over all 200 10 vectors. As expected, Argmax flow 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 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 f flew: y and wa 1 ked:wal k 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 functional representation of these words.
CONCLUSION
In this paper, we propose TRUFL, flow based dequantiser with truncated support for stochastically embedding categorical data in a continuous space. This allows us to employ density-based models like normalising 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 rugo-regressive models that directly model the categorical space and latent variable models.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENTS
We would like to thank Yikang Shen and Christos Tsirigotis for their input and suggestions during the course of this work. Phillip Lippe also gracious in helping understand the specifics of the CatNF codebase and running the experiments.
REFERENCES
Yoshua Bengio, Réjean Ducharme, Pascal Vincent, and Christian Janvin. A neural probabilistic language model. The journal of machine learning research, 3:1137-1155, 2003.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 1okkiv//300 1308.3432, 2013.
John Adrian Bondy, Uppaluri Siva Ramachandra Murty, et al. Graph theory with applications, volume 290. Macmillan London, 1976.
Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from continuous space, arXiv preprint arXiv:1511.063 2015.
Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel. Variational lossy autoencoder. arXiv preprint arXiv: 10Xiv:/611.022 1.02731, 2016.
Xuelu Chen, Michael Boratko, Muhao Chen, Shib Sankar Dasgupta, Xiang Lorraine Li, and Andrew McCallum. Probabilistic box embeddings for uncertain knowledge graph reasoning. arXiv preprini arXiv:2104.04597, 2021.
Shib Dasgupta, Michael Boratko, Dongxu Zhang, Luke Vilnis, Xiang Li, and Andrew McCallum. Improving local identifiability in probabilistic box embeddings. Advances in Neural Information Processing Systems, 33:182-195 192, 2020.
Tim R Davidson, Luca Falorsi, Nicola De Cao, Thomas Kipf, and Jakub M Tomezak. Hyperspherical variational auto-encoders. arXiv preprint arXiv: 1ee::88040000 000400891, 2018.
Michael Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit reparameterization gradients. arXiv preprini arXiv: 18X.v:1805.08445 2018.
Mathieu Germain, Karol Gregor, lain Murray, and Hugo Larochelle. Made: Masked autoencoder for distribution estimation. In International Conference on Machine Learning, pp. 881-889 PMLR, 2015.
Joshua Goodman. Classes for fast maximum entropy training. In 2001 IEEE International Conference on Acoustics, Speech, and Signal Processing. Proceedings (Cat. No. O1CH37221), volume 1, pp- 561-564. IEEE, 2001.
Jonathan Ho. Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. In International Conference on Machine Learning, pp. 2722 2730. PMLR, 2019.
Emiel Hoogeboom, Jorn Peters, Rianne van den Berg, and Max Welling, Integer discrete flows and lossless compression. Advances in Neural Information Processing Systems, 12:1213-12144
2019.
Emiel Hoogeboom, Taco Cohen, and Jakub Mikolaj Tomczak. Learning discrete distributions by dequantization. In Third Symposium on Advances in Approximate Bayesian Inference, 2020.
Emiel Hoogehoom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Towards non-autoregressive language models. arXiv preprint arXiv:2/02.05372 2021.
Chin-We Mein-Wei Huang, Ahmed Touati, Laurent Dinh, Michal Drozdzal, Mohammad Havaei, Laurent Charlin, and Aaron Courville. Learnalle explicit density for continuous latent space and variational 2017.
inference. arXiv preprint
10

---

Published as a conference paper at ICLR 2022
Chin-War Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. In International Conference on Machine Learning, pp.
PMLR, 2018
John J Irwin, Teague Sterling, Michael M Mysinger, Erin S Bolstad, and Ryan Coleman. Zinc: a free tool discover chemistry for biology. Journal of chemical information and modeling, 52(7): 1757-1768, 2012.
Priyank Jaini, Ivan Kobyzev, Yaoliang Yu, and Marcus Brubaker. Tails of lipschitz triangular flows. In International Conference on Machine Learning, pp. 4673 PMLR, 2020.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In International conference on machine learning, pp. PMLR, 2018.
Yoon Kim, Chris Dyer, and Alexander M Rush. Compound probabilistic context-free grammars for grammar induction. arXiv preprint arXiv:1906.1022: 2019.
Diederik P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions.
arXiv preprint arviv: 0000000339 2018.
Diederik P Kingma and Max Welling, Stochastic gradient vb and the variational auto-encoder. In Second International Conference Learning Representations, ICLR, volume 19, pp. 121, 2014.
Xiang Li, Luke Vilnis, Dongxu Zhang, Michael Boratko, and Andrew McCallum. Smoothing the ge-
ometry of probabilistic box embeddings. In Infernational Conference on Learning Representations. 2018.
Alexandra Lindt and Emiel Hoogeboom. Discrete denoising flows. In ICMI. Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, 2021.
Phillip Lippe and Efstratios Gavves. Categorical normalizing flows via continuous transformations. arXiv preprint arXiv:2006.0979) 09790, 2020.
Jenny Liu, Aviral Kumar, Jimmy Ba, Jamie Kiros, and Kevin Swersky. Graph normalizing flows. arXiv preprint 1r05v:1905.133 2019.
Tengfei Ma. Jie Chen, and Cao Xiao. Constrained generation of semantically valid graphs via regularizing variational autoencoders. arXiv preprint 180.02280, 2018.
Kaushalya Madhawa, Katushiko Ishiguro, Kosuke Nakago, and Motoki Abe. Graphnvp: An invertible flow model for generating molecular graphs. arXiv preprint arXiv: /905. /600, 2019.
Mitchell Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building a large annotated corpus of english: The penn treebank. 1993.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv: 1609:076000773 2016.
Tomas Mikolov, Armand Joulin, Sumit Chopra, Michael Mathieu, and Mare' Aurelio Ranzato. Learning longer memory in recurrent neural networks. arXiv preprint arXiv: 14X2::34...7753. 2014.
Andriy Mnih and Geoffrey E. Hinton. scalable hierarchical distributed language model. Advances in neural information processing systems, 21:1081-1088 2008
Frederic Morin and Yoshua Bengio. Hierarchical probabilistic neural network language model. In International workshop on artificial intelligence and statistics, pp. 246-252 PMLR, 2005.
Kevin P Murphy. Machine learning: a probabilistic perspective. MIT press, 2012.
Didrik Nielsen and Ole Winther. Closing the dequantization gap: Pixelenn as single-layer flow. arXiv preprint arXiv:2002.025444 2020.
Didrik Nielsen, Priyank Jaini, Emiel Hoogeboom, Ole Winther, and Max Welling. Survae flows: Surjections to bridge the gap between vaes and flows. Advances in Neural Information Processing Systems, 33, 2020.
11

---

Published as a conference paper at ICLR 2022
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv:1912.02762, 2019.
Jeffrey Pennington, Richard Socher, and Christopher Manning. Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP). pp. 1533,1543, 2014.
Daniil Polykovskiy and Dmitry Vetrov. Deterministic decoding for discrete data in variational autoencoders. In International Conference Artificial Intelligence and Statistics, pp. 3046-3056. PMLR, 2020.
Yikang Shen, Shawn Tan, Chrisopher Pal, and Aaron Courville. Self-organized hierarchical softmax. arXiv preprint wixiv:1777.08.. 2017.
Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang. and Jian Tang. Graphaf: flow-based autoregressive model for molecular graph generation. arXiv preprint arX1v:2001.09382, 2020.
Lucas Theis, Aäron van den Oord, and Matthias Bethge. note on the evaluation of generative models. arXiv preprint arXiv:1511.0184/ 2015.
Jakub Tomezak and Max Welling. Vae with vampprior. In International Conference Artificial Intelligence and Statistics, pp. 1214-1223. PMLR, 2018.
Dustin Tran, Keyon Vafa, Kumar Agrawal, Laurent Dinh, and Ben Poole. Discrete flows: Invertible generative models of discrete data. Advances Neural Information Processing Systems, 32: 14719-14728, 2019.
Alexandre Verine, Benjamin Negrevergne, Fabrice Rossi, and Yann Chevaleyre. On the expressivity of bi-lipschitz normalizing flows. arXiv preprint arXiv:2107.07232, 2021.
Jiacheng Xu and Greg Durrett. Spherical latent spaces for stable variational autoencoders. arXiv preprint ar 1ix: 100. 2018
Zichao Yang, Zhiting Hu, Ruslan Salakhutdinov, and Taylor Kong-Kirkpatrick. Improved variational autoencoders for text modeling using dilated convolutions. In International conference on machine learning, pp. 3881-3890 PMLR, 2017.
Zachary Ziegler and Alexander Rush. Latent normalizing flows for discrete sequences. In Interna- tional Conference on Machine Learning, pp. 7673-7682. PMLR, 2019.
12

---

Published as a conference paper at ICLR 2022
A
APPENDIX
A.1
GRADIENT ESTIMATES OF a AND b
b)2;a,
b-a 0,
a < otherwise
(10)
Let
a (b a) u Uniform (0, 1)
(11)
(12)
(13)
b
[.E=[ E Uniforr((.)) Vaf(a + (b a).
(14)
E Uniform((.1) [[1 (V+f(()))
(15)
[.R=[[(=)] = E Uniffrrm((). [v(=[(()))
(16)
A.2 SOFT DEQUANTIZATION
Let denote the data (e.g. a sentence) and L denote the dimensionality of the data; Lee x E C and L Z+. x can be dequantized by sampling 2j ~ ((2z)æz)) for I.... L. When {g(2)12;; = e E C) has the same support, we call it soft dequantization. We write p(zz|zz) and p(2) to denote the quantizer and the prior over the dequantized data. We consider the following ELBO
Σ log {(2g(æ,)
p(0z)1z) + log logr(=)
f(x, L) = E (195))))))
(17)
Theorem (Maximizer of ELBO, dequantization). Assume q(I, L) > 0 for all x E C² and L E Z+.
and that 11 LL are conditionally independent given L. Then
By(a (log ((([L)) L)] > Egte, [A]
(18)
as long as the densities q(2j xz c) for different categories a E C have the same support. That is, soft dequantization with shared support is suboptimal.
Proof. The optimal quantizer is
20 Σ
a(0z(ız) x ((z5)a(2)123))
where a(x;) x
Σ aaa' = =z L)
(19)
L=1j'=1
Plugging it into the ELBO gives
log {(zد|ع))
1
p(2)
L(x, L) = BII ((1gla))
+ logp(z) {(2b(1,)
= ((1
log ((x,) + log
II, ((2;)
(20)
where ((2,) =Z)))=or ((a)))(zz)())
On the other hand, we can rewrite the negentropy of the data as
((x[LL)
(og Ey(22,)
Σ
log ((x;) + log
(21)
11 g(x,))
13

---

Published as a conference paper at ICLR 2022
Now to compare the second term with that of the ELBO we let ((|[)) -
((yloz))))))))
For simplicity, we denote by T=l22) I1 ((2g)1())))) (z;)) the transition kernel applied to either ((x(L) or [T (xz)). Then
log q(1) L) I1 {
g(x) (z)((x)
L)
log
(22)
[I ((03)) T(2)²)
[T ((x,)) T()aa)
log
(23)
[1l ((x))1)))))))))))))))) T(])))
5 By(2,L)
loog
By(x|==L)
(24)
(<x)))
9(2)) ((=[L)
log ((([L)
-Eg(z,L)
log
Bo(a,L)
(25)
II, 9(2,)
which is greater or equal 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 #1,..., implies there exist some x² and x² s.t. II, ( ) + ((=*|L) for i E {1, 2]. Now, since T( (x) > 0 over the shared support by assumption, for almost all and I., g( 2, L) > 0 for i E (1, That is, with non-zero probability the argument to the convex function log has a different value from its mean. This implies the inequality strict and concludes the proof.
B FURTHER RELATED WORK
Consequences of a binary encoding The binary encoding scheme of Argmax flow bears some resemblance to prior work on hierachical softmax schemes reduce the computational footprint of the large softmax layer (Goodman, 2001; Morin & Bengio, 2005; Mnih & Hinton, 2008; Shen al., 2017). Specifically, Mnih Hinton (2008) proposed an analogous scheme for scaling up the softmax layer by representing words as leaves in tree, and binary encodings as traversal of such tree. However, the partitioning of the [log, K" -dimensional space into K equal categories will only be exact if is power of 2. Otherwise, K 2 2og, K) orthants of the embedding space will 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 arbitrary, the similarity between words/categories is not reflected in the embedding space as is a learned distributed representation.
C EXPERIMENTS & IMPLEMENTATION DETAILS
We incorporated the codebases from the following sources:
1. Lippe & Gavves (2020): htt00//:../pg/g///gt/...co///o//
2. Hoogeboom et al. htt.:///tttt////ttaoa...oo/aaaaa
flows
3. Kim et al. (2019): httppp///tttaa..oom///aaacctttep
C. TOY EXAMPLE: Two CATEGORY JOINT DISTRIBUTION
In this simple synthetic example, demonstrate case in which TRUFL provides advantage CatNF and Argmax flow. Our goal is to model joint probability over 2 categorical random variables, each with categories. This joint distribution is described in Table 5.
14

---

Published as a conference paper at ICLR 2022
0.
75-
10
25
80-
00
25-
589
(a) CatNF
(b) TRUFL
(c) Argmax Flow
Figure Embeddings (z:z( of toy example for each method.
Table Left: The joint distribution for æ1 and 22 for the toy example. Right: Log probability over the data of the model, and the breakdown of the ELBO.
D2
Cat.
0
1
2
3
Method logp(z)
E [log p(+|2)]
log g(alx) p(=)
0
1
0
1
0
CatNF -2.000
-1.064
0.967
21
R
Argmax
-1.915
-0.687
1.263
2
0
0
TRUFL.
-1.983
0.000
2.098
3
0
Notice that under this distribution, p(21|02 = 0) = p(21|22 = 2) and p(x2|=1 = 0) = p(22 [ = 2).
For the purposes of this subsection will say categories 0 and 2 have equal functionality, while categories and 3 are similar in functionality, but not equal.
To model this distribution, perform variational dequantisation using the above methods, and learn a prior over the latent variables. Conditioned on I;, the dequanti- sation method produces distribution over dimensional space - (( (z;)). It is this distribution we plot in Figure 5.
the discrete variables T1,
In both CatNF and TRUFL,, the distribution q<=//>; and ((2,/1; = 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 the distributions of these two categories will have to be modelled by the prior, which evidently resulted in higher log probability in this particular case. In our implementation of Argmax flow for this task, the training eventually resulted 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 latent space. When this happens, the closeness in the components again result in decoding becoming non-determiiistic.. 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 see that the support and distribution for æ; and v; overlap the most, while the overlap with x; = and x; = is not 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.2 HYPERPARAMETERS FOR LARGE GRAPH COLOURING TASK
C.3 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', use
15

---

Published as a conference paper at ICLR 2022
Param.
Value
batch.size
128
encoding dim
2
coupling www.nn..ummm..mmmmmmmm.mm....... num num.foows optimizer
10
4
learning.rate
3e-4
Table New arguments for the LARGE graph colouring task we use in our experiments.
nomenclature from Ziegler & Rush (2019). Autoregression in time is governed by an LSTM, and autoregression hidden is governed by MADE model (Germain et al., 2015), conditioned on said LSTM. We use mixture of 4 logisties per-dimension (Ho et al., 2019) in order to model the multi-modality of the distribution over 2, which Ziegler & Rush (2019) suggests is prevalent text data for characters.
16


Reviews:
[{"review_id": "ozWMkHg-8dn", "reviewer": "Reviewer_zxth", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is mostly well written and easy to follow, although I have some reservations about the motivation:\n\n1. I don't think the authors properly motivate the use of the normalizing flow to begin with. Is it really better to have the flow as compared to increasing capacity for m and s? I think this is a very natural question, which is neither posed nor answered in the paper; and that needs answering to better motivate the paper.\n\n2. Part of the motivation is enabling the dequantized variable's dimension to not depend (logarithmically) on the number of categories. It is actually not clear to me that this dependence is particularly harmful. For example, one might expect that categories which are semantically close to each other have corresponding nearby dequantized intervals. However, it might not be possible to achieve this: consider a case with 3 categories, A, B, and C which are though of \"semantically equidistant\". There is no way to have 3 real intervals satisfying this requirement, and I suspect this might make learning the functions $m$ and $s$ harder. In contrast, if in the same example a 2-dimensional dequantized space was used, one could actually obtain the desired symmetry. I think this comment bears some philosophical resemblance to kernel methods, where the intuition is that simple classifiers might do better on high-dimensional representations. In other words, the fact that one can obtain a one-dimensional dequantization does not automatically imply that one should, and I believe the authors should make a stronger point for this choice.\n\n3. Finally, while dequantization is a sensible approach, I think that it should be at least empirically verified that dequantizing and using a continuous model actually outperforms using a flexible discrete model to begin with. In fairness to the authors though, I think this is more of an issue I have with dequantization itself rather than the paper being reviewed.\n\nAs for novelty, the method is novel as far as I am aware.\n\nFinally, I did not find the experiments particularly convincing: while the proposed method does seem to be slightly better than the baselines at the evaluated tasks, the improvements seems small and do not in my view compensate for the issues I have with the motivation. Additionally, only averaged results over runs are reported, and given how close the numbers are along with the lack of error bars, it is hard to assess whether the improvements are actually significant.\n\nMinor remarks:\n\n-In algorithm 1, $\\hat{z}_t$ is used in the denominator for the change-of-volume term, but $\\hat{z}_t$ is not defined.\n\n-The notation in equation 7 should be changed to either $\\tilde{f}(z;m,s)$ or $\\tilde{f}(z|m,s)$ to more clearly differentiate between the different types of inputs.\n\n-Table 2 should be referenced in the text in section 5.2.\n\n==========================================================================================================\n\nUPDATE AFTER REBUTTAL\n\n==========================================================================================================\n\nI have increased my score after seeing the authors' updates regarding motivation and the promised ablations regarding point 3 above.", "summary_of_the_paper": "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.", "summary_of_the_review": "While this paper proposes a simple idea for dequantization, I am not convinced the proposed method is needed in the first place, and would need to see much stronger empirical evidence to be convinced otherwise.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "RnqJqu2czbq", "reviewer": "Reviewer_Pdpw", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Learning how to handle discrete data in continuous deep neural networks is an important task, and contributions in this direction are very useful. However, I am a bit concerned with the authors' claim that TRUFL better achieves the two goals mentioned in their introduction: (i) easily learnable, and (ii) lossless. First, it is unclear whether TRUFL is more \"lossless\" than CatNF in practice. While TRUFL uses truncation to zero out the probability mass beyond a bounded range, there could still be major overlap between the latent space of different categories. And compared to that, truncating the low-probability regions might be less effective in making the dequantization lossless. To make the dequantizations more lossless, another simple choice could be to make distributions on z farther away from each other (Figure 1 CatNF, making the three gaussian farther away from each other).\n\nIn fact, I can imagine in certain cases we don't need lossless dequantization, as some variation of the data might be irrelevant to a task.\n\nWhile the proposed TRUFL model may still perform lossy dequantization, according to the claim in the introduction, the benefit of TRUFL could be it's easily learnable. However, due to the similarity with CatNF and the additional need for rejection sampling, it is not very clear to me why TRUFL can be learned easily. To summarize, it is unclear to me why truncating the latent space could simplify the learning problem itself.\n\nSince the proposed approach uses rejection sampling, it would be nice to also compare the inference time of different models, just to see how much the sampling procedure influences the model's overall efficiency.", "summary_of_the_paper": "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. ", "summary_of_the_review": "I tend to vote for rejection because it is unclear to me why truncating the latent distribution could necessarily improve the dequantization quality. Although the authors provide two explanations in the introduction, they are not very well justified. Apart from the truncation technique, the proposed model doesn't seem to differ too much from existing approaches.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1ih6joVD9Yq", "reviewer": "Reviewer_QH7t", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Pros:** The paper is very well written and the method is explained clearly, concisely and in appropriate detail. The use of examples, figures and algorithm boxes at appropriate places makes the paper easy to follow and understand. Furthermore, the problem considered in the paper is a valid and interesting problem for modelling discrete data using NFs. Multiple papers have come out in the past year or so addressing this problem, each with their own sets of solutions (and restrictions) and the present paper proposes a new method that alleviates some of the drawbacks of CatNFs and Argmax flows. The experimental analysis by the authors consider a diverse set of problems from graph coloring, molecular generation, and language modelling with reasonable performances. \n\n**Cons:** I found the empirical results to be mixed for TRUFL as compared to Argmax flows and CatNF. While, this in itself not a huge issue, I am interested to understand what might explain this considering that the method basically tries to address the drawbacks of both these approaches. I also wonder if it might be possible to use synthetic examples to bring out these advantages for TRUFL.", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper is well written and explored. The empirical analysis is reasonable. Overall, I found the paper quite interesting.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a variational dequantization method for categorical data, based on flows with learned truncated support. The problem has been studied before, but the paper makes it clear how the proposed method differs from existing ones. The method is empirically evaluated on a large variety of diverse tasks.

The reviews were initially borderline. In general, the reviewers did not identify major quality of technical issues with the paper, and appreciated the clarity of writing. On the other hand, the reviewers were not fully convinced by the motivation or the empirical performance of the proposed method. After discussion with the authors, some concerns were allayed (especially regarding motivation) and all three reviewers decided to recommend weak acceptance.

Seeing as there are no major technical or quality issues with the paper, and the paper is clearly written and well executed, I'm leaning towards recommending acceptance, although some doubts remain about the significance of the contribution.
Decision: Accept (Poster)