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

Original: Abstract
Revised: Abstract

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

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

Original: This is done through constructing normalizing flows on convex polytopes parameterized using a simple homeomorphism with an efficient log determinant Jacobian.
Revised: This is done through constructing normalizing ﬂows on convex polytopes parameterized using a simple homeomorphism with an 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 Voronoi dequantization allows automatically learning quantization boundaries in a multidimensional space.
Revised: Firstly, a Voronoi dequantization allows automatically learning quantization boundaries in a 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 near-constant computation cost for likelihood evaluation regardless of the number of mixture components.
Revised: Secondly, a Voronoi mixture model has near-constant computation cost for likelihood evaluation regardless of the number of mixture components.

Original: Empirically, we show improvements over existing methods across a range of structured data modalities.
Revised: Empirically, we show improvements 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 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 a wide variety of domains such as images [7, 23], graphs [29], invariant distributions [2, 25] and molecular structures [42].
Revised: Across a variety of modeling approaches, the family of normalizing ﬂows stands out as a large amount of structure can be incorporated into the model, aiding its usage in modeling a wide variety of domains such as images [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: 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.
Revised: This restriction 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 entirely composable with any other normalizing flow.
Revised: In this work we incorporate discrete structure into standard normalizing ﬂows, while being entirely composable with any other normalizing ﬂow.

Original: Specifically, we propose a homeomorphism between the unbounded domain and convex polytopes, which are defined through a learnable tessellation of the domain, i.e. a set of disjoint subsets that together fully cover the domain.
Revised: Speciﬁcally, we propose a homeomorphism between the unbounded domain and convex polytopes, which are deﬁned through a learnable tessellation of the 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 results in an efficient formulation of the resulting change in density.
Revised: This homeomorphism is cheap to compute, has a cheap inverse, and results in an efﬁcient formulation of the resulting change in density.

Original: In other words, this transformation 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 tessellation naturally allows defining a dequantization method for discrete variables.
Revised: Firstly, the learned tessellation naturally allows deﬁning a dequantization method for discrete variables.

Original: This allows likelihood-based models that normally only act on continuous spaces, such as normalizing flows, to work directly on discrete data in a flexible choice of embedding space with the ability to capture relational structure.
Revised: This allows likelihood-based models that normally only act on continuous spaces, such as normalizing ﬂows, to work directly on discrete 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 full tessellation 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 full tessellation deﬁnes a mixture model with disjoint components.

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

Original: 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.
Revised: Following semi-discrete optimal transportation [37], which deﬁnes couplings between discrete and continuous measures, we refer to our models as semi-discrete normalizing ﬂows that learn transformations between discrete and continuous random variables.

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Revised: f

Original: [Figure]
Revised: [Equation]

Original:  Preliminaries
Revised:  Preliminaries

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

Original: The relationship between the original and transformed density functions have a closed form expression,
Revised: The relationship between the original 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 close to 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 close to some target distribution.

Original: Many different choices of f have been discussed in the literature, leading to different trade-offs and use cases [24, 36].
Revised: Many different choices of f have been discussed in the literature, leading to 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 different supports, 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 different supports, i.e. regions where probability is non-zero.

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

Original: To the best of our knowledge, the use of general support-modifying invertible transformations have not been discussed extensively in the literature.
Revised: To the best of our knowledge, the use of general support-modifying invertible transformations have not been discussed extensively in the literature.

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

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

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

Original: Let A y denote the subset corresponding to y and let q ( x | y ) be the dequantization model which has 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 a bounded support in A y .

Original: Then any density model p ( x ) with support over R D satisfies log p ( y ) ≥ E x ∼
Revised: Then any density model p ( x ) with support over R D satisﬁes log p ( y )  E x ⇠

Original: [Equation]
Revised: [Equation]

Original:  Thus with an appropriate choice of dequantization, maximizing the likelihood under the density model p ( x ) is equivalent to maximizing the likelihood under the discrete model p ( y ) .
Revised: ( x | y ) [Equation] Thus with an appropriate choice of dequantization, maximizing the likelihood under the density model 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 dependent on the type of discrete data (ordinal vs non-ordinal) and the number of discrete values |Y| in the non-ordinal case.
Revised: A y are entirely dependent on the type of discrete data (ordinal vs non-ordinal) and the number of discrete values |Y| in the non-ordinal case.

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

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

Original: 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].
Revised: Disjoint mixture models Building mixture models is one of the simplest methods for creating more ﬂexible distributions from simpler one, and mixture models can typically be shown to be universal density estimators [12].

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

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

