Research map / Probabilistic and scientific modeling / Variational and generative sampling
Batch and match: black-box variational inference with a score-based divergence
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes batch and match (BaM), a black-box VI method that minimizes a regularized Monte Carlo score-based divergence over full-covariance Gaussians with closed-form proximal updates, with exponential convergence proof for Gaussian targets.
Problem
Black-box variational inference based on stochastic ELBO maximization converges slowly due to high-variance gradients and hyperparameter sensitivity, especially for full-covariance Gaussian families.
Approach
Define an affine-invariant score-based divergence (weighted Fisher divergence with M = Cov(q)) that needs only the target score, and minimize its batch Monte Carlo estimate plus a KL proximal regularizer; the match step has a closed-form solution via a quadratic matrix equation, avoiding SGD.
Key results
- On Gaussian targets, BaM converges orders of magnitude faster than ADVI in gradient evaluations across dimensions D=4 to 256.
- For the deep generative model with B=300, BaM converges an order of magnitude or more faster than ADVI and GSM in wallclock time for CIFAR-10 image reconstruction.
- Theorem: for Gaussian targets in the infinite-batch limit, the mean and covariance errors decay exponentially at rate controlled by the regularization and initialization, for any regularization level.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use BaM instead of ADVI | Full-covariance Gaussian VI · Synthetic Gaussian targets, D=4, 16, 64, 256 | gradient evaluations to convergence: orders of magnitude faster | Helped |
| Increase BaM batch size | BaM · Synthetic Gaussian targets | convergence speed | Helped |
| Use BaM instead of ADVI | Full-covariance Gaussian VI · Sinh-arcsinh targets with normal tails and varying skew | forward KL | Mixed |
| Use BaM instead of ADVI | Full-covariance Gaussian VI · Sinh-arcsinh targets with normal tails and varying skew | convergence speed | Helped |
| Use GSM on a highly skewed target | GSM · Sinh-arcsinh target, s=1.8 | reverse KL | Hurt |
| Use the Score method on a highly skewed target | Score · Sinh-arcsinh targets with high skew | convergence | Hurt |
| Use BaM instead of ADVI | Full-covariance Gaussian VI · Sinh-arcsinh targets with varying tail heaviness | gradient evaluations to convergence | Helped |
| Use BaM instead of ADVI | Full-covariance Gaussian VI · Hierarchical Bayesian models: arK, gp-pois-regr, eight-schools-centered | relative posterior mean error | Helped |
| Increase BaM batch size | BaM · Hierarchical Bayesian models | convergence speed and stability | Helped |
| Use BaM with a very small batch size (B=10) | BaM · CIFAR-10 deep generative model posterior reconstruction | reconstruction convergence | Hurt |
| Use BaM with B=300 instead of ADVI or GSM | BaM · CIFAR-10 deep generative model posterior reconstruction | gradient evaluations to convergence: an order of magnitude (or more) faster | Helped |
| Use a full-covariance approximation instead of a factorized Gaussian | BaM and ADVI versus AVI · CIFAR-10 deep generative model posterior reconstruction | reconstruction error | Helped |
Surprising findings
- The score-based divergence is typically more sensitive to the learning rate than ELBO-based gradient methods.
- For small batch sizes relative to dimension, BaM performs poorly in the deep generative experiment (B=10).
Limitations
- Convergence proof only covers Gaussian targets in the infinite-batch limit; finite-batch and non-Gaussian convergence is empirical.
- Non-Gaussian targets required a decaying learning rate for BaM to converge.
- Performance degrades when batch size is small relative to latent dimension.
Experimental setup
Datasets: CIFAR-10 posteriordb (arK, gp-pois-regr, eight-schools-centered) synthetic Gaussian and sinh-arcsinh targets
Benchmarks: posteriordb CIFAR-10 deep generative model reconstruction
Metrics: forward KL divergence reverse KL divergence relative posterior mean error relative posterior SD error reconstruction MSE
Baselines: ADVI GSM Fisher divergence ADVI score-based ADVI
Builds on: automatic differentiation VI Gaussian score matching stochastic proximal point method
Closest papers by meaning
- EigenVI: score-based variational inference with orthogonal function expansions (NeurIPS 2024 · similarity 0.90)
- Nearly Dimension-Independent Convergence of Mean-Field Black-Box Variational Inference (NeurIPS 2025 · similarity 0.89)
- Efficient Mixture Learning in Black-Box Variational Inference (ICML 2024 · similarity 0.89)
- Least squares variational inference (NeurIPS 2025 · similarity 0.88)
- Provably Scalable Black-Box Variational Inference with Structured Variational Families (ICML 2024 · similarity 0.87)
- Denoising Diffusion Variational Inference: Diffusion Models as Expressive Variational Posteriors (AAAI 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.