Robust Langevin Sampling under State-Dependent Heavy-Tailed Noise with Applications to LLM Fine-Tuning

M. Berk Sahin; Ege C. Kaya; Arda Fazla; Behzad Sharif; Abolfazl Hashemi

cs.LGOct 9, 2026

Robust Langevin Sampling under State-Dependent Heavy-Tailed Noise with Applications to LLM Fine-Tuning

Authors: M. Berk Sahin, Ege C. Kaya, Arda Fazla, Behzad Sharif, Abolfazl Hashemi

Organizations: Elmore Family School of Electrical and Computer Engineering, Purdue University · Weldon School of Biomedical Engineering, Purdue University

Abstract

Existing convergence guarantees for non-log-concave sampling with stochastic gradient Langevin Monte Carlo (SG-LMC) typically assume globally Lipschitz gradients and uniformly bounded noise variance. Both assumptions can fail even in elementary sampling problems, and relaxing them better reflects practical applications such as Bayesian fine-tuning of large language models (LLMs). We relax both assumptions through generalized smoothness and our proposed state-dependent heavy-tailed oracle with finite $p$ -th moments, $1<p\leq 2$ . The oracle permits infinite variance when $p<2$ and includes the Blum-Gladyshev (BG-0) condition. We study these relaxations separately and jointly, proposing three methods. Under globally Lipschitz gradients, robust LMC (R-LMC) handles the proposed oracle through geometric-median centering and residual clipping. With exact gradients, adaptive LMC (A-LMC) controls discretization error under generalized smoothness. Robust adaptive LMC (RA-LMC) addresses both relaxations through robust estimation and a noise-aware step size. For all three methods, we establish non-asymptotic convergence guarantees in Fisher information (FI) without assuming dissipativity or a functional inequality. Synthetic experiments demonstrate FI convergence and the benefit of each mechanism. Bayesian fine-tuning experiments with LLMs and vision-language models support our oracle model and show that RA-LMC matches or outperforms state-of-the-art baselines in accuracy and uncertainty quantification.

Figures & tables

Figure 1: Robust estimation and adaptive steps improve FI convergence under heavy-tailed noise and generalized smoothness, respectively, while RA-LMC improves convergence when both challenges arise by integrating robust estimation with a noise-aware step size. SG, R, and A denote stochastic gradient averaging, proposed robust estimator and adaptive steps, respectively. Panels (a)–(c) use a Gaussian mixture with a BG-

p

oracle of infinite variance, a radial quartic target with exact gradients, and a radial cubic target with the same oracle, respectively. Each method’s curve shows the median across 10 seeds, with its shaded region indicating the interquartile range.

Figure 2: Panel (a) shows minibatch gradient error norms measured at step 2,500 of Qwen3-4B LoRA fine-tuning on OpenBookQA. Panel (b) shows synthetic norms generated from a one-sided Lévy-stable reference construction with an infinite second moment. The similarity between the histograms suggests heavy-tailed behavior in the stochastic gradient errors.

Figure 3: Panels (a) and (b) show the empirical first moment of minibatch gradient-error norms versus distance from the final LoRA iterate during fine-tuning of Qwen3-4B on OpenBookQA and Qwen3-VL-4B-Instruct on SLAKE. Each point averages error norms from 32 independent minibatches of size four. The increase in the first moment with distance motivates the state-dependent term in the BG-

p

oracle.

	OpenBookQA				SLAKE
Method	ACC $\uparrow$	NLL $\downarrow$	Brier $\downarrow$	ECE $\downarrow$	ACC $\uparrow$	NLL $\downarrow$	Brier $\downarrow$	ECE $\downarrow$
MLE	$91.400$ $\pm 0.849$	$0.617$ $\pm 0.050$	$0.151$ $\pm 0.012$	$0.070$ $\pm 0.007$	$89.155$ $\pm 0.934$	$1.070$ $\pm 0.286$	$0.203$ $\pm 0.017$	$0.103$ $\pm 0.010$
MAP	$91.500$ $\pm 0.757$	$0.583$ $\pm 0.051$	$0.147$ $\pm 0.014$	$0.069$ $\pm 0.011$	$89.366$ $\pm 0.423$	$0.954$ $\pm 0.247$	$0.198$ $\pm 0.012$	$0.099$ $\pm 0.006$
TempScale	$91.400$ $\pm 0.849$	$0.293$ $\pm 0.003$	$0.138$ $\pm 0.003$	$0.031$ $\pm 0.010$	$89.155$ $\pm 0.934$	$0.278$ $\pm 0.018$	$0.167$ $\pm 0.011$	$0.036$ $\pm 0.007$
MC Dropout	$91.250$ $\pm 1.182$	$0.545$ $\pm 0.044$	$0.151$ $\pm 0.013$	$0.066$ $\pm 0.008$	$89.366$ $\pm 0.841$	$0.832$ $\pm 0.177$	$0.197$ $\pm 0.017$	$0.098$ $\pm 0.009$
Deep Ensemble	$91.600$ $\pm 0.163$	$0.459$ $\pm 0.053$	$0.141$ $\pm 0.002$	$0.058$ $\pm 0.003$	$\underline{90.000}$ $\pm 0.586$	$0.660$ $\pm 0.046$	$0.177$ $\pm 0.007$	$0.084$ $\pm 0.006$

