What part is amortized?

The shared network learns a reusable mapping from observations to initial variational parameters.

Additional optimization adapts those parameters for the particular observation under the current generative model.

Does refinement eliminate all variational error?

No. The chosen family can still be restrictive, and finite or imperfect optimization can leave a residual gap.

AI Glossary / Probabilistic machine learning / Advanced

Semi-amortized variational inference

Definition

Semi-amortized variational inference combines an inference network's fast prediction of variational parameters with additional optimization for each individual observation. The network supplies a starting approximation, and local refinement improves it within the chosen variational family. In an end-to-end differentiable implementation, training can account for how the initialization affects the refined result.

What it means

Variational inference approximates a difficult posterior with a simpler distribution. Fully amortized inference uses one shared network to predict that distribution's parameters for every input. This is fast at prediction time, but the network's answer can differ from the best local approximation available in the same family.

Fully local optimization instead adjusts parameters separately for each observation, potentially at substantial cost. Semi-amortized inference combines these approaches. Its name describes sharing part of the inference work across observations, not using only some labels. It is distinct from semisupervised learning, despite the similar prefix.

How it works

For observation x, the inference network outputs initial variational parameters lambda_0. These might specify a Gaussian mean and log variance. A limited sequence of local updates then improves an evidence lower bound, or ELBO, while conditioning on x and the current generative model.

The final parameters lambda_K define the refined variational distribution. In the semi-amortized VAE approach of Kim and colleagues, the optimization steps are differentiated through so the encoder and generative model learn using the refined objective. Merely attaching an optimizer at test time is a related refinement strategy but does not reproduce that training procedure.

This can reduce the amortization gap, the loss caused by the network failing to produce the best member of the selected family. It cannot automatically remove the family's approximation limitations. Iterative updates add computation, and finite steps, noisy gradients or unsuitable step sizes can leave optimization error or even worsen a particular estimated objective.

Three practical examples

1. Text latent-variable modeling: initialize a document's latent distribution with an encoder, then refine it before evaluating the training objective. Check whether the latent representation is useful rather than assuming refinement prevents every form of collapse.

2. Image modeling: spend extra local inference effort on an encoded observation. Compare inference quality and runtime using the same generative model and clearly stated refinement budget.

3. Offline posterior analysis: refine many observations whose initial network approximations are available. Separate gains from local optimization from gains caused by changing the variational family or retraining the model.

Think of it this way

Think of a shared estimating tool producing a first draft, followed by a short case-specific revision. Training the tool through the revision process teaches it to produce drafts that improve effectively. The analogy explains initialization and refinement. It does not imply that the revision finds the exact answer or that every additional revision is beneficial.

Worked example

Consider an analytically simple model used only for teaching: z has prior Normal(0,2), and x given z is Normal(z,2), where the second number is variance. Observe x=4. The exact posterior is Normal(2,1). Choose the variational family q(z)=Normal(m,1), so only its mean is adjustable.

Here ELBO(m)=C-(m-2) squared/2, where C is the log marginal likelihood. Suppose an encoder initially returns m_0=0. Its gap from C is 2. Gradient ascent with step size 0.5 uses gradient 2-m: the first update gives m_1=1 and the second gives m_2=1.5. The gaps become 0.5 and 0.125 respectively.

Each update has form m_next=0.5m+1, so after two steps m_2=0.25m_0+1.5. The derivative of the final ELBO with respect to m_0 is (2-m_2) × 0.25=0.125 at this initialization. This illustrates the gradient path through refinement. The toy model has an exact solution and would not require this method in practice; its purpose is to make the local updates and remaining gap explicit.

Illustrative case study

Fictional case: Larch Language Lab evaluates a latent document model with a shared encoder. The team finds that extra local updates improve the per-document objective on some held-out inputs. Its first comparison uses more refinement steps for the proposed method than for its baseline without reporting the extra runtime.

Reviewers request comparisons at stated computation budgets and checks of the final variational parameters. They also distinguish training through the updates from refinement added only after training. This fictional case illustrates fair evaluation, not a claim that semi-amortization always yields better representations or lower operating costs.

Common misconceptions and mistakes

An improved ELBO for a fixed generative model means a tighter variational approximation under the relevant conditions, not automatically better real-world predictions. If the model changes between comparisons, objective differences can reflect both inference and model changes.

Refining a diagonal Gaussian does not make it capable of representing every correlated or multimodal posterior. A larger number of steps can increase time and memory requirements, and unstable inner optimization can complicate outer training. Stopping gradient flow through the refinement changes the training calculation. Finally, a good initialization reduces work only if the local procedure and stopping criteria actually exploit it.

Business use

Consider semi-amortized inference when posterior approximation quality matters enough to justify additional per-case computation. Measure the incremental benefit of refinement against latency, memory and engineering cost. An offline scientific analysis may tolerate a budget that an interactive service cannot.

Document the variational family, initialization network, local objective, optimizer and number of updates used in both training and deployment. Assess difficult or shifted inputs as well as typical ones. Keep uncertainty claims tied to the approximation being used; local optimization can improve an approximation without certifying that the model or its posterior uncertainty is correct.

Check your understanding

Why is the final toy approximation still imperfect after two improving updates?

Answer: Its mean is 1.5 rather than the exact posterior mean 2. Its ELBO remains 0.125 below C despite the improvement from the starting gap of 2.

Related terms

Variational autoencoder
Amortized inference
variational inference
inference network
amortization gap
evidence lower bound
differentiable optimization
posterior collapse

Learn next

Reproduce the quadratic local updates and differentiate through them. Then distinguish approximation, amortization and optimization errors before comparing encoder-only and refined inference at matched computation budgets.

Sources

Kim and colleagues, Semi-Amortized Variational Autoencoders: https://proceedings.mlr.press/v80/kim18e/kim18e.pdf

The encoder-initialized local refinement and end-to-end differentiation approach were checked in the original paper. The conjugate Gaussian calculation and fictional case are original illustrations. Independent review reconfirmed initialization and differentiated refinement in the original paper.

Previous: Semaphore (programming) Next: Semi-gradient temporal-difference learning