Batch and match: black-box variational inference with a score-based divergence
Diana Cai, Chirag Modi, Loucas Pillaud-Vivien, Charles C. Margossian, Robert M. Gower, David M. Blei, Lawrence K. Saul
ICML)
stat.ML, cs.AI, cs.LG, stat.CO
2024-02-23
BaM replaces ELBO maximization with a score-based divergence and closed-form proximal updates, needing an order of magnitude fewer gradients than ADVI.
The default recipe in black-box variational inference (BBVI) is to maximize the ELBO with reparameterized SGD. Two complaints recur: noisy gradient estimates and brittle sensitivity to the learning rate. Both get worse in high dimensions and with richer variational families, which is awkward because full-covariance Gaussians are exactly what you want when latent variables are correlated. Stan, PyMC, and Pyro all ship this recipe. One prior escape route, Gaussian score matching (GSM, 2023), avoids SGD for Gaussian families but minimizes no explicit divergence and leans on heuristics for non-Gaussian targets. The question here is concrete: change the divergence, keep the full-covariance Gaussian family, and solve every update in closed form.
BaM (batch and match) optimizes a score-based divergence: the expected squared gap between the scores (gradients of log densities) of target and approximation, averaged under the approximation and weighted by the inverse of its covariance. Only the unnormalized target's score is needed. The divergence is nonnegative, zero exactly when the two densities match, and invariant to affine changes of coordinates, a property plain Fisher divergence lacks.
Each iteration alternates two steps.
The closed form is the point. No linearization anywhere; this is a stochastic proximal point algorithm, a family known for staying stable when hyperparameters move. Limiting cases: λ→0 does nothing; B=1 with λ→∞ recovers GSM; infinite batch with λ→∞ nails a Gaussian target in one step. For Gaussian targets in the infinite-batch limit, the normalized mean and covariance errors decay like (1−δ)^t with δ = λβ/(1+λ), for every λ > 0 and every initialization, with no tying the step size to an unknown Hessian eigenvalue as gradient methods must. Each update costs O(D³), dropping to O(D²B + B³) when B ≪ D via a low-rank solver.
Baselines: ADVI (Stan's ELBO implementation) and GSM.
| Setting | Baseline | Result |
| Gaussian targets, D = 4/16/64/256, forward KL | ADVI | Orders of magnitude fewer gradient evaluations |
| sinh-arcsinh targets (varying skew and tails) | ADVI, GSM | Faster than ADVI; GSM and the Score variant diverge at skew 1.8 while BaM stays stable |
| Three posteriordb Bayesian models (D = 7/10/13, HMC reference) | ADVI, GSM | Lower relative mean error; higher relative SD error on eight-schools |
| CIFAR-10 deep generative model, latent D = 256 | ADVI, GSM | More than an order of magnitude faster at B = 300 |
Exact error values live in log-scale convergence curves; the paper reports no numeric tables. One hard number from the CIFAR-10 experiment: with a budget of 3,000 gradient evaluations, ADVI's best setting is B = 10, T = 300, and BaM matches its reconstruction MSE with B = 300, T = 10. The structural difference is parallelism: BaM's B gradients per iteration are independent and can run in parallel, while ADVI's are mostly sequential, and the wallclock curves show BaM winning on that experiment.
For anyone running ADVI-style inference in Stan, PyMC, or Pyro, this is a drop-in candidate for the same black-box setting (you supply the log of an unnormalized density and its gradients): fewer gradient evaluations, closed-form updates, no SGD learning-rate search, and batches that parallelize naturally on GPUs. The divergence itself has independent value: unlike KL it is computable for unnormalized models, so it can serve as a goodness-of-fit diagnostic. Honest framing: this is an algorithmic upgrade inside the Gaussian family. The approximation ceiling is still set by whether a Gaussian can fit the target at all.
The authors list theirs: convergence is proven only for Gaussian targets with infinite batches; only Gaussian families are supported; no data subsampling, so large-N Bayesian problems are untouched. Reading closely adds more. λ needs less babysitting in theory than in practice: constant λ = BD worked for Gaussian targets, but non-Gaussian targets required a decay schedule BD/(t+1). On skewed targets BaM settles at a worse forward KL than ADVI (reverse KL comparable), since heavy skew is out of a Gaussian's reach. The method needs batch sizes comparable to the dimension (B = 10 with D = 256 underperforms), which raises per-iteration memory and matrix-sqrt cost. And on eight-schools the relative SD error ends higher, a reminder that a full-covariance Gaussian does not automatically win everywhere.