DVAE#: Discrete Variational Autoencoders with Relaxed Boltzmann Priors
Boltzmann machines are powerful distributions that have been shown to be an effective prior over binary latent variables in variational autoencoders (VAEs). However, previous methods for training discrete VAEs have used the evidence lower bound and not the tighter importance-weighted bound. We propose two approaches for relaxing Boltzmann machines to continuous distributions that permit training with importance-weighted bounds. These relaxations are based on generalized overlapping transformations and the Gaussian integral trick. Experiments on the MNIST and OMNIGLOT datasets show that these relaxations outperform previous discrete VAEs with Boltzmann priors. An implementation which reproduces these results is available at https://github.com/QuadrantAI/dvae.
# Response to author feedback  Having read the authors' response I am happy with my original assessment that this would be a good addition to the conference and vote for including this submission. The authors' discussion of the potential reasons for the relatively poor performance of the Gaussian integral trick method in their feedback in response to a comment in my review seem reasonable and from my reading they also adequately address the points raised by the other reviewers.  ---  # Summary and relation to previous work  This submission builds upon several previous pieces of work proposing methods for including discrete latent variables in variational autoencoder (VAE) models (Kingma and Welling, 2014). A key aspect of these works is retaining the ability to train the models with low-variance reparameterisation trick based gradient estimates of the variational objective by relaxing the discrete latent variables with associated continuous valued variables. Of particular significance to this submission are the discrete VAE (dVAE) (Rolfe, 2016) and dVAE++ (Vahdat et al., 2018) models which use a Boltzmann machine (BM) prior on the discrete latent variables and construct a differentiable proxy variational objective by introducing continuous variables zeta corresponding to relaxations of the discrete variables $z$, with $\zeta$ depending on $z$ via a *smoothing* conditional distribution $r(\zeta | z)$. The generative process in the decoder model is specified such that generated outputs $x$ are conditionally independent of the discrete variables $z$ given the continuous variables $\zeta$.  An issue identified with the (differentiable proxy) variational objective used in both the dVAE and dVAE++ approaches is that it is not amenable to being formulated as an importance-weighted bound, with importance-weighted objectives for continuous VAE models having been found to give significant improvements in training performance (Burda et al., 2015). In this submission the authors suggest an alternative dVAE formulation they term dVAE# which is able to use an importance weighted objective. This requires evaluating (/estimating) the log marginal density on the $\zeta$ variables. In general evaluation of the marginal density on $\zeta$ is intractable as it requires marginalising $z$ from joint density on $z$ and $\zeta$, with the exhaustive summation over all $z$ configurations in general intractable to compute.   The authors main contribution is to propose two alternative approaches to overcoming this issue. In the first the smoothing conditional density $r(\zeta | z)$ is assumed to take the same form  as the 'overlapping smoothing transformations' previously proposed in the dVAE++ paper. In this case as the marginalisation over $z$ is intractable, the authors propose fitting a mean field approximation to the BM prior on $z$ and using this factorial approximation to approximate the marginalisation over $z$ to compute approximations to the log marginal density on $\zeta$ and its gradient. The second suggestion is to use a particularly tractable form of continuous relaxation proposed previously for sampling from BM models (Zhang et al., 2012), termed the  'Gaussian integral trick', which allows exact marginalisation over $z$, at the expense over a more complex geometry resulting marginal distribution on the $\zeta$ variables (and so difficulties in choosing an appropriate variational approximation).  The authors also propose generalisations of the 'overlapping smoothing transformations' previously employed in the dVAE++ paper exploit an idea for applying the reparameterization trick to mixture distributions (Graves, 2016) to allow more general forms for the transformations while still remaining reparameterizable, with in particular several transformations which it is suggested will give lower variance gradient estimates being proposed. The proposed dVAE# approach is empirically validated to generally give significant improvements over the previous dVAE and dVAE++ approaches in terms of test set log likelihoods on two benchmark binarized image generative model datasets (MNIST and OMNIGLOT) across a range of different model architectures for each method. The benefit of using an importance weighted objective is validated by a generally increased performance as more importance samples are used and the mean-field approach using a particular 'power-function transformation' instance of their generalised smoothing transformation found to generally give better performance than the other suggested alternatives.  # Evaluation  This submission mainly proposes fairly incremental extensions to the existing dVAE and dVAE++ approaches, however the suggested improvements are motivated well and explained clearly, and the extensive experiments appear to show consistent significant test-set log-likelihood gains for a particular instance of the proposed approach (power-function transform with mean-field approximation) compared to competing dVAE and dVAE++ approaches on a wide range of architectures and across two different datasets. Further the inclusion of the negative results for their being significant gains to using what would seem a perhaps natural alternative of a Gaussian integral trick relaxation, given the ability to sidestep the need to introduce a mean-field approximation in this case, is also an important and interesting empirical contribution. It would be helpful however to have more discussion of potential reasons why the Gaussian integral approach seems to generally not perform as well. Overall though the originality of the proposal is not massive, the quality of the communication of the ideas present and relatively strong empirical validation of the proposals mean I feel the submission would be a worthwhile inclusion in the conference.  # Minor comments / typos / suggestions:  Equation 5: How to get from the second to third line is not necessarily immediately obvious - might be worth noting this relies (I think) on the linearity of the gradient of the augmented energy function with respect to z.  L154: 'covariance matrix becomes diagonally dominant' -> I think this should be precision not covariance matrix?  Digit grouping: in a few places commas are used for thousands digit grouping - generally a thin space should be preferred as a comma can be ambiguous as to whether a decimal point is intended in some locales.  Figures 1 and 2: plots difficult to read at print size due to light weight of text / lines - it would help to plot with fonts, line weights etc. specifically set for small figure size rather than what looks like downscaling a larger plot.  References: several references e.g. [3], [4]. [8], [9], [21] list arXiv pre-print versions when published versions are available. More of a nitpick, but it would also be nicer if all references consistently used initials or full first names. Hamiltonian in reference [41] should be capitalised.