Original: D IRECT O PTIMIZATION THROUGH arg max FOR D ISCRETE V ARIATIONAL A UTO -E NCODER
Revised: D IRECT O PTIMIZATION THROUGH arg max FOR D ISCRETE V ARIATIONAL A UTO -E NCODER

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Reparameterization of variational auto-encoders is an effective method for reducing the variance of their gradient estimates.
Revised: Reparameterization of variational auto-encoders is an effective method for reducing the variance of their gradient estimates.

Original: However, when the latent variables are discrete, a reparameterization is problematic due to discontinuities in the discrete space.
Revised: However, when the latent variables are discrete, a reparameterization is problematic due to discontinuities in the discrete space.

Original: In this work, we extend the direct loss minimization technique to discrete variational auto-encoders.
Revised: In this work, we extend the direct loss minimization technique to discrete variational auto-encoders.

Original: We ﬁrst reparameterize a discrete random variable using the arg max function of the Gumbel-Max perturbation model.
Revised: We ﬁrst reparameterize a discrete random variable using the arg max function of the Gumbel-Max perturbation model.

Original: We then use direct optimization to propagate gradients through the non-differentiable arg max using two perturbed arg max operations
Revised: We then use direct optimization to propagate gradients through the non-differentiable arg max using two perturbed arg max operations

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Models with discrete latent variables drive extensive research in machine learning applications, such as language classiﬁcation and generation (Yogatama et al., 2016;
Revised: Models with discrete latent variables drive extensive research in machine learning applications, such as language classiﬁcation and generation (Yogatama et al., 2016;

Original: Hu et al., 2017; Shen et al., 2018), molecular synthesis (Kusner et al., 2017), or game solving (Mena et al., 2018).
Revised: Hu et al., 2017; Shen et al., 2018), molecular synthesis (Kusner et al., 2017), or game solving (Mena et al., 2018).

Original: Compared to their continuous counterparts, discrete latent variable models can decrease the computational complexity of inference calculations, for instance, by discarding alternatives in hard attention models (Lawson et al., 2017), they can improve interpretability by illustrating which terms contributed to the solution (Mordatch & Abbeel, 2017; Yogatama et al., 2016), and they can facilitate the encoding of inductive biases in the learning process, such as images consisting of a small number of objects (Eslami et al., 2016) or tasks requiring intermediate alignments (Mena et al., 2018).
Revised: Compared to their continuous counterparts, discrete latent variable models can decrease the computational complexity of inference calculations, for instance, by discarding alternatives in hard attention models (Lawson et al., 2017), they can improve interpretability by illustrating which terms contributed to the solution (Mordatch & Abbeel, 2017; Yogatama et al., 2016), and they can facilitate the encoding of inductive biases in the learning process, such as images consisting of a small number of objects (Eslami et al., 2016) or tasks requiring intermediate alignments (Mena et al., 2018).

Original: Finally, in some cases, discrete latent variables are natural choices, including when modeling datasets with discrete classes (Rolfe, 2016;
Revised: Finally, in some cases, discrete latent variables are natural choices, including when modeling datasets with discrete classes (Rolfe, 2016;

Original: Jang et al., 2016; Maddison et al., 2016).
Revised: Jang et al., 2016; Maddison et al., 2016).

Original: Models involving discrete latent variables are hard to train, with the key issue being to estimate the gradients of the resulting non-differentiable objectives.
Revised: Models involving discrete latent variables are hard to train, with the key issue being to estimate the gradients of the resulting non-differentiable objectives.

Original: While one can use unbiased estimators, such as REINFORCE (Williams, 1992), their variance is typically high (Paisley et al., 2012; Mnih & Gregor, 2014; Titsias, 2015; Gu et al., 2015; Mnih & Rezende, 2016; Tucker et al., 2017).
Revised: While one can use unbiased estimators, such as REINFORCE (Williams, 1992), their variance is typically high (Paisley et al., 2012; Mnih & Gregor, 2014; Titsias, 2015; Gu et al., 2015; Mnih & Rezende, 2016; Tucker et al., 2017).

Original: In variational auto-encoders (VAEs) with continuous latent variables, the reparameterization trick provides a successful alternative (Kingma & Welling, 2013; Rezende et al., 2014).
Revised: In variational auto-encoders (VAEs) with continuous latent variables, the reparameterization trick provides a successful alternative (Kingma & Welling, 2013; Rezende et al., 2014).

Original: However, it cannot be directly applied to non-differentiable objectives.
Revised: However, it cannot be directly applied to non-differentiable objectives.

Original: Recent work (Maddison et al., 2016; Jang et al., 2016) uses a relaxation of the discrete VAE objective, where latent variables follow a Gumbel-Softmax distribution.
Revised: Recent work (Maddison et al., 2016; Jang et al., 2016) uses a relaxation of the discrete VAE objective, where latent variables follow a Gumbel-Softmax distribution.

Original: The Gumbel-Softmax approach is motivated by the connection of the Gumbel-Max perturbation model and the Gibbs distribution (Papandreou & Yuille, 2011;
Revised: The Gumbel-Softmax approach is motivated by the connection of the Gumbel-Max perturbation model and the Gibbs distribution (Papandreou & Yuille, 2011;

Original: Tarlow et al., 2012; Hazan et al., 2013; Maddison et al., 2014).
Revised: Tarlow et al., 2012; Hazan et al., 2013; Maddison et al., 2014).

Original: The
Revised: The

Original: Gumbel-Softmax distribution relaxes the non-differentiable Gumbel-Max trick for sampling from the categorical distribution by replacing the arg max operation with a softmax operation.
Revised: Gumbel-Softmax distribution relaxes the non-differentiable Gumbel-Max trick for sampling from the categorical distribution by replacing the arg max operation with a softmax operation.

Original: The reformulation results in a continuous objective function, which allows the use of the reparameterization trick.
Revised: The reformulation results in a continuous objective function, which allows the use of the reparameterization trick.

Original: (Mena et al., 2018) provide an extension of this approach for discrete latent structures, namely distributions over latent matchings.
Revised: (Mena et al., 2018) provide an extension of this approach for discrete latent structures, namely distributions over latent matchings.

Original: Our work proposes minimizing the non-differentiable objective, by extending the direct loss minimization technique to generative models (McAllester et al., 2010; Song et al., 2016).
Revised: Our work proposes minimizing the non-differentiable objective, by extending the direct loss minimization technique to generative models (McAllester et al., 2010; Song et al., 2016).

Original: Since categorical variables are represented by the Gibbs distribution, we start from the arg max formulation of the Gibbs distribution.
Revised: Since categorical variables are represented by the Gibbs distribution, we start from the arg max formulation of the Gibbs distribution.

Original: We subsequently derive an optimization method that directly propagates (biased) gradients through reparameterized arg max .
Revised: We subsequently derive an optimization method that can propagate (biased) gradients through reparameterized arg max .

Original: The direct differentiation of the resulting expectation is estimated by the difference of gradients of two arg max -perturbations.
Revised: The gradient of the resulting expectation is estimated by the difference of gradients of two arg max -perturbations.

Original: We begin by introducing the notation, the VAE formulation and the equivalence between the Gibbs distribution and the Gumbel-Max perturbation model in Section 3.
Revised: We begin by introducing the notation, the VAE formulation and the equivalence between the Gibbs distribution and the Gumbel-Max perturbation model in Section 3.

