Original: Semi-Discrete Normalizing Flows through Differentiable Tessellation
Revised: Semi-Discrete Normalizing Flows through Differentiable Tessellation

Revised: Anonymous Author(s)

Revised: Afﬁliation email

Original: Abstract
Revised: Abstract

Original: Mapping between discrete and continuous distributions is a difficult task and manyhave had to resort to heuristical approaches.
Revised: Mapping between discrete and continuous distributions is a difﬁcult task and manyhave had to resort to heuristical approaches.

Original: We propose a tessellation-based approach that directly learns quantization boundaries in a continuous space, completewith exact likelihood evaluations.
Revised: We propose a tessellation-based approach that directly learns quantization boundaries in a continuous space, completewith exact likelihood evaluations.

Original: This is done through constructing normalizingflows on convex polytopes parameterized using a simple homeomorphism withan efficient log determinant Jacobian.
Revised: This is done through constructing normalizingﬂows on convex polytopes parameterized using a simple homeomorphism withan efﬁcient log determinant Jacobian.

Original: We explore this approach in two application settings, mapping from discrete to continuous and vice versa.
Revised: We explore this approach in two application settings, mapping from discrete to continuous and vice versa.

Original: Firstly, a
Revised: Firstly, a

Original: Voronoi dequantization allows automatically learning quantization boundaries ina multidimensional space.
Revised: Voronoi dequantization allows automatically learning quantization boundaries ina multidimensional space.

Original: The location of boundaries and distances between regions can encode useful structural relations between the quantized discrete values.
Revised: The location of boundaries and distances between regions can encode useful structural relations between the quantized discrete values.

Original: Secondly, a Voronoi mixture model has constant computation cost for likelihoodevaluation regardless of the number of mixture components.
Revised: Secondly, a Voronoi mixture model has constant computation cost for likelihoodevaluation regardless of the number of mixture components.

Original: Empirically, we showimprovements over existing methods across a range of structured data modalities.
Revised: Empirically, we showimprovements over existing methods across a range of structured data modalities.

Original:  Introduction
Revised:  Introduction

Original: Likelihood-based models have seen increasing usage across multiple data modalities.
Revised: Likelihood-based models have seen increasing usage across multiple data modalities.

Original: Across a varietyof modeling approaches, the family of normalizing flows stands out as a large amount of structurecan be incorporated into the model, aiding its usage in modeling a wide variety of domains such asimages [7, 23], graphs [29], invariant distributions [2, 25] and molecular structures [42].
Revised: Across a varietyof modeling approaches, the family of normalizing ﬂows stands out as a large amount of structurecan be incorporated into the model, aiding its usage in modeling a wide variety of domains such asimages [7, 23], graphs [29], invariant distributions [2, 25] and molecular structures [42].

Original: However,the majority of works focus on only continuous functions and continuous random variables.
Revised: However,the majority of works focus on only continuous functions and continuous random variables.

Original: Thisrestriction 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.
Revised: Thisrestriction can make it difﬁcult to apply such models to distributions with implicit discrete structures,e.g. distributions with discrete symmetries, multimodal distributions, distributions with holes.

Original: In this work, we incorporate discrete structure into standard normalizing flows, while being entirelycomposable with any other normalizing flow.
Revised: In this work, we incorporate discrete structure into standard normalizing ﬂows, while being entirelycomposable with any other normalizing ﬂow.

Original: Specifically, we propose a homeomorphism betweenthe unbounded domain and convex polytopes, which are defined through a learnable tessellation ofthe  domain.
Revised: Speciﬁcally, we propose a homeomorphism betweenthe unbounded domain and convex polytopes, which are deﬁned through a learnable tessellation ofthe domain, i.e. a set of disjoint subsets that together fully cover the domain.

Original: This homeomorphism is cheap to compute, has a cheap inverse, and we additionallyprovide an efficient and exact formulation of the resulting change in density.
Revised: This homeomorphism ischeap to compute, has a cheap inverse, and results in an efﬁcient formulation of the resulting changein density.

Original: In other words, thistransformation is highly scalable—at least computationally.
Revised: In other words, this transformation is highly scalable—at least computationally.

Original: Our method has the potential to be useful for a variety of applications.
Revised: Our method has the potential to be useful for a variety of applications.

Original: Firstly, the learned tessellationnaturally allows defining a dequantization method for discrete variables.
Revised: Firstly, the learned tessellationnaturally allows deﬁning a dequantization method for discrete variables.

Original: This allows likelihood-basedmodels that normally only act on continuous spaces, such as normalizing flows, to work directly ondiscrete data in a flexible choice of embedding space with the ability to capture relational structure.
Revised: This allows likelihood-basedmodels that normally only act on continuous spaces, such as normalizing ﬂows, to work directly ondiscrete data in a ﬂexible choice of embedding space with the ability to capture relational structure.

Original: Secondly, if we take each convex polytope as the support of a single mixture component, then the fulltessellation defines a mixture model with disjoint components.
Revised: Secondly, if we take each convex polytope as the support of a single mixture component, then the fulltessellation deﬁnes a mixture model with disjoint components.

Original: This allows us to scale to mixturesof normalizing flows while retaining the compute cost of a single model.
Revised: This allows us to scale to mixturesof normalizing ﬂows while retaining the compute cost of a single model.

Original: Following semi-discreteoptimal transportation [37], which defines couplings between discrete and continuous measures, werefer to our models as semi-discrete normalizing flows that learn transformations between discreteand continuous random variables.
Revised: Following semi-discreteoptimal transportation [37], which deﬁnes couplings between discrete and continuous measures, werefer to our models as semi-discrete normalizing ﬂows that learn transformations between discreteand continuous random variables.

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: f

Original: [Equation]
Revised: [Equation]

Original:  Preliminaries
Revised:  Preliminaries

Original: Normalizing Flows.
Revised: Normalizing Flows.

Original: This family of generative models typically includes any model that makes useof invertible transformations f to map samples between distributions.
Revised: This family of generative models typically includes any model that makes useof invertible transformations f to map samples between distributions.

Original: The relationship between theoriginal and transformed density functions have a closed form expression,
Revised: The relationship between theoriginal and transformed density functions have a closed form expression,

Original: [Equation]
Revised: [Equation]

Original: For instance, p z can be a simple base distribution while f is learned so that the resulting p x is closeto some target distribution.
Revised: For instance, p z can be a simple base distribution while f is learned so that the resulting p x is closeto some target distribution.

Original: Many different choices of f have been discussed in the literature, leadingto different trade-offs and use cases [24, 36].
Revised: Many different choices of f have been discussed in the literature, leadingto different trade-offs and use cases [24, 36].

Original: Furthermore, if the domain and codomain of f are different, then p x and p z can have differentsupports, i.e. regions where probability is non-zero.
Revised: Furthermore, if the domain and codomain of f are different, then p x and p z can have differentsupports, i.e. regions where probability is non-zero.

Original: Currently, existing designs of f that have thisproperty act independently for each dimension.
Revised: Currently, existing designs of f that have thisproperty act independently for each dimension, such as the logit transform [7].

Original: The use of general support-modifying invertibletransformations have not been discussed extensively in the literature.
Revised: To the best of ourknowledge, the use of general support-modifying invertible transformations have not been discussedextensively in the literature.

Original: Dequantization.
Revised: Dequantization.

