📊 statistics

The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity

This paper introduces the Subgradient Tamed Unadjusted Langevin Algorithm (SG-TULA), a novel sampling method that handles non-smooth, non-convex potentials with superlinear gradient growth without smoothing, achieving improved non-asymptotic convergence bounds and demonstrating competitive performance in pretraining LLMs compared to standard optimizers like AdamW and Muon.

Original authors: Iosif Lytras, Nikolaos Makras, Sotirios Sabanis

Published 2026-08-07

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Iosif Lytras, Nikolaos Makras, Sotirios Sabanis

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 deepest valley in a vast, foggy mountain range. This is a problem that shows up everywhere in science and technology, from training the brains of artificial intelligence to figuring out the structure of crystals. The goal is to find the absolute lowest point, the "global minimum," where the energy is lowest and the system is most stable.

To do this, scientists often use a method inspired by physics called the "Langevin algorithm." Think of it like sending a hiker into the fog. The hiker has a map (the gradient) that tells them which way is downhill, but the map is a bit noisy. So, the hiker takes steps downhill but also gets a little push from a random gust of wind (noise). This helps them bounce out of small, shallow dips (local minima) and keep searching for the deepest valley. Usually, this works great if the mountains are smooth and the slopes are predictable.

But what happens if the terrain is a nightmare? What if the ground is jagged and sharp (non-smooth), the slopes get steeper and steeper the further you go (superlinear growth), and the landscape is full of weird, confusing bumps that aren't actually valleys (non-convex)? In these chaotic conditions, the standard hiker's map often breaks. The steps can become so huge that the hiker flies off the map entirely, or they get stuck on a sharp corner where the map gives no direction. This is the exact problem that modern AI training faces: the "mountains" of neural networks are rough, steep, and full of traps.

This paper introduces a new, tougher hiker named SG-TULA (Subgradient Tamed Unadjusted Langevin Algorithm). Instead of trying to smooth out the jagged rocks or pretending the steep cliffs don't exist, SG-TULA carries a special pair of "taming boots." These boots automatically adjust the hiker's stride. If the slope gets too steep, the boots shorten the step to prevent a fall. If the ground is sharp, the boots find a safe path along the edge. The authors prove mathematically that this new hiker can successfully navigate these chaotic, jagged, and infinitely steep landscapes to find the deepest valley, all while providing a precise map of how long the journey will take.

The Problem: When the Map Breaks

Most optimization algorithms, the tools used to train AI, rely on a few comforting assumptions. They assume the landscape is smooth (no sharp edges), the slopes aren't too crazy (linear growth), and the terrain is generally bowl-shaped (convex). But real-world AI models, like the ones powering chatbots, break all these rules. Their "landscapes" are full of sharp corners (from activation functions like ReLU), the slopes can explode to infinity (superlinear growth), and the terrain is a mess of hills and valleys (non-convex).

When you try to use a standard algorithm on this kind of terrain, it's like sending a hiker with giant, rigid boots into a field of cactus and cliffs. The hiker might take a step that is too big, fly off the edge, or get stuck because the ground is too rough to get a grip. To fix this, people have tried "smoothing" the ground (making the cactus soft) or "clipping" the steps (forcing the hiker to walk slowly). But smoothing is computationally expensive, and clipping is a blunt instrument that can distort the path.

The Solution: SG-TULA

The authors propose SG-TULA, a new algorithm designed specifically for this messy reality. Instead of smoothing the ground or blindly clipping steps, SG-TULA uses a technique called "taming."

Imagine the hiker's stride is controlled by a smart leash. As the hiker approaches a cliff where the slope gets dangerously steep, the leash gently but firmly pulls back, shortening the step just enough to keep the hiker safe. If the ground is jagged, the hiker uses a "subgradient," which is like feeling the ground with a cane to find a safe direction even when the map is broken.

The key innovation is that this "taming" happens automatically and continuously. It doesn't require a pre-set limit (like a speed bump) that might stop the hiker from moving fast enough when it's safe. Instead, the step size rescales itself based on how steep the ground is right at that moment. This allows the algorithm to handle:

Non-smoothness: Jagged edges where the slope changes instantly.
Superlinear growth: Slopes that get infinitely steep.
Non-convexity: Landscapes full of confusing bumps and false valleys.

What They Found

The authors didn't just invent this algorithm; they proved it works. They derived mathematical guarantees showing that SG-TULA will eventually find the target distribution (the deepest valley) and stay there.

