Variational inference

Overview

A latent-variable or Bayesian model has an exact posterior that is often intractable. Variational inference picks a tractable family Q and makes Q close to the true posterior. Every variant below maximises a lower bound on the marginal log-likelihood, the evidence lower bound (ELBO). The bound uses a factorised or otherwise restricted Q 1903.00114 §6.2 eq. 23. The gap between the bound and the true log-likelihood is a KL divergence, so tightening the bound shrinks that gap 1903.00114 §6.2 eq. 24.

The bound has the same form in each setting, with a different latent variable: graph memberships for block models 1903.00114 §6.2, network weights for Bayesian neural networks 2006.12024 §2.2 §2.4, noisy images for diffusion models 2107.00630 §3.5 eq. 11 and latent guidance trajectories for a recursive reasoner 2605.19376 §2.2.

The shared invariant: the exact bound never exceeds the true marginal log-likelihood, because the gap is a KL divergence 1903.00114 §6.2 eq. 24. In model choice, the bound can stand in for the marginal likelihood 1903.00114 §7.1. GRAM's training surrogate is not the exact bound: it truncates gradients, which the paper calls biased 2605.19376 §2.2. Do not treat its training loss as the bound.

The block-model and topic-model claims rest on one review's classification of the literature 1903.00114 tab. 1. The other three cards are one paper each. Their numbers are author-reported, and none of the papers reports results that compare directly with another card's.

Variants

Variational EM for stochastic block models

Each step alternates between two updates. The E step tightens Q, and the M step updates the model parameters. Stochastic optimisation and belief propagation are named as alternatives 1903.00114 §6.2. The review lists variational inference as the method for a long list of SBM papers, for binary, valued, directed and undirected graphs 1903.00114 tab. 1. Card: A Review of Stochastic Block Models and Extensions for Graph Clustering.

Variational bound as a marginal likelihood for choosing K

The review lists a variational lower bound among the options for choosing the number of groups K, used in place of the marginal likelihood [arxiv:1903.00114 §7.1]. The review says BIC is intractable for these models because it needs the observed-data likelihood 1903.00114 §7.1. Two of the papers it tabulates select K this way 1903.00114 tab. 1.

Variational methods for longitudinal SBMs

Four time-varying SBM papers are tabulated with variational inference, using BIC or ICL to choose the number of groups 1903.00114 tab. 2. The review gives no separate description of their updates.

Variational Bayes for topic models

Topic models are inferred mostly by Bayesian methods: collapsed Gibbs sampling or variational methods 1903.00114 §10.1. The review tabulates a variational topic model by Blei et al. and a variational-plus-SMC longitudinal model by Blei and Lafferty. Neither has a card in this wiki [arxiv:1903.00114 tab. 3].

Classification VEM for text and graphs

For models that combine a block model with topic models, one variant alternates a variational step with a greedy update of the group labels 1903.00114 §11.2.

Bayes by Backprop for neural-network weights

Each weight gets a factorised Gaussian with mean μ and scale σ = softplus(ρ), and the bound is optimised over μ and ρ 2006.12024 §2.4 alg. 1. A weight sample is written as μ + σ ⊙ ε, so gradients flow through μ and σ. The survey says this reparameterisation has lower variance than the score-function estimator 2006.12024 §2.4. For a factorised Gaussian, the survey also describes local reparameterisation, which samples the pre-activation instead 2006.12024 §2.4. Its one experiment applies the method to every convolution weight of a LeNet. That gives 98.99% test accuracy against 99.92% for a vanilla network, from one run 2006.12024 §3.1. The survey notes this variant has twice the parameters of the base network [arxiv:2006.12024 §3]. Card: Bayesian Neural Networks: An Introduction and Survey.

MC dropout as variational inference

Dropout can be read as variational inference with a Bernoulli posterior over units. The KL term is approximated by treating the scaled Bernoulli as a mixture of two narrow Gaussians, and the survey does not give those variances 2006.12024 §2.4. At test time, predictive moments come from averaging stochastic forward passes 2006.12024 §3. The survey says the method adds no parameters and has the same per-pass cost as a vanilla network 2006.12024 §3. On a toy regression, its variance is larger than a GP's out of distribution and too large inside the data 2006.12024 §3 fig. 7.

Closed-form mean-field and low-rank Gaussians

Before gradient methods, variational treatments of one hidden layer had closed forms. One is a factorised Gaussian over the weights. Another is a full-covariance Gaussian whose covariance is low-rank plus diagonal, which scales linearly in the number of hidden units 2006.12024 §2.3.1.

Diffusion models trained on a continuous-time bound

The variational diffusion model writes the bound on -log p(x) as a reconstruction term, a prior term and a diffusion term 2107.00630 §3.5 eq. 11. In continuous time the diffusion term becomes a noise-prediction error weighted by the schedule's derivative 2107.00630 eq. 17. The schedule is itself learned 2107.00630 §3.2. The bound depends on the schedule only through its two endpoint signal-to-noise ratios [arxiv:2107.00630 §5.1] eq. 18. The authors report 2.65 bits per dimension on CIFAR-10 without augmentation, a variational bound, the best of the models in their table 2107.00630 tab. 1. Card: Variational Diffusion Models.

Recursive reasoner with amortised latent guidance

GRAM adds a Gaussian latent at each high-level transition. A prior head sees only the current state. A posterior head also sees the target during training 2605.19376 §2.2 app. B.1. Training maximises an ELBO over latent trajectories, with a KL weight that is set per task 2605.19376 §2.2 app. B.2. Gradients are truncated to the last transition of each supervision step, which the paper calls biased but memory-efficient [arxiv:2605.19376 §2.2] app. A.2. At inference the prior is used. N trajectories are drawn and the most frequent answer is returned 2605.19376 §2.3 alg. 2.

