Scalable Bayesian Learning with posteriors
Authors: Samuel Duffield
Kaelan Donatella
Johnathan Chiu
Phoebe Klett
Daniel Simpson
Venue: ICLR 2025 Poster (2025)
Keywords: Bayesian deep learning
PyTorch
Variational Inference
MCMC
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: Introduces a flexible PyTorch package for minibatch-first scalable Bayesian learning, with unification of SGMCMC and deep ensemble and applications to LLMs.

Abstract: Although theoretically compelling, Bayesian learning with modern machine learning models is computationally challenging since it requires approximating a high dimensional posterior distribution. In this work, we (i) introduce **_posteriors_**, an easily extensible PyTorch library hosting general-purpose implementations making Bayesian learning accessible and scalable to large data and parameter regimes; (ii) present a tempered framing of stochastic gradient Markov chain Monte Carlo, as implemented in posteriors, that transitions seamlessly into optimization and unveils a minor modification to deep ensembles to ensure they are asymptotically unbiased for the Bayesian posterior, and (iii) demonstrate and compare the utility of Bayesian approximations through experiments including an investigation into the cold posterior effect and applications with large language models.

_**posteriors**_ repository: https://github.com/normal-computing/posteriors

Published as a conference paper at ICLR 2025
SCALABLE BAYESIAN LEARNING WITH POSTERIORS
Samuel Duffield Normal Computing
Kaelan Donatella
Johnathan Chiu Normal Computing
Normal Computing
Phoebe Klett Normal Computing
Daniel Simpson Normal Computing
ABSTRACT
Although theoretically compelling, Bayesian learning with modern machine learn- ing models is computationally challenging since it requires approximating high dimensional posterior distribution. In this work, we (i) introduce posterio orr, an easily extensible Py Torch library hosting general-purpor ing Bayesian learning accessible and scalable to large data and parameter regimes; (ii) present a tempered framing of stochastic gradient Markov chain Monte Carlo, implemented in posterl s., that transitions seamlessly into optimization and unveils a minor modification to deep ensembles to ensure they are asymptotically unbiased for the Bayesian posterior, and (iii) demonstrate and compare the utility of Bayesian approximations through experiments including an investigation into the cold posterior effect and applications with large language models.
implementations mak-
posteriors: github wwwm/oormal-coooeoo cowp://ormaa-oooopptttttt/otttee
INTRODUCTION
Bayesian learning is a framework for updating and inferring unknown model parameters in the presence of data. It is based on Bayes' theorem:
p(0 VI:N) x p(0)p(m) 10) x p(0 1INN-1)(UUN 0),
(I)
offline
online
which provides a principled way to update our belief's about the parameters @ e R as data is observed once V1.N or sequentially yv. Bayesian learning represents uncertainty about the parameters using the posterior distribution p(0 VI.N), which combines prior beliefs p(0) with the likelihood or generating process p(y () for the data. This contrasts with optimization approaches, which store a single point estimate of the parameters 6" arg max, p(31-N 0). In our view, for modern machine learning models and pipelines, Bayesian learning offers three key advantages, which are visualized in Figure I: (a) improved generalization and out-of-distribution predictions, (b) coherent online learning of new information (without catastrophic forgetting), and (c) the ability to decompose predictive uncertainty into alcatoric (natural or explained uncertainty in the data) and epistemic (unexplained uncertainty that diminishes with more training data) components. It's important to note that in the more classical setting of statistical models with interpretable parameters, there additional benefits to be gained from Bayesian learning- these include propagation of uncertainty (Duffield et al., 2024) and model validation (Gelman et al., 2020). In this work, we introduce posterin orss python package designed to be minibatch-first and easily applicable to large-scale models with PyTorc (Paszke et al.. 2019) and Hugging Face (Wolf et al., 2020) whilst also being compatible with classical Bayesian models through composition with Pyro (Bingham et al. 2019, Appendix G).
Generalization a key metric for machine learning models, quantifying the ability of the model to perform tasks beyond its training set. Bayesian predictions offer a compelling approach to mitigate overfitting by averaging over posterior distribution of plausible models rather than relying on single fit. In various settings, Bayesian methods have demonstrated improved out-of-distribution performance (Neal, 2012: Blundell et al., 2015; Liu et al., 2020).
"Correspondence sam@normal www.orrmmacmmuu....mmmmmmmmmmmme comput ing.at

---

Published as a conference paper at ICLR 2025
Uncertaia
out-of-distribation
Learns new data
Raistemic
Confident in-distribution
Rememhers previous data
Decompose predictive uncertainty
Figure I: Pictorial representation of the benefits of Bayesian learning. Left: Averaging over multiple plausible fits to the data improves out-o--distribution generalisation. Center: Adapting to new online data without forgetting previous data. Right: Decomposing predictive uncertainty, with epistemic uncertainty providing an improved indicator for out-of-distribttinn detection.
Another vital consideration for modern machine learning pipelines is that of online continual learning, where the model continuously updates and learns from new data in real-time without needing to be retrained from scratch. Bayes' theorem (equation 1) provides coherent framework for online learning, where new data can be sequentially incorporated into the model by using the posterior distribution the prior receipt of the next data batch. When implemented exactly, this is equivalent to learning all data yI.N at once equation I, and all data points are exchangeable and treated identically. This contrasts naive online optimization, which typically results in the catastrophic forgetting (Goodfellow et al., 2013; Kirkpatrick et al., 2017) of previously learnt tasks.
In addition, it essential for statistical models to auditable and reliable, especially when they are met with data that is far outside the training distribution. With an optimization approach, the model is only able to provide predictions in the form of total uncertainty for the predictive distribution over test inputs x², via the entropy H [p(y) xz 6*)] == Epcula.8 flog p(y ix, 0")]. This is in contrast to Bayesian methods, which can decompose the total uncertainty (TU) of the posterior predictive distribution p(y x) Ep(0lvn:x) (pl xr, 0)] into two components (Wimmer et al., 2023; Sale et al., 2023; Hofman et al., 2024):
TU H(p(y x)), AU == Eppelnn [H[p(y x. O)],, EU TU AU.
Aleatoric uncertainty (AU) represents natural uncertainty in the data, such as synonyms sentence starts in natural language generation. Epistemic uncertainty (EU) represents uncertainty in the model parameters and diminishes with more data (as p(0 91:N) 5() (*) ) N o0) and therefore is better indication of model uncertainty in predictions. This decomposition is crucial for understanding model behaviour and making informed decisions in high-stakes applications (Kendall and Gal, 2017), although is not unique to Bayesian models but rather the broader concept of second-orde uncertainty (Osband al., 2024; Sale al., 2023; Hofman et al., 2024).
In this work, we concisely survey techniques for Bayesian learning with view scalability in both large data and large parameter regimes (with an extensive survey found Appendix A). We then make the following contributions:
1. We introduce pesteriors: functional. minibatch-firs PyTorch (Paszke et al., 2019) package for scalable and extensible Bayesian learning.
2. We formalize gradient descent methods as the low-temperature limit of stochastic gradient Markov chain Monte Carlo (SGMCMC) methods, thus implying simple modification of deep ensembles parallel SGMCMC that ensures convergence to the Bayesian posterior.
3. We provide experiments demonstrating the broad utility of posterion eiteerirrs and the acces- sibility of the aforementioned benefits. We investigate the generalization performance of host of approximate Bayesian methods and the prominence of the cold posterior effect (Wenzel et al., 2020; Izmailov et al., 2021), composing poster peerioon viors with LoRA (Hu et al., 2021) mitigate forgetting in a continual learning task with Llama (Touvron al., 2023) and out-of-distributiio detection via epistemic uncertainty with Llama (AI@Meta, Meta, 2024).
We conclude with discussion of related work and an outlook on future avenues for research and impact.
2

---

Published as a conference paper at ICLR 2025
2 BAYESIAN LEARNING
We now review the most prominent Bayesian learning methods for large-scale problems, all of which are implemented in posteriors. An in-depth survey with full details and discussions on caveats and computational bottlenecks can be found in Appendix A.
Laplace approximation is motivated by a second-order Taylor expansion around the mode of the posterior distribution 0MAP = arg max; p(( y1:N). This gives Gaussian approximation p(0 y1:N) = N(@ MAP. Σ). In the supervised learning setting, the covariance Σ is typically chosen (Daxberger et al., 2021) to be approximation to the Fisher information matrix F()) pp(e.ylle (V) pog x, xzy)] which is intractable as we do not have access to the true distribution over imputs p(z). Instead, two approximations (Kunstner et al., 2019) are used based on the training
data p(z, = 11 Σ 5(x,y x²,y") The first we denote the empirical Fisher Fε(0) = Ep( (Vo log y) logp((
J and is easy to compute. The second we denote the conditional Fisher Fc(O) = Bp(a)p(yla: 0) VV logp(0 y) Ve log p(0 y) and can be computed efficiently for common neural network likelihoods via an equivalence with the Generalized Gauss- Newton matrix (GGN, Martens 2020) with further details in Appendix A.I.
In its most common form variational inference (VI, Blei et al. 2017) also forms a Gaussian approximation p(0 y1:N) R go (0) = N( μr, E). However, VI directly learns the mean and covariance Ø = (n, 2) through an optimization routine to minimize the Kullback-Leibler objective (0) p(0 VI.N)], whose gradient can be approximated with Monte Carlo samples and tricks including reparameterzation (Rezende et al., 2014) and thic-thhe-landiggg lan-ingg (Roeder et al., 2017).
An important caveat to note for Gaussian approximations is that the full covariance matrix requires (dd) memory and often (dd) runtime for the required inversion and sampling operations. In the large parameter regime, this is prohibitive, and further approximations Ritter et al., 2018; Daxberger et al., 2021). Additionally, we can reduce memory at inference time by constructing uncertainty directly in function space by linearizing the forward model (Foong et al., 2019; Immer et al., 2021). If the likelihood has the form p(y x, 0) p(y fe(x)) and we have a posterior approximation N(O jr, Σ) then the linearized forward distribution becomes [(f(z) (z), (x) VVVVVVV ef (z)) where Vefu (z) is the Jacobian of fz(z) evaluated at This avoids having to sample many large parameter vectors.
required (MacKay, 1992;
Stochastic gradient Markov chain Monte Carlo (SGMCMC, Ma et al. 2015; Nemeth and Fearnhead
2021) instead forms Monte Carlo approximation p(0 yi:n) R N 5(0 0k). The samples {0+)
k=1 are typically collected by (approximately) evolving a stochastic differential equation (SDE) see Section 4 and Appendix B and collecting samples along the trajectory. Traditional full-batch Markov chain Monte Carlo (MCMC) schemes (Andrieu et al., 2003) use Metsopolis-Hastings step to ensure convergence of the Monte Carlo approximation is preserved through discretization. In the minibatch setting, the Metropolis- SeatroooissHHattinss ratio cannot be estimated easily. Instead, suppose have access to an unbiased stochastic gradient. Then, as noted in Welling and Teh (2011), the bias in the trajectory decreases quadratically O(e)) for learning rate ér. Therefore, adopting suitably decreasing learning rate schedule Σ of Ee oc (Robbins and Monro, 1951) will ensure an asymptotically unbiased Monte Carlo approximation.
3 POSTERICRS: UNCERTAINTY QUANTIFICATION WITH PYTORCH
posteriors is designed be a comprehensive library for uncertainty quantification in deep learning models. The key features outlining the posteriors rors philosophy are:
Composabilityy posteriors is written in PyTorc (Paszke al., 2019) and has compati- bility with a broad range of tools including the Llama (Touvron et al., 2023) and Mistral (Jiang et al., 2023) models, lightnims (Falcon and The PyTorch Lightning team, 2019) for convenient logging and device management, optimizers from torchopt (Ren et al., 2023) and probabilistic programs from pyro (Bingham et al., 2019), see Appendix G.
Extensible: The transform framework (Figure 3) adopted by posterion is very general and allows for the easy adoption of new algorithms. Additionally, posteriors supports
3

---

Published as a conference paper at ICLR 2025
Load classifier #ode1
Load cotaleader
nun,cate Lent smndetaalooerrddatanee
Tog rng.pesteriin (param, batchi:
latels batch
wtt.c/ciioeeaalllleeett.... params, tog_pest_vst
-terch.tm. fantional. creas.entropp logits, labela) postertors- dinganorma log_snob/paraas) nua_sata
18
11
12
reture log post_val, lngits Reture value and auctltary info
13
Dery Eminalle Senil SGMENCC Paslld SCMCMCC
transform posteriars.v..dae buits!
torchept.aderfi bonpanacurer///uuuuu
svop any postertors algoritta
stated transfora, Inttiparnes 28 21 V for tatch in datalonder:
Figure 2: Trajectories of various sampling methods for toy mul- timodal posterior. Deep ensemble concentrates the modes, serial SGMCMC struggles to transfer be- tween modes, parallel SGMCMC combines the benefits both.
state, aux transfora, spdate/sttte basch)
Figure posteriors code snippet to train a classifier with variational inference. posteriors recommends nor- malising the log posterior across the batch with scale inde- pendent of batch size or N. Scaling the prior and temperature by N ensures the posterior is still correctly targeted.
arbitrary likelihoods rather than being restricted to hard-coded regression or classification as is common in other libraries (Daxberger et al., 2021; Detommaso et al., 2023).
Functional:
posteriors adopts a functional API via PyTorch's s (Paszke et al., 2019) functional module. As championed by the JAX (Bradbury et al., 2018) ecosystem, the functional approach makes for code that easier to test and compose with other functions. Importantly for costeriors, functional programming is also closer to the mathematical description, which particularly useful for Bayesian modelling.
Unlike most common Bayesian software (Carpenter et al., 2017; Cabezas et al., 2024; Bingham et al., 2019), costeriors os is minibatch-first, thus allowing for efficient computation in large datasets. Additionally, flexible sub-spacing is provided for scaling to large models.
Scalable:
Swappable: The transform framework (Figure 3) allows users to switch between approaches seamlessly, allowing them to experiment and find the best method for the use case.
This combination of features makes posteriors unique amongst the PyToorch ecosystem which hosts the majority of the popular pretrained Hugging Face models (Wolf al., 2020).
3.1 SUPPORTED METHODS AND COMPLEXITIES
All methods discussed in Section 2 are included in pesteriors. Flexible Laplace approximations are supported based on both the empirical Fisher and conditional Fisher (via the GGN matrix, see Section 2 and Appendix A). Laplace approximations rely on an optimization routine to find the MAP followed by (bbd) time and memory operation (for batch size b) construct diagonal Fisher information matrices which ultimately require O(d) memory. For very large models this can be memory bottleneck meaning the batchsize may need to be reduced for the Laplace construction in comparison the optimization (although the Laplace construction requires only single epoch).
Variational inference for fitting Gaussian distribution with diagonal covariance costs O(kd) time and memory per iteration where k is the number of samples from the variational distribution which is often set to although larger & reduces the variances the gradient estimation. Dense covariance versions of both Laplace and VI are also supported but require (dd) memory and O(d²) for sampling, however, this can still be useful for machine learning models when combined with subspace approaches (Mangrulkar et al., 2022; Daxberger et al., 2021).
As we will see Section 4, SGMCMC is closely related to stochastic gradient descent and has the same per-iteration complexity. Although, key difference is the final output is an ensemble of
4

