Natural Gradient VI: Guarantees for Non-Conjugate Models

TL;DR

The paper tackles the lack of theoretical guarantees for Natural Gradient Variational Inference (NGVI) in non-conjugate models by establishing a relative-smoothness framework with respect to a KL-induced mirror map. It introduces Proj-SNGD, a non-Euclidean projected NGVI algorithm, and proves its global non-asymptotic convergence to stationary points; under log-concave likelihoods, hidden convexity yields fast global convergence to the optimum. The results reveal how geometric properties of the variational objective govern convergence in challenging inference tasks and provide practical guidance for mean-field Gaussian NGVI. Overall, the work extends NGVI theory to non-conjugate regimes, enabling principled, stable, and efficient Bayesian inference in complex models with explicit rates.

Abstract

Stochastic Natural Gradient Variational Inference (NGVI) is a widely used method for approximating posterior distribution in probabilistic models. Despite its empirical success and foundational role in variational inference, its theoretical underpinnings remain limited, particularly in the case of non-conjugate likelihoods. While NGVI has been shown to be a special instance of Stochastic Mirror Descent, and recent work has provided convergence guarantees using relative smoothness and strong convexity for conjugate models, these results do not extend to the non-conjugate setting, where the variational loss becomes non-convex and harder to analyze. In this work, we focus on mean-field parameterization and advance the theoretical understanding of NGVI in three key directions. First, we derive sufficient conditions under which the variational loss satisfies relative smoothness with respect to a suitable mirror map. Second, leveraging this structure, we propose a modified NGVI algorithm incorporating non-Euclidean projections and prove its global non-asymptotic convergence to a stationary point. Finally, under additional structural assumptions about the likelihood, we uncover hidden convexity properties of the variational loss and establish fast global convergence of NGVI to a global optimum. These results provide new insights into the geometry and convergence behavior of NGVI in challenging inference settings.

Figures (8)

• Figure 1: Convergence of SNGD on Poisson regression with different initialization. The function value gap refers to the difference between the objective value at iteration $t$ and the optimal value; performance is averaged over 10 runs for each algorithm. In the left panel, SNGD is initialized with $\sigma_0^2=2$. We observe instability at the initial iteration and slow convergence with step sizes 0.5 and 0.3. In contrast, with a different initialization $\sigma_0^2=0.4$ shown in the right panel, we observe stable behavior and fast convergence for larger step sizes.
• Figure 2: SNGD (left panel) and Proj-SNGD (right panel) applied to Poisson regression problem with different step-sizes and initialized with $\sigma_0^2 = 2.$Proj-SNGD is used with $U=4$ and $D=25$. Non-Euclidean projection improves convergence and sensitivity of SNGD. The left panel here is the same as the left panel in \ref{['fig:poisson 0']}. The right panel shows that projection fixes SNGD even when starting with the same (large) initial point $\sigma_0^2=2.$
• Figure 3: Euclidean and non-Euclidean algorithms on MNIST dataset. Left: Objective during optimization with tuned step size. Right: Number of iterations before the objective falls below $\ell(\omega) \leq 700$ for different initial step-sizes $\gamma_0$. Non-Euclidean algorithms show consistently better performance, tolerate larger step-sizes and are more robust to step-size tuning.
• Figure 4: Euclidean and non-Euclidean algorithms on Madelon dataset. Left: Objective during optimization with tuned step size. Right: Number of iterations before the objective falls below $\ell(\omega) \leq 3000$ for different initial step sizes $\gamma_0$. Non-Euclidean algorithms show consistently better performance, tolerate larger step sizes and are more robust to step size tuning.
• Figure 5: Euclidean and non-Euclidean algorithms on CIFAR-10 dataset. Left: Objective during optimization with tuned step size. Right: Number of iterations before the objective falls below $\ell(\omega) \leq 35000$ for different initial step sizes $\gamma_0$. Euclidean algorithms are more robust to step size tuning and achieve faster convergence in the initial phase of optimization, but non-Euclidean algorithms exhibit faster convergence after 3000 epochs.

Theorems & Definitions (61)

• Remark
• Definition 2.1: Bregman Divergence
• Definition 3.1
• Lemma 3.1: Bonnet's and Price's Gradients
• Proposition 1
• Theorem 3.2: Sufficient Conditions for Relative Smoothness
• Example 1
• Definition 3.2: Hidden Convexity; fatkhullin2023stochastic
• Proposition 2
• Proposition 3