Original: In order to model discrete data with density models, a standard approach is tocombined with dequantization methods.
Revised: In order to model discrete data with density models, a standard approach is tocombined with dequantization methods.

Original: These methods provide a correspondence between discretevalues and continuous subspaces, where a discrete variable y ∈ Y is randomly placed within a disjointsubset of R D .
Revised: These methods provide a correspondence between discretevalues and convex subsets, where a discrete variable y 2 Y is randomly placed within a disjointsubset of R D .

Original: There is also a correspondence between the likelihood values of the discrete randomvariable and the dequantized random variable [44].
Revised: There is also a correspondence between the likelihood values of the discrete randomvariable and the dequantized random variable [44].

Original: Let A y denote the subspace corresponding to y and let q ( x | y ) be the dequantization model whichhas a bounded support in A y .
Revised: Let A y denote the subset corresponding to y and let q ( x | y ) be the dequantization model which has abounded support in A y .

Original: Then any density model p ( x ) with support over R D satisfieslog p 
Revised: Then any density model p ( x ) with support over R D satisﬁeslog p (

Original: [Equation]
Revised: [Equation]

Original: Thus with an appropriate choice of dequantization, maximizing the likelihood under the densitymodel p ( x ) is equivalent to maximizing the likelihood under the discrete model p ( y ) .
Revised: Thus with an appropriate choice of dequantization, maximizing the likelihood under the densitymodel p ( x ) is equivalent to maximizing the likelihood under the discrete model p ( y ) .

Original: In existing dequantization methods, the value of D and the choice of subsets
Revised: In existing dequantization methods, the value of D and the choice of subsets

Original: A y are entirely dependenton the type of discrete data (ordinal vs non-ordinal) and the number of discrete values |Y| in thenon-ordinal case.
Revised: A y are entirely dependenton the type of discrete data (ordinal vs non-ordinal) and the number of discrete values |Y| in thenon-ordinal case.

Original: Furthermore, the subsets A y do not interact with one another and are fixed duringtraining.
Revised: Furthermore, the subsets A y do not interact with one another and are ﬁxed duringtraining.

Original: In contrast, we conjecture that important relations between discrete values should bemodeled as part of the parameterization of A y , and that it’d be useful to be able to automatically learnthe boundaries of A y based on gradients.
Revised: In contrast, we conjecture that important relations between discrete values should bemodeled as part of the parameterization of A y , and that it’d be useful to be able to automatically learnthe boundaries of A y based on gradients.

Original: Disjoint mixture models.
Revised: Disjoint mixture models.

Original: Building mixture models is one of the simplest methods for creatingmore flexible distributions from simpler one, and mixture models can typically be shown to beuniversal density estimators [12].
Revised: Building mixture models is one of the simplest methods for creatingmore ﬂexible distributions from simpler one, and mixture models can typically be shown to beuniversal density estimators [12].

Original: However, in practice this often requires a large number of mixturecomponents, which quickly becomes computationally expensive.
Revised: However, in practice this often requires a large number of mixturecomponents, which quickly becomes computationally expensive.

Original: This is because evaluating thelikelihood under a mixture model requires evaluating the likelihood of each component.
Revised: This is because evaluating thelikelihood under a mixture model requires evaluating the likelihood of each component.

Original: To alleviatethis issue, Dinh et al.
Revised: To alleviatethis issue, Dinh et al.

Original: [8] recently proposed to use components with disjoint support.
Revised: [8] recently proposed to use components with disjoint support.

Original: In particular, let{ A k } Kk =1 be disjoint subsets of R D , such that each mixture component is defined on one subset and has support restricted to that particular subset.
Revised: In particular, let{ A k } Kk =1 be disjoint subsets of R D , such that each mixture component is deﬁned on one subset and has support restricted to that particular subset.

Original: The likelihood of the mixture model then simplifies top
Revised: The likelihood of the mixture model then simpliﬁes top

Original: [Equation]
Revised: [Equation]

Original: (3) where g : R D → { 1 , . . .
Revised: (3) where g : R D ! { 1 , . . .

Original: , K } is a set identification function that satisfies x ∈ A g ( x ) .
Revised: , K } is a set identiﬁcation function that satisﬁes x 2 A g ( x ) .

Original: This frameworkallows building large mixture models while no longer having a compute cost that scales with K .
Revised: This framework 73 allows building large mixture models while no longer having a compute cost that scales with K .

Original: Incontrast to variational approaches with discrete latent variables, the use of disjoint subsets providesan exact (log-)likelihood and not a lower bound.
Revised: Incontrast to variational approaches with discrete latent variables, the use of disjoint subsets providesan exact (log-)likelihood and not a lower bound.

Original:  Voronoi Tessellation for Normalizing Flows
Revised:  Voronoi Tessellation for Normalizing Flows

Original: We first discuss how we parameterize each subset as a convex polytope, and in the next part, discussconstructing distributions on each convex polytope using a differentiable homeomorphism.
Revised: We ﬁrst discuss how we parameterize each subset as a convex polytope, and in the next part, discussconstructing distributions on each convex polytope using a differentiable homeomorphism.

Original: Parameterizing disjoint subsets We separate the domain into subsets through a Voronoi tessellation[47].
Revised: Parameterizing disjoint subsets We separate the domain into subsets through a Voronoi tessellation[48].

Original: This induces a correspondence between each subset and a corresponding anchor point , whichprovides a differentiable parameterization of the tessellation for gradient-based optimization.
Revised: This induces a correspondence between each subset and a corresponding anchor point , whichprovides a differentiable parameterization of the tessellation for gradient-based optimization.

Original: Let X = { x 1 , . . .
Revised: Let X = { x 1 , . . .

Original: , x K } be a set of anchor points in R D .
Revised: , x K } be a set of anchor points in R D .

Original: The Voronoi cell for each anchor point is
Revised: The Voronoi cell for each anchor point is

Original: [Equation]
Revised: [Equation]

Original: i.e. it defines a subspace containing all points which have the anchor points as its nearest neighbor.
Revised: i.e. it deﬁnes a subset containing all points which have the anchor point as their nearest neighbor.

Revised: Together, these subsets { V k } Kk =1 form a Voronoi tessellation of R D .

Original: This subspace can equivalently be expressed in the form of a convex polytope, 
Revised: Each subset can equivalently beexpressed in the form of a convex polytope, V

Original: [Equation]
Revised: [Equation]

Original:  } , where a i 
Revised: 6 = k } , where a i = 2( x i  x k ) T , b

Original: [Equation]
Revised: i =

Original:  x [Equation]
Revised: k x i k

Original: [Equation]
Revised: [Equation]

Original: For simplicity, we also include box constraints so that all Voronoi cells are bounded in all directions.
Revised: For simplicity, we also include box constraints so that all Voronoi cells are bounded in all directions.

Original: [Equation]
Revised: [Equation]

Original: Thus, the learnable parameters of this Voronoi tessellation are the anchor points { x 1 , . . .
Revised: Thus, the learnable parameters of this Voronoi tessellation are the anchor points { x 1 , . . .

Original: , x K } ,and the box constraints c l , c r ∈ R D .
Revised: , x K } ,and the box constraints c l , c r 2 R D .

Original: These will be trained using gradients from an invertibletransformation and resulting distribution defined in each Voronoi cell, which we discuss next.
Revised: These will be trained using gradients from an invertibletransformation and resulting distribution deﬁned in each Voronoi cell, which we discuss next.

