Variational auto-encoders (VAEs) offer a tractable approach when performing approximate inference in otherwise intractable generative models. However, standard VAEs often produce latent codes that are disperse and lack interpretability, thus making the resulting representations unsuitable for auxiliary tasks (e.g. classification) and human interpretation. We address these issues by merging ideas from variational auto-encoders and sparse coding, and propose to explicitly model sparsity in the latent space of a VAE with a Spike and Slab prior distribution. We derive the evidence lower bound using a discrete mixture recognition function thereby making approximate posterior inference as computational efficient as in the standard VAE case. With the new approach, we are able to infer truly sparse representations with generally intractable non-linear probabilistic models. We show that these sparse representations are advantageous over standard VAE representations on two benchmark classification tasks (MNIST and Fashion-MNIST) by demonstrating improved classification accuracy and significantly increased robustness to the number of latent dimensions. Furthermore, we demonstrate qualitatively that the sparse elements capture subjectively understandable sources of variation.
This paper proposes an extension of VAEs with sparse priors and posteriors to learn sparse interpretable representations. Training is made tractable by computing the analytic KL for spike and slab distributions, and using a continuous relaxation for the spike variable. The technique is evaluated on MNIST, Fashion MNIST, and CelebA where it learns sparse representations with reasonable log-likelihood compared to Gaussian priors/posteriors, but improved classification accuracy and interpretability of the representation.
While this paper is clear and well written, the novelty of the approach is limited. In particular, this is a straightforward application of vanilla VAEs with a different prior/posterior. The authors missed a bunch of related work and their main theoretical contributions are known in the literature (KL for spike and slab distributions, effective continuous relaxations for Bernoulli variables). The experiments are interesting but the authors should compare to more baselines with alternative priors (e.g. stick breaking VAEs, VampPrior, epitomic VAEs, discrete VAEs).
Strengths
+ Well written, clear, and self-contained paper. Figures are nice and polished.
+ Thorough experiments studying the effect of sparsity on the representation
Weaknesses
- No discussion/comparison to other VAE approaches that incorporate sparsity into the latents: Eptimoic VAEs (2017), discrete VAEs with binary or categorical latents are sparse (see: Discrete VAEs, Concrete/Gumbel-Softmax, VQ-VAE, output-interpretable VAEs), stick breaking VAEs, structured VAEs for the Beta-Bernoulli process (Singh, Ling, et al., 2017). Missing citation to foundational work on sparse coding from Olshausen and Field (1996). 
- Lack of novelty: The analytic KL term for spike and slab priors has been derived before in Discrete VAEs (Rolfe, 2017) and in work on weight uncertainty (Yarin Gal's thesis, Blundell et al. 2016). Continous relaxations like the one used for the spike variable has been presented in earlier work (Concrete distributon, Gumbel-Softmax, Discrete VAEs).
Minor comments:
- Eq. 1, shape for B should be MxJ
- Cite Rezende & Mohamed for VAEs along w/ Kingma & Welling
- Definition of VAE is overly-restrictive. Typically a VAE is the combo of variational inference with an amortized inference network (and optionally reparameterization gradients). Saying that VAE implies Gaussian prior and Gaussian posterior is far too restrictive.
- VLB is a non-standard acronym, use ELBO for evidence lower bound
- I'm surprised that VAEs perform so poorly as latent dim increases. I'd expect it to just prune latent dimensions. Do you have an explanation for why performance drops for VAEs? Are they overfitting?
- VAEs with Gaussian p(x|z) are typically harder to train and more sensitive to hyperparameters than Bernoulli p(x|z). Could you repeat your experiments using the more common binarized MNIST so that numbers are comparable to prior work?
- If the goal is to learn representations with high information, then beta-VAEs or InfoVAEs should be compared (see analysis in Alemi et al., 2017). The number of dimensions may matter less for classification than the rate of the VAE. To analyze this further, you could plot the rate (KL(q(z|x) || p(z)) vs. the classification accuracy for all your models.
- Fig 4: consider adding in plots of continuous interpolation of the latent dimension (as in beta-VAE, TC-VAE, etc.)
- Would be interested to see how much class information is stored in the value vs. the pattern of non-zeroes in the latent representation (as done in Understanding Locally Competitive networks from Srivasta et al. 2014).
- Not at all expected as this came out after your submission, but would be nice to compare to a similar recent paper: https://www.biorxiv.org/content/early/2018/08/23/399246