Scalable Bayesian Learning with posteriors

## Abstract

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 1 1 1 posteriors library: https://github.com/normal-computing/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.

## 1 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(\theta\mid y_{1:N})\propto\underbrace{p(\theta)p(y_{1:N}\mid\theta)}_{\text{ offline}}\propto\underbrace{p(\theta\mid y_{1:N-1})p(y_{N}\mid\theta)}_{\text{ online}},$ (1) $$

which provides a principled way to update our beliefs about the parameters $\theta\in\operatorname{\mathbf{R}}^{d}$ as data is observed once $y_{1:N}$ or sequentially $y_{N}$. Bayesian learning represents uncertainty about the parameters using the posterior distribution $p(\theta\mid y_{1:N})$, which combines prior beliefs $p(\theta)$ with the likelihood or generating process $p(y\mid\theta)$ for the data. This contrasts with optimization approaches, which store a single point estimate of the parameters $\theta^{*}=\operatorname*{arg\,max}_{\theta}p(y_{1:N}\mid\theta)$. In our view, for modern machine learning models and pipelines, Bayesian learning offers three key advantages, which are visualized in Figure 1: (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 aleatoric (natural or explained uncertainty in the data) and epistemic (unexplained uncertainty that diminishes with more training data) components.

Generalization is 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 a posterior distribution of plausible models rather than relying on a single fit. Bayesian methods have demonstrated improved out-of-distribution performance [, , ].

Another vital consideration for modern machine learning pipelines is that of online or continual learning, where the model continuously updates and learns from new data in real-time without needing to be retrained from scratch. Bayes’ theorem (1) provides a coherent framework for online learning, where new data can be sequentially incorporated into the model by using the posterior distribution as the prior on receipt of the next data batch. When implemented exactly, this is equivalent to learning all data $y_{1:N}$ at once (1), and all data points are exchangeable and treated identically. This contrasts to naive online optimization, which typically results in the catastrophic forgetting [, ] of previously learnt tasks.

In addition, it is essential for statistical models to be 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\mid x,\theta^{*})]:=-\operatorname{\mathbb{E}}_{p(y\mid x,\theta^{*})}[
\log p(y\mid x,\theta^{*})]$. This is in contrast to Bayesian methods, which can decompose the total uncertainty (TU) of the posterior predictive distribution $p(y\mid x)=\operatorname{\mathbb{E}}_{p(\theta\mid y_{1:N})}[p(y\mid x,\theta)]$ into two components [, , ]:

$$ $\displaystyle\text{TU}:=H[p(y\mid x)],\quad\text{AU}:=\operatorname{\mathbb{E} }_{p(\theta\mid y_{1:N})}[H[p(y\mid x,\theta)]],\quad\text{EU}$ $\displaystyle:=\text{TU}-\text{AU}.$ (2) $$

Aleatoric uncertainty (AU) represents natural uncertainty in the data, such as synonyms or sentence starts in natural language generation. Epistemic uncertainty (EU) represents uncertainty in the model parameters and diminishes with more data (as $p(\theta\mid y_{1:N})\to\delta(\theta\mid\theta^{*})$ as $N\to\infty$) and therefore is a better indication of model uncertainty in predictions. This decomposition is crucial for understanding model behaviour and making informed decisions in high-stakes applications [], although is not unique to Bayesian models but rather the broader concept of second-order uncertainty [, , ].

Figure: Figure 1: Pictorial representation of the benefits of Bayesian learning. Left: Averaging over multiple plausible fits to the data improves out-of-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-distribution detection.

Exact Bayesian learning is computationally challenging due to the intractability of the high-dimensional posterior distribution. This has led to the development of various approximate Bayesian learning methods, each with its own trade-offs in terms of flexibility, accuracy and computational cost. These methods have been implemented in many software packages e.g. [, , ]. Yet, these packages mostly focus on the traditional Bayesian setting of relatively small data and small parameter regimes and are not designed for the scale of many machine learning applications.

In this work, we concisely survey techniques for Bayesian learning with a view to scalability in both large data and large parameter regimes (with an extensive survey found in Appendix A). We then make the following contributions:

- 1.
We introduce posteriors; a functional, minibatch-first PyTorch [] 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 a simple modification of deep ensembles to parallel SGMCMC that ensures convergence to the Bayesian posterior.
- 3.
We provide experiments demonstrating the broad utility of posteriors and the accessibility of the aforementioned benefits. We investigate the generalization performance of a host of approximate Bayesian methods and the prominence of the cold posterior effect [, ], composing posteriors with LoRA [] to mitigate forgetting in a continual learning task with Llama 2 [] and out-of-distribution detection via epistemic uncertainty with Llama 3 [].

We conclude with a discussion of related work and an outlook on future avenues for research and impact.

## 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.

A Laplace approximation is motivated by a second-order Taylor expansion around the mode of the posterior distribution $\theta_{\text{MAP}}=\operatorname*{arg\,max}_{\theta}p(\theta\mid y_{1:N})$. This gives a Gaussian approximation $p(\theta\mid y_{1:N})\approx\operatorname{\mathbf{N}}(\theta\mid\theta_{\text{
MAP}},\hat{\Sigma})$. In the supervised learning setting, the covariance $\hat{\Sigma}$ is typically chosen [] to be an approximation to the Fisher information matrix $\operatorname{\mathrm{F}}(\theta)=\operatorname{\mathbb{E}}_{p(x,y\mid\theta)}
[\nabla^{2}_{\theta}\log p(\theta\mid x,y)]$ which is intractable as we do not have access to the true distribution over inputs $p(x)$. Instead, two approximations [] are used based on the training data $\hat{p}(x,y)=N^{-1}\sum_{i=1}^{N}\delta(x,y\mid x^{i},y^{i})$. The first we denote the empirical Fisher $\operatorname{\mathrm{F}}_{\text{E}}(\theta)=\operatorname{\mathbb{E}}_{\hat{p
}(x,y)}[\nabla_{\theta}\log p(\theta\mid y)\nabla_{\theta}\log p(\theta\mid y)
^{\top}]$ and is easy to compute. The second we denote the conditional Fisher $\operatorname{\mathrm{F}}_{\text{C}}(\theta)=\operatorname{\mathbb{E}}_{\hat{p
}(x)p(y\mid x,\theta)}[\nabla_{\theta}\log p(\theta\mid y)\nabla_{\theta}\log p
(\theta\mid y)^{\top}]$ and can be computed efficiently for common neural network likelihoods via an equivalence with the Generalized Gauss-Newton matrix (GGN) [] with further details in A.1.

In its most common form variational inference (VI) [] also forms a Gaussian approximation $p(\theta\mid y_{1:N})\approx q_{\phi}(\theta)=\operatorname{\mathbf{N}}(\theta
\mid\mu,\Sigma)$. However, VI directly learns the mean and covariance $\phi=(\mu,\Sigma)$ through an optimization routine to minimize the Kullback-Leibler objective $\text{KL}[q_{\phi}(\theta)\mid\mid p(\theta\mid y_{1:N})]$, whose gradient can be approximated with Monte Carlo samples and tricks including reparameterzation [] and stick-the-landing [].

