Research map / Probabilistic and scientific modeling / Variational and generative sampling
Parameter Expanded Stochastic Gradient Markov Chain Monte Carlo
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 PX-SGMCMC, which reparameterizes each BNN weight matrix as a product of expanded matrices (W = PVQ) so SGHMC/SGLD explores the posterior more diversely at no extra inference cost.
Problem
SGMCMC samplers for Bayesian neural networks suffer from limited sample diversity, which hurts uncertainty estimation and OOD robustness, and existing fixes need preconditioning, multiple chains, or extra approximations.
Approach
Reparameterize each weight matrix (and normalization parameters) as a product of square expanded matrices so that SGLD/SGHMC dynamics become implicitly preconditioned, enlarging exploration step sizes; the expanded factors are merged back into a single matrix at inference.
Key results
- On CIFAR-10 (no augmentation), PX-SGHMC reduces ERR to 0.121 vs 0.133 for SGHMC and NLL to 0.388 vs 0.422.
- On CIFAR-10-C (averaged over corruptions), PX-SGHMC achieves ERR 0.275 vs 0.294 for SGHMC.
- With data augmentation and T=0.01, PX-SGHMC improves TinyImageNet ERR to 0.498 from 0.538 for SGHMC; cost analysis shows only modest training overhead (383k vs 274k sampling params for c=d=1).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply parameter expansion to SGHMC | SGHMC on CIFAR-10 · CIFAR-10 test | ERR: 0.133\PM{0.002} -> 0.121\PM{0.002} | Helped |
| Apply parameter expansion to SGHMC | SGHMC on CIFAR-10 · CIFAR-10 test | NLL: 0.422\PM{0.005} -> 0.388\PM{0.005} | Helped |
| Apply parameter expansion to SGHMC | SGHMC on CIFAR-10 and distribution shifts · CIFAR-10, CIFAR-10.1, CIFAR-10.2, STL, CIFAR-10-C | ERR | Helped |
| Apply parameter expansion to SGHMC | SGHMC for OOD detection · SVHN and LSUN OOD detection | AUROC: 0.790\PM{0.051} -> 0.832\PM{0.014} (SVHN); 0.860\PM{0.018} -> 0.884\PM{0.007} (LSUN) | Helped |
| Apply parameter expansion to SGHMC | SGHMC with data augmentation and T=0.01 · CIFAR-10, CIFAR-100, TinyImageNet | ERR: 0.071\PM{0.001} -> 0.069\PM{0.001}; 0.319\PM{0.002} -> 0.290\PM{0.004}; 0.538\PM{0.003} -> 0.498\PM{0.004} | Helped |
| Add expanded parameterization (ablation) | SGHMC with constant step size schedule · CIFAR-10 test | ERR: 0.135\PM{0.003} -> 0.115\PM{0.002} | Helped |
| Add expanded parameterization (ablation) | SGHMC with constant step size schedule · CIFAR-10 test | AMB: 0.209\PM{0.001} -> 0.232\PM{0.003} | Helped |
| Replace constant with cyclical step size schedule (ablation) | SGHMC · CIFAR-10 test | ERR: 0.135\PM{0.003} -> 0.133\PM{0.002} | Helped |
| Replace constant with cyclical step size schedule (ablation) | PX-SGHMC · CIFAR-10 test | ERR: 0.115\PM{0.002} -> 0.121\PM{0.002} | Hurt |
| Use low-rank plus diagonal expanded matrices with r=p/8 (ablation) | PX-SGHMC · CIFAR-10 test | ERR: 0.131 -> 0.126 | Helped |
| Vary low-rank plus diagonal expanded matrix rank from r=p/8 to r=p/4 (ablation) | PX-SGHMC · CIFAR-10 test | ERR: 0.126 -> 0.122 | Helped |
| Vary low-rank plus diagonal expanded matrix rank from r=p/4 to r=p (ablation) | PX-SGHMC · CIFAR-10 test | AMB: 0.215 -> 0.242 | Helped |
Surprising findings
- Ablation shows the expanded parameterization contributes more to performance gains than the cyclical step size schedule.
- A low-rank plus diagonal variant of the expanded matrices still outperforms SGHMC while reducing memory overhead.
Limitations
- Requires extra training memory and computation for the expanded parameters.
- Expanded matrices increase the number of sampling-time parameters (e.g., 274k to 383k).
Experimental setup
Datasets: CIFAR-10 CIFAR-100 TinyImageNet SVHN LSUN
Benchmarks: CIFAR-10 CIFAR-10.1 CIFAR-10.2 STL CIFAR-10-C SVHN LSUN CIFAR-100 TinyImageNet
Metrics: ERR NLL ECE AUROC TNR95 TNR99 ensemble ambiguity variance of predictions parameter distance
Baselines: SGLD pSGLD SGHMC SGNHT HMC Learning to Explore (L2E) Deep ensembles
Builds on: SGLD SGHMC deep linear networks cyclical step size schedule
Closest papers by meaning
- Structured Stochastic Gradient MCMC (ICML 2022 · similarity 0.92)
- Learning to Explore for Stochastic Gradient MCMC (ICML 2024 · similarity 0.91)
- Meta-Learning For Stochastic Gradient MCMC (ICLR 2019 · similarity 0.91)
- High-Order Stochastic Gradient Thermostats for Bayesian Learning of Deep Models (AAAI 2016 · similarity 0.89)
- Microcanonical Langevin Ensembles: Advancing the Sampling of Bayesian Neural Networks (ICLR 2025 · similarity 0.89)
- A Study of Bayesian Neural Network Surrogates for Bayesian Optimization (ICLR 2024 · similarity 0.88)
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.