Original: Invertible mapping onto the Voronoi cell Within the normalizing flow framework, defining aprobability distribution on a bounded support is equivalent to defining an invertible transformationthat maps from the unbounded support onto the appropriate support.
Revised: Invertible mapping onto subset Within the normalizing ﬂowframework, deﬁning a probability distribution on a bounded support is equivalent to deﬁning an invertible transformation that mapsfrom the unbounded support onto the appropriate support.

Revised: Figure 2illustrates the transformation of mapping onto a shaded region.

Original: Given a cell V k for somek ∈ { 1 , . . .
Revised: Given a cell V k for some k 2 { 1 , . . .

Original: , K } , we construct an invertible mapping f 
Revised: , K } , we construct an invertiblemapping f k : R D !

Original: k R D → V k by following 2 steps.
Revised: V k by following 2 steps.

Original: Letx ∈ R D .
Revised: Let x 2 R D .

Original: If x = x k , the anchor point of V k , we simply set f k ( x ) = x . Otherwise:1.
Revised: First,if x = x k , the anchor point of V k , we simply set f k ( x ) = x . Otherwise:1.

Original: Determine where the ray starting from x k in the direction of x intersects with the boundary of V k .
Revised: Determine where the ray starting from x k in the direction of x intersects with the boundary of V k .

Original: First define the direction δ k ( x ) ≜ x − x k ∥ x − x k ∥ and the ray x ( λ ) ≜ x k + λ δ k ( x ) , with λ > 0 .
Revised: First deﬁne the direction  k ( x ) , x  x k k x  x k k and the ray x (  ) , x k +   k ( x ) , with  > 0 .

Original: Since V kis a convex polytope, we can frame this as a linear programming problem.
Revised: Since V kis a convex polytope, we can frame this as a linear programming problem.

Original: [Equation]
Revised: [Equation]

Original: where V k is the closure of V k .
Revised: where V k is the closure of V k .

Original: We discuss solving this efficiently in Section 3.1.
Revised: We discuss solving this efﬁciently in Section 3.1.

Original: Let λ ∗ be the solution.
Revised: Let  ⇤ be the solution.

Original: This solution exists if V k is bounded, which is always true due to the use ofbox constraints.
Revised: This solution exists if V k is bounded, which is always true due to the use of [Figure] box constraints.

Original: Then x ( λ ∗ ) will be the point of intersection.
Revised: Then x (  ⇤ ) will be the point of intersection.

Original: This first step solves for the farthest point in the Voronoi cell in the direction of x .
Revised: This ﬁrst step solves for the farthest point in the Voronoi cell in the direction of x .

Original: Using thisknowledge, we can now map all points that lie on this ray onto the Voronoi cell.2.
Revised: Using thisknowledge, we can now map all points that lie on this ray onto the Voronoi cell.

Original: Apply an invertible transformation such that any point on the ray { x ( λ ) : λ > 0 } is mapped ontothe line segment{ x k
Revised: Apply an invertible transformation such that any point on the ray { x (  ) :  > 0 } is mapped onto the line segment

Original:  + α ( x ( λ ∗ ) − x k ) : α ∈ (0 , 1) } .
Revised: { x k + ↵ ( x (  [Equation] )  x k ) : ↵ 2 (0 , 1) } .

Original: There are many possible choices for designing this transformation.
Revised: There are many possible choices for designing this transformation.

Original: An intuitive choice is to use amonotonic transformation of the relative distance from x to the anchor point x k .
Revised: An intuitive choice is to use amonotonic transformation of the relative distance from x to the anchor point x k .

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original:  where α k is an appropriate invertible squashing function from R + to (0 , 1) .
Revised: ) where ↵ k is an appropriate invertible squashing function from R + to (0 , 1) .

Original: In our experiments, weuse α k ( h )
Revised: In our experiments, weuse ↵ k ( h )

Original: =
Revised: =

Original: softsign ( γ k h ) where γ k is a learned cell-dependent scale.
Revised: softsign (  k h ) where  k is a learned cell-dependent scale.

Original: Other choices should alsowork, such as a monotonic neural network; however, depending on the application, we may need tocompute α − 1 k for computing the inverse mapping, so it’s preferable to have an analytical inverse.
Revised: Other choices should alsowork, such as a monotonic neural network; however, depending on the application, we may need tocompute ↵  1 k for computing the inverse mapping, so it’s preferable to have an analytical inverse.

Original: 3.1 Remarks and Propositions
Revised: 3.1 Remarks and Propositions

Original: Box constraints.
Revised: Box constraints.

Original: There can be continuity problems if a Voronoi cell is unbounded, as the solutionto Equation (8) does not exist if x ( λ ∗ ) diverges.
Revised: There can be continuity problems if a Voronoi cell is unbounded, as the solutionto Equation (8) does not exist if x (  ⇤ ) diverges.

Original: Furthermore, when solving Equation (7), it can bedifficult to numerically distinguish between an unbounded cell and one whose boundaries are veryfar away.
Revised: Furthermore, when solving Equation (7), it can bedifﬁcult to numerically distinguish between an unbounded cell and one whose boundaries are veryfar away.

Original: It is for these reasons that we introduce box constraints (Equation 6) in the formulation of
Revised: It is for these reasons that we introduce box constraints (Equation 6) in the formulation of

Original: Voronoi cells which allows us to sidestep these issues for now.
Revised: Voronoi cells which allows us to sidestep these issues for now.

Original: Solving for λ ∗ .
Revised: Solving for  ⇤ .

Original: Equation (7) can be solved numerically, but this approach is prone to numericalerrors and requires implicit differentiation through convex optimization solutions [1].
Revised: Equation (7) can be solved numerically, but this approach is prone to numericalerrors and requires implicit differentiation through convex optimization solutions [1].

Original: We insteadnote that the solution of Equation (7) can be expressed in closed form, since it is always going to bethe intersection of the ray x ( λ ) with the nearest linear constraint.
Revised: We insteadnote that the solution of Equation (7) can be expressed in closed form, since it is always going to bethe intersection of the ray x (  ) with the nearest linear constraint.

Original: Let a T i x
Revised: Let a T i x

Original: = b i be the plane that represents one of the linear constraints in Equation (6), which areexpressed in Equation (4).
Revised: = b i be the plane that represents one of the linear constraints in Equation (6), which areexpressed in Equation (4).

Original: Let λ i be the intersection of this plane with the ray, i.e. it is the solutionto a T
Revised: Let  i be the intersection of this plane with the ray, i.e. it is the solutionto a T

Original: i x ( λ i ) =
Revised: i x (  i ) =

Original: b i , then the solution is simply the smallest positive λ i , which satisfies all the linearconstraints:
Revised: b i , then the solution is simply the smallest positive  i , which satisﬁes all the linearconstraints:

Original: [Equation]
Revised: [Equation]

Original: There are a total of K + 2 D − 1 linear constraints, including the Voronoi cell boundaries and boxconstraints, which can be computed fully in parallel.
Revised: There are a total of K + 2 D  1 linear constraints, including the Voronoi cell boundaries and boxconstraints, which can be computed fully in parallel.

Original: Automatic differentiation.

Original: Note that the above formulation of the solution λ ∗ also providesgradients ∗ dδ k ( x ) through automatic differentiation, since it only involves primitive operations. Thisallows end-to-end differentiation of the mapping f k , and will allow us to learn the parameters of f k ,i.e. parameters of α k and the Voronoi tessellation.
Revised: This also allows end-to-end differentiation ofthe mapping f k via automatic differentiation, providing the ability to learn the parameters of f k , i.e.parameters of ↵ k and the Voronoi tessellation.

Original: Homeomorphism with continuous density We can show that f k is a bijection, and both f k andf − 1 k are continuous.
Revised: Homeomorphism with continuous density We can show that f k is a bijection, and both f k andf  1 k are continuous.

Original: This allows us to use f k within the normalizing flows framework, as a mappingbetween a distribution p x defined on R D and the transformed distribution p z on V k .
Revised: This allows us to use f k within the normalizing ﬂows framework, as a mappingbetween a distribution p x deﬁned on R D and the transformed distribution p z on V k .

Original: Furthermore, the
Revised: Furthermore, the

Original: Jacobian is continuous almost everywhere.
Revised: Jacobian is continuous almost everywhere.

Original: Proofs are in Appendix A.
Revised: Proofs are in Appendix A.

Original: Proposition 1 The mapping f k : R D 
Revised: Proposition 1 The mapping f k : R D !

Original: → V k as defined in the 2-step procedure is a homeomorphism.
Revised: V k as deﬁned in the 2-step procedure is a homeomorphism.

Original: Proposition 2 If p x ( x ) is continuous, then the transformed density p z ( f k ( x )) is continuous a.e.
Revised: Proposition 2 If p x ( x ) is continuous, then the transformed density p z ( f k ( x )) is continuous a.e.

Original: Efficient computation of the log det Jacobian As f k is a mapping in D dimensions, computingthe log determinant Jacobian for likelihood computation can be costly and will scale poorly with D ifcomputed naïvely.
Revised: Efﬁcient computation of the log det Jacobian As f k is a mapping in D dimensions, computingthe log determinant Jacobian for likelihood computation can be costly and will scale poorly with D ifcomputed naïvely.

Original: Instead, we note that the Jacobian of f k is highly structured.
Revised: Instead, we note that the Jacobian of f k is highly structured.

Original: Intuitively, becausef k depends only on the direction δ k
Revised: Intuitively, becausef k depends only on the direction  k

Original: ( x ) and the distance away from x k , it only has two degrees offreedom regardless of D .
Revised: ( x ) and the distance away from x k , it only has two degrees offreedom regardless of D .

Original: In fact, the Jacobian of f k can be represented as a rank-2 update on ascaled identity matrix.
Revised: In fact, the Jacobian of f k can be represented as a rank-2 update on ascaled identity matrix.

Original: This allows us to use the matrix determinant lemma to reformulate the logdeterminant Jacobian in a compute- and memory-efficient form.
Revised: This allows us to use the matrix determinant lemma to reformulate the logdeterminant Jacobian in a compute- and memory-efﬁcient form.

Original: We summarize this in a proposition.
Revised: We summarize this in a proposition.

Original: Proposition 3 Let the transformation f k ( x ) and all intermediate quantities be as defined in Section 3for some given input x .
Revised: Proposition 3 Let the transformation f k ( x ) and all intermediate quantities be as deﬁned in Section 3for some given input x .

Original: Then the Jacobian factorizes as
Revised: Then the Jacobian factorizes as

Original: [Equation]
Revised: [Equation]

Original: 10) for some c ∈ R , u i ∈ R D , v i ∈ R D , and its log determinant has the form
Revised: for some c 2 R , u i 2 R D , v i 2 R D , and its log determinant has the formlog

Original: [Equation]
Revised: [Equation]

Original: +
Revised:

Original: [Equation]
Revised: [Equation]

Revised: Ordinal

Revised: Simplex

Revised: BinaryArgmax

Original: (11) where w ij = c − 1 v T i u j .
Revised: Voronoi where w ij = c  1 v T i u j .

Original: This expression for the log determinant only requires inner products betweenvectors in R D .
Revised: This expression for the log determinant only requires inner products betweenvectors in R D .

Original: To reduce notational clutter, exact formulas for c, u 1 , u 2 , v 1 , v 2 are in Appendix A.Ordinal
Revised: To reduce notational clutter, exact formulas for c, u 1 , u 2 , v 1 , v 2 are in Appendix A.

Original: Simplex BinaryArgmax Voronoi All vectors used in computing Equation (11) are either readily available as intermediate quantitiesafter computing f k ( x ) or are gradients of scalar functions and can be efficiently computed throughreverse-mode automatic differentiation.
Revised: All vectors used in computing Equation (11) are either readily available as intermediate quantitiesafter computing f k ( x ) or are gradients of scalar functions and can be efﬁciently computed throughreverse-mode automatic differentiation.

Original: The only operations on these vectors are dot products, andno large D -by- D matrices are ever constructed.
Revised: The only operations on these vectors are dot products, andno large D -by- D matrices are ever constructed.

Original: Compared to explicitly constructing the Jacobianmatrix, this is more efficient in both compute and memory, and can readily scale to large values of D .
Revised: Compared to explicitly constructing the Jacobianmatrix, this is more efﬁcient in both compute and memory, and can readily scale to large values of D .

Original:  Application settings
Revised:  Application settings

Original: We discuss two applications of our method to likelihood-based modeling of discrete and continuousdata.
Revised: We discuss two applications of our method to likelihood-based modeling of discrete and continuousdata.

Original: The first is a new approach to dequantization which allows training a model of discrete datausing density models that normally only act on continuous variables such as normalizing flows.
Revised: The ﬁrst is a new approach to dequantization which allows training a model of discrete datausing density models that normally only act on continuous variables such as normalizing ﬂows.

Original: Compared to existing dequantization methods [7, 18, 38], the Voronoi dequantization is not restrictedto any fixed-dimensional continuous space and can benefit from learning similarities between discretevalues.
Revised: Compared to existing dequantization methods [7, 18, 38], the Voronoi dequantization is not restrictedto ﬁxed dimensions and can beneﬁt from learning similarities between discrete values.

Original: In fact, existing approaches can be seen as special cases of a Voronoi dequantization withfixed equidistant anchor points.
Revised: In fact,existing approaches can be seen as special cases of a Voronoi dequantization.

Original: The second is a formulation of disjoint mixture models, where each Voronoi cell represents thesubspace of each component.
Revised: The second is a formulation of disjoint mixture models, where each component lies strictly withineach subset of a Voronoi tessellation.

Original: To the best of our knowledge, disjoint mixture models have not beenexplored significantly in the literature and have not been successfully applied in more than a coupledimensions.
Revised: To the best of our knowledge, disjoint mixture models have notbeen explored signiﬁcantly in the literature and have not been successfully applied in more than acouple dimensions.

Original: Compared to an existing method [8], the Voronoi mixture model is not restricted toacting individually for each dimension.
Revised: Compared to an existing method [8], the Voronoi mixture model is not restrictedto acting individually for each dimension.

Original: Voronoi dequantization Let y be a discrete random variable with finite support Y .
Revised: Voronoi dequantization Let y be a discrete random variable with ﬁnite support Y .

Original: Then we canuse tessellation to define subsets for dequantizing y as long as there are at least as many Voronoicells—equivalently, anchors points—as the number of discrete values, i.e. K ≥ |Y| .
Revised: Then we canuse tessellation to deﬁne subsets for dequantizing y as long as there are at least as many Voronoicells—equivalently, anchors points—as the number of discrete values, i.e. K  |Y| .