The author-reported Sudoku-Extreme accuracy goes from 86.6% at N=1 to 96.1% at N=20. The authors' reproduction of TRM gives 83.6% 2605.19376 tab. 8. The ablation says that guidance without stochasticity fails completely (0.0 on Sudoku-Extreme and N-Queens). Stochasticity without guidance gives 96.0 on Sudoku-Extreme but 62.1 on N-Queens 2605.19376 tab. 3. Card: Generative Recursive Reasoning.

When to use which

The cards do not rank these variants on a common benchmark. The rows below follow from the constraints each paper reports.

Gaps: the review's Table 1 marks some rows "Unknown" for inference, so do not assume a variational method for those papers 1903.00114 tab. 1. The BNN survey gives no hyperparameters for its experiments and reports no seeds 2006.12024 §3. GRAM's N=20 comparisons use 20 forward passes against single-pass baselines, so state that when comparing 2605.19376 tab. 10.

Implementation sketch

The loop below is the block-model variant. Its update functions are model-specific, so they are left as stubs. The review gives the bound and the KL gap, but no update formulas 1903.00114 §6.2.

# Variational EM loop for a latent-variable graph model  [§6.2] [eq. 23]
# Input:  adjacency Y (n, n), number of groups K, iteration limit
# Output: factorised Q over latent memberships; parameters theta, C

def variational_em(Y, K, max_iter):
    # Start from a factorised Q and initial parameters  [§6.2] [eq. 23]
    Q = init_factorised(n, K)
    theta, C = init_parameters(K)
    for it in range(max_iter):
        # E step: tighten the lower bound by updating Q  [§6.2] [eq. 24]
        Q = update_q(Q, Y, theta, C)
        # M step: update the parameters given Q  [§6.2]
        theta, C = update_params(Q, Y)
        # Stopping rule: not stated in the review
        if bound_converged(Q, Y, theta, C):
            break
    return Q, theta, C

When the latent variable is continuous and the model is differentiable, the usual step is a reparameterised gradient on the ELBO. This is the Bayes by Backprop form. The survey gives the same estimator in its algorithm 1.

# One reparameterised gradient step on the ELBO  [§2.4] [alg. 1]
# Input:  mu, rho (variational parameters); batch; dataset_size; S draws
# Output: updated mu, rho

def elbo_step(mu, rho, batch, dataset_size, lr, S):
    # Positive scale from an unconstrained parameter  [§2.4]
    sigma = softplus(rho)
    elbo = 0.0
    for s in range(S):
        # Noise is drawn independently of the parameters  [§2.4]
        eps = random_normal(mu.shape)
        # Weight sample as a differentiable function of mu and rho  [§2.4]
        w = mu + sigma * eps
        # Rescale the minibatch likelihood to the full data set  [§2.4]
        scale = dataset_size / len(batch)
        joint = log_prior(w) + scale * log_likelihood(batch, w)
        # Average over draws, not sum  [alg. 1]
        elbo += (joint - log_q(w, mu, sigma)) / S
    # Gradient ascent on the bound  [§2.4]
    mu = mu + lr * grad(elbo, mu)
    rho = rho + lr * grad(elbo, rho)
    return mu, rho

Notes on where the variants differ:

• Averaging. The survey's cost sums over samples, but its algorithm divides by the sample count. Use the mean, because a sum scales the gradient by the sample count 2006.12024 §2.4 alg. 1.
• Minibatch scaling. The survey's minibatch objective has an N/M factor but sums over the full index range. The standard estimator sums over the batch, as the sketch does 2006.12024 §2.4.
• Discrete diffusion term. Compute exp(g_t - g_s) - 1 with expm1. The paper says 32-bit floats are enough with this form 2107.00630 §4 app. J.
• Schedule training. The paper trains the schedule's interior to reduce the variance of the estimate 2107.00630 §5.3 app. I.2. The official code does not appear to implement that step, so check before reproducing the variance results 2107.00630 app. I.2.
• Training versus inference in GRAM. Training uses the posterior guidance in the cycles without gradient too, as the paper's algorithm shows. At inference, use the prior 2605.19376 alg. 1 alg. 2.

Contradictions

No two cards report incompatible numbers for the same setting. Three points need care rather than resolution:

• Variance direction. The BNN survey says factorised VI underestimates predictive variance, and it cites a textbook for this rather than deriving it 2006.12024 §2.6. Its own toy regression shows Bayes by Backprop under-estimating variance outside the training range, while MC dropout over-estimates it 2006.12024 §3 fig. 7.
• Bound versus reported number. The diffusion results are variational bounds, not exact likelihoods, and the authors expect tighter estimates from importance sampling 2107.00630 tab. 1 caption.
• Comparability. The review cautions that its cross-paper comparisons mix different networks, sizes and criteria 1903.00114 §11.4 §12.1. The same applies across the four cards, which evaluate different tasks.

Papers

• A Review of Stochastic Block Models and Extensions for Graph Clustering: a survey that classifies SBM papers by inference method, including variational EM, and by how they select the number of groups.
• Bayesian Neural Networks: An Introduction and Survey: a survey of Bayesian neural networks covering Bayes by Backprop, MC dropout and HMC, with two small experiments.
• Variational Diffusion Models: a diffusion model trained on a continuous-time variational bound with a learned noise schedule.
• Generative Recursive Reasoning: a recursive reasoner whose latent steps are sampled and trained with amortised variational inference.