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.
Department of Mathematics & Statistics, Boston University · Questrom School of Business, Boston University · Faculty of Computing & Data Sciences, Boston University