An important caveat to note for Gaussian approximations is that the full covariance matrix requires $\operatorname{\mathcal{O}}(d^{2})$ memory and often $\operatorname{\mathcal{O}}(d^{3})$ runtime for the required inversion and sampling operations. In the large parameter regime, this is prohibitive, and further approximations are required [, , ]. Additionally, we can reduce memory at inference time by constructing uncertainty directly in function space by linearizing the forward model [, ]. If the likelihood has the form $p(y\mid x,\theta)=p(y\mid f_{\theta}(x))$ and we have a posterior approximation $\operatorname{\mathbf{N}}(\theta\mid\mu,\Sigma)$ then the linearized forward distribution becomes $\operatorname{\mathbf{N}}(f(x)\mid f_{\mu}(x),\nabla_{\theta}f_{\mu}(x)^{\top}
\Sigma\nabla_{\theta}f_{\mu}(x))$ where $\nabla_{\theta}f_{\mu}(x)$ is the Jacobian of $f_{\theta}(x)$ evaluated at $\mu$. This avoids having to sample many large parameter vectors.

Stochastic gradient Markov chain Monte Carlo (SGMCMC) [, ] instead forms a Monte Carlo approximation $p(\theta\mid y_{1:N})\approx N^{-1}\sum_{k=1}^{K}\delta(\theta\mid\theta_{k})$. The samples $\{\theta_{k}\}_{k=1}^{K}$ 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 [] use a Metropolis-Hastings step to ensure convergence of the Monte Carlo approximation is preserved through discretization. In the minibatch setting, the Metropolis-Hastings ratio cannot be estimated easily. Instead, suppose we have access to an unbiased stochastic gradient. Then, as noted in [], the bias in the trajectory decreases quadratically $\operatorname{\mathcal{O}}(\epsilon_{t}^{2})$ for learning rate $\epsilon_{t}$. Therefore, adopting a decreasing Robbins-Monro [] learning rate schedule $\sum\epsilon_{t}^{2}<\sum\epsilon_{t}=\infty$ will ensure an asymptotically unbiased Monte Carlo approximation.

Figure: Figure 2: Trajectories of various sampling methods for a toy multimodal posterior. Deep ensemble concentrates on the modes, serial SGMCMC struggles to transfer between modes, parallel SGMCMC combines the benefits of both.
Refer to caption: extracted/5627454/figs/double_well_schematic.png

## 3 posteriors : Uncertainty quantification with PyTorch

posteriors is designed to be a comprehensive library for uncertainty quantification in deep learning models. The key features outlining the posteriors philosophy are:

- •
Composability: posteriors is written in PyTorch [] and has compatibility with a broad range of tools including the Llama [] and Mistral [] models, lightning [] for convenient logging and device management, optimizers from torchopt [] and probabilistic programs from pyro [].
- •
Extensible: The transform framework (Figure 3) adopted by posteriors is very general and allows for the easy adoption of new algorithms. Additionally, posteriors supports arbitrary likelihoods rather than being restricted to hard-coded regression or classification as is common in other libraries [, ].
- •
Functional: posteriors adopts a functional API via PyTorch’s [] functional module. As championed by the JAX [] ecosystem, the functional approach makes for code that is easier to test and compose with other functions. Importantly for posteriors, functional programming is also closer to the mathematical description, which is particularly useful for Bayesian modelling.
- •
Scalable: Unlike most common Bayesian software [, , ], posteriors is minibatch-first, thus allowing for efficient computation in large datasets. Additionally, flexible sub-spacing is provided for scaling to large models.
- •
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.

In addition to the extensible suite of Bayesian methods, posteriors also supplies a host of general-purpose 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 [] and experiment with other second-order methods [, ].

### 3.1 Limitations

posteriors supports arbitrary likelihoods and loss functions; this means that methods that utilise the specific form of the loss function (normally as classification or regression), e.g. the deterministic last layer approach in [] or the probit approximations from [], are currently out of scope. Similarly, posteriors does not assume any specific structure of the model, which makes model-aware techniques such as DVI [] or KFAC [] difficult to provide directly; however, the functional approach should allow easy composition with bespoke model-aware code. Finally, we opt not to include a fit convenience function [, ] in favour of the more flexible and transparent update framework in Figure 3.

## 4 Stochastic gradient descent, SGMCMC and deep ensembles

This section introduces a tempered version of SGMCMC, as implemented in posteriors, and demonstrates how this provides a unified transition between stochastic gradient descent (SGD) and SGMCMC methods. We then show how this points to a small modification of deep ensembles [] to ensure convergence to the posterior. Comprehensive details on SGMCMC are found in Appendix A.3.

### 4.1 A unified framework for SGD and SGMCMC

Consider targeting a tempered target distribution $\pi_{\operatorname{\mathcal{T}}}(z)\propto\pi(z)^{\frac{1}{N\operatorname{
\mathcal{T}}}}$ for temperature $\operatorname{\mathcal{T}}\in[0,\infty)$, which can be considered as a synthetic hyperparameter controlling the variance of $\pi_{\operatorname{\mathcal{T}}}(z)$ [, ]. Here $z=(\theta,\omega)$ where $\omega$ are any auxiliary parameters such as momenta [] or thermostat []. The distribution $\pi(z)$ is thus an extended target with $p(\theta\mid y_{1:N})$ marginal in $\theta$.

An SDE is invariant for $\pi_{\operatorname{\mathcal{T}}}(z)$ if and only if it has the following form []

$$ $\displaystyle dz$ $\displaystyle=\operatorname{\mathcal{T}}^{-1}[\operatorname{\mathrm{D}}(z)+ \operatorname{\mathrm{Q}}(z)]N^{-1}\nabla\log\pi(z)dt+\nabla\cdot[ \operatorname{\mathrm{D}}(z)+\operatorname{\mathrm{Q}}(z)]dt+\sqrt{2 \operatorname{\mathrm{D}}(z)}dw,$ $\displaystyle\implies dz$ $\displaystyle=[\operatorname{\mathrm{D}}(z)+\operatorname{\mathrm{Q}}(z)]N^{-1 }\nabla\log\pi(z)dt+\operatorname{\mathcal{T}}\nabla\cdot[\operatorname{ \mathrm{D}}(z)+\operatorname{\mathrm{Q}}(z)]dt+\sqrt{2\operatorname{\mathcal{T }}\operatorname{\mathrm{D}}(z)}dw,$ (3) $$

