Convergence of projected stochastic natural gradient variational inference for various step size and sample or batch size schedules

Abstract: Stochastic natural gradient variational inference (NGVI) is a popular and efficient algorithm for Bayesian inference. Despite empirical success, the convergence of this method is still not fully understood. In this work, we define and study a projected stochastic NGVI when variational distributions form an exponential family. Stochasticity arises when either gradients are intractable expectations or large sums. We prove new non-asymptotic convergence results for combinations of constant or decreasing step sizes and constant or increasing sample/batch sizes. When all hyperparameters are fixed, NGVI is shown to converge geometrically to a neighborhood of the optimum, while we establish convergence to the optimum with rates of the form $\mathcal{O}\left(\frac{1}{T^ρ} \right)$, possibly with $ρ\geq 1$, for all other combinations of step size and sample/batch size schedules. These rates apply when the target posterior distribution is close in some sense to the considered exponential family. Our theoretical results extend existing NGVI and stochastic optimization results and provide more flexibility to adjust, in a principled way, step sizes and sample/batch sizes in order to meet speed, resources, or accuracy constraints.

• The paper demonstrates that projected stochastic NGVI with constant and adaptive schedules achieves geometric convergence near the optimum under variance constraints.
• It introduces a novel variance-like term in Bregman geometry that quantifies noise from stochastic gradient estimates to inform precise convergence bounds.
• Empirical results validate that increasing batch sizes or decreasing step sizes lead to exact convergence and improved asymptotic accuracy in Bayesian inference.

Convergence Guarantees for Projected Stochastic Natural Gradient Variational Inference

Introduction

The paper "Convergence of projected stochastic natural gradient variational inference for various step size and sample or batch size schedules" (2604.00683) systematically investigates the theoretical and practical convergence properties of stochastic Natural Gradient Variational Inference (NGVI) within exponential families. Central to the analysis is the interplay between step size and sample/batch size schedules, in both fixed and adaptive regimes, with a focus on projection-based algorithms in Bregman geometry. The study formulates new non-asymptotic convergence results for NGVI in scenarios where either gradient computation is based on intractable expectations or large finite sums, crucial for modern large-scale Bayesian inference.

Background: NGVI, Mirror Descent, and Bregman Geometry

NGVI exploits the Riemannian geometry of the parameter space by preconditioning updates with the Fisher Information matrix. In exponential families, this is equivalent to mirror descent with Bregman divergence as the proximity measure, leveraging a duality between natural and expectation parameters. The projected stochastic version considered here restricts iterations to convex subsets via Bregman projection, ensuring algorithmic stability and feasibility under various parameter constraints (e.g., positive-definite covariance matrices for Gaussians).

Two major sources of stochasticity in practice are considered: Monte Carlo approximation (sampling from the variational distribution) and data subsampling (random batches in large datasets). The analysis relies on relative strong convexity and relative smoothness, generalized to the non-Euclidean geometry induced by the exponential family.

Algorithmic Formulation and Convergence Rates

The core algorithm operates by alternating mirror descent steps (using stochastic gradient estimators) and Bregman projections. The main theoretical contributions are non-asymptotic convergence bounds under various scheduling scenarios:

• Constant Step Size and Batch Size: Iterates exhibit geometric convergence to a neighborhood around the optimum. The size of this neighborhood scales inversely with batch size and directly with step size.
• Increasing Batch Size or Decreasing Step Size: Exact convergence to the optimum is achieved, with rates of $\mathcal{O}\left(\frac{1}{T^{\rho}}\right)$ for $\rho \geq 1$ depending on the scheduling strategy. For example, polynomially increasing sample sizes ($N_t = (t+1)^\gamma$) yield a $\mathcal{O}\left(\frac{1}{T^\gamma}\right)$ convergence rate.

A novel variance-like term, tailored for mirror descent in Bregman geometry, quantifies the "noise" induced by stochastic gradient estimates and directly informs the convergence bounds.

Figure 1: Demonstration of geometric convergence for fixed step and batch size schedules, showing the mean Bregman divergence reduction across iterations.

Posterior Recoverability and Extension Conditions

A key theoretical advance is the introduction of the Linearly Extended Recoverability Condition (LERC), which substantially generalizes recoverability requirements found in prior work (e.g., $\pi \in \mathcal{Q}$). LERC allows posterior distributions to belong to a richer exponential family connected by a linear operator, enabling the derived rates to apply even when the posterior is not exactly recoverable by variational family—an important relaxation for practical settings.

Under LERC, the variational objective retains relative strong convexity and smoothness, ensuring validity of the main convergence results.

Practical Implementation: Estimator Variance Bounds

The paper details practical stochastic gradient estimators:

• Bonnet–Price Estimators: Efficient, unbiased gradient estimators for Gaussian variational families, exploiting closed-form derivatives of expectations.
• Subsampling-Based Estimators: For models with large datasets and conjugate structure (e.g., Bayesian linear regression), the variance of estimator noise can be tightly controlled.

Explicit variance bounds are established for both cases, confirming theoretical rates and providing actionable guidance for practitioners regarding sample size and step size selection.

Figure 2: Empirical behavior with $N=100$ and multiple step size values; convergence neighborhood shrinks with smaller $\eta$.

Empirical Results

Extensive numerical experiments validate the theory, comparing schedules with constant/increasing sample sizes and constant/decreasing step sizes in Gaussian and regression settings (with real and synthetic data). The findings include:

• Fixed schedules yield fast geometric convergence but plateau within a variance-controlled neighborhood; increasing batch sizes or decreasing step sizes enable tighter convergence to the optimum.
• Schedule choice affects transient performance and trade-offs between computational costs and accuracy. Larger batch sizes improve asymptotic accuracy but increase computational expense.

Projection operations are shown to stabilize optimization in non-log-concave regimes (e.g., heavy-tailed posteriors, robust regression), preventing pathological parameter values such as degenerate covariances.

Figure 3: Comparative performance across NGVI schedules by iterations; adaptive scheduling accelerates precise convergence.

Figure 4: Schedule $\eta_t = 5\times 10^{-3}$; impact on convergence speed and neighborhood size as predicted by theory.

Related Work and Theoretical Implications

Compared to prior literature ([hanzely2021], [wu2024]), the paper advances tighter variance definitions and broader applicability by decoupling convergence guarantees from explicit posterior recoverability and extending results to Bregman-projected domains. The theoretical framework seamlessly connects stochastic mirror descent advances in optimization to variational inference, supplying a robust foundation for large-scale, non-conjugate Bayesian inference.

The results open pathways for algorithmic scheduling strategies driven by convergence rate requirements, resource constraints, and desired approximation accuracy, formalizing the trade-off landscape for practitioners.

Future Directions and Applications

The methodology is amenable to generalized exponential families, such as $q$-exponential and $\lambda$-exponential classes, enabling heavy-tailed and mixture-based approximate inference. Integrating these families can facilitate robust Bayesian modeling and extend the benefits of NGVI to broader classes of models where standard conjugacy does not hold.

Anticipated developments include adaptive learning algorithms that dynamically tune step size and batch size during training, leveraging real-time variance estimates to optimize convergence behavior in heterogeneous data environments.

Conclusion

This paper provides a rigorous, general framework for the convergence analysis of projected stochastic NGVI. By leveraging Bregman geometry, mirror descent duality, and novel variance analysis, it delivers actionable theoretical guarantees across a spectrum of scheduling regimes. The practical implications span algorithmic design, computational efficiency, and model robustness, establishing a foundation for principled variational inference in contemporary statistical and machine learning applications.