Revisiting Reweighted Wake-Sleep

Abstract:  Discrete latent-variable models, while applicable in a variety of settings, can often be difficult to learn. Sampling discrete latent variables can result in high-variance gradient estimators for two primary reasons: 1) branching on the samples within the model, and 2) the lack of a pathwise derivative for the samples. While current state-of-the-art methods employ control-variate schemes for the former and continuous-relaxation methods for the latter, their utility is limited by the complexities of implementing and training effective control-variate schemes and the necessity of evaluating (potentially exponentially) many branch paths in the model. Here, we revisit the Reweighted Wake Sleep (RWS; Bornschein and Bengio, 2015) algorithm, and through extensive evaluations, show that it circumvents both these issues, outperforming current state-of-the-art methods in learning discrete latent-variable models. Moreover, we observe that, unlike the Importance-weighted Autoencoder, RWS learns better models and inference networks with increasing numbers of particles, and that its benefits extend to continuous latent-variable models as well. Our results suggest that RWS is a competitive, often preferable, alternative for learning deep generative models.

Reviews and author responses:

Note that our claim is not based only on the GMM experiment. It is also backed up by results from training (i) a VAE with continuous latent variable on MNIST data (compared against IWAE since VIMCO is not needed) and (ii) the AIR model on moving MNIST data (compared against VIMCO; VQ-VAE not applicable).

I second the reviewer suggestion of real-world experiments. Improvements on toy data-sets like GMMs do not necessarily transfer over to real world data. And if the authors make the claim that  
"Our results suggest that RWS is a competitive, often preferable, alternative for learning deep generative models", then they should back it up with results that match or improve state-of-the-art generative models like VQ-VAE/VIMCO in bits/dim on large scale, real data-sets.

There are two reasons why we think RWS should be used based on solid (although simple) theoretical reasoning. First, in discrete latent variable models, we don’t need reparameterization, continuous relaxation or control variates. Second, in all models, RWS can be helpful in light of the “tighter variational bound is not better” effect. The GMM and AIR experiments support both points. The continuous VAE (on MNIST data) gives further evidence for the second point.

Why not compare to RBM and DVAE? We agree that even more evidence would be good. However:
- For RBMs, this is an entirely different class of models (the joint density p(z, x) can be evaluated only up to a normalizing constant, instead of directly), which is not learnable using RWS or other ELBO-maximizing approaches (learning RBMs requires contrastive divergence or similar). 
- For DVAE, in addition to being slightly different in focus due to branching, it is also orthogonal in another way. DVAE can be used in conjunction with IWAE to tighten the bound. We show through the continuous VAE experiment that RWS can help. It might be interesting to see whether RWS can be used on the continuously relaxed model defined in DVAE in order to improve DVAE further.

Why not test on real-world data? We agree that the transfer from synthetic to real-world data is difficult. Our paper is methodological and RWS is in its core a statistical method for maximizing the log marginal likelihood and minimizing the KL divergence from p to q. Whether a model works on a real dataset is not a function of the learning and inference algorithms, but rather the particular generative model and inference network. This is also true for other learning and inference algorithms (like VAE, IWAE, REBAR, RELAX, DVAE, etc.).

Main idea:
This paper studies a problem of the importance weighted autoencoder (IWAE) pointed out by  Rainforth 18, that is, tighter lower bounds arising from increasing the number of particles improve the learning of the generative model, but worsen the learning of the inference network. The authors show that the reweighted wake-sleep algorithm (RWS) doesn't suffer from this issue. Moreover, as an alternative to control variate scheme and reparameterization trick, RWS doesn't suffer from high variance gradients, thus it is particularly useful for discrete latent variable models.   
To support the claim, they conduct three experiments: 1) on ATTEND, INFER, REPEAT, a generative model with both discrete and continuous latent variables; 2) on MNIST with a continuous latent variable model; 3) on a synthetic GMM.

Clarity issues:
1. "branching" has been used many times, but AFAIK, this seems not a standard terminology. What do "branching on the samples", "conditional branching", "branching paths" mean?
2. zero-forcing failure mode and delta-WW: I find this part difficult to follow. For example, the following sentence 
"the inference network q(z|x) becomes the posterior for this model which, in this model, also has support at most {0, . . . , 9} for all x". 
However, this failure mode seems an interesting finding, and since delta-WW outperforms other methods, it deserves a better introduction. 

Questions:
1. In Fig 1 (right), how do you estimate KL(q(z|x) || p(z|x))?
2. In Sec 4.2, why do you say IWAE learns a better model only up to a point (K = 128) and suffers from diminishing returns afterwards?  
3. In Fig 4, why WS doesn't achieve a better performance when K increasing?

Experiments:
1. Since the motivating story is about discrete latent variable models, better baselines should be compared, e.g. RBM, DVAE, DVAE++, VQ-VAE etc. 
2. All experiments were on either on MNIST or synthetic data, at least one large scale experiment on discrete data should be made to verify the performance of RWS. 


" the resulting discreteness cannot be used for directing the control flow"

At least RBM doesn't need any continuous relaxation in training. 
To test on those discrete applications which can be continuously relaxed during training is also important. It offers us a better understanding of when RWS should be applied.  I don't think your reasons are valid to not compare with RBM, DVAE etc. 

I still insist on real dataset/tasks, e.g. semantic segmentation, since there is always a gap between synthetic world and real world.

This manuscript investigates the performance of Reweighted Wake-Sleep (RWS) framework for learning deep generative models with discrete latent variables. It gives a clear introduction to variational autoencoder based models for scenarios with discrete latent variables, including IWAE and also models based on continuous relaxations of discrete variables. The paper performs several experiments, which suggest that RWS is more appropriate for discrete latent variables than other methods such as IWAE. Especially, increasing the number of particles, unlike IWAE, always enhances the performance of RWS.

While this paper investigates an important problem, and also offers interesting observations, it lacks a rigorous analysis of why the RWS performance is consistently better than IWAE. More precisely, the propositions should be stated in more formal language and they should be accompanied with a minimal rigorous justification.

Clarity issues:

By stochastic branching we refer to the evaluation of generative models where discrete latent variables are used to select which part of the model is going to be evaluated next. For example, in AIR, this decides when the program halts, in GMM the cluster index decides which likelihood function is evaluated. Another example of this are probabilistic context-free grammars (PCFGs) where discrete variables are used to describe which production rule is used (for example https://arxiv.org/abs/1806.07832). This is in contrast with modeling approaches where the discrete latent variable is merely an input to a neural network that doesn’t distinguish it from a non-discrete latent variable since it does not explicitly use the discreteness to model distinct modes of the data. (also see our general comment)

We will clarify the “zero-forcing” failure mode and delta-WW in the updated manuscript.

Questions:

To estimate KL(q(z|x) || p(z|x)), we take the difference of the log likelihood estimated by a 5000-particle IWAE bound and the ELBO estimated by 5000 Monte Carlo samples.

The statement that “IWAE learns a better model only up to a point” is justified by the IWAE curve in the middle of Figure 2: the decreasing slope indicates that improvements in marginal log probability decrease with increasing numbers of particles. This is even more pronounced in Figure 1, where IWAE performance decreases for k > 10.

In Fig 4, WS actually does achieve better performance as K increases - the final value of the learning curve goes down, although only very slightly.

Experiments:

Regarding experiments, RBM/DVAE/++/# and VQ-VAE allow learning models with discrete latent variables in general; however, the resulting discreteness cannot be used for directing the control flow of a generative model (see also response to AnonReviewer3 and our general comment).
    - In the DVAE family of algorithms, learning in discrete latent variable models is achieved by a continuous relaxation. This prevents using these variables as hard branching conditions.
    - In the VQ-VAE algorithm, the discrete latent variable is explicitly designed to be used to select an embedding and it is deterministic. This limits the use of a discrete latent variable (cannot be used to model a cluster identity or stopping of a while loop). 

Even though we do not have experiments on large-scale real-world datasets, AIR is a non-trivial model, and using it can be seen as a large-scale experiment - taking several days (and several GPUs) to obtain results summarized in Figure 1. Similarly, to the best of our knowledge, ours is the first reported result of an MNIST model trained with IWAE with 512 particles.


The key formal justification is relatively straightforward: RWS, unlike IWAE, does not suffer from the “tighter bounds” problem. On the contrary, RWS uses self-normalized importance sampling to estimate the gradient with respect to \phi. Both the asymptotic bias and variance of a self-normalized importance sampling estimator decrease linearly in number of particles. This means that increasing number of particles improves our gradient estimator and thus the optimization procedure.

We will explain this in more detail in the updated manuscript.

Baselines: As set out in our overall response, we aim to show that RWS is a better choice for inference in models that have stochastic control flow, where the choice from the discrete latent variables matters explicitly. In our GMM example, the cluster identity is such a choice, and in AIR, the stopping condition for the loop is another such choice. The work done in DVAE++, DVAE#, and other such approaches do not really handle this general class of problems well---by typically requiring enumeration of all possible branches and choices.

GMM: We will include a more detailed description of how defensive sampling ameliorates issues discovered in the GMM experiments in the updated manuscript.

Theoretical Rigour: We will include a more comprehensive discussion of the theoretical basis of why RWS is better than IWAE in the updated manuscript. Briefly, the justification for why RWS does not suffer from the "tighter bounds" problem is due to RWS's use of self-normalised importance sampling to compute the gradient of proposal parameters---resulting in both the asymptotic bias and variance decreasing linearly with number of samples.

Empirical Rigour: Our experiments strongly support our hypotheses:
  a. Unlike IWAE, RWS performs better with more particles, both in terms of the generative model and inference network, and
  b. It allows for effective and easy application to models where the choice from the discrete random variables affects model expansion or computation---something that requires expensive enumeration with continuous relaxations, or extremely finicky and unreliable construction with control-variate methods.

We thank the reviewers for their time and for appreciating our arguments about why RWS is preferable to IWAE-based approaches as we increase number of particles.

We would like to re-emphasize an important implication from our paper: RWS is a good inference algorithm for models that have _stochastic control flow_: i.e. models where latent variables can be instantiated dynamically, and different branching paths explored, based on the choice from a random variable[4-7]. Note that this is orthogonal to a large number of recent work that employ discrete latent variables in deep generative models[8-10]---where either the discrete variable is transformed via continuous relaxations, or, marginalized out entirely. The crucial difference is that none of these approaches address models where the model execution itself is determined via the discrete choices (if-statements or for-loops), as opposed to simply passing it through to a neural network---what is done with the discrete choice matters.

This is best illustrated in the domain of universal probabilistic programs (like Pyro [1]) which can contain arbitrary continuous and discrete latent variables, and where latent variables are instantiated dynamically and defined by running the program (or generative model). Stochastic control flow is a feature of such models [2] and allows the definition of expressive models, with potentially infinite number of latent variables [3], as mentioned in the discussion.

Universal (or higher-order) probabilistic programs form the largest family of samplable distributions and thus are a powerful tool to model data. Amortized inference and model parameter learning in such probabilistic programs, however, is typically only done using variational methods in the VAE/IWAE family of algorithms (as summarized in our manuscript). We’re trying to say: RWS is a simple and often superior algorithm to use in this model family.

Our point about probabilistic programs, hard selection and stochastic branching is illustrated by our choice of experiments (GMM and AIR). However, we will more strongly emphasize this point in the updated manuscript.

[1] http://pyro.ai/
[2] http://pyro.ai/examples/intro_part_i.html#Universality:-Stochastic-Recursion,-Higher-order-Stochastic-Functions,-and-Random-Control-Flow
[3] Quote from [2]
    "For example, we can construct recursive functions that terminate their recursion
     nondeterministically, provided we take care to pass pyro.sample unique sample names
     whenever it’s called."
[4] GMM models in this work
[5] Tree-structured latent variables in https://arxiv.org/abs/1806.07832
[6] Memory-based models in https://arxiv.org/abs/1709.07116
[7] AIR-like models in this work and https://arxiv.org/abs/1806.01794
[8] DVAE++ - https://arxiv.org/abs/1802.04920
[9] DVAE#  - https://arxiv.org/abs/1805.07445
[10] VQ-VAE - https://arxiv.org/abs/1711.00937


This paper conducts an extensive set of experiments on RWS and compares it against a set of benchmarks such as GMM and IWAE. The main contribution of the paper is the fact revealed by these experiments, that RWS learns better models and inference networks with increasing numbers of particles, and that its benefits extend to continuous latent variable models as well. The performance of RWS will increase significantly if we increase the number of particles. 

The experimental part is written in an inspiring way, and I enjoyed reading it. However, there should be stronger baselines incorporated. for example, https://arxiv.org/abs/1805.07445. Also, I think the authors could try to emphasize more on the shortcomings of RWS discovered by the GMM experiments, and how defensive importance sampling fixes it. There are several other parts in the paper that indicates interesting facts, diving deeper into it could possibly lead to more interesting findings.

In all, I would consider these comparison results important to be somewhere in the literature, but because its lack of rigorous analysis and explanation for the observations, I personally think these observations alone are not novel enough to be an ICLR paper. 


Meta-review: The paper presents a well conducted empirical study of the Reweighted Wake Sleep (RWS) algorithm (Bornschein and Bengio, 2015). It shows that it performs consistently better than alternatives such as Importance Weighted Autoencoder (IWAE) for the hard problem of learning deep generative models with discrete latent variables acting as a stochastic control flow. 
The work is well-written and extracts valuable insights supported by empirical observations: in particular the fact that increasing the number of particles improves learning in RWS but hurts in IWAE, and the fact that RWS can also be successfully applied to continuous variables.
The reviewers and AC note the following weaknesses of the work as it currently stands:  a) it is almost exclusively empirical and while reasonable explanations are argued, it does not provide a formal theoretical analysis justifying the observed behaviour b) experiments are limited to MNIST and synthetic data, confirmation of the findings on larger-scale real-world data and model would provide a more complete and convincing evidence. 
The paper should be made stronger on at least one (and ideally both) of these accounts.



Decision: rejected-papers