Reparameterization of variational auto-encoders is an effective method for reducing the variance of their gradient estimates. However, when the latent variables are discrete, a reparameterization is problematic due to discontinuities in the discrete space. In this work, we extend the direct loss minimization technique to discrete variational auto-encoders. We first reparameterize a discrete random variable using the arg max function of the Gumbel-Max perturbation model. We then use direct optimization to propagate gradients through the non-differentiable arg max using two perturbed arg max operations.
This paper proposes combining the Gumbel-max trick and "direct loss optimization" for variance reduction in VAEs with discrete latent variables. This is a natural combination (in hindsight), since the Gumbel-max trick turns sampling into non-differentiable optimization, and direct loss optimization provides a way to optimize the expected value of a non-differentiable loss. The paper is well-written for the most part and is backed by good experimental results. However it like some of the mathematical details and some of the exposition could be greatly improved.
I think there are several mistakes in the reasoning presented in the proof of Theorem 1 (see detailed comments below). Theorem 1 in the current paper seems to me to be a special case of Theorem 1 in (Song 2016), where the expectation over data is replaced by an expectation over the Gumbel variable gamma. If I've understood correctly, it seems like it would be more correct and concise to simply cite that paper with some explanatory comments.
The word "direct" occurs quite a lot in the paper. It sometimes seemed misplaced. For example for "The direct differentiation of the resulting expectation" in the introduction, in what sense is the differentiation direct, and what would non-direct differentiation be?
In section 3, that's not the meaning of the term "exponential family".
The re-use of theta and psi as both model parameters and the log probability density / distribution is unnecessarily confusing.
A small point, but in "the challenge in generative learning is to reparameterize and optimize (2)", the authors assume that q has analytic expression for the second KL term in (1). That's often the case but definitely not always. Also, even if this KL term has an analytic expression, it is not always better to use it (see Duvenaud "Sticking the landing...").
In (3), the usual notation is P(x = i) where x is the random variable and i is its possible value, whereas in (3) the random variable z^{\phi + \gamma} appears on the right of the equals sign.
The first paragraph of section 4.1 and (4) and (5) are just a simple application of the law of total expectation, and it would be simpler and clearer to state that.
"gradient of the decoder" should be "gradient of the decoder log probability" (or log prob density depending on preference). Similarly with "the decoder is a smooth function". The decoder is a conditional probability distribution (at least according to my understanding of conventional usage).
In the first equation in the proof of Theorem 1, it seems as though the authors are using the standard change of variables formula for integrals. However the new variable \hat{\gamma} depends on \hat{z} through \theta, so I don't see how it's valid to ignore the max in the way the present paper does. One way to see that something is wrong is the fact that the integrand on LHS has \hat{z} as a bound variable only, whereas the integrand on RHS has \hat{z} as both a bound variable (inside the max) and a free variable (since \theta depends on \hat{z}, though strangely that is not written in the equation). What is the value of \hat{z} used for \theta on RHS?
There's a missing [] after \partial_\epsilon in the third line of the paragraph starting "We turn to prove Equation (8)".
In the same line, I don't see why the two expectations are equal. It seems to me that the differentiation w.r.t. epsilon ignores the fact that changing epsilon occasionally changes z^{\epsilon \theta + \phi_v + \gamma} in a discontinuous way. The term being differentiated has both a continuous-in-epsilon component and a piecewise-constant-in-epsilon component, and the latter appears to have been ignored. While the gradient of a piecewise constant function is zero almost everywhere, the occasional large changes (which could be thought of as delta functions) still can make a large contribution to the overall expression once we take the expectation. To look at it another way, if the reasoning here is correct, why can't the same argument be used on the RHS of (8), first to take the derivative inside the expectation and subsequently to compute the derivative as zero, since the inner term is a piecewise constant function of v? Yet clearly the RHS of (8) is not always zero.
Around "However when we approach the limit, the variance of the estimate increases...", I think it would be extremely helpful to explain that for small epsilon, we occasionally obtain a large gradient (and otherwise zero), while for large epsilon we often obtain a moderate non-zero gradient. That gives some insight into the effect of epsilon, and why the variance is larger for small epsilon.
Any reason not to plot the bias in right Figure 1, which is ostensibly about the bias-variance trade-off?
I didn't follow the meaning of the diagram or caption for left Figure 1.
In (11), I wasn't sure whether S included the supervised examples or not (i.e. whether S_1 was disjoint from or a subset of S). If disjoint, shouldn't the KL term be included, or the expectation-over-gamma term be changed to use ground truth z? I guess I was unclear on the form of loss used for the supervised data, and unclear on the motivation for this choice.
In the last sentence of section 5.1, should "chain rule" be "variance reparameterization trick"?
In section 5.2, it would be helpful to mention what mean field means in terms of the variational distribution q (namely q(z | x) = \prod_i q(z_i | x) ). Also, the term "mean field" is not conventionally used for general distributions (such as the decoder here) as far as I'm aware, only for variational distributions. "Conditionally independent" might be clearer.
What does "for which we can approximate z^{...} efficiently" refer to?
In section 6.1, what is the annealing rate? Also, the minimal epsilon is set to 0.1. Is epsilon changed as training progresses according to some schedule?
"The main advantage of our framework is that it seamlessly integrates semi-supervised learning" seems like an overstatement. Wouldn't semi-supervised learning be relatively straightforward to incorporate into any form of VAE? And why not just use log p(x, z) for updating the decoder parameters and log q(z | x) for updating the encoder parameters?
How many labeled examples were used for the CelebA semi-supervised learning?
Some bibliography typos. For example, no capitalization throughout (e.g. "gumbel" instead of "Gumbel"). Also lots of arxiv preprints cited when published papers exist (e.g. Jang 2017 should be ICLR 2017 not arxiv preprint).