Stochastic Variational Inference

Ordinary variational inference fits an approximate distribution to a posterior by climbing the ELBO (the evidence lower bound) with gradient ascent, but computing that gradient exactly requires summing a term over every single data point in the dataset on every single update. With a million rows, that is a million-row pass just to take one step. Stochastic variational inference (SVI) asks: what if each step only looked at a small random handful of rows, scaled up to stand in for the whole dataset?

The analogy: polling instead of a full census

To estimate an election result, you do not need to survey every voter — a well-chosen random sample of a few hundred people, scaled up to the size of the electorate, gives an estimate that is unbiased and, on average, points the right direction, even though any single poll is noisy. SVI treats each gradient update the same way: grab a random minibatch, compute the update as if the whole dataset looked like this minibatch repeated many times, and take a small step. Individual steps are noisy polls, not censuses, but averaged over many steps they walk to the same place a full census would have reached.

The mechanics

For a dataset of $N$ points and a minibatch of size $M$, the exact ELBO gradient with respect to a variational parameter (say the mean $m$ of a mean-field Gaussian approximation) sums a per-point term over all $N$ points. SVI approximates that sum by averaging over the minibatch and rescaling by $N/M$:

$$\hat{g} = \frac{N}{M} \sum_{i \in \text{minibatch}} \nabla_m \, \ell_i(m)$$

In words: take the average gradient contribution from the points you actually looked at, then multiply by how many times bigger the full dataset is than your sample. This $\hat{g}$ is an unbiased estimator of the true full-data gradient — its average over many random minibatches equals the true gradient exactly, even though any one minibatch's estimate is off. Because gradient ascent only needs the right direction on average, a shrinking step size (the Robbins-Monro schedule, e.g. $\rho_t \propto 1/t$) lets the noise wash out over iterations while the parameter still converges.

Worked example 1: one update by hand

Suppose $N = 1{,}000{,}000$ data points, minibatch $M = 4$, fixed noise variance $\sigma^2=1$, current mean estimate $m_t = 0$. The minibatch values are $x = \{2.0, 2.4, 1.6, 2.2\}$, and the per-point gradient of the log-likelihood term with respect to $m$ is $(x_i - m)/\sigma^2$. Average over the minibatch: $(2.0+2.4+1.6+2.2)/4 = 2.05$. Scale by $N/M = 250{,}000$: $\hat{g} = 250{,}000 \times 2.05 = 512{,}500$. That huge number is exactly why SVI uses a tiny step size $\rho_t$ — with $\rho_t = 4\times 10^{-7}$, the update is $m_{t+1} = 0 + 4\times10^{-7} \times 512{,}500 \approx 0.205$, nudging $m$ toward the data mean of about $2.05$ without ever touching the other 999,996 points.

Worked example 2: cost comparison

Full-batch VI on $N=1{,}000{,}000$ points needs one full pass (1,000,000 evaluations) per gradient step; to take 500 steps costs 500 million evaluations. SVI with $M=100$ costs only $100 \times 500 = 50{,}000$ evaluations for the same 500 steps — 10,000 times cheaper — and because each step is an unbiased estimate, it still converges to essentially the same fitted distribution, just with a noisier, wigglier path there.

Watch how a noisy, jittery descent path still lands near the same minimum as a smooth one — that jitter is exactly what a minibatch gradient estimate looks like next to a full-batch one.

What this means in practice

SVI is why variational autoencoders, topic models, and Bayesian neural networks can be fit on datasets with millions of rows using ordinary minibatch training loops — the same infrastructure built for stochastic gradient descent on a supervised loss works unchanged, because the ELBO gradient estimator plugs into exactly that machinery. The price is that convergence diagnostics get harder: a noisy ELBO trace can look like it has plateaued when it has actually converged, or look like it is still moving when it is just sampling noise.