---

Published as a conference paper at ICLR 2025
parameters rather than the single point estimate. As such memory requirements are ((Md) where M is the number of output samples, as with deep ensembles (Lakshminarayanan et al., 2017). Further, as suitable decay in autocorrelation is desired between samples the number of iterations is typically longer for SGMCMC than optimization, although this can be parallelized, see Section 4.2.
In addition, posteriors hosts extended Kalman methods (Ollivier, 2018) (which has the same complexity as VI, Zhang et al. 2018), unified optimization via torchopt (Ren et al., 2023) and has roadmap of methods be added within the general framework including SVGD (Liu and Wang, 2016) and more advanced SGMCMC (Ma et al., 2015; Zhang et al., 2020). There is also a host of general-purposs utility functions not found elsewhere in the PyTorch ecosystem. This includes linearized forward calls, Fisher/GGN matrices, their diagonal counterparts, fast Fisher/GGN vector products and the conjugate gradient method, which provides the tools to implement Hessian-free optimization (Martens et al., 2010) and experiment with other second-order methods (Titsias, 2024; Duffield and Singh, 2022).
3.2 LIMITATIONS
posteriors supports arbitrary likelihoods and loss functions; as a result methods that utilise the specific form of the loss function (normally as classification or regression), e.g. the deterministic last layer approach in Harrison et al. (2024) or the probit approximations from Daxberger al. (2021), are currently out of scope.
Similarly, the current methods in posteriors oors do not assume any specific structure of the model or its parameters, see Figure 3. However, this makes it more challenging to support model-aware techniques such as DVI (Wu et al., 2018) or K-FAC (Martens and Grosse, 2015), although the benefits of such approaches may warrant an inclusion in the library going forward. Specifically for K-FAC, the Fisher information matrix is assumed to exhibit a local block structure that admits cheap inversion and sampling via Krotecker-factors. Initially, K-FAC was only developed for feedforward layers (Martens and Grosse, 2015), before being extended to convolutional (Grosse and Martens, 2016) and recurrent (Martens al., 2018) networks. More recently, it was generalized to weight-sharing networks (Eschenhagen et al., 2024) (including transformers and graph neural networks). Developing flexible K-FAC submodule within posteriors' functional framework a high priority. Such module should be suitably model-agnostie and capable of being applied across various inference methods (e.g.. optimization, Laplace. VI). This advancement would enable scalable inference with non-diagonal Fisher information matrices, broadening the library's s capabilities.
Convergence diagnostics (e.g. Cowles and Carlin 1996) represent a key component of MCMC workflow, however are not included directly in posteriors due to the easy composition with Pyro's 's (Bingham et al., 2019) extensive suite of convergence diagnostics. Similarly, posterio ors an inference library rather than probabilistic programming language, therefore does not include tools such as distribution definitions which can be utilized through composition with Pyro as demonstrated in Appendix G. Finally, we opt not to include fit convenience function (Chollet et al., 2015; Daxberger et al., 2021) in favour of the more flexible and transparent update framework in Figure 3.
4 STOCHASTIC GRADIENT DESCENT, SGMCMC AND DEEP ENSEMBLES
This section describes a tempered version of SGMCMC that introduces an additional temperature parameter that allows us view SGMCMC generalization of stochastic gradient descent (SGD) and variants. This viewpoint has guided the implementation of poster interroos allowing users access to SGMCMC within a familiar functional framework of popular optimization libraries such as torchopt (Ren et al., 2023) and optax (DeepMind et al., 2020), thus making Bayesian learning scalable and maximally accessible to the broader machine learning community. In the second half of this section, we show how this points to small modification of deep ensembles (Lakshminarayanal et al., 2017) to ensure convergence to the posterior, and therefore how posteriors can be used to easily convert deep ensemble pipeline into one of (parallel) Bayesian learning. Comprehensive details on SGMCMC are found in Appendix A.3.
The tempered connection between SGMCMC and SGD has been discussed in previous works for spe- cific optimizers/samplec e.g. Wenzel et al. (2020), however here we give a complete characterization of the connection based on the framework in Ma et al. (2015) which dictates the implementations in
5

---

Published as a conference paper at ICLR 2025
posteriors. Additionally, to our knowledge, the corollary unifying deep ensembles and parallel SGMCMC is novel.
The minibatch-first nature of posteriors alongside the reframing of stochastic gradient optimiza- tion as SGMCMC results in an implementation with access to both approaches within unified syntax as well as enabling the codesign of new algorithms, all within familiar optimization-style framework.
4.1 A UNIFIED FRAMEWORK FOR SGD AND SGMCMC
Consider targeting tempered target distribution कт (z) π (2) for temperature T [0, oc). which can be considered as synthetic hyperparameter controlling the variance of (2) (Wenzel al., 2020; Ding et al., 2014). Here 2 (0, w) where w are any auxiliary parameters such as momenta (Chen et al., 2014) or thermostat (Ding al., 2014). The distribution T(=) is thus an extended target with p(0 yi.N) marginal in 0.
An SDE is invariant for (2) if and only it has the following form (Ma et al., 2015)
de T (D) + Q()) (z) V [D=) + Q(=)]dt 22D(z)da))
dz = (D)) +
V [D() + Q(=)] 27 D(z)du, (3)
where the second line comes from rescaling (2)) (z) - TD;;), (Q(z)
Setting T N regains the original target, but conveniently. the scale of the drift term V log (2) no longer grows with N as (0) log + Σ log p(y 0) + const. This normalization of the loss or log posterior is the convention for optimization and results in more consistent learning rate selection across problems and is recommended in posteriors see Figure 3.
Intuitively, the low-temperature limit T 0 converts the target distribution TT2) into one with all probability mass concentrated on (local) maxima and transparently the SDE equation 3 into an ODE
T dz [D(2) + Q(=) log क (z)dt,
which represents exactly variant of continuous time gradient ascent N- (z).
All SGMCMC algorithms in posteri ors feature the temperature T as a parameter with T = 0 resulting in a valid SGD variant (specific examples can be found in Appendix C). The extensibility and strict support for minibatches allow a unified posteriors' implementation and new procedures immediately apply to both SGMCMC and SGD (e.g. step-size adaptation, which has in classical Bayesian statistics utilized Metropolis acceptance rates which do not apply to SGMCMC).
4.2 PARALLEL SGMCMC (AS BAYESIAN DEEP ENSEMBLES)
Deep ensembles (Lakshminarayanann et al., 2017) represent simple to implement and parallelizable route to uncertainty quantification via multiple runs of gradient descent with different random seeds (for initialization and batch shuffling). Thus, the arguments above point to the formulation of Bayesian deep ensembles, enabled by posteriore equation 4 typically used for deep ensemble training into a Bayesian version that is stationary for the posterior distribution by setting T N N- i.c. adding back the noise term in equation 3. Assuming D and Q are chosen to be independent of = this becomes
that converts the gradient descent ODE
dz = [D QO log (()dtt + V2N Ddw.
(5)
The core difference between SGMCMC and the Bayesian deep ensemble or parallel SGMCMC implementation above is that in SGMCMC, samples are collected from a single trajectory, whereas in parallel SGMCMC, multiple trajectories are simulated with only the last sample being retained (although multiple samples per trajectory can also be saved (Margossian et al., 2021)). Parallel SGMCMC provides an asymptotically unbiased approximation T(2) and therefore p(0 y1:N) ) as long as the criteria discussed in Appendix A.3 are met for single chain.
"qquivaeenly we rescale dt - T dt as Brownian motion (dw) scales with vt, see (1.27) in Pavliotis (2014)

---

Published as a conference paper at ICLR 2025
We note that in the large data regime, the additional noise term in equation 5 versus equation 4 becomes very small, and so existing deep ensembles (Lakshminarayanan et al., 2017) can already be seen approximating the posterior in some sense (with a small bias).
The functional framework of posteriors (in combination with torch. func) allows easy parallelization of the above unified SGMCMC and SGD algorithms into a parallel SGMCMC implementation that generalizes deep ensembles with minimal code change.
Figure 2 visualizes the behaviour of the techniques discussed in this section for a two-dimensional multi-modal double well posterior. We see clearly that although a deep ensemble finds all modes, it offers no within-mode uncertainty. In contrast, serial SGMCMC explores well within modes but struggles to transition between modes. Parallel SGMCMC combines the benefits of both methods to provide a more representative posterior approximation. In more realistic high dimensional, large N settings runtime and memory requirements become key consideration. Serial SGMCMC requires less total computation since runs only a single trajectory. However, the trajectory must be significantly longer and cannot be parallelized; therefore, serial SGMCMC is perhaps not well suited to modern computational architectures as the deep ensemble or parallel SGMCMC approaches. However, parallel SGMCMC naturally shares many of the caveats of deep ensembles. Notably, the resource requirements are significant to parallelize trajectories, although we note the opportunity to collect multiple samples per trajectory (Margossian et al., 2021). Another point of note is that communication costs between cores are expensive and hyperparameter tuning may require alternative approaches the serial setting and typically need more care (than optimization) to preserve the correct target distribution across trajectories.
5 EXPERIMENTS
We now present experiments investigating the Bayesian learning benefits motivated in Section The functional and swappable nature of posteriors makes these experiments particularly accessible across a range of Bayesian learning algorithms. Notably, its scalability and composability with the PyTorch ecosystem allow us to apply Bayesian learning to a host of pre-trained HuggingFace models (Wolf et al., 2020). Reproducible code for all experiments (and more) can be found the examples folder of the posterion ors repository.
5.1 COLD POSTERIOR EFFECT
This experiment fully utilizes posteriors swappability between algorithms to investigate the generalization capabilities of Bayesian learning. As seen in Figure 3, within a posteriors implementation the majority of the code is algorithm agnostic and we can easily swap approaches by only changing the build of the rras formfrrm object and its configuration parameters. Extensive experimental details can be found in Appendix D.
The cold posterior effect (Wilson and Izmailov, 2020; Wenzel et al., 2020; Aitchison, 2020; Izmailov et al., 2021) is a phenomenon that has been observed for approximate Bayesian machine learning models where increased predictive performance has been observed when targeting (2 y1:N) x p(0 V1.N) 1 (as discussed in Section 4 with T TN 1) for T which appears odds with the Bayesian paradigm (Aitchison, 2020).
Following Wenzel et al. (2020); Izmailov et al. (2021), we train a CNN-LSTM model with 2.7 million parameters on the IMDB (Maas et al., 2011) 1) for binary classification of positive or negative reviews. As in Izmailov et al. (2021), we use a diagonal Gaussian prior with variance 1/40 for all methods.
The test set loss performance across range of temperatures for Laplace, VI and SGMCMC ap- proaches are depicted in Figure 4, alongside the baseline optimization MAP approach. See Ap- pendix D for the same plots for accuracy and also optimization for the MLE (which overfits and performs significantly worse than MAP). There are several takeaways from Figure 4, which we describe now.
The Bayesian methods, particularly VI and SGHMC can significantly outperform gradient descent for the MAP when evaluating the test data. A simple Laplace approximation performs badly, however, improved performance is gained via the linearization discussed in Section 2. Laplace GGN, which is equivalent to Fe in this case (and in most machine learning models, see Appendix A.
7

---

Published as a conference paper at ICLR 2025
MAP
MAP
Sortal SCTTMIC
100
Laploce ET
Lineatied
Laplace
Laplace OGN Lineurized
365
305
Temperture
Torrperuture
Terperature
Figure 4: Investigation into the cold-posterior effect for range of posteriors algorithms for CNN-LSTM model (Wenzel et al., 2020) on the IMDB dataset (Maas et al., 2011). Non-linearized Laplace EF and Laplace GGN are indistinguishable (left panel). All Gaussian approximations are diagonal. All approaches display error bars with one standard deviation over random seeds.
improves over Fg, which in agreement with Kunstner et al. (2019). We observe that the reverse true for VI and that the linearization is actually detrimental. We posit that this is due to the VI optimization acting in parameter space without knowledge of forward linearization. We observe strong cold posterior effect for the Gaussian approaches with significantly superior performance achieved at T < 1. However, for the (non-Gaud Gaussian) SGMCMC approaches, only a very mild cold posterior approach is observed. This is evidence in agreement with Izmailov et al. (2021) that the cold posterior effect is more result of the crudeness of posterior approximations than the of Bayesian learning itself equation 1. Finally, we observe that the parallel SGHMC approach (Bayesian deep ensemble) slightly outperforms the single trajectory serial approach, relatably Figure 2. And further, reducing the temperature deteriorates performance (observing from Section 4 and Appendix that traditional deep ensemble is recovered as parallel SGHMC with 7 = 0),
5.2 CONTINUAL LEARNING WITH LORA
The purpose of this experiment is to apply a simple Bayesian technique made accessible by posteriors, namely, an empirical Fisher Laplace approximation, mitigate catastrophic forget- ting in an online setting with large-scale model. Here, the composability of posteriors, well as the flexibility the functional framework, is vital make use of parameter-efficien fine-tuning (PEFT) (Mangrulkar et al., 2022) and the Llama (Touvron et al., 2023) large language model from the PyTorch (Paszke et al., 2019) ecosystem. Extensive details can be found in Appendix E.
Task
SGD
40
Relative Test Loss
Laplace
Test Loss
30
20-
Task
10.
Office
50
100
150
200
Training Epoch
10 Episode
1617111811119
Figure 5: Continual learning with Llama 2. The online SGD and Laplace methods train one book after another, whilst the Offline SGD ap- proach sees all books simultaneously. represent- ing the network's 's capacity. Vertical dashed lines represent episode changes.
Figure 6: Continual performance averaged over episodes seen thus far. We use relative loss from baseline (SGD trained on single episode) for fairer averaging across episodes.
We consider the problem of episodal continual learning (Kirkpatrick et al., 2017; Nguyen et al., 2017) where each episode's data, Vnn is one of 20 books from the PG19 (Rae et al., 2019) dataset with 15%
8

---

Published as a conference paper at ICLR 2025
of each book held out for testing. We compare simple online gradient descent with approximation to the Bayesian online update equation p(0 y1:)) x p(θ /1.m-1)((H yv. [0) in Figures 5 and Our baseline implements AdamW via torchopt 's integration within the posteriors unified API. The Laplace posterior from the previous episode is then used as the prior for the next episode following the suggested amendment to the EWC method (Kirkpatrick et al., 2017) recommended in (Huszár, 2018). We also train an offline version of the same model with access to all books to indicate the model's total capacity (although unrealistic of an online pipeline).
We use LoRA (Hu et al., 2021) to fine-tune the query, key and output weight matrices in the last attention layer of the 7B parameter model Llama model (Touvron et al., 2023). LoRA computes low-rank upproximation to the weight updates and is implemented in the PEFT library Mangrulkar al. (2022), which integrates easily into functional posteriors ors code. We use standard choices for the rank (r = 8) and scaling (a 32) parameters.
Figure shows test loss for each episode, over all four episodes. The Laplace method maintains low loss in early tasks, while SGD forgets. For example, in the top the Laplace approximation stays low throughout training demonstrating that it continues to perform well on task even though is now being trained on data from tasks 2-4. In contrast, continually applying gradient descent quickly decreases the model's performance on task 1, Figure shows the difference of test loss from bascline, averaged over all episodes seen thus far. We see that the Laplace approximation approach maintains low loss averaged over all tasks.
30-
English
5.3 BAYESIAN LLAMA
Haae Hamee
We now show how posterior can used to fine-tune a very recent large language model provide the capability decompose uncertainty (equation 2) and provide improved aut-of-distribution detection. The flexible and scalable design of posterio fors allows us to easily use subspace methods for fine-tuning large pre-trained model in parallel using the novel Bayesian generalization of deep ensembles introduced in Section 4.2. Additionally, poster. oirr' functional approach allows us to easily and efficiently map over the ensemble in forward calls at inference time. Extensive experimental details can be found in Appendix F.
20-
15-
We fine-tune the last attention layer of the 8B Llama 3 model (AI@ Meta, 2024) (resulting in 218 million trainable parameters) on the TQA (Kembhavi et al., 2017) dataset, which consists of sci- entific textbooks. We fine-tune baseline with SGD and a size 10 ensemble trained in parallel with posteriors" SGHMC. We also compare against the base Llama model without fine-tuning. We evaluate the models by asking them to complete 100 scientific statements (in-distribution) with the last word removed. We then repeat the experiment with the statements translated into Samoan, an out-of-distribution language for TQA (Kembhavi et al., 2017).
30
Panilea
20-
Pamild 10
We can see from Figure that although total uncertainty pro- vides some indication for classifying the Samoan text as out-of- distribution, the epistemic uncertainty of parallel SGMHC offers a clearer split. Table validates this with the AUROC metric (Kuhn al., 2023), whilst also displaying the categorical loss values for the next token of the English statements to verify a successful fine-tune. Overall, we can see a significant advantage in -uitt--distriuution detection via the Bayesian decom- position of uncertainty, whereas the point-cstimate-B based methods (MLE/MAP) cannot decompose uncertainty to differentiate epistemic uncertainty from aleatoric or semantic uncertainty which is key for open-ended natural language generation (Kuhn al., 2023).
0,12
0.04
Figure Distributions of predictive uncertainties with Llama 3.
In Appendix F, we include results for serial SGHMC, which interestingly fails to provide improved epistemic uncertainty quantification. We posit that this has do with serial SGHMC staying isolated single mode of the posterior in this case, and therefore, averaging along the trajectory provides little over an SGD approach in Figure 2.

---

Published as a conference paper at ICLR 2025
Table Out-of-distributioo detection with Bayesian Llama
Base Llama 3
SGD
Parallel SGHMC
TU
TU
TU
EU
AUROCT f
0.62
0.80
0.90
0.95
Loss on English statements
4.68
4.20
4.40
RELATED WORK
Section 2 (and more thoroughly Appendix A) summarizes much of the related work regarding Bayesian methodology. Additional comprehensive reviews can be found in Wang and Yeung (2020) with Fortuin (2022) focussing on prior specifications (which we do not into detail here but is handled with broad flexibility in posteriors) and Izmailov et al. (2021) providing extensive experiments with full-batch MCMC to give closer indication of the behaviour of the true Bayesian posterior.
There is a broad selection of software for uncertainty quantification; however, to our knowledge and experience, none satisfy all five points in Section Stan (Carpenter et al., 2017) is widely used within the field of Bayesian statistics however does not support minibatching and is not easily extensible. Pyro (Bingham et al., 2019)
broad framework that supports range of models and inference techniques; however, is more heavyweight than post costeriors and harder to compose with existing PyTorch code. BlackJAX (Cabezas et al., 2024) (and optax (DeepMind et al., 2020)) heavily inspired the poster r ors functional approach. However, BlackJAX also has primary focus full-batch inference whilst posteria ors is minibatch- first and, like other JAX (Phan et al., 2019; Detommaso al., 2023; Duffield, 2021) packages, does not compose with the Py PyTorch ecosystem.
7 CONCLUSION AND OUTLOOK
We have presented new software package posteriors for scalable and extensible uncertainty quantification with PyTorch (Paszke et al., 2019) models. Through numerical experiments, we have demonstrated how posterin ors can be used to achieve range of benefits over optimization-aan approaches and composing with PyTorch toolkits. In particular, we highlighted that the cold posterior effect mostly occurs Gaussian approximations to the Bayesian posterior as well as demonstrating the benefits of Bayesian learning in the very practical and large-scale setting of free-form text generation with large language models.
We are excited by the future prospects for posteriors. Its functional API makes it easy maintain and extend; we already have a long list of methods and approaches to add. poster lors is fully open source, and contributions. The framework from Section 4 allows for simultaneous research into optimization and sampling methods. Combining that with posteric ors' broad range of utility functions (that, to our knowledge, are not available elsewhere in the Py Torch ecosystem), means that implementation and research into second-orded techniques (Martens al., 2010; Martens, 2020) and improved discretization methods (Leimkuhler et al., 2018) represent accessible avenues for future work in both optimization and SGMCMC.
keen to foster community for feedback for feature requests and encourage
REFERENCES
AI@ Meta. Llama model card. 2024. URL blob/maing http://aa/MMMMMMMMMMMMM md.
1aaeetttlllllla/aaa 11/mmm3/
L. Aitchison. A statistical theory of cold posteriors in deep neural networks. arXiv preprint 2xXi::2008.05122 2020.
A. Alexos, A. J. Boyd, and Mandt. Structured stochastic gradient MCMC. International
Conference Machine Learning, pages 414 434. PMLR, 2022.
10

---

Published as a conference paper at ICLR 2025
C. Andrieu, N. De Freitas, A. Doucet, and M. I. Jordan. An introduction to MCMC for machine learning. Machine learning,
2003.
S. Bieringer, G. Kasieczka, M. F. Steffen, and M. Trabs. AdamMCMC: Combining Metropolis Adjusted Langevin with Momentum-based Optimization. arXiv preprint arXi::2312. 2023.
E. Bingham, J. Chen, M. Jankowiak, Obermeyer, N. Pradhan, T. Karaletsos, Singh, A.
Szerlip, P. Horsfall, and N. D. Goodman. Pyro: Deep universal probabilisttie programming. J. Mach. html.
Learn. Res., 20:28:1-288 28:6, 2019. URL http: wtw.:.or...ooo........oooooeee wwtppaaaapppaaapeeeaaepppp http:://///....ooooppppss///////
D. M. Blei, A. Kucukelbir, and D. McAuliffe. Variational inference: review for statisticians. Journal of the American statistical Association, 112(518)859-877,, 2017.
C. Blundell, J. Corebise, K. Kavukcuoglu, and D. Wierstra. Weight uncertainty in neural network.
In International conference on machine learning, pages 1613-1622. PMLR, 2015.
J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, VanderPlas, S. Wanderman-Milne, and Q. Zhang. JAX: composable transformations of Python+Nummyy programs, 2018. URL ht / /iithb.. com/ htt/gooooleaaaaae/
A. Cabezas, A. Corenflos, J. Lao, R. Louf, A. Carnec, K. Chaudhari, R. Cohn-Gordon, J. Coullon, W. Deng, S. Duffield, G. Durán-Martin M. Elantkowski, D. Foreman-Macken M. Gregori, C. Iguaran, R. Kumar, M. Lysy, K. Murphy, J. C. Orduz, K. Patel, X. Wang, and R. Zinkov. BlackJAX: Composable Bayesian inference in JAX, 2024.
B. Carpenter, A. Gelman, M. D. Hoffman, D. Lee, B. Goodrich, M. Betancourt, M. Brubaker, Guo, P. Li, and A. Riddell. Stan: A probabilistic programming language. Journal of statistical software, 76(1), 2017.
T. Chen, E. Fox, and C. Guestrin. Stochastic gradient Hamiltonian Monte Carlo. In International conference on machine learning, pages 1683-1691. PMLR, 2014.
Chollet et Keras. https:://// /tt:////tttttt http:/tttuoocoo.eeeeeeeeeeeeeeee wtt. fccolcchollll let wetlleetteee...eeee keras, 2015.
M. Cowles and B. P. Carlin. Markov chain monte carlo convergence diagnostics: a comparative review. Journal of the American statistical Association, 91(434):883-90) 1996.
E. Daxberger, A. Kristiadi, A. Immer, R. Eschenhagen, M. Bauer, and P. Hennig. Laplace redux- effortless Bayesian deep learning. Advances in Neural Information Processing Systems, 34: 20089-20103, 2021.
DeepMind, I. Babuschkin, K. Baumli, A. Bell, S. Bhupatiraju, J. Bruce, P. Buchlovsky, D. Bud- den, T. Cai, A. Clark, I. Danihelka, A. Dedieu, C. Fantacci, J. Godwin, C. Jones, R. Hemsley, T. Hennigan, M. Hessel, Hou, S. Kapturowski. T. Keck, I. Kemaev, M. King, M. Kunesch, L. Martens, H. Merzic, V. Mikulik, T. Norman, G. Papamakarios, J. Quan, R. Ring, F. Ruiz, A. Sanchez, Sartran, Schneider, E. Sezener, S. Spencer, Srinivasan, M. Stanojević, W. Stokowiec, L. Wang, G. Zhou, and F. Viola. The DeepMind JAX Ecosystem, 2020. URL com/ google deepmind.
http::/gitt
G. Detommaso, A. Gasparin, M. Donini, M. Seeger, A. G. Wilson, and C. Archambeau. Fortuna: A library for uncertainty quantification in deep learning. arXiv preprint arXiv:2302.04019, 2023.
N. Ding, Y. Fang, R. Babbush, C. Chen, D. Skeel, and H. Neven. Bayesian sampling using stochastic gradient thermostats. Advances in neural information processing systems, 27, 2014.
S. Duffield. mocat: All things Monte Carlo, written in JAX., 2021. SamDuffield/mocee
Duffield and Singh. Quasi-Newton Sequential Monte Carlo, arXiv preprint arXiv:2211. 2022.
S. Duffield, S. Power, and L. Rimella. A state-space perspective modelling and inference for online skill rating. Journal of the Royal Statistical Society Series C: Applied Statistics. 73(5): 1262-1282, 2024.
11

---

Published as a conference paper at ICLR 2025
G. Duran-Martin, A. Kara, and K. Murphy, Efficient online Bayesian inference for neural bandits. In International Conference on Artificial Intelligence and Statistics, pages 6002-6021. PMLR, 2022.
R. Eschenhagen, A. Immer, R. Turner, F. Schneider, and P. Hennig. Kronecker- factoree approximate curvature for modern neural network architectures. Advances in Neural Information Processing Systems, 36, 2024.
W. Falcon and The PyTorch Lightning team. PyTorch Lightning, Mar. 2019. URL https: // github. com/
A. Foong, Li, J. M. Hermández-Lobat,, and E. Turner. 'In-Retween' Uncertainty in Bayesian Neural Networks. arXiv preprint arXiv:/906. 1537, 2019.
V. Fortuin. Priors in Bayesian deep learning: A review. International Statistical Review, 90(3): 563-591, 2022.
Y. Gal and Z. Ghahramani. Dropout as a Bayesian approximation: Representing model uncertainty in deep learning. In international conference machine learning, pages 1050-1059. PMLR, 2016.
Y. Gal, J. Hron, and A. Kendall. Concrete dropout. Advances in neural information processing systems, 30, 2017.
A. Gelman, A. Vehtari, D. Simpson, C. C. Margossian, B. Carpenter, Y. Yao. L. Kennedy, J. Gabry, P.-C. Bürkner, and M. Modrák. Bayesian workflow. arXiv preprint arXiv:201 0.01808, 2020.
I. J. Goodfellow, M. Mirza, D. Xiao, Courville, and Bengio. An empirical investigation of catastrophic forgetting in gradient-based neural networks. arXiv preprint arXiv:1312.6211, 2013.
J. Gorham and L. Mackey, Measuring sample quality with Stein's method. Advances in neural information processing systems, 28, 2015.
R. Grosse and J. Martens. kronecker-factoree approximate fisher matrix for convolution layers. In International Conference Machine Learning, pages 573-582, PMLR, 2016.
J. Harrison, J. Willes, and Snoek. Variational Bayesian last layers. arXiv preprint arXiv:2404. 11599, 2024.
J. Hartigan and J. Hartigan. Asymptotic normality of posterior distributions. Bayes theory, pages 107-118, 1983.
M. D. Hoffman, A. Gelman, et al. The No-U-Tumn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. J. Mach. Learn. Res., 15(1):1593-1623, 2014.
P. Hofman, Y. Sale, and E. Hüllermeier. Quantifying aleatoric and epistemic uncertainty with proper scoring rules. arXiv preprint arXiv:2-444122115 2024.
E. J. Hu, Y. Shen, P. Wallis, Allen- Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. LoRA: Low-Rank Adaptation of Large Language Models, 2021.
F. Huszár. Note the quadratic penalties in elastic weight consolidation. Proceedings of the National Academy of Sciences, 1 15(11 ):E2496-E2297, 2018.
A. Immer, M. Korzepa, and M. Bauer. Improving predictions of Bayesian neural nets via local linearization. In International conference on artificial intelligence and statistics, pages 703-711. PMLR, 2021.
P. Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A. G. Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv: 1803: 0840000007...00 2018.
Izmailov, S. Vikram, D. Hoffman, and A. G. G. Wilson. What are Bayesian neural network posteriors really like? In International conference machine learning, pages 1629-4640 PMLR, 2021.
A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. S. Chaplot, D. d. L. Casas, F. Bressand, G. Lengyel, G. Lample, L.. Saulnier, et al. Mistral 7b. arXiv preprint 2rXiv:2310.06822 2023.
12

---

Published as a conference paper at ICLR 2025
A. Kembhavi, M. Seo, D. Schwenk, J. Choi, A. Farhadi, and H. Hajishirzi. Are you smarter than sixth grader? textbook question answering for multimodal machine comprehension. In Proceedings of the IEEE Conference Computer Vision and Pattern recognition, pages 4999-5007, 2017.
A. Kendall and Gal. What uncertainties do we need Bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint 1411x:141..6980. 2014.
J. Kirkpatrick, R. Pascanu, N. Rabinowitz, J. Veness, G. Desjardins, A. A. Rusu, K. Milan, J. Quan, et al. Overcoming catastrophic forgetting in neural networks.
T. Ramalho, A. Proceedings of the national academy sciences, 114(13):321-3352, 2017.
L. Kuhn, Y. Gal, and S. Farquhar. Semantic uncertainty: Linguistic invariances for uncertainty estimation in natural language generation. arXiv preprint arXiv:2302.09664 2023.
Kunstner, P. Hennig, and Balles. Limitations of the empirical Fisher approximation for natural gradient descent. Advances in neural information processing systems, 32, 2019.
B. Lakshminarayanan, A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
B. Leimkuhler, C. Matthews, and J. Weare. Ensemble preconditioning for Markov chain Monte Carlo simulation. Statistics and Computing, 28:277-290, 2018
C. Li, C. Chen, D. Carlson, and L.. Carin. Preconditioned stochastic gradient Langevin dynamics for deep neural networks. In Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
J. Liu, Lin, S. Padhy, D. Tran, T. Bedrax Weiss, and B. Lakshminarayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. Advances in neural information processing systems, 33:7498-7511 2020.
Q. Liu and D. Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances neural information processing systems, 29, 2016.
I. Loshchilov and F. Hutter. Decoupled weight decay regularization, In International Conference on Learning Representations, 2018.
Y.-A. Ma, T. Chen, and E. Fox. A complete recipe for stochastic gradient MCMC. Advances in neural information processing systems, 28, 2015.
A. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguisties: Human language technologies, pages 142-150, 2011.
D. J. MacKay. A practical bayesian framework for backpropagation networks. Neural computation, 4(3):448-47,, 1992.
W. Maddox, Izmailov, T. Garipov, D. Vetrov, and A. G. Wilson. simple baseline for Bayesian uncertainty in deep learning. Advances in neural information processing systems, 32, 2019.
Mangrulkar, Gugger, L. Debut, Y. Belkada, Paul, and Bossan. PEFT: State-of of the--art Parameter-Efficien Fine- Tine-TTnnng methods. https: /ttss///gtttbbb/t// wwwp//ggggggoggneeeeeeoooooo/ facpgpeffee.eeeeeeeeeee// ft 2022.
C. C. Margossian, M. D. Hoffiman, P. Sountsov, L. Riou-Durand, A. Vehtari, and A. Gelman. Nested R: Assessing the convergence of Markov chain Monte Carlo when running many short chains. arXiv preprint arXiv:2110.13017. 2021.
J. Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76 2020.
13

---

Published as a conference paper at ICLR 2025
J. Martens and R. Grosse. Optimizing neural networks with Kronecker-factored approximate curva- ture. In International conference on machine learning, pages 2408-2442 2417. PMLR, 2015.
J. Martens, J. Ba, and M. Johnson. Kronecker-factored curvature approximations for recurrent neural networks, International Conference on Learning Representations, 2018.
J. Martens et al. Deep learning via Hessian- frriinn--reeeee optimization. In Icml, volume 27, pages 735-742, 2010.
R. M. Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
C. Nemeth and P. Fearnhead. Stochastic gradient Markov chain Monte Carlo. Journal of the American Statistical Association, 1166533):433-450, 2021.
C. V. Nguyen, Y. Li, T. D. Bui, and R. E. Turner. Variational continual learning. arXiv preprint 17Xiv:1710.10661 2017.
Y. Ollivier. Online natural gradient as kalman filter. 2018.
I. Osband, Z. Wen, S. M. Asghari, V. Dwaracherla, M. Ibrahimi, X. Lu, and B. Van Roy. Epistemic neural networks. Advances Neural Information Processing Systems, 36, 2024.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein,
L. Antiga, et al. Pytorch: An imperative style, serfoperformacce deep learning library. Advances in neural information processing systems, 32, 2019.
G. A. Pavliotis. Stochastic processes and applications. Texts in Applied Mathematics, 60, 2014.
D. Phan, N. Pradhan, and M. Jankowiak. Composable Effects for Flexible and Accelerated Proba- bilistic Programming in NumPyro. arXiv preprint arXiv:1912. 1912.
2019.
W. Rae, A. Potapenko, S. M. Jayakumar, C. Hillier, and T. P. Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL htt.httss///axiv.oorr arb/ 1911. 05507.
J. Ren, X. Feng, B. Liu, X. Pan, Y. Fu, L. Mai, and Y. Yang. TorchOpt, Nov. 2023. URL https: //gitt/ttt //github. wwtmm//etapp//.to htt.oppttpoooo....oooooooooo.
J. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, pages 1278-1286. PMLR, 2014.
H. Ritter, A. Botev, and D. Barber. scalable Laplace approximation for neural networks. In 6th international conference on learning representations, ICLR 2018 2018-conference track proceedings, volume 6. International Conference on Representation Learning, 2018.
H. Robbins and S. Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400-407, 1951.
G. Roeder, Y. Wu, and D. K. Duvenaud. Sticking the landing: Simple. estimators for variational inference. Advances in Neural Information Processing Systems, 30, 2017.
gradient
Y. Sale, V. Bengs, M. Caprio, and E. Hüllermeier. Second-order uncertainty quantification: A distance-based approach. arXiv preprint arXiv:2312.00995, 2023.
W. Smith, J. Everhart, W. Dickson, W. C. Knowler, and R. S. Johannes. Using the adap learning algorithm forecast the onset of diabetes mellitus. In Proceedings of the annual symposium on computer application in medical care, page 261. American Medical Informatics Association, 1988.
I. Sutskever, J. Martens, G. Dahl, and G. Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning, pages 1139-1144 147. PMLR, 2013.
14

---

Published as a conference paper at ICLR 2025
M. Titsias. Optimal Preconditioning and Fisher Adaptive Langevin Sampling. Advances in Neural Information Processing Systems, 36, 2024.
H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint ar Xtx::2377.002888 2023.
H. Wang and D.-Y. Yeung. A survey on Bayesian deep learning. ACM computing surveys (esur), 53 (5):1-37, 2020.
M. Welling and Y. W. Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681-688. Citeseer, 2011.
F. Wenzel, K. Roth, B. Veeling, J. Swiatkowski, L. Tran, S. Mandt, J. Snoek, T. Salimans, R. Jenatton, and S. Nowozin. How Good is the Bayes Posterior in Deep Neural Networks Really? In International Conference on Machine Learning, pages 10248-10259, PMLR, 2020.
A. G. Wilson and P. Izmailov. Bayesian deep learning and probabilistic perspective of generalization. Advances in neural information processing systems, 33:4697-4708, 2020.
L. Wimmer, Y. Sale, P. Hofman, B. Bischl, and E. Hüllermeier. Quantifying aleatoric and epistemic uncertainty in machine learning: Are conditional entropy and mutual information appropriate measures? In Uncertainty in Artificial Intelligence, pages 2282-2292. PMLR, 2023.
Wolf, L. Debut, Sanh, Chaumond, C. Delangue, A. Moi, Cistac, C. Ma, Jernite, Plu, C. Xu, T. Le Scao, S. Gugger, M. Drame, Q. Lhoest, and A. M. Rush. Transformers: State-of-the. Art Natural Language Processing. pages 28-45. Association for Computational Linguistics, Oct. 2020. URL https: / /ww.. a web. org/ ant.hoooyyyyooo.. o20plllyyyy22200 emn p- pedemos...
A. Wu, Nowozin, E. Meeds, R. E. Turner, M. Hernandez-Lobato, and A. Gaunt. Deterministic variational inference for robust Bayesian neural networks. arXiv preprint arXiv:1810.03958, 2018.
A. X. Yang, M. Robeyns, X. Wang, and L. Aitchison. Bayesian low-rank adaptation for large language models, 2024.
G. Zhang, Sun, D. Duvenaud, and R. Grosse. Noisy natural gradient as variational inference. In International conference on machine learning, pages 3362-5861. PMLR, 2018.
R. Zhang, C. L. Zhang, C. Chen, and A. G. Wilson. Cyclical stochastic gradient meme for bayesian deep learning. International Conference on Learning Representations, 2020.
A SURVEY OF SCALABLE BAYESIAN LEARNING
We now provide a more thorough review of methods for scalable Bayesian learning, extending Section 2.
A. LAPLACE APPROXIMATION
As discussed, the Laplace approximation (MacKay. 1992; Daxberger et al., 2021) forms Gaussian approximation to the posterior distribution p(0 Vi:N) = N(# OMAP, Σ) based on the second order Taylor expansion around the posterior mode 0MA arg max p(0 y1
1o@p(0 "y.N)) z (0 (MAP) T V log (OMAP V1:N)
(MAP) const,
Rd log N( OMAP, Σ) + const.
where Vo l(g V1:N) = 0
AMAP-
The Taylor expansion implies the use of the negative inverse Hessian Vg log p(0MAp V as the covariance E. however this is not guaranteed to be positive semidefinite and is therefore invalid
15

---

Published as a conference paper at ICLR 2025
covariance matrix. Instead, we can draw from the Bernstein-von Mises theorem (Hartigan and Hartigan, 1983):
p(0 yI:N) N-90c
N BMAP, N F(OMAP)
with Fisher information matrix (which is guaranteed to be positive semidefinite):
F(0) = E log log p(0 v)] p(ile) Vo p(O) v)Vo log 1)
(6)
In the supervised learning setting with inputs x, the Fisher information matrix becomes F(0) ylo) [ logp(0 p(0 x, v)] which intractable as we do not have access to the distribution over
inputs p(æ)). Two popular approximations the training data p(æ,, y) N-1 Σ 0(2, y xt x²,v"). The empirical Fisher:
(O = [V Vologp(0 y) logp(0 1 (1) 1,
and conditional Fisher:
Fc( (
V) log p(0) y) }.
with the conditional Fisher generally recommended (Kunstner al., 2019).
The integral with respect to ply x. in equation 7 may initially seem difficult to compute. Fortunately, things simplified in the common machine learning setting where the likelihood has the form ply x, 0) ply fa (z)) where p(y =) is the probability density function for an exponential family distribution with natural parameter z. That is
logp(y 2) = T(x) log / TT()ddy.
for some T(y). This setting includes the common cases of classification (softmax with cross-entropy loss) and regression (mean squared error loss). Then have exact equivalence between (0) and the Generalized Gauss-Newton matrix (GGN, Martens (2020)) matrix:
G(O) = Egt o (x)) HI(2, VaJo())
where (x) is the Jacobian the forward function and
( y) = ( V ==)==)((+))
is the Hessian of the loss function evaluated at = fo(r).
CAVEATS AND COMPUTATIONAL BOTTLENECKS
Laplace approximation training consists of first determining the posterior mode OMAP through gradient descent. This is followed by single additional epoch calculate the approximate covariance. It is, thus, cheap to compute in the large data regime. However, in the large parameter regime, storing the covariance matrix with (d2) elements becomes memory bottleneck inverting and sampling operations, which typically have runtime ( ). As such, approximations are used, such as diagonal (MacKay, 1992), K-FAC (Ritter al., 2018) and low-rank (Daxberger et al., 2021) with reduced (d) memory. The main caveat for Laplace approximations is the accuracy of a Gaussian approximation under the extreme non-linearities neural network posterior (Izmailov et al., 2021), which applies even further to reduced memory approximations.
well the required
A.2 VARIATIONAL INFERENCE
Variational inference seeks a parameterized distribution go (0) that minimizes some distance from the target posterior distribution p(0 y1:N). ). Most commonly (Blei et al., 2017) we assume a Gaussian
variational distribution (0) N(O p4, Σ) with Φ (n,, 3. and the distance be minimized the KL divergence:
0 < KL(ga(0) |1 p(0) Vi:N)] )] = log p(vIN) + (g) log (0)
log ())
(8)
p(0)
NELBO(
Thus, the problem becomes of optimization minimize the negative evidence lower bound objective NE which can be achieved with stochastic gradient descent alongside the reparam- eterization trick (Rezende al., 2014) to obtain unbiased gradients of the expectation.
16

---

Published as a conference paper at ICLR 2025
CAVEATS AND COMPUTATIONAL BOTTLENECKS
Notably, the NELBO () objective can be estimated stochastically by sampling from (o(0) and even minibatching the data V1:N- Therefore, training is scalable in the large data regime. The caveat regarding the crudeness of a Gaussian approximation to the posterior also applies here, except that variational inference minimizes formal divergence from the posterior rather than relying on Bernstein-von Mises asymptotics. However, the large parameter regime, the NELBO( objective requires high dimensional integration at every optimization step, which can suffer from high variance even with the available tricks (Blei et al., 2017).
A.3 STOCHASTIC GRADIENT MARKOV CHAIN MONTE CARLO
completely different paradigm to posterior approximation is to construct a Monte Carlo approxima- tion p(0 VI.N) s - k 5(0 0z). This most commonly achieved by evolving a stochastic differential equation (SDE) and collecting samples along the trajectory. Following Ma et al. (2015),
an SDE has (2)) as a stationary distribution (and therefore provides an asymptotically unbiased Monte Carlo approximation) if and only if it is of the form
+ dz = [D²) + Q(z) log (2) + V PD((2) + Q (2
(9)
with symmetric positive-semiddeinite D=z) and skew symmetric Q(z) - Q(z) Here 2 - (0,w) where w are any auxiliary parameters such as momenta and thermostat (Ma et al., 2015; Ding et al., 2014). The distribution ((2) is thus extended target with p(0 V1:N) marginal in 6. The matrices D==) and (=) are often chosen as independent of z, meaning the matrix divergence term vanishes.
The exact simulation of such non-linear SDE is intractable, and such, approximate discretization schemes such as the Fuler-Maruyama (Welling and Teh, 2011) method are employed. Traditional full- batch Markov chain Monte Carlo (MCMC) schemes (Andrieu al., 2003) use a Metropolis-Hastmme step to ensure convergence of the Monte Carlo approximation is preserved through discretization. In the minibatch setting, the Eatropolis-HHstting ratio cannot be estimated easily. Instead, we can
use an unbiased stochastic gradient Vogg (2) ~ N(. V log (2). B(=)), then as noted in Welling and Teh (2011) 1) the stochasticity arising from minibatching, (z)). decreases like OEe)) for learning rate et. A Robbins-Monro (Robbins and Monro, 1951) learning rate schedule Σ E < > ft oo
ensures an asymptotically unbiased Monte Carlo approximation (Welling and Teh, 2011). Additional enhancements can be made if an approximation to B(z) is available (Ma et al., 2015).
We detail specific implementations of SGMCMC in Section B.
CAVEATS AND COMPUTATIONAL BOTTLENECKS
The stochastic gradient formulation of MCMC permits easy application to the minibatch and large data regime. However, the decreasing learning rate requirement results in diffusive behaviour (Alexos al., 2022; Li et al., 2016; Bieringer et al., 2023) where the discretized SDE puts excessive weight on the noise term du over the informative gradient term. This is problematic since exploration is very much desired for MCMC. Often, a constant learning rate used in practice (Leimkuhler et al., 2018) with an asymptotic bias incurred. In the large parameter regime, the formulation of Monte Carlo approximation becomes a memory bottleneck. Typical lower dimensional MCMC procedures take many thousands of samples, but due to memory constraints, only a significantly smaller sample size available for large neural network parameters, resulting in crude approximation. A notable further caveat is that, in contrast the Gaussian approximations, the resulting Monte Carlo approximation does not provide access to a density with pointwise evaluations and cannot be used as a prior for subsequent Bayesian updates in online regime equation 1.
A.4 NOTABLE MENTIONS
Perhaps the most common route to uncertainty quantification in neural networks is dropout (Gal and Ghahramani, 2016; Gal et al., 2017), where model parameters or neurons are randomly removed from each forward call. This approach can be effective and has Bayesian interpretation in the posterior probability space of inclusion/cxxlusion, but this is naturally more restricted than the posterior approximations above.
17

