Revisiting Reweighted Wake-Sleep for Models with Stochastic Control Flow
Stochastic control-flow models (SCFMs) are a class of generative models that involve branch- ing on choices from discrete random vari- ables. Amortized gradient-based learning ofSCFMs is challenging as most approaches tar- geting discrete variables rely on their contin- uous relaxations—which can be intractable inSCFMs, as branching on relaxations requires evaluating all (exponentially many) branch- ing paths. Tractable alternatives mainly com- bine REINFORCE with complex control-variate schemes to improve the variance of na ̈ıve esti- mators. Here, we revisit the reweighted wake- sleep (RWS)  algorithm, and through ex- tensive evaluations, show that it outperforms current state-of-the-art methods in learningSCFMs. Further, in contrast to the importance weighted autoencoder, we observe that RWSlearns better models and inference networks with increasing numbers of particles. Our re- sults suggest that RWS is a competitive, often preferable, alternative for learning SCFMs.