Original: To alleviate this issue, Dinh et al.
Revised: To alleviate this 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 to 
Revised: The likelihood of the mixture model then simpliﬁes to p

Original: [Equation]
Revised: [Equation]

Original:  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 framework allows building large mixture models while no longer having a compute cost that scales with K .
Revised: This framework allows building large mixture models while no longer having a compute cost that scales with K .

Original: In contrast to variational approaches with discrete latent variables, the use of disjoint subsets provides an exact (log-)likelihood and not a lower bound.
Revised: In contrast to variational approaches with discrete latent variables, the use of disjoint subsets provides an 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, discuss constructing 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, discuss constructing distributions on each convex polytope using a differentiable homeomorphism.

Original: Parameterizing disjoint subsets We separate the domain into subsets through a Voronoi tessellation [48].
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 , which provides a differentiable parameterization of the tessellation for gradient-based optimization.
Revised: This induces a correspondence between each subset and a corresponding anchor point , which provides 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 subset containing all points which have the anchor point as their nearest neighbor.
Revised: i.e. it deﬁnes a subset containing all points which have the anchor point as their nearest neighbor.

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

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

Original: [Equation]
Revised: [Equation]

Original:  , where a 
Revised: } , where a i = 2( x i  x [Equation] [Equation] [Equation]

Revised: =

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][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 invertible transformation and resulting distribution defined in each Voronoi cell, which we discuss next.
Revised: These will be trained using gradients from an invertible transformation and resulting distribution deﬁned in each Voronoi cell, which we discuss next.

Original: [Figure] Invertible mapping onto subset Within the normalizing flow framework, defining a probability distribution on a bounded support is equivalent to defining an invertible transformation that maps from the unbounded support onto the appropriate support.
Revised: Invertible mapping onto subset Within the normalizing ﬂow framework, deﬁning a probability distribution on a bounded support is equivalent to deﬁning an invertible transformation that maps from the unbounded support onto the appropriate support.

Original: Figureillustrates the transformation of mapping onto a shaded region.
Revised: Figureillustrates the transformation of mapping onto a shaded region.

Original: Given a cell V k for some k ∈ { 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 invertible mapping f

Original: k
Revised: k

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

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

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

Original: Otherwise: 
Revised: 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 k is convex, we can frame this as a linear programming problem
Revised: Since V k is convex, we can frame this as a linear programming problem.

Original: [Equation]
Revised: [Equation]

Original:  where V k is the closure of V k .
Revised: (7) 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 of box constraints.
Revised: This solution exists if V k is bounded, which is always true due to the use of 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 this knowledge, we can now map all points that lie on this ray onto the Voronoi cell.
Revised: Using this knowledge, 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 ( λ ) : λ >
Revised: Apply an invertible transformation such that any point on the ray { x (  ) :

Original:  is mapped onto theline segment
Revised: > 0 } is mapped onto the

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

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 a monotonic transformation of the relative distance from x to the anchor point x k .
Revised: An intuitive choice is to use a monotonic transformation of the relative distance from x to the anchor point x k .

Original: [Equation]
Revised: [Equation]

Original: (8) 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, we use α k ( h ) =
Revised: In our experiments, we use ↵ k ( h ) =

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 also work, such as a monotonic neural network; however, depending on the application, we may need to compute α − 1 k for computing the inverse mapping, so it’s preferable to have an analytical inverse.
Revised: Other choices should also work, such as a monotonic neural network; however, depending on the application, we may need to compute ↵  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 There can be continuity problems if a Voronoi cell is unbounded, as the solution to Equation (8) does not exist if x ( λ ∗ ) diverges.
Revised: Box constraints There can be continuity problems if a Voronoi cell is unbounded, as the solution to Equation (8) does not exist if x (  ⇤ ) diverges.

Original: Furthermore, when solving Equation (7), it can be difficult to numerically distinguish between an unbounded cell and one whose boundaries are very far away.
Revised: Furthermore, when solving Equation (7), it can be difﬁcult to numerically distinguish between an unbounded cell and one whose boundaries are very far 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 numerical errors and requires implicit differentiation through convex optimization solutions [1].
Revised: Equation (7) can be solved numerically, but this approach is prone to numerical errors and requires implicit differentiation through convex optimization solutions [1].

Original: 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 x ( λ ) with the nearest linear constraint.
Revised: 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 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 are expressed in Equation (4).
Revised: = b i be the plane that represents one of the linear constraints in Equation (6), which are expressed in Equation (4).

Original: Let λ i be the intersection of this plane with the ray, i.e. it is the solution to a T
Revised: Let  i be the intersection of this plane with the ray, i.e. it is the solution to 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 linear constraints:
Revised: b i , then the solution is simply the smallest positive  i , which satisﬁes all the linear constraints:

Original: [Equation]
Revised: [Equation]

Original: where λ i =
Revised: where  i =

Original: b i − a T i x k a T
Revised: b i  a T i x k a T

Original: i δ k 
Revised: i  k (

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: There are a total of K + 2 D − 1 linear constraints, including the Voronoi cell boundaries and box constraints, 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 box constraints, which can be computed fully in parallel.

Original: This also allows end-to-end differentiation of the 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.
Revised: This also allows end-to-end differentiation of the 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 and f − 1 k are continuous.
Revised: Homeomorphism with continuous density We can show that f k is a bijection, and both f k and f  1 k are continuous.

Original: This allows us to use f k within the normalizing flows framework, as a mapping between 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 mapping between 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, computing the log determinant Jacobian for likelihood computation can be costly and will scale poorly with D if computed naïvely.
Revised: Efﬁcient computation of the log det Jacobian As f k is a mapping in D dimensions, computing the log determinant Jacobian for likelihood computation can be costly and will scale poorly with D if computed 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, because f k depends only on the direction δ k
Revised: Intuitively, because f k depends only on the direction  k

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

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

Original: This allows us to use the matrix determinant lemma to reformulate the log determinant Jacobian in a compute- and memory-efficient form.
Revised: This allows us to use the matrix determinant lemma to reformulate the log determinant 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 Sectionfor some given input x .
Revised: Proposition 3 Let the transformation f k ( x ) and all intermediate quantities be as deﬁned in Sectionfor 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 ∈
Revised: 10) for some c 2 R , u i 2 R D , v

Original: R D , v i ∈ R D , and its log determinant has the form 
Revised: i 2 R D , and its log determinant has the form log  [Equation]  [Equation]

Revised:

Original: [Equation]
Revised: [Equation]

Original: Ordinal
Revised: Ordinal

Original: Simplex
Revised: Simplex

Original: BinaryArgmax
Revised: BinaryArgmax

Original: Voronoi where w ij = c − 1 v T
Revised: Voronoi where w ij = c  1 v T

Original: i u j .
Revised: i u j .

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

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

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

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

Original: Compared to explicitly constructing the Jacobian matrix, this is more efficient in both compute and memory, and can readily scale to large values of D .
Revised: Compared to explicitly constructing the Jacobian matrix, 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 continuous data.
Revised: We discuss two applications of our method to likelihood-based modeling of discrete and continuous data.

Original: The first is a new approach to dequantization which allows training a model of discrete data using 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 data using 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 restricted to fixed dimensions and can benefit from learning similarities between discrete values.
Revised: Compared to existing dequantization methods [7, 18, 38], the Voronoi dequantization is not restricted to ﬁ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.
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 component lies strictly within each subset of a Voronoi tessellation.
Revised: The second is a formulation of disjoint mixture models, where each component lies strictly within each subset of a Voronoi tessellation.

Original: 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.
Revised: To the best of our knowledge, disjoint mixture models have not been explored signiﬁcantly in the literature and have not been successfully applied in more than a couple dimensions.

Original: Compared to an existing method [8], the Voronoi mixture model is not restricted to acting individually for each dimension.
Revised: Compared to an existing method [8], the Voronoi mixture model is not restricted to 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 can use tessellation to define subsets for dequantizing y as long as there are at least as many Voronoi cells—equivalently, anchors points—as the number of discrete values, i.e. K ≥ |Y| .
Revised: Then we can use tessellation to deﬁne subsets for dequantizing y as long as there are at least as many Voronoi cells—equivalently, anchors points—as the number of discrete values, i.e. K  |Y| .

Original: By assigning each discrete value to a Voronoi cell, we can then define a dequantization model q ( x | y ) by first sampling from a base distribution z ∼ q ( z | y ) in R D and then applying the mapping x = f k ( z ) from
Revised: By assigning each discrete value to a Voronoi cell, we can then deﬁne a dequantization model q ( x | y ) by ﬁrst 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 ) efficiently and train the dequantization alongside the density model p ( x ) on R D .
Revised: We can then obtain probabilities q ( x | y ) efﬁciently and train the dequantization alongside the density model p ( x ) on R D .

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 write y =
Revised: We can write y =

Original: g
Revised: g

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

Original: From Equation (4), it is easy to see that g ( x ) is the nearest neighbor operation, g ( x ) =
Revised: From Equation (4), it is easy 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 anchor points 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 anchor points and Voronoi cells more room to change over the course of training.

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

Original: We illustrate this in Figure 3.
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 mixture model can be constructed by defining distributions on each Voronoi cell.
Revised: Then a disjoint mixture model can be constructed by deﬁning distributions on each Voronoi cell.

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

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

Original: Let x be a point in R D , our Voronoi mixture model assigns the density log 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 density [Equation]

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

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

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 can easily be composed with normalizing flow layers where, in addition to the change of variable due to f − 1 k , we also apply the change in density resulting from choosing one out of K components.
Revised: This can easily be composed with normalizing ﬂow layers where, in addition to the change of variable due to f  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 Invertible mappings have been proposed for discrete data, where discrete values are effectively rearranged from a factorized distribution.
Revised: Normalizing ﬂows for discrete data 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: [47] use quantized ordinal transformations, while Tran et al.
Revised: [47] use quantized ordinal transformations, while Tran et al.

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

Original: 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 is sufficient to transform any joint discrete distribution into a factorized distribution.
Revised: 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 is sufﬁcient to transform any joint discrete distribution into a factorized distribution.

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

Original: 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].
Revised: Dequantization methods Within the normalizing ﬂows literature, the act of adding noise was originally used for ordinal data as a way to combat numerical issues [7, 46].