---

Published as a conference paper at ICLR 2025
Stochastic weight averaging (SWA, Izmailov et al. (2018)) and SWA-Gaussian (SWAG, Maddox et al. (2019)) are similar in spirit to SGMCMC as they use samples from a single trajectory to generate point Gaussian approximation to the posterior respectively. Another option for formulating a Gaussian distribution over parameters is to consider each minibatch as a small-scale Bayesian update; local linearization of this update is the paradigm of the extended Kalman filter (Duran-Martin et al., 2022), which also has natural gradient descent (Ollivier, 2018) and variational inference interpretations (Zhang et al., 2018). These methods are easy to implement but again suffer from the caveats of Gaussian posterior approximations above.
B SGMCMC ALGORITHMS
Recall the generic SDE Ma et al. (2015) from Section 4 for targeting the tempered (T) x (()) where we set T = N to target the Bayesian posterior
d= [)) + Q(=) 'V log (=) TV (D(2) + Q(=)] + (2TD(2)duu +
for symmetric sesiiive-eemidefinite (22) and skew symmetric Q(=) Q(=) With (0.w)
for any auxiliary variables w and π (0,, w) dw = (0) = p(0 81:N).
The matrix divergence V M(2)) E Rd is defined Ma et al. (2015) as
Σ a dz,
M(z))as>
and vanishes in the common case that the preconditioners are set to be independent of i.e. D(2) D and Q(z) = Q.
We now cover the most popular SGMCMC algorithms and those implemented in posteriors, all of which are also found in Ma et al. (2015).
B. SGLD
The simplest SGMCMC algorithm is stochastic gradient Langevin dynamics (Welling and Teh, 201 I), which contains auxiliary variables (i.e. = = 0) and is represented by D and 0. This gives continuous-timg SGLD:
d# N V V log (() + dww.
And Suler-Maruyane discretization:
Ox+1 = On + + V log (Os) + 2E
Ca N(C 0.).
B.2 SGHMC
Stochastic gradient Hamiltonian Monte Carlo Chen et al. (2014) extends = (0, m) to include momenta with Gaussian target distribution
log T (0, m) = log
202" m m + const.
Choosing D and Q = ( gives continuous-time SGHMC
:
0 aT
d0 mddt.
dm = N ' log (8)dt ao ax- mmdtt + V2T adw,
and Fuler-Maruyam discretization:
0x+1 = 0g + 60 -mmz
m m. eN-1 log (0)) 60 2amm /26 T ToEE,
GA N I)
18