Table 1: OpenBookQA and SLAKE test results (mean

\pm

std. dev.). Accuracy (ACC, in percent) is higher-is-better; negative log-likelihood (NLL), Brier score, and expected calibration error (ECE) are lower-is-better. Bold: the best. Underline: the second best.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Potential $V(x)$	Geometry
State-dependent heavy tails	$\left\lVert x\right\rVert^{2}/2-\log\cosh(2x_{1})$	Globally $3$ -smooth
Generalized smoothness	$(\left\lVert x\right\rVert^{2}-1)^{2}/4$	$\gamma=2/3$
Combined setting	$\left\lVert x\right\rVert^{3}/3-\left\lVert x\right\rVert^{2}/2$	$\gamma=1/2$

Appendix

Table B.1: Target potentials used in the experiments of Fig. 1 and the supplementary figures. We note that the two radial potentials are not globally smooth.

Figure B.1: Fisher information of the time-averaged laws

\bar{\mu}_{\tau_{N}}

after

N

updates for the three settings and step sizes of Figure 1 . For each trajectory, we sample a time uniformly from

[0,\tau_{N}]

, where

\tau_{N}

is the time reached after its

N

th update and may vary across trajectories for adaptive methods.

Figure B.2: Fisher information of the time-averaged laws

\bar{\mu}_{T}

across the three settings of Figure 1 , using larger step sizes

h=0.04

0.04

, and

0.32

in panels (a)–(c), respectively. For each trajectory, a time is sampled uniformly from

[0,T]

to estimate the time-averaged law.

Figure B.3: Fisher information of the time-averaged laws

\bar{\mu}_{\tau_{N}}

after

N

updates for the same settings and step sizes as Figure B.2 . For each trajectory, a time is sampled uniformly from

[0,\tau_{N}]

, where

\tau_{N}

is the time reached after its

N

th update and may vary across trajectories for adaptive methods.

Figure B.4: Minibatch gradient-error norm histograms across Qwen3-4B LoRA fine-tuning on OpenBookQA. Columns show steps

0

1250

2500

3750

, and

5000

. The top row contains empirical norms and the bottom row contains samples generated from the fitted positive

\beta

-stable reference model in ( 130 ). Both rows use the same 80 bins and axes, and fitted

\beta

values are

(0.200,0.200,0.200,0.607,0.837)

Figure B.5: Minibatch gradient-error norm histograms across Qwen3-VL-4B-Instruct LoRA fine-tuning on SLAKE. The rows, columns, bins, and density normalization follow Figure B.4 . Fitted

\beta

values are

(0.287,0.200,0.578,0.730,0.972)

Figure B.6: Empirical

p

-th moments of minibatch gradient-error norms versus distance from the final LoRA iterate during Qwen3-4B fine tuning on OpenBookQA. Panels report

p\in\{1,1.25,1.5,1.75,2\}

. Each panel contains eight estimates at each of 50 checkpoints, and every estimate averages

R_{i,k}^{p}

over 32 independently sampled minibatch estimates with batch size of 4 as in ( 132 ).

Figure B.7: Empirical

p

-th moments of minibatch gradient-error norms versus distance from the final LoRA iterate during Qwen3-VL-4B-Instruct fine tuning on SLAKE. The moment orders, checkpoints, and grouped estimates follow Figure B.6 .

Explore similar work

The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity
Aug 6, 2026Iosif Lytras, Nikolaos Makras, Sotirios SabanisStochastic Gradient Langevin Dynamics Langevin Dynamics
School of Mathematics, University of Edinburgh, UK · National Technical University of Athens, Greece · Athena/Archimedes Research Centre, Greece
Improved Guarantees for Langevin Monte Carlo with Average Smoothness
May 29, 2026Arnak S. Dalalyan, Avetik KaragulyanMarkov Chain Monte Carlo Stochastic Gradient Langevin Dynamics
CREST, ENSAE Paris, Institut Polytechnique de Paris and MBZUAI · L2S, CNRS, Centrale Supélec, Université Paris-Saclay
Robust Stochastic Gradient Posterior Sampling with Lattice Based Discretisation
Feb 17, 2026Zier Mensch, Lars Holdijk, Samuel Duffield +4Posterior Sampling Markov Chain Monte Carlo
Institute of Physics, University of Amsterdam, Netherlands · Department of Physics, National Taiwan University, Taiwan · Department of Computer Science, University of Oxford, United Kingdom +4

Aug 6, 2026J/K move · Enter open · S save

The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity

Iosif Lytras, Nikolaos Makras, Sotirios Sabanis

School of Mathematics, University of Edinburgh, UK · National Technical University of Athens, Greece · Athena/Archimedes Research Centre, Greece

Stochastic Gradient Langevin Dynamics Langevin Dynamics Nonconvex Optimization

Open paper AskPDF arXiv

We study the problem of sampling from target distributions whose potentials are simultaneously non-smooth, subject to superlinear gradient growth, and non-convex. We introduce the Subgradient Tamed Unadjusted Langevin Algorithm (SG-TULA), a discretisation of the Langevin diffusion that operates directly on subgradients, without relying on computationally demanding smoothing procedures. To handle the superlinear regime, taming techniques are employed to produce a stable, explicit scheme. We derive non-asymptotic convergence bounds in Wasserstein-2 distance, with all constants tracked explicitly in terms of dimension and inverse temperature, improving upon the currently known rates for subgradient-based Langevin algorithms. We further provide excess risk estimates for the associated optimisation problem. We verify the assumptions, with explicit constants, for the regularized pretraining potential of a LLM in the GPT-2 lineage and the boosted coordinate-wise variant of SG-TULA pretrains the former competitively against finetuned AdamW and Muon, for which no comparable non-asymptotic guarantees are presently available.

Abstract

Explore similar work

The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity

Improved Guarantees for Langevin Monte Carlo with Average Smoothness

Robust Stochastic Gradient Posterior Sampling with Lattice Based Discretisation