Semi-Discrete Normalizing Flows through Differentiable Tessellation
Authors: Ricky T. Q. Chen
Brandon Amos
Maximilian Nickel
Venue: NeurIPS 2022 Accept (2022)
Keywords: deep probabilistic modeling
normalizing flows
dequantization
disjoint mixture modeling

Abstract: Mapping between discrete and continuous distributions is a difficult task and many have had to resort to heuristical approaches. We propose a tessellation-based approach that directly learns quantization boundaries in a continuous space, complete with exact likelihood evaluations. This is done through constructing normalizing flows on convex polytopes parameterized using a simple homeomorphism with an efficient log determinant Jacobian. We explore this approach in two application settings, mapping from discrete to continuous and vice versa. Firstly, a Voronoi dequantization allows automatically learning quantization boundaries in a multidimensional space. The location of boundaries and distances between regions can encode useful structural relations between the quantized discrete values. Secondly, a Voronoi mixture model has near-constant computation cost for likelihood evaluation regardless of the number of mixture components. Empirically, we show improvements over existing methods across a range of structured data modalities.

Semi-Discrete Normalizing Flows through Differentiable Tessellation
Ricky Q. Chen
Brandon Amos
Maximilian Nickel Meta AI maxn@meta, com
Meta AI rtqichen@meta.com com
Meta AI bamosOmetta ta com
Abstract
Mapping between discrete and continuous distributions is difficult task and many have had to resort to heuristical approaches. We propose daseellation-bael ap- proach that directly learns quantization boundaries in continuous space, complete with exact likelihood evaluations. This is done through constructing normalizing flows on convex polytopes parameterized using a simple homeomorphism with an efficient log determinant Jacobian. We explore this approach in two application settings, mapping from discrete to continuous and vice versa. Firstly, a Voronoi dequantization allows automatically learning quantization boundaries in multidi- mensional space. The location of boundaries and distances between regions can encode useful structural relations between the quantized discrete values. Secondly, a Voronoi mixture model has near-constant computation cost for likelihood eval- uation regardless of the number of mixture components. Empirically, show improvements over existing methods across a range of structured data modalities.
Introduction
Likelihood-based models have seen increasing usage across multiple data modalities. Across variety of modeling approaches, the family of normalizing flows stands out as a large amount of structure can be incorporated into the model, aiding its usage in modeling wide variety of domains such as images [7, 23), graphs [29], invariant distributions [2, 25] and molecular structures [42]. However, the majority of works focus only continuous functions and continuous random variables. This restriction can make it difficult to apply such models to distributions with implicit discrete structures, e.g. distributions with discrete symmetries, multimodal distributions, distributions with holes.
In this work, we incorporate discrete structure into standard normalizing flows, while being entirely composable with any other normalizing flow. Specifically, we propose a homeomorphism between the unbounded domain and convex polytopes, which are defined through a learnable tessellation of the domain, i.e. set of disjoint subsets that together fully cover the domain. This homeomorphism is cheap to compute, has a cheap inverse, and results an efficient formulation of the resulting change in density. In other words, this transformation is highly scalable-at at least computationally.
Our method has the potential to be useful for variety of applications. Firstly, the learned tessellation naturally allows defining a dequantization method for discrete variables. This allows likelihood- based models that normally only act on continuous spaces, such as normalizing flows, to work directly on discrete data in flexible choice of embedding space with the ability to capture relational structure. Secondly, we take each convex polytope as the support single mixture component, then the full tessellation defines a mixture model with disjoint components. This allows to scale to mixtures of normalizing flows while retaining the compute cost of single model. Following semi-discrete optimal transportation [37], which defines couplings between discrete and continuous measures, we refer to our models as semi-discrete normalizing flows that learn transformations between discrete and continuous random variables.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