The Speed: They calculated exactly how fast the algorithm converges. In the worst-case scenario of these messy landscapes, the error shrinks at a rate of roughly the step size to the power of 1/4. While this sounds slow, it is actually a significant improvement over previous methods for this specific type of problem, which were often slower or didn't work at all.
The Constants: They tracked every single number in their formulas, showing exactly how the number of dimensions (the size of the problem) and the "temperature" (how much noise is added) affect the speed.
Real-World Test: To prove it wasn't just theory, they tested SG-TULA on a small language model called "nanochat" (a simplified version of the famous GPT-2). They compared it against two of the best standard optimizers, AdamW and Muon.
- At a depth of 12 layers, SG-TULA achieved the best results in terms of "bits per byte" (a measure of how well the model predicts text) and a score called CORE.
- At a deeper 24 layers, it stayed competitive, performing just as well as the top contenders.
- Crucially, they did this without the benefit of "scaling laws" (rules of thumb that usually help tune these models), showing the algorithm is robust even when you have to tune it by hand.

Why It Matters

This paper is a big deal because it bridges a gap between theory and practice. For years, we've known that AI training happens in messy, non-smooth, superlinear environments, but our mathematical guarantees were stuck in a "perfect world" where everything is smooth and convex. SG-TULA provides the first rigorous proof that we can navigate this chaos safely and efficiently.

It shows that we don't need to pretend the world is smooth to solve our problems. We can build tools that respect the jagged, steep, and complex nature of real data. By proving that this "tamed" approach works, the authors give us a new, reliable way to train the next generation of AI models, ensuring they can find the best solutions even when the path is full of sharp turns and steep drops.

Technical Summary: The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity

1. Problem Statement

The paper addresses the challenge of sampling from target distributions $\pi_\beta(dx) \propto e^{-\beta u(x)} dx$ and solving the associated optimization problem $\theta^\star \in \arg\min u(\theta)$ in regimes where the potential $u: \mathbb{R}^d \to \mathbb{R}$ simultaneously exhibits three difficult properties:

Non-smoothness: The potential is not differentiable everywhere (e.g., due to ReLU activations or $\ell_1$ penalties), requiring the use of subgradients $h \in \partial u$ .
Superlinear Gradient Growth: The magnitude of the subgradient grows polynomially ( $|h(x)| \leq m + K|x|^p$ for $p \geq 1$ ), which causes standard Euler–Maruyama discretizations of the Langevin diffusion to become numerically unstable (moments diverge in finite time).
Non-convexity: The potential is not globally convex, though it satisfies specific structural conditions (semi-convexity and strong convexity at infinity).

Existing literature typically addresses these issues in isolation (e.g., smoothing for non-smoothness, taming for superlinear growth, or functional inequalities for non-convexity). The paper argues that a non-asymptotic guarantee for the simultaneous presence of all three has been missing.

2. Methodology

The authors propose the Subgradient Tamed Unadjusted Langevin Algorithm (SG-TULA). This is an explicit Euler–Maruyama discretization of the overdamped Langevin diffusion that operates directly on measurable subgradient selections $h \in \partial u$ without relying on computationally expensive smoothing procedures (such as Moreau–Yosida regularization).

The Algorithm:
The iterates are defined as:
$\theta^\lambda_{n+1} = \theta^\lambda_n - \lambda h^\lambda(\theta^\lambda_n) + \sqrt{2\lambda \beta^{-1}} \xi_{n+1}$
where $\xi_{n+1}$ are i.i.d. standard Gaussian vectors. The core innovation is the taming function $h^\lambda$ , which stabilizes the drift:
$h^\lambda(x) = \frac{h(x) - \mu x}{1 + \lambda^{1/2}|h(x) - \mu x|} + \mu x$
This mechanism rescales the subgradient continuously based on its magnitude, preventing explosion without requiring a fixed threshold (unlike gradient clipping). The paper analyzes both a uniform-wise variant (SG-TULA $_u$ ) and a coordinate-wise variant (SG-TULA $_c$ ).

Theoretical Framework:
The analysis relies on three primary assumptions on $u$ :

A1 (Semi-convexity): $u + \frac{L}{2}|\cdot|^2$ is convex. This allows for gradient discontinuities on sets of measure zero (e.g., ReLU) while providing a one-sided Lipschitz bound.
A2 (Strong Convexity at Infinity): Outside a compact set, the potential is strongly convex, ensuring dissipativity and contraction of the diffusion.
A3 (Polynomial Growth): Subgradients satisfy $|h(x)| \leq m + K|x|^p$ .

The proof strategy involves:

