Amortized Variational Inference: When and Why?

TL;DR

This paper investigates the conditions under which amortized variational inference (A-VI) can match the optimal solutions of factorized variational inference (F-VI), providing theoretical insights and practical examples across different models.

Contribution

It derives necessary and sufficient conditions for A-VI to attain F-VI's optimal solution, and characterizes models where the amortization gap can or cannot be closed.

Findings

A-VI can match F-VI in simple hierarchical models.

Conditions for closing the amortization gap are identified and verified.

Examples include hidden Markov models where the gap cannot be closed.

Abstract

In a probabilistic latent variable model, factorized (or mean-field) variational inference (F-VI) fits a separate parametric distribution for each latent variable. Amortized variational inference (A-VI) instead learns a common inference function, which maps each observation to its corresponding latent variable's approximate posterior. Typically, A-VI is used as a step in the training of variational autoencoders, however it stands to reason that A-VI could also be used as a general alternative to F-VI. In this paper we study when and why A-VI can be used for approximate Bayesian inference. We derive conditions on a latent variable model which are necessary, sufficient, and verifiable under which A-VI can attain F-VI's optimal solution, thereby closing the amortization gap. We prove these conditions are uniquely verified by simple hierarchical models, a broad class that encompasses many models in machine learning. We then show, on a broader class of models, how to expand the domain of AVI's inference function to improve its solution, and we provide examples, e.g. hidden Markov models, where the amortization gap cannot be closed.