p(æ)
p(()))
Figure We propose an invertible mapping f between R&D and a convex polytope, which is parameterized based on differentiable Voronoi tessellation of DRD This mapping adds discrete structure into normalizing flows, and its inverse f and log determinant Jacobian can both be efficiently computed.
Preliminaries
Normalizing Flows This family of generative models typically includes any model that makes use of invertible transformations f to map samples between distributions. The relationship between the original and transformed density functions have closed form expression.
Pa p=((())) (a() det @f(m) De
For instance, Ps can be simple base distribution while is learned that the resulting Pa is close to some target distribution. Many different choices of f have been discussed in the literature, leading different trade-offs and cases [24, 36].
Furthermore, if the domain and codomain of f are different, then Ps and Pz can have different supports, i.e. regions where probability is non-zero. Currently, existing designs of that have this property act independently for each dimension, such as the logit transform [7]. To the best of our knowledge, the use of general support-modifying invertible transformations have not been discussed extensively in the literature.
Dequantization In order to model discrete data with density models, a standard approach is to combined with dequantization methods. These methods provide a correspondence between discrete values and convex subsets, where discrete variable y E Y' randomly placed within a disjoint subset of R There is also correspondence between the likelihood values of the discrete random variable and the dequantized random variable [44].
Let Ay denote the subset corresponding to y and let g(øly)) be the dequantization model which has bounded support in Ay. Then any density model p(z) with support over RD, satisfies
log)(y) 2 Bm-q((l))) (1
[log("(mE (IEEA))(()) - log ((=]v)) = Em-a(=/vv) p(z) log (<>|v)] (2)
Thus with an appropriate choice of dequantization, maximizing the likelihood under the density model p(x) equivalent maximizing the likelihood under the discrete model p(y).
In existing dequantization methods, the value of and the choice of subsets A, are entirely dependent on the type of discrete data (ordinal non-ordinal) and the number of discrete values |>/ in the non-ordinal case. Furthermore, the subsets A, do not interact with one another and are fixed during training. In contrast, we conjecture that important relations between discrete values should be modeled as part of the parameterization of Ay. and that it' be useful to be able to automatically learn the boundaries of Ag based on gradients.
Disjoint mixture models Building mixture models is one of the simplest methods for creating more flexible distributions from simpler one, and mixture models can typically be shown to be universal density estimators [12]. However, in practice this often requires a large number of mixture components, which quickly becomes computationally expensive. This because evaluating the likelihood under a mixture model requires evaluating the likelihood of each component. To alleviate this issue, Dinh et [8] recently proposed to use components with disjoint support. particular, let {Ap kk= be disjoint subsets of D such that each mixture component is defined on one subset and

---

has support restricted to that particular subset. The likelihood of the mixture model then simplifies to
K
K
p()) = Σ p(a)k))((k) = Σ 1 p(2
(k) = p(z)k p(=))p(kk g(x))
(3)
where g RD_ (I,..., K} is set identification function that satisfies æ E Agta)- This framework allows building large mixture models while no longer having a compute cost that scales with K. In contrast to variational approaches with discrete latent variables, the use of disjoint subsets provides an exact blkeelhod and not a lower bound.
Voronoi Tessellation for Normalizing Flows
We first discuss how parameterize each subset as convex polytope, and in the next part, discuss constructing distributions on each convex polytope using a differentiable homeonorphism.
Parameterizing disjoint subsets We separate the domain into subsets through a Voronoi tessellation [48]. This induces a correspondence between each subset and a corresponding anchor point, which provides a differentiable parameterization of the tessellation for gradient- based optimization.
be set of anchor points R., The Voronoi cell for each anchor point Let X ......... kk
V& - fxr E E
||
21 c =1| , l...,,.... K}:
(4)
Lee it defines a subset containing all points which have the anchor point as their nearest neighbor. Together, these subsets ] form a Voronoi tessellation of D. Each subset can equivalently be expressed in the form of a convex polytope,
V&= {o E RR a. x < b, vi + k), where a, 2(x; xxx) b, || || 11
(5)
For simplicity, we also include box constraints so that all Voronoi cells are bounded all directions.
V = fx E R
vi < (!) <<<<<h
(6)
Voronoi cell
constraints
Thus, the learnable parameters of this Voronoi tessellation are the anchor points {ø1,.... Φk}. and the box constraints C₁; C. RD. These will be trained using gradients from invertible transformation and resulting distribution defined in each Voronoi cell, which we discuss next.
Invertible mapping onto subset Within the normalizing flow framework, defining a probability distribution on bounded sup- port is equivalent defining invertible transformation that maps from the unbounded support onto the appropriate support. Figure 2 illustrates the transformation of mapping onto shaded region.
Input
Anchor
Given cell V2 for some k E (1....: K], we construct an invertible mapping S& RDD - V& by following 2 steps. Let x RD, First, æ Sks the anchor point of Ve, we simply set (x) a. Otherwise:
Ray
Intersection
x22
Output
Figure An illustration.
Determine where the ray starting from ae the direction of x intersects with the boundary of Vi-
First define the direction 5n (x)
æ-26 and the ray z<)) é æ Aōx(z)) with 0. Since V/x |
convex, we can frame this as a linear programming problem.
max A s.t. xxa) E V2
where Ve is the closure Vg. We discuss solving this efficiently in Section 3.1.
Let A be the solution. This solution exists if V. is bounded, which is always true due to the use of box constraints. Then a=(*) will be the point of intersection.
This first step solves for the farthest point the Voronoi cell in the direction of a. Using this knowledge, we can now map all points that lie this ray onto the Voronoi cell.
3

---

