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

TL;DR

This paper introduces a scalable primal-dual optimization framework for large-scale mean-field variational inference, improving convergence and robustness over existing methods, with theoretical guarantees and empirical validation.

Contribution

It develops a novel primal-dual algorithm for MFVI reformulated as a constrained finite-sum problem, including a block-preconditioned extension for enhanced efficiency and robustness.

Findings

Achieves $O(1/T)$ convergence to stationary points.

Demonstrates linear convergence under strong convexity.

Outperforms existing stochastic variational inference methods in speed and quality.

Abstract

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, P$^2$D-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 P$^2$D-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.