Original: Later on, appropriate dequantization approaches have been shown to lower bound the log-likelihood of a discrete model [16, 44].
Revised: Later on, appropriate dequantization 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 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 a Gumbel base distribution.
Revised: 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 a Gumbel base distribution.

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

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

Original: Intuitively, after fixing one of the logits, the softmax operation is an invertible transformation on the bounded domain ∆ K − 1 = { x ∈ R K : (cid:80)
Revised: Intuitively, after ﬁxing one of the logits, the softmax operation is an invertible transformation on the bounded 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 be broken into K subsets, each corresponding to a particular discrete value.
Revised: The ( K  1) -simplex can then be broken 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: which effectively increases the number of dimensions by one compared to the simplex approach.

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

Original: In order to make their approach work when K is large, [18] propose binarizing the data as a preprocessing step.
Revised: In order to make their approach work when K is large, [18] propose binarizing the data as a preprocessing step.

Original: In contrast, Voronoi dequantization 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 the constraint that subsets are disjoint, and instead work with general mixture models with unbounded support, relying on the conditional model p ( y | x ) being sufficiently weak so that the task of modeling is forced onto a flow-based prior.
Revised: 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 | x ) being sufﬁciently weak so that the task of modeling is forced onto a ﬂow-based prior.

Original: They have achieved good performance on a number of tasks, similar to general approaches that combine normalizing flows with variational inference [19, 33, 50].
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 ) gained from using disjoint subsets.
Revised: However, they lose the computational savings and the deterministic decoder p ( y | x ) gained from using disjoint subsets.

