Original: A RAD APPROACH TO DEEP MIXTURE MODELS
Revised: A RAD APPROACH TO DEEP MIXTURE MODELS

Original: Laurent Dinh
Revised: Laurent Dinh

Original: Google Brain laurentdinh@google.com
Revised: Google Brain laurentdinh@google.com

Original: Jascha Sohl-Dickstein
Revised: Jascha Sohl-Dickstein

Original: Google Brain jaschasd@google.com
Revised: Google Brain jaschasd@google.com

Original: Razvan Pascanu
Revised: Razvan Pascanu

Original: DeepMind razp@google.com
Revised: DeepMind razp@google.com

Original: Hugo Larochelle
Revised: Hugo Larochelle

Original: Google Brain hugolarochelle@google.com
Revised: Google Brain hugolarochelle@google.com

Original: A BSTRACT
Revised: A BSTRACT

Original: Flow based models such as R EAL NVP are an extremely powerful approach to density estimation.
Revised: Flow based models such as R EAL NVP are an extremely powerful approach to density estimation.

Original: However, existing ﬂow based models are restricted to transforming continuous densities over a continuous input space into similarly continuous distributions over continuous latent variables.
Revised: However, existing ﬂow based models are restricted to transforming continuous densities over a continuous input space into similarly continuous distributions over continuous latent variables.

Original: This makes them poorly suited for modeling and representing discrete structures in data distributions, for example class membership or discrete symmetries.
Revised: This makes them poorly suited for modeling and representing discrete structures in data distributions, for example class membership or discrete symmetries.

Original: To address this difﬁculty, we present a normalizing ﬂow architecture which relies on domain partitioning using locally invertible functions, and possesses both real and discrete valued latent variables.
Revised: To address this difﬁculty, we present a normalizing ﬂow architecture which relies on domain partitioning using locally invertible functions, and possesses both real and discrete valued latent variables.

Original: This Real and Discrete (R AD ) approach retains the desirable normalizing ﬂow properties of exact sampling, exact inference, and analytically computable probabilities, while at the same time allowing simultaneous modeling of both continuous and discrete structure in a data distribution.
Revised: This Real and Discrete (R AD ) approach retains the desirable normalizing ﬂow properties of exact sampling, exact inference, and analytically computable probabilities, while at the same time allowing simultaneous modeling of both continuous and discrete structure in a data distribution.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Latent generative models are one of the prevailing approaches for building expressive and tractable generative models.
Revised: Latent generative models are one of the prevailing approaches for building expressive and tractable generative models.

Original: The generative process for a sample x can be expressed as
Revised: The generative process for a sample x can be expressed as

Original: [Equation]
Revised: [Equation]

Original: where z is a noise vector, and g a parametric generator network (typically a deep neural network).
Revised: where z is a noise vector, and g a parametric generator network (typically a deep neural network).

Original: This paradigm has several incarnations, including variational autoencoders (Kingma & Welling, 2014; Rezende et al., 2014), generative adversarial networks (Goodfellow et al., 2014), and ﬂow based models (Baird et al., 2005; Tabak & Turner, 2013; Dinh et al., 2015; 2017; Kingma & Dhariwal, 2018; Chen et al., 2018; Grathwohl et al., 2019).
Revised: This paradigm has several incarnations, including variational autoencoders (Kingma & Welling, 2014; Rezende et al., 2014), generative adversarial networks (Goodfellow et al., 2014), and ﬂow based models (Baird et al., 2005; Tabak & Turner, 2013; Dinh et al., 2015; 2017; Kingma & Dhariwal, 2018; Chen et al., 2018; Grathwohl et al., 2019).

Original: The training process and model architecture for many existing latent generative models, and for all published ﬂow based models, assumes a unimodal smooth distribution over latent variables z .
Revised: The training process and model architecture for many existing latent generative models, and for all published ﬂow based models, assumes a unimodal smooth distribution over latent variables z .

Original: Given the parametrization of g as a neural network, the mapping to x is a continuous function.
Revised: Given the parametrization of g as a neural network, the mapping to x is a continuous function.

Original: This imposed structure makes it challenging to model data distributions with discrete structure – for instance, multi-modal distributions, distributions with holes, distributions with discrete symmetries, or distributions that lie on a union of manifolds (as may approximately be true for natural images, see Tenenbaum et al., 2000).
Revised: This imposed structure makes it challenging to model data distributions with discrete structure – for instance, multi-modal distributions, distributions with holes, distributions with discrete symmetries, or distributions that lie on a union of manifolds (as may approximately be true for natural images, see Tenenbaum et al., 2000).

Original: Indeed, such cases require the model to learn a generator whose input Jacobian has highly varying or inﬁnite magnitude to separate the initial noise source into different clusters.
Revised: Indeed, such cases require the model to learn a generator whose input Jacobian has highly varying or inﬁnite magnitude to separate the initial noise source into different clusters.