Original: By assigningeach discrete value to a Voronoi cell, we can then define a dequantization model q ( x | y ) by firstsampling from a base distribution z ∼ q ( z | y ) in R D and then applying the mapping x = f k ( z ) from
Revised: By assigning 171 each discrete value to a Voronoi cell, we can then deﬁne a dequantization model q ( x | y ) by ﬁrst 172 sampling from a base distribution z ⇠ q ( z | y ) in R D and then applying the mapping x = f k ( z ) from

Original: Section 3 to construct a distribution over V k .
Revised: Section 3 to construct a distribution over V k .

Original: We can then obtain probabilities q ( x | y ) efficientlyusing Proposition 3 and train the dequantization alongside the density model p ( x ) on R D .
Revised: We can then obtain probabilities q ( x | y ) efﬁciently andtrain the dequantization alongside the density model p ( x ) on R D .

Original: Gradientsthrough f k provide signal for training parameters of the tessellation.

Original: Sampling from the model p ( y | x ) is straightforward and deterministic after sampling x .
Revised: Sampling from the model p ( y | x ) is straightforward and deterministic after sampling x .

Original: We can writey =
Revised: We can writey =

Original: g ( x ) where g is the set identification function satisfying x ∈ V g ( x ) .
Revised: g ( x ) where g is the set identiﬁcation function satisfying x 2 V g ( x ) .

Original: From Equation (4), it iseasy to see that g ( x ) is the nearest neighbor operation, g ( x ) =
Revised: From Equation (4), it iseasy to see that g ( x ) is the nearest neighbor operation, g ( x ) =

Original: arg min k ∥ x − x k ∥ .
Revised: arg min k k x  x k k .

Original: We are free to choose the number of dimensions D , where a smaller D assigns less space for each
Revised: We are free to choose the number of dimensions D , where a smaller D assigns less space for each

Original: Voronoi cell induces dequantized distributions that are easier to fit, while a larger D allows the anchorpoints and Voronoi cells more room to change over the course of training.
Revised: Voronoi cell induces dequantized distributions that are easier to ﬁt, while a larger D allows the anchorpoints and Voronoi cells more room to change over the course of training.

Original: When the anchor pointsare fixed to specific positions, we can recover the disjoint subsets used by prior methods as specialcases.
Revised: When the anchor pointsare ﬁxed to speciﬁc positions, we can recover the disjoint subsets used by prior methods as specialcases.

Original: We illustrate this in Figure 2.
Revised: We illustrate this in Figure 3.

Original: Voronoi mixture models Let { V k } be a Voronoi tessellation of R D .
Revised: Voronoi mixture models Let { V k } be a Voronoi tessellation of R D .

Original: Then a disjoint mixturemodel can be constructed by defining distributions on each Voronoi cell.
Revised: Then a disjoint mixturemodel can be constructed by deﬁning distributions on each Voronoi cell.

Original: Here we make use of theinverse mapping f − 1 k : V k 
Revised: Here we make use of theinverse mapping f  1 k : V k !

Original: → R D (discussed in Appendix B) so that we only need to parameterizedistributions over R D .
Revised: R D (discussed in Appendix B) so that we only need to parameterizedistributions over R D .

Original: Let x be a point in R D , our Voronoi mixture model assigns the densitylog p mix ( x ) = log p comp ( f − 1 k = g ( x ) ( x ) |
Revised: Let x be a point in R D , our Voronoi mixture model assigns the densitylog p mix ( x ) = log p comp [Equation]

Original: k = g ( x ))
Revised: [Equation]

Original: + log (cid:12)(cid:12)(cid:12) det ∂f − 1 k ( x ) ∂ x (cid:12)(cid:12)(cid:12)
Revised:

Original: [Equation]
Revised: [Equation]

Original: where p comp can be any distribution over R D , including another disjoint mixture model.
Revised: where p comp can be any distribution over R D , including another disjoint mixture model.

Original: This caneasily be composed with normalizing flow layers where, in addition to the change of variable due tof − 1 k , we also apply the change in density resulting from choosing one out of K components.
Revised: This caneasily be composed with normalizing ﬂow layers where, in addition to the change of variable due tof  1 k , we also apply the change in density resulting from choosing one out of K components.

Original:  Related Work
Revised:  Related Work

Original: Normalizing flows for discrete data.
Revised: Normalizing ﬂows for discrete data.

Original: Invertible mappings have been proposed for discrete data,where discrete values are effectively rearranged from a factorized distribution.
Revised: Invertible mappings have been proposed for discrete data,where discrete values are effectively rearranged from a factorized distribution.

Original: In order to parameterize the transformation, Hoogeboom et al.
Revised: In order to parameterize the transformation, Hoogeboom et al.

Original: [17], van den Berg et al.
Revised: [17], van den Berg et al.

Original: [46] use quantized ordinaltransformations, while Tran et al.
Revised: [47] use quantized ordinaltransformations, while Tran et al.

Original: [45] takes a more general approach of using modulo addition onone-hot vectors.
Revised: [45] takes a more general approach of using modulo addition onone-hot vectors.

Original: These approaches suffer from gradient bias due to the need to use discontinuousoperations and do not have universality guarantees since it’s unclear whether simple rearrangement issufficient to transform any joint discrete distribution into a factorized distribution.
Revised: These approaches suffer from gradient bias due to the need to use discontinuousoperations and do not have universality guarantees since it’s unclear whether simple rearrangement issufﬁcient to transform any joint discrete distribution into a factorized distribution.

Original: In contrast, thedequantization approach provides a universality guarantee since the lower bound in Equation (2) istight when p ( x ) 1 [ x ∈ A y ] ∝ q ( x | y ) , with a proportionality equal to p ( y ) .
Revised: In contrast, thedequantization approach provides a universality guarantee since the lower bound in Equation (2) istight when p ( x ) 1 [ x 2 A y ] / q ( x | y ) , with a proportionality equal to p ( y ) .

Original: Dequantization methods.
Revised: Dequantization methods.

Original: Within the normalizing flows literature, the act of adding noise wasoriginally used for ordinal data as a way to combat numerical issues [7].
Revised: Within the normalizing ﬂows literature, the act of adding noise wasoriginally used for ordinal data as a way to combat numerical issues [7, 46].

Original: Later on, appropriatedequantization approaches have been shown to lower bound the log-likelihood of a discrete model [16,44].
Revised: Later on, appropriatedequantization approaches have been shown to lower bound the log-likelihood of a discrete model [16,44].

Original: For non-ordinal data, many works have proposed simplex-based approaches.
Revised: For non-ordinal data, many works have proposed simplex-based approaches.

Original: Early works onrelaxations [20, 32] proposed continuous distribution on the simplex that mimic the behaviors ofdiscrete random variables; however, these were only designed for the use with a Gumbel basedistribution.
Revised: Early works onrelaxations [20, 32] proposed continuous distribution on the simplex that mimic the behaviors ofdiscrete random variables; however, these were only designed for the use with a Gumbel basedistribution.

Original: Potapczynski et al.
Revised: Potapczynski et al.

Original: [38] extend this to a Gaussian distribution—although it is not hard tosee this can work with any base distribution—by designing invertible transformations between R Dand the probability simplex with K vertices, with D = K − 1 , where K is the number of classes of adiscrete random variable.
Revised: [38] extend this to a Gaussian distribution—although it is not hard tosee this can work with any base distribution—by designing invertible transformations between R Dand the probability simplex with K vertices, with D = K  1 , where K is the number of classes of adiscrete random variable.

Original: Intuitively, after fixing one of the logits, the softmax operation is an invertible transformation on thebounded domain ∆ K − 1 = { x ∈ R K : (cid:80)
Revised: Intuitively, after ﬁxing one of the logits, the softmax operation is an invertible transformation on thebounded domain  K  1 = { x 2 R K : P

Original: Ki x i = 1 , x i > 0 } .
Revised: Ki x i = 1 , x i > 0 } .

Original: The ( K − 1) -simplex can then bebroken into K subsets, each corresponding to a particular discrete value.
Revised: The ( K  1) -simplex can then bebroken into K subsets, each corresponding to a particular discrete value.

Original: [Equation]
Revised: [Equation]

Original: More recently, Hoogeboom et al.
Revised: More recently, Hoogeboom et al.

Original: [18] proposed ignoring the simplex constraint and simply use
Revised: [18] proposed ignoring the simplex constraint and simply use

Original: [Equation]
Revised: [Equation]

Original:  which effectively increases the number of dimensions by one compared to the simplex approach.
Revised: 14) which effectively increases the number of dimensions by one compared to the simplex approach.

Original: However, bothsimplex-based approaches force the dimension of the continuous space D to scalewith K .
Revised: However, both approaches force the dimension of the continuous space D to scale with K .

Original: In order to make simplex-based dequantization work when K is large, they propose reducingall discrete variables to a set of binary variables before applying dequantization.
Revised: In orderto make their approach work when K is large, [18] propose binarizing the data as a preprocessingstep.

Original: In contrast, Voronoidequantization has full flexibility in choosing D regardless of K .
Revised: In contrast, Voronoi dequantization has full ﬂexibility in choosing D regardless of K .

Original: Related models for discrete data Among other related methods, a recent work proposed normalizing flows on truncated supports [43] but had to resort to rejection sampling for training.
Revised: Related models for discrete data Among other related methods, a recent work proposed normalizing ﬂows on truncated supports [43] but had to resort to rejection sampling for training.

Original: Furthermore,their subsets are not disjoint by construction.
Revised: Furthermore,their subsets are not disjoint by construction.

Original: Prior works [28, 31] also proposed removing theconstraint that subsets are disjoint, and instead work with general mixture models with unboundedsupport, relying on the conditional model p ( y | x ) being sufficiently weak so that the task of modelingis forced onto a flow-based prior.
Revised: Prior works [28, 31] also proposed removing theconstraint that subsets are disjoint, and instead work with general mixture models with unboundedsupport, relying on the conditional model p ( y | x ) being sufﬁciently weak so that the task of modelingis forced onto a ﬂow-based prior.

Original: Similar to general approaches that combine normalizing flows withvariational inferenceapproaches [19, 33, 49], they have achieved good performance on a number oftasks.
Revised: They have achieved good performance on a number of tasks,similar to general approaches that combine normalizing ﬂows with variational inference [19, 33, 50].

Original: However, they lose the computational savings and the deterministic decoder p ( y | x ) gainedfrom using disjoint subsets.
Revised: However, they lose the computational savings and the deterministic decoder p ( y | x ) gained fromusing disjoint subsets.

Original: On the other hand, quantization based on nearest neighbor have beenused for learning discrete latent variables [34, 40], but no likelihoods are constructed, the boundariesare not explicitly differentiated through, and the model relies on training with heuristical updates.
Revised: On the other hand, quantization based on nearest neighbor have been used forlearning discrete latent variables [34, 40], but no likelihoods are constructed, the boundaries are notexplicitly differentiated through, and the model relies on training with heuristical updates.

Original: Disjoint mixture models.
Revised: Disjoint mixture models.

Original: The computational savings from using disjoint subsets was pointed outby
Revised: The computational savings from using disjoint subsets was pointed outby

Original: Dinh et al.
Revised: Dinh et al.

Original: However, their method only works in each dimension individually.
Revised: However, their method only works in each dimension individually.

Original: They transformsamples using a linear spline, which is equivalent to creating subsets based on the knots of the splineand applying a linear transformation within each subset.
Revised: They transformsamples using a linear spline, which is equivalent to creating subsets based on the knots of the splineand applying a linear transformation within each subset.

