Towards Model-Agnostic Posterior Approximation for Fast and Accurate Variational Autoencoders

TL;DR

The paper tackles the inefficiency and instability of joint VAE training by proposing a model-agnostic posterior approximation (MAPA) grounded in an empiricalized model. MAPA yields a deterministic, index-based posterior that is independent of the true prior and likelihood, making inference faster and robust to non-identifiability. A proof-of-concept MAPA-based lower bound for the empiricalized model enables training with fewer forward passes while maintaining or improving density estimation on synthetic data. Empirical results show MAPA can capture posterior trends and reduce computation relative to standard VAE/IWAE baselines, and the authors outline a scalable roadmap for high-dimensional data and broader applicability.

Abstract

Inference for Variational Autoencoders (VAEs) consists of learning two models: (1) a generative model, which transforms a simple distribution over a latent space into the distribution over observed data, and (2) an inference model, which approximates the posterior of the latent codes given data. The two components are learned jointly via a lower bound to the generative model's log marginal likelihood. In early phases of joint training, the inference model poorly approximates the latent code posteriors. Recent work showed that this leads optimization to get stuck in local optima, negatively impacting the learned generative model. As such, recent work suggests ensuring a high-quality inference model via iterative training: maximizing the objective function relative to the inference model before every update to the generative model. Unfortunately, iterative training is inefficient, requiring heuristic criteria for reverting from iterative to joint training for speed. Here, we suggest an inference method that trains the generative and inference models independently. It approximates the posterior of the true model a priori; fixing this posterior approximation, we then maximize the lower bound relative to only the generative model. By conventional wisdom, this approach should rely on the true prior and likelihood of the true model to approximate its posterior (which are unknown). However, we show that we can compute a deterministic, model-agnostic posterior approximation (MAPA) of the true model's posterior. We then use MAPA to develop a proof-of-concept inference method. We present preliminary results on low-dimensional synthetic data that (1) MAPA captures the trend of the true posterior, and (2) our MAPA-based inference performs better density estimation with less computation than baselines. Lastly, we present a roadmap for scaling the MAPA-based inference method to high-dimensional data.

Figures (13)

• Figure 1: Intuition behind MAPA: nearby points score highly under each other's posteriors---distant points do not. Bottom row: four samples drawn from $p_z(\cdot) = \mathcal{U}[0, 1]$. Top row: the generative process (i.e. $z_i$, $f_\theta(z_i)$, $\epsilon_i$, and the resultant $x_i$), visualized on-top of the joint distribution of the observed and latent variables (gray contours) for two different functions that yield the same marginal distribution (in blue). In "Variant 1," $f_\theta(z) = (0.5 - z)^2$, and in "Variant 2," $f_\theta(z) = 0.25 \cdot z^2$yacoby2020failure. The posteriors of each of the four observations $x_i$ are visualized in the green density plots.
• Figure 2: (a) MAPA is more accurate than baselines; it better matches the true data distribution (lower test KL) than all baselines (VAE, IWAE). (b) MAPA is more efficient; it requires fewer NN-passes per training step. Details and full results in \ref{['sec:ll-results', 'sec:num-passes']}.
• Figure 3: MAPA captures trend of ground-truth posterior. Each panel compares the log-posteriors given a randomly-chosen $x$ vs. $z^\text{GT}$. The black, red, and blue represent the true posterior of the original model, the true posterior of the empiricalized model, and the MAPA, respectively. Details in \ref{['sec:mapa-trends-inline']} and additional plots in \ref{['sec:mapa-trends']}.
• Figure 4: MAPA is robust to model non-identifiability on "Absolute-Value" Example. Top-row: Two generative models with different $f_\theta(\cdot)$ that yield the same $p_x(\cdot; \theta)$. Under them, the left and right columns compare the MAPA approximation to the ground-truth posteriors for each of the two models, respectively, on the same $x$'s. MAPA captures the trend in both equally well.
• Figure 5: MAPA is robust to model non-identifiability on "Circle" Example. Top-row: Two generative models with different $f_\theta(\cdot)$ that yield the same $p_x(\cdot; \theta)$. Under them, the left and right columns compare the MAPA approximation to the ground-truth posteriors for each of the two models, respectively, on the same $x$'s. MAPA captures the trend in both equally well.