Original: In Section 4.1, we use these components to reformulate the discrete VAE objective, with respect to the arg max prediction operation.
Revised: In Section 4.1, we use these components to reformulate the discrete VAE objective, with respect to the arg max prediction operation.

Original: We subsequently state and prove the main result that allows us to directly optimize through the arg max function in Section 4.2.
Revised: We subsequently state and prove the main result that allows us to differentiate through the arg max function in Section 4.2.

Original: In Section 5, we extend this result to mixed discrete-continuous VAEs and to semi-supervised VAE objectives.
Revised: In Section 5, we extend this result to mixed discrete-continuous VAEs and to semi-supervised VAE objectives.

Original: Finally, we demonstrate the effectiveness of our approach on image generation.
Revised: Finally, we demonstrate the effectiveness of our approach on image generation.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Variational inference has been extensively studied in machine learning, see Blei et al.
Revised: Variational inference has been extensively studied in machine learning, see Blei et al.

Original: (2017) for a review paper.
Revised: (2017) for a review paper.

Original: In our work, we consider variational Bayes bounds with a discrete latent space.
Revised: In our work, we consider variational Bayes bounds with a discrete latent space.

Original: Many approaches to optimizing the variational Bayes objective, that are based on samples from the distribution, can be seen as applications of the REINFORCE gradient estimator (Williams, 1992).
Revised: Many approaches to optimizing the variational Bayes objective, that are based on samples from the distribution, can be seen as applications of the REINFORCE gradient estimator (Williams, 1992).

Original: These estimators are unbiased, but without a carefully chosen baseline, their variance tends to be too high for the estimator to be useful and considerable work has gone into ﬁnding effective baselines (Paisley et al., 2012).
Revised: These estimators are unbiased, but without a carefully chosen baseline, their variance tends to be too high for the estimator to be useful and considerable work has gone into ﬁnding effective baselines (Paisley et al., 2012).

Original: Other methods use various techniques to reduce the estimator variance (Ranganath et al., 2014; Mnih & Gregor, 2014;
Revised: Other methods use various techniques to reduce the estimator variance (Ranganath et al., 2014; Mnih & Gregor, 2014;

Original: Gu et al., 2015).
Revised: Gu et al., 2015).

Original: Reparameterization is an effective method to reduce the gradient estimate variance in generative learning.
Revised: Reparameterization is an effective method to reduce the gradient estimate variance in generative learning.

Original: Kingma & Welling (2013) have shown its effectiveness in auto-encoding variational Bayes (also called variational auto-encoders, VAEs) for continuous latent spaces.
Revised: Kingma & Welling (2013) have shown its effectiveness in auto-encoding variational Bayes (also called variational auto-encoders, VAEs) for continuous latent spaces.

Original: Rezende et al.
Revised: Rezende et al.

Original: (2014) demonstrated its effectiveness in deep latent models.
Revised: (2014) demonstrated its effectiveness in deep latent models.

Original: The success of these works led to reparameterization approaches in discrete latent spaces.
Revised: The success of these works led to reparameterization approaches in discrete latent spaces.

Original: Rolfe (2016) and Vahdat et al.
Revised: Rolfe (2016) and Vahdat et al.

Original: (2018) represent the marginal distribution of each binary latent variable as a continuous variable in the unit interval.
Revised: (2018) represent the marginal distribution of each binary latent variable as a continuous variable in the unit interval.

Original: This reparameterization allows the backpropogation of gradients through the continuous representation.
Revised: This reparameterization allows the backpropogation of gradients through the continuous representation.

Original: These works are restricted to binary random variables, and as a by-product, it encourages highdimensional representations for which inference is exponential in the dimension size.
Revised: These works are restricted to binary random variables, and as a by-product, it encourages highdimensional representations for which inference is exponential in the dimension size.

Original: In contrast, our work reparameterizes the discrete Gibbs latent model, using a Gumbel-Max perturbation model and directly propagates gradients through the reparameterized objective.
Revised: In contrast, our work reparameterizes the discrete Gibbs latent model, using a Gumbel-Max perturbation model and directly propagates gradients through the reparameterized objective.

Original: Maddison et al.
Revised: Maddison et al.

Original: (2016) and Jang et al.
Revised: (2016) and Jang et al.

Original: (2016) recently introduced a novel distribution, the Concrete distribution or the Gumbel-Softmax, that continuously relaxes discrete random variables.
Revised: (2016) recently introduced a novel distribution, the Concrete distribution or the Gumbel-Softmax, that continuously relaxes discrete random variables.

Original: Replacing every discrete random variable in a model with a Concrete random variable, results in a continuous model, where the reparameterization trick is applicable.
Revised: Replacing every discrete random variable in a model with a Concrete random variable, results in a continuous model, where the reparameterization trick is applicable.

Original: These works are close to ours, with a few notable differences.
Revised: These works are close to ours, with a few notable differences.

Original: They use the Gumbel-Softmax function and their model is smooth and reparameterization may use the chain rule to propagate gradients.
Revised: They use the Gumbel-Softmax function and their model is smooth and reparameterization may use the chain rule to propagate gradients.

Original: Similar to our setting, the GumbelSoftmax operation results in a biased estimate of the gradient.
Revised: Similar to our setting, the GumbelSoftmax operation results in a biased estimate of the gradient.

Original: Different from our setting, the softmax operation relaxes the variational Bayes objective and results in a non-tight representation.
Revised: Different from our setting, the softmax operation relaxes the variational Bayes objective and results in a non-tight representation.

Original: Our work uses the Gumbel-Max perturbation model, which is an equivalent representation of the Gibbs distribution.
Revised: Our work uses the Gumbel-Max perturbation model, which is an equivalent representation of the Gibbs distribution.

Original: With that, we do not relax the variational Bayes objective, while our arg max prediction remains non-differentiable and we cannot naively use the chain rule to propagate gradients.
Revised: With that, we do not relax the variational Bayes objective, while our arg max prediction remains non-differentiable and we cannot naively use the chain rule to propagate gradients.

Original: Instead, we develop a direct optimization method to propagate gradients through arg max operation using the difference of gradients of two max-perturbations.
Revised: Instead, we develop a direct optimization method to propagate gradients through arg max operation using the difference of gradients of two max-perturbations.

Original: Our gradient estimate is biased, except for its limit.
Revised: Our gradient estimate is biased, except for its limit.

Original: Differentiating through arg max prediction was previously done in discriminative learning, in the context of direct loss minimization (McAllester et al., 2010; Song et al., 2016).
Revised: Differentiating through arg max prediction was previously done in discriminative learning, in the context of direct loss minimization (McAllester et al., 2010; Song et al., 2016).

Original: Unfortunately, direct loss minimization cannot be applied to generative learning, since it does not have a posterior distribution around its arg max prediction.
Revised: Unfortunately, direct loss minimization cannot be applied to generative learning, since it does not have a posterior distribution around its arg max prediction.

Original: We apply the Gumbel-Max perturbation model to transform the arg max prediction to the Gibbs distribution.
Revised: We apply the Gumbel-Max perturbation model to transform the arg max prediction to the Gibbs distribution.

Original: This also allows us to overcome the “general position” assumption in (McAllester et al., 2010; Song et al., 2016) using our “prediction generating function”.
Revised: This also allows us to overcome the “general position” assumption in (McAllester et al., 2010; Song et al., 2016) using our “prediction generating function”.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: To model the data generating distribution, we consider samples S = { x 1 , ..., x m } originating from some unknown underlying distribution.
Revised: To model the data generating distribution, we consider samples S = { x 1 , ..., x m } originating from some unknown underlying distribution.

