[Paper Review] Amortized Inference Regularization

Rui Shu, Hung Bui+3 more|arXiv (Cornell University)|May 23, 2018

Gaussian Processes and Bayesian Inference33 references48 citations

TL;DR

The paper introduces amortized inference regularization (AIR) to constrain the inference model in VAEs, showing that smoothing the amortized posterior can improve test-set log-likelihood and reduce inference/generative gaps, and extends AIR to IWAE.

ABSTRACT

The variational autoencoder (VAE) is a popular model for density estimation and representation learning. Canonically, the variational principle suggests to prefer an expressive inference model so that the variational approximation is accurate. However, it is often overlooked that an overly-expressive inference model can be detrimental to the test set performance of both the amortized posterior approximator and, more importantly, the generative density estimator. In this paper, we leverage the fact that VAEs rely on amortized inference and propose techniques for amortized inference regularization (AIR) that control the smoothness of the inference model. We demonstrate that, by applying AIR, it is possible to improve VAE generalization on both inference and generative performance. Our paper challenges the belief that amortized inference is simply a mechanism for approximating maximum likelihood training and illustrates that regularization of the amortization family provides a new direction for understanding and improving generalization in VAEs.

Motivation & Objective

• Motivate regularizing the amortized inference to improve generalization in VAEs rather than solely expanding the inference family.
• Propose mechanisms (denoising and weight normalization) to enforce smoothness in the amortized inference model.
• Analyze theoretical implications of AIR on the posterior and generator.
• Demonstrate AIR effectiveness on VAE and IWAE across multiple datasets.

Proposed method

• Reinterpret the ELBO as a data-dependent regularized maximum likelihood objective and introduce AIR as a capacity-constrained inference family.
• Propose two AIR instantiations: (1) Denoising Variational Autoencoder (DVAE) that injects Gaussian noise into the encoder during regularization; (2) Weight-Normalized Inference (WNI) that enforces smoothness via weight normalization.
• Show that the optimal DVAE inference model converges to a kernel-regression form and that the regularization strength is controlled by the noise level σ (and via a convex combination with standard ELBO).
• Extend AIR to IWAE by defining a regularizer R_k that accounts for multiple importance samples and show how AIR interacts with k to modulate regularization.
• Provide theoretical results (Lemmas/Theorems) on smoothness and regularization strength and present empirical results across multiple datasets.

Experimental results

Research questions

• RQ1 Does restricting the capacity of the amortized inference model (via AIR) improve test-set log-likelihood and reduce inference and generator gaps in VAEs?
• RQ2 How do DVAE and weight-normalized encoders compare to standard VAE in terms of test performance and regularization strength?
• RQ3 How does AIR interact with importance-weighted autoencoders (IWAE) and the number of importance samples k?
• RQ4 What theoretical insights describe how a smoothed amortized inference affects the generative model?
• RQ5 Can AIR provide practical, robust regularization across datasets with varying data regimes?

Key findings

• AIR improves test-set log-likelihood and reduces the inference gap on MNIST, OMNIGLOT, and Caltech Silhouettes compared to standard VAE.
• DVAE and WNI-VAE generally outperform VAE on all three datasets across multiple metrics (negative log-likelihood, amortized ELBO, inference gap).
• Increasing the regularization strength via σ (DVAE) or H (WNI) transitions models from overfitting to underfitting, with well-chosen values yielding 0.5–1.0 nat improvements in test log-likelihood on some datasets.
• In IWAE settings, DIWAE and WNI-IWAE consistently outperform IWAE with no AIR in terms of test log-likelihood and often show reduced inference gaps.
• AIR’s regularization effect attenuates as k increases in IWAE, but remains beneficial for practical k (e.g., k=8, k=64), with higher k mitigating underfitting caused by strong AIR.
• Theoretical propositions show that a smoother amortized inference leads to a convex combination weighting of training samples in the generator, suggesting a natural regularization mechanism for the generator.