where the second line comes from rescaling $\operatorname{\mathrm{D}}(z),\operatorname{\mathrm{Q}}(z)\to\operatorname{
\mathcal{T}}\operatorname{\mathrm{D}}(z),\operatorname{\mathcal{T}}
\operatorname{\mathrm{Q}}(z)$ (^1^11Equivalently we can rescale $dt\to\operatorname{\mathcal{T}}dt$ as Brownian motion $dw$ scales with $\sqrt{t}$, see (1.27) in []).

Setting $\mathcal{T}=N^{-1}$ regains the original target, but conveniently, the scale of the drift term $N^{-1}\nabla\log\pi(z)$ no longer grows with $N$ as $\log\pi(\theta)=\log p(\theta)+\sum_{n=1}^{N}\log p(y_{n}\mid\theta)+\text{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 $\operatorname{\mathcal{T}}\to 0$ converts the target distribution $\pi(z)$ into one with all probability mass concentrated on (local) maxima and transparently the SDE (3) into an ODE

$$ $\displaystyle\operatorname{\mathcal{T}}\to 0\implies dz=[\operatorname{\mathrm {D}}(z)+\operatorname{\mathrm{Q}}(z)]N^{-1}\nabla\log\pi(z)dt,$ (4) $$

which represents exactly a variant of continuous time gradient ascent on $N^{-1}\log\pi(z)$.

In Appendix C, we give specific examples demonstrating the tempered equivalence between vanilla SGD and stochastic gradient Langevin dynamics (SGLD) [] as well as SGD with momenta and stochastic gradient Hamiltonian Monte Carlo (SGHMC) [].

### 4.2 Parallel SGMCMC (as Bayesian deep ensembles)

Deep ensembles [] represent a 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 a formulation of Bayesian deep ensembles by converting the gradient descent SDE (4) typically used for deep ensemble training into a Bayesian version that is stationary for the posterior distribution by setting $\operatorname{\mathcal{T}}=N^{-1}$, i.e. adding back the noise term in (3). Assuming $\operatorname{\mathrm{D}}$ and $\operatorname{\mathrm{Q}}$ are chosen to be independent of $z$ this becomes

$$ $\displaystyle dz=[\operatorname{\mathrm{D}}+\operatorname{\mathrm{Q}}]N^{-1} \nabla\log\pi(z)dt+\sqrt{2N^{-1}\operatorname{\mathrm{D}}}dw.$ (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 []). Parallel SGMCMC provides an asymptotically unbiased approximation to $\pi(z)$ and therefore $p(\theta~{}\mid~{}y_{1:N})$ as long as the criteria discussed in A.3 are met for a single chain with the main one being the decaying learning rate which we are more accepting of in the parallel regime since the exploration requirements are divided across trajectories rather relying on a single trajectory. Parallel MCMC is also a bright avenue of research for full-batch inference [, ], and we expect these techniques to become even more prevalent with further advances in parallel hardware.

We note that in the large data regime, the additional noise term in (5) versus (4) becomes very small, and so existing deep ensembles [] can already be seen as approximating the posterior in some sense (with a small bias).

Figure 3 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 a key consideration. Serial SGMCMC requires less total computation since it runs only a single trajectory. However, the trajectory must be significantly longer and cannot be parallelized; therefore, serial SGMCMC is perhaps not as well suited to modern computational architectures as the deep ensemble or parallel SGMCMC approaches.

## 5 Experiments

We now present some experiments with posteriors that demonstrate the three key benefits of the Bayesian paradigm for machine learning as discussed in Section 1. Reproducible code for all experiments can be found in the examples folder of the [posteriors](https://github.com/normal-computing/posteriors/tree/main) repository.

### 5.1 Cold Posterior Effect

The purpose of this experiment is to investigate the generalization capabilities of the Bayesian methods discussed in 2. We use posteriors’ unified framework, Figure 3, to easily swap between Bayesian algorithms. Of particular note, we observe that the cold posterior effect is very prominent in Gaussian approximations whilst being mild for SGMCMC which can significantly outperform gradient descent approaches. Extensive experimental details can be found in Appendix D.

The cold posterior effect [, , , ] is a phenomenon that has been observed for approximate Bayesian machine learning models where increased predictive performance has been observed when targeting $p_{T}(\theta\mid y_{1:N})\propto p(\theta\mid y_{1:N})^{\frac{1}{T}}$ (as discussed in 4 with $\operatorname{\mathcal{T}}=TN^{-1}$) for $T<1$, which appears at odds with the Bayesian paradigm [].

Figure: Figure 4: Investigation into the cold-posterior effect for a range of posteriors algorithms for a CNN-LSTM model [] on the IMDB dataset []. Non-linearized Laplace EF and Laplace GGN are indistinguishable. All Gaussian approximations are diagonal. All approaches display error bars with one standard deviation over 5 random seeds.
Refer to caption: extracted/5627454/figs/cold-posterior/laplace_loss.png

Following [, ], we train a CNN-LSTM model with 2.7 million parameters on the IMDB [] for binary classification of positive or negative reviews. As in [], we use a diagonal Gaussian prior with variance $1/40$ for all methods.

The test set loss performance across a range of temperatures for Laplace, VI and SGMCMC approaches are depicted in Figure 4, alongside the baseline optimization MAP approach. See Appendix 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 on the test data. A simple Laplace approximation performs badly, however, improved performance is gained via the linearzation discussed in 2. Laplace GGN, which is equivalent to $F_{C}$ in this case (and in most machine learning models, see Appendix A.1), improves over $F_{E}$, which is in agreement with []. We observe that the reverse is 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 a strong cold posterior effect for the Gaussian approaches with significantly superior performance achieved at $T<1$. However, for the (non-Gaussian) SGMCMC approaches, only a very mild cold posterior approach is observed. This is evidence in agreement with [] that the cold posterior effect is more a result of the crudeness of posterior approximations than the of Bayesian learning itself (1). Finally, we observe that the parallel SGHMC approach (Bayesian deep ensemble) slightly outperforms the single trajectory serial approach, relatably Figure 3. And further, reducing the temperature deteriorates performance (observing from Section 4 and Appendix C that traditional deep ensemble is recovered as parallel SGHMC with $T=0$).

### 5.2 Continual learning with LoRA

The purpose of this experiment is to apply a simple Bayesian technique, in particular, an empirical Fisher Laplace approximation to mitigate catastrophic forgetting in an online learning problem. It also demonstrates the composability of posteriors with parameter-efficient fine-tuning (PEFT) [] and a large language model in Llama 2 [] from the PyTorch [] ecosystem. Extensive experimental details can be found in Appendix E.

We consider the problem of episodal continual learning [, ] where each episode’s data, $y_{n}$, is one of 20 books from the PG19 [] dataset with 15% of each book held out for testing. We compare simple online gradient descent with an approximation to the Bayesian online update (1) $p(\theta\mid y_{1:n})\propto p(\theta\mid y_{1:n-1})p(y_{n}\mid\theta)$ in Figure 6 and 6. Our baseline uses AdamW [] with a constant learning rate, and we reset all hyperparameters (including momentum) between each episode. 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 [] recommended in []. 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).

Figure: Figure 5: Continual learning with Llama 2. The online SGD and Laplace methods train one book after another, whilst the Offline SGD approach sees all books simultaneously, representing the network’s capacity. Vertical dashed lines represent episode changes.
Refer to caption: extracted/5627454/figs/continual_lora/plot_A.png

We use LoRA [] to fine-tune the query, key and output weight matrices in the last attention layer of the 7B parameter model Llama 2 model []. LoRA computes a low-rank approximation to the weight updates and is implemented in the PEFT library [], which integrates easily into posteriors code. We use standard choices for the rank ($r=8$) and scaling ($\alpha=32$) parameters.

Figure 6 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 row, the Laplace approximation stays low throughout training demonstrating that it continues to perform well on task 0 even though it is now being trained on data from tasks 1-3. In contrast, continually applying gradient descent quickly decreases the model’s performance on task 0. Figure 6 shows the difference of test loss from baseline, averaged over all episodes seen thus far. We see that the Laplace approximation approach maintains low loss averaged over all tasks.

### 5.3 Bayesian Llama 3

**Table 1: Out-of-distribution detection with Bayesian Llama 3.**
|  | Base Llama 3 | SGD | Parallel SGHMC |  |
| --- | --- | --- | --- | --- |
|  | TU | TU | TU | EU |
| AUROC $\uparrow$ | 0.62 | 0.80 | 0.90 | 0.95 |
| Loss on English statements $\downarrow$ | 4.68 | 4.20 | 4.40 |  |

The purpose of this experiment is to show how we can use posteriors to fine-tune a very recent large language model to provide the capability to decompose uncertainty (2) and provide improved out-of-distribution detection. Extensive experimental details can be found in Appendix F.

We fine-tune the last attention layer of the very recent 8B Llama 3 model [] (resulting in 218 million trainable parameters) on the TQA [] dataset, which consists of scientific textbooks. We fine-tune a baseline with SGD and a size 10 ensemble trained in parallel with posteriors’ SGHMC. We also compare against the base Llama 3 model without fine-tuning. We then 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 []).

Figure: Figure 7: Distributions of predictive uncertainties with Llama 3.
Refer to caption: extracted/5627454/figs/bayesllama/uncertainties.png

We can see from Figure 7 that although total uncertainty provides some indication for classifying the Samoan text as out-of-distribution, the epistemic uncertainty of parallel SGMHC offers a clearer split. Table 1 validates this with the AUROC metric [], whilst also displaying the categorical loss values for the next token of the English statements to verify a successful fine-tune.

Indeed, these are important results for mission-critical tasks. Expanding upon these techniques can result in higher reliability of language models as they can inform the user when the model feels high uncertain about its provided answers.

In Appendix F, we include results for serial SGHMC, which interestingly fails to provide improved epistemic uncertainty quantification. We posit that this has to do with serial SGHMC staying isolated in a single mode of the posterior in this case, and therefore, averaging along the trajectory provides little over an SGD approach as in Figure 3.

## 6 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 [] with [] focussing on prior specifications (which we do not go into detail here but is handled with broad flexibility in posteriors) and [] providing extensive experiments with full-batch MCMC to give a 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 3. Stan [] is widely used within the field of Bayesian statistics however does not support minibatching and is not easily extensible. Pyro [] is a broad framework that supports a range of models and inference techniques; however, it is more heavyweight than posteriors and harder to compose with existing PyTorch code. BlackJAX [] (and optax []) heavily inspired the posteriors functional approach. However, BlackJAX also has a primary focus on full-batch inference whilst posteriors is minibatch-first and, like other JAX [, , ] packages, does not compose with the PyTorch ecosystem.

## 7 Conclusion and outlook

We have presented a new software package posteriors for scalable and extensible uncertainty quantification with PyTorch [] models. Through numerical experiments, we have demonstrated how posteriors can be used to achieve a range of benefits over optimization-based approaches and composing with PyTorch toolkits. In particular, we highlighted that the cold posterior effect mostly occurs in 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 to maintain and extend; we already have a long list of methods and approaches to add. posteriors is fully open source, and we are keen to foster a community for feedback for feature requests and encourage contributions. The framework from Section 4 allows for simultaneous research into optimization and sampling methods. Combining that with posteriors’ broad range of utility functions (that, to our knowledge, are not available elsewhere in the PyTorch ecosystem), means that implementation and research into second-order techniques [, ] and improved discretization methods [] represent accessible avenues for future work in both optimization and SGMCMC.

## Acknowledgments and Disclosure of Funding

We thank Patrick Coles, Antonio Martinez and Samuel Power for helpful comments on the manuscript. We also thank Dmitri Saberi for support in the development of posteriors.

## Appendix A Survey of scalable Bayesian learning

We now provide a more thorough review of methods for scalable Bayesian learning, extending Section 2.

### A.1 Laplace approximation

As discussed, the Laplace approximation [, ] forms a Gaussian approximation to the posterior distribution $p(\theta\mid y_{1:N})\approx\operatorname{\mathbf{N}}(\theta\mid\theta_{\text{
MAP}},\hat{\Sigma})$ based on the second order Taylor expansion around the posterior mode $\theta_{\text{MAP}}=\operatorname*{arg\,max}_{\theta}p(\theta\mid y_{1:N})$:

$$ $\displaystyle\log p(\theta\mid y_{1:N})$ $\displaystyle\approx(\theta-\theta_{\text{MAP}})^{\top}\nabla^{2}_{\theta}\log p (\theta_{\text{MAP}}\mid y_{1:N})(\theta-\theta_{\text{MAP}})+\text{const},$ $\displaystyle\approx\log\operatorname{\mathbf{N}}(\theta\mid\theta_{\text{MAP} },\hat{\Sigma})+\text{const}.$ $$

where $\nabla_{\theta}\log p(\theta\mid y_{1:N})=0$ at $\theta=\theta_{\text{MAP}}$.

The Taylor expansion implies the use of the negative inverse Hessian $-\nabla^{2}_{\theta}\log p(\theta_{\text{MAP}}\mid y_{1:N})^{-1}$ as the covariance $\hat{\Sigma}$, however this is not guaranteed to be positive semidefinite and is therefore invalid as a covariance matrix. Instead, we can draw from the Bernstein-von Mises theorem []:

$$ $p(\theta\mid y_{1:N})\underset{N\to\infty}{\to}\operatorname{\mathbf{N}}\left( \theta\mid\theta_{\text{MAP}},N^{-1}\operatorname{\mathrm{F}}(\theta_{\text{ MAP}})^{-1}\right),$ $$

with Fisher information matrix (which is guaranteed to be positive semidefinite):