Original: We explain the generation process of a parameterized model p θ ( x ) , by minimizing its log-loss when marginalizing over its hidden variables z .
Revised: We explain the generation process of a parameterized model p θ ( x ) , by minimizing its log-loss when marginalizing over its hidden variables z .

Original: Using variational
Revised: Using variational

Original: Bayes, we upper bound the log-loss of an observed data point
Revised: Bayes, we upper bound the log-loss of an observed data point

Original: [Equation]
Revised: [Equation]

Original: Typically, the model distribution p θ ( x | z ) is a member of the exponential family p θ ( x | z ) =
Revised: Typically, the model distribution p θ ( x | z ) is a member of the exponential family p θ ( x | z ) =

Original: e − θ ( x,z ) .
Revised: e − θ ( x,z ) .

Original: When considering a discrete latent space, i.e., z ∈ { 1 , ..., k } , the approximated posterior distribution follows the Gibbs distribution law q φ ( z | x ) ∝ e φ ( x,z ) .
Revised: When considering a discrete latent space, i.e., z ∈ { 1 , ..., k } , the approximated posterior distribution follows the Gibbs distribution law q φ ( z | x ) ∝ e φ ( x,z ) .

Original: The challenge in generative learning is to reparameterize and optimize
Revised: The challenge in generative learning is to reparameterize and optimize

Original: [Equation]
Revised: [Equation]

Original: In our work, we reparameterize the variational bound using the equivalence between Gibbs models and Gumbel-Max perturbation models.
Revised: In our work, we reparameterize the variational bound using the equivalence between Gibbs models and Gumbel-Max perturbation models.

Original: Gumbel-Max perturbation models allow an alternative representation of Gibbs distributions q φ ( z | x ) ∝ e φ ( x,z ) that is based on the extreme value statistics of Gumbel-distributed random variables.
Revised: Gumbel-Max perturbation models allow an alternative representation of Gibbs distributions q φ ( z | x ) ∝ e φ ( x,z ) that is based on the extreme value statistics of Gumbel-distributed random variables.

Original: Let γ be a random function that associates random variable γ ( z ) for each z
Revised: Let γ be a random function that associates random variable γ ( z ) for each z

Original: = 1 , ..., k .
Revised: = 1 , ..., k .

Original: When the random perturbations follow the zero mean Gumbel distribution law, whose probability density function is g ( γ ) =
Revised: When the random perturbations follow the zero mean Gumbel distribution law, whose probability density function is g ( γ ) =

Original: e − ( γ + c
Revised: e − ( γ + c

Original: + e − ( γ + c ) ) for the Euler constant c ≈ 0 .
Revised: + e − ( γ + c ) ) for the Euler constant c ≈ 0 .

Original: 57 , we obtain the following identity between Gibbs models and Gumbel-Max perturbation models 1 (cf.
Revised: 57 , we obtain the following identity between Gibbs models and Gumbel-Max perturbation models 1 (cf.

Original: Kotz & Nadarajah (2000))
Revised: Kotz & Nadarajah (2000))

Original: e φ ( x,z )
Revised: e φ ( x,z )

Original: [Equation]
Revised: [Equation]

Original:  where z
Revised: ] , where z

Original: [Equation]
Revised: [Equation]

Original: For completeness, a proof for this statement appears in Appendix A.
Revised: For completeness, a proof for this statement appears in Appendix A.

Original:  R EPARAMETERIZATION AND DIRECT OPTIMIZATION
Revised:  R EPARAMETERIZATION AND DIRECT OPTIMIZATION

Original: In the following section, we reformulate the discrete VAE objective using the Gumbel-Max perturbation model and present an unbiased estimator for its gradient.
Revised: In the following section, we reformulate the discrete VAE objective using the Gumbel-Max perturbation model and present an unbiased estimator for its gradient.

Original: We then derive an optimization method that directly propagates gradients through the reparameterized arg max function.
Revised: We then derive an optimization method that directly propagates gradients through the reparameterized arg max function.

Original: 4. VAE S WITH G UMBEL -M AX PERTURBATION MODELS
Revised: 4. VAE S WITH G UMBEL -M AX PERTURBATION MODELS

Original: Perturbation models allow an alternative representation of Gibbs distributions q φ ( z | x ) ∝ e φ ( x,z ) .
Revised: Perturbation models allow an alternative representation of Gibbs distributions q φ ( z | x ) ∝ e φ ( x,z ) .

Original: Using the Gumbel-Max perturbation model in Equation (3), the negative log likelihood in Equation (2) takes the form
Revised: Using the Gumbel-Max perturbation model in Equation (3), the negative log likelihood in Equation (2) takes the form

Original: [Equation]
Revised: [Equation]

Revised: These qualities result from applying the law of total expectation, while realizing the probability events for z φ + γ .

Original: The last equality in Equation (4) holds since we restrict to the space to z= z φ + γ  .
Revised: The last equality in Equation (4) holds since we restrict to the space to z φ + γ = z .

Original: The last equality in Equation (5) holds by the linearity of the expectation, i.e., (cid:80) kz
Revised: The last equality in Equation (5) holds by the linearity of the expectation, i.e., (cid:80) kz

Original: =1 E γ ∼ g
Revised: =1 E γ ∼ g

Original: [ 1 z=z φ + γ  θ ( x, z φ + γ )] =
Revised: [ 1 z φ + γ = z θ ( x, z φ + γ )] =

Original: E γ ∼ g
Revised: E [Equation]

Original:  [ (cid:80) kz
Revised: ∼ g [ (cid:80) kz

Original: =1 1 z  =
Revised: =1 1 z φ + γ =

Original: zφ + γ θ ( x, z φ + γ )] and the fact that (cid:80)kz =1 1 z
Revised: z θ ( x, z φ + γ )] and the fact that (cid:80)

Original: =
Revised: [Equation]

Original: z φ + γ = 1 . The gradient of the decoder  θ ( x, z φ + γ ) with respect to its parameters is derived by the chain rule, as the decoder is a smooth function of its parameters for any x, z .
Revised: The gradient of the decoder log probability θ ( x, z φ + γ ) with respect to its parameters is derived by the chain rule.

Original: The main challenge is to evaluate the gradient of E γ ∼ g
Revised: The main challenge is to evaluate the gradient of E γ ∼ g

Original: [ θ ( x, z φ + γ )] with respect to the encoder parameters, since the chain rule does not propagate through the arg max function z φ + γ .
Revised: [ θ ( x, z φ + γ )] with respect to the encoder parameters, since the chain rule does not propagate through the arg max function z φ + γ .

Original: 4. D IRECT OPTIMIZATION THROUGH arg max
Revised: 4. D IRECT OPTIMIZATION THROUGH arg max

Original: Our main result is presented in Theorem 1 and shows how to compute the gradient of the reparameterized discrete VAE, i
Revised: Our main result is presented in Theorem 1 and shows how to compute the gradient of the reparameterized discrete VAE, i

Original: [Equation]
Revised: [Equation]

Original: e., E γ ∼ g
Revised: e., E γ ∼ g

Original: [ θ ( x, z φ + γ )] with respect to the encoder parameters v .
Revised: [ θ ( x, z φ + γ )] with respect to the encoder parameters v .

Original: In the following we omit γ ∼ g for brevity.
Revised: In the following we omit γ ∼ g for brevity.