Original: 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.
Revised: 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.

Original: Disjoint mixture models The computational savings from using disjoint subsets was pointed out by Dinh et al.
Revised: The computational savings from using disjoint subsets was pointed out by 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 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.
Revised: 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.

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

Original: [Figure]
Revised: [Figure]

Original: relations values dering given

Original: [Figure]
Revised: [Figure]

Original: for more discussion).
Revised: for more discussion).

Original: The use of monotonic splines have previously been combined with normalizing flows [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, 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 in high dimensions.
Revised: 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 in high dimensions.

Original:  Experiments
Revised:  Experiments

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

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

Original: For this reason, we just use simple coupling blocks with affine transformations [6, 7] as a base flow model in our experiments, unless stated otherwise.
Revised: For this reason, we just use simple coupling blocks with afﬁne transformations [6, 7] as a base ﬂow model in our experiments, unless stated 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: 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.
Revised: 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.

Revised: Open source code is available at https://github.com/facebookresearch/semi-discrete-flow .

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 compare against Discrete Flows [45].
Revised: We start by considering synthetic distributions over two discrete variables, to qualitatively compare against Discrete Flows [45].

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

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

Original: We see that Discrete Flows has trouble fitting 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 learn a flow model on the combined R 4 , parameterized by multilayer perceptrons (MLPs).
Revised: We then learn a ﬂow model on the combined R 4 , parameterized by multilayer perceptrons (MLPs).

Original: In Figurewe visualize the learned Voronoi tessellation and samples from our model.
Revised: In Figurewe visualize the learned Voronoi tessellation and samples from our model.

Original: The learned tessellation seems to group some of discrete values that occur frequently together, to reduce the number of modes.
Revised: The learned tessellation seems 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 number of classes.
Revised: We experiment with complex data sets where each discrete variable can have a varying number of classes.

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

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

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

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

Original: Most of these data sets have a combination of both ordinal and non-ordinal variables, and we 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, and we 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 short on more challenging data sets such as the USCensus90 , the largest data set we consider with 2.
Revised: [Table]

Original: [Table] [Table][Table]
Revised: [Table] [Table]

Revised: 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.

Original: million 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 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.
Revised: We do see a 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.

Original: 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 USCensus data 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 same name, which exhibit complex spatial dependencies between the board pieces, and the USCensus data 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 invariance to specific group symmetries into the density model.
Revised: An appeal of using normalizing ﬂows in continuous space is the ability to incorporate invariance to speciﬁc group symmetries into the density model.

Original: For instance, this can be done by ensuring the ordinary differential equation is equivariant [2, 25, 42] in a continuous normalizing flow [4, 13] framework.
Revised: For instance, this can be done by ensuring the 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 a bipartition of discrete variables.
Revised: This invariance cannot be explicitly modeled by Discrete Flows as they require an ordering or a bipartition of discrete variables.

Original: 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.
Revised: We preprocessed a data set of retail market basket data from an anonymous Belgian retail store [3] and a data set of anonymized trafﬁc accident data [11], which contain sets of discrete variables each with 765 and 269 values, respectively.

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

Original: Meanwhile, our Voronoi dequantization has no problems staying competitive as its embedding space is freely decoupled from the number of discrete values.
Revised: Meanwhile, our Voronoi dequantization has no problems staying 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 used the open source code provided by Hoogeboom et al.
Revised: Here we used the open source code provided by Hoogeboom et al.

Original: [18] with the exact same autoregressive flow model and optimizer setups.
Revised: [18] with the exact same autoregressive ﬂow model and optimizer setups.

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

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 more constrained, 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 more constrained, 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 6 shows results of training a simple normalizing flow on 2D data, which has trouble completely separating all the modes.
Revised: Figure 6 shows results of training a simple normalizing ﬂow on 2D data, which has trouble completely separating all the modes.

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

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

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

Original: We add a disjoint mixture component after a few layers of normalizing flows, with each component modeled by a normalizing Data Baseline Voronoi
Revised: We add a disjoint mixture component after a few layers of normalizing ﬂows

Original: [Table][Table]
Revised: [Table]

Original: flow conditioned on a embedding of the component index.
Revised: with each component modeled by a normalizing ﬂow conditioned on a embedding of the component index.

Original: The number of mixture components ranges from 16 to
Revised: The number of mixture components ranges from 16 to

Original: [Table]
Revised: [Table]

Original: 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].
Revised: 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 the disjoint mixture, as well as a strong but computationally costly density model FFJORD [13].