$$ $\displaystyle\operatorname{\mathrm{F}}(\theta)=\operatorname{\mathbb{E}}_{p(y \mid\theta)}[\nabla^{2}_{\theta}\log p(\theta\mid y)]=\operatorname{\mathbb{E} }_{p(y\mid\theta)}[\nabla_{\theta}\log p(\theta\mid y)\nabla_{\theta}\log p( \theta\mid y)^{\top}].$ (6) $$

In the supervised learning setting with inputs $x$, the Fisher information matrix becomes $\operatorname{\mathrm{F}}(\theta)=\operatorname{\mathbb{E}}_{p(x,y\mid\theta)}
[\nabla^{2}_{\theta}\log p(\theta\mid x,y)]$ which is intractable as we do not have access to the distribution over inputs $p(x)$. Two popular approximations use the training data $\hat{p}(x,y)=N^{-1}\sum_{i=1}^{N}\delta(x,y\mid x^{i},y^{i})$. The empirical Fisher:

$$ $\operatorname{\mathrm{F}}_{\text{E}}(\theta)=\operatorname{\mathbb{E}}_{\hat{p }(x,y)}[\nabla_{\theta}\log p(\theta\mid y)\nabla_{\theta}\log p(\theta\mid y) ^{\top}],$ $$

and conditional Fisher:

$$ $\operatorname{\mathrm{F}}_{\text{C}}(\theta)=\operatorname{\mathbb{E}}_{\hat{p }(x)p(y\mid x,\theta)}[\nabla_{\theta}\log p(\theta\mid y)\nabla_{\theta}\log p (\theta\mid y)^{\top}],$ (7) $$

with the conditional Fisher generally recommended [].

The integral with respect to $p(y\mid x,\theta)$ in (7) may initially seem difficult to compute. Fortunately, things are simplified in the common machine learning setting where the likelihood has the form $p(y\mid x,\theta)=p(y\mid f_{\theta}(x))$ where $p(y\mid z)$ is the probability density function for an exponential family distribution with natural parameter $z$. That is

$$ $\log p(y\mid z)=z^{\top}T(y)-\log\int z^{\top}T(y)dy,$ $$

for some $T(y)$. This setting includes the common cases of classification (softmax with cross-entropy loss) and regression (mean squared error loss). Then we have an exact equivalence between $F_{C}(\theta)$ and the Generalized Gauss-Newton matrix (GGN) [] matrix:

$$ $G(\theta)=\mathbb{E}_{\hat{p}(x,y)}[\nabla_{\theta}f_{\theta}(x)^{\top}H_{L}(x ,y)\nabla_{\theta}f_{\theta}(x)],$ $$

where $\nabla_{\theta}f_{\theta}(x)$ is the Jacobian of the forward function and

$$ $H_{L}(x,y)=[-\nabla^{2}_{z}p(y\mid z)]_{z=f_{\theta}(x)},$ $$

is the Hessian of the loss function evaluated at $z=f_{\theta}(x)$.

#### Caveats and computational bottlenecks

Laplace approximation training consists of first determining the posterior mode $\theta_{\text{MAP}}$ through gradient descent. This is followed by a single additional epoch to 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 $\operatorname{\mathcal{O}}(d^{2})$ elements becomes a memory bottleneck as well as the required inverting and sampling operations, which typically have runtime $\operatorname{\mathcal{O}}(d^{3})$. As such, approximations are used, such as diagonal [], KFAC [] and low-rank [] with reduced $\operatorname{\mathcal{O}}(d)$ memory. The main caveat for Laplace approximations is the accuracy of a Gaussian approximation under the extreme non-linearities in a neural network posterior [], which applies even further to reduced memory approximations.

### A.2 Variational inference

Variational inference seeks a parameterized distribution $q_{\phi}(\theta)$ that minimizes some distance from the target posterior distribution $p(\theta\mid y_{1:N})$. Most commonly [] we assume a Gaussian variational distribution $q_{\phi}(\theta)=\operatorname{\mathbf{N}}(\theta\mid\mu,\Sigma)$ with $\phi=(\mu,\Sigma)$ and the distance to be minimized is the KL divergence:

$$ $\displaystyle 0\leq\text{KL}[q_{\phi}(\theta)\mid\mid p(\theta\mid y_{1:N})]= \log p(y_{1:N})+\underbrace{\operatorname{\mathbb{E}}_{q_{\phi}(\theta)}\left[ \log\frac{q_{\phi}(\theta)}{p(\theta)}-\log p(y_{1:N}\mid\theta)\right]}_{ \text{NELBO}(\phi)}.$ (8) $$

Thus, the problem becomes one of optimization to minimize the negative evidence lower bound objective $\text{NELBO}(\phi)$, which can be achieved with stochastic gradient descent alongside the reparameterization trick [] to obtain unbiased gradients of the expectation.

#### Caveats and computational bottlenecks

