# Scalable Mean-Field Variational Inference via Preconditioned Primal-Dual Optimization

arXiv: https://arxiv.org/abs/2602.07632 · HTML: https://mlsanity.com/paper/2602.07632/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 54 references that are in the map.
- First posted on arXiv 7 Feb 2026.

## Builds on (most cited first, 26 of 26)

- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Federated Learning Based on Dynamic Regularization](https://mlsanity.com/paper/2111.04263/) (2021)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [Advances in Variational Inference](https://mlsanity.com/paper/1711.05597/) (2017)
- [Handbook of Convergence Theorems for (Stochastic) Gradient Methods](https://mlsanity.com/paper/2301.11235/) (2023)
- [FedPD: A Federated Learning Framework with Optimal Rates and Adaptivity to Non-IID Data](https://mlsanity.com/paper/2005.11418/) (2020)
- [Adam-mini: Use Fewer Learning Rates To Gain More](https://mlsanity.com/paper/2406.16793/) (2024)
- [Adam Can Converge Without Any Modification On Update Rules](https://mlsanity.com/paper/2208.09632/) (2022)
- [Theory and Experiments on Vector Quantized Autoencoders](https://mlsanity.com/paper/1805.11063/) (2018)
- [BayeSeg: Bayesian Modeling for Medical Image Segmentation with Interpretable Generalizability](https://mlsanity.com/paper/2303.01710/) (2023)
- [Challenges and Opportunities in High-dimensional Variational Inference](https://mlsanity.com/paper/2103.01085/) (2021)
- [Structured Stochastic Variational Inference](https://mlsanity.com/paper/1404.4114/) (2014)
- [Batch and match: black-box variational inference with a score-based divergence](https://mlsanity.com/paper/2402.14758/) (2024)
- [Using Large Ensembles of Control Variates for Variational Inference](https://mlsanity.com/paper/1810.12482/) (2018)
- [Robust, Accurate Stochastic Optimization for Variational Inference](https://mlsanity.com/paper/2009.00666/) (2020)
- [Faster Stochastic Variational Inference using Proximal-Gradient Methods with General Divergence Functions](https://mlsanity.com/paper/1511.00146/) (2015)
- [Mean-field Variational Inference via Wasserstein Gradient Flow](https://mlsanity.com/paper/2207.08074/) (2022)
- [On the Convergence of Coordinate Ascent Variational Inference](https://mlsanity.com/paper/2306.01122/) (2023)
- [Approximation Based Variance Reduction for Reparameterization Gradients](https://mlsanity.com/paper/2007.14634/) (2020)
- [Understanding Stochastic Natural Gradient Variational Inference](https://mlsanity.com/paper/2406.01870/) (2024)
- [Principled Parallel Mean-Field Inference for Discrete Random Fields](https://mlsanity.com/paper/1511.06103/) (2015)
- [Preconditioning for Accelerated Gradient Descent Optimization and Regularization](https://mlsanity.com/paper/2410.00232/) (2024)
- [Asynchronous Stochastic Variational Inference](https://mlsanity.com/paper/1801.04289/) (2018)
- [Design Criteria for SGD Preconditioners: Local Conditioning, Noise Floors, and Basin Stability](https://mlsanity.com/paper/2511.19716/) (2025)
- [Dynamic Learning Rate Decay for Stochastic Variational Inference](https://mlsanity.com/paper/2412.15745/) (2024)
- [Bayesian Variational Inference for Mixed Data Mixture Models](https://mlsanity.com/paper/2507.16545/) (2025)

## Built on by (most cited first, 1 of 1)

- [Stability of Finite-Batch Particle Mean-Field Variational Inference Beyond Strong Convexity](https://mlsanity.com/paper/2608.11486/) (2026)

## Abstract (arXiv)

In this work, we investigate the large-scale mean-field variational inference (MFVI) problem from a mini-batch primal-dual perspective. By reformulating MFVI as a constrained finite-sum problem, we develop a novel primal-dual algorithm based on an augmented Lagrangian formulation, termed primal-dual variational inference (PD-VI). PD-VI jointly updates global and local variational parameters in the evidence lower bound in a scalable manner. To further account for heterogeneous loss geometry across different variational parameter blocks, we introduce a block-preconditioned extension, P2D-VI, which adapts the primal-dual updates to the geometry of each parameter block and improves both numerical robustness and practical efficiency. We establish convergence guarantees for both PD-VI and P2D-VI under properly chosen constant step size, without relying on conjugacy assumptions or explicit bounded-variance conditions. In particular, we prove O(1/T) convergence to a stationary point in general settings and linear convergence under strong convexity. Numerical experiments on synthetic data and a real large-scale spatial transcriptomics dataset demonstrate that our methods consistently outperform existing stochastic variational inference approaches in terms of convergence speed and solution quality.