Quasi Monte Carlo methods enable extremely low-dimensional deep generative models

TL;DR

QLVMs tackle the challenge of obtaining interpretable, extremely low-dimensional embeddings for high-dimensional data by replacing encoder-based variational inference with quasi-Monte Carlo marginal likelihood estimation. They sample latent codes from a fixed prior $p(\boldsymbol{z})$ using randomly shifted lattices (e.g., Fibonacci for 2D, Korobov for 3D) and map $\boldsymbol{z}$ through a decoder with a periodic first layer, allowing direct evaluation of $\log p_\theta(\boldsymbol{x})$ via a log-sum-exp bound. Empirically, 2D QLVMs outperform 2D VAEs and 2D IWAEs in marginal log-likelihood and reconstructions across several datasets, while enabling density estimation, clustering, and geodesic visualization in the latent space; however, performance gains diminish on high-complexity data and the method is computationally intensive and scale-limited for higher latent dimensions. Overall, QLVMs offer a transparent, data-generating approach to learning interpretable, low-dimensional embeddings, facilitating exploratory analyses that are difficult to validate in higher-dimensional latent spaces.

Abstract

This paper introduces quasi-Monte Carlo latent variable models (QLVMs): a class of deep generative models that are specialized for finding extremely low-dimensional and interpretable embeddings of high-dimensional datasets. Unlike standard approaches, which rely on a learned encoder and variational lower bounds, QLVMs directly approximate the marginal likelihood by randomized quasi-Monte Carlo integration. While this brute force approach has drawbacks in higher-dimensional spaces, we find that it excels in fitting one, two, and three dimensional deep latent variable models. Empirical results on a range of datasets show that QLVMs consistently outperform conventional variational autoencoders (VAEs) and importance weighted autoencoders (IWAEs) with matched latent dimensionality. The resulting embeddings enable transparent visualization and post hoc analyses such as nonparametric density estimation, clustering, and geodesic path computation, which are nontrivial to validate in higher-dimensional spaces. While our approach is compute-intensive and struggles to generate fine-scale details in complex datasets, it offers a compelling solution for applications prioritizing interpretability and latent space analysis.

Figures (19)

• Figure 1: (A) On each batch, latent samples are generated by uniformly shifting a lattice of $M$ points (Batch 1 grid and batch B grid denote examples of samples that may appear during training). (B) These samples are fed through a shared decoder, $f_\theta$, to produce reconstructed datapoints (e.g. MNIST digits). During training, a log-sum-exp ($\textrm{LSE}$) reduction operator is applied to a vector of log probabilities, $\log p = \log p(\boldsymbol{x}_i ~\vert~ \tilde{\boldsymbol{z}}_j)$, to calculate \ref{['eq:mc-lower-bound']} up to a constant. (C) To approximate the latent posterior, $p(\boldsymbol{z} ~\vert~ \boldsymbol{x}_i)$, the same vector of log probabilities is normalized to form a discrete approximation over the $m$ lattice points. This employs Bayes' rule as described in the main text.
• Figure 2: (A) Marginal log likelihood estimates on heldout test data from 2D models (orange) across three datasets (QMC estimate for QLVMs, ELBO for VAE and IWAE). The QLVM bound is higher than the ELBO and IWAE bound across all 2D models. Error bars indicate 1 standard deviation across 10 random seeds. (B) Reconstructions for a 2-D QLVM and VAE, with the QLVM displaying higher-quality reconstructions. (C) Samples from the prior of a 2-D QLVM, VAE, and IWAE. QLVMs show greater sample quality and sample diversity.
• Figure 3: (A) VAE and IWAE lower bounds on marginal log likelihood (orange dots) for 2D latent spaces vs. the QMC estimate based on VAE and IWAE decoders (blue dots) from \ref{['fig:model_performance']}. Black horizontal line indicates the 2D QLVM bound on the marginal log likelihood. The gap between the blue dots and the black line represents the additional benefit of training with QMC samples. (B) Empirical vs. variational posteriors for MNIST, Celeb-A, and zebra finch vocalizations. Heatmap indicates 99% of probability mass in the empirical posterior (estimated by a dense lattice), while the red dots and ellipses indicate means and $\pm$3 standard deviation of the variational posterior. This panel is best viewed on screen with magnification.
• Figure 4: Performance vs. computational cost curves for 2D QLVMs (red), VAEs (green), and IWAEs (purple) on MNIST, Celeb-A, and zebra finch datasets. Curves closer to the lower left quadrant of each plot indicate a more favorable tradeoff. See §\ref{['subsec:results-computational-cost']} and \ref{['sec:pareto_fronts']} for comparison details.
• Figure 5: (A) Latent embeddings of MNIST, colored by digit. (B) Aggregated posterior of MNIST with mean-shift centroids (red) overlaid. (C) Reconstructions of centroids in (B). (D) Smoothed Frobenius norm of decoder Jacobian on MNIST, with centroids from (B) overlaid. (E) Latent embeddings of gerbil vocalizations, colored by mean frequency of the vocalization. (F) Aggregated posterior of gerbil vocalizations, with mean-shift centroids (red) overlaid. (G) Reconstructions of centroids from (F). (H) Smoothed log Frobenius norm of decoder Jacobian on gerbil vocalizations.