攤銷推理：為何訓練網絡推理勝過每次執行最佳化

Variational inference has a scaling problem. For every new data point, it asks you to solve a fresh optimisation problem: find the approximate posterior that best explains that point. If you have a million data points, you face a million optimisation problems. Amortised inference removes that bottleneck with a simple trade. Instead of optimising a separate set of variational parameters for each observation, you train one neural network that maps any observation directly to its approximate posterior. You pay once, during training, and then inference for a new point becomes a single forward pass.

That trade is the reason variational autoencoders work at scale, and it is also the reason they sometimes fail in a distinctive way. This article explains the maths behind amortisation, separates the two gaps it introduces, and shows when the bargain is worth making.

The problem amortisation solves

Start with a latent variable model. An observation x is generated from a latent variable z through a likelihood pθ(x∣z) and a prior p(z) . The posterior pθ(z∣x) is usually intractable, so variational inference approximates it with a distribution q(z) chosen to maximise the evidence lower bound, or ELBO:

log⁡pθ(x)≥Eq(z)[log⁡pθ(x∣z)]−KL(q(z)∥p(z))

In classical mean-field variational inference, q has its own parameters for each data point. Call them λi for observation xi . Inference means running an optimiser until λi converges, separately for every i . That is fine for a few hundred points. It is hopeless for streaming data, for a model that must answer in milliseconds, or for a data set with millions of examples.

Amortised inference replaces the per-point parameters with a shared function. An inference network, usually called an encoder, with parameters ϕ , outputs the variational parameters directly:

qϕ(z∣x)=N(z;μϕ(x),σϕ2(x))

Here μϕ(x) and σϕ(x) are the outputs of a neural network fed with x . There is no longer a separate optimisation per point. The encoder is trained once, on the whole data set, by maximising the average ELBO:

L(θ,ϕ)=1N∑i=1N(Eqϕ(z∣xi)[log⁡pθ(xi∣z)]−KL(qϕ(z∣xi)∥p(z)))

Gradients flow through the sampling step using the reparameterisation trick: draw ϵ∼N(0,I) and set z=μϕ(x)+σϕ(x)⊙ϵ . The network therefore learns, from examples, how to infer.

Why this is a bargain, not a free lunch

Amortisation trades flexibility for speed. A per-point optimiser can, in principle, find the best member of the variational family for each observation. An encoder is constrained: two similar inputs must produce similar posteriors, because the same network weights produce both. If the true posterior changes sharply as x changes, a smooth encoder may fail to track it.

That loss has a name. The total gap between the true log marginal likelihood and the ELBO achieved by an amortised model splits into two parts:

log⁡pθ(x)−Lamortised=KL(q∗∥pθ(z∣x))⏟approximation gap+KL(qϕ∥q∗)⏟amortisation gap

Here q is the best approximation within the chosen variational family, found by per-point optimisation. The approximation gap is the price of restricting yourself to that family at all, for example to Gaussians when the true posterior is multimodal. The amortisation gap is the extra price of using a shared encoder instead of optimising per point. It is zero only if the encoder can represent q for every x .

This decomposition matters because the two gaps call for different fixes. A large approximation gap means you need a richer variational family, such as a normalising flow for q . A large amortisation gap means you need a more expressive encoder, more training, or a hybrid: start from the encoder output and then run a few steps of per-point refinement at test time. That hybrid, sometimes called semi-amortised inference, recovers much of the lost accuracy while keeping most of the speed, because the encoder provides a warm start close to the optimum.

A concrete example: the variational autoencoder

The variational autoencoder is amortised inference in its most familiar form. The encoder maps an image to a mean and a variance. The decoder maps a sampled latent code back to a reconstruction. Training maximises the ELBO above, with the KL term pulling the encoder output towards a standard Gaussian prior and the reconstruction term pushing it to preserve information about x .

A minimal implementation sketch makes the structure clear:

# encoder: x -> (mu, logvar); decoder: z -> x_reconstructed
mu, logvar = encoder(x)
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
z = mu + eps * std
x_hat = decoder(z)
recon = -log_likelihood(x, x_hat)
kl = -0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp(), dim=1)
loss = (recon + kl).mean()
loss.backward()

Notice what is missing from this code: there is no inner optimisation loop over z , and no per-example parameter. At test time, inference is one forward pass through the encoder. That is why a trained autoencoder can encode a new image instantly, while a classical variational fit would need to optimise from scratch.

The same structure appears far beyond images. In topic modelling, an encoder can map a document to a topic distribution. In neuroscience, an encoder can map a neural recording to a latent state. In simulation-based inference, an encoder can map an observed summary statistic to a posterior over simulator parameters, after being trained on simulated pairs. Whenever the same inference problem recurs many times, amortisation is the natural design.

When amortisation helps, and when it hurts

Amortisation helps most when three conditions hold. First, inference is repeated: many observations, streaming data, or a real-time requirement. Second, the mapping from observation to posterior is learnable: similar inputs genuinely have similar posteriors, so a network can generalise across them. Third, the cost of training is acceptable: you can afford to fit the encoder on a large training set before deployment.

It hurts, or at least underperforms, in the opposite cases. If you have only a handful of observations, training an encoder is wasteful and a per-point fit will be more accurate. If the posterior is highly multimodal or discontinuous in x , a smooth encoder may average across modes and produce a posterior that sits where no true posterior mass lives. If the test distribution shifts away from the training distribution, the encoder extrapolates badly, while a per-point optimiser would simply re-optimise on the new point. Out-of-distribution inputs are therefore a classic failure mode for amortised inference, and a good reason to monitor the ELBO on new data rather than trusting the encoder blindly.

A related pitfall is posterior collapse, where the encoder ignores x and returns the prior for every input. This happens when the decoder is powerful enough to reconstruct well without the latent code, so the KL term wins and the encoder learns nothing. The symptom is easy to spot: the encoder outputs barely vary across inputs, and the KL term is near zero. Common remedies include weakening the decoder, annealing the KL weight during early training, or using a structured latent space that the decoder cannot easily bypass.

How to tell which gap is biting

If you suspect amortisation is costing accuracy, measure it directly. Take a trained encoder, freeze the generative model, and then refine qϕ(z∣x) per point with a few optimisation steps. If the ELBO improves substantially, the amortisation gap is large and the encoder is the bottleneck. If refinement barely helps, the variational family itself is the limit, and a richer q or a different model is needed.

A practical checklist for using amortised inference well:

• Compare against a per-point baseline on a small subset before scaling up, so you know the size of the amortisation gap.
• Check encoder outputs across diverse inputs: if they barely vary, suspect posterior collapse.
• Test on held-out and shifted inputs, not just the training distribution, because encoder generalisation is the whole gamble.
• Consider semi-amortised refinement at test time when accuracy matters more than latency: a few gradient steps from the encoder output often buys back most of the gap.
• Remember that the encoder is a model of inference, not of the data: it inherits the biases and blind spots of its training set.

Takeaway

Amortised inference is the idea that inference itself can be learned. Instead of solving an optimisation problem for every new observation, you train a network to predict the solution, and then inference becomes a forward pass. The cost is the amortisation gap: the difference between what the best per-point approximation could achieve and what a shared encoder actually achieves. When inference is repeated, the mapping is learnable, and the gap is small or can be closed by a few refinement steps, that trade is one of the best bargains in modern probabilistic modelling. When it is not, the honest answer is to optimise per point and pay the price in compute.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial