RELTA-SGLD: Relative-Growth Localized Taming for Nonconvex Stochastic-Gradient Langevin Learning

AI summary

Overview

Research area: Stochastic optimization and sampling theory — specifically tamed stochastic-gradient Langevin dynamics (SGLD) for nonconvex learning, with non-asymptotic Wasserstein-accuracy guarantees.

Technical level: Advanced. The paper relies on Foster–Lyapunov drift arguments, Lyapunov–minorization ergodicity theory, and Wasserstein discretization analysis.

One-sentence scope: The paper proposes RELTA-SGLD, a taming scheme whose denominator is calibrated by the relative growth of the quadratic stochastic-gradient term $Q$ to the inward drift $I$ and localized by a threshold, and proves first-order stationary accuracy in both $W_1$ and $W_2$ along with supporting experiments on Fashion-MNIST and a quartic potential.

What This Paper Is About

Taming schemes make explicit SGLD updates stable when the stochastic-gradient oracle grows faster than linearly, but standard tamed SGLD (e.g., TUSLA) uses a globally active $\sqrt{\lambda}$-scale denominator that suppresses the learning drift everywhere, including regions where no stabilization is needed. This paper asks how much taming the one-step Lyapunov balance actually requires and where it should turn on, and answers by designing a denominator from the tail ratio $Q/I$ that is inactive below a calibrated threshold. The goal is to retain stability in the far tail while leaving ordinary learning dynamics nearly untouched, and to obtain sharper stationary Wasserstein step-size orders than comparable stochastic-gradient tamed schemes.

Key Contributions

1. A relative-growth, localized taming rule (RELTA-SGLD). The denominator is $D_\lambda(\theta)=\sqrt{1+(K\lambda S_{s_0}(\theta))^2}$ with threshold score $S_{s_0}(\theta)=\left[|\theta|^{2r}-s_0\right]_+$, where $s_0 \ge 0$ is a threshold calibrated in practice from a train-only untamed pilot run (the empirical $0.8$-quantile $R_0$ of recorded parameter-norm snapshots, with $s_0=R_0^{2r}$; this choice is used in both Fashion-MNIST experiments). The design replaces the $\sqrt{\lambda}$-scale of TUSLA with a $\lambda$-scale calibrated by the exterior $Q/I$ relative-growth law.

2. A stability theory for the quadratic and polynomial moment cases. Theorem 3.1 establishes a quadratic Foster–Lyapunov drift whenever $K > C_{\rm rel}/2$; Theorem 3.2 extends this to every fixed even order $2\ell$ under the explicit coefficient condition $K > \eta/2$, giving $R_\lambda V_{2\ell}(\theta) \le (1-a_\ell \lambda)V_{2\ell}(\theta)+b_\ell \lambda$ with $V_{2\ell}(\theta)=1+|\theta|^{2\ell}$. Corollary 3.3 gives uniform moment bounds and Corollary 3.5 gives a unique invariant law $\pi_\lambda^{\rm SG}$ with $\int V_{2\ell},\mathrm{d}\pi_\lambda^{\rm SG} \le b_\ell/a_\ell$.

3. First-order stationary accuracy in both $W_1$ and $W_2$. Under explicit nonconvex SGLD assumptions, RELTA achieves first-order stationary accuracy simultaneously in $W_1$ and $W_2$, improving the half-order $W_1$ and quarter-order $W_2$ step-size accuracies available for comparable stochastic-gradient tamed schemes. The finite-time excess-risk complexity is $\widetilde{O}(\varepsilon^{-1})$, up to the fixed-temperature Gibbs bias.

4. Two structural improvements over the $\sqrt{\lambda}$-scale denominator, plus experimental validation. Before the tail regime, the leading $O(\sqrt{\lambda})$ attenuation is replaced by second-order $O(\lambda^2)$ relative attenuation and becomes exactly zero where $S_{s_0}(\theta)=0$; in the far tail, the one-step radial return is $-\eta\sqrt{\lambda},\theta$ for TUSLA versus $-\frac{\eta}{K}\theta$ for RELTA, a coefficient independent of $\lambda$. Experiments on Fashion-MNIST under active stabilization pressure show RELTA improves the mean learning metrics over untamed SGLD and TUSLA and remains competitive with a tuned AdamW reference; in an ordinary-training regime it maintains nearly untamed learning dynamics; a stationary sampling experiment in a quartic potential confirms the predicted first-order $W_1/W_2$ behavior and its separation from the $\sqrt{\lambda}$-scale baseline.

Main Findings

• TUSLA trades stability for over-suppression. The TUSLA denominator $D_{\mathrm{TUSLA},\lambda}(\theta)=1+\sqrt{\lambda},|\theta|^{2r}$ controls the quadratic Euler term largely in isolation, does not exploit the inward drift at the same state, and produces an $O(\sqrt{\lambda})$ relative perturbation on bounded sets where stabilization is not needed.

• The far-tail return degenerates under TUSLA but not under RELTA. As $\lambda \to 0$, the TUSLA radial return per step is of order $\sqrt{\lambda}$, whereas the $Q/I$-calibrated return remains linear in $|\theta|$ with a coefficient independent of $\lambda$.

• The required taming scale is set by relative, not absolute, growth. Lemma 2.1 gives $I(\theta) \ge c_I|\theta|^{2r+2}$ and $Q(\theta) \le C_Q|\theta|^{4r+2}$ outside a large ball, so balancing the two state-dependent terms of the one-step Lyapunov identity requires only a $\lambda|\theta|^{2r}$-scale denominator rather than TUSLA's larger $\sqrt{\lambda}|\theta|^{2r}$ scale. In a radial polynomial class, $Q(\theta)/I(\theta) \asymp S_{s_0}(\theta)$, so $\lambda S_{s_0}$ is the matching tail scale.

• Attenuation is second order in the localized regime. For the conventional and relative-growth denominators the bounds are $1-1/D_\lambda^{\mathrm T}(\theta) \le \sqrt{\lambda}|\theta|^{2r}$ and $1-1/D_\lambda^{\mathrm{QI}}(\theta) \le \frac{K^2}{2}\lambda^2 S_{s_0}(\theta)^2$ respectively; RELTA leaves the update completely unchanged where $S_{s_0}(\theta)=0$.

• Polynomial moment stability holds at every fixed even order. Theorem 3.2 and Corollary 3.3 give $\mathbb{E}V_{2\ell}(\theta_n) \le (1-a_\ell\lambda)^n,\mathbb{E}V_{2\ell}(\theta_0)+b_\ell/a_\ell$ and $\sup_{n\ge 0}\mathbb{E}V_{2\ell}(\theta_n) < \infty$, and Corollary 3.4 gives finite-time moment control for the mean-oracle chain (3.10) uniformly over $\lambda \in (0,\lambda_0]$ and $n\lambda \le T$.

• Ergodicity via a small-set argument. Corollary 3.5 shows the RELTA chain admits a unique invariant probability measure; the ball $C_R={|\theta|\le R}$ is a small set because the Gaussian increment $2\lambda\beta^{-1}I_d$ produces densities with a common positive lower bound on a fixed bounded ball.

• Oracle assumptions are weaker than exact-gradient sampling. The analysis requires only the averaged drift to be $C^2$, allows the samplewise oracle to be nonsmooth, and uses radial tail dissipativity rather than the pairwise convexity at infinity assumed by mTULA. The likelihood ratio and entropy analyses of tamed Langevin samplers in the literature are described as complementary to this setting.

• Experimental outcomes. On Fashion-MNIST under active stabilization pressure, RELTA improves the mean learning metrics over untamed SGLD and TUSLA and stays competitive with a tuned AdamW reference. In the ordinary-training regime, its lighter localized denominator minimally perturbs the original update. The quartic-potential sampling experiment confirms the predicted first-order $W_1/W_2$ behavior and separation from the $\sqrt{\lambda}$-scale baseline. The specific numeric values of these metrics are not reported in the available content.

Methodology in Plain English

The authors start from the one-step Lyapunov identity for the weighted function $V_2(\theta)=1+|\theta|^2$, which contains a negative inward term $-2\lambda I(\theta)/D_\lambda(\theta)$, a positive Euler term $\lambda^2 Q(\theta)/D_\lambda(\theta)^2$, and a state-independent Gaussian term $2\lambda\beta^{-1}d$. Because the positive Euler term carries $D_\lambda^{-2}$ while the inward term carries only $D_\lambda^{-1}$, stability reduces to making the denominator large enough that the inward term absorbs the Euler remainder. The paper writes this condition as $\lambda Q(\theta)/D_\lambda(\theta) \lesssim I(\theta)$ and, from an upper bound on $Q$ in terms of $I$ and the tail score $S_{s_0}$, derives the sufficient scale $D_\lambda(\theta)\gtrsim \lambda S_{s_0}(\theta)$. A square-root form smooths the transition between the inactive and far-tail regimes. The threshold is set from the empirical $0.8$-quantile of parameter norms recorded in an untamed pilot run, so no untamed tail excursions contaminate the calibration.

The theory then proceeds in three layers: a quadratic drift bound (Theorem 3.1) obtained by bounding $Q$ relative to $I(1+S_{s_0})$ and choosing the coefficient so that $2-C_{\rm rel}/K-C_{\rm rel}\lambda_0$ stays positive; an extension to all fixed even orders $2\ell$ (Theorem 3.2) using the effective radial coefficient $a_\lambda(t)=\lambda t^{2r}/D_{\lambda,R_0}(t)$, the condition $K>\eta/2$ that makes $|1-\eta/K|<1$, and a Young-inequality treatment of the Gaussian increment; and a small-set plus drift argument yielding a unique invariant law. Separately, a mean-oracle chain with the same denominator but the conditional mean $h(\theta)=\mathbb{E}[H(\theta,Z)]$ is introduced to isolate stochastic-gradient error from taming and time-discretization error. The final Wasserstein analysis compares the RELTA chain against the semigroup $P_t$ of the target diffusion $\mathrm{d}Z_t = b(Z_t)\mathrm{d}t + \sqrt{2\beta^{-1}},\mathrm{d}B_t$ with $b=-h$, using the same fixed-time approximation estimates for the $W_2$ and $W_1$ arguments and the SGLD target structure to transfer finite-time estimates to a stationary first-order bound. The detailed $W_1/W_2$ derivation in Section 4 is not reproduced in the available content.

