Natural Gradient VI: Guarantees for Non-Conjugate Models

TL;DR

This paper advances the theoretical understanding of stochastic Natural Gradient Variational Inference (NGVI) for non-conjugate models, establishing convergence guarantees and uncovering hidden convexity properties.

Contribution

It derives conditions for relative smoothness, proposes a modified NGVI algorithm with convergence guarantees, and reveals hidden convexity in the variational loss for non-conjugate models.

Findings

Proposed a modified NGVI algorithm with global convergence.

Established conditions for relative smoothness of the variational loss.

Uncovered hidden convexity leading to fast convergence to global optima.

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.