🔢 mathematics

SGD for Variational Inference: Tackling Unbounded Variance via Preconditioning and Dynamic Batching

This paper bridges the gap between stochastic optimization theory and Black-Box Variational Inference by proving the existence of ELBO solutions and establishing convergence guarantees for Minibatch Projected SGD with dynamic batching and preconditioning under the Blum-Gladyshev condition, which accounts for the unbounded variance inherent in BBVI gradients.

Original authors: Hippolyte Labarrière, Cesare Molinari, Silvia Villa, Lorenzo Rosasco

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Hippolyte Labarrière, Cesare Molinari, Silvia Villa, Lorenzo Rosasco

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the lowest point in a vast, foggy valley (this is the goal of Variational Inference: finding the best approximation of a complex probability distribution). You can't see the whole valley at once, so you have to take steps based on the ground beneath your feet.

In the world of machine learning, this is done using an algorithm called Stochastic Gradient Descent (SGD). Think of SGD as a hiker taking small steps downhill. Usually, we assume the ground is somewhat predictable: if you take a step, the slope doesn't change wildly.

However, in Black-Box Variational Inference (BBVI), the ground is treacherous. The "slope" (the gradient) you measure is incredibly noisy. In fact, the paper argues that the noise isn't just random; it gets wildly louder the further you are from the target. Standard hiking rules (mathematical assumptions) say the noise should stay within a certain limit, but here, the noise grows quadratically with your distance from the goal. It's like trying to walk down a hill where the wind gets exponentially stronger the further you are from the bottom.

The Paper's Solution: A Smarter Hiking Strategy

The authors, Hippolyte Labarrière and colleagues, propose two main tools to help the hiker survive this noisy terrain and actually reach the bottom:

1. The "Dynamic Batching" Backpack

Usually, a hiker looks at one patch of ground to decide where to step. In machine learning, this is called a "batch size of 1."

The Problem: If the ground is super noisy, looking at just one spot gives you a terrible idea of the slope.
The Fix: The paper suggests looking at more ground (increasing the batch size) as you get closer to the bottom or as you progress.
The Analogy: Imagine you are in a foggy forest. When you are far away, you might just peek through the trees. But as you get closer to the destination, you stop and scan a wider area to make sure you aren't walking off a cliff. By taking more samples (looking at more ground) over time, you smooth out the noise.

2. The "Preconditioning" Compass

Sometimes, the valley isn't just steep; it's shaped weirdly. Maybe it's a long, narrow canyon. If you take steps of the same size in every direction, you might bounce back and forth against the canyon walls instead of moving forward.

The Problem: The noise in the "location" (where you are) and the "scale" (how wide the distribution is) behaves differently. One might be very noisy, while the other is calm.
The Fix: The authors use a preconditioning matrix. Think of this as a smart compass that tells you, "Hey, the ground is slippery on the left, so take tiny steps there. The ground is solid on the right, so you can take big steps."
The Result: This balances the noise. It prevents the hiker from being thrown off course by the loudest part of the noise.

What Did They Prove?

The paper makes two major claims, backed by rigorous math:

The Destination Exists: Before you can find the bottom of the valley, you have to be sure a bottom actually exists. In many previous papers, researchers just assumed the solution existed. These authors proved that for a broad class of distributions (elliptic location-scale families, which include Gaussians and Laplace distributions), a solution definitely exists, provided the target function grows fast enough.
The Hiker Will Arrive: They proved that if you use their specific combination of Dynamic Batching (looking at more ground as you go) and Preconditioning (adjusting your step size based on the terrain), the algorithm is guaranteed to converge to the solution.
- They showed this works for both "finite time" (how fast you get there in a set number of steps) and "asymptotic" (what happens if you walk forever).
- Crucially, they showed that even though the noise is unbounded (it can get huge), these two techniques tame it effectively.

The "Real World" Test

To prove their theory works, they ran a simulation with a high-dimensional problem (200 dimensions, which is like a valley with 200 different directions to move).

The Result: The standard hiking method (plain SGD) was slow and shaky.
The New Method: The method using their smart compass (preconditioning) and dynamic backpack (batching) reached the bottom much faster and more smoothly.
Key Insight: The paper highlights that simply taking more steps isn't enough; you have to take smarter steps by adjusting how much data you look at and how you scale your movement.

Summary

In simple terms, this paper says: "We know the standard way of solving these complex probability problems is mathematically shaky because the noise gets too loud. We proved that a solution exists, and we showed that by using a 'smart compass' to balance the steps and a 'dynamic backpack' to gather more data as you go, you can reliably find the answer even in the noisiest, most chaotic environments."

Technical Summary: SGD for Variational Inference: Tackling Unbounded Variance via Preconditioning and Dynamic Batching

Problem Statement
Black-Box Variational Inference (BBVI) approximates intractable posterior distributions by optimizing the Evidence Lower Bound (ELBO) using Stochastic Gradient Descent (SGD). A fundamental challenge in this setting is that the stochastic gradients of the ELBO inherently exhibit unbounded variance. Specifically, the variance grows quadratically with the distance from the optimum, violating the standard assumption of uniformly bounded variance required for classical SGD convergence analysis. Instead, these gradients satisfy the weaker Blum-Gladyshev (BG) condition. While recent literature has established that BBVI gradients satisfy the BG condition, there remains a gap between the theoretical convergence guarantees for SGD under BG and the practical implementation of BBVI, particularly regarding the use of minibatching and preconditioning in general settings.

Methodology
The paper focuses on the optimization of the Negative ELBO within the broad class of elliptic location-scale families of distributions (which include Gaussian, Laplace, and Student-T densities). The authors analyze Minibatch Projected Stochastic Gradient Descent (PSGD) equipped with two key mechanisms:

Preconditioning: A diagonal scaling matrix $\Lambda$ is applied to the gradient updates to adapt step sizes to the specific variance characteristics of the location ( $\mu$ ) and scale ( $\Sigma$ ) parameters.
Dynamic Batching: The minibatch size ( $N_k$ ) is allowed to vary over iterations, either remaining constant or increasing, to manage the trade-off between gradient variance reduction and computational cost.

The analysis is conducted under the BG condition, where the expected squared norm of the gradient estimate is bounded by a quadratic function of the distance to the optimum. The authors derive convergence rates for both finite-time horizons and asymptotic regimes, considering both convex and non-convex settings (the latter under the assumption of bounded iterates).

Key Contributions
The paper makes three primary theoretical contributions:

Existence of a Minimizer: Under minimal growth assumptions on the log-joint density $\log q(x, z)$ , the authors prove the existence of a global minimizer for the Negative ELBO. This addresses a foundational property that is frequently assumed a priori in the literature but rarely proven for general elliptic families.
Convergence Guarantees for Minibatch PSGD: The authors establish comprehensive convergence rates for Minibatch PSGD with preconditioning and dynamic batching under the BG condition. They provide finite-time bounds and asymptotic rates (almost sure convergence) that demonstrate how dynamic batching combined with preconditioning can systematically handle the unbounded variance inherent in BBVI.
Generalization of SGD Theory: The work extends existing convergence theory for SGD under the BG condition to encompass minibatch estimators and preconditioning matrices. It explicitly quantifies the role of the minibatch size and the preconditioning matrix in mitigating variance, showing that these mechanisms allow for rigorous guarantees even in complex, unbounded variance settings.

Results

Finite-Time Rates: For a fixed budget of stochastic gradient evaluations $E$ , the authors derive convergence rates of order $O(1/\sqrt{E})$ (arbitrarily close to) by appropriately tuning the step size and batch size. They demonstrate that different strategies (decaying step size with constant batch, decaying step size with increasing batch, or constant step size with increasing batch) can achieve these rates, with the optimal choice depending on computational constraints (e.g., memory availability vs. parallelization capabilities).
Preconditioning Efficacy: Theoretical analysis shows that preconditioning (specifically scaling the step size for the $\Sigma$ parameter based on the kurtosis of the base distribution) significantly improves the convergence constants. Numerical experiments on Bayesian Logistic Regression with high-dimensional data ( $d=200$ ) confirm that adapted scaling and minibatching lead to faster convergence in terms of both the number of gradient evaluations and CPU time, particularly by reducing the frequency of costly Singular Value Decomposition (SVD) operations required for the projection step.
Asymptotic Convergence: Under summability conditions on the step sizes and batch sizes, the algorithm achieves almost sure convergence. The paper provides specific policies for step size and batch size schedules that satisfy these conditions while maintaining optimal convergence rates.

Significance
The paper bridges the gap between stochastic optimization theory and the practical demands of BBVI. By proving the existence of solutions and establishing convergence guarantees under the relaxed BG condition, the work provides a rigorous theoretical foundation for using standard SGD-based methods in BBVI without relying on restrictive assumptions like exponential families or conjugate models. The authors argue that dynamic batching and preconditioning are principled mechanisms for managing the unbounded variance that characterizes variational inference problems. While the analysis is currently limited to elliptic location-scale families and relies on convexity (or bounded iterates for non-convex cases), the framework offers a pathway for extending theoretical guarantees to broader inference tasks and optimization strategies in the future.