Original: To make the encoder parameters v explicit in our notation, we denote the encoder function by φ v ( x, z ) .
Revised: To make the encoder parameters v explicit in our notation, we denote the encoder function by φ v ( x, z ) .

Original: Instrumental to our approach, is a novel “prediction generating function”.
Revised: Instrumental to our approach, is a novel “prediction generating function”.

Original: [Equation]
Revised: [Equation]

Original: The proof of Theorem 1 is composed from three steps:
Revised: The proof of Theorem 1 is composed from three steps:

Original: We prove that G ( v, (cid:15) ) is a smooth function of v, (cid:15) .
Revised: We prove that G ( v, (cid:15) ) is a smooth function of v, (cid:15) .

Original: Therefore, the Hessian of G ( v, (cid:15) ) exists and it is symmetric, namely
Revised: Therefore, the Hessian of G ( v, (cid:15) ) exists and it is symmetric, namely

Original: [Equation]
Revised: [Equation]

Original: We show that encoder gradient is apparent in the Hessian:
Revised: We show that encoder gradient is apparent in the Hessian:

Original: [Equation]
Revised: [Equation]

Original: We derive our update rule as the complement representation of the Hessian:
Revised: We derive our update rule as the complement representation of the Hessian:

Original: [Equation]
Revised: [Equation]

Original: The reparameterized gradient computation appears in Equation (10) and it is formally derived in
Revised: The reparameterized gradient computation appears in Equation (10) and it is formally derived in

Original: Theorem 1. Theorem 1.
Revised: Theorem 1. Theorem 1.

Original: Assume φ v ( x, z ) is a smooth function of v .
Revised: Assume φ v ( x, z ) is a smooth function of v .

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: First, we prove that G ( v, (cid:15) ) is a smooth function.
Revised: First, we prove that G ( v, (cid:15) ) is a smooth function.

Original: Recall, g ( γ ) is the zero mean Gumbel probability density function.
Revised: Recall, g ( γ ) is the zero mean Gumbel probability density function.

Original: Applying a change of variable ˆ γ ( z ) =
Revised: Applying a change of variable ˆ γ ( z ) =

Original: (cid:15)θ ( x, ˆ z ) + φ v ( x, ˆ z ) + γ
Revised: (cid:15)θ ( x, ˆ z ) + φ v ( x, ˆ z ) + γ

Original: (ˆ z ) , we obtain
Revised: (ˆ z ) , we obtain

Original: [Equation]
Revised: [Equation]

Original: Since g (ˆ γ − (cid:15)θ − φ v ) is a smooth function of (cid:15) and φ v ( x, z ) and φ v ( x, z ) is a smooth function of v , we conclude that G ( v, (cid:15) ) is a smooth function of v, (cid:15) .
Revised: Since g (ˆ γ − (cid:15)θ − φ v ) is a smooth function of (cid:15) and φ v ( x, z ) and φ v ( x, z ) is a smooth function of v , we conclude that G ( v, (cid:15) ) is a smooth function of v, (cid:15) .

Original: Therefore, the Hessian of G ( v, (cid:15) ) exists and symmetric, i.e., ∂ v ∂ (cid:15)
Revised: Therefore, the Hessian of G ( v, (cid:15) ) exists and symmetric, i.e., ∂ v ∂ (cid:15)

Original: G ( v, (cid:15) ) =
Revised: G ( v, (cid:15) ) =

Original: ∂ (cid:15) ∂ v G ( v, (cid:15) ) .
Revised: ∂ (cid:15) ∂ v G ( v, (cid:15) ) .

Original: We thus proved Equation (7).
Revised: We thus proved Equation (7).

Original: To prove Equations (8) and (9) we differentiate under the integral, both with respect to (cid:15) and with respect to v .
Revised: To prove Equations (8) and (9) we differentiate under the integral, both with respect to (cid:15) and with respect to v .

Original: We are able to differentiate under the integral, since g (ˆ γ − (cid:15)θ − φ v ) is a smooth function of (cid:15) and v and its gradient is bounded by an integrable function (cf.
Revised: We are able to differentiate under the integral, since g (ˆ γ − (cid:15)θ − φ v ) is a smooth function of (cid:15) and v and its gradient is bounded by an integrable function (cf.

Original: Folland (1999), Theorem 2.27).
Revised: Folland (1999), Theorem 2.27).

Original: We turn to prove Equation (8).
Revised: We turn to prove Equation (8).

Original: We begin by noting that max ˆ z { (cid:15)θ ( x, ˆ z )
Revised: We begin by noting that max ˆ z { (cid:15)θ ( x, ˆ z )

Original: + φ v ( x, ˆ z )+ γ
Revised: + φ v ( x, ˆ z )+ γ

Original: (ˆ z ) } equals to (cid:15)θ ( x, z (cid:15)θ + φ v + γ )+ φ v ( x, z (cid:15)θ + φ v + γ )+ γ ( z (cid:15)θ + φ v + γ ) .
Revised: (ˆ z ) } equals to (cid:15)θ ( x, z (cid:15)θ + φ v + γ )+ φ v ( x, z (cid:15)θ + φ v + γ )+ γ ( z (cid:15)θ + φ v + γ ) .

Original: Therefore, by differentiating under the integral, ∂ (cid:15) G ( v, (cid:15) ) = 
Revised: Therefore, by differentiating under the integral, ∂ (cid:15) G ( v, (cid:15) ) = E [Equation] [Equation]

Original: E γ [ ∂ (cid:15) (cid:15)θ ( x, z (cid:15)θ + φ v + γ ) + φ v ( x, z (cid:15)θ + φ v + γ ) + γ ( z (cid:15)θ + φ v + γ )] =
Revised: x, [Equation] [Equation] [Equation]

Revised: =

Original: E γ [ θ ( x, z (cid:15)θ + φ v + γ )] .
Revised: E [Equation] [Equation] ( x, z (cid:15)θ + φ v + γ )] .

Original: We obtain Equation (8) by differentiating under the integral, now with respect to v , and setting (cid:15) = 0 .
Revised: We obtain Equation (8) by differentiating under the integral, now with respect to v , and setting (cid:15) = 0 .

Original: Finally, we turn to prove Equation (9).
Revised: Finally, we turn to prove Equation (9).

Original: By differentiating under the integral ∂ v G ( v, (cid:15) )
Revised: By differentiating under the integral ∂ v G ( v, (cid:15) )

Original: =
Revised: =

Original: E γ [ ∇ v φ v ( x, z (cid:15)θ + φ v + γ )] .
Revised: E γ [ ∇ v φ v ( x, z (cid:15)θ + φ v + γ )] .

Original: Equation (9) is attained by taking the derivative with respect to (cid:15)
Revised: Equation (9) is attained by taking the derivative with respect to (cid:15)

Original: =on both sides.
Revised: =on both sides.

Original: The theorem follows by combining Equation (7) when (cid:15)
Revised: The theorem follows by combining Equation (7) when (cid:15)

Original: = 

Revised: =

Original: , i.e., ∂ v ∂ (cid:15) G ( v, 0) =
Revised: , i.e., ∂ v ∂ (cid:15) G ( v, 0) =

Original: ∂ (cid:15) ∂ v G ( v, 0) with the equalities in Equations (8) and (9).
Revised: ∂ (cid:15) ∂ v G ( v, 0) with the equalities in Equations (8) and (9).

Original: The gradient estimate in Theorem 1 is unbiased in the limit (cid:15) → 0 .
Revised: The gradient estimate in Theorem 1 is unbiased in the limit (cid:15) → 0 .

Original: However, when we approach the limit, the variance of the estimate increases, as the gradient involves the term 1 (cid:15) .
Revised: However, when we approach the limit, the variance of the estimate increases, as the gradient involves the term 1 (cid:15) .

Original: In practice we use (cid:15) ≥ 0 . 1 which means that the gradient estimate is biased.
Revised: In practice we use (cid:15) ≥ 0 . 1 which means that the gradient estimate is biased.

Original: This bias-variance tradeoff is evaluated in Figure 1, where we computed the gradients 1000 times for a ﬁxed (cid:15) to evaluate the standard deviation for each coordinate of the gradient.
Revised: This bias-variance tradeoff is evaluated in Figure 1, where we computed the gradients 1000 times for a ﬁxed (cid:15) to evaluate the standard deviation for each coordinate of the gradient.

Original: We report in the graph the average of these standard deviations.
Revised: We report in the graph the average of these standard deviations.

Original: The above theorem closely relates to the direct loss minimization technique (cf.
Revised: The above theorem closely relates to the direct loss minimization technique (cf.

Original: McAllester et al.
Revised: McAllester et al.

Original: (2010); Song et al.
Revised: (2010); Song et al.

Original: (2016)), which, in our setting, can be used to compute the gradient of E x θ ( x, z φ ) .
Revised: (2016)), which, in our setting, can be used to compute the gradient of E x θ ( x, z φ ) .

Original: The direct loss minimization predicts a single z φ for a given x and, therefore, cannot generate a posterior distribution on all z = 1 , ..., k , i.e., it lacks a generative model that exists in Gumbel-Max perturbation models.
Revised: The direct loss minimization predicts a single z φ for a given x and, therefore, cannot generate a posterior distribution on all z = 1 , ..., k , i.e., it lacks a generative model that exists in Gumbel-Max perturbation models.

Original:  E XTENSIONS 5. S EMI SUPERVISED GENERATIVE MODELS
Revised:  E XTENSIONS 5. S EMI SUPERVISED GENERATIVE MODELS

Original: The main advantage in our framework is that learning a discrete VAE using Gumbel-Max reparameterization is intimately related to predicting a discrete latent label through the arg max operation z φ + γ .
Revised: The main advantage in our framework is that learning a discrete VAE using Gumbel-Max reparameterization is intimately related to predicting a discrete latent label through the arg max operation z φ + γ .

Original: Therefore, semi-supervised VAEs are naturally integrated into our reparameterization framework using any loss function.
Revised: Therefore, semi-supervised VAEs are naturally integrated into our reparameterization framework using any loss function.

Original: Formally, assume that a subset of the data is labeled, i.e.,
Revised: Formally, assume that a subset of the data is labeled, i.e.,

Original: S 1 = { ( x 1 , z 1 ) , ..., ( x m 1 , z m 1 ) } .
Revised: S 1 = { ( x 1 , z 1 ) , ..., ( x m 1 , z m 1 ) } .

Original: In semi-supervised learning, we may add to the learning objective the loss function (cid:96) ( z, z φ + γ ) , for any ( x, z ) ∈ S 1 , to better control the prediction of the latent space.
Revised: In semi-supervised learning, we may add to the learning objective the loss function (cid:96) ( z, z φ + γ ) , for any ( x, z ) ∈ S 1 , to better control the prediction of the latent space.

Original: The semi-supervised discrete VAEs objective function is(cid:88)
Revised: The semi-supervised discrete VAEs objective function is

Original: [Equation]
Revised: [Equation]

Original: The supervised component is explicitly handled by Theorem 1 and optimizing a semi-supervised discrete VAEs is straight forward in our framework.
Revised: The supervised component is explicitly handled by Theorem 1 and optimizing a semi-supervised discrete VAEs is straight forward in our framework.

Original: Our supervised component is intimately related to direct loss minimization (McAllester et al., 2010; Song et al., 2016), which, in our setting, minimizes (cid:80) ( x,z ) ∈ S 1 (cid:96) ( z, z φ ) .
Revised: Our supervised component is intimately related to direct loss minimization (McAllester et al., 2010; Song et al., 2016), which, in our setting, minimizes (cid:80) ( x,z ) ∈ S 1 (cid:96) ( z, z φ ) .

Original: Compared to direct loss minimization, our work adds random perturbation γ to the encoder and thus overcomes the “general position” assumption of direct loss minimization.
Revised: Compared to direct loss minimization, our work adds random perturbation γ to the encoder and thus overcomes the “general position” assumption of direct loss minimization.

Original: This addition allows us to introduce the “prediction generating function”, which greatly simpliﬁes our proof.
Revised: This addition allows us to introduce the “prediction generating function”, which greatly simpliﬁes our proof.

Original: In addition, the added random perturbation γ allows us to use a generative n × k direct
Revised: In addition, the added random perturbation γ allows us to use a generative n × k direct

Original: GSM
Revised: GSM

Original: REBAR
Revised: REBAR

Original: RELAX 20 ×103.
Revised: RELAX 20 ×103.

Original: - - 20 ×126.
Revised: - - 20 ×126.

Original: 1 ×145.
Revised: 1 ×145.

Original: model to prediction, namely, we can randomly generate different explanations z φ + γ while the direct loss minimization allows a single explanation in the form of z φ .
Revised: model to prediction, namely, we can randomly generate different explanations z φ + γ while the direct loss minimization allows a single explanation in the form of z φ .

Original: Semi-supervised generative models require the discrete space to have a semantic meaning.
Revised: Semi-supervised generative models require the discrete space to have a semantic meaning.

Original: For example, in generating images of digits, the image may be represented both by the discrete digit class (e.g., 0 , 1 , ... ) and the continuous style (e.g., bold, tilted,...).
Revised: For example, in generating images of digits, the image may be represented both by the discrete digit class (e.g., 0 , 1 , ... ) and the continuous style (e.g., bold, tilted,...).

Original: Following (Kingma et al., 2014; Jang et al., 2016)
Revised: Following (Kingma et al., 2014; Jang et al., 2016)

Original: we use a mixture of discrete and continuous latent variables in our semi-supervised setting, to capture both the discrete class information and the continuous class style.
Revised: we use a mixture of discrete and continuous latent variables in our semi-supervised setting, to capture both the discrete class information and the continuous class style.

Original: The discrete latent space is represented by the Gumbel-Max perturbation model while the continuous latent space is represented by a Gaussian random variable.
Revised: The discrete latent space is represented by the Gumbel-Max perturbation model while the continuous latent space is represented by a Gaussian random variable.

Original: An illustration of our mixed discrete-continuous VAE architecture appears in ﬁgure 1.
Revised: An illustration of our mixed discrete-continuous VAE architecture appears in ﬁgure 1.

Original: The gradient of the continuous component of the encoder is computed by the chain rule.
Revised: The gradient of the continuous component of the encoder is computed by reparameterization trick.

Original: The gradient of the discrete component is computed according to
Revised: The gradient of the discrete component is computed according to

Original: Theorem 1.
Revised: Theorem 1.

Original: 5. D ISCRETE PRODUCT SPACES
Revised: 5. D ISCRETE PRODUCT SPACES

Original: Discrete VAEs are being applied also to discrete product spaces (Maddison et al., 2016; Jang et al., 2016).
Revised: Discrete VAEs are being applied also to discrete product spaces (Maddison et al., 2016; Jang et al., 2016).

Original: A single discrete random variable z ∈ { 1 , .., k } may not represent the variability of the generative process.
Revised: A single discrete random variable z ∈ { 1 , .., k } may not represent the variability of the generative process.

Original: Therefore, one may represent the latent space with a discrete product space z
Revised: Therefore, one may represent the latent space with a discrete product space z

Original: = ( z 1 , ..., z n ) , where z i ∈ { 1 , ..., k } .
Revised: = ( z 1 , ..., z n ) , where z i ∈ { 1 , ..., k } .

Original: We note that Theorem 1 holds without any change, if we match the structures of θ ( x, z ) , φ ( x, z ) and γ ( z ) , i.e., we have an independent Gumbel random variable for each z =
Revised: We note that Theorem 1 holds without any change, if we match the structures of θ ( x, z ) , φ ( x, z ) and γ ( z ) , i.e., we have an independent Gumbel random variable for each z =

Original: ( z 1 , ..., z n ) .
Revised: ( z 1 , ..., z n ) .

Original: However, this may be computationally inefﬁcient, since the number of
Revised: However, this may be computationally inefﬁcient, since the number of

Original: Gumbel random variables in this case is exponential in n .
Revised: Gumbel random variables in this case is exponential in n .

Original: The gradient in Theorem 1 requires two arg max -perturbations:
Revised: The gradient in Theorem 1 requires two arg max -perturbations:

Original: z φ + γ that requires the encoder output and z (cid:15)θ + φ + γ that requires both the encoder and the decoder outputs.
Revised: z φ + γ that requires the encoder output and z (cid:15)θ + φ + γ that requires both the encoder and the decoder outputs.

Original: Similarly to Maddison et al.
Revised: Similarly to Maddison et al.

Original: (2016) and Jang et al.
Revised: (2016) and Jang et al.

Original: we use a mean ﬁeld decomposition of the encoder,  φ ( x, z ) = (cid:80) ni
Revised: we use a conditionally independent encoder, i.e., φ ( x, z ) = (cid:80) ni

Original: =

Revised: =

Original: φ i ( x, z i ) for which we can compute z φ + γ efﬁciently.
Revised: φ i ( x, z i ) for which we can compute z φ + γ efﬁciently.

Original: In contrast to their approach, we use a mean ﬁeld approximation for the decoder for which we can approximate z (cid:15)θ + φ + γ efﬁciently.
Revised: In contrast to their approach, we use an approximation for the decoder for which we can approximate z (cid:15)θ + φ + γ efﬁciently.

Original: Our mean-ﬁeld approximation considers a low-dimensional random perturbation γ ( z ) = (cid:80) ni
Revised: In discrete product spaces we use low-dimensional random perturbation γ ( z ) = (cid:80) ni

Original: =

Revised: =

Original: γ i ( z i ) .
Revised: γ i ( z i ) .

Original: Thus instead of exponential number of random variables we use a linear number of random variables (linear in k and n ).
Revised: Thus instead of exponential number of random variables we use a linear number of random variables (linear in k and n ).

Original: To compute z φ + γ efﬁciently, we ﬁrst note that the encoder decomposes according to its dimensions, i.e., φ ( x, z )
Revised: To compute z φ + γ efﬁciently, we ﬁrst note that the encoder decomposes according to its dimensions, i.e., φ ( x, z )

Original: = (cid:80) ni
Revised: = (cid:80) ni

Original: =

Revised: =

Original: φ i ( x, z i ) .
Revised: φ i ( x, z i ) .

Original: Thus the perturb-max argument z φ + γ =
Revised: Thus the perturb-max argument z φ + γ =

Original: ( z φ + γ, ..., z φ + γ n ) also decomposes, i.e., z φ + γ
Revised: ( z φ + γ, ..., z φ + γ n ) also decomposes, i.e., z φ + γ

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original:  Therefore, our mean ﬁeld approximation is exact for the arg max -perturbation z
Revised: [Equation] i (ˆ z ) } . Therefore, our approach is exact for the arg max -perturbation z

Original: φ + γ
Revised: φ + γ

Original: Ourmean-ﬁeld approximation is not necessarily exact for the arg max -perturbation z
Revised: Our approximation is not necessarily exact for the arg max -perturbation z

Original: (cid:15)θ + φ
Revised: (cid:15)θ + φ

Original: + γ since the decoder does not decompose according to its dimensions.
Revised: + γ since the decoder does not decompose according to its dimensions.

Original: To be able to compute z (cid:15)θ + φ + γ efﬁciently, we use the fact we can compute z φ + γ efﬁciently and apply a low dimensional approximation for the decoder ˜ θ ( x, z )
Revised: To be able to compute z (cid:15)θ + φ + γ efﬁciently, we use the fact we can compute z φ + γ efﬁciently and apply a low dimensional approximation for the decoder ˜ θ ( x, z )

Original: = (cid:80) ni
Revised: = (cid:80) ni

Original: =1 ˜ θ i
Revised: =1 ˜ θ i

Original: ( z i ) , where
Revised: ( z i ) , where

Original: ˜ θ
Revised: ˜ θ

Original: i
Revised: i

Original: ( z i )
Revised: ( z i )

Original: =
Revised: =

Original: θ ( z γ 1 + φ 1 , ..., z i , ..., z γ n + φ n ) .
Revised: θ ( z γ 1 + φ 1 , ..., z i , ..., z γ n + φ n ) .

Original: With this in mind, we approximate z (cid:15)θ + φ + γ using z (cid:15)
Revised: With this in mind, we approximate z (cid:15)θ + φ + γ using z (cid:15)

Original: ˜ θ + φ + γ , and its coordinates are computed efﬁciently by 
Revised: ˜ θ + φ + γ , and its coordinates are computed efﬁciently by z (cid:15) ˜ θφ + γ i = arg max

Original: [Equation]
Revised: ˆ z =1 ,...,k { (cid:15) ˜ θ i ( x, z i ) + φ i ( x, ˆ z ) + γ i (ˆ z ) } .

Original: accuracy bound #labels direct
Revised: accuracy bound #labels direct

Original: GSM direct
Revised: GSM direct

Original: GSM 95.51% 91.78% 104.
Revised: GSM 95.51% 91.78% 104.

Original: Lastly, the KL-divergence in Equation (1) can be computed efﬁciently when using the decomposed encoder φ ( x, z ) =
Revised: Lastly, the KL-divergence in Equation (1) can be computed efﬁciently when using the decomposed encoder φ ( x, z ) =

Original: (cid:80) ni
Revised: (cid:80) ni

Original: =

Revised: =

Original: φ i ( x, z i ) .
Revised: φ i ( x, z i ) .

Original: In this case q ( z | x ) ∝ e φ ( x,z ) and q ( z i | x ) ∝ e φ i ( x,z i ) , therefore q ( z | x ) = (cid:81) ni
Revised: In this case q ( z | x ) ∝ e φ ( x,z ) and q ( z i | x ) ∝ e φ i ( x,z i ) , therefore q ( z | x ) = (cid:81) ni

Original: =1 q ( z | x ) and KL ( q φ ( z | x ) || p
Revised: =1 q ( z | x ) and KL ( q φ ( z | x ) || p

Original: θ ( z )) =
Revised: θ ( z )) =

Original: (cid:80) ni
Revised: (cid:80) ni

Original: =1 KL ( q φ i
Revised: =1 KL ( q φ i

Original: ( z i | x ) || p θ ( z i )) , where p θ ( z i ) is the marginal probability of p θ ( z ) with respect to its i − th entry.
Revised: ( z i | x ) || p θ ( z i )) , where p θ ( z i ) is the marginal probability of p θ ( z ) with respect to its i − th entry.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We start by comparing our method to the state-of-the-art (Maddison et al., 2016; Jang et al., 2016;
Revised: We start by comparing our method to the state-of-the-art (Maddison et al., 2016; Jang et al., 2016;

Original: Tucker et al., 2017; Grathwohl et al., 2017).
Revised: Tucker et al., 2017; Grathwohl et al., 2017).

Original: We also investigate ourmean-ﬁeld approximation for discrete product spaces.
Revised: We also investigate our approximation for discrete product spaces.

Original: We conclude with a set of experiments that demonstrate the effectiveness of our approach in semi-supervised learning of VAEs and the importance of weak supervision in image generation.
Revised: We conclude with a set of experiments that demonstrate the effectiveness of our approach in semi-supervised learning of VAEs and the importance of weak supervision in image generation.

Original: 6. D ISCRETE VAE S
Revised: 6. D ISCRETE VAE S

Original: We begin our experiments by comparing the test loss bound of our direct optimization biased gradient estimator with the biased gradient estimator Gumbel-Softmax (GSM) of Maddison et al.
Revised: We begin our experiments by comparing the test loss bound of our direct optimization biased gradient estimator with the biased gradient estimator Gumbel-Softmax (GSM) of Maddison et al.

Original: (2016) and Jang et al.
Revised: (2016) and Jang et al.

Original: (2016), and the unbiased estimators REBAR (Tucker et al., 2017) and RELAX (Grathwohl et al., 2017).
Revised: (2016), and the unbiased estimators REBAR (Tucker et al., 2017) and RELAX (Grathwohl et al., 2017).

Original: We performed these experiments using the binarized MNIST dataset (Salakhutdinov & Murray, 2008), and the standard 50,000/10,000 split into training/testing sets.
Revised: We performed these experiments using the binarized MNIST dataset (Salakhutdinov & Murray, 2008), and the standard 50,000/10,000 split into training/testing sets.

Original: Following
Revised: Following

Original: Jang et al.
Revised: Jang et al.

Original: we set our learning rate to 1 e − 3 and the annealing rate to 1 e − 5 and we used their annealing schedule every 1000 step, setting the minimal (cid:15) to be 0 .
Revised: we set our learning rate to 1 e − 3 and the annealing rate to 1 e − 5 and we used their annealing schedule every 1000 step, setting the minimal (cid:15) to be 0 .

Original: When considering a latent space with n random variables z = ( z 1 , ..., z n ) where z i ∈ { 1 , ..., k } , our encoder has a (784 , n × k ) linear layer and the decoder has a ( n × k, 784) linear layer.
Revised: When considering a latent space with n random variables z = ( z 1 , ..., z n ) where z i ∈ { 1 , ..., k } , our encoder has a (784 , n × k ) linear layer and the decoder has a ( n × k, 784) linear layer.

Original: We applied REBAR and RELAX only to discrete binary random variables 2 .
Revised: We applied REBAR and RELAX only to discrete binary random variables 2 .

Original: The results appear in Figure 2.
Revised: The results appear in Figure 2.

Original: For the 20 × 2 and 20 ×architectures, we applied ourmean-ﬁeld approximation that is described in Section 5.2.
Revised: For the 20 × 2 and 20 ×architectures, we applied our approximation that is described in Section 5.2.

Original: Next, we explore our mean-ﬁeld approximation.
Revised: Next, we explore our discrete product space approximation.

Original: We compare a latent space with n random variables z
Revised: We compare a latent space with n random variables z

Original: = ( z 1 , ..., z n ) with z i ∈ { 1 , ..., k } with a latent space that has a single random variable z ∈ { 1 , ...,n k  } .
Revised: = ( z 1 , ..., z n ) with z i ∈ { 1 , ..., k } with a latent space that has a single random variable z ∈ { 1 , ..., k n } .

Original: When n = 3 , k = 2 , a discrete VAE with a single categorial random variable that has 2 3 possible values achieved a test log-loss bound of 168 .
Revised: When n = 3 , k = 2 , a discrete VAE with a single categorial random variable that has 2 3 possible values achieved a test log-loss bound of 168 .

Original: 57 , while 3 binary variables achieved test log-loss bound of 177 . 9 .
Revised: 57 , while 3 binary variables achieved test log-loss bound of 177 . 9 .

Original: When n = 6 , k = 2 , a single categorial random variable achieved a test log-loss bound of 147 .
Revised: When n = 6 , k = 2 , a single categorial random variable achieved a test log-loss bound of 147 .

Original: 59 , while 6 binary variables achieved test log-loss bound of 161 . 09 .
Revised: 59 , while 6 binary variables achieved test log-loss bound of 161 . 09 .

Original: One can see that there is a gap when using ourmean-ﬁeld approximation that increases with n .
Revised: One can see that there is a gap when using our approximation that increases with n .

Original: The main advantage of our framework is that it seamlessly integrates semi-supervised learning.
Revised: The main advantage of our framework is that it seamlessly integrates semi-supervised learning.

Original: For these experiments, we used the mixed continuous discrete architecture in Figure 1, where the architecture of φ consists of (784 , 400) linear layer, followed by a ReLU, and a (400 , 200) linear layer.
Revised: For these experiments, we used the mixed continuous discrete architecture in Figure 1, where the architecture of φ consists of (784 , 400) linear layer, followed by a ReLU, and a (400 , 200) linear layer.

Original: The discrete latent space consists of z d ∈ { 1 , ..., 10 } and its encoder φ d consists of a (200 , 100) linear layer, a ReLU, and a (100 , 10) linear layer.
Revised: The discrete latent space consists of z d ∈ { 1 , ..., 10 } and its encoder φ d consists of a (200 , 100) linear layer, a ReLU, and a (100 , 10) linear layer.

Original: The continuous latent space considers k = 10 , c = 20 in Figure 1, and its encoder φ c consists of a (200 , 100) linear layer, a (100 , 66)
Revised: The continuous latent space considers k = 10 , c = 20 in Figure 1, and its encoder φ c consists of a (200 , 100) linear layer, a (100 , 66)

Original: linear layer followed by ReLU and dropout and a (66 , 40) linear layer to estimate the mean and variance of c − dimensional Gaussian random variables z 1 , ..., z k .
Revised: linear layer followed by ReLU and dropout and a (66 , 40) linear layer to estimate the mean and variance of c − dimensional Gaussian random variables z 1 , ..., z k .

Original: w/o glasses glasses woman man woman man w/o smile smile w/o smile smile w/o smile smile w/o smile smile
Revised: w/o glasses glasses woman man woman man w/o smile smile w/o smile smile w/o smile smile w/o smile smile

Original: Following Kingma et al.
Revised: Following Kingma et al.

Original: (2014), we conducted a quantitive experiment with weak supervision on
Revised: (2014), we conducted a quantitive experiment with weak supervision on

Original: MNIST with 50 / 100 / 300 / 600 labeled examples out of the 50 , 000 training examples.
Revised: MNIST with 50 / 100 / 300 / 600 labeled examples out of the 50 , 000 training examples.

Original: For labeled examples, we set the perturbed label z (cid:15)θ + φ + γ + (cid:96) to be the true label.
Revised: For labeled examples, we set the perturbed label z (cid:15)θ + φ + γ + (cid:96) to be the true label.

Original: This is equivalent to using the indicator function over the space of correct predictions.
Revised: This is equivalent to using the indicator function over the space of correct predictions.

Original: A comparison of our method with
Revised: A comparison of our method with

Original: Gumbel-Softmax appears in Figure 3.
Revised: Gumbel-Softmax appears in Figure 3.

Original: We can see that direct loss minimization combined with
Revised: We can see that direct loss minimization combined with

Original: VAE improves the performance even with weak supervision, e.g., with only 50 examples direct loss minimization with VAE achieves better accuracy than semi-supervised Gumbel-Softmax VAE.
Revised: VAE improves the performance even with weak supervision, e.g., with only 50 examples direct loss minimization with VAE achieves better accuracy than semi-supervised Gumbel-Softmax VAE.

Original: We note that we cannot compare the objective function of both methods, as our objective considers direct loss minimization, however, we are able to compare the test log-loss bound, see Figure 3.
Revised: We note that we cannot compare the objective function of both methods, as our objective considers direct loss minimization, however, we are able to compare the test log-loss bound, see Figure 3.

Original: Supervision in generative models also greatly helps to control discrete semantics within images.
Revised: Supervision in generative models also greatly helps to control discrete semantics within images.

Original: We learn to generate images using k = 8 discrete attributes of the CelebA dataset (cf.
Revised: We learn to generate images using k = 8 discrete attributes of the CelebA dataset (cf.

Original: Liu et al.
Revised: Liu et al.

Original: while using our semi-supervised VAE.
Revised: while using our semi-supervised VAE.

Original: We applied our loss function over all training labels.
Revised: We applied our loss function over all training labels.

Original: For this task, we use convolutional layers for both the encoder and the decoder, except the last two layers of the continuous latent model which are linear layers that share parameters over the 8 possible representations of the image.
Revised: For this task, we use convolutional layers for both the encoder and the decoder, except the last two layers of the continuous latent model which are linear layers that share parameters over the 8 possible representations of the image.

Original: In Figure 4, we show generated images with discrete semantics turned on/off (with/without glasses, with/without smile, woman/man).
Revised: In Figure 4, we show generated images with discrete semantics turned on/off (with/without glasses, with/without smile, woman/man).

Original:  D ISCUSSION AND FUTURE WORK
Revised:  D ISCUSSION AND FUTURE WORK

Original: In this work, we use the Gumbel-Max trick to reparameterize discrete VAEs using the arg max prediction operator and develop a gradient estimation method that enables direct optimization of the resulting problem by propagating gradients through the non-differentiable arg max function.
Revised: In this work, we use the Gumbel-Max trick to reparameterize discrete VAEs using the arg max prediction operator and develop a gradient estimation method that enables direct optimization of the resulting problem by propagating gradients through the non-differentiable arg max function.

Original: We show that this approach compares favorably to state-of-the-art methods, and extend it to semisupervised learning and image attribute generation.
Revised: We show that this approach compares favorably to state-of-the-art methods, and extend it to semisupervised learning and image attribute generation.

Original: These results can be taken in a number of different directions.
Revised: These results can be taken in a number of different directions.

Original: Our gradient estimation is practically biased, while REINFORCE is an unbiased estimator.
Revised: Our gradient estimation is practically biased, while REINFORCE is an unbiased estimator.

Original: Our methods may beneﬁt from the REBAR/RELAX framework, which directs our biased gradients towards the unbiased gradient (Tucker et al., 2017).
Revised: Our methods may beneﬁt from the REBAR/RELAX framework, which directs our biased gradients towards the unbiased gradient (Tucker et al., 2017).

Original: There are also open problems when ﬁtting this approach to structured latent spaces (Mena et al., 2018; Jin et al., 2018). Such spaces better ﬁt to generate a complex scene, consisting of several objects and their interactions.
Revised: There are also open problems when ﬁtting this approach to structured latent spaces (Mena et al., 2018; Jin et al., 2018), as well as estimating its KL-divergence (Roeder et al., 2017).

Original: However, the direct optimization technique does not immediately apply in this setting, when using a linear number of Gumbel random variables.

Original: There are also optimization-related questions that arise from our work: the interplay of (cid:15) and the learning rate is unexplored and might be correlated.
Revised: There are also optimization-related questions that arise from our work: the interplay of (cid:15) and the learning rate is unexplored and might be correlated.

Original: The number of stochastic gradient steps, interleaving Gumbel perturbation with batch samples, might also beneﬁt from a rigorous investigation.
Revised: The number of stochastic gradient steps, interleaving Gumbel perturbation with batch samples, might also beneﬁt from a rigorous investigation.

Original: The direct optimization approach we present is general and may be applied beyond VAEs, including reinforcement learning and attention models.
Revised: The direct optimization approach we present is general and may be applied beyond VAEs, including reinforcement learning and attention models.

Original: Further investigation in this direction is required.
Revised: Further investigation in this direction is required.

Original: A G UMBEL -M AX PERTURBATION MODEL AND THE G IBBS DISTRIBUTION
Revised: A G UMBEL -M AX PERTURBATION MODEL AND THE G IBBS DISTRIBUTION

Original: Theorem 2.
Revised: Theorem 2.

Original: Gumbel & Lieblein (1954); Luce (1959); McFadden (1973) Let γ be a random function that associates random variable γ ( z ) for each z
Revised: Gumbel & Lieblein (1954); Luce (1959); McFadden (1973) Let γ be a random function that associates random variable γ ( z ) for each z

Original: = 1 , ..., k whose distribution follows the zero mean
Revised: = 1 , ..., k whose distribution follows the zero mean

Original: Gumbel distribution law, i.e., its probability density function is g ( t ) =
Revised: Gumbel distribution law, i.e., its probability density function is g ( t ) =

Original: e − ( t + c + e − ( t + c ) ) for the Euler constant c ≈ 0 .
Revised: e − ( t + c + e − ( t + c ) ) for the Euler constant c ≈ 0 .

Original: Then 
Revised: Then e

Original: [Equation]
Revised: [Equation]

Revised: ] , where z

Original: [Equation]
Revised: [Equation]

Original: z
Revised: +

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Let G ( t ) =
Revised: Let G ( t ) =

Original: e − e − ( t + c ) be the Gumbel cumulative distribution function.
Revised: e − e − ( t + c ) be the Gumbel cumulative distribution function.

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: Since g ( t ) = e − ( t + c ) G ( t ) it holds that
Revised: Since g ( t ) = e − ( t + c ) G ( t ) it holds that

Original: [Equation]
Revised: [Equation]

Original:  where 1 Z = (cid:82) e − ( t + c ) (cid:81)
Revised: (15) where 1 Z = (cid:82) e − ( t + c ) (cid:81)

Original: k
Revised: k

Original: ˆ z =1 G ( t − φ (ˆ z ))
Revised: ˆ z =1 G ( t − φ (ˆ z ))

Original: dt is independent of z .
Revised: dt is independent of z .

Original: Since P γ ∼ g
Revised: Since P γ ∼ g

Original: [ z = z φ + γ ] is a distribution then Z must equal to (cid:80)
Revised: [ z = z φ + γ ] is a distribution then Z must equal to (cid:80)

Original: k
Revised: k

Original: ˆ z =1 e φ ( x, ˆ z ) . w [Figure] glasses woman man woman man w/o smile smile w/o smile smile w/o smile smile w/o smile smile
Revised: ˆ z =1 e φ ( x, ˆ z ) w/o glasses glasses woman man woman man w/o smile smile w/o smile smile w/o smile smile w/o smile smile