Implicit Reparameterization Gradients
By providing a simple and efficient way of computing low-variance gradients of continuous random variables, the reparameterization trick has become the technique of choice for training a variety of latent variable models. However, it is not applicable to a number of important continuous distributions. We introduce an alternative approach to computing reparameterization gradients based on implicit differentiation and demonstrate its broader applicability by applying it to Gamma, Beta, Dirichlet, and von Mises distributions, which cannot be used with the classic reparameterization trick. Our experiments show that the proposed approach is faster and more accurate than the existing gradient estimators for these distributions.
• We develop implicit reparameterization gradients that provide unbiased estimators for continuous distributions with numerically tractable CDFs. This allows many other important distributions to be used as easily as the Normal distribution in stochastic computation graphs.
• We show that the proposed gradients are both faster and more accurate than alternative approaches.
• We demonstrate that our method can outperform existing stochastic variational methods at training the Latent Dirichlet Allocation topic model in a black-box fashion using amortized inference.
• We use implicit reparameterization gradients to train VAEs with Gamma, Beta, and von Mises latent variables instead of the usual Normal variables, leading to latent spaces with interesting alternative topologies.
EDITED AFTER REBUTTAL:  I appreciate the authors' effort to make the notation consistent and improve the experiments by taking into account the computational complexity.  There are a couple of comments that were not addressed in the rebuttal, but I still recommend acceptance. For completeness, I am listing these comments below: + Experiments without optimizing the hyperparameters in the LDA experiment. + Experiments to showcase truncated distributions. + Comparison / Discussion about the approach of Roeder et al. for mixtures (they also integrate out the discrete component).  ------------------  This paper proposes a method to reparameterize distribution that are considered "non-reparameterizable", such as the Gamma or vonMises, as well as truncated distributions and mixtures. The proposed approach is compared to other baselines such as RSVI and is shown to outperform them.  The paper addresses a relevant problem and the proposed solution is clear and elegant.  My first comment is that this method has already been proposed in the paper "Pathwise Derivatives Beyond the Reparameterization Trick" [Jankowiak & Obermeyer, ICML 2018]. However, I still think that there is value in the Implicit Reparameterization Gradients, for three reasons. First, this paper approximates the CDF of the gamma without resorting to Taylor approximations. Second, this paper has more experiments, including experiments with the vonMises distribution. Third, accepted ICML papers where not available at the time of the NIPS deadline, so [Jankowiak & Obermeyer 2018] should not be taken into account. That said, the citation should be included.  My two main criticisms about this paper are: 1. The notation is not consistent. 2. The experiments can be improved; e.g., by showing the evolution of the marginal lower bound.  1. The notation in the paper is not consistent, please revise and improve it. For example, the parameterization of a distribution (say, q) in terms of some parameters (say, phi) is sometimes written as q(z|phi) and sometimes as q_phi(z). Please use one. More comments on the notation: + Eq. 9: Replace u with epsilon. Variable u has not been defined anywhere else. + In the mixture example (line 116), it should be q_phi(z) instead of p(z|phi). + In the LDA experiment, alpha is both the hyperparameter of the Dirichlet (line 214) and the topics (Fig 2a). Please remove the "alpha" in Fig 2a.  2. Regarding experiments: + In both the LDA and VAE experiments, can you show the evolution of the variational lower bound w.r.t. wall-clock time? The paper states that each step RSVI step is 15% slower, but the information about convergence of each method is missing.  + In section 6.1, the paper states that "The results presented on Fig. 1 show that the implicit gradient is faster and has lower variance than RSVI". In Fig. 1, however, only the computational time for the vonMises distribution is reported, but not for the Dirichlet. Is it also faster in the Dirichlet case? Does that hold for any value of the shape augmentation? I think this should be added to the paper. + In the LDA experiment, I think it would be better not to optimize the hyperparameters, because the goal is to assess how good the different variational approximations are for a fixed model. If the model is not fixed (by fitting alpha), then the comparisons get more blurry. That said, I wouldn't expect any significant difference with the results reported in the paper, but it would be nice to check that.  Finally, I also have two more comments. The paper explains how to reparameterize mixtures and truncated distributions, but it does not include experiments with those. It would also be nice to see an experiment with truncated distributions. In addition, mixtures can be reparameterized by following the approach of "sticking the landing" [Roeder et al., 2017]; is there any reason why either method should be preferred?