Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes
Scalable Bayesian Learning with posteriors
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.
Introduces the posteriors PyTorch library for scalable Bayesian learning, frames SGMCMC as a tempered generalization of SGD that yields a parallel-SGMCMC (Bayesian deep ensemble) modification, and demonstrates it on cold posterior, continual LoRA learning with Llama 2, and OOD detection with Llama 3.
Problem
Bayesian learning with modern large models requires approximating high-dimensional posteriors, which is computationally expensive and lacks accessible, scalable, minibatch-capable software in the PyTorch ecosystem.
Approach
A tempered SGMCMC framework in which temperature interpolates between SGMCMC (T=1/N) and gradient descent (T=0), so SGD variants and samplers share one API; a small modification of deep ensembles (adding the noise term, parallel SGMCMC) makes them asymptotically unbiased for the posterior.
Key results
- On IMDB with a CNN-LSTM, Gaussian approximations (VI, Laplace) show a strong cold posterior effect with better test loss at T<1, while SGHMC shows only a mild effect.
- In continual LoRA learning on Llama 2 (7B) over 20 PG19 books, the Laplace method keeps low loss on early tasks while SGD forgets.
- For OOD detection with Llama 3 8B on Samoan statements, parallel SGHMC epistemic uncertainty reaches AUROC 0.95 versus 0.80 for SGD total uncertainty and 0.62 for base Llama 3.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use parallel SGHMC epistemic uncertainty instead of point-estimate uncertainty | Llama 3 8B fine-tuned on TQA · TQA Samoan OOD detection | AUROC: 0.80 -> 0.95 | Helped |
| Use parallel SGHMC total uncertainty instead of SGD total uncertainty | Llama 3 8B fine-tuned on TQA · TQA Samoan OOD detection | AUROC: 0.80 -> 0.90 | Helped |
| Fine-tune with parallel SGHMC instead of SGD | Llama 3 8B · TQA English statements | loss: 4.20 -> 4.40 | Hurt |
| Use Gaussian posterior approximations at temperature T<1 instead of T=1 | Laplace and variational inference on the CNN-LSTM · IMDB test set | loss | Helped |
| Use linearized Laplace instead of non-linearized Laplace | Laplace on the CNN-LSTM · IMDB test set | loss | Helped |
| Use conditional Fisher/GGN instead of empirical Fisher | Laplace on the CNN-LSTM · IMDB test set | loss | Helped |
| Linearize the forward model for variational inference | Variational inference on the CNN-LSTM · IMDB test set | loss | Hurt |
| Use parallel SGHMC instead of serial SGHMC | SGHMC on the CNN-LSTM · IMDB test set | loss | Helped |
| Lower temperature from the parallel-SGHMC setting toward T=0 | Parallel SGHMC on the CNN-LSTM · IMDB test set | loss | Hurt |
| Use empirical-Fisher Laplace instead of online SGD | Llama 2 with LoRA rank 8 · PG19 continual learning | test loss | Helped |
Surprising findings
- Strong cold posterior effect appears mainly for Gaussian approximations, while non-Gaussian SGMCMC shows only a mild effect, suggesting the effect stems from crude posterior approximations.
- Serial SGHMC failed to give improved epistemic uncertainty for Llama 3, attributed to staying in a single posterior mode.
Limitations
- Methods exploiting specific loss forms (deterministic last layer, probit approximations) and model-aware structure (K-FAC, DVI) are not yet supported.
- Parallel SGMCMC requires substantial resources to run many trajectories and hyperparameter tuning needs more care than in serial settings.
- Convergence diagnostics are not included and rely on composition with Pyro.
Experimental setup
Datasets: IMDB PG19 TQA
Benchmarks: IMDB test set PG19 continual episodes Samoan OOD statements
Metrics: test loss test accuracy AUROC categorical loss on next token
Baselines: MLE MAP Laplace GGN Laplace EF variational inference serial SGHMC SGD AdamW base Llama 3
Builds on: stochastic gradient Langevin dynamics SGHMC deep ensembles LoRA Bernstein-von Mises Laplace approximation
Closest papers by meaning
- Scalable Nonparametric Sampling from Multimodal Posteriors with the Posterior Bootstrap (ICML 2019 · similarity 0.88)
- Posterior Refinement Improves Sample Efficiency in Bayesian Neural Networks (NeurIPS 2022 · similarity 0.88)
- Bayesian Low-Rank Learning (Bella): A Practical Approach to Bayesian Neural Networks (AAAI 2025 · similarity 0.87)
- Bayesian Uncertainty for Gradient Aggregation in Multi-Task Learning (ICML 2024 · similarity 0.87)
- A Bayesian Approach to Data Point Selection (NeurIPS 2024 · similarity 0.86)
- Recursive PAC-Bayes: A Frequentist Approach to Sequential Prior Updates with No Information Loss (NeurIPS 2024 · similarity 0.86)
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.