--- title: Adaptive Weighted SGLD url: https://www.emergentmind.com/topics/adaptive-weighted-stochastic-gradient-langevin-dynamics-sgld type: topic --- # Adaptive Weighted SGLD Adaptive Weighted Stochastic Gradient Langevin Dynamics (SGLD) refers to a class of scalable stochastic gradient Markov Chain Monte Carlo (SG-MCMC) algorithms where both the magnitude and direction of gradient-driven updates and injected noise are dynamically modulated, typically to address poor mixing, pathological curvature, and saddle points in high-dimensional, non-convex Bayesian inference problems. These algorithms encompass time-rescaling (stepsize adaptation), state-dependent preconditioning, adaptive drift mechanisms, and data-dependent importance weighting, aiming to improve posterior sampling efficiency, stability, and bias control in domains such as Bayesian deep learning and large-scale probabilistic modeling. ## 1. Langevin Diffusion, Discretization, and Gradient Noise At the core of SGLD is the discretization of the overdamped Langevin diffusion targeting a posterior law, $$ p(\theta \mid \mathcal{D}) \propto \exp(-U(\theta)), $$ where $U(\theta) = -\log p(\mathcal{D} \mid \theta) - \log p(\theta)$ encodes the energy landscape. The continuous dynamics, $$ d\theta_t = -\nabla U(\theta_t)\,dt + \sqrt{2}\,dW_t, $$ are replaced by the Euler-Maruyama scheme using noisy stochastic gradients $G_n \approx \nabla U(\theta_n)$: $$ \theta_{n+1} = \theta_n - h\,G_n + \sqrt{2h}\,\xi_n, \qquad \xi_n\sim\mathcal{N}(0,I). $$ Adaptive methods alter $h$ (stepsize), employ curvature-dependent preconditioners, or introduce state-dependent modifications to this standard iterative update [2511.11666]. ## 2. Adaptive Time-Rescaling and Stepsize Selection SA-SGLD, built upon the SamAdams framework [2511.11666], utilizes adaptive time-rescaling to modulate the stepsize based on a monitor function, typically the exponential average of the local gradient norm: - Auxiliary time $\tau$ and state-dependent Sundman transform $\psi(\zeta)$ control $$ dt/d\tau = \psi(\zeta_\tau). $$ - Monitor $\zeta_\tau$ aggregates local curvature information via $g(\theta) = \|\nabla_\theta U(\theta)\|^2 + \delta$. - Rescaled stepsize $$ \Delta t_{n+1} = \psi(\zeta_{n+1})\,\Delta\tau. $$ The update becomes $$ \theta_{n+1} = \theta_n - \Delta t_{n+1} G_n + \sqrt{2\Delta t_{n+1}}\,\epsilon_{n+1}. $$ This adaptive shrinking/enlargement of step length prevents overshooting in steep regions and improves mixing in flat regions, avoiding bias since only physical time is rescaled, not the geometry of the parameter space [2511.11666]. ## 3. State-Dependent Preconditioning and Gradient Weighting Preconditioned or weighted SGLD variants employ matrix-valued or diagonal preconditioners to adaptively scale both the gradient and noise: - RMSProp/Adam-style diagonal preconditioning uses running moment estimates [1906.04324, 1512.07666]: $$ V_t = \alpha V_{t-1} + (1-\alpha)(g_t)^2,\quad G_t = \mathrm{diag}(1/(\lambda+\sqrt{V_t})). $$ $$ \theta_{t+1} = \theta_t - \frac{\epsilon_t}{2}G_t \hat{\nabla} U + \sqrt{\epsilon_t}G_t^{1/2}\xi_t. $$ - More sophisticated blockwise or quasi-diagonal approximations (e.g., K-FAC, QDOP) further exploit model structure for efficient curvature adaptation [1806.02855, 1712.01076]. - Importance-weighted gradient sampling assigns non-uniform inclusion probabilities to data points according to their gradient magnitude or Lipschitz constant, and adaptively tunes batch size for variance control [2210.16189]. ## 4. Adaptive Drift, Dynamic Biases, and Escape from Saddles Adaptive drift-infused SGLD methods introduce a learnable bias term, which is dynamically adjusted according to running gradient statistics: - Momentum-SGLD (MSGLD) and Adam-style SGLD (ASGLD) maintain exponential running averages and second moments [2009.09535, 1805.09416]: $$ m_t = \beta_1 m_{t-1} + (1-\beta_1)\nabla\tilde{U}(\theta_{t-1}),\quad V_t = \beta_2 V_{t-1} + (1-\beta_2)(\nabla\tilde{U})^2. $$ $$ \theta_{t+1} = \theta_t - \epsilon[\nabla\tilde{U}(\theta_t)+a\,A_t] + \sqrt{2\epsilon\tau}\,\eta_{t+1}. $$ - This explicit drift modification enhances rapid escape from saddle points, improves multimodal traversal, and provably preserves ergodicity under suitable smoothness/dissipativity assumptions. - Escape from strict saddle points occurs within $O(\log d)$ iterations in adaptive schemes, essentially independent of problem dimension [1805.09416]. ## 5. Contour Flattening, Importance Sampling, and Weighted Ergodic Estimation CSGLD and related contour/adaptive weighting methods flatten the energy landscape by dynamically learning a piecewise-exponential weight function on the energy axis, transforming the sampling density: - The target is sampled from $\omega_\Psi(x) \propto \pi(x)/\Psi(U(x))^\zeta$, with $\Psi(u)$ adapted via a stochastic approximation procedure. - Ergodic averages of test functions must be reweighted to correct for the bias introduced by contour flattening [2010.09800]: $$ \hat{\mu}_k = \frac{\sum_{i=1}^k \theta_i(J_{\widetilde{U}}(x_i))^\zeta f(x_i)}{\sum_{i=1}^k \theta_i(J_{\widetilde{U}}(x_i))^\zeta}. $$ - Under stability and mean-field convergence theory, the learned weights approach their true bin probabilities, yielding provably small bias in posterior expectation estimation. ## 6. Thermostats, Friction Adaptation, and Superconvergence The use of auxiliary thermostat variables (Nosé–Hoover or SGNHT) provides robust temperature control for SG-MCMC, dynamically offsetting unknown gradient noise: - The friction variable $\xi$ is tuned so that the system's kinetic energy matches the target thermostat temperature; no prior estimate of gradient noise variance is required [1505.06889]. - Symmetric splitting schemes such as BADODAB achieve fourth-order “superconvergence” for configurational averages in the high-friction regime, enabling much larger stable timesteps than standard SGLD. - For Bayesian inference, this adaptation yields exact marginal sampling in the limit $\Delta t \to 0$, with dramatically reduced numerical overhead and improved efficiency. ## 7. Empirical Evaluation, Mixing, Generalization, and Calibration Adaptive weighted SGLD methods have been systematically benchmarked across curated toy problems and modern architectures: - In high-curvature potentials, adaptive stepsize SGLD (SA-SGLD) shows unbiased well traversal, stable contraction in steep regions, and superior mixing in plateaux [2511.11666]. - On image classification with Bayesian neural networks (BNNs), adaptive weighting outperforms SGLD under sharp horseshoe priors, achieving lower negative log-likelihood, higher accuracy, reduced expected calibration error, and greater robustness to stepsize escalation [2511.11666]. - Preconditioned SGLD demonstrates up to 3× improvement in effective sample size, accelerated convergence vs. unweighted SGLD, and competitive or superior generalization accuracy relative to SGD, Adam, and AdaGrad [1906.04324, 1512.07666]. - Adaptive drift and contour-flattening methods provide substantially enhanced multimodal exploration and reliable uncertainty quantification, especially in highly non-convex and multimodal landscapes [2009.09535, 2010.09800, 2210.16189]. | Variant | Mechanism | Key Empirical Findings | |------------------------|-----------------------------|-----------------------------------------| | SA-SGLD [2511.11666] | Time-rescaled stepsize | Superior mixing, unbiased sampling | | pSGLD [1512.07666] | Diag. preconditioning | Faster convergence, ESS up to 3× SGLD | | ASGLD [1906.04324] | 1st/2nd moment cumulants | Speed matches Adam/AdaGrad, gen. ≈ SGD | | CSGLD [2010.09800] | Adaptive contour weighting | Avoids multimodal local traps | | SGNHT [1505.06889] | Friction/thermostat | O(Δt⁴) bias, stepsizes 2–10× vs. SGLD | In summary, adaptive weighted SGLD methodologies incorporate curvature-sensing, state-dependent stepsize modulation, adaptive gradient and noise preconditioning, and advanced dynamical control to yield scalable, robust, and unbiased posterior sampling in high-dimension Bayesian inference, enabling marked improvements in mixing, exploration, and downstream predictive performance over classical SGLD. Source: https://www.emergentmind.com/topics/adaptive-weighted-stochastic-gradient-langevin-dynamics-sgld