Research map / Probabilistic and scientific modeling / Variational and generative sampling
Learning to Explore for Stochastic Gradient MCMC
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes L2E, a meta-learned SGMCMC sampler that learns the gradients of the kinetic energy with neural networks and is meta-trained with evolution strategies on a BMA predictive meta-loss, generalizing to unseen datasets for BNN posterior sampling.
Problem
Efficient exploration of high-dimensional multi-modal BNN posteriors with SGMCMC, which requires many sampling steps and hand-tuned design choices, and poor generalization of existing meta-learned samplers.
Approach
Instead of learning the diffusion and curl matrices (as in Meta-SGMCMC), parameterize the gradients of the kinetic energy function with MLPs applied per-coordinate, keep D and Q fixed as in SGHMC, and meta-train with a Bayesian-model-average predictive loss using antithetic evolution-strategies gradient estimates over diverse tasks.
Key results
- On CIFAR-10 HMC-fidelity, L2E reaches 0.946 agreement vs 0.920 for Deep Ensembles and 0.910 for cSGMCMC.
- On CIFAR-10 OOD detection with SVHN as OOD, L2E achieves AUROC 0.944 vs 0.837 (DE) and 0.858 (cSGMCMC).
- ESS/s on CIFAR-10 is 82.97 for L2E versus 56.61 for cSGMCMC, while L2E is the only method with cR^2<1.1 consistently high across datasets (0.968 on CIFAR-10).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use meta-learned kinetic-energy gradients with the BMA meta-objective | SGMCMC · CIFAR-10 and CIFAR-100 | Agreement with HMC predictive distribution: 0.920$\spm{0.001}$ -> \textbf{0.946}$\spm{0.002}$; 0.751$\spm{0.001}$ -> \textbf{0.771}$\spm{0.003}$ | Helped |
| Use L2E instead of CSGMCMC | Bayesian neural network inference · CIFAR-10-C | Average accuracy: 0.773 -> 0.755 | Hurt |
| Use L2E instead of CSGMCMC | Bayesian neural network inference · fashion-MNIST ESS/s | ESS/s: 219.85$\spm{6.64}$ -> 136.31$\spm{0.42}$ | Hurt |
| Use L2E instead of CSGMCMC | Bayesian neural network inference · CIFAR-10, CIFAR-100, and Tiny-ImageNet | ESS/s: 56.61$\spm{2.51}$ -> \textbf{82.97}$\spm{0.57}$; 41.14$\spm{1.51}$ -> \textbf{63.81}$\spm{0.00}$; 2.43$\spm{0.01}$ -> \t | Helped |
| Use L2E instead of DE | Bayesian neural network inference · CIFAR-10-C | Average accuracy: 0.791 -> 0.755 | Hurt |
| Use L2E instead of Meta-SGMCMC | Bayesian neural network inference · CIFAR-10 NT test set | accuracy: 78.12$\spm{0.035}$ -> \textbf{79.21}$\spm{0.203}$ | Helped |
| Use BMA meta-loss instead of cross-entropy meta-loss (ablation) | L2E · CIFAR-10 | loss barrier between collected parameters | Helped |
| Learn gradients of kinetic energy instead of diffusion and curl matrices (ablation) | Meta-SGMCMC · Image classification experiments | classification accuracy and predictive diversity | Helped |
| Parameterize the kinetic energy function rather than its gradients (ablation) | L2E · Image classification experiments | predictive accuracy and exploration | No clear effect |
Surprising findings
- L2E performs worst under heavy CIFAR-10-C corruption (average 0.755 vs 0.791 for Deep Ensembles), mirroring HMC, suggesting high fidelity to HMC hurts covariate-shift robustness.
- Meta-SGMCMC performs poorly on most tasks despite being the closest prior work, with cR^2<1.1 of only 0.160 on CIFAR-10.
Limitations
- L2E is not robust under covariate shift (CIFAR-10-C), where it is worse than Deep Ensembles and cSGMCMC at higher corruption intensities.
- The learned kinetic energy parameterization requires an assumption on the existence of an underlying energy function for convergence guarantees; an alternative version needs extra gradient computation.
- Analysis of exploration is limited to magnitude of network outputs rather than their direction.
Experimental setup
Datasets: MNIST fashion-MNIST EMNIST MedMNIST CIFAR-10 CIFAR-100 Tiny-ImageNet SVHN CIFAR-10-C IMDB
Benchmarks: fashion-MNIST CIFAR-10 CIFAR-100 Tiny-ImageNet CIFAR-10-C SVHN OOD IMDB MNIST 1-D synthetic regression
Metrics: accuracy NLL ECE AUROC agreement with HMC total variation ESS per second split R-hat (cR^2 < 1.1)
Baselines: Deep Ensembles cSGMCMC Meta-SGMCMC HMC
Builds on: SGHMC Meta-SGMCMC cyclical SGMCMC learned optimizers
Closest papers by meaning
- Parameter Expanded Stochastic Gradient Markov Chain Monte Carlo (ICLR 2025 · similarity 0.91)
- Meta-Learning For Stochastic Gradient MCMC (ICLR 2019 · similarity 0.91)
- Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning (ICLR 2020 · similarity 0.91)
- A Stochastic Approach to Bi-Level Optimization for Hyperparameter Optimization and Meta Learning (AAAI 2025 · similarity 0.88)
- Scalable Bayesian Learning with posteriors (ICLR 2025 · similarity 0.86)
- A Study of Bayesian Neural Network Surrogates for Bayesian Optimization (ICLR 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.