Original: Such variations imply a challenging optimization problem due to large changes in curvature.
Revised: Such variations imply a challenging optimization problem due to large changes in curvature.

Original: This shortcoming can be critical as several problems of interest are hypothesized to follow a clustering structure, i.e. the distributions is concentrated along several disjoint connected sets (Eghbal-zadeh et al., 2018).
Revised: This shortcoming can be critical as several problems of interest are hypothesized to follow a clustering structure, i.e. the distributions is concentrated along several disjoint connected sets (Eghbal-zadeh et al., 2018).

Original: A standard way to address this issue has been to use mixture models (Yeung et al., 2017; Richardson & Weiss, 2018; Eghbal-zadeh et al., 2018) or structured priors (Johnson et al., 2016).
Revised: A standard way to address this issue has been to use mixture models (Yeung et al., 2017; Richardson & Weiss, 2018; Eghbal-zadeh et al., 2018) or structured priors (Johnson et al., 2016).

Original: In order to efﬁciently parametrize the model, mixture models are often formulated as a discrete latent variable models (Hinton & Salakhutdinov, 2006; Courville et al., 2011; Mnih & Gregor, 2014; van den Oord
Revised: In order to efﬁciently parametrize the model, mixture models are often formulated as a discrete latent variable models (Hinton & Salakhutdinov, 2006; Courville et al., 2011; Mnih & Gregor, 2014; van den Oord

Original: [Figure]
Revised: [Figure]

Original: et al., 2017), some of which can be expressed as a deep mixture model (Tang et al., 2012; Van den
Revised: et al., 2017), some of which can be expressed as a deep mixture model (Tang et al., 2012; Van den

Original: Oord & Schrauwen, 2014; van den Oord & Dambre, 2015).
Revised: Oord & Schrauwen, 2014; van den Oord & Dambre, 2015).

Original: Although the resulting exponential number of mixture components with depth in deep mixture models is an advantage in terms of expressivity, it is an impediment to inference, evaluation, and training of such models, often requiring as a result the use of approximate methods like hard- E M or variational inference (Neal & Hinton, 1998).
Revised: Although the resulting exponential number of mixture components with depth in deep mixture models is an advantage in terms of expressivity, it is an impediment to inference, evaluation, and training of such models, often requiring as a result the use of approximate methods like hard- E M or variational inference (Neal & Hinton, 1998).

Original: In this paper we combine piecewise invertible functions, with discrete auxiliary variables selecting which invertible function applies, to describe a deep mixture model.
Revised: In this paper we combine piecewise invertible functions with discrete auxiliary variables, selecting which invertible function applies, to describe a deep mixture model.

Original: This framework enables a probabilistic model’s latent space to have both real and discrete valued units, and to capture both continuous and discrete structure in the data distribution.
Revised: This framework enables a probabilistic model’s latent space to have both real and discrete valued units, and to capture both continuous and discrete structure in the data distribution.

Original: It achieves this added capability while preserving the exact inference, exact sampling, exact evaluation of log-likelihood, and efﬁcient training that make standard ﬂow based models desirable.
Revised: It achieves this added capability while preserving the exact inference, exact sampling, exact evaluation of log-likelihood, and efﬁcient training that make standard ﬂow based models desirable.

Original:  M ODEL DEFINITION
Revised:  M ODEL DEFINITION

Original: We aim to learn a parametrized distribution p X ( x ) on the continuous input domain R d by maximizing log-likelihood.
Revised: We aim to learn a parametrized distribution p X ( x ) on the continuous input domain R d by maximizing log-likelihood.

Original: The major obstacle to training an expressive probabilistic model is typically efﬁciently evaluating log-likelihood
Revised: The major obstacle to training an expressive probabilistic model is typically efﬁciently evaluating log-likelihood.

Original: . 2. P ARTITIONING
Revised: 2. P ARTITIONING

Original: If we consider a mixture model with a large number | K | of components, the likelihood takes the form
Revised: If we consider a mixture model with a large number | K | of components, the likelihood takes the form

Original: [Equation]
Revised: [Equation]

Revised: p X | K ( x | k ) .

Original: In general, evaluating the likelihood requires computing probabilities for all | K | components.
Revised: In general, evaluating the likelihood requires computing probabilities for all | K | components.

Original: However, following a strategy similar to Rainforth et al.
Revised: However, following a strategy similar to Rainforth et al.

Original: (2018), if we partition the domain R d into disjoint subsets
Revised: (2018), if we partition the domain R d into disjoint subsets

Original: A k for 1 ≤ k ≤ | K | such that ∀ i (cid:54) =
Revised: A k for 1 ≤ k ≤ | K | such that ∀ i (cid:54) =

