Amortised Inference and the Amortisation Gap

Classical variational inference fits a separate set of variational parameters for every data point. A thousand images means a thousand little optimisation problems, each run to convergence. That does not scale to a training loop, so the variational autoencoder replaced it with a single neural network that maps \(x\) straight to the parameters of \(q(z \mid x)\) in one forward pass (Kingma and Welling, 2014, Auto-Encoding Variational Bayes, arXiv:1312.6114). The cost of inference stops growing with the dataset. The cost is that the encoder is now a function approximator, and function approximators make errors.

Two gaps, not one

For a single data point the difference between the true log evidence and the ELBO is exactly \(\mathrm{KL}(q(z \mid x) \,\|\, p(z \mid x))\). That total inference gap splits cleanly in two (Cremer, Li and Duvenaud, 2018, Inference Suboptimality in Variational Autoencoders, arXiv:1801.03558):

Here \(q^{*}\) is the best member of the variational family for this particular \(x\), found by direct optimisation. The approximation gap is the price of the family: a diagonal Gaussian cannot bend into a curved or multimodal posterior no matter how well you fit it. The amortisation gap is the price of the encoder: even inside the family, one shared network has to produce good parameters for every point it will ever see, and it will not.

The decomposition matters because the two gaps call for opposite fixes. A large approximation gap says make the family richer, with a normalising flow on top of the Gaussian (Rezende and Mohamed, 2015, Variational Inference with Normalizing Flows, arXiv:1505.05770). A large amortisation gap says the family is fine and the encoder is not finding its best member. Cremer and colleagues found the second case more common than the field had assumed, and also found that a richer family shrinks the amortisation gap as a side effect, because a flexible posterior is a more forgiving target to hit approximately.

Measuring it in practice

The measurement is a diagnostic you can run on any trained VAE, and it takes a few minutes:

1. Freeze the decoder. Take a held-out batch.
2. Record the ELBO using the encoder's output as-is. Call it \(\mathcal{L}[q_\phi]\).
3. For each \(x\) in the batch, initialise \((\mu, \log \sigma)\) from the encoder and run a few hundred steps of gradient ascent on the ELBO with respect to those parameters alone. Call the result \(\mathcal{L}[q^{*}]\).
4. Estimate \(\log p(x)\) with a large-sample importance-weighted bound, typically 5,000 samples.

Step 3 minus step 2 is the amortisation gap. Step 4 minus step 3 is the approximation gap. If the first number is large, your encoder is the bottleneck, and no amount of decoder capacity will help.

Closing the gap without giving up speed

The obvious repair is to stop amortising, which is also the expensive one. The useful middle ground keeps the encoder as an initialiser and adds a short refinement. Semi-amortised VAEs run a handful of stochastic variational inference steps on the encoder's output and backpropagate through those steps, so the encoder learns to produce parameters that are a good starting point rather than a good answer (Kim et al., 2018, Semi-Amortized Variational Autoencoders, arXiv:1802.02550). Iterative inference models go further and learn the refinement itself, feeding the current estimate and its ELBO gradient into a network that emits the update (Marino, Yue and Mandt, 2018, Iterative Amortized Inference, arXiv:1807.09356). Both buy accuracy with forward passes.

When it breaks

The gap is invisible in the training curve. A model with a large amortisation gap still shows a smoothly decreasing loss. The loss is the ELBO, and the ELBO is exactly the quantity the gap is hidden inside. You only see it by running the refinement above.

Refinement at training time changes what the encoder learns. Semi-amortised training requires differentiating through an inner optimisation loop, which means second-order terms, higher memory, and a training step several times more expensive. Teams often approximate those terms away, and the approximation quietly reintroduces some of the gap.

A refined posterior at test time is not free. Neural compression systems exploit this deliberately, spending encoder-side optimisation per image because the bitstream is what matters and encode time is not, but the same trade is a bad one for an interactive system.

Amortisation interacts with collapse. An encoder that cannot represent a sharp posterior for an unusual \(x\) will produce something close to the prior for it, which looks exactly like posterior collapse on that example and gets misdiagnosed as a decoder problem. Check the amortisation gap before reaching for KL annealing.