Natural Gradient VI: Guarantees for Non-Conjugate Models

AI summary

Overview

Research area: Machine learning theory — specifically the convergence theory of stochastic natural gradient variational inference (NGVI), and stochastic mirror descent (SMD) in non-Euclidean geometry.

Technical level: Advanced. The paper assumes familiarity with exponential families, Fisher information geometry, Bregman divergences, relative smoothness, and non-convex stochastic optimization theory.

Scope: The paper derives sufficient conditions for the negative ELBO to be relatively smooth under mean-field Gaussian variational families, proposes a projected variant of stochastic natural gradient descent (Proj-SNGD), and proves non-asymptotic convergence guarantees for non-conjugate likelihoods such as logistic and Poisson regression.

What This Paper Is About

Variational inference approximates a Bayesian posterior by turning inference into an optimization problem over a tractable family of distributions, and natural gradient variational inference (NGVI) does this using a geometry-aware update that follows steepest descent in KL divergence rather than Euclidean distance. Prior theory only guaranteed convergence for conjugate models, where the objective is well-conditioned and convex-like; for non-conjugate models such as logistic or Poisson regression, the objective becomes non-convex and the earlier guarantees break down. This paper characterizes the loss landscape in that harder regime and proves convergence rates for a modified NGVI algorithm.

Key Contributions

1. Relative smoothness conditions. The authors derive sufficient conditions under which the negative ELBO is smooth relative to the Bregman geometry induced by KL divergence on a compact parameter set, and give explicit constants (polynomial in problem dimension) for a broad class of non-conjugate mean-field models, including logistic regression.

2. A projected algorithm (Proj-SNGD). They introduce a projected variant of stochastic NGVI that keeps updates inside a compact parameter domain using non-Euclidean projections, implemented efficiently in the mean-field Gaussian case via entry-wise clipping of the mean and variance parameters.

3. Convergence to a stationary point. Analyzing Proj-SNGD as stochastic mirror descent, they prove a non-asymptotic convergence rate of O(1/√T) to a stationary point using the Bregman Forward-Backward Envelope (BFBE) as the stationarity criterion.

4. Hidden convexity and fast global convergence. Under the additional assumption that the log-likelihood is concave in the latent variable z, they uncover hidden convexity of the variational objective and a Polyak-Łojasiewicz (PL) inequality, enabling a fast O(1/T) global convergence rate.

Main Findings

• The KL divergence term is well-behaved: The KL divergence between the mean-field Gaussian variational distribution and the standard Gaussian prior is exactly 1-1 smooth relative to the mirror map A*, since the Hessians of the two functions coincide.

• Non-conjugate objectives are relatively smooth, not Euclidean-smooth: For the log-likelihood term, the smoothness constant on the bounded set Ω̃ is L = O(d D² U (L₁ + L₂ U) + d D³ (L₁ + L₂ U)), where U bounds |μᵢ|, D bounds the diagonal variances between D⁻¹ and D, and L₁, L₂ bound the first and second derivatives of the negative log-likelihood.

• Logistic regression instantiation: With L₁ = ‖x‖∞ and L₂ = ‖x‖∞², the objective is smooth relative to A* with parameter O(d D² ‖x‖∞ (U + D)(1 + U ‖x‖∞)).

• Contrast with conjugate models: For conjugate likelihoods the objective is 1-smooth and 1-strongly convex relative to A*, but for non-conjugate models the lower curvature parameter α is typically negative, confirming genuine non-convexity. The relative smoothness constant scales polynomially with the size of the bounded set and with dimension d.

• Hidden convexity under log-concave likelihood: When log p(D|z) is concave in z (as in logistic and Poisson regression), the objective restricted to Ω̃ is hidden convex with modulus μ_C = (4U² + 4D + 1)^(−1/2) and μ_H = 1.

• PL inequality: Under concavity of the log-likelihood and with a stationary point ω* in Ω̃, that point is a global minimum and the objective satisfies ‖∇ℓ(ω)‖² ≥ 2μ_C²(ℓ(ω) − ℓ*), which is a sufficient condition for global convergence.

• Explicit convergence rate for Proj-SNGD: With constant step size γ = min{1/(2L), √(λ₀/(V²LT))} and λ₀ = ℓ(ω₀) − ℓ*, the expected BFBE criterion satisfies E[E_{3L}(ω̄_T)] ≤ 18 L λ₀ / T + 9 √(L V² λ₀ / T).

• Projection is necessary, not cosmetic: The authors argue that SNGD can leave the bounded region Ω̃ (as illustrated in the Poisson regression experiments), which would invalidate both the relative smoothness and hidden convexity guarantees. The projection step reduces to clipping μᵢ into [−U, U] and Σᵢᵢ into [1/D, D] in the mean-field Gaussian case.

• Illustrative experiment: Figure 1 shows SNGD on Poisson regression averaged over 10 runs; with initialization σ₀² = 2 the method is unstable in the first iteration and converges slowly with step sizes 0.5 and 0.3, whereas with σ₀² = 0.4 the behavior is stable and converges faster at those step sizes.

• Not reported in the provided content: Named benchmark datasets, dataset sizes, comparison baselines, and the detailed statement of the O(1/T) global convergence theorem (Section 4.3) and full experimental results (Section 5) are outside the truncated text, so no specific figures for those can be given.

Methodology in Plain English

The authors reframe natural gradient variational inference as stochastic mirror descent over the expectation parameters of an exponential family, where the mirror map is the convex conjugate A* of the log-partition function. Because the Bregman divergence induced by A* equals the KL divergence between variational distributions, this reformulation lets them import tools from non-convex mirror descent theory.

They then compute the Hessian of the variational objective with respect to the expectation parameters. To do this for the log-likelihood term, they use Bonnet's and Price's gradient identities, which express gradients of expectations under a Gaussian in terms of expectations of derivatives of the integrand. Higher-order derivatives are controlled using Stein's Lemma, and relative smoothness is established by proving positive semidefiniteness of the appropriate matrices. Since these bounds only hold on a bounded region, they define a compact set Ω̃ (bounded means and bounded-away-from-zero, bounded-above variances) and design the algorithm to project back onto it after each step — a projection that decomposes into simple element-wise clipping in the mean-field Gaussian case.

For the global convergence result, they switch tools: they show the objective is convex after an invertible reparameterization (hidden convexity), which is inherited from the fact that strong convexity of a function transfers to its Gaussian expectation under the Cholesky parameterization. Hidden convexity then yields a PL inequality, which they combine with a mild stochastic gradient variance assumption (Assumption 4.1, previously introduced by Hanzely and Richtárik) to convert the PL condition into an actual function-value convergence rate.

Why This Matters

This work closes a gap between the empirical popularity of natural-gradient methods in variational inference and their theoretical justification. Prior guarantees required conjugate models, which exclude most practical likelihoods; the authors extend the analysis to the non-conjugate case that practitioners actually use.

Real-world applications implied by the paper's own framing include:

• Bayesian neural networks, where NGVI is used for training and prior work has shown geometry-aware updates to be faster and more stable.
• Probabilistic filtering and Kalman filtering, cited as a domain where NGVI is applied.
• Regression models with non-conjugate likelihoods, specifically logistic regression and Poisson regression, which are the running examples throughout the paper.
• Reinforcement learning, where natural policy gradient methods share the same natural-gradient structure, and where value functions are highly non-convex but admit structural properties analogous to those studied here.

Industry relevance: the projected algorithm has a cheap implementation (entry-wise clipping) in the mean-field Gaussian case, and the paper's stability motivation — that unprojected NGVI can diverge from the region where its guarantees hold — speaks directly to practitioners who observe unstable natural-gradient training runs.

Future Directions

• Removing or relaxing the bounded-domain requirement. The guarantees hold on Ω̃, and the optimal point must lie in that set (or at least in its relative interior for the PL inequality). Determining how to choose U and D adaptively, or proving guarantees on unbounded domains, is a natural open problem.

• Extending beyond mean-field Gaussian and diagonal covariance. The paper commits to mean-field parameterization with a diagonal Σ; whether the relative smoothness and hidden convexity results extend to full-covariance or structured variational families is not addressed by the provided content.

• Tighter constants and dimension dependence. The smoothness constant scales polynomially in d and in the size of the bounded set. Whether this is improvable — the paper notes a stronger, tight characterization exists for the univariate case in the appendix — remains open.

• Validating the theory experimentally. The provided content only shows an illustrative Poisson regression plot for SNGD; the paper's fuller experimental section (Section 5) would need to establish whether Proj-SNGD's projected steps translate into practical gains on larger models, since projection introduces bias.

Target Audience

Theoretical machine learning researchers working on variational inference, information geometry, or stochastic mirror descent; PhD-level readers in optimization who need convergence guarantees for non-Euclidean algorithms in non-convex settings. Practitioners using natural gradient methods for Bayesian deep learning or probabilistic modeling may benefit from the stability insight behind projection, but the paper is not written as an implementation guide and assumes substantial mathematical background.

Authors’ 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.