Original: j A i ∩ A j = ∅ and | K | (cid:83) k
Revised: j A i ∩ A j = ∅ and | K | (cid:83) k

Original: =
Revised: =

Original: A k = R d , constrain the support of p X | K
Revised: A k = R d , constrain the support of p X | K

Original: ( x | k ) to A k (i.e. ∀ x / ∈
Revised: ( x | k ) to A k (i.e. ∀ x / ∈

Original: A k , p X | K ( x | k ) = 0 ), and deﬁne a set identiﬁcation function f K ( x ) such that ∀ x ∈ R d , x
Revised: A k , p X | K ( x | k ) = 0 ), and deﬁne a set identiﬁcation function f K ( x ) such that ∀ x ∈ R d , x

Original: ∈ A f K ( x ) , we can write the likelihood as
Revised: ∈ A f K ( x ) , we can write the likelihood as

Original: [Equation]
Revised: [Equation]

Original: This transforms the problem of summation to a search problem x
Revised: This transforms the problem of summation to a search problem x

Original: (cid:55)→ f K ( x ) .
Revised: (cid:55)→ f K ( x ) .

Original: This can be seen as the inferential converse of a stratiﬁed sampling strategy (Rubinstein & Kroese, 2016).
Revised: This can be seen as the inferential converse of a stratiﬁed sampling strategy (Rubinstein & Kroese, 2016).

Original: 2. C HANGE OF VARIABLE FORMULA
Revised: 2. C HANGE OF VARIABLE FORMULA

Original: The proposed approach will be a direct extension of ﬂow based models (Rippel & Adams, 2013;
Revised: The proposed approach will be a direct extension of ﬂow based models (Rippel & Adams, 2013;

Original: Dinh et al., 2015; 2017; Kingma & Dhariwal, 2018).
Revised: Dinh et al., 2015; 2017; Kingma & Dhariwal, 2018).

Original: Flow based models enable log-likelihood  evaluation by relying on the change of variable formula
Revised: Flow based models enable log-likelihood [Figure] [Figure] evaluation by relying on the change of variable formula

Original: [Equation]
Revised: [Equation]

Original: with f Z a parametrized bijective function from R d onto R d and (cid:12)(cid:12)(cid:12)
Revised: with f Z a parametrized bijective function from R d onto R d and (cid:12)(cid:12)(cid:12)

Original: ∂f Z ∂ x T (cid:12)(cid:12)(cid:12) the absolute value of the determinant of its Jacobian.
Revised: ∂f Z ∂ x T (cid:12)(cid:12)(cid:12) the absolute value of the determinant of its Jacobian.

Original: As also proposed in Falorsi et al.
Revised: As also proposed in Falorsi et al.

Original: (2019), we relax the constraint that f Z be bijective, and instead have it be surjective onto R d and piecewise invertible.
Revised: (2019), we relax the constraint that f Z be bijective, and instead have it be surjective onto R d and piecewise invertible.

Original: That is, we require f Z |
Revised: That is, we require f Z |

Original: A k ( x ) be an invertible function, where f Z | A k ( x ) indicates f Z ( x ) restricted to the domain A k .
Revised: A k ( x ) be an invertible function, where f Z | A k ( x ) indicates f Z ( x ) restricted to the domain A k .

Original: Given a distribution p Z,K ( z , k ) =
Revised: Given a distribution p Z,K ( z , k ) =

Original: p K | Z ( k | z ) p Z ( z ) such that ∀
Revised: p K | Z ( k | z ) p Z ( z ) such that ∀

Original: ( z , k ) , z / ∈ f Z ( A k ) ⇒
Revised: ( z , k ) , z / ∈ f Z ( A k ) ⇒

Original: p Z,K = 0 , we can deﬁne the following generative process:
Revised: p Z,K = 0 , we can deﬁne the following generative process:

Original: [Equation]
Revised: [Equation]

Original: If we use the set identiﬁcation function f K associated with A k , the distribution corresponding to this stochastic inversion can be deﬁned by a change of variable formula
Revised: If we use the set identiﬁcation function f K associated with A k , the distribution corresponding to this stochastic inversion can be deﬁned by a change of variable formula

Original: [Equation]
Revised: [Equation]

Original: Because of the use of both Real and Discrete stochastic variables, we call this class of model R AD .
Revised: Because of the use of both Real and Discrete stochastic variables, we call this class of model R AD .

Original: The particular parametrization we use on is depicted in Figure 2.
Revised: The particular parametrization we use on is depicted in Figure 2.

Original: We rely on piecewise invertible functions that allow us to deﬁne a mixture model of repeated symmetrical patterns, following a
Revised: We rely on piecewise invertible functions that allow us to deﬁne a mixture model of repeated symmetrical patterns, following a

Original: Z (0)
Revised: Z (0)

Original: K (0)
Revised: K (0)

Original: Z (1)
Revised: Z (1)

Original: K (1)
Revised: K (1)

Original: Z (2)
Revised: Z (2)

Original: K (2)
Revised: K (2)

Original: X (a) Sampling.
Revised: X (a) Sampling.

Original: Z (0)
Revised: Z (0)

Original: K (0)
Revised: K (0)

Original: Z (1)
Revised: Z (1)

Original: K (1)
Revised: K (1)

Original: Z (2)
Revised: Z (2)

Original: K (2)
Revised: K (2)

Original: X (b) Inference.
Revised: X (b) Inference.

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

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

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

Original: method of folding the input space .
Revised: method of folding the input space .

Original: Note that in this instance the function f
Revised: Note that in this instance the function f

Original: K is implicitly deﬁned by f Z , as the discrete latent corresponds to which invertible component of the piecewise function x falls on.
Revised: K is implicitly deﬁned by f Z , as the discrete latent corresponds to which invertible component of the piecewise function x falls on.

Original: So far, we have deﬁned a mixture of | K | components with disjoint support.
Revised: So far, we have deﬁned a mixture of | K | components with disjoint support.

Original: However, if we factorize p Z,K as p Z · p K | Z , we can apply another piecewise invertible map to Z to deﬁne p Z as another mixture model.
Revised: However, if we factorize p Z,K as p Z · p K | Z , we can apply another piecewise invertible map to Z to deﬁne p Z as another mixture model.

Original: Recursively applying this method results in a deep mixture model (see Figure 3).
Revised: Recursively applying this method results in a deep mixture model (see Figure 3).

Original: Another advantage of such factorization is in the gating network p K | Z , as also designated in (van den
Revised: Another advantage of such factorization is in the gating network p K | Z , as also designated in (van den

Original: Oord & Dambre, 2015).
Revised: Oord & Dambre, 2015).

Original: It provides a more constrainted but less sample wasteful approach than rejection sampling (Bauer & Mnih, 2019) by taking into account the untransformed sample z before selecting the mixture component k .
Revised: It provides a more constrainted but less sample wasteful approach than rejection sampling (Bauer & Mnih, 2019) by taking into account the untransformed sample z before selecting the mixture component k .

Original: This allows the model to exploit the distribution p Z in different regions A k in more complex ways than repeating it as a patternm as illustrated in Figure 4.
Revised: This allows the model to exploit the distribution p Z in different regions A k in more complex ways than repeating it as a patternm as illustrated in Figure 4.

Original: The function from the input to the discrete variables, f K ( x ) , contains discontinuities.
Revised: The function from the input to the discrete variables, f K ( x ) , contains discontinuities.

Original: This presents the danger of introducing discontinuities into log p X ( x ) , making optimization more difﬁcult.
Revised: This presents the danger of introducing discontinuities into log p X ( x ) , making optimization more difﬁcult.

Original: However, by carefully imposing boundary conditions on the gating network, we are able to exactly counteract the effect of discontinuities in f K , and cause log p X ( x ) to remain continuous with respect to the parameters.
Revised: However, by carefully imposing boundary conditions on the gating network, we are able to exactly counteract the effect of discontinuities in f K , and cause log p X ( x ) to remain continuous with respect to the parameters.

Original: This is discussed in detail in Appendix A.
Revised: This is discussed in detail in Appendix A.

Original: [Figure]
Revised: [Figure]

Original: hypothesis

Original: [Figure]

Original: [Figure]
Revised: [Figure]

Original:  E XPERIMENTS 3. P ROBLEMS
Revised:  E XPERIMENTS 3. P ROBLEMS

Original: We conduct a brief comparison on six two-dimensional toy problems with R EAL NVP to demonstrate the potential gain in expressivity R AD models can enable.
Revised: We conduct a brief comparison on six two-dimensional toy problems with R EAL NVP to demonstrate the potential gain in expressivity R AD models can enable.

Original: Synthetic datasets of 10 , 000 points each are constructed following the manifold hypothesis and/or the clustering hypothesis .
Revised: Synthetic datasets of 10 , 000 points each are constructed following the manifold hypothesis and/or the clustering hypothesis .

Original: We designate these problems as: grid Gaussian mixture , ring Gaussian mixture , two moons , two circles , spiral , and many moons (see Figure 5).
Revised: We designate these problems as: grid Gaussian mixture , ring Gaussian mixture , two moons , two circles , spiral , and many moons (see Figure 5).

Original: 3. A RCHITECTURE
Revised: 3. A RCHITECTURE

Original: For the R AD model implementation, we use the piecewise linear activations deﬁned in Appendix A in a coupling layer architecture (Dinh et al., 2015; 2017) for f Z where, instead of a conditional linear transformation, the conditioning variable x 1 determines the parameters of the piecewise linear activation on x 2 to obtain z 2 and k 2 , with z 1 = x 1 (see Figure 6).
Revised: For the R AD model implementation, we use the piecewise linear activations deﬁned in Appendix A in a coupling layer architecture (Dinh et al., 2015; 2017) for f Z where, instead of a conditional linear transformation, the conditioning variable x 1 determines the parameters of the piecewise linear activation on x 2 to obtain z 2 and k 2 , with z 1 = x 1 (see Figure 6).

Original: For the gating network p K | Z , the gating logit neural network s ( z ) take as input z
Revised: For the gating network p K | Z , the gating logit neural network s ( z ) take as input z

Original: = ( z 1 , z 2 ) .
Revised: = ( z 1 , z 2 ) .

Original: We compare with a R EAL NVP model using only afﬁne coupling layers.
Revised: We compare with a R EAL NVP model using only afﬁne coupling layers.

Original: p Z is a standard Gaussian distribution.
Revised: p Z is a standard Gaussian distribution.

Original: As both these models can easily approximately solve these generative modeling tasks provided enough capacity, we study these model in a relatively low capacity regime, where we can showcase the potential expressivity R AD may provide.
Revised: As both these models can easily approximately solve these generative modeling tasks provided enough capacity, we study these model in a relatively low capacity regime, where we can showcase the potential expressivity R AD may provide.

Original: Each of these models uses six coupling layers, and each coupling layer uses a one-hidden-layer rectiﬁed network with a tanh output activation scaled by a scalar parameter as described in Dinh et al.
Revised: Each of these models uses six coupling layers, and each coupling layer uses a one-hidden-layer rectiﬁed network with a tanh output activation scaled by a scalar parameter as described in Dinh et al.

Original: For R AD , the logit network s ( · ) also uses a one-hidden-layer rectiﬁed neural network, but with linear output.
Revised: For R AD , the logit network s ( · ) also uses a one-hidden-layer rectiﬁed neural network, but with linear output.

Original: In order to fairly compare with respect to number of parameters, we provide R EAL NVP seven times more hidden units per
Revised: In order to fairly compare with respect to number of parameters, we provide R EAL NVP seven times more hidden units per

Original: [Figure]
Revised: [Figure]

Original:  hidden layer than R AD , which uses 8 hidden units  hidden layer.
Revised: where the data distribution has low density. Here, these spaces clusters illustrat ing the challenges that come with modeling multimodal data as one continuous hidden layer than R AD , which uses 8 hidden units per hidden layer.

Original:  are trained using stochastic gradient ascent with A DAM (Kingma & Ba, 2015) on the log-likelihood with a batch size of 500 for 50 , 000 steps.
Revised: For each level, p K | Z and f Z are trained using stochastic gradient ascent with A DAM (Kingma & Ba, 2015) on the log-likelihood with a batch size of 500 for 50 , 000 steps.

Original: 3. R ESULTS
Revised: 3. R ESULTS

Original: In each of these problems, R AD is consistently able to obtain higher log-likelihood than R EAL NVP.
Revised: In each of these problems, R AD is consistently able to obtain higher log-likelihood than R EAL NVP.

Original: R AD
Revised: R AD

Original: R EAL NVP
Revised: R EAL NVP

Original: Grid Gaussian mixture − 1 .− 2 .
Revised: Grid Gaussian mixture − 1 .− 2 .

Original: Ring Gaussian mixture 3 .1 .
Revised: Ring Gaussian mixture 3 .1 .

Original: Two moons − 1 .− 1 .
Revised: Two moons − 1 .− 1 .

Original: Two cicles − 1 .− 2 .
Revised: Two cicles − 1 .− 2 .

Original: Spiral 0 .− 0 .
Revised: Spiral 0 .− 0 .

Original: Many moons − 0 .−
Revised: Many moons − 0 .−

Original: S AMPLING AND G AUSSIANIZATION
Revised: S AMPLING AND G AUSSIANIZATION

Original: We plot the samples (Figure 7) of the described R AD and R EAL NVP models trained on these problems.
Revised: We plot the samples (Figure 7) of the described R AD and R EAL NVP models trained on these problems.

Original: In the described low capacity regime, R EAL NVP fails by attributing probability mass over spaces where the data distribution has low density.
Revised: In the described low capacity regime, R EAL NVP fails by attributing probability mass over spaces where the data distribution has low density.

Original: This is consistent with the mode covering behavior of maximum likelihood.
Revised: This is consistent with the mode covering behavior of maximum likelihood.

Original: However, the particular inductive bias of R EAL NVP is to prefer modeling the data as one connected manifold.
Revised: However, the particular inductive bias of R EAL NVP is to prefer modeling the data as one connected manifold.

Original: This results in the unwanted probability mass being distributed along the space between clusters.
Revised: This results in the unwanted probability mass being distributed along the space between clusters.

Original: Flow-based models often follow the principle of Gaussianization (Chen & Gopinath, 2001), i.e. transforming the data distribution into a Gaussian.
Revised: Flow-based models often follow the principle of Gaussianization (Chen & Gopinath, 2001), i.e. transforming the data distribution into a Gaussian.

Original: The inversion of that process on a Gaussian distribution would therefore approximate the data distribution.
Revised: The inversion of that process on a Gaussian distribution would therefore approximate the data distribution.

Original: We plot in Figure 8 the inferred
Revised: We plot in Figure 8 the inferred

Original: Gaussianized variables z (5) for both models trained on the ring Gaussian mixture problem.
Revised: Gaussianized variables z (5) for both models trained on the ring Gaussian mixture problem.

Original: The
Revised: The

Original: Gaussianization from R EAL NVP leaves some area of the standard Gaussian distribution unpopulated.
Revised: Gaussianization from R EAL NVP leaves some area of the standard Gaussian distribution unpopulated.

Original: These unattended areas correspond to unwanted regions of probability mass in the input space.
Revised: These unattended areas correspond to unwanted regions of probability mass in the input space.

Original: R AD suffers signiﬁcantly less from this problem.
Revised: R AD suffers signiﬁcantly less from this problem.

Original: An interesting feature is that R AD seems also to outperform R EAL NVP on the spiral dataset.
Revised: An interesting feature is that R AD seems also to outperform R EAL NVP on the spiral dataset.

Original: One hypothesis is that the model successfully exploits some non-linear symmetries in this problem.
Revised: One hypothesis is that the model successfully exploits some non-linear symmetries in this problem.

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

Original: 3.3. F OLDING
Revised: 3.3. F OLDING

Original: We take a deeper look at the Gaussianization process involved in both models.
Revised: We take a deeper look at the Gaussianization process involved in both models.

Original: In Figure 9 we plot the inference process of z (5) from x for both models trained on the two moons problem.
Revised: In Figure 9 we plot the inference process of z (5) from x for both models trained on the two moons problem.

Original: As a result of a folding process similar to that in Montufar et al.
Revised: As a result of a folding process similar to that in Montufar et al.

Original: (2014), several points which were far apart in the input space become neighbors in z (5) in the case of R AD .
Revised: (2014), several points which were far apart in the input space become neighbors in z (5) in the case of R AD .

Original: We further explore this folding process using the visualization described in Figure 10.
Revised: We further explore this folding process using the visualization described in Figure 10.

Original: We verify that the non-linear folding process induced by R AD plays at least two roles: bridging gaps in the distribution of probability mass, and exploiting symmetries in the data.
Revised: We verify that the non-linear folding process induced by R AD plays at least two roles: bridging gaps in the distribution of probability mass and exploiting symmetries in the data.

Original: We observe that in the case of the ring Gaussian mixture (Figure 11a), R AD effectively uses foldings in order to bridge the different modes of the distribution into a single mode, primarily in the last layers of the transformation.
Revised: We observe that in the case of the ring Gaussian mixture (Figure 11a), R AD effectively uses foldings in order to bridge the different modes of the distribution into a single mode, primarily in the last layers of the transformation.

Original: We contrast this with R EAL NVP (Figure 11b) which struggles to combine these modes under the standard Gaussian distribution using bijections.
Revised: We contrast this with R EAL NVP (Figure 11b) which struggles to combine these modes under the standard Gaussian distribution using bijections.

Original: In the spiral problem (Figure 12), R AD decomposes the spiral into three different lines to bridge (Figure 12a) instead of unrolling the manifold fully, which R EAL NVP struggles to do (Figure 12b).
Revised: In the spiral problem (Figure 12), R AD decomposes the spiral into three different lines to bridge (Figure 12a) instead of unrolling the manifold fully, which R EAL NVP struggles to do (Figure 12b).

Original: In both cases, the points remain generally well separated by labels, even after being pushed through a
Revised: In both cases, the points remain generally well separated by labels, even after being pushed through a

Original: R AD layer (Figure 11a and 12a).
Revised: R AD layer (Figure 11a and 12a).

Original: This enables the model to maximize the conditional log-probability log( p K | Z ) .
Revised: This enables the model to maximize the conditional log-probability log( p K | Z ) .

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We introduced an approach to tractably evaluate and train deep mixture models using piecewise invertible maps as a folding mechanism.
Revised: We introduced an approach to tractably evaluate and train deep mixture models using piecewise invertible maps as a folding mechanism.

Original: This allows exact inference, exact generation, and exact evaluation of log-likelihood, avoiding many issues in previous discrete variables models.
Revised: This allows exact inference, exact generation, and exact evaluation of log-likelihood, avoiding many issues in previous discrete variables models.

Original: This method can easily be combined with other ﬂow based architectural components, allowing ﬂow based models to better model datasets with discrete as well as continuous structure.
Revised: This method can easily be combined with other ﬂow based architectural components, allowing ﬂow based models to better model datasets with discrete as well as continuous structure.

Original: A CKNOWLEDGEMENTS
Revised: A CKNOWLEDGEMENTS

Original: The authors would like to thank Kyle Kastner, Johanna Hansen, Harm De Vries, Ben Poole,Prajit
Revised: The authors would like to thank Kyle Kastner, Johanna Hansen, Harm De Vries, Ben Poole,

Original:  Ramachandran, Dustin Tran, David Grangier, George J. Tucker, Matt D. Hoffman, Daniel Duckworth, Anna Huang,  Arvind Neelakantan, Dale Schuurmans, Graham Taylor, Bart van Merrienb ¨ oer, Daniel Duckworth, Vincent Dumoulin, Marc G. Bellemare, and Ross Goroshin for valuable discussion and feedbacks.
Revised: Prajit Ramachandran, Dustin Tran, David Grangier, George J. Tucker, Matt D. Hoffman, Daniel Duckworth, Anna Huang, Fabian Pedregosa, Arvind Neelakantan, Dale Schuurmans, Graham Taylor, Bart van Merrienb ¨ oer, Daniel Duckworth, Vincent Dumoulin, Marc G. Bellemare, and Ross Goroshin for valuable discussion and feedbacks.

Original: A C ONTINUITY
Revised: A C ONTINUITY

Original: The standard approach in learning a deep probabilistic model has been stochastic gradient descent on the negative log-likelihood.
Revised: The standard approach in learning a deep probabilistic model has been stochastic gradient descent on the negative log-likelihood.

Original: Although the model enables the computation of a gradient almost everywhere, the log-likelihood is unfortunately discontinuous.
Revised: Although the model enables the computation of a gradient almost everywhere, the log-likelihood is unfortunately discontinuous.

Original: Let’s decompose the log-likelihood
Revised: Let’s decompose the log-likelihood

Original: [Equation]
Revised: [Equation]

Original: There are two sources of discontinuity in this expression: f K is a function with discrete values (therefore discontinuous) and ∂f Z ∂ x T is discontinuous because of the transition between the subsets A k , leading to the expression of interest
Revised: There are two sources of discontinuity in this expression: f K is a function with discrete values (therefore discontinuous) and ∂f Z ∂ x T is discontinuous because of the transition between the subsets A k , leading to the expression of interest

Original: [Equation]
Revised: [Equation]

Original: which takes a role similar to the log-Jacobian determinant, a pseudo log-Jacobian determinant .
Revised: which takes a role similar to the log-Jacobian determinant, a pseudo log-Jacobian determinant .

Original: Let’s build from now on the simple scalar case and a piecewise linear function
Revised: Let’s build from now on the simple scalar case and a piecewise linear function

Original: [Equation]
Revised: [Equation]

Original: or f Z ( x ) =
Revised: or f Z ( x ) =

Original: α 1 min( x + β, 0) + α 3 max( x − β, 0) − α 2 min(max( x, − β ) , β ) (with α k > 0 and β > = 0 , see ﬁgure13a), then A 1 =] − ∞ , − β
Revised: α 1 min( x + β, 0) + α 3 max( x − β, 0) − α 2 min(max( x, − β ) , β ) (with α k > 0 and β > = 0 , see ﬁgure13a), then A 1 =] − ∞ , − β

Original: [ , A 2 = [ − β, β ] , A 3 =] β, + ∞ [ and f K ( x ) = 1 ( x > − β ) + 1 ( x > β ) .
Revised: [ , A 2 = [ − β, β ] , A 3 =] β, + ∞ [ and f K ( x ) = 1 ( x > − β ) + 1 ( x > β ) .

Original: In this case, s ( z ) =
Revised: In this case, s ( z ) =

Original: (cid:0) log (cid:0) p K | Z (cid:0) k
Revised: (cid:0) log (cid:0) p K | Z (cid:0) k

Original: | z (cid:1)(cid:1)(cid:1) k ≤ N + C 1 | K | can be seen as a vector valued function.
Revised: | z (cid:1)(cid:1)(cid:1) k ≤ N + C 1 | K | can be seen as a vector valued function.

Original: We can attempt at parametrizing the model such that the pseudo log-Jacobian determinant becomes continuous with respect to β by expressing the boundary condition at x = β
Revised: We can attempt at parametrizing the model such that the pseudo log-Jacobian determinant becomes continuous with respect to β by expressing the boundary condition at x = β

Original: [Equation]
Revised: [Equation]

Original:  ⇒ s ( − α 2 β ) 2 + log( α 2 ) =
Revised: ( β + ) (cid:1) ⇒ s ( − α 2 β ) 2 + log( α 2 ) =

Original: s ( − α 2 β ) 3 + log( α 3 ) .
Revised: s ( − α 2 β ) 3 + log( α 3 ) .

Original: If we deﬁne Ω 2 , 3 = (cid:20) 

Revised: If we deﬁne Ω 2 , 3 = (cid:20)

Original: Ω (cid:21) and Ω 1 , 2 = (cid:20) Ω  (cid:21) , with Ω =(cid:20) −  −(cid:21) , then this boundary condition can be enforced, together with a similar one at x = − β , by replacing the function s with
Revised: Ω (cid:21) and Ω 1 , 2 = (cid:20) Ω  (cid:21) , with Ω =(cid:20) −  −(cid:21) , then this boundary condition can be enforced, together with a similar one at x = − β , by replacing the function s with

Original: [Equation]
Revised: [Equation]

Original: Another type of boundary condition can be found at between the non-invertible area and the invertible area z
Revised: Another type of boundary condition can be found at between the non-invertible area and the invertible area z

Original: = α 2 β , as ∀ z > α 2 β, p 3 | Z (3 | z ) = 1 , therefore
Revised: = α 2 β , as ∀ z > α 2 β, p 3 | Z (3 | z ) = 1 , therefore

Original: [Equation]
Revised: [Equation]

Original: Since the condition ∀ k < 3 , p K | Z (cid:0) k | z ) → 0 when z → ( α 2 β ) − will lead to an inﬁnite loss barrier at x = − β, another way to enforce this boundary condition is by adding linear pieces (Figure 13b): f Z
Revised: Since the condition ∀ k < [Equation] p K | Z (cid:0) k | z ) → 0 when z → ( α 2 β ) − will lead to an inﬁnite loss barrier at x = − β

Original: [Equation]
Revised: [Equation]

Original:  = 
Revised: another way to enforce this boundary condition is by adding linear pieces (Figure 13b): f Z ( x ) = 

Original: [Equation]
Revised: [Equation]

Original: The inverse is deﬁned as
Revised: The inverse is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: In order to know the values of s at the boundaries ±
Revised: In order to know the values of s at the boundaries ±

Original: α 2 β
Revised: α 2 β

Original: [Equation]
Revised: [Equation]

Original: we can use the logit function
Revised: we can use the logit function

Original: [Equation]
Revised: [Equation]

Original: where ( a s , b s ) ∈ ( R 3 ) 2 .
Revised: where ( a s , b s ) ∈ ( R 3 ) 2 .

Original: Given those constraints, the model can then be reliably learned through gradient descent methods.
Revised: Given those constraints, the model can then be reliably learned through gradient descent methods.

Original: Note that the resulting tractability of the model results from the fact that the discrete variables k is only interfaced during inference with the distribution p K | Z , unlike discrete variational autoencoders approaches (Mnih & Gregor, 2014; van den Oord et al., 2017) where it is fed to a deep neural network.
Revised: Note that the resulting tractability of the model results from the fact that the discrete variables k is only interfaced during inference with the distribution p K | Z , unlike discrete variational autoencoders approaches (Mnih & Gregor, 2014; van den Oord et al., 2017) where it is fed to a deep neural network.

Original: Similar to Rolfe (2017), the learning of discrete variables is achieved by relying on the the continuous component of the model, and, as opposed as other approaches (Jang et al., 2017;
Revised: Similar to Rolfe (2017), the learning of discrete variables is achieved by relying on the the continuous component of the model, and, as opposed as other approaches (Jang et al., 2017;

Original: Maddison et al., 2017; Grathwohl et al., 2018; Tucker et al., 2017), this gradient signal extracted is exact and closed form. 
Revised: Maddison et al., 2017; Grathwohl et al., 2018; Tucker et al., 2017), this gradient signal extracted is exact and closed form. [Figure] [Figure] [Figure]

Original: B I NFERENCE PROCESSES
Revised: B I NFERENCE PROCESSES

Original: We plot the remaining inference processes of R AD and R EAL NVP on the remaining problems not plotted previously: grid Gaussian mixture (Figure 14), two circles (Figure 15), two moons (Figure 16) and many moons (Figure 17).
Revised: We plot the remaining inference processes of R AD and R EAL NVP on the remaining problems not plotted previously: grid Gaussian mixture (Figure 14), two circles (Figure 15), two moons (Figure 16) and many moons (Figure 17).

Original: We also compare the ﬁnal results of the Gaussianization processes on both models on the different toy problems in Figure 18. 
Revised: We also compare the ﬁnal results of the Gaussianization processes on both models on the different toy problems in Figure 18. [Figure] [Figure] [Figure]

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

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]