2. Apply an invertible transformation such that any point on the ray (a(A): À > 0} is mapped onto the line segment { + (xx(x) xx) a E (0, 1)}.
There are many possible choices for designing this transformation. An intuitive choice is to use monotonic transformation of the relative distance from æ to the anchor point Fik-
[x)))) mp Ork
| ||-= (x(Aa)
(x*) 27k)
where drk is an appropriate invertible squashing function from Rt to (0, 1). In experiments, use ag (h) = softsign(?%t where "i is a learned cell-dependent scale. Other choices should also work, such monotonic neural network; however, depending on the application, we may need to compute 0g
-1 for computing the inverse mapping, so it's preferable to have an analytical inverse.
3.1 Remarks and Propositions
Box constraints There can be continuity problems if a Voronoi cell is unbounded, as the solution to Equation (8) does not exist if x(x)) diverges. Furthermore, when solving Equation (7), it can be difficult to numerically distinguish between an unbounded cell and one whose boundaries are very far away. It is for these reasons that we introduce box constraints (Equation 6) in the formulation of Voronoi cells which allows us to sidestep these issues for now.
Solving for X Equation (7) can be solved numerically, but this approach is prone to numerical errors and requires implicit differentiation through convex optimization solutions [1]. We instead note that the solution of Equation (7) can be expressed in closed form, since it is always going to be the intersection of the ray æ(A with the nearest linear constraint.
Let a. æ b; be the plane that represents one of the linear constraints in Equation (6), which are expressed in Equation (4), Let X; be the intersection this plane with the ray, L.e is the solution to a. (xx)) = bg, then the solution is simply the smallest positive A,. which satisfies all the linear constraints:
b,-a,0x aT62(=)
x = minfAe X; > 0}, where X
There are total of K + 2D linear constraints, including the Voronoi cell boundaries and box constraints, which can be computed fully in parallel. This also allows end-to-end differentiation of the mapping fe via automatic differentiation, providing the ability to learn the parameters of fx. i.e. parameters of ag and the Voronoi tessellation.
Homeomorphism with continuous density We can show that fa is a bijection, and both fa and fx continuous. This allows us use fü within the normalizing flows framework, as mapping between distribution pz defined on Rº and the transformed distribution ps on Vx. Furthermore, the Jacobian continuous almost everywhere. Proofs are in Appendix A.
Proposition / The mapping fx RRD V₁ defined in the 2-step procedure is homeomorphism.
Proposition 2 If Pa (xx) is continuous, then the transformed density pa (Jx(z)) is continuous a.e.
Efficient computation of the log det Jacobian As f& is mapping D) dimensions, computing the log determinant Jacobian for likelihood computation can be costly and will scale poorly with D if computed naïvely. Instead, we note that the Jacobian of f& is highly structured. Intuitively, because fü depends only on the direction (x) and the distance away from Tk. it only has two degrees of freedom regardless of D. In fact, the Jacobian of fu can be represented as a rank-2 update on a scaled identity matrix. This allows us to use the matrix determinant lemma to reformulate the log determinant Jacobian in a compute- and memory-
form. We summarize this in a proposition.
Proposition 3 Let the transformation fa((x) and all intermediate quantities be as defined in Section for some given input æ. Then the Jacobian factorizes as
af((#) eT + Tnte 42u 8x
(10)
for some E R uz E IRD,, 04 E RDD and its log determinant has the form
log del (fx(2)
log W1nl + log 1 + W22 - wigea
+ log
(11)
Ox
Ttwn

---

Ordinal
Simplex
BinaryArgman
Voronoi
Figure 3: Existing dequantization methods can be seen special cases of Voronoi dequantization with fixed anchor points. In additional to decoupling the dimension of the embedding space from the number of discrete values, the Voronoi dequantization can learn to model the similarities of discrete values through the positions and boundaries of their Voronoi cells.
where Wij = =c--u.uu; Uj. This expression for the log determinant only requires inner products between vectors in RP_ To reduce notational clutter, exact formulas for c, u1, 42, V1, 12 are in Appendix A, All vectors used in computing Equation (11) either readily available intermediate quantities after computing (x) or are gradients of scalar functions and can be efficiently computed through reverse-mode automatic differentiation. The only operations on these vectors are dot products, and no large D-by-D matrices are ever constructed. Compared to explicitly constructing the Jacobian matrix, this more efficient in both compute and memory, and can readily scale to large values of D.
4 Application settings
We discuss two applications of our method to likelihood-based modeling of discrete and continuous data. The first is new approach dequantization which allows training model of discrete data using density models that normally only act on continuous variables such as normalizing flows. Compared existing dequantization methods [7, 18, 38], the Voronoi dequantization is not restricted to fixed dimensions and can benefit from learning similarities between discrete values. In fact, existing approaches can be seen as special cases of Voronoi dequantization.
The second is a formulation of disjoint mixture models, where each component lies strictly within each subset of Voronoi tessellation. To the best of our knowledge, disjoint mixture models have not been explored significantly in the literature and have not been successfully applied in more than a couple dimensions. Compared an existing method [8], the Voronoi mixture model is not restricted to acting individually for each dimension.
Voronoi dequantization Let y be discrete random variable with finite support V Then use tessellation to define subsets for dequantizing y as long as there are at least as many Voronoi cells--quivaleently anchors points-as the number of discrete values, i.e. K 2 | By assigning each discrete value to a Voronoi cell, we can then define a dequantization model ((>ly) by first sampling from base distribution Section to construct a distribution over Vg-- We can then obtain probabilities g(xly) efficiently and train the dequantization alongside the density model p(æ) on R
q( ((yly) in IR and then applying the mapping fa(z) from
Sampling from the model p(yla) (e) is straightforward and deterministic after sampling æ. We can write y g(x) where g is the set identification function satisfying E Vg(a)- From Equation (4), it is easy to see that g(x) ) is the nearest neighbor operation, g(x)) = arg min; || æxll-
We free choose the number of dimensions D, where smaller D assigns less space for each Voronoi cell induces dequantized distributions that are easier to fit, while a larger D allows the anchor points and Voronoi cells more room change over the course of training. When the anchor points are fixed to specific positions, we can recover the disjoint subsets used by prior methods as special cases. We illustrate this in Figure 3.
Voronoi mixture models Let (Vi. be a Voronoi tessellation of RDD. Then a disjoint mixture model can be constructed by defining distributions each Voronoi cell. Here make of the inverse mapping f-1 1 VV + R (discussed in Appendix B) so that we only need to parameterize distributions over I2 Let be a point in DR Voronoi mixture model assigns the density
log Pmix (x) = log Peomp (f (x) l: = g(x)) + log del "
+ llogpkk = g(x)) (12)
5

---

where Peomp can be any distribution over R", including another disjoint mixture model. This can easily be composed with normalizing flow layers where, in addition the change of variable due to fx
we also apply the change in density resulting from choosing one out of K components.
5 Related Work
Normalizing flows for discrete data Invertible mappings have been proposed for discrete data, where discrete values are effectively rearranged from a factorized distribution. In order param- eterize the transformation, Hoogeboom et al. [17], van den Berg et al. [47] use quantized ordinal transformations, while Tran et al. [45] takes a more general approach of using modulo addition on
one-hot vectors. These approaches suffer from gradient bias due to the need to use discontinuous operations and do not have universality guarantees since it's unclear whether simple rearrangement sufficient to transform any joint discrete distribution into factorized distribution. In contrast, the dequantization approach provides a universality guarantee since the lower bound in Equation (2) is ((=ly), with proportionality equal to p(y).
tight when p(x) [mE
Dequantization methods Within the normalizing flows literature, the act of adding noise was originally used for ordinal data as a way to combat numerical issues [7, 46]. Later appropriate dequantization approaches have been shown to lower bound the log-likelihood of discrete model [16, 44]. For non-ordinal data, many works have proposed simplex-based approaches. Early works on relaxations [20, 32] proposed continuous distribution on the simplex that mimic the behaviors of discrete random variables; however, these were only designed for the use with Gumbel base distribution. Potapczynski et al. [38] extend this to a Gaussian distribution alttbbu although it is not hard to this can work with any base distribution- designing invertible transformations between RDD and the probability simplex with K vertices, with D = K - 1, where is the number of classes of a discrete random variable.
Intuitively, after fixing one of the logits, the softmax operation is an invertible transformation on the bounded domain A K = (a E RRK Σ æ; = 1, æ. > 0}. The (K - 1)-simplex can then be broken into K subsets, each corresponding to particular discrete value.
A
simplex
(ır E A
AK- r* > L, Vi k).
(13)
More recently, Hoogeboom et al. [18] proposed ignoring the simplex constraint and simply use
Amrmaa
(x
R K I Vi + k}.
(14)
which effectively increases the number of dimensions by one compared to the simplex approach. However, both approaches force the dimension of the continuous space D scale with K. In order to make their approach work when K is large, [18] propose binarizing the data as a preprocessing step. In contrast, Voronoi dequantization has full flexibility in choosing regardless of K.
Related models for discrete data Among other related methods, a recent work proposed normaliz- ing flows truncated supports [43] but had resort to rejection sampling for training. Furthermore, their subsets are not disjoint by construction. Prior works [28, 31] also proposed removing the constraint that subsets are disjoint, and instead work with general mixture models with unbounded support, relying on the conditional model p(y (z) being sufficiently weak so that the task of modeling forced onto a flow-based prior. They have achieved good performance on number of tasks, similar to general approaches that combine normalizing flows with variational inference [19, 33, 50]. However, they lose the computational savings and the deterministic decoder p(yl x) gained from using disjoint subsets. On the other hand, quantization based on nearest neighbor have been used for learning discrete latent variables [34, 40], but no likelihoods are constructed, the boundaries are not explicitly differentiated through, and the model relies on training with heuristical updates.
Disjoint mixture models The computational savings from using disjoint subsets was pointed out by Dinh et al. [8]. However, their method only works in each dimension individually. They transform samples using a linear spline, which is equivalent to creating subsets based on the knots of the spline and applying a linear transformation within each subset. Furthermore, certain parameterizations of the spline can lead to discontinuous density functions, whereas our disjoint mixture has density function that is continuous almost everywhere (albeit exactly zero on the boundaries; see Section 7

---

Target PMF
Discrete Flow
Voronoi Flow
Dequantized samples
Figure 4: Quantized 2D data. Voronoi dequantization can model complex relations between discrete values. No knowledge or- dering is given to the models.
Figure 5: The model learns to cluster discrete values with similar probability values.
for more discussion). The use of monotonic splines have previously been combined with normalizing flows [10], but interestingly, since the splines in Dinh et al. [8] are not enforced to be monotonic, the full transformation is not bijective and acts like disjoint mixture model. Ultimately, their experiments were restricted to two-dimensional syntheic data sets, and it remained an interesting research question whether disjoint mixture models can be successfully applied high dimensions.
Experiments
We experimentally validate our semi-discrete approach of combining Voronoi tessellation with seelihood-based modeling a variety of data domains: discrete-valuee valued UCI data sets, itemset modeling, language modeling, and disjoint mixture modeling for continuous UCI data. The goal of these experiments is not to show state-of-the-ar results in these domains but to showcase the relative merit of Voronoi tessellation compared to existing methods. For this reason, we just use simple coupling blocks with affine transformations [6, 7] stated otherwise. When comparing to Discrete Flows [45], we use their bipartite layers. Details regarding preprocessing for data sets can be found in Appendix C, and detailed experimental setup is in Appendix D. All tables, when shown, display standard deviations across three random seeds. Open source code is available at https: https:///gtiiii /ggtttb....... comm/facabookreseerre semii wttp//smmmmdiddde-eeee-eeee discrete-ffow.
base flow model in experiments, unless
6.1 Quantized 2D data
We start by considering synthetic distributions over two discrete variables, to qualitatively compare against Discrete Flows [45]. Figure shows the target probability mass functions, which created by a 2D distribution where each dimension is quantized into 91 discrete values. Though the data is ordinal, knowledge of ordering given the models. We that Diserete Flows has trouble fitting these, because it is difficult to rearrange the target distribution into a factorized distribution. For our model, we dequantize each discrete variable with Voronoi tessellation in IR2. We then learn a flow model on the combined RR, parameterized by multilayer perceptrons (MLPs). In Figure we visualize the learned Voronoi tessellation and samples from our model. The learned tessellation seems to group some of discrete values that occur frequently together, to reduce the number of modes.
6.2 Discrete-valued UCI data
We experiment with complex data sets where each discrete variable can have varying number of classes. Furthermore, these discrete variables may have hidden semantics. To this end, we use unprocessed data sets from the UCI database [9]. The only pre-processing we perform to remove variables that only have one discrete value. We then take 80% as train, 10% validation, and 10% as the test set. Most of these data sets have combination of both ordinal and non-ordinal variables, and expect the non-ordinal variables to exhibit relations that unknown (e.g. spatial correlations).
7

---

Table 1: Discrete UCI data sets. Negative log- ligeliikldood results on the test sets in nats.
Method
Connect4
Forests
Mushroom
Nursery
Pokerllands
USCensus90
Voronoi Deq.
12.92.00r
14.20.o.05
9.06.0.0
9.27 1000
19.86
24. 19 1012
Simplex Deq.
13.46
16.58
9.26=mum:
9.50
+9000
28.09 +..0
BiggayArrmax Deq.
13.71 1:004
16.73 20.17
9.53 0.00
9.49
19.90
27.23 1:002
Discrete Flow
19.80
21.91 +0.0
22.06=m.o
9.53
19.8
55.62
Table 2: Permutation-invarian díserete itemset modeling.
Table 3: Language modeling.
Model (Dequantization)
Retail (nats)
Accidents (nats)
Dequartization
text8 (bpc)
eawik8 (bpc)
Veronoi (D=2)
CNF (Voronoi)
9.44+2234
7.81+284
Veronoi (D=t)
1.37
CNF (Simplex)
24.16+02:
19.19
Varonai (D=6)
1,37
1.40
CNF "BinaryArgmax
10.47
6.72+m23
Voronai (D=8)
1,36gmm
1.39
Determinantal Point Process
20.35
15.78
BinaryArgmax [18]
1.38
1.42
Ordinal [18]
1.43
1.44
We see that Discrete Flows can be competitive with dequantization approaches, but can also fall short on more challenging data sets such as the USCensus90, the largest data set we consider with 2.7 million examples and 68 different discrete variables of varying types. For dequantization methods, simplex [38] and binary argmax [18] approaches are mostly on par. We do see non-negligible gap in performance between these baselines and Voronoi dequantization for most of the data sets, likely due to the ability to learn semantically useful relations between the values of each discrete variable. For instance, the Connect4 data set contains positions of a two-player board game with the same name, which exhibit complex spatial dependencies between the board pieces, and the USCensus90 data set contains highly correlated variables and is often used to test clustering algorithms.
6.3 Permutation-invarian itemset modeling
An appeal of using normalizing flows continuous space is the ability to incorporate invariance to specific group symmetries into the density model. For instance, this can be done by ensuring the ordinary differential equation is equivariant [2, 25, 42] a continuous normalizing flow [4, 13] framework. Here we focus on invariance with respect to permutations, i.e. sets of discrete variables. This invariance cannot be explicitly modeled by Discrete Flows as they require an ordering or a bipartition of discrete variables. We preprocessed a data set of retail market basket data from an anonymous Belgian retail store [3] and a data set of anonymized traffic accident data [11], which contain sets of discrete variables each with 765 and 269 values, respectively. Without the binarization trick, simplex dequantization must use a large embedding dimensions and performs worse than a determinental point process baseline [26]. Meanwhile, Voronoi dequantization has no problems staying competitive as its embedding space is freely decoupled from the number of discrete values.
6.4 Language modeling
Language modeling a widely used benchmark for discrete models [18, 28, 45J. Here we used the open source code provided by Hoogeboom et al. [18] with the exact same autoregressive flow model and optimizer setups. The only difference is replacing their binary argmax with our Voronoi dequantization. Results are shown in Table 3, where Generally, we find D = 2 to be too low and can stagnate training since the Voronoi cells are more constrained, while larger values of D improve upon argmax dequantization.
tried out multiple embedding dimensions D.
6.5 Voronoi mixture models
We test disjoint mixture modeling with Voronoi tessellation on continuous data sets. Figure 6 shows results of training a simple normalizing flow on data, which has trouble completely separating all the modes. Adding in a Voronoi mixture allows us to better fit these multimodal densities.
8

---

Data
Baseline
Voronoi
Samples
Data
Baseline
Voronoi
Samples
Figure 6: Tessellation done in transformed space; nonlinear boundaries are shown.
Table 4: Disjoint mixture modeling. NLL on the test sets in nats.
*Bascline results from [13, 35].
Method
POWER
GAS
HEPMASS
MINIBOONE
BSDS300
Real NVP
-0.17
-0,33
18.71
13,55
-1.3.28
MAF*
+0.24
=10008
+7.73
12.24
-154.93
FFJORD
-0.462001
12
14.92 +0.08
10.43
-157 a0
Base Coupling Flow
+0.44
-11.75+0oz
16.78
10.87
-155. 1450.14
Voronoi Disjoint Mixture
-0.52
-12.63
16.16.1
10.24
-156.59
To test Voronoi mixture models on higher dimensions, we apply our approach to data sets preprocessed by Papamakarios al. [35], which range from to 63 dimensions. We add disjoint mixture component after a few layers of normalizing flows, with each component modeled by normalizing flow conditioned on a embedding of the component index. The number of mixture components ranges from 16 to 64, with complete sweeps and hyperparameters detailed in Appendix D. For comparison, we also show results from the baseline coupling flow that uses the same number of total layers as the disjoint mixture, as well as a strong but computationally costly density model FFJORD [13]. From Table 4, we see that the disjoint mixture model approach allows to increase complexity quite a bit, with almost no additional cost compared to the baseline flow model.
7 Conclusion and Discussion
We combine Voronoi tessellation with normalizing flows construct invertible transformation that has learnable discrete structure. This acts as a learnable mapping between discrete and continuous distributions. We propose two applications of this method: Voronoi dequantization that maps discrete values into a learnable convex polytope, and a Voronoi mixture modeling approach that has around the same compute cost single component. We showcased the relative merit of our approach across a range of data modalities with varying hidden structure. Below, we discuss some limitations and possible future directions.
Diminishing density boundaries. The distributions within each Voronoi cell necessarily go zero on the boundary between cells due to the use of a homeomorphism from an unbounded domain. This is different and undesirable property as opposed to the partitioning method used by Dinh et al. [8]. However, alternatives could require more compute cost in the form of solving normalization constants. Balancing expressiveness and compute cost a delicate problem that sits the core of probabilistic machine learning.
Design of homeomorphisms and tessellations. Our proposed homeomorphism is simple and computationally scalable, but this comes the cost of smoothness and expressiveness. As depicted in Figure 1, the transformed density has a non-smooth probability density function. This non-smoothness exists when the ray intersects with multiple boundaries. This may make optimization more difficult as the gradients of the log- likelihoodd objective can be discontinuous. Additionally, our use of Euclidean distance can become problematic in high dimensions, this metric can cause almost all points to be nearly equidistant, resulting in behavior where all points lie seemingly very close to the boundary. Improvements on the design of homeomorphisms bounded domains could help alleviate these problems. Additionally, more flexible tessellations -ucncsssssshhhhee as the Laguerre tessellation and additional concepts from semi-discrete optimal transport [14, 27, 37] may adapted improve semi-discrete normalizing flows.

---

References
[1] Agrawal, A., Amos, B., Barratt, S., Boyd, S., Diamond, S., and Kolter, Z. Differentiable convex optimization layers. Advances in Neural Information Processing Systems, 2019.
[2] Biloš, M. and Günnemann, S. Scalable normalizing flows for permutation invariant densities. In International Conference on Machine Learning, pp. 957-967. PMLR, 2021.
[3] Brijs, T., Swinnen, G., Vanhoof, K., and Wets, G. Using association rules for product assortment decisions: case study. In Knowledge Discovery and Data Mining, pp. 254-260, 1999.
[4] Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 2018.
[5] Chen, R. Q. torchdiffeq, 2018. URL htttps://gitee
[6] Dinh, L.,, Krueger, D., and Bengio, Y. NICE: Non-linear independent components estimation.
arXiv preprint arXiv: 2014.
[7] Dinh, L., Sohl-Dickstein, J., and Bengio, S. Density estimation using real NVP. International
Conference Learning Representations, 2017.
[8] Dinh, L., Sohl-Dickstein, J., Larochelle, H., and Pascanu, R. A rad approach to deep mixture models. arXiv preprint arXiv:1903.07714, 2019.
[9] Dua, D. and Graff, C. UCI machine learning repository, 2017. URL http:///rrccciee.iiie uci. edu/ml.
[10] Durkan, C., Bekasov, A., Murray, I., and Papamakarios, G. Neural spline flows. Advances in Neural Information Processing Systems, 32:7511-7522 2019.
[11] Geurts, K., Wets, G., Brijs, T., and Vanhoof, K. Profiling high frequency accident locations using association rules. In Proceedings of the 82nd Annual Transportation Research Board, Washington DC. (USA), Jamuary 12-16, pp. 18pp, 2003.
[12]
Goodfellow, I., Bengio, Y., and Courville, Deep learning. MIT press, 2016.
[13] Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., and Duvenaud, D. FFJORD: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv.1810.01367, 2018. 2018.
[14] Gruber, P. M. Optimum quantization and its applications. Advances in Mathematics, 186(2): 456 4666-497,22 2004.
[15] Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint 16Xiv:1606.0888 18 2016.
[16] Ho, J., Chen, X., Srinivas, A., Duan, Y., and Abbeel, Flow++: Improving flow-
generative models with variational dequantization and architecture design. In International Conference on Machine Learning, pp. 2722-273 2730.. PMLR, 2019.
[17] Hoogeboom, E., Peters, J. W., Berg, R. d., and Welling, M. Integer discrete flows and lossless compression. Advances in Neural Information Processing Systems, 2019.
[18] Hoogeboom, F., Nielsen, D., Jaini, P., Forré, P., and Welling, M. Argmax flows and multinomial diffusion: Towards auoreevoreeeessvve language models. Advances in Neural Information Processing Systems, 2021.
[19] Huang, C.-W., Dinh, L.., and Courville, A. Augmented normalizing flows: Bridging the gap between generative flows and latent variable models. arXiv preprint arXiv:2002.07 2020.
[20] Jang, E., Gu, S., and Poole, B. Categorical reparameterization with gumbel-softmas Interna- tional Conference on Learning Representations, 2017.
[21] Kim, H., Papamakarios, G., and Mnih, A. The lipschitz constant of self-attention. In Interna- tional Conference on Machine Learning, pp. 5562-5571. PMLR, 2021.
[22] Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2015.
[23] Kingma, D. P. and Dhariwal, P. Glow: Generative flow with invertible 1x1 convolutions.
Advances in Neural Information Processing Systems, 2018.
[24] Kobyzev, I., Prince, S., and Brubaker, M. Normalizing flows: An introduction and review of current methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020.
10

---

[25] Köhler, J., Klein, L., and Noé, F. Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning. pp. 5361-5370. PMLR, 2020.
[26] Kulesza, A. and Taskar, B. Determinantal point processes for machine learning. Foundations and Trendy@ in Machine Learning, 5(2-3):123-286, 2012. ISSN 1935-8237. doi: 10.1561 2200000044.
[27] Lévy, B. A numerical algorithm for 12 semi-discrete optimal transport in 3d. ESAIM: Mathe- matical Modelling and Numerical Analysis, 19(6):1633-177 2015.
[28] Lippe, P. and Gavves, E. Categorical normalizing flows via continuous transformations. Inter- national Conference on Learning Representations, 2021.
[29] Liu, J., Kumar, A., Ba, J., Kiros, J., and Swersky, K. Graph normalizing flows. Advances in Neural Information Processing Systems, 2019.
[30] Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. International Conference on Learning Reprexentations, 2019.
[31] Ma, X., Zhou, C., Li, X.. Neubig. G., and Hovy, E. Flowseq: Non-autoregresssvv conditional sequence generation with generative flow. arXiv preprini (:Xiv::90.00244:: 2019.
[32] Maddison, C. J., Mnih, A., and Teh, Y. W. The concrete distribution: continuous relaxation of discrete random variables. International Conference on Learning Representations, 2017.
[33] Nielsen, D., Jaini, P., Hoogeboom, E., Winther, O., and Welling, M. Survae flows: Surjections bridge the gap between vaes and flows. Advances in Neural Information Processing Systems, 33, 2020.
[34] Oord, A. v. d., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. arXiv preprint arXiv:1711. 1.00977, 2017.
[35] Papamakarios, G., Pavlakou, T., and Murray, I. Masked autoregressive flow for density estima- tion. Advances in neural information processing systems, 30, 2017.
[36] Papamakarios, G., Nalisnick, E., Rezende, D. J., Mohamed, S.,, and Lakshminarayanan, B. Normalizing flows for probabilistic modeling and inference. arXiv preprint 192.0022. 2019.
[37] Peyré, G., Cuturi, M., et al. Computational optimal transport. Center for Research in Economics and Statistics Working Papers, (2017-86), 2017.
[38] Potapezynski, A., Laniz, G., and Cunningham, J. P. Invertible gaussian reparameteriza- tion: Revisiting the gumbel-softmax. arXiv preprint 1rXiv:1912.09535 2019.
[39] Ramachandran, P., Zoph, B., and Le, Q. Searching for activation functions. arXiv preprint 11010:171,.05941, 2017.
[40] Razavi, A., van den Oord, A., and Vinyals, O. Generating diverse high-fidelity images with Vq-vae-2. In Advances in neural information processing systems, pp. 14866-14877 14876, 2019. [41] Rudin, W. et al. Principles of mathematical analysis, volume 3. McGraw-hill New York, 1964.
[42] Satorras, V. G., Hoogeboom, F.,, Fuchs, F. B., Posner, I., and Welling, M. F(n) equivariant normalizing flows. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021.
[43] Tan, S., Huang, C.-W., Sordoni, A., and Courville, A. Learning to dequantise with truncated flows. In International Conference on Learning Representations, 2022.
[44] Theis, L.., Oord, A. v. d., and Bethge, M. A note on the evaluation of generative models. International Conference on Learning Representations, 2016.
[45] Tran, D., Vafa, K., Agrawal, K., Dinh, L.,, and Poole, B. Discrete flows: Invertible generative models of discrete data. Advances in Neural Information Processing Systems, 32:14719-11428,, 2019.
[46] Uria, B., Murray, I., and I Larochelle, H. Rnade: The real-valued neural autoregressive density- estimator. Advances in Neural Information Processing Systems, 26, 2013.
[47] van den Berg, R., Gritsenko, A. A., Dehghani, M., Sønderby, C. K., and Salimans, T. Idf++: Analyzing and improving integer discrete flows for lossless compression. In International Conference Learning Representations, 2020.

---

[48] Voronoi, G. Nouvelles applications des paramètres continus à la théorie des formes quadratiques. premier mémoire. sur quelques propriétés des formes quadratiques positives parfaites. Journal für die reine und angewandte Mathematik (Crelles Journal), 1908(133):97-102) 1908.
[49] Zhang, G., Wang, C., Xu, B., and Grosse, R. Three mechanisms of weight decay regularization. International Conference Learning Representations, 2019.
[50] Ziegler, Z. and Rush, A. Latent normalizing flows for discrete sequences. In International Conference on Machine Learning, pp. 7673-7682. PMLR, 2019.
12


Reviews:
[{"review_id": "fQ0VPbqPdoB", "reviewer": "Reviewer_xCLG", "summary": "The paper applies Voronoi tessellation to allow the application of normalizing flows on discrete data. The main contributions are Voronoi dequantization, which is a new mapping method between discrete values and continuous space, and a disjoint mixture model, which uses an invertible mapping between unbounded support and convex polytopes. Experimental result on a variety of datasets shows that the proposed method outperforms existing baselines. ", "questions": "- Since the experiments mainly compared with other flow models, it would be interesting to see how the proposed methods compare with state-of-the-art methods. Is normalizing flow outperformed by other methods even with the Voronoi dequantization and mixture modeling, or can it actually have better performance? ", "limitations": "The authors have discussed limitations, along with possible ways to address them, in the conclusion. \n", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The Voronoi tessellation idea is simple, novel, and intuitive. Additionally, it addresses an important problem of modeling discrete data with normalizing flows.\n- The writing is very clear. All mathematical details are concise yet complete. The authors also provide enough context to distinguish their contribution from related works. \n- The proposed method for dequantization and disjoint mixture modeling significantly improves performance with little additional computational cost. \n\nWeaknesses:\nThere are no major flaws that I can think of. The paper can be further improved if the authors can provide a few alternatives to the current invertible mapping to convex polytopes and compare performance and computation costs. Other design choices can be explored as mentioned in the conclusion section, but the current content seems sufficient for a paper.\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "VL4P_W_FLdw", "reviewer": "Reviewer_cFmy", "summary": "The paper proposes a new method for modelling discrete data using normalizing flows by extending previous works on dequantization based method (like argmax flows) using a \"tessellation\" (not introduced or defined in the paper) based approach that automatically learns the boundaries for quantization. This is explore both for discrete to continuous and continuous to discrete.", "questions": "Please see in the strengths and weakness section.", "limitations": "Limitations have been discussed.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n- The paper describes sound extensions to normalizing flows for modelling discrete data and explores novel ways of incorporating tessellation based approaches in this context.\n- The experiments are done on a variety of datasets including language modelling etc to study the empirical performance of the method. \n- Significance: Getting good performance of generative models on discrete data is challenging and thus the problem is quite hard. This paper explores an interesting extension to previous methods to tackle this problem.\n\n**Weakness:**\n- Unfortunately, the paper is not clear on technical details. Important concepts are assumed to be known by the reader instead of properly introducing them e.g. tessellation, anchor points, Voronoi tessellation etc. The paper will greatly benefit from having a brief notations and definitions section. I also believe the paper will benefit greatly from having intuitive explanations accompany technical details to make it easier to follow.\n\n- Similarly, several appropriate references are missing e.g. for dequantization, the paper by Uria, 2013 needs to be cited. Similarly, statements/claims as in line 46-48 need to be supported with relevant references. \n\n- It was not immediately clear to me why is it useful to go from continuous case to the discrete case? What applications will this be important in? Apart from testing this on simple UCI datasets, I do not see any specific experiments for this.\n\n- I also found the experiments fairly limited in scope. There are no experiments to show the capability of the model on modeling categorical data with large number of classes. The language modeling task have small vocabularies. I'd recommended to evaluate the model on word-based language models with large vocabularies.\n\n- Similar to RAD flows, I believe the method will suffer from exploding gradients near the boundaries. How is the training stabalised to address this behaviour?", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "GMO4KBkmh58", "reviewer": "Reviewer_qSwN", "summary": "The paper presents a tessellation based approach to mapping between continuous and discrete distributions which learns quantization boundaries in continuous space. The approach has several desirable mathematical and computational properties that give rise to more efficient scaling in dimensions than previous approaches. The approach is used in two ways--Voronoi dequantization and Voronoi mixtures--which are demonstrated on data and shown to be competitive with baselines. The authors concluded by noting some limitations and directions for extension. ", "questions": "Following the above discussion, I would be interested in understanding the advantage of the Voronoi approaches. ", "limitations": "The authors do mention a couple limitations and directions for future work, which I applaud. There is no reason to discuss social impact for this work. ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is a mostly clear and interesting presentation of an interesting approach to bridging discrete and continuous distributions. Strengths include: clear mathematical and general exposition, an interesting mix of mathematical and computational considerations, and interesting empirical tests. There are not many weaknesses, though I found myself wanting to understand the advantages of the Voronoi approaches better. The experiments did what the authors proposed (show competitiveness), but it would have been nice to have stronger intuition. ", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "BZAr-5CbXVh", "reviewer": "Reviewer_vHMS", "summary": "A dequantization method is presented that is able to map between discrete and continuous distributions using continuous, invertible normalizing flows. Two applications are presented, namely Voronoi dequantization and Voronoi mixture models.", "questions": "In Figure 5, the caption could directly explain what \"Samples\" refers to. Line 325 also mentions how the method could potentially make gradient-based optimization difficult, has this been tried in any experiments? It would be interesting to see such results, or these kinds of applications with the framework.", "limitations": "The limitations are clear.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The results are well founded, and figures visually explain the concepts well. Though the text is often a bit convoluted, where the text explanation tends to be over-complex and hard to understand for non-experts in the field.\n\nAlthough it is great how much space is given to explaining the background and catching the reader up, more results and experiments would perhaps be even more intriguing. \n\nCould it be that no Appendix was present with the submission, even though it was referred to?", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept