Stability of Finite-Batch Particle Mean-Field Variational Inference Beyond Strong Convexity

Technical Summary: Stability of Finite-Batch Particle Mean-Field Variational Inference Beyond Strong Convexity

Problem Statement
The paper addresses the stability and non-asymptotic convergence of the Finite-Batch Particle Algorithm for Mean-Field Variational Inference (PAVI). While Mean-Field Variational Inference (MFVI) approximates a target density $p(x) \propto e^{-V(x)}$ by minimizing relative entropy over product measures, practical implementations rely on particle approximations, stochastic batch evaluations of projected drifts, and time discretization.

Existing analyses, such as those by Du et al., typically assume the target potential $V$ is globally strongly convex. This assumption ensures contractivity of the dynamics but is often violated in modern applications. This paper investigates the stability of PAVI when $V$ is globally smooth (Lipschitz gradient) but not necessarily strongly convex. The challenge lies in quantifying how the loss of contractivity (due to non-convex regions) affects the error bounds of the fully discrete, finite-batch algorithm.

Methodology and Framework
The authors model the PAVI algorithm as a fully discrete stochastic approximation of the projected Wasserstein dynamics. The core methodological innovation is the introduction of a curvature defect metric to quantify the failure of strong convexity.

1. Curvature Defect: Instead of assuming $\nabla^2 V \succeq \alpha I$, the paper defines a defect function:
   $$d_\alpha(x, y) = \left[ \alpha \|x - y\|^2 - \langle \nabla V(x) - \nabla V(y), x - y \rangle \right]_+$$
   This measures the additive loss in the one-step Euler contraction estimate. The analysis assumes a uniform bound $d_\alpha(x, y) \le \beta$, which allows for non-convexity as long as the "negative curvature" is bounded.

2. Stationary Comparison Array: To derive non-asymptotic bounds, the authors construct a stationary comparison process. This process consists of an array of particles whose row-wise marginals are optimal couplings to the true MFVI minimizer $q^\star$. Crucially, the product of the row empirical measures of this stationary array is a random measure, distinct from the deterministic product law $q^\star$. The analysis explicitly controls the discrepancy between the projected drift evaluated at this random empirical law and the drift at $q^\star$.

3. Decomposition of Errors: The proof utilizes a one-step recursion that separates the total error into distinct components:

   • Initialization error (transient contraction).
   • Product-empirical approximation error (finite particle number $N$).
   • Finite-batch drift error (stochastic sampling of size $B$).
   • Time discretization error (step size $h$).
   • Accumulated curvature defects along the coupled trajectories.

Key Contributions and Results

• Non-Asymptotic Stability Bound: The primary result (Corollary 7.3) provides a Wasserstein stability bound for the finite-batch scheme. Under the uniform defect bound $d_\alpha \le \beta$, the root-mean-square Wasserstein distance between the particle iterate $q_{X_n}$ and any MFVI minimizer $q^\star$ satisfies:
  $$\left( \mathbb{E} W_2^2(q_{X_n}, q^\star) \right)^{1/2} \lesssim (1 - c\alpha h)^{n/2} \left( \mathbb{E} W_2^2(q_{X_0}, q^\star) \right)^{1/2} + \left(1 + \frac{\kappa}{\alpha}\right)\varepsilon_N(q^\star) + \sqrt{\frac{\beta}{\alpha}} + \sqrt{\frac{h\Gamma^\star}{\alpha B}} + \sqrt{\frac{L}{\alpha}\sqrt{h(h\Gamma^\star + m)}}$$
  Here, $\varepsilon_N$ is the empirical approximation error, $\kappa$ measures the sensitivity of the projected drift to perturbations in the product law, and $\Gamma^\star$ is the second moment of the gradient under $q^\star$.

• Separation of Parameters: The bound explicitly separates the effects of particle number $N$, batch size $B$, and step size $h$. It demonstrates that increasing $B$ is beneficial only until the batch error becomes comparable to the particle and discretization errors.

• Existence and Moment Bounds: The paper establishes the existence of MFVI minimizers under the uniform defect condition and derives explicit moment bounds for the stationary marginals without assuming log-concavity. It proves that the uniform defect implies quadratic coercivity of the potential. The authors explicitly note that while existence is guaranteed, uniqueness is not implied; the set of minimizers may contain multiple distinct elements when $\beta > 0$.

• Dimension-Independent Sensitivity: The authors introduce coordinatewise defect estimates and show that under sparse interaction structures, the sensitivity constant $\kappa$ can be dimension-independent, mitigating the curse of dimensionality in the error bound.

• Non-Convex Benchmark: A novel arbitrary-dimensional smooth non-convex benchmark is constructed with a closed-form unique MFVI minimizer. This allows for the direct measurement of algorithmic errors without contamination from optimization uncertainty in the reference solution.

• Limitations on Drift Growth: The paper clarifies that the explicit Euler scheme used in PAVI is unstable for potentials with superlinearly growing gradients (polynomial growth), even if they are strongly convex. It argues that tamed or implicit schemes would be required for such regimes, as standard taming would destroy the unbiasedness of the finite-batch estimator essential to the proof.

Significance and Claims
The paper provides a non-asymptotic stability analysis of finite-batch particle MFVI that relaxes the strong convexity assumption. By quantifying the "curvature defect," the work separates the numerical requirements of Lipschitz continuity from the geometric mechanisms of contraction. This analysis builds upon and extends prior work on local contraction and non-convex Langevin dynamics.

The authors emphasize that their results are modest regarding convergence to a specific minimizer when $\beta > 0$; in such cases, the algorithm remains stable within an $O(\sqrt{\beta/\alpha})$ neighborhood of the set of minimizers, rather than converging to a single point. The analysis highlights that the non-uniqueness of the MFVI solution in non-convex settings is quantitatively controlled by the same defect scale.

The numerical experiments validate the theoretical scaling laws, showing that the batch error scales as $B^{-1/2}$ and the particle error scales near $N^{-1/2}$ for smooth benchmarks, while the terminal error increases monotonically with the magnitude of the curvature defect. The paper concludes that the curvature-defect viewpoint offers a rigorous framework for understanding the trade-offs between accuracy, computational cost, and the geometry of the target potential in variational inference.