Original: From
Revised: From

Original: Table
Revised: Table

Original: [Table]
Revised: [Table]

Original: 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 flow model [Table]
Revised: 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.

Original:  Conclusion and Discussion
Revised:  Conclusion and Discussion

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

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

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

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

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

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

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

Original: This is a different and undesirable property as opposed to the partitioning method used by Dinh et al.
Revised: This is a different and undesirable property as opposed to the partitioning method used by Dinh et al.

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

Original: Balancing expressiveness and compute cost is a delicate problem that sits at the core of probabilistic machine learning.
Revised: Balancing expressiveness and compute cost is a delicate problem that sits at the core of probabilistic machine learning.

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

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

Original: As depicted in
Revised: As depicted in

Original: Figure 1, the transformed density has a non-smooth probability density function.
Revised: Figure 1, the transformed density has a non-smooth probability density function.

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

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

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

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

Original: Additionally, more flexible tessellations—such as the Laguerre tessellation—and additional concepts from semi-discrete optimal transport [14, 27, 37] may be adapted to improve semi-discrete normalizing flows. A Proofs Proposition 1 The mapping f k : R D → V k as defined in the 2-step procedure is a homeomorphism.
Revised: Additionally, more ﬂexible tessellations—such as the Laguerre tessellation—and additional concepts from semi-discrete optimal transport [14, 27, 37] may be adapted to improve semi-discrete normalizing ﬂows.

Original: Proof.

Original: Let x ∈ R D and x k be a given anchor point corresponding to a Voronoi cell V k .

Original: If x ̸ = x k , then x is uniquely represented by the tuple (∆ , δ ) where ∆ = ∥ x − x k ∥ and δ = x − x k ∥ x − x k ∥ , since x = x k + ∆ δ .

Original: Since δ uniques defines the ray { x k + λδ ; λ > 0 } and x ( λ ∗ ) , then because α k in

Original: Equation (8) is a bijection in ∆ , f k is a bijection.

Original: For x ̸ = x k , the continuity of f k follows from Rudin et al.

Original: [41, Theorem 4.7] since ∆ and x ( λ ∗ ) are continuous in x and α k is continuous in ∆ .

Original: Then since f k ( x ) → x k as x → x k from all directions, this justifies the choice of setting f k ( x ) =

Original: x k when x = x k .

Original: Finally, by the invariance of domain theorem, since V k is an open set in R D , f k is an open map and the inverse f − 1 k is continuous, and we can conclude f is a homeomorphism between

Original: R D and V k .

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

Original: Proof.

Original: See proof of Proposition 3 below for the form of the Jacobian.

Original: For the case where x ̸ = x k , all quantities in Equations (15) to (22) are continuous with respect to x .

Original: Hence the Jacobian ∂f k ( x ) ∂ x is continuous, and since it is always full-rank, then the composition (cid:12)(cid:12)(cid:12) det

Original: ∂f k ( x ) ∂ x (cid:12)(cid:12)(cid:12) is continuous [41,

Original: Theorem 4.7] and so is the product p x ( x ) (cid:12)(cid:12)(cid:12) det ∂f k ( x ) ∂ x (cid:12)(cid:12)(cid:12) [41, Theorem 4.9].

Original: Proposition 3 Let the transformation f k ( x ) and all intermediate quantities be as defined in Sectionfor some given input x .

Original: Then the Jacobian factorizes as

Original: [Equation]

Original: for some c ∈ R , u i ∈ R D , v i ∈ R D , and its log determinant has the form

Original: [Equation]

Original: where

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: (21) w

Original: [Equation]

Original: for i, j ∈ {

Original: [Equation]

Original: [Equation]

Original: Proof.

Original: We first write the Jacobian of f k in the form of cI + u 1 v T 1 + u 2 v T 2 where c is a scalar, and u 1 , u 2 , v 1 , v 2 are vectors of size D .

Original: Define the shorthand ∆ ≜

Original: ∥ x − x k ∥ , ∆ ∗ ≜

Original: ∥ x ( λ ∗ ) − x k ∥ , and ˜∆ ≜ ∆ / ∆ ∗ .

Original: To simplify notation, we use the short-hands α k

Original: = α k ( ˜∆) , δ k

Original: = δ k ( x ) .

Original: Then

Original: the

Original: Jacobian follows

Original: [Figure]

Original: [Table]

Original: [Equation]

Original: Now that ∂f k ( x ) ∂x is in the form of cI + u 1 v T 1 + u 2 v T 2 , where

Original: [Equation]

Original: Applying the matrix determinant lemma twice

Original: [Equation]

Original: we can show that

Original: [Equation]

Original: We simplify this by defining the scaled dot products,

Original: [Equation]

Original: Then

Original: [Equation]

Original: Intermediate steps above used the following gradient identities.

Original: [Equation]

Original: B The inverse mapping

Original: We use the inverse mapping f − 1 k : V k

Original: → R D for training disjoint mixture models, so we next describe how to compute this.

Original: Let z

Original: =

Original: f k ( x ) .

Original: Conveniently, since both x and z lie on the ray { x ( λ ) : λ > 0 } , we know δ k ( x ) =

Original: δ k ( z ) .

Original: So the first step is the same as the forward procedure: we solve for λ ∗ and x ( λ ∗ ) .

Original: Following this, we then recover x by inverting Step 2 of the forward procedure.

Original: This inverse transformation is given by

Original: [Equation]

Original: Equation (50) is an element-wise division.

Original: Since ˜ α will the same in all dimensions, we can simply pick a dimension in Equation (51).

Original: In our experiments, the inverse α − 1 k can be computed analytically, though since it is just a scalar function, simple methods like bisection can also work when the inverse is not known analytically.

Original: Lastly, Equation (53) follows from the observation that δ k ( x ) =

Original: δ k ( z ) .

Original: Log determinant of the inverse.

Original: We can also use Proposition 3 to compute the log determinant of the inverse transform without needing to recompute f k ( x ) .

Original: The only difference is a sign: log

Original: [Equation]

Original: det ∂f k ( x ) ∂x (cid:12)(cid:12)(cid:12) .

Original: The required quantities, ∆ , x ( λ ∗ ) , and δ k ( x ) , are readily available after computing x = f − 1 k

Original: ( z ) .

Original: The gradients with respect to quantities of x can be expressed using gradients with respect to quantities of z ,

Original: [Equation]

Original: and

Original: [Equation]

Original: which are accessible through automatic differentiation.

Original: C Data sets C.1 UCI Data sets

Original: The main preprocessing we did was to (i) remove the “label” attribute from each data set, and (ii) remove attributes that only ever take on one value.

Original: Apart from this, the USCensus90 data set contains a unique identifier for each row, which was removed.

Original: Descriptions for all data set are below.

Original: Connect4 [webpage] This data set contains all legal 8-ply positions in the game of Connect Four in which neither player has won yet, and in which the next move is not forced.

Original: The original task was to predict which player would win, which has been removed during preprocessing.

Original: There are a totaldiscrete variables (one for each location on the board), each with 3 possible discrete values (taken by player 1, taken by player 2, blank).

Original: This data set was randomly split into 54045 training examples, 6755 validation examples, and 6757 test examples.

Original: Forests [webpage] This data set contains cartographic variables regarding forests including four wilderness areas located in the Roosevelt National Forest of northern Colorado.

Original: These areas represent forests with minimal human-caused disturbances.

Original: The original task was to predict the forest cover type, which has been removed during preprocessing.

Original: There are a total of 54 discrete variables, with 10 being the highest number of discrete values.

Original: This data set was randomly split into 464809 training examples, 58101 validation examples, and 58102 test examples.

Original: Mushroom [webpage] This data set includes descriptions of hypothetical samples corresponding to 23 species of gilled mushrooms in the Agaricus and Lepiota Family.

Original: The original task was to predict whether each species is edible, which has been removed during preprocessing.

Original: There are a total ofdiscrete variables, with 12 being the highest number of discrete values.

Original: This data set was randomly split into 6499 training examples, 812 validation examples, and 813 test examples.

Original: Nursery [webpage] This data set contains attributes of applicants to nursery schools, during a period when there was excessive enrollment to these schools in Ljubljana, Slovenia, and the rejected applications frequently needed an objective explanation.

Original: All data have been completely anonymized.

Original: The original task was to predict whether an applicant would be recommended for acceptance by hierarchical decision model, which has been removed during preprocessing.

Original: There are a total ofdiscrete variables, with 5 being the highest number of discrete values.

Original: This data set was randomly split into 10367 training examples, 1296 validation examples, and 1297 test examples.

Original: PokerHands [webpage] This data set contains poker hands consisting of five playing cards drawn from a standard deck of 52.

Original: Each card is described using two attributes (suit and rank), for a total of 10 predictive attributes.

Original: There is one Class attribute that describes the “poker hand”.

Original: The original task was to predict the poker hand class (pairs, full house, royal flush, etc.), which has been removed during preprocessing.

Original: There are a total of 10 discrete variables, with 13 being the highest number of discrete values.

Original: This data set was randomly split into 820008 training examples, 102501 validation examples, and 102501 test examples.

Original: USCensus90 [webpage] This data set contains a portion of the data collected as part of thecensus in the United States, with the data completely anonymized.

Original: There are a total of 68 discrete variables, with 18 being the highest number of discrete values.

Original: This data set was randomly split into 2212456 training examples, 122914 validation examples, and 122915 test examples.

Original: C.2 Itemset Data sets

Original: These data sets were taken from the Frequent Itemset Mining Data set Repository [webpage].

Original: Each row is interpreted as a set of items with no emphasis on the ordering of items.

Original: Retail [3] This data set contains anonymized retail market basket data from an anonymous Belgian retail store.

Original: We first removed rows with less than 4 items, then randomly sampled a subset ofitems for every row.

Original: Items that appear in less 300 rows were dropped from the data set.

Original: The final preprocessed data set contains 765 distinct items.

Original: This data set was randomly split into 24280 training examples, 3035 validation examples, and 3036 test examples.

Original: Accidents [11] This data set contains contains anonymized traffic accident data.

Original: Data on traffic accidents are obtained from the National Institute of Statistics (NIS) for the region of Flanders (Belgium) for the period 1991-2000.

Original: We first removed rows with less than 4 items, then randomly sampled a subset of 4 items for every row.

Original: This subsampling occurred 10 times if a row has 16 or more items, 5 times if the row has 8 to 15 items, and once if the row has 4-7 items.

Original: Items that appear in less 300 rows were dropped from the data set.

Original: The final preprocessed data set contains 213 distinct items.

Original: This data set was randomly split into 270129 training examples, 33766 validation examples, and 33767 test examples.

Original: D Experimental Details

Original: All experiments were run on a single NVIDIA V100 GPU.

Original: Detailed hyperparameter sweeps are below.

Original: We used the validation set to choose hyperparameters as well as to perform early stopping.

Original: 2D synthetic data sets Continuous data sets were quantized into 91 bins for each coordinate.

Original: For

Original: Voronoi dequantization, we dequantized each coordinate into an embedding space of 2 dimensions, with 91 Voronoi cells.

Original: The dequantization model is parameterized by 4 layers of coupling blocks, each with a 2 hidden layer MLP with 256 hidden units each, where the Swish activation function was used [39].

Original: The flow model is similarly parameterized but with 16 layers of coupling blocks.

Original: Each block alternated between 4 different partitioning schemes: maksing out the first half, masking out the second half, masking out the odd indices, and masking out the even indices.

Original: We trained with the

Original: Adam optimizer [22] with a learning rate of 1e-3 .

Original: Discrete-valued UCI data sets For Voronoi dequantization

Original: , we dequantized each coordinate into an embedding space of 4 or 6 dimensions, with the number of Voronoi cells set to the highest number of discrete values over all discrete variables.

Original: The dequantization model is parameterized by 4 layers of coupling blocks, each with a 2 hidden layer MLP with 256, 512, or 1024 hidden units each, where the Swish activation function was used [39].

Original: The flow model is similarly parameterized but with 16 or 32 layers of coupling blocks.

Original: Each block alternated between 4 different partitioning schemes: maksing out the first half, masking out the second half, masking out the odd indices, and masking out the even indices.

Original: We trained with the Adam optimizer [22] with a learning rate sweep over { 1e-3 , 5e-4 , 1e-4 }.

Original: Itemset data sets For Voronoi dequantization, we dequantized each coordinate into an embedding space of 6 dimensions, with the number of Voronoi cells set to the number of items in the data set.

Original: We used a continuous normalizing flow (CNF) with the ordinary differential equation (ODE) defined using a Transformer archiecture and a L 2 -distance based multihead attention layer [21] and the GeLU activation function [15].

Original: No positional embeddings were provided to the model to ensure the model is equivariant to permutations.

Original: We composed 12 CNF layers, each defined using a Transformer model that has 2 or 3 layers of alternating multihead attention and fully connected residual connections.

Original: To solve the ODE and train our model, we used the dopri5 solver from the torchdiffeq library [5] with atol=rtol=1e-5 .

Original: We trained with the AdamW optimizer [30, 49] with a learning rate of 1e-3 and weight decay of 1e-6 .

Original: For the Voronoi dequantization, we set D =6 for both data sets, though it may be possible to improve performance by tuning D .

Original: Character-level language modeling We used the provided hyperparameters from the open source repository [URL].

Original: Some parts of code had to be adapted for our usage, but model architecture and optimizer remained largely the same.

Original: Continuous-valued UCI data sets Disjoint mixture modeling is the inverse of a dequantization method.

Original: We use a flow model which maps the data nonlinearly, z = f ( x ) .

Original: Then a Voronoi mixture model partitions the space and assigns each z an index through the set identification function k

Original: = g ( z ) .

Original: The probabilities p ( k ) are parameterized through a softmax and learned.

Original: We then apply the inverse transformation from Section 3 to map from V k to R D .

Original: A conditional normalizing flow is then used to model each mixture model p ( z | k ) .

Original: We set the base distribution to be a Gaussian with standard deviation 0 .

Original: 2 , as this helps concentrate the density around the anchor points at initialization.

Original: We tuned the number of flow layers before and after the mixture layer, with the total number of layers in { 16 , 32 , 48 , 64 } .

Original: We also tuned the number of mixtures in { 8 , 16 , 32 , 64 , 128 } .

Original: Each coupling block uses a neural network with 3 hidden layers of dimension 64 with either the GeLU or Swish [39] activation function.

Original: We trained with a batch size of 1048 and the Adam optimizer [22] with a learning rate sweep over { 1e-3 , 5e-3 }.