Establishing the existence and uniqueness of the invariant measure.
Proving a Log-Sobolev inequality (LSI) for the target measure $\pi_\beta$ under these assumptions.
Deriving exponential contraction rates for the continuous SDE towards $\pi_\beta$ .
Establishing uniform moment bounds for the discrete algorithm via the taming mechanism.
Combining discretization error estimates with contraction rates to bound the Wasserstein-2 distance.

3. Key Contributions and Results

Non-Asymptotic Sampling Bounds:
The paper derives explicit non-asymptotic convergence bounds in the Wasserstein-2 distance ( $W_2$ ). For stepsize $\lambda < 1/(4\mu)$ and $n$ iterations:
$W_2(\pi_\beta, \mathcal{L}(\theta^\lambda_n)) \leq 2 C_E \lambda^{1/4} + e C_D e^{-n\lambda/T_0}$

Rate: The discretization error converges at a rate of $O(\lambda^{1/4})$ .
Constants: All constants ( $C_E, C_D, T_0$ ) are tracked explicitly in terms of dimension $d$ and inverse temperature $\beta$ . The dimension dependence is polynomial ( $O((d/\beta)^{p/2})$ ).
Comparison: This improves upon the state-of-the-art $O(\lambda^{1/8})$ rate for subgradient-based Langevin algorithms under similar non-smooth, superlinear, and non-convex conditions (specifically improving upon Johnston et al. [2025a]). It matches the $O(\lambda^{1/4})$ rate found in smooth non-convex settings but extends it to the non-smooth, superlinear regime.

Optimization Guarantees:
The sampling bounds are transferred to the optimization problem to bound the excess risk:
$\mathbb{E}[u(\theta^\lambda_n)] - u(\theta^\star) \leq C_{T1} W_2(\pi_\beta, \mathcal{L}(\theta^\lambda_n)) + C_{T2}$
where $C_{T2}$ represents the "temperature gap" (bias due to finite $\beta$ ). The complexity to achieve an $\epsilon$ -accuracy is $O(\epsilon^{-4} \log(1/\epsilon))$ iterations with stepsize $\lambda = O(\epsilon^4)$ .

Empirical Verification:
The authors verify their assumptions and algorithm on two specific problems:

LLM Pretraining (Nanochat): They apply a boosted coordinate-wise variant (SG-TheoPouLA) to pretrain a GPT-2 lineage model. The paper verifies that the potential satisfies the required curvature bounds (specifically $\nabla^2 u(\theta) \succeq -c_\star(1+|\theta|)^{12L+2} I$ ). Empirically, at depth 12, SG-TheoPouLA attains the lowest validation bits-per-byte and the highest CORE score among the tested optimizers under specific configurations, performing competitively against finetuned AdamW and Muon. At depth 24, Muon leads in validation bits-per-byte, while SG-TheoPouLA follows closely and remains competitive with AdamW on CORE, all within evaluation noise.
Sparse Phase Retrieval: The paper demonstrates that the sparse phase retrieval problem (quartic energy + $\ell_1$ penalty) satisfies all three challenging conditions (superlinear growth with $p=3$ , non-convexity, and gradient discontinuities on coordinate hyperplanes).

4. Significance and Claims

The paper claims to provide the first non-asymptotic convergence guarantees for a Langevin-based algorithm that simultaneously handles non-smoothness, superlinear gradient growth, and non-convexity without resorting to smoothing approximations.

Theoretical Impact: It bridges a gap in the literature where previous works required either smoothness (to use standard taming) or convexity (to ensure stability). The $O(\lambda^{1/4})$ rate is presented as an improvement over existing subgradient methods in this specific regime.
Practical Impact: The taming technique is shown to be computationally efficient (overhead significantly smaller than a matrix multiplication) and effective in deep learning contexts. The empirical results on LLM pretraining suggest that the theoretical framework applies to modern deep architectures, where gradients are polynomial in weights and non-convex.
Limitations: The paper notes that the exponential dependence on the inverse temperature $\beta$ in the constants is largely an artifact of the proof strategy (routing optimization guarantees through sampling). The authors suggest that under a Polyak–Łojasiewicz condition, this dependence could become polynomial. Additionally, the complexity $O(\epsilon^{-4})$ is noted to be inherited from the specific regularity assumptions on the drift, and higher-order smoothness would theoretically allow for better rates (e.g., $O(\lambda^{1/2})$ ).

In summary, the work establishes a rigorous theoretical foundation for using subgradient Langevin dynamics in the "wild" conditions of deep learning (non-smooth, non-convex, superlinear) and provides empirical evidence that the proposed taming scheme is a viable, competitive alternative to standard optimizers like AdamW and Muon.