---

Published as a conference paper at ICLR 2025
B.3 SGNHT
The stochastic gradient Nosé-Hoover thermostat (SGNHT) algorithm (Ding et al., 2014) uses ideas from thermodynamies (Leimkuhler al., 2018) introduce an additional scalar variable that automates the selection of the a friction parameter in SGHMC. We have = (0, m, <) targeting" 2
1 T
d a)² const.
log x (0, m. E) log (1()
20
2
Choosing
0 all 0
1
0
D(0, m, m,
Q(0, m(0
1
0
m
0/
m
//d
gives continuous-time SGNHT
d0 a -2 mdl,
dm N V log π (f)dl Eo 3mdl V2T adw,
dg (a 2d m T m T)dt,
and Euler-Maruyama discretization:
B 02 60 m+:
mk+1 = m. + eN eN 1 log co 2 + VE Tack.
Sk ~ N(s 0,
Ck+1 & + (o 2d m. -T).
Although SGNHT has the same tuning parameters (c, o, a)) as SGHMC, we note the a parameter only controls the noise level for the m update and is absent from the deterministic dynamics, with even a representing reasonable choice.
C SGMCMC FOR OPTIMIZATION
Let us consider gradient descent (with momenta) (Sutskever et al., 2013), parameterization from PyTorch, for minimising (0)::
0 +++1 02 7m7k
m μu, - (1-r)VU(BL).
Here vanilla SGD 02+1 02 7 VU (0x) obtained with μr 0 and 0. And also SGHMC B.2 with T for (0)) N log (0) = - - log y1:N):
0x+1 02 + Ca 2mp,
mx+1 = (1 - co -2 aa)m - vOx().)
We see an equivalence between SGD and SGHMC (T=
= 0) (Wenzel et al., 2020) with
7 = 1 co -2, 1-
-2
1- T, a
o 2 7(1-)
(10)
We note that the transition from SGMCMC to optimization and back is not ubiquitous amongst methods. Indeed, the Murer-Maruyama SGNHT sampling algorithm B.3, does not result functional optimization algorithm at the zero temperature limit as the value thermostat can only increase and results in divergant dynamics. Also, the popular Adam (Kingma and Ba, 2014) optimizer does not fit the general sampling SDE equation due to its use of elementwise operations, which do not have a convenient matrix-veetor-produue form. posteriors framework makes these mathematical considerations transparent and an easily accessible avenue for research.
Note that the SGNHT target distribution for & is erroncously claimed as z4 (g ax2 in Ma et al. (2015) and Nemeth and Fearnhead (2021)
19

---

Published as a conference paper at ICLR 2025
D EXPERIMENT DETAILS: COLD POSTERIOR EFFECT
We now give additional details for the experiment into generalization capabilities of the Bayesian methods in posteriors and the cold posterior effect, from Section 5.1. Fully reproducible code can be found at posteriors/ examp ls/ imdb.
We here display more details on the results displayed in Figure 4. Figure 8 is exactly the as Figure 4, but with the test loss for the MLE, approach added (no prior), which severely overfits. Figure reports all accuracies which agree with the takeaways discussed in Section 5.1.
1.3
DCN Liwetedd
SCaveE
11 aucted
800
02 Tenpesuuee
0.2
042
15 Terpertzure
Terpentuue
Figure 8: Cold posterior test losses (lower is better) for range of posteriors algorithms and temperatures for the CNN-LSTM model (Wenzel et al., 2020) on the IMDB dataset (Maas et al., 2011) I) from Section 5.1 (same Figure 4 with MLE added). Laplace EF is indistinguishable from Laplace GGN.
0.55
0.0
6.8
6.58
MAP
Te4
500
501P
Tenperstur
Tarperituoe
Torpuratoo:
Figure 9: Cold posterior test accuracies (higher better) for range of posteriors algorithms and temperatures for the CNN-LSTM model (Wenzel et al., 2020) on the IMDB dataset (Maas et al., 2011) from Section 1.
We exactly the same model CNN-LSTM from Wenzel et al. (2020); Izmailov et al. (2021), which consists of an embedding layer to convert input sequences into dense vectors, a ID convolutional layer to capture local features, ReLU activation and max-pooling layers, LSTM layer capture long-term dependencies, and a fully connected layer to produce the final classification logits. The model has 2.7 million trainable parameters. Following Izmailov et al. (2021) Gaussian prior with all variances set to 1/40 (aside from MLE which has no prior or equivalently infinite prior variances).
diagonal
We train on the IMDB Maas et al. (2011) dataset for binary classification of positive/negative sentiment. We follow the default 50-50 split for the dataset with 25 thousand samples for training and 25 thousand samples for testing. All methods used batch size 32.
All approaches were averaged over 5 random seeds where for parallel SGHMC 5 ensembles were bootstrapped from 35 chains each with their own random seed.
We now detail the hyperparameters used for each method (except for serial SGHMC, all methods were trained for 30 epochs):
MLE: We train using the AdamW optimizer (Loshchilov and Hutter, 2018) with all hyper-
parameters set to the detaults from TorchOpt (Ren et al., 2023) (learning rate 10-3)
MAP: Same exact settings as MLE but with finite prior variances (set to 1/40 as with the other methods).
20

---

Published as a conference paper at ICLR 2025
Laplace: We take the trained MAP parameters as the mean of the Gaussian distribution. The diagonal covariances are calculated with posteriors and we compare both empirical Fisher and GGN (which is equivalent to the conditional Fisher for model). For test set evaluation, we sample 50 samples in parameter space and 10 samples in logit space for the linearized approach.
Variational Inference: As above, train using the AdamW optimizer (Loshchilov and Hutter, 2018) with all default parameters. We apply the stick-the-landing variance reduction (Roeder al., 2017) and parameterization (Rezende et al., 2014) tricks. We use single sample Monte Carlo estimate at each step. We train the variational variances in log space avoid negative variances and initialise all log standard deviations -3. For test set evaluation, we sample 50 samples in parameter space and 10* samples in logit space for the linearized approach.
Serial SGHMC: We train with learning rate 0.1 and friction a We run for 60 epochs (which is longer as many samples are collected along a single trajectory). We apply a burn-in of 20000 thousand iterations and then collect samples every 1000 iterations, resulting in final collection of 27 samples.
Parallel SGHMC: Setup was exactly the same as serial except only 30 epochs were trained, and only the final parameters were taken. This was repeated across 35 different random seeds (for data shuffling and SGHMC noise) in parallel, for testing we bootstrap sampled 5 ensembles of size 15 from the 35 chains.
All cold posterior simulations over 5 random seeds) take ~ 1 day to run.
on an NVIDIA A100, and all simulations (including repeats
E EXPERIMENT DETAILS: CONTINUAL LEARNING WITH LORA
In this section, we detail details for the continual learning experiment from Section 5.2. Fully reproducible code can be found at posteriors/eeee
wwteelee//oonniieeeeeeeeeeeeeeee continual lora.
Our dataset Rae et al. (2019) for the continual learning experiment, a collection of long books, is divided into N episodes of train and test data. In the results reported in Figures and 6, we use book per episode, holding out the last 15%. We select the first twenty books from Rae et al. (2019) that satisfy the following criterion:
1. Published in 1900 later
2. Is least 100,000 tokens, and at most 1,000,000 tokens
We use LoRA (Hu et al., 2021) on Llama 2 (Touvron et al., 2023) to reduce the dimension of the inference task to 200k trainable parameters. We test two baseline methods. We fine-tune the model on each episode's 'ss new data, a standard approach known to result in catastrophic forgetting, and we also implement a static offline baseline that sees all data every episode. This represents the LoRA network's total capacity but is computationally infeasible for online learning.
The prior for the next episode is the posterior from the previous episode, becoming a quadratic penalty in the loss function during gradient descent. Whereas the original work catastrophic forgetting (Kirkpatrick et al., 2017) suggested using multiple penalties (incorporating data from all previous episodes into the prior), we use single penalty following Huszár (2018). We use only the last episode's posterior, reminiscent of exact Bayesian updates.
For each episode of data, we perform fine-tuning on the current episode's data and test for all data episodes seen thus far. The Laplace method additionally updates the posterior based on new MAP estimates of weights, and the latest Fisher information. The posterior is used as the prior for the next episode. These steps are handled by posteriors.
Additional Experiment Settings:
1. We use rank 8 and 32 for LoRA, the standard settings (Hu al., 2021), and fine- tune three weight matrices in the last layer (key, query, and output projections). following the literature (Yang et al., 2024).
21

---

Published as a conference paper at ICLR 2025
Table Out-of-distribution metries with Bayesian Llama 3.
Base Llama 3
SGD TU
Serial SGHMC
Parallel SGHMC
TU
TU
EU
TU
EU
AUROCA
0.62
0.80 + 0.02
0.85
0.17
0.90
0.95
Loss on English statements +
4.68
4.20 0.14
4.28
4.40
Loss on Samoan statements |
11.41
9.87 + 0.05
9.98
9.83
2. We stride over the texts so that all tokens have 2048 context tokens, given that there are 2048 previous tokens condition on.
Continual LoRA simulations were NVIDIA A100; simulations over the 20 books take 6 hours to run.
F EXPERIMENT DETAILS: BAYESIAN LLAMA
Here we describe further details for our final experiment fine-tuning the 8B Llama 3 model (AI Meta, 2024). Fully reproducible code can be found at poster eeteerrrrrs/ree ///eeemmlll/baayaa 1aamyyssslllma.m. a.a..
As previously mentioned, experiments involve using 10 sets of ensemble models for the last attention layer of the Llama3 model, resulting in 218 million trainable parameters. We train via SGHMC sampling for 20 epochs over the TQA dataset (Kembhavi al., 2017). For evaluation, generated random set 100 scientific factual statements from ChatGPT, which we also validated by hand, Additional hyperparameter and training details
follows:
We set the learning rate to be 10-3, We set alpha and beta parameters to 10-2 and respectively, with the momenta initialized to O.
2. For the textbook content, tokenize with stride length of 300 and stride overlap of size 100 while using batch size of 10. During training, we mask the loss on the first 250 tokens and only consider the loss on the last 50 tokens.
3. All layers frozen except for the final attention layer.
30
Sena. SCIDMC TU
Panallet SCHICC EU
Pandlle sonme:
SCID TU
0,2
0044
Figure 10: Distributions of predictive uncertaintics with Llama 3 (Serial SGHMC included).
In Table 2, we extend Table 1 to include the loss results for predicting the last word of the statements translated into Samoan. Where we can clearly see the Samoan text is significantly out distribution. We also include one standard deviation enror from 4 random seeds for SGD.
In Figure 10, we extend Figure include results for serial SGHMC. We see that serial SGHMC fails to provide a useful measure of epistemic uncertainty, which, discussed in Section 5.3, posit is due to serial SGHMC struggling to escape local modes and, therefore, lacking diversity in predictions and even further, as the collected parameters are very similar (in such high dimension), this gives a misguided indication low epistemic uncertainty (a large agreement in produced logits) that results in
performance than using TU.
Bayesian Llama simulations were run on an NVIDIA A100; training of the SGD and ensemble approaches take 16 hours whilst evaluation over the 100 statements is fast.
22

---

Published as a conference paper at ICLR 2025
Pyso (NLITS)
(SGUMM,
postoioon (SCHMMC Full Buch)
Full Busc)
0.5
1.0
-0.55
-0.50
-0.25
0.00
0.25
-0.25
pragror
0.50
0.25
0.00
0.25
0.0
0.5
1.0
0.25
0.00
0.25
0.50
0.75
0.25
0.00
0.25
0.30
smtmm Insulin
tmii
ciabetes ciabere_pedigne
Figure 11: Full batch marginal distributions for Bayesian logistic regression the Pima Indians dataset. All methods generated 5000 samples and provided approximations that are broadly in agreement.
G BAYESIAN LOGISTIC REGRESSION: COMPARISON AND COMPOSITION WITH PYRO
In this section, we the probabilistic programming language (PPL) Pyro (Bingham et al., 2019) construct a more classical Bayesian logistic regression model on the Pima Indians dataset (Smith al., 1988) which is a standard benchmark in computational statistics, see e.g. Titsias (2024). The model is constructed in Pyro and seamlessly converted into a log posterior function for scalable inference with posteri ors. These simulations demonstrate appropriate convergence of the eoster eppsterii ors in smaller-scale setting, compare efficiency to other frameworks and demonstrate seamless composition with PPL in Pyro.
We compare posteriors' SGHMC and VI implementations with Pyro's 's automated NUTS (Hoff- man et al., 2014) implementation (which represents gold standard) and JAX (Bradbury al., 2018) implementation with BlackJAX (Cabezas et al., 2024). All serial MCMC methods were run for 5000 samples with similar warmup and thinning configurations (noting that NUTS has an adaptive thinning procedure). VI was ran for 8000 iterations which sufficient for convergence. For parallel SGHMC we run 5000 samples independently for 2000 iterations.
We observe, in Figure 11, that all methods (including mini-batched versions, Figure 12) are broadly in agreement and have roughly converged to the same posterior distribution. In Table 3, we compare the fidelity of the posterior approximations by visualising the marginal distributions and comparing quantitatively via kernelized Stein discrepancy (Gorham and Mackey, 2015). We also display the runtimes the approaches although in this case the model small (8 dimensions and 768 data points) and therefore the Python overheads are relatively significant over model calls (as demonstrated by only marginal speedup from mini- batchbb and this is more for PyTorch over JAX (although this rapidly deteriorates in the case of larger, more expensive model calls, in our experience). Nevertheless, we see that posteriors competitive with Pyro (Bingham al., 2019). Importantly, these simulations demonstrate that posteriors programming languages Pyro (Bingham et al., 2019) perform efficient and scalable mini-batch inference in traditional Bayesian models.
easily be composed one of the most popular probabilistic
23

---

Published as a conference paper at ICLR 2025
Pyno (NUTS 5)
(SCHIMC,
postrtoos (V), llacc Smwr52)
poserion (Penlla SCHMM Buich Sin--1)
0.0
0.5
10
1.5
2.0
-0.5
0.
0.5
-0.5
0.0
0.5
progaant
boodd
skir.
0.0
-0.3
0.0
ScRUTs (roaiin
bmi
dishetoc_pafigrze
ug2
Figure 12: Mini batch marginal distributions for Bayesian logistic regression. Otherwise, same setup as Figure 11 with full batch Pyro retained as a baseline. Again all methods provide relatively consistent posterior approximations.
Table 3: Metrics for Bayesian Logistic Regression. Kernelized Stein discrepancy (KSD) Gorham and Mackey (2015) measures the distance between the samples provided by the algorithm and the true posterior via kernel function (in this case standard Gaussian). All results are averaged over 10 random seeds with one standard deviation displayed. "The displayed parallel SGHMC time represents the time for a single chain that could be obtained with sufficient parallel resources.
KSD |
Time (s)
Pyro (NUTS Hoffman et (2014))
2.680 0.001
15.38 + 0.74
BlackJAX (SGHMC, Full Batch) posteriors (VI, Full Batch) posteriors (VI, Batch Size=32)
2.689 + 0.013
5.29 + 0.74
2.623 0.001
7.04
2.516 1 0.003
6.22
posteriors (SGHMC, Full Batch)
2.693 0.024
14.37
posteriors (SGHMC, Batch Size=32)
2,686 0.012
0.33
posteriors (Parallel SGHMC, Batch Size=32)
2,690 0.001
0.79 + 0.04
24


Reviews:
[{"review_id": "VvqXKnLJiS", "reviewer": "Reviewer_zc1P", "summary": "The paper introduces the `posteriors` pytorch library for programmatic approximate Bayesian inference with a focus on deep neural networks as models, which further contains some advances, like a variant of stochastic gradient Markov chain Monte Carlo. The paper assesses the implemented methods on a variety of problems.", "questions": "- Would it be possible to extend the library to non-diagonal parametric posterior approximations in the future?\n- Could the authors clarify their SGMCMC contribution in relation to recent prior work on such algorithms?\n- How were priors chosen in the experiments on the cold posterior effect? For VI and Laplace, these can have quite a strong effect and I couldn't find any information on that.\n- How would Figure 7 look like for the MAP as a comparison? I have seen claims that large models are already fairly well calibrated when looking at confidence itself.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Software release that allows other researchers and practitioners to build on the proposed and baseline methods. Overall appears to contain MCMC-based methods, variational inference, and Laplace approximations, all of which are relevant and practical methods.\n- Clearly written motivation for Bayesian inference in general\n- Interface seems very clear and flexible due to functional paradigm\n- Additional contribution in form of a new SGMCMC algorithm?\n- Convincing experiments on modern architectures like lora and llama.", "weaknesses": "- Only diagonal variational and Laplace approximations, which are known to have certain issues and are generally outperformed by more structured posterior approximations. For example, Laplace approximations perform better across the board when using layer-wise Kronecker-factored structure (Ritter et al, ICLR 2018, https://discovery.ucl.ac.uk/id/eprint/10080902/1/kflaplace.pdf). Apparently, extending the library to such approximations seems out of scope? (lines 209-210)\n- I find the interweaved SGMCMC contribution hard to understand and it appears to break the story a bit. Maybe it misses parts of the appendix to be understandable by itself.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "X4lpr7JNtg", "reviewer": "Reviewer_cMYJ", "summary": "I have increased my score from 5 to 6. \n----- \nThis paper presents a pytorch library called `posteriors` for general-purpose Bayesian learning tasks, which is featured by (1) composability, (2) extensibility, (3) functionality, (4) scalability, and (5) swappability. The implementation includes a tempered version of Parallel SGMCMC, which is shown to combine the favorable properties of deep ensembles and serial SGMCMC. Empirically, the library has been tested for cold posteriors, continual learning with LORA, and finetuning of Llamma3 with Bayesian learning, demonstrating the benefits of Bayesian paradigm for machine learning.", "questions": "- figure 4: only 4 lines are plotted but 5 lines are included in the legend box?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is in general well-written and the empirical results are presented clearly. I like the key features of the proposed library which are lacking in existing libraries. Below are the key strengths. \n\n- The key features of the proposed library as sketched in section 3 are important for modern Bayesian deep learning. \n\n- The experiments cover various scenarios in current deep learning landscape. \n\n- the library is open-sourced and supports customization.", "weaknesses": "I think the paper can be improved by better scoping the problems that the library aims to solve, the inference approaches provided by the library, and the practical tradeoffs between the proposed Bayesian approach versus the non-Bayesian approaches (e.g. SGD, LoRA and Deep Ensemles and their various combinations). \n\nIn particular, the following questions should be addressed:\n\n\n- What problems can the proposed library solve? Examples include disentangling various sources of uncertainties, continualing learning, finetuning, standard classification / regression,  classical Bayesian inference problems e.g. time series, hierarchical modeling? Detailing these problems will help readers understand what set this approach different from the existing ones such as Stan and BlackJAX. And perhaps a table that outlines the differences among these approaches can be helpful. \n\n- What are the key inference algorithms currently supported in the library?\n\n- What is the practical time and memory complexity of the current implementations when compared against the existing non-Bayesian approaches? \n\n-- What are the limitations / scenarios that the current library does not support?\n\nOverall I think the proposed library is very promising. And the throughout the paper some of the above points have been touched upon briefly. However, a comprehensive discussion is lacking, which is critical to the potential target users of the library. I will consider raising my scores if the above points are addressed during rebuttal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JaxKjMjYsW", "reviewer": "Reviewer_pqmW", "summary": "This paper introduces \"posteriors\", a PyTorch-based library designed to enhance scalability and usability of Bayesian inference for deep learning. The focus is on making Bayesian learning feasible for large datasets and models, often challenging due to the computational demands of approximating posterior distributions. Authors proposes a novel approach by reframing stochastic gradient MCMC as a form of gradient descent, creating a unified pathway between them. Experimental results highlight improvements in generalization, continual learning, and out-of-distribution detection, facilitated by scalable, minibatch-focused Bayesian techniques.", "questions": "- could the authors clarify how \"posteriors\" manages scalability in large parameter regimes, particularly regarding memory management with large covariance matrices?\n- how does \"posteriors\" compare in ease of use and computational efficiency with similar Bayesian packages like Pyro or BlackJAX?\n- Given that parallel SGMCMC shows promising results in certain tasks, do the authors anticipate specific scenarios or types of models where it may not perform well, and if so, why?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the paper is mostly clear and concise, presents a straightforward approach, keeping primary discussions in the main sections and moving technical details to the appendix.\n- the paper provides extensive experimentation, including large datasets and Bayesian inference with large-scale models, showing the library's potential across applications\n- the library is composable, extendable, and scalable; and promotes compatibility with pytorch and other popular libraries, and could therefore have a relevant impact in the community\n- the paper unifies stochastic gradient MCMC with traditional gradient descent, enabling efficient transitions between sampling and optimization\n- the library is open-source, and provides a functional API that aligns with pytorc , and so facilitates user contributions", "weaknesses": "- the clarity of writing in some technical sections could be improved. Certain descriptions of key concepts, particularly the tempered SGMCMC framework (section 4), would benefit from a more step-by-step breakdown. I felt it was a bit too rushed.\n- the main text could further emphasize how \"posteriors\" differentiates itself technically, showing specific implementations or optimizations unique to this library (in a similar spirit to Figure 3). The current presentation could be misinterpreted as implementing previously established methods rather than providing a novel platform for Bayesian inference\n- although the experiments are robust, comparisons to other popular Bayesian libraries are quite limited. Explicitly comparing computational efficiency or accuracy with other Bayesian packages would better position \"posteriors\" in the current landscape.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "En0LHjoz0g", "reviewer": "Reviewer_zqMe", "summary": "This paper proposes `posteriors`, a new PyTorch library for scalable Bayesian learning. Unlike previous libraries, `posteriors` is functional: The user can freely code their own minibatch-wise loss function and `posteriors` will use it to compute the posterior. Various Bayesian deep learning methods are available, such as the (linearized) Laplace approximation, variational inference, and stochastic-gradient MCMC.\n\nThe key selling points of `posteriors` is its composability, extensibility, swappablilty, and its functional nature. For instance, performing Bayesian inference on parameter-efficient fine-tuning in LLMs is a breeze since `posteriors` does not need to know the underlying model.\n\nMoreover, the authors discussed a parallel deep ensemble technique and applied it to the Llama-3-8b model, effectively enabling this latest LLM model to detect out-of-distribution data. Further experiment results on inference techniques available in `posteriors`, such as the linearized Laplace for continual learning, were also presented in this paper.", "questions": "The authors mentioned that `posteriors` is unable to provide model-aware techniques such as KFAC-Laplace. I believe that the future of Bayesian foundation models requires model-aware approaches: E.g. the current Laplace-LoRA with KFAC covariance (Yang et al., ICLR 2023) is insufficient since KFAC combined with LoRA has the same costs (at least memory-wise) as KFAC on the original weight matrices, thus negating the benefits of PEFT. (To see this, notice that KFAC on LoRA implies Kronecker factors of $n \\times n$, $k \\times k$, $k \\times k$, and $m \\times m$ for $n \\times m$ original weights.) How does `posteriors` handles these limitations?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "I like the approach of `posteriors` in making Bayesian inference easier to perform, even in highly sophisticated models like LLMs. Indeed, by focusing in composability and by being functional-first means that `posteriors` is not restricted to a small class of models.\n\nMoreover, I'm impressed by how clean and general the code of `posteriors` from the users' point-of-view (Fig. 3). I believe that `posteriors` is a useful addition to the community, enabling easy-to-do Bayesian inference on various models.", "weaknesses": "While the main contribution, i.e. the `posteriors` package, is great, the main weakness of this work is the presentation in the paper itself.\n\nThe story-telling of the paper is rather haphazard, badly motivated, and rather disjoint. The main purpose of this paper is to describe `posteriors`, to report the details of the package, and to show that it is useful. However, right from the introduction, the authors fail to motivate why do the community needs `posteriors` while other libraries such as `laplace-torch` exists, which also support sophisticated models like [LLMs](https://aleximmer.github.io/Laplace/huggingface_example/) and [reward models](https://aleximmer.github.io/Laplace/reward_modeling_example/). Indeed, the authors only motivate and compare `posteriors` with packages for traditional Bayesian inference (Line 86) and other functional libraries like Pyro in the Related Work section. Please note that I'm not saying that `posteriors` should not exist, but the motivation and comparison should be made stronger to make the paper more appealing to the readers, who are, ultimately, potential users.\n\nIn Section 3, I am underwhelmed by the fact that the authors did not describe their library in detail. As a library paper, I would expect the paper to be structured as in e.g. [PyTorch's paper](https://arxiv.org/pdf/1912.01703). I suggest the authors expand Section 3 more and reduce unnecessary discussion like Sec. 4, which ultimately does not contribute to the description of the library itself.\n\nIndeed, Section 4 is rather disjoint from the rest of the paper. It is much more suitable for a method paper. If the main purpose, as the authors intended, is to showcase `posteriors`' extendability, then this section will be much better used to show _how_ to implement the method with `posteriors`, not to propose the method itself. \n\nI am also confused about the selection of the experiments, especially Sec. 5.1. There, the authors only studied the cold posterior effects, which reads more like a validation of the parallel MCMC method proposed in Sec. 4. Again, it reads like a method paper, and thus disjointed from the rest of the paper.\n\nSecs. 5.2 and 5.3 are fine in terms of the problems. But they would be much stronger if the authors showcased _how_ `posteriors` play a part in making this possible. For example, the code needed to construct such a Bayesian LLM is very short with `posteriors`. The failure to do so made the paper less coherent, and ultimately, it is hard for a reader to fully understand the benefits of `posteriors`.\n\nLastly, I noticed quite a bit of typo/misformatting in the paper, e.g. incorrect use of `\\citep` and `\\citet` (Line 157 etc), incorrect reference formatting like \"A.1\" instead of \"Appendix A.1\" in line 124.\n\nIn sum, I believe the paper needs several more writing iterations to make the story coherent, thus fully showcasing the potential of `posteriors`, which I do really like.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)