Notably, the $\text{NELBO}(\phi)$ objective can be estimated stochastically by sampling from $q_{\phi}(\theta)$ and even minibatching the data $y_{1: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 a formal divergence from the posterior rather than relying on Bernstein-von Mises asymptotics. However, in the large parameter regime, the $\text{NELBO}(\phi)$ objective requires high dimensional integration at every optimization step, which can suffer from high variance even with the available tricks [].

### A.3 Stochastic gradient Markov chain Monte Carlo

A completely different paradigm to posterior approximation is to construct a Monte Carlo approximation $p(\theta\mid y_{1:N})\approx N^{-1}\sum_{k=1}^{K}\delta(\theta\mid\theta_{k})$. This is most commonly achieved by evolving a stochastic differential equation (SDE) and collecting samples along the trajectory. Following [], an SDE has $\pi(z)$ as a stationary distribution (and therefore provides an asymptotically unbiased Monte Carlo approximation) if and only if it is of the form

$$ $\displaystyle dz=[\operatorname{\mathrm{D}}(z)+\operatorname{\mathrm{Q}}(z)] \nabla\log\pi(z)dt+\nabla\cdot[\operatorname{\mathrm{D}}(z)+\operatorname{ \mathrm{Q}}(z)]dt+\sqrt{2\operatorname{\mathrm{D}}(z)}dw.$ (9) $$

with symmetric positive-semidefinite $\operatorname{\mathrm{D}}(z)$ and skew symmetric $\operatorname{\mathrm{Q}}(z)=-\operatorname{\mathrm{Q}}(z)^{\top}$. Here $z=(\theta,\omega)$ where $\omega$ are any auxiliary parameters such as momenta and thermostat [, ]. The distribution $\pi(z)$ is thus an extended target with $p(\theta\mid y_{1:N})$ marginal in $\theta$. The matrices $\operatorname{\mathrm{D}}(z)$ and $\operatorname{\mathrm{Q}}(z)$ are often chosen as independent of $z$, meaning the matrix divergence term vanishes.

The exact simulation of such a non-linear SDE is intractable, and as such, approximate discretization schemes such as the Euler-Maruyama [] method are employed. Traditional full-batch Markov chain Monte Carlo (MCMC) schemes [] use a Metropolis-Hastings step to ensure convergence of the Monte Carlo approximation is preserved through discretization. In the minibatch setting, the Metropolis-Hastings ratio cannot be estimated easily. Instead, we can use an unbiased stochastic gradient $\widehat{\nabla\log\pi}(z)\sim\operatorname{\mathbf{N}}(\cdot\mid\nabla\log\pi
(z),\mathrm{B}(z))$, then as noted in [] the stochasticity arising from minibatching, $\mathrm{B}(z)$, decreases like $\operatorname{\mathcal{O}}(\epsilon_{t}^{2})$ for learning rate $\epsilon_{t}$. A Robbins-Monro [] learning rate schedule $\sum\epsilon_{t}^{2}<\sum\epsilon_{t}=\infty$ ensures an asymptotically unbiased Monte Carlo approximation []. Additional enhancements can be made if an approximation to $\mathrm{B}(z)$ is available [].

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 [, , ] where the discretized SDE puts excessive weight on the noise term $dw$ over the informative gradient term. This is problematic since exploration is very much desired for MCMC. Often, a constant learning rate is used in practice [] with an asymptotic bias incurred. In the large parameter regime, the formulation of a 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 is available for large neural network parameters, resulting in a crude approximation. A notable further caveat is that, in contrast to 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 an online regime (1).

### A.4 Notable mentions

Perhaps the most common route to uncertainty quantification in neural networks is dropout [, ], where model parameters or neurons are randomly removed from each forward call. This approach can be effective and has a Bayesian interpretation in the posterior probability space of inclusion/exclusion, but this is naturally more restricted than the posterior approximations above.

Stochastic weight averaging (SWA) [] and SWA-Gaussian (SWAG) [] are similar in spirit to SGMCMC as they use samples from a single trajectory to generate a point or 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 [], which also has natural gradient descent [] and variational inference interpretations []. These methods are easy to implement but again suffer from the caveats of Gaussian posterior approximations above.

## Appendix B SGMCMC algorithms

Recall the generic SDE [] from Section 4 for targeting the tempered $\pi_{\operatorname{\mathcal{T}}}(z)\propto\pi(z)^{\frac{1}{N\operatorname{
\mathcal{T}}}}$ where we set $\operatorname{\mathcal{T}}=N^{-1}$ to target the Bayesian posterior

$$ $\displaystyle dz$ $\displaystyle=[\operatorname{\mathrm{D}}(z)+\operatorname{\mathrm{Q}}(z)]N^{-1 }\nabla\log\pi(z)dt+\operatorname{\mathcal{T}}\nabla\cdot[\operatorname{ \mathrm{D}}(z)+\operatorname{\mathrm{Q}}(z)]dt+\sqrt{2\operatorname{\mathcal{T }}\operatorname{\mathrm{D}}(z)}dw,$ $$

for symmetric positive-semidefinite $\operatorname{\mathrm{D}}(z)$ and skew symmetric $\operatorname{\mathrm{Q}}(z)=-\operatorname{\mathrm{Q}}(z)^{\top}$. With $z=(\theta,\omega)$ for any auxiliary variables $\omega$ and $\int\pi(\theta,\omega)d\omega=\pi(\theta)=p(\theta\mid y_{1:N})$.

The matrix divergence $\nabla\cdot\mathrm{M}(z)\in\operatorname{\mathbf{R}}^{d}$ is defined [] as

$$ $[\nabla\cdot\mathrm{M}(z)]_{i}=\sum_{j=1}^{d}\frac{\partial}{\partial z_{j}}[ \mathrm{M}(z)]_{ij},$ $$

and vanishes in the common case that the preconditioners are set to be independent of $z$, i.e. $\operatorname{\mathrm{D}}(z)=D$ and $\operatorname{\mathrm{Q}}(z)=Q$.

We now cover the most popular SGMCMC algorithms and those implemented in posteriors (all of which are also found in []).

### B.1 SGLD

The simplest SGMCMC algorithm is stochastic gradient Langevin dynamics [], which contains no auxiliary variables (i.e. $z=\theta$) and is represented by $\operatorname{\mathrm{D}}=\mathbb{I}$ and $\operatorname{\mathrm{Q}}=0$. This gives continuous-time SGLD:

$$ $d\theta=N^{-1}\nabla\log\pi(\theta)dt+\sqrt{2\operatorname{\mathcal{T}}}dw.$ $$

And Euler-Maruyama discretization:

$$ $\theta_{k+1}=\theta_{k}+\epsilon N^{-1}\nabla\log\pi(\theta_{k})+\sqrt{2 \epsilon\operatorname{\mathcal{T}}}\zeta_{k},$ $\zeta_{k}\sim\operatorname{\mathbf{N}}(\zeta\mid 0,\mathbb{I}).$ $$

### B.2 SGHMC

Stochastic gradient Hamiltonian Monte Carlo [] extends $z=(\theta,m)$ to include momenta $m$ with Gaussian target distribution

$$ $\log\pi(\theta,m)=\log\pi(\theta)-\frac{1}{2\sigma^{2}}m^{\top}m+\text{const}.$ $$

Choosing
$\operatorname{\mathrm{D}}=\begin{pmatrix}0&0\\
0&\alpha\mathbb{I}\end{pmatrix}$
and
$\operatorname{\mathrm{Q}}=\begin{pmatrix}0&-\mathbb{I}\\
\mathbb{I}&0\end{pmatrix}$
gives continuous-time SGHMC

$$ $\displaystyle d\theta$ $\displaystyle=\sigma^{-2}mdt,$ $\displaystyle dm$ $\displaystyle=N^{-1}\nabla\log\pi(\theta)dt-\alpha\sigma^{-2}mdt+\sqrt{2 \operatorname{\mathcal{T}}\alpha}dw,$ $$

and Euler-Maruyama discretization:

$$ $\displaystyle\theta_{k+1}$ $\displaystyle=\theta_{k}+\epsilon\sigma^{-2}m_{k},$ $\displaystyle m_{k+1}$ $\displaystyle=m_{k}+\epsilon N^{-1}\nabla\log\pi(\theta)-\epsilon\sigma^{-2} \alpha m_{k}+\sqrt{2\epsilon\operatorname{\mathcal{T}}\alpha}\zeta_{k},$ $\zeta_{k}\sim\operatorname{\mathbf{N}}(\zeta\mid 0,\mathbb{I}).$ $$

### B.3 SGNHT

The stochastic gradient Nosé-Hoover thermostat (SGNHT) algorithm [] uses ideas from thermodynamics [] to introduce an additional scalar variable $\xi$ that automates the selection of the $\alpha$ friction parameter in SGHMC. We have $z=(\theta,m,\xi)$ targeting
(^2^22Note that the SGNHT target distribution for $\xi$ is erroneously claimed as $\frac{1}{2d}(\xi-\alpha)^{2}$ in [] and [])

$$ $\log\pi(\theta,m,\xi)=\log\pi(\theta)-\frac{1}{2\sigma^{2}}m^{\top}m-\frac{d}{ 2}(\xi-\alpha)^{2}+\text{const}.$ $$

Choosing

$$ $\operatorname{\mathrm{D}}(\theta,m,\xi)=\begin{pmatrix}0&0&0\\ 0&\alpha\mathbb{I}&0\\ 0&0&0\end{pmatrix},\quad\operatorname{\mathrm{Q}}(\theta,m,\xi)=\begin{pmatrix }0&-\mathbb{I}&0\\ \mathbb{I}&0&m/d\\ 0&-m^{\top}/d&0\end{pmatrix},$ $$

gives continuous-time SGNHT

$$ $\displaystyle d\theta$ $\displaystyle=\sigma^{-2}mdt,$ $\displaystyle dm$ $\displaystyle=N^{-1}\nabla\log\pi(\theta)dt-\xi\sigma^{-2}mdt+\sqrt{2 \operatorname{\mathcal{T}}\alpha}dw,$ $\displaystyle d\xi$ $\displaystyle=(\sigma^{-2}d^{-1}m^{\top}m-\operatorname{\mathcal{T}})dt,$ $$

and Euler-Maruyama discretization:

$$ $\displaystyle\theta_{k+1}$ $\displaystyle=\theta_{k}+\epsilon\sigma^{-1}m_{k},$ $\displaystyle m_{k+1}$ $\displaystyle=m_{k}+\epsilon N^{-1}\nabla\log\pi(\theta)-\epsilon\sigma^{-2} \xi_{k}m_{k}+\sqrt{2\epsilon\operatorname{\mathcal{T}}\alpha}\zeta_{k},$ $\zeta_{k}\sim\operatorname{\mathbf{N}}(\zeta\mid 0,\mathbb{I}),$ $\displaystyle\xi_{k+1}$ $\displaystyle=\xi_{k}+\epsilon\left(\sigma^{-2}d^{-1}m_{k}^{\top}m_{k}- \operatorname{\mathcal{T}}\right).$ $$

Although SGNHT has the same tuning parameters $(\epsilon,\sigma,\alpha)$ as SGHMC, we note the $\alpha$ parameter only controls the noise level for the $m$ update and is absent from the deterministic dynamics, with even $\alpha=0$ representing a reasonable choice.

## Appendix C SGMCMC for optimization

Let us consider gradient descent (with momenta) [] (parameterization from PyTorch) for minimising $U(\theta)$:

$$ $\displaystyle\theta_{k+1}$ $\displaystyle=\theta_{k}+\gamma m_{k},$ $\displaystyle m_{k+1}$ $\displaystyle=\mu m_{k}-(1-\tau)\nabla U(\theta_{k}).$ $$

Here vanilla SGD $\theta_{k+1}=\theta_{k}-\gamma\nabla U(\theta_{k})$ is obtained with $\mu=0$ and $\tau=0$.

And also SGHMC B.2 with $\operatorname{\mathcal{T}}=0$ for $U(\theta)=-N^{-1}\log\pi(\theta)=-N^{-1}\log p(\theta\mid y_{1:N})$:

$$ $\displaystyle\theta_{k+1}$ $\displaystyle=\theta_{k}+\epsilon\sigma^{-2}m_{k},$ $\displaystyle m_{k+1}$ $\displaystyle=(1-\epsilon\sigma^{-2}\alpha)m_{k}-\epsilon\nabla U(\theta_{k}).$ $$

We see an equivalence between SGD and SGHMC ($T=0$) [] with

$$ $\begin{array}[]{rcl}\gamma&=&\epsilon\sigma^{-2},\\ \mu&=&1-\epsilon\sigma^{-2}\alpha,\\ \tau&=&1-\epsilon,\end{array}\quad\iff\quad\begin{array}[]{rcl}\epsilon&=&1- \tau,\\ \sigma^{-2}&=&\gamma(1-\tau)^{-1},\\ \alpha&=&(1-\mu)\gamma^{-1}.\end{array}$ (10) $$

We note that the transition from SGMCMC to optimization and back is not ubiquitous amongst methods. Indeed, the Euler-Maruyama SGNHT sampling algorithm B.3, does not result in a functional optimization algorithm at the zero temperature limit as the value thermostat can only increase and results in divergant dynamics. Also, the popular Adam [] optimizer does not fit the general sampling SDE (9) due to its use of elementwise operations, which do not have a convenient matrix-vector-product form. posteriors’ framework makes these mathematical considerations transparent and an easily accessible avenue for research.

## Appendix 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 5.1. Fully reproducible code can be found at [posteriors/examples/imdb](https://github.com/normal-computing/posteriors/blob/main/examples/imdb).

We here display more details on the results displayed in Figure 4. Figure 8 is exactly the same as Figure 4, but with the test loss for the MLE, approach added (no prior), which severely overfits. Figure 9 reports all accuracies which agree with the takeaways discussed in 5.1.

Figure: Figure 8: Cold posterior test losses (lower is better) for a range of posteriors algorithms and temperatures for the CNN-LSTM model [] on the IMDB dataset [] from 5.1 (same as Figure 4 with MLE added). Laplace EF is indistinguishable from Laplace GGN.
Refer to caption: extracted/5627454/figs/cold-posterior/laplace_loss_with_mle.png

Figure: Figure 9: Cold posterior test accuracies (higher is better) for a range of posteriors algorithms and temperatures for the CNN-LSTM model [] on the IMDB dataset [] from 5.1.
Refer to caption: extracted/5627454/figs/cold-posterior/laplace_accuracy_with_mle.png

We use exactly the same model CNN-LSTM from [, ], which consists of an embedding layer to convert input sequences into dense vectors, a 1D convolutional layer to capture local features, ReLU activation and max-pooling layers, an LSTM layer to capture long-term dependencies, and a fully connected layer to produce the final classification logits. The model has 2.7 million trainable parameters. Following [] we use a diagonal Gaussian prior with all variances set to 1/40 (aside from MLE which has no prior or equivalently infinite prior variances).

We train on the IMDB [] 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 [] with all hyperparameters set to the detaults from TorchOpt [] (learning rate $10^{-3}$).
- •
MAP: Same exact settings as MLE but with finite prior variances (set to 1/40 as with the other methods).
- •
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 our model). For test set evaluation, we sample 50 samples in parameter space and $10^{4}$ samples in logit space for the linearized approach.
- •
Variational Inference: As above, we train using the AdamW optimizer [] with all default parameters. We apply the stick-the-landing variance reduction [] and parameterization [] tricks. We use a single sample Monte Carlo estimate at each step. We train the variational variances in log space to avoid negative variances and initialise all log standard deviations to $-3$. For test set evaluation, we sample 50 samples in parameter space and $10^{4}$ samples in logit space for the linearized approach.
- •
Serial SGHMC: We train with learning rate $0.1$ and friction $\alpha=1$. 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 a 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 were run on an NVIDIA A100, and all simulations (including repeats over 5 random seeds) take $\sim 1$ day to run.