Original: Furthermore, certain parameterizations ofthe spline can lead to discontinuous density functions, whereas our disjoint mixture has a density{+densityfunction that is continuous almost everywhere (albeit exactly zero on the boundaries; see Section 
+}
Revised: Furthermore, certain parameterizations ofthe spline can lead to discontinuous density functions, whereas our disjoint mixture has a {+densityfunction that is continuous almost everywhere (albeit exactly zero on the boundaries; see Section 
+}

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: function that is continuous almost everywhere (albeit exactly zero on the boundaries; see Section 7for more discussion).
Revised: for more discussion).

Original: The use of monotonic splines have previously been combined with normalizingflows [10], but interestingly, since the splines in Dinh et al.
Revised: The use of monotonic splines have previously been combined with normalizingﬂows [10], but interestingly, since the splines in Dinh et al.

Original: [8] are not enforced to be monotonic,the full transformation is not bijective and acts like a disjoint mixture model.
Revised: [8] are not enforced to be monotonic,the full transformation is not bijective and acts like a disjoint mixture model.

Original: Ultimately, theirexperiments were restricted to two-dimensional syntheic data sets, and it remained an interestingresearch question whether disjoint mixture models can be successfully applied in high dimensions.
Revised: Ultimately, theirexperiments were restricted to two-dimensional syntheic data sets, and it remained an interestingresearch question whether disjoint mixture models can be successfully applied in high dimensions.

Original:  Experiments
Revised:  Experiments

Original: We experimentally validate our semi-discrete approach of combining Voronoi tessellation withlikelihood-based modeling on a variety of data domains: discrete-valued UCI data sets, itemsetmodeling, language modeling, and disjoint mixture modeling for continuous UCI data.
Revised: We experimentally validate our semi-discrete approach of combining Voronoi tessellation withlikelihood-based modeling on a variety of data domains: discrete-valued UCI data sets, itemsetmodeling, language modeling, and disjoint mixture modeling for continuous UCI data.

Original: The goal ofthese experiments is not to show state-of-the-art results in these domains but to showcase the relativemerit of Voronoi tessellation compared to existing methods.
Revised: The goal ofthese experiments is not to show state-of-the-art results in these domains but to showcase the relativemerit of Voronoi tessellation compared to existing methods.

Original: For this reason, we just use simplecoupling blocks with affine transformations [6, 7] as a base flow model in our experiments, unlessstated otherwise.
Revised: For this reason, we just use simplecoupling blocks with afﬁne transformations [6, 7] as a base ﬂow model in our experiments, unlessstated otherwise.

Original: When comparing to Discrete Flows [45], we use their bipartite layers.
Revised: When comparing to Discrete Flows [45], we use their bipartite layers.

Original: Detailsregarding preprocessing for data sets can be found in Appendix C, and detailed experimental setup isin Appendix D. All tables, when shown, display standard deviations across three random seeds.
Revised: Detailsregarding preprocessing for data sets can be found in Appendix C, and detailed experimental setup isin Appendix D. All tables, when shown, display standard deviations across three random seeds.

Original: 6.1 Quantized 2D data
Revised: 6.1 Quantized 2D data

Original: We start by considering synthetic distributions over two discrete variables, to qualitatively compareagainst Discrete Flows [45].
Revised: We start by considering synthetic distributions over two discrete variables, to qualitatively compareagainst Discrete Flows [45].

Original: Figure 3 shows the target probability mass functions, which are createdby a 2D distribution where each dimension is quantized into 91 discrete values.
Revised: Figure 4 shows the target probability mass functions, which are createdby a 2D distribution where each dimension is quantized into 91 discrete values.

Original: Though the data isordinal, no knowledge of ordering is given to the models.
Revised: Though the data isordinal, no knowledge of ordering is given to the models.

Original: We see that Discrete Flows has troublefitting these, because it is difficult to rearrange the target distribution into a factorized distribution.
Revised: We see that Discrete Flows has troubleﬁtting these, because it is difﬁcult to rearrange the target distribution into a factorized distribution.

Original: For our model, we dequantize each discrete variable with a Voronoi tessellation in R 2 .
Revised: For our model, we dequantize each discrete variable with a Voronoi tessellation in R 2 .

Original: We then learna flow model on the combined R 4 , parameterized by multilayer perceptrons (MLPs).
Revised: We then learna ﬂow model on the combined R 4 , parameterized by multilayer perceptrons (MLPs).

Original: In Figure 4we visualize the learned Voronoi tessellation and samples from our model.
Revised: In Figure 5we visualize the learned Voronoi tessellation and samples from our model.

Original: The learned tessellationseems to group some of discrete values that occur frequently together, to reduce the number of modes.
Revised: The learned tessellationseems to group some of discrete values that occur frequently together, to reduce the number of modes.

Original: 6.2 Discrete-valued UCI data
Revised: 6.2 Discrete-valued UCI data

Original: We experiment with complex data sets where each discrete variable can have a varying numberof classes.
Revised: We experiment with complex data sets where each discrete variable can have a varying numberof classes.

Original: Furthermore, these discrete variables may have hidden semantics.
Revised: Furthermore, these discrete variables may have hidden semantics.

Original: To this end, we useunprocessed data sets from the UCI database [9].
Revised: To this end, we useunprocessed data sets from the UCI database [9].

Original: The only pre-processing we perform is to removevariables that only have one discrete value.
Revised: The only pre-processing we perform is to removevariables that only have one discrete value.

Original: We then take 80% as train, 10% as validation, and 10% asthe test set.
Revised: We then take 80% as train, 10% as validation, and 10% asthe test set.

Original: Most of these data sets have a combination of both ordinal and non-ordinal variables, andwe expect the non-ordinal variables to exhibit relations that are unknown ( e.g. spatial correlations).
Revised: Most of these data sets have a combination of both ordinal and non-ordinal variables, andwe expect the non-ordinal variables to exhibit relations that are unknown ( e.g. spatial correlations).

Original: {+We see that Discrete Flows can be competitive with dequantization approaches, but can also fall shorton more challenging data sets such as the USCensus90 , the largest data set we consider with 2.
+}
Revised: {+We see that Discrete Flows can be competitive with dequantization approaches, but can also fall shorton more challenging data sets such as the USCensus90 , the largest data set we consider with 2.
+}

Original: [Table]
Revised: [Table]

Original: [Table] 
Revised: [Table] [Table] 42 Ordinal [18] [Table]

Original: [Table]
Revised: [Table]

Original: We see that Discrete Flows can be competitive with dequantization approaches, but can also fall shorton more challenging data sets such as the USCensus90 , the largest data set we consider with 2.7million examples and 68 different discrete variables of varying types.
Revised: million examples and 68 different discrete variables of varying types.

Original: For dequantization methods,simplex [38] and binary argmax [18] approaches are mostly on par.
Revised: For dequantization methods,simplex [38] and binary argmax [18] approaches are mostly on par.

Original: We do see a non-negligible gapin performance between these baselines and Voronoi dequantization for most of the data sets, likelydue to the ability to learn semantically useful relations between the values of each discrete variable.
Revised: We do see a non-negligible gapin performance between these baselines and Voronoi dequantization for most of the data sets, likelydue to the ability to learn semantically useful relations between the values of each discrete variable.

Original: For instance, the Connect4 data set contains positions of a two-player board game with the samename, which exhibit complex spatial dependencies between the board pieces, and the USCensus90data set contains highly correlated variables and is often used to test clustering algorithms.
Revised: For instance, the Connect4 data set contains positions of a two-player board game with the samename, which exhibit complex spatial dependencies between the board pieces, and the USCensus90data set contains highly correlated variables and is often used to test clustering algorithms.

Original: 6.3 Permutation-invariant itemset modeling
Revised: 6.3 Permutation-invariant itemset modeling

Original: An appeal of using normalizing flows in continuous space is the ability to incorporate invarianceto specific group symmetries into the density model.
Revised: An appeal of using normalizing ﬂows in continuous space is the ability to incorporate invarianceto speciﬁc group symmetries into the density model.

Original: For instance, this can be done by ensuringthe ordinary differential equation is equivariant [2, 25, 42] in a continuous normalizing flow [4, 13]framework.
Revised: For instance, this can be done by ensuringthe ordinary differential equation is equivariant [2, 25, 42] in a continuous normalizing ﬂow [4, 13]framework.

Original: Here we focus on invariance with respect to permutations, i.e. sets of discrete variables.
Revised: Here we focus on invariance with respect to permutations, i.e. sets of discrete variables.

Original: This invariance cannot be explicitly modeled by Discrete Flows as they require an ordering or abipartition of discrete variables.
Revised: This invariance cannot be explicitly modeled by Discrete Flows as they require an ordering or abipartition of discrete variables.

Original: We preprocessed a data set of retail market basket data from ananonymous Belgian retail store [3] and a data set of anonymized traffic accident data [11], whichcontain sets of discrete variables each with 765 and 269 values, respectively.
Revised: We preprocessed a data set of retail market basket data from ananonymous Belgian retail store [3] and a data set of anonymized trafﬁc accident data [11], whichcontain sets of discrete variables each with 765 and 269 values, respectively.

Original: Without the binarizationtrick, simplex dequantization must use a large embedding dimensions and performs worse than adeterminental point process baseline [26].
Revised: Without the binarizationtrick, simplex dequantization must use a large embedding dimensions and performs worse than adeterminental point process baseline [26].

Original: Meanwhile, our Voronoi dequantization has no problemsstaying competitive as its embedding space is freely decoupled from the number of discrete values.
Revised: Meanwhile, our Voronoi dequantization has no problemsstaying competitive as its embedding space is freely decoupled from the number of discrete values.

Original: 6.4 Language modeling
Revised: 6.4 Language modeling

Original: Language modeling a widely used benchmark for discrete models [18, 28, 45].
Revised: Language modeling a widely used benchmark for discrete models [18, 28, 45].

Original: Here we usedthe open source code provided by Hoogeboom et al.
Revised: Here we usedthe open source code provided by Hoogeboom et al.

Original: [18] with the exact same autoregressive flowmodel and optimizer setups.
Revised: [18] with the exact same autoregressive ﬂowmodel and optimizer setups.

Original: The only difference is replacing their binary argmax with our Voronoidequantization.
Revised: The only difference is replacing their binary argmax with our Voronoidequantization.

Original: Results are shown in Table 3, where we tried out multiple embedding dimensions D .
Revised: Results are shown in Table 3, where we tried out multiple embedding dimensions D .

Original: Generally, we find D = 2 to be too low and can stagnate training since the Voronoi cells are moreconstrained, while larger values of D improve upon argmax dequantization.
Revised: Generally, we ﬁnd D = 2 to be too low and can stagnate training since the Voronoi cells are moreconstrained, while larger values of D improve upon argmax dequantization.

Original: 6.5 Voronoi mixture models
Revised: 6.5 Voronoi mixture models

Original: We test disjoint mixture modeling with Voronoi tessellation on continuous data sets.
Revised: We test disjoint mixture modeling with Voronoi tessellation on continuous data sets.

Original: Figure 5 showsresults of training a simple normalizing flow on 2D data, which has trouble completely separating allthe modes.
Revised: Figure 6 showsresults of training a simple normalizing ﬂow on 2D data, which has trouble completely separating allthe modes.

Original: Adding in a Voronoi mixture allows us to better fit these multimodal densities.Data
Revised: Adding in a Voronoi mixture allows us to better ﬁt these multimodal densities.

Original: Baseline [Table] [Table] To test Voronoi mixture models on higher dimensions, we apply our approach to data sets preprocessedby Papamakarios et al.
Revised: To test Voronoi mixture models on higher dimensions, we apply our approach to data sets preprocessedby Papamakarios et al.

Original: [35], which range from 6 to 63 dimensions.
Revised: [35], which range from 6 to 63 dimensions.

Original: We add a disjoint mixturecomponent after a few layers of normalizing flows, with each component modeled by a normalizingflow conditioned on a embedding of the component index.
Revised: We add a disjoint mixturecomponent after a few layers of normalizing ﬂows, with each component modeled by a normalizing [Table] [Table] ﬂow conditioned on a embedding of the component index.

Original: The number of mixture components rangesfrom 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 thedisjoint mixture, as well as a strong but computationally costly density model FFJORD [13].
Revised: The number of mixture components rangesfrom 16 to 64 , with complete sweeps and hyperparameters detailed in Appendix D. For comparison,we also show results from the baseline coupling ﬂow that uses the same number of total layers as thedisjoint mixture, as well as a strong but computationally costly density model FFJORD [13].

Original: From
Revised: From

Original: Table 4, we see that the disjoint mixture model approach allows us to increase complexity quirte a bit,with almost no additional cost compared to the baseline flow model.
Revised: Table [Table] we see that the disjoint mixture model approach allows us to increase complexity quite a bit,with almost no additional cost compared to the baseline ﬂow model [Table]

Original:  Conclusion and Discussion
Revised:  Conclusion and Discussion

Original: We combine Voronoi tessellation with normalizing flows to construct a new invertible transformationthat has learnable discrete structure.
Revised: We combine Voronoi tessellation with normalizing ﬂows to construct a new invertible transformationthat has learnable discrete structure.

Original: This acts as a learnable mapping between discrete and continuousdistributions.
Revised: This acts as a learnable mapping between discrete and continuousdistributions.

Original: We propose two applications of this method: a Voronoi dequantization that mapsdiscrete values into a learnable convex polytope, and a Voronoi mixture modeling approach thathas around the same compute cost as a single component.
Revised: We propose two applications of this method: a Voronoi dequantization that mapsdiscrete values into a learnable convex polytope, and a Voronoi mixture modeling approach thathas around the same compute cost as a single component.

Original: We showcased the relative merit of ourapproach across a range of data modalities with varying hidden structure.
Revised: We showcased the relative merit of ourapproach across a range of data modalities with varying hidden structure.

Original: Below, we discuss somelimitations and possible future directions.
Revised: Below, we discuss somelimitations and possible future directions.

Original: Diminishing density on boundaries.
Revised: Diminishing density on boundaries.

Original: The distributions within each Voronoi cell necessarily goto zero on the boundary between cells due to the use of a homeomorphism from an unboundeddomain.
Revised: The distributions within each Voronoi cell necessarily goto zero on the boundary between cells due to the use of a homeomorphism from an unboundeddomain.

Original: This is a different and undesirable property as opposed to the partitioning method used by[8].
Revised: This is a different and undesirable property as opposed to the partitioning method used by[8].

Original: However, alternatives could require more compute cost in the form of solving normalizationconstants.
Revised: However, alternatives could require more compute cost in the form of solving normalizationconstants.

Original: Balancing expressiveness and compute cost is a delicate problem that sits at the core oftractable probabilistic modeling.
Revised: Balancing expressiveness and compute cost is a delicate problem that sits at the core oftractable probabilistic modeling.

Original: Design of homeomorphisms and tessellations.
Revised: Design of homeomorphisms and tessellations.

Original: Our proposed homeomorphism is simple andscalable, but this comes at the cost of smoothness and expressiveness.
Revised: Our proposed homeomorphism is simple andscalable, but this comes at the cost of smoothness and expressiveness.

Original: As depicted in Figure 1, thetransformed density has a non-smooth probability density function.
Revised: As depicted in Figure 1, thetransformed density has a non-smooth probability density function.

Original: This non-smoothness existswhen the ray intersects with multiple boundaries.
Revised: This non-smoothness existswhen the ray intersects with multiple boundaries.

Original: This may make optimization more difficult as thegradients of the log-likelihood objective can be discontinuous.
Revised: This may make optimization more difﬁcult as thegradients of the log-likelihood objective can be discontinuous.

Original: Additionally, our use of Euclideandistance can become problematic in high dimensions, as this metric can cause almost all pointsto be nearly equidistant, resulting in a behavior where all points lie seemingly very close to theboundary.
Revised: Additionally, our use of Euclideandistance can become problematic in high dimensions, as this metric can cause almost all pointsto be nearly equidistant, resulting in a behavior where all points lie seemingly very close to theboundary.

Original: Improvements on the design of homeomorphisms to bounded domains could help alleviatethese problems.
Revised: Improvements on the design of homeomorphisms to bounded domains could help alleviatethese problems.

Original: Additionally, more flexible tessellations—such as the Laguerre tessellation—andadditional concepts from semi-discrete optimal transport [14, 27, 37] may be adapted to improvesemi-discrete normalizing flows.
Revised: Additionally, more ﬂexible tessellations—such as the Laguerre tessellation—andadditional concepts from semi-discrete optimal transport [14, 27, 37] may be adapted to improvesemi-discrete normalizing ﬂows.

Original: Checklist
Revised: Checklist

Original: For all authors
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and

Original: contributions and scope?
Revised: scope?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes]
Revised: Yes

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes]
Revised: [Yes]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?
Revised: them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [Yes]
Revised: [Yes]

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [No]
Revised: [No]

Original: (b) Did you specify all the training details (e
Revised: (b) Did you specify all the training details (e

Original: g., data splits, hyperparameters,
Revised: g., data splits, hyperparameters,

Original: how they
Revised: how they

Original: were chosen)?
Revised: were chosen)?

Original: [Yes]
Revised: [Yes]

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi-
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experi-

Original: ments multiple times)?
Revised: ments multiple times)?

Original: [Yes]
Revised: [Yes]

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., type
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., type

Original: of GPUs, internal cluster, or cloud provider)?
Revised: of GPUs, internal cluster, or cloud provider)?

Original: [Yes]
Revised: [Yes]

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [No]
Revised: [No]

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No]
Revised: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

Original: [Yes]
Revised: [Yes]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable
Revised: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Original: information or offensive content?
Revised: information or offensive content?

Original: [No]
Revised: [No]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: [N/A]
Revised: [N/A]