In variational inference (VI), the practitioner approximates a high-dimensional distribution π with a simple surrogate one, often a (product) Gaussian distribution. However, in many cases of practical interest, Gaussian distributions might not capture the correct radial profile of π, resulting in poor coverage. In this work, we approach the VI problem from the perspective of optimizing over these radial profiles. Our algorithm radVI is a cheap, effective add-on to many existing VI schemes, such as Gaussian (mean-field) VI and Laplace approximation. We provide theoretical convergence guarantees for our algorithm, owing to recent developments in optimization over the Wasserstein space--the space of probability distributions endowed with the Wasserstein distance--and new regularity properties of radial transport maps in the style of Caffarelli (2000).
Figures & tables
Compute stochastic gradient ∇λF(Tλ(k))
λ(k+1)=ProjR+J+1,∥⋅∥Q(λ(k)−hQ−1∇λF(Tλ(k)))
Algorithm 1 radVI
Figure 1: Convergence of radVI for various target distributions. See Table 1 for final-iterate comparisons between GVI and LA.
Isotropic targets
Method
Gaussian
Laplace
Logistic
Student- t
LA
2.45×10−4
20.00
1.6×103
25.87
GVI
7.34×10−4
8.24
3.96
1.99
radVI
1.15×10−4
5.37×10−2
1.84×10−1
1.19×10−1
Table 1: Estimated squared Wasserstein distance between various VI solutions for learning isotropic targets.
Figure 2: Comparing learned radial profiles of radVI versus other approximation methods for learning the Student- t distribution in the isotropic ( top ) and anisotropic case (bottom) .
Figure 3: In the case where π is an isotropic Gaussian with d=50 , we verify that radVI is robust to the choice of α .
Figure 4: Top: Comparing whitening methods for learning the anisotropic logistic distribution, with and without radVI . Bottom: Visual comparison of true target samples, those generated by LA, and ours (LA+ radVI ).
d=25
Parameter
E[z2]=4
E[x12]≈7.389
P(∣z∣>2)≈0.317
GVI
0.274 ±4×10−3
1.12 ±1.4×10−2
0
GVI+radVI
2.41±2×10−2
5.96±1.0×10−1
0.214±3×10−3
d=50
Parameter
E[z2]=4
E[x12]≈7.389
P(∣z∣>2)≈0.317
GVI
0.328 ±5×10−3
1.61 ±2×10−2
0
Table 2: In d∈{25,50} , we compare the performance of GVI and GVI+ radVI for parameter estimation.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Isotropic targets
Method
Gaussian
Laplace
Logistic
Student- t
LA
1.48×10−4
43.04
7450
68.89
GVI
5.07×10−4
18.34
8.67
5.21
radVI
3.71×10−4
7.67×10−2
1.96×10−1
1.89×10−1
Appendix
Table 3: Estimated Wasserstein distance between various VI solutions for learning isotropic targets in d=100 .
Figure 5: Comparing learned radial profiles of radVI versus other approximation methods for the isotropic Student- t distribution.
Figure 6: Comparing learned radial profiles of radVI versus other approximation methods for the isotropic logistic distribution.
Figure 7: Visual comparison of true target samples, those generated by LA, and ours (LA+ radVI ), for learning the anisotropic Student- t distribution.
Figure 8: Left: Comparing whitening methods for learning the anisotropic Laplace distribution, with and without radVI . Right: Visual comparison of true target samples, those generated by GVI, and ours (GVI+ radVI ).
For approximating a target distribution given only its unnormalized log-density, stochastic gradient-based variational inference (VI) algorithms are a popular approach. For example, Wasserstein VI (WVI) and black-box VI (BBVI) perform gradient descent in measure space (Bures-Wasserstein space) and parameter space, respectively. Previously, for the Gaussian variational family, convergence guarantees for WVI have shown superiority over existing results for black-box VI with the reparametrization gradient, suggesting the measure space approach might provide some unique benefits. In this work, however, we close this gap by obtaining identical state-of-the-art iteration complexity guarantees for both. In particular, we identify that WVI's superiority stems from the specific gradient estimator it uses, which BBVI can also leverage with minor modifications. The estimator in question is usually associated with Price's theorem and utilizes second-order information (Hessians) of the target log-density. We will refer to this as Price's gradient. On the flip side, WVI can be made more widely applicable by using the reparametrization gradient, which requires only gradients of the log-density. We empirically demonstrate that the use of Price's gradient is the major source of performance improvement.
Kyurae Kim, Qiang Fu, Yi-An Ma +2
University of Pennsylvania, Philadelphia, U.S. · Yale University, New Haven, U.S. · University of California San Diego, San Diego, U.S. +1
When approximating an intractable density via variational inference (VI) the variational family is typically chosen as a simple parametric family that very likely does not contain the target. This raises the question: Under which conditions can we recover characteristics of the target despite misspecification? In this work, we extend previous theoretical results on robust VI with location-scale families under target symmetries in two substantial ways: (1) We open them up to a wider range of divergences by providing sufficient conditions for exact recovery of the target mean and correlation matrix when using the forward Kullback-Leibler divergence and α-divergences. (2) By doing so, we find that we can drop the restrictive assumption of a log-concave target made in previous work, allowing us to give guarantees for a wider range of targets, including multi-modal ones. In our experiments, we show how our guarantees can serve as guidelines for the choice of the variational family and α-value and we illustrate on a diverse set of examples how and why optimization can fail in the absence of our sufficient conditions.
Particle-based variational inference (ParVI) methods approximate an intractable target distribution by evolving an ensemble of interacting samples. Existing approaches rely predominantly on kernel-based repulsion (e.g., SVGD), which suffers from variance collapse in high dimensions and mode collapse on multimodal targets -- pathologies caused by the absence of global transport structure. We introduce entropic transport descent (ETD), a ParVI family that frames each particle update as an entropy-regularized optimal transport problem. Derived from the JKO proximal scheme by lifting to the space of couplings and relaxing via the KL chain rule, each ETD iteration reduces to a Sinkhorn computation. The resulting transport plan provides global coordination, guiding each particle to nearby high-density proposals and naturally preserving multimodal structure. ETD can operate entirely score-free, requiring only pointwise evaluations of the unnormalized target density. Experiments on variance-collapse diagnostics, Bayesian logistic regression, neural networks, and molecular Boltzmann distributions show that ETD matches or outperforms SVGD, AGF-SVGD, and SGLD, with the largest gains in high-dimensional and multimodal settings.
Vincent Pacelli, Akash Ratheesh, Evangelos Theodorou
School of Aerospace Engineering Georgia Institute of Technology Atlanta, GA 30332