Robust Stochastic Gradient Posterior Sampling with Lattice Based Discretisation
Authors: Zier Mensch, Lars Holdijk, Samuel Duffield, Maxwell Aifer, Patrick J. Coles, Max Welling, Miranda C. N. Cheng
Organizations: Institute of Physics, University of Amsterdam, Netherlands · Department of Physics, National Taiwan University, Taiwan · Department of Computer Science, University of Oxford, United Kingdom · Normal Computing Corporation, New York, New York, USA · Amsterdam Machine Learning Lab, University of Amsterdam, Netherlands · Institute for Mathematics, Academia Sinica, Taiwan · Korteweg-de Vries Institute for Mathematics, University of Amsterdam, Netherlands
Stochastic-gradient MCMC methods enable scalable Bayesian posterior sampling but often suffer from sensitivity to minibatch size and gradient noise. To address this, we propose Stochastic Gradient Lattice Random Walk (SGLRW), an extension of the Lattice Random Walk discretization. Unlike conventional Stochastic Gradient Langevin Dynamics (SGLD), SGLRW introduces stochastic noise only through the off-diagonal elements of the update covariance; this yields greater robustness to minibatch size while retaining asymptotic correctness. Furthermore, as comparison we analyze a natural analogue of SGLD utilizing gradient clipping. Experimental validation on Bayesian regression and classification demonstrates that SGLRW remains stable in regimes where SGLD fails, including in the presence of heavy-tailed gradient noise, and matches or improves predictive performance.
Figures & tables
Figure 1: Comparing SGLD (left) and SGLRW (right) discretisations of Langevin dynamics, we can observe that the lattice based discretisation suppresses large parameter jumps that occur due to minibatch noise, resulting in more stable sampling.
Figure 2: Multimodal univariate target with exact gradient corrupted by synthetic α -stable noise ( α=1.5 ) of increasing scale. We observe that as we increase the noise scale SGLD quickly fails while SGLRW remains stable.
Figure 3: Mean-squared error (MSE) of the posterior covariance as a function of the step size δt , shown for different batch sizes for 50 -dimensional Bayesian linear regression.
Hyperparameters
KL Divergence
B
δt
SGLD
SGLRW
Clipped SGLD
8
10−3
19.889
6.060
18.184
8
10−4
0.483
0.202
0.777
16
10−3
7.155
2.317
8.530
16
10−4
0.175
0.070
0.204
32
10−3
2.441
0.729
3.540
Table 1: Kullback–Leibler (KL) divergence between the true posterior and the empirical Gaussian fit of the samples, shown for different samplers (SGLD, SGLRW, Clipped SGLD), minibatch size B , and base learning rate δt . The Monte Carlo reference KL divergence, which quantifies the intrinsic sampling variability when estimating the analytic posterior, is 0.055201 . Bold indicates the lowest KL divergence for a given (B,δt) pair.
Figure 4: Covariance difference matrices Σest−Σtrue for Bayesian linear regression at stepsize δt=10−3 , shown across increasing minibatch sizes B . Top: Clipped SGLD. Middle: SGLD. Bottom: SGLRW. Each panel visualizes the deviation of the empirical posterior covariance from the analytic posterior covariance; the Frobenius norm (Frob) reports the total error magnitude. We observe that the error in the diagonal terms of the estimated covariance matrices is lower for SGLRW than SGLD and Clipped-SGLD
Figure 5: Relative improvement of SGLRW over clipped SGLD at increased learning-rate scale ( η0=1.5×10−4 ). Heatmaps show percentage differences in negative log-likelihood ( left ) and expected calibration error ( right ) across training-set sizes and minibatch sizes.
Hyperparameters
KL Divergence
B
δt
SGLD
SGLRW
Clipped SGLD
1
100
inf
8.3504
9.3560
1
10−1
16.6812
6.0144
6.4732
1
10−2
10.3472
3.9594
4.2549
2
100
inf
7.6706
9.1187
2
10−1
8.9809
5.0197
5.5856
Table 2: Inferred Kullback–Leibler divergence for the logistic regression problem. The KL divergence is measured between Gaussian distributions fitted to the empirical mean and covariance of a gold-standard reference sample and those obtained by each algorithm under the specified hyperparameters. Bold indicates the lowest KL divergence for a given (B,δt) pair.
Figure 6: Predictive accuracy and negative log-likelihood (NLL) as a function of minibatch size for a large training-set size ( N=25,000 ), using base learning schedule with η0=7.5×10−6 .
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Mean-squared error (MSE) of the posterior covariance as a function of the stepsize δt , shown for different batch sizes. Top: 50 -dimensional Bayesian linear regression. Bottom: Bayesian logistic regression on the breast cancer dataset ( Wolberg et al., 1993 ) .
Figure 8: Multimodal univariate target with exact gradient corrupted by synthetic noise of increasing scale. Left: Gaussian noise ( α=2 ). Right: heavy-tailed noise drawn from an α -stable distribution with α=1.5 .
Figure 9: Predictive accuracy heatmaps for sentiment classification using fixed OPT-350M embeddings. Top: Clipped SGLD. Bottom: SGLRW. Each panel corresponds to a minibatch size B , with columns showing training-set size N and rows indicating learning-rate scale. Values report mean test accuracy across chains, with standard deviation shown below each entry.
Figure 10: NLL heatmaps for sentiment classification using fixed OPT-350M embeddings. Top: Clipped SGLD. Bottom: SGLRW. Each panel corresponds to a minibatch size B , with columns showing training-set size N and rows indicating learning-rate scale. Values report mean NLL values across chains, with standard deviation shown below each entry.
Figure 11: ECE heatmaps for sentiment classification using fixed OPT-350M embeddings. Top: Clipped SGLD. Bottom: SGLRW. Each panel corresponds to a minibatch size B , with columns showing training-set size N and rows indicating learning-rate scale. Values report mean ECE values across chains, with standard deviation shown below each entry.
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≤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.
M. Berk Sahin, Ege C. Kaya, Arda Fazla +2
Elmore Family School of Electrical and Computer Engineering, Purdue University · Weldon School of Biomedical Engineering, Purdue University
Tuning algorithms such as stochastic gradient descent (SGD) and stochastic gradient Langevin dynamics (SGLD) for approximate sampling and uncertainty quantification remains challenging, particularly in the practically relevant settings when the batch size is large or the model is misspecified. Existing theory that provides tuning guidance relies on continuous-time limits or strong statistical assumptions, which can become quantitatively inaccurate in these regimes. We address these shortcomings by proposing new discrete-time approximations to SG(L)D with and without momentum, which enables accurate predictions of the stationary covariance, iterate average covariance, and integrated autocorrelation time. Moreover, we prove quantitative, non-asymptotic error bounds showing that these estimates are sufficiently accurate for practical tuning and uncertainty quantification. Numerical experiments demonstrate that our theory yields improved tuning guidance across a range of models and data-generating distributions where existing approaches fail, including when using the β-divergence rather than log-loss to obtain statistically robust inferences.
Yu Wang, Jie Ding, Jonathan H. Huggins
Department of Mathematics & Statistics, Boston University · Questrom School of Business, Boston University · Faculty of Computing & Data Sciences, Boston University
We propose Decentralized Proximal Stochastic Gradient Langevin Dynamics (DE-PSGLD), a decentralized Markov chain Monte Carlo (MCMC) algorithm for sampling from a log-concave probability distribution constrained to a convex domain. Constraints are enforced through a shared proximal regularization based on the Moreau-Yosida envelope, enabling unconstrained updates while preserving consistency with the target constrained posterior. We establish non-asymptotic convergence guarantees in the 2-Wasserstein distance for both individual agent iterates and their network averages. Our analysis shows that DE-PSGLD converges to a regularized Gibbs distribution and quantifies the bias introduced by the proximal approximation. We evaluate DE-PSGLD for different sampling problems on synthetic and real datasets. As the first decentralized approach for constrained domains, our algorithm exhibits fast posterior concentration and high predictive accuracy.
Mohammad Rafiqul Islam, Lingjiong Zhu
Department of Mathematics, Florida State University, Tallahassee, Florida, United States of America