Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space

2023/12/05 by Yiheng Jiang, Jiang, Yiheng, Sinho Chewi +3 · 3 citations
Mathematics · #FOS: Computer and information sciences #FOS: Mathematics #Machine Learning (cs.LG) #Markov Chains and Monte Carlo Methods #Optimization and Control (math.OC) #Statistical Methods and Inference #Statistics Theory (math.ST)

paper · pdf · doi:10.48550/arxiv.2312.02849

openalex publication_date 2023/12/05 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

We develop a theory of finite-dimensional polyhedral subsets over the Wasserstein space and optimization of functionals over them via first-order methods. Our main application is to the problem of mean-field variational inference, which seeks to approximate a distribution π over ℝ^d by a product measure π^⋆. When π is strongly log-concave and log-smooth, we provide (1) approximation rates certifying that π^⋆ is close to the minimizer π^⋆_\diamond of the KL divergence over a polyhedral set P_\diamond, and (2) an algorithm for minimizing KL(⋅‖π) over P_\diamond based on accelerated gradient descent over \R^d. As a byproduct of our analysis, we obtain the first end-to-end analysis for gradient-based algorithms for MFVI.

Cited by

• Mode Collapse of Mean-Field Variational Inference
• Rotated Mean-Field Variational Inference and Iterative Gaussianization
• Gaussian mixture layers for neural networks

Related

• Variational Analysis in the Wasserstein Space
• On Representations of Mean-Field Variational Inference
• Sliced Wasserstein Variational Inference
• Mean-field variational inference with the TAP free energy: Geometric and statistical properties in linear models
• Wasserstein Contraction of Coordinate Ascent Variational Inference