Variational & Flow Models advanced 7 min read 6 flashcards

Importance-Weighted Bounds and the SNR Problem

How averaging importance weights inside the logarithm gives a strictly tighter bound on the evidence, and the result that this tighter bound degrades the encoder's gradient signal-to-noise ratio as the sample count grows.

The standard ELBO takes the expectation of a log. The importance-weighted bound takes the log of an expectation, which by Jensen is larger and therefore closer to \(\log p(x)\):

\[ \mathcal{L}_K = \mathbb{E}_{z_{1:K} \sim q}\left[\log \frac{1}{K}\sum_{k=1}^{K} \frac{p(x, z_k)}{q(z_k \mid x)}\right] \]

At \(K = 1\) this is the ordinary ELBO. As \(K\) grows the bound tightens monotonically and converges to \(\log p(x)\) (Burda, Grosse and Salakhutdinov, 2016, Importance Weighted Autoencoders, arXiv:1509.00519). On binarised MNIST with one stochastic layer, training with \(K = 50\) reached 84.78 nats against 86.76 for the single-sample VAE of the same architecture. For a decade this was taken as settled: tighter is better.

What the tighter bound actually changes

The IWAE objective relaxes the pressure on \(q\). A single-sample ELBO is maximised only when \(q\) matches the true posterior, so the encoder is forced to be accurate. With \(K\) samples, a few good draws can carry the estimate, so \(q\) only has to cover the posterior rather than match it. The implicit posterior being optimised is no longer \(q\) itself but a reweighted version of it, which is a more expressive object than the Gaussian you parameterised.

That is a real gain for the decoder. It is a loss for the encoder, and the reason is a variance argument rather than a modelling one.

Why the encoder's gradient degrades

Write the gradient estimator for the inference parameters \(\phi\) using \(M\) independent copies of the \(K\)-sample estimate. As \(K\) grows, the expected gradient with respect to \(\phi\) shrinks like \(1/K\), because the objective is becoming less sensitive to \(q\) by construction. Its standard deviation shrinks only like \(1/\sqrt{K}\). The signal-to-noise ratio is the ratio of those, so

\[ \mathrm{SNR}_\phi = O\!\left(\sqrt{M/K}\right), \qquad \mathrm{SNR}_\theta = O\!\left(\sqrt{MK}\right) \]

The generative parameters \(\theta\) get a cleaner signal as \(K\) rises; the inference parameters get a dirtier one, and in the limit the encoder's updates are pure noise (Rainforth et al., 2018, Tighter Variational Bounds are Not Necessarily Better, arXiv:1802.04537). The bound is tighter and the encoder is worse. Both statements are true at once, which is why the result was surprising.

The fixes

Three of them come from the same paper and differ only in which objective each network sees. PIWAE trains the encoder on a small-\(K\) bound while the decoder gets a large-\(K\) one. MIWAE fixes a total budget \(MK\) and spends more of it on independent copies than on samples per copy. CIWAE takes a convex combination of the ELBO and the IWAE bound. All three include IWAE as a special case, and all three trade bound tightness for gradient quality.

The sharper fix reparameterises the estimator instead. The doubly reparameterised gradient estimator rewrites the \(\phi\) gradient so the problematic score-function term cancels, restoring an SNR that improves with \(K\) rather than degrading (Tucker et al., 2019, Doubly Reparameterized Gradient Estimators for Monte Carlo Objectives, arXiv:1810.04152). It costs one extra backward pass through the encoder and is close to free otherwise, which makes it the default worth reaching for.

When it breaks

Tightness is not the objective you care about. If you want a representation, you want an accurate encoder, and the IWAE bound explicitly stops asking for one. Reach for it when density estimation is the goal and for evaluation, not when the latent is a product deliverable.

Evaluation and training are different questions. Using a large-\(K\) bound to report held-out likelihood is uncontroversial and standard, because no gradients flow. The SNR problem is a training-time phenomenon only.

The weights degenerate in high dimensions. Importance sampling suffers when the proposal is poor, and \(q\) in a large latent space usually is. The effective sample size can collapse to one or two out of fifty, at which point the extra \(K\) buys almost nothing while costing \(K\) decoder evaluations. Log the effective sample size before you raise \(K\).

Memory scales with \(K\). Every sample needs its own decoder forward pass held for the backward pass. \(K = 50\) on a large decoder is a fifty-fold activation-memory increase, and the usual response, gradient checkpointing, gives back much of the speed.

References and further reading

Every source this page cites, in the order it cites them. All of them open in a new tab.

  1. Burda, Grosse and Salakhutdinov, 2016, Importance Weighted Autoencoders, arXiv:1509.00519 arxiv.org
  2. Rainforth et al., 2018, Tighter Variational Bounds are Not Necessarily Better, arXiv:1802.04537 arxiv.org
  3. Tucker et al., 2019, Doubly Reparameterized Gradient Estimators for Monte Carlo Objectives, arXiv:1810.04152 arxiv.org
Check yourself

6 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track