Why This Matters

Tamed SGLD is a practical fix for a practical failure mode — training with gradient oracles that grow faster than linearly, where Euler-type SGLD updates can lose moment control or diverge numerically. The paper's contribution is showing that this fix can be applied more surgically: turn it on where the tail actually is, and calibrate its strength from the drift available at the same state rather than from the absolute size of the gradient.

Impact on research. The results sharpen stationary Wasserstein step-size orders for stochastic-gradient tamed SGLD (first order in both $W_1$ and $W_2$, versus the half-order $W_1$ and quarter-order $W_2$ available for comparable schemes) and supply a design principle — the $Q/I$ balance — that can be transferred to other tamed or clipped samplers. The assumptions are tailored to SGLD rather than exact-gradient sampling, using radial tail dissipativity instead of pairwise convexity at infinity, and requiring only $C^2$ averaged drift with potentially nonsmooth samplewise oracles.

Real-world applications.

• Large-scale nonconvex model training using mini-batch stochastic-gradient oracles, where stability under heavy-tailed or superlinear gradient noise matters.
• Bayesian posterior sampling with Langevin-type algorithms, where stationary accuracy controls the discretization bias of the sampled distribution.
• Robust training of neural networks whose gradients blow up in some parameter regions, the setting originally motivating TUSLA.
• Any pipeline that needs an untamed update preserved in ordinary operating regions while still guarding against rare large-norm excursions.

Industry relevance. The threshold is calibrated from a train-only pilot run and an empirical $0.8$-quantile of recorded parameter norms, which is a cheap, automatable procedure. The claim of near-untamed dynamics in ordinary training regimes is directly relevant to teams that want stability insurance without paying an accuracy or convergence-speed tax on the common case.

Future Directions

1. Closing the truncation in the Wasserstein analysis. The available content ends partway through Section 4; the full constants and the complete first-order $W_1$ and $W_2$ derivations, and how the SGLD target structure supplies the long-time transfer, remain to be worked through.

2. Tighter necessity statements for the denominator scale. The authors present $D_\lambda(\theta)\gtrsim \lambda S_{s_0}(\theta)$ as a sufficient balance condition, with sharpness shown only for the radial polynomial-growth class. Whether the $\lambda S_{s_0}$ scale is necessary for broader oracle families is left open.

3. Extension beyond the radial polynomial model. The analysis assumes the radial-leading oracle decomposition with $|\mathsf{G}(\theta,u)|\le\mathcal{K}(u)(1+|\theta|^q)$, $q<2r+1$. Whether the localized relative-growth principle extends to nonsmooth or non-radially dominant oracles is not settled.

4. Sharper comparison with functional-inequality and entropy analyses. The paper describes recent KL, total-variation, and Wasserstein results for tamed Langevin samplers under exact-drift formulations as complementary. Reconciling those sharper sampling guarantees with the explicit nonconvex SGLD setting — and with the fixed-temperature Gibbs bias that bounds the $\widetilde{O}(\varepsilon^{-1})$ excess-risk complexity — is a natural next step.

Target Audience

Researchers in stochastic optimization, Langevin sampling, and machine-learning theory who work on nonconvex learning guarantees, tamed discretization schemes, or Wasserstein-accuracy bounds for stochastic-gradient methods. Practitioners of large-scale neural-network training will benefit from the algorithm design and the threshold-calibration procedure, though the stability and accuracy proofs require comfort with Lyapunov drift conditions, minorization-based ergodicity, and stochastic differential equation approximations.

Authors’ abstract

We introduce RELTA-SGLD, a taming scheme that stabilizes superlinear stochastic-gradient updates while reducing unnecessary suppression of the original learning drift. A threshold determines where the taming turns on, while a relative-growth principle derived from the one-step Lyapunov stability condition determines the required taming strength. Together, they produce a lighter $λ$-scale denominator and preserve a nonvanishing far-tail return. As a consequence, we prove polynomial moment stability and first-order stationary accuracy in both $W_1$ and $W_2$ for nonconvex SGLD with superlinearly growing stochastic-gradient oracles, improving the corresponding half-order and quarter-order bounds for comparable stochastic-gradient tamed schemes. On Fashion-MNIST under active stabilization pressure, RELTA improves the mean learning metrics over both untamed SGLD and TUSLA and remains competitive with a tuned AdamW reference. In an ordinary-training regime, its lighter localized denominator reduces unnecessary perturbation of the original update and maintains nearly untamed learning dynamics.