## Appendix E Experiment details: Continual learning with LoRA

In this section, we detail details for the continual learning experiment from 5.2.
Fully reproducible code can be found at [posteriors/examples/continual_lora](https://github.com/normal-computing/posteriors/tree/main/examples/continual_lora).

Our dataset [] 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 6 and 6, we use 1 book per episode, holding out the last $15\%$. We select the first twenty books from [] that satisfy the following criterion:

- 1.
Published in 1900 or later
- 2.
Is at least $100,000$ tokens, and at most $1,000,000$ tokens

We use LoRA [] on Llama 2 [] 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 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 on catastrophic forgetting [] suggested using multiple penalties (incorporating data from all previous episodes into the prior), we use a single penalty following []. 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 on 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 $r=8$ and $\alpha=32$ for LoRA, the standard settings [], and fine-tune three weight matrices in the last layer (key, query, and output projections), following the literature [].
- 2.
We stride over the texts so that all tokens have 2048 context tokens, given that there are 2048 previous tokens to
condition on.

Continual LoRA simulations were run on an NVIDIA A100; simulations over the 20 books take $\sim 6$ hours to run.

## Appendix F Experiment details: Bayesian Llama 3

Here we describe further details for our final experiment fine-tuning the 8B Llama 3 model []. Fully reproducible code can be found at [posteriors/examples/bayes_llama3](https://github.com/normal-computing/posteriors/tree/main/examples/bayes_llama3).

As previously mentioned, our 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 []. For evaluation, we generated a random set of 100 scientific factual statements from ChatGPT, which we also validated by hand. Additional hyperparameter and training details are as follows:

- 1.
We set the learning rate to be $10^{-3}$. We set alpha and beta parameters to $10^{-2}$ and 0, respectively, with the momenta initialized to 0.
- 2.
For the textbook content, we tokenize with a stride length of 300 and a stride overlap of size 100 while using a 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 are frozen except for the final attention layer.

Figure: Figure 10: Distributions of predictive uncertainties with Llama 3 (Serial SGHMC included).
Refer to caption: extracted/5627454/figs/bayesllama/uncertainties_appendix.png

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 of distribution. We also include one standard deviation error from 4 random seeds for SGD.

**Table 2: Out-of-distribution metrics with Bayesian Llama 3.**
|  | Base Llama 3 | SGD | Serial SGHMC | Parallel SGHMC |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | TU | TU | TU | EU | TU | EU |
| AUROC $\uparrow$ | 0.62 | 0.80 $\pm$ 0.02 | 0.85 | 0.17 | 0.90 | 0.95 |
| Loss on English statements $\downarrow$ | 4.68 | 4.20 $\pm$ 0.14 | 4.28 | 4.40 |  |  |
| Loss on Samoan statements $\downarrow$ | 11.41 | 9.87 $\pm$ 0.05 | 9.98 | 9.83 |  |  |

In Figure 10, we extend Figure 7 to include results for serial SGHMC. We see that serial SGHMC fails to provide a useful measure of epistemic uncertainty, which, as discussed in Section 5.3, we 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 of low epistemic uncertainty (a large agreement in produced logits) that results in worse performance than using TU.

Bayesian Llama 3 simulations were run on an NVIDIA A100; training of the SGD and ensemble approaches take $\sim 16$ hours to run whilst evaluation over the 100 statements is fast.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: - Would it be possible to extend the library to non-diagonal parametric posterior approximations in the future?
- Could the authors clarify their SGMCMC contribution in relation to recent prior work on such algorithms?
- 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.
- 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.
soundness: 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.
- Clearly written motivation for Bayesian inference in general
- Interface seems very clear and flexible due to functional paradigm
- Additional contribution in form of a new SGMCMC algorithm?
- Convincing experiments on modern architectures like lora and llama.
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.
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)
- 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.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: - figure 4: only 4 lines are plotted but 5 lines are included in the legend box?
soundness: 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. 

- The key features of the proposed library as sketched in section 3 are important for modern Bayesian deep learning. 

- The experiments cover various scenarios in current deep learning landscape. 

- the library is open-sourced and supports customization.
summary: I have increased my score from 5 to 6. 
----- 
This 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.
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). 

In particular, the following questions should be addressed:


- 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. 

- What are the key inference algorithms currently supported in the library?

- What is the practical time and memory complexity of the current implementations when compared against the existing non-Bayesian approaches? 

-- What are the limitations / scenarios that the current library does not support?

Overall 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.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - could the authors clarify how "posteriors" manages scalability in large parameter regimes, particularly regarding memory management with large covariance matrices?
- how does "posteriors" compare in ease of use and computational efficiency with similar Bayesian packages like Pyro or BlackJAX?
- 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?
soundness: 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.
- the paper provides extensive experimentation, including large datasets and Bayesian inference with large-scale models, showing the library's potential across applications
- 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
- the paper unifies stochastic gradient MCMC with traditional gradient descent, enabling efficient transitions between sampling and optimization
- the library is open-source, and provides a functional API that aligns with pytorc , and so facilitates user contributions
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.
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.
- 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
- 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.

---

confidence: 4
contribution: 4
initial_rating: 6
presentation: 2
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?
soundness: 3
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.

Moreover, 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.
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.

The 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.

Moreover, 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.
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.

The 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.

In 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.

Indeed, 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. 

I 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.

Secs. 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`.

Lastly, 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.

In 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.

Meta-review: This paper introduces a new PyTorch library, `posteriors,` for approximate Bayesian inference with deep neural networks and foundation models. Unlike existing frameworks like Pyro or BlackJAX, this library focuses on Bayesian inference on deep models with a prescriptive list of possible inference algorithms (SGMCMC, VI, and Laplace). The authors include experiments that demonstrate models that can be easily implemented with their library.

The library has some limitations, notably that it does not allow for non-diagonal covariance approximations. This omission limits the package's applicability, though the advantage is that all other methods are linear time/memory. Moreover, the presentation of this paper could be improved to motivate the package better, describe limitations, and provide usable code examples for potential users.

Nevertheless, this package will be of interest to the community. It will increase the adoption of Bayesian inference methods, mainly when used with parameter-efficient fine-tuning and LLMs. Given that the authors made a significant effort during the revision to improve the manuscript, I recommend accepting this paper. However, I strongly encourage the authors to continue to improve the manuscript before the camera-ready deadline to address reviewers' suggestions - in particular with the inclusion of code examples - to best facilitate their library's interest and adoption.

Decision: Accept (Poster)