Temperature Optimization for Bayesian Deep Learning

## Abstract

Abstract The Cold Posterior Effect (CPE) is a phenomenon in Bayesian Deep Learning
(BDL), where tempering the posterior to a cold temperature often improves the
predictive performance of the posterior predictive distribution (PPD).
Although the term ‘CPE’ suggests colder temperatures are inherently better,
the BDL community increasingly recognizes that this is not always the case.
Despite this, there remains no systematic method for finding the optimal
temperature beyond grid search. In this work, we propose a data-driven
approach to select the temperature that maximizes test log-predictive density,
treating the temperature as a model parameter and estimating it directly from
the data. We empirically demonstrate that our method performs comparably to
grid search, at a fraction of the cost, across both regression and
classification tasks. Finally, we highlight the differing perspectives on CPE
between the BDL and Generalized Bayes communities: while the former primarily
focuses on predictive performance of the PPD, the latter emphasizes calibrated
uncertainty and robustness to model misspecification; these distinct
objectives lead to different temperature preferences.

## 1 Introduction

The last two decades have seen substantial advances in the development of machine
learning tools. This has led to deployment across a range of application
domains, from image recognition and natural language processing to autonomous
vehicles and medical diagnostics, to name a few. Driven in no small part by the
technological advances that have enabled large-scale training of high-fidelity
models, these models obtain impressive generalization performance and accuracy
that are not readily explained with classical statistical wisdom.

One promising avenue that has been explored for explaining and enhancing
performance, as well as providing much-needed robustness guarantees, is
Bayesian deep learning (BDL). However, extending classical Bayesian
techniques to the machine learning setting is a formidable task—the scale of
both the data and model size render naïve extensions intractable. Combined with
the fact that neural network models are typically singular (Wei et al., ),
the mismatch between the size of the data and the number of parameters makes
many of the classical asymptotics inappropriate when applied to problems of
interest. Furthermore, a number of phenomena have been observed when analyzing
deep learning models that were previously considered curiosities or corner-cases
in the classical setting.

A prominent example of this is the so-called cold posterior effect
(CPE). Named in analogy with statistical physics, tempered posteriors are obtained by raising the posterior density to the power of an artificial inverse-temperature parameter. The
CPE refers to improved generalization performance of the posterior
predictive density (PPD) in both regression (Adlam et al., ) and
classification (Wenzel et al., ) tasks when the temperature $T$ is taken to
be cold with $0<T<1$. This peculiar effect is frequently viewed as a hack to
improve generalization performance, and has led to a plethora of works that
attempt to explain or ‘fix’ the CPE, see Section 5 for a literature
review.

#### Contribution.

There is growing recognition in the BDL community that, despite the term ‘CPE’,
colder temperatures do not always result in better predictive performance for
the PPD (Adlam et al., ; Zhang et al., ). Unfortunately, the common approach of
temperature tuning via grid search is computationally expensive, as it requires
extra posterior sampling for multiple $T$ across the grid(^1^11Consider ResNet20-CIFAR10 for example. Our approach requires only 35 minutes to find the optimal $T$, and an additional 327 minutes to compute the posterior at that optimal $T$. In contrast, a grid search across 10 temperatures takes a substantial 3270 minutes, resulting in a 9x speedup with our method.).
Sampling from a cold posterior is challenging and requires careful tuning of the
sampler with temperature diagnostics (Wenzel et al., ). To address these
issues, we propose a data-driven method to select an appropriate temperature. To
the best of our knowledge, no such dedicated tool exists without appealing to
intermediate approximations (e.g., variational inference
(Laves et al., )). Our method only requires maximizing a likelihood
function to find a suitable temperature. This can usually be done as part of the
sampler warm-up phase, and importantly, does not require any extra posterior
sampling.

## 2 Tempered Posteriors and the Cold Posterior Effect

Let $\{q(x,y),p(x,y|\theta),p(\theta)\}$ form a triplet representing the
truth-model-prior, where the data-generating mechanism is
$q(x,y)=q(y|x)q(x)$, the model $p(x,y|\theta)=p(y|x,\theta)q(x)$
is indexed by $\theta\in\Theta\subseteq\mathbb{R}^{d}$ representing
neural network weights, and the prior on $\theta$ is $p(\theta)$. Here, we focus
on a supervised learning setup with a training dataset
$\mathcal{D}=\{(x_{i},y_{i})\}_{i=1}^{n}$ containing $n$ observations drawn
from $q(x,y)$. The standard Bayesian update is derived from Bayes’ theorem,
resulting in the standard posterior distribution,

$$ $p(\theta|\mathcal{D})\propto p(\theta)\prod_{(x,y)\in\mathcal{D}}p(y|x,\theta).$ (1) $$

By introducing a ‘temperature’ parameter, we obtain a family of tempered
posteriors $p_{\beta}(\theta|\mathcal{D})$ by raising the likelihood and prior
to the power of the inverse temperature $\beta\coloneq\frac{1}{T}$
for $\beta>0$,

$$ $p_{\beta}(\theta|\mathcal{D})\propto p(\theta)^{\beta}\prod_{(x,y)\in\mathcal{ D}}p(y|x,\theta)^{\beta}.$ (2) $$

The CPE describes a phenomenon in Bayesian deep learning where the PPD,

$$ $p_{\beta}(y|x,\mathcal{D})\coloneq\mathbb{E}_{p_{\beta}(\theta|\mathcal{D})}[p (y|x,\theta)]=\int p(y|x,\theta)p_{\beta}(\theta|\mathcal{D})\,\textrm{d}\theta,$ (3) $$

which is constructed from Bayesian model averaging, can achieve better
performance (in terms of the test LPD defined below) in regression and
classification by artificially tempering the posterior to $\beta>1$
(Wenzel et al., ). While our construction in (2)
follows the convention of tempering both the likelihood and the prior
(Wenzel et al., ; Fortuin et al., ), improvements in performance have also
been observed when only tempering the likelihood
(Aitchison, ; Bachmann et al., ; Kapoor et al., ).

Define the test log predictive density (LPD) of a predictive density $\hat{p}(y|x)$ as

$$ $\operatorname{LPD}(\hat{p}(y|x))\coloneq\mathbb{E}_{q(x,y)}\log\hat{p}(y|x),$ $$

where the hat over $p$ indicates dependence on the training data $\mathcal{D}$, and
the expectation is only taken over the ‘new’ observation $(x,y)$. Note that the
test LPD is sometimes referred to as the ‘test log-likelihood’ in the
literature, and the negative test LPD is often called ‘negative log-likelihood
(NLL)’. However, we will avoid these terms as they can sound ambiguous. It can
be shown that a predictive density $\hat{p}(y|x)$ with a higher test LPD is
closer to the truth $q(y|x)$ in the sense that the Kullback-Leibler (KL) divergence
$D_{\mathrm{KL}}(q(y|x)\|\hat{p}(y|x))$ is smaller. Our primary interest in this work is to select $\beta$ such that the test LPD of (3) is high.

### 2.1 Outline

Following the reasoning that the aleatoric uncertainty in the data can be
quantified by $\beta$ (Adlam et al., ; Kapoor et al., ), we advocate a
likelihood-based approach to select $\beta$ directly from the data; details are
presented in Section 3. We empirically show that our method can
select a $\beta$ that is near-optimal according to the LPD metric. Experimental
results are presented in Section 4.

We then conclude with a detailed discussion aimed at dispelling some common
misconceptions surrounding the CPE. In particular, we review recent work in the BDL
community that explains why colder temperatures do not always lead to better
performance in terms of the test LPD metric. We also address the converse hypothesis sometimes encountered outside the BDL literature—that warmer temperatures are often better. This can hold, but only for metrics that differ fundamentally from test LPD.
Specifically, in Sections 5 and 6,
we examine existing literature on CPE from the perspectives of the BDL and
Generalized Bayes (GB) communities, respectively. We emphasize that these two
communities prioritize different objectives: in BDL, the primary focus is on
maximizing the predictive performance of the PPD, while in GB, the emphasis lies
on robustness to model misspecification and calibrated uncertainty. Furthermore,
the types of statistical models and datasets explored in each community are
notably distinct. In BDL, large-scale neural networks are trained on extensive,
nearly noiseless datasets, whereas in GB, the datasets are typically smaller,
noisier, and the models are classical regular statistical models. Consequently,
it is not surprising that each community arrives at different recommendations
regarding temperature tuning.

## 3 Temperature Selection for Test LPD

It is commonly believed that the PPD (3) at $\beta=1$ is optimal for
test LPD when the model is well-specified
(Adlam et al., ; Aitchison, ). However, this is not always the case
(see Appendix A for a counter-example). In fact,
Zhang et al. () argues that in the likelihood tempering case, $\beta=1$
can only be optimal if the expected Gibbs training loss remains unchanged with
the inclusion of new data. We follow the reasoning in Zhang et al. () and
reach a similar conclusion for posterior tempering; see
Appendix B for details. This motivates the development of an
efficient method to select the optimal $\beta$.

If we could have some theoretical grasp on how temperature affects the test LPD
of the PPD in (3), that is, $\operatorname{LPD}(p_{\beta}(y|x,\mathcal{D}))$, then
it might suggest a methodology for temperature selection. A little known result
to the BDL and GB communities is the following from singular learning theory
(Watanabe, , Lemma 3), which applies to both regular and
singular models $p(y|x,\theta)$:

$$ $\displaystyle-\mathbb{E}_{\mathcal{D}}\operatorname{LPD}(p_{\beta}(y|x, \mathcal{D}))$ $\displaystyle=-\mathbb{E}_{\mathcal{D}}\mathbb{E}_{q(x,y)}\log p_{\beta}(y|x, \mathcal{D})$ $\displaystyle=-\mathbb{E}_{q(x,y)}\log p(y|x,\theta_{\dagger})+\left[\frac{ \lambda-\nu(\beta)}{\beta}+\nu(\beta)\right]\frac{1}{n}+o\left(\frac{1}{n}\right)$ (4) $$

where $\theta_{\dagger}\in\operatorname*{arg\,min}_{\theta}D_{\mathrm{KL}}(q(y|x)\|p(
y|x,\theta))$. Here
$\lambda$ and $\nu$ are strictly positive numbers, respectively known as the
learning coefficient and singular fluctuation. They are invariants of the
underlying truth-model-prior triplet. Note that $\lambda$ is independent of $\beta$
while $\nu$ is a (complex) function of $\beta$. The functional dependence of $\nu$
on $\beta$ is unknown in the current singular learning theory literature.

Some comments on (4) are in order. This result allows for both misspecification ($\theta_{\dagger}$ is not necessarily such that the
KL divergence between the truth and $p(y|x,\theta_{\dagger})$ is zero) and singular models singular models $p(y|x,\theta)$,
including neural networks, as well as classical models satisfying standard
regularity conditions. On the other hand, it is
an asymptotic result in the sample size $n$, and hence the prior plays no role. We also remark that the relation in (4) pertains to the
average (negative) test LPD of the PPD in (3), with the
average taken over the training set $\mathcal{D}$, as indicated by the notation
$\mathbb{E}_{\mathcal{D}}$. More precise interpretation of (4) can be divided into two settings:

#### Regular models.

For well-specified regular models,
$\lambda=\nu(\beta)=d/2$ for all $\beta$
(Watanabe, ), and the second term,
$\left[\frac{\lambda-\nu(\beta)}{\beta}+\nu(\beta)\right]\frac{1}{n}$,
reduces to $\frac{d}{2n}$. This implies
that for large $n$, temperature has little impact (though it may appear in
higher-order terms in the expansion). This aligns with recent work by
McLatchie et al. (), which arrives at a similar conclusion using
different techniques. Under model misspecification, the second term can
alternatively be expressed as $\frac{\beta}{n}\mathbb{E}_{\mathcal{D}}V(n)$
(Watanabe, ), where $V(n)$ represents the functional variance and
now depends on the temperature in a non-trivial manner.

#### Singular models.

The asymptotic relation in (4) also
applies to singular models, e.g., neural network models (Wei et al., ).
However, the theoretical values of $\lambda$ are generally unknown, except for a
few simple models, such as one-layer tanh or reduced rank regression, in the
well-specification setting
(Yamazaki & Watanabe, ; Aoyagi & Watanabe, ; Rusakov & Geiger, ; Zwiernik, ).
Similarly, the singular fluctuation and its temperature dependence are unknown,
even in the well-specification setting.

Since our setup involves modern deep learning models, we are dealing with
singular models, where $\lambda$ and $\nu$ are unknown. While methods to
estimate $\lambda$ and $\nu$ from training data do exist
(Lau et al., ; Watanabe, ), they require posterior sampling
over neural network weights. In principle, we could use these sample-based
estimates to select the optimal $\beta$, but this would require posterior
sampling at multiple temperatures and is no better than a grid search. As this
approach is computationally challenging for deep learning applications, we are
motivated to propose a data-driven technique for efficiently determining the
optimal temperature, which we detail in the next section.

### 3.1 Selecting Temperature using the Tempered Model

In this section, we introduce a method to select $\beta$ that corresponds to
high test LPD in both regression and classification tasks. Our approach is
grounded in the insight that the tempered posterior can be reframed as the
posterior from an alternative model-prior pair (Zeno et al., ; Zhang et al., ).
Following these works, we define the tempered model,

$$ $p(y|x,\theta,\beta)\coloneq\frac{p(y|x,\theta)^{\beta}}{\int p(y^{\prime}|x, \theta)^{\beta}\textrm{d}y^{\prime}}.$ (5) $$

From here, the tempered posterior in (2) can be equivalently expressed
as

$$ $p_{\beta}(\theta|\mathcal{D})\propto\tilde{p}(\theta|\beta)\prod_{(x,y)\in \mathcal{D}}p(y|x,\theta,\beta),$ (6) $$

where the ‘rest of the terms’,
$\tilde{p}(\theta|\beta)\propto p(\theta)^{\beta}\prod_{x\in\mathcal{D}}\int p(
y^{\prime}|x,\theta)^{\beta}\textrm{d}y^{\prime}$,
can be viewed as a prior on $\theta$ with the ‘normalizing constant’ being a
function of $x$. Note that $\tilde{p}(\theta|\beta)$ can be seen as an
input-dependent prior (Zeno et al., ) but we suppress the dependence on $x$ in
the notation. The (inverse) temperature $\beta$ has an intuitive interpretation
of controlling the ‘spikiness’ of the tempered model. We give two examples with
Gaussian (regression) and categorical (classification) models:

#### Regression.

Given an arbitrary scalar function $\mu(x;\theta)$ and a fixed, known variance
$\sigma^{2}$, we have a Gaussian model:
$p(y|x,\theta)=\mathcal{N}(y|\mu(x;\theta),\sigma^{2})$.
This leads to

$$ $p(y|x,\theta,\beta)=\mathcal{N}\left(y\left|\,\mu(x;\theta),\frac{\sigma^{2}}{ \beta}\right.\right),\quad\tilde{p}(\theta|\beta)\propto p(\theta)^{\beta} \left(\dfrac{2\pi\sigma^{2}}{\beta(2\pi\sigma^{2})^{\beta}}\right)^{n/2}.$ $$

Therefore, the temperature effectively scales the model and prior variance. For
brevity, we suppress the dependency of $p(y|x,\theta,\beta)$ on the
fixed $\sigma^{2}$ in the notation.

#### Classification.

In $K$-class classification, we have $p(y|x,\theta)=f_{y}(x;\theta)$ for
$y\in\{1,\ldots,K\}$, where $f_{y}(\cdot)$ denotes the $y$-th entry of a
softmax output $f$. This leads to

$$ $p(y|x,\theta,\beta)=\dfrac{\exp(\beta f_{y}(x;\theta))}{\sum^{K}_{k=1}\exp( \beta f_{k}(x;\theta))},\quad\tilde{p}(\theta|\beta)\propto p(\theta)^{\beta} \prod_{x\in\mathcal{D}}\sum_{k}[f_{k}(x;\theta)]^{\beta}.$ $$

This tempered model is also known as the tempered softmax (Hinton et al., ; Agarwala et al., ). For large $\beta$,
the tempered model will concentrate most of the mass in one class, and the
converse will encourage a more uniform distribution of mass across all classes. For
$\beta>1$, the prior $\tilde{p}(\theta|\beta)$ will also favor $f$ that
concentrates mass in one class.

As $\beta$ can be used to capture aleatoric uncertainty in the data
(Adlam et al., ; Kapoor et al., ), we propose selecting $\beta$ using a
maximum likelihood estimator for the tempered model in
(5):

$$ $\hat{\theta}^{*},\hat{\beta}^{*}\coloneq\operatorname*{arg\,max}_{\theta,\beta }\frac{1}{n}\sum_{(x,y)\in\mathcal{D}}[\log p(y|x,\theta,\beta)].$ (7) $$

Standard consistency results imply that, under mild conditions, provided the
model is regular and *some* tempered model is well-specified,
choosing parameters according to (7) recovers the optimal model as $n\to\infty$
(van der Vaart, , Chapter 5.2). In practice, we optimize this using SGD,
and stop when the log-likelihood shows no further improvement on a validation
set. To ensure that $\beta$ remains positive, we reparameterize it as
$\exp(\log\beta)$ and optimize with respect to $\log\beta$. Further
implementation details are provided in Appendix C. We also
discussed several variants of our proposed method in
Appendix D.

### 3.2 Supporting theory

Ideally, we would like to theoretically show that the temperature selection
method in (7) produces high test LPD of (3).
However, it turns out that the theoretical guarantees are more natural for a
related object which is suggested by the reformulation of the tempered posterior
in (6). Namely, we can consider an
alternative PPD that is the expectation of the tempered model:

$$ $\mathbb{E}_{p_{\beta}(\theta|\mathcal{D})}p(y|x,\theta,\beta)=\int p(y|x, \theta,\beta)p_{\beta}(\theta|\mathcal{D})\textrm{d}\theta.$ (8) $$

This is also the object of study in Adlam et al. () and is to be contrasted
with the PPD in (3). We emphasize that our temperature selection
method can be used with either (3) or (8), and we will
compare the resulting performance of our method for these two PPDs in the experiments in
Section 4.

To study the theoretical properties of the temperature selection method for (8), we consider the objective at
the population level, leading to

$$ $\theta^{*},\beta^{*}\coloneq\operatorname*{arg\,max}_{\theta,\beta}\mathbb{E}_ {q(x,y)}\log p(y|x,\theta,\beta).$ $$

We justify $\beta^{*}$ in the case of Gaussian linear regression, but our
temperature methodology can be applied in far more general settings as we
demonstrate in Section 4. We first compute the test LPD constructed
with (8) and show that $\beta^{*}$ approximately maximizes a lower
bound of the test LPD.

###### Lemma 3.1 .

Consider a linear regression model
$p(y|x,\theta)=\mathcal{N}(y|x^{\top}\theta,\sigma^{2})$ with a
$d$-dimensional input $x$ and known variance $\sigma^{2}$, and a prior
$p(\theta)=\mathcal{N}(\theta|0,\sigma^{2}_{p})$ with finite variance
$\sigma^{2}_{p}$. Let
$\bm{X}\coloneq(x_{1},\ldots,x_{n})^{\top}\in\mathbb{R}^{n\times d}$
and
$\bm{\Sigma}\coloneq(\bm{X}^{\top}\bm{X}+\frac{\sigma^{2}}{\sigma^{2}_{p}}\bm{I
})^{-1}$.
The test LPD of the PPD in (8) at a fixed $\beta$ is bounded from
below:

$$ $\operatorname{LPD}(\mathbb{E}_{p_{\beta}(\theta|\mathcal{D})}[p(y|x,\theta, \beta)])>\mathbb{E}_{q(x,y)}\log p(y|x,\hat{\theta}_{\textrm{MAP}},\beta)- \frac{1}{2}\mathbb{E}_{q(x,y)}\log{(1+x^{\top}\bm{\Sigma}x)},$ $$

where $\hat{\theta}_{\textrm{MAP}}\coloneq\bm{\Sigma}\bm{X}^{\top}\bm{y}$
is the *maximum-a-posteriori* solution of the posterior
$p_{\beta}(\theta|\mathcal{D})$ at $\beta=1$ and
$\bm{y}\coloneq(y_{1},\ldots y_{n})^{\top}\in\mathbb{R}^{n}$.

The proof can be found in Appendix E. As our goal is to select
$\beta$ that maximizes test LPD, a reasonable strategy is to optimize the lower
bound presented in Lemma 3.1 with respect to $\beta$. As the second
term in the lower bound is independent of $\beta$, this is therefore equivalent
to maximizing the first term,
$\mathbb{E}_{q(x,y)}\log p(y|x,\hat{\theta}_{\text{MAP}},\beta)$. Note that this
objective function requires $\hat{\theta}_{\text{MAP}}$, which is often
unavailable in closed-form. A straightforward solution is to replace it with an
estimate from an iterative optimizer, before optimizing again with respect to
$\beta$. This comes at the cost of two optimization runs (and the effort to tune
their hyperparameters). Instead, we propose maximizing the empirical version of
$\mathbb{E}_{q(x,y)}\log p(y|x,\theta,\beta)$ with respect to $\theta$ and $\beta$
simultaneously, leading us back to (7). We study the efficacy of
this method empirically in the next section.

###### Lemma 3.1 .

## 4 Experiments

We now illustrate the behavior of the PPDs, (3) and (8),
across different values of $\beta$ on a suite of benchmark datasets for both
regression and classification tasks. We refer to (3) as SM-PD and
(8) as TM-PD, where SM and TM stand for standard model and tempered
model, respectively. Notably, our data-driven procedure for selecting $\beta$ is
agnostic to its downstream use in either SM-PD or TM-PD. We demonstrate that the
proposed method performs well in both according to the test LPD metric. While
our method aims to maximize test LPD and is supported by the theory discussed in
Section 3.2, we also evaluate the point prediction of the PPD.
Specifically, we use $\hat{y}=\mathbb{E}_{p_{\beta}(y|x,\mathcal{D})}[y]$ to compute
mean squared error (MSE)
$\frac{1}{|\mathcal{D_{\mathrm{test}}}|}\sum_{(x,y)\in\mathcal{D_{\mathrm{test}
}}}(y-\hat{y})^{2}$ for regression
tasks, and $\hat{y}=\arg\max_{y}p_{\beta}(y|x,\mathcal{D})$ to compute accuracy
$\frac{1}{|\mathcal{D_{\mathrm{test}}}|}\sum_{(x,y)\in\mathcal{D_{\mathrm{test}
}}}\mathds{1}[y=\hat{y}]$ for
classification tasks. By definition, these point predictions $\hat{y}$ are
identical across SM-PD and TM-PD, and their results are consolidated in our
reporting. We compare our method against a grid search over
$\beta\in\{0.1,0.3,1,3,\ldots,1000\}$, corresponding to increments of
roughly 0.5 on the $\log_{10}$ scale. The optimal $\beta$ from the grid search
is selected based on test LPD, MSE, or accuracy on a validation set(^2^22In contrast, our method tracks the log-likelihood of $p(y|x,\theta,\beta)$ on a validation set), and the optimal $\beta$ may differ depending on the metric
used.

To illustrate our method, we follow the experimental setup of
Wenzel et al. (). Regression tasks are conducted using a one-layer ReLU
network on UCI datasets (Concrete, Energy, Naval), while classification is
performed using a CNN on MNIST and a ResNet20 on CIFAR10, both with and without
data augmentation (DA). For the prior on neural network weights, we restrict
$p(\theta)$ to a zero-mean isotropic Gaussian, as this is a common choice for
achieving state-of-the-art performance (Izmailov et al., ). Further details
on the model and prior can be found in Appendix G.

We employ the stochastic-gradient Markov chain Monte Carlo (SGMCMC) algorithm
from (Wenzel et al., ) with a cyclical step-size scheduler
(Zhang et al., ) to sample from the posterior (details in
Appendix H). Hyperparameters are carefully tuned based on the temperature
diagnostics in Wenzel et al. () to ensure sampler convergence to the
posterior with the specified $\beta$; see Appendix I and
Appendix J for diagnostics and hyperparameters. This approach
contrasts with the common practice of tuning hyperparameters for predictive
performance. For CIFAR10, we collect 30 samples per run, and 100 for all other
datasets. Each SGMCMC run is repeated five times with different initializations
to generate five sets of posterior samples.

### 4.1 Results

#### Comparison between SM-PD and TM-PD.

The test LPD across different temperatures for SM-PD and TM-PD are shown in
Figure 1; see Appendix K for
additional figures. For Concrete, Energy, Naval and MNIST, we observe that TM-PD
generally outperforms SM-PD. This is perhaps not surprising, as we expect the
former to better account for aleatoric uncertainty. For Concrete and MNIST, the
peaks of SM-PD and TM-PD also coincide. Surprisingly, TM-PD does not outperform
SM-PD in the CIFAR10 examples.

In terms of the efficacy of our temperature selection method, we find that it
can generally recover $\beta$ with a good test LPD. An interesting observation
here is that our method tends to select the optimal $\beta$ for TM-PD in the
regression examples, and optimal $\beta$ for SM-PD in the classification
examples.

We also report MSE (for regression) and accuracy (for classification) to assess
the point predictions of the PPDs. As the point predictions are identical across
SM-PD and TM-PD by definition, we consolidate the results in
Table 1. In general, we observe that SM-PD and TM-PD outperform
both SGD and the PPD at $\beta=1$ across both metrics. Furthermore, our method
can select temperatures that achieves similar, if not better, performance than
that computed from the grid search.

Figure: Figure 1: Test LPD plotted against inverse temperature $\beta$. We compare two types of PPD: SM-PD (green) as defined in (3) and TM-PD (blue) as defined in (8). Zoomed-in versions of these curves are also provided in Figure 4 and Figure 5, respectively. In each example, we have five evaluations of $\hat{\beta}^{*}$ from our method. Each of these $\hat{\beta}^{*}$ has a corresponding test LPD computed with SM-PD (red circle) and TM-PD (red cross). Some of the red crosses in the CIFAR-10 examples are out of range. Solid lines and shaded areas represent the mean $\pm$ standard error across five repetitions. The vertical dotted lines indicate the PPD at $\beta=1$. Higher test LPD is better.
Refer to caption: x1.png

**Table 1: MSE (for regression) and accuracy (for classification) of the point predictions of the PPDs at $\beta=1$, $\beta=\hat{\beta}^{*}$ (our method), and the optimal $\beta$ obtained from grid search. We also include the predictions from SGD as a baseline. The results for SM-PD and TM-PD are consolidated, as both produce identical point predictions by definition. The presented values are means $\pm$ standard error across five repetitions, with the best value among the four methods boldfaced. The PPDs at $\beta\neq 1$ generally outperform SGD and the PPD at $\beta=1$, and our method can achieve comparable, if not better, performance than the grid search.**
| Method | MSE $\downarrow$ ($\times 10^{-3}$) | Accuracy $\uparrow$ |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Concrete | Energy | Naval | MNIST | CIFAR10 | CIFAR10 (DA) |  |
| SGD | $106\pm 20$ | $1.7\pm 0.2$ | $0.149\pm 0.063$ | $99.05\pm 0.11$ | $84.9\pm 3.4$ | $91.9\pm 0.3$ |
| $\beta=1$ | $\mathbf{75\pm 7}$ | $2.1\pm 0.1$ | $0.045\pm 0.005$ | $99.28\pm 0.04$ | $89.1\pm 0.3$ | $88.4\pm 0.2$ |
| $\beta=\hat{\beta}^{*}$ | $82\pm 4$ | $1.6\pm 0.1$ | $0.032\pm 0.009$ | $\mathbf{99.38\pm 0.03}$ | $\mathbf{89.9\pm 0.2}$ | $\mathbf{92.8\pm 0.2}$ |
| Grid | $76\pm 8$ | $\mathbf{1.4\pm 0.1}$ | $\mathbf{0.027\pm 0.008}$ | $99.32\pm 0.05$ | $\mathbf{89.9\pm 0.2}$ | $\mathbf{92.8\pm 0.4}$ |

#### Optimal temperature across data augmentation strength.

It has been frequently observed that data augmentation is one of the key ingredients
for observing CPE, and the optimal temperature may depend on the strength of
data augmentation (Bachmann et al., ). In view of this, we conducted an
experiment to determine how the optimal temperature differs across different
levels of augmentation. Here we only focus on SM-PD and present the
results in Figure 2. We observe a subtle and gradual increase in
the optimal temperature with the strength of data augmentation. Moreover, our
method can also recover temperatures that produce good test LPD and accuracy
under different augmentation levels. Note that, in addition to the well-known
CPE observed in CIFAR-10 with data augmentation, it also manifests in a milder
form in our CIFAR-10 experiments without data augmentation. This contrasts with
previous studies reporting the absence of CPE in the same setup
(Izmailov et al., ).

Figure: Figure 2: Test LPD and accuracy of CIFAR-10 plotted against inverse temperature $\beta$ under various levels of data augmentation (color). The lines and shaded areas represent the mean $\pm$ standard error across five repetitions. There are five dots for each colored curve, and each of these dots corresponds to a repetition of $\hat{\beta}^{*}$ from our method. The vertical dotted lines indicate the PPD at $\beta=1$. There is a subtle shift of peaks from left to right as the augmentation strength increases. Higher test LPD and accuracy indicate better performance.
Refer to caption: x2.png

## 5 CPE through the lens of Bayesian Deep Learning

Since the publication of Wenzel et al. (), there have been numerous attempts
to explain CPE, with particular focus on predictive metrics such as test LPD,
accuracy (for classification), and MSE (for regression). In addition to the
arguments presented in Section 3, we provide a summary of the
popular insights into CPE from the BDL literature here.

#### Poor posterior approximation.

One explanation for the presence or absence of CPE is inadequate posterior
approximation, e.g., an inappropriate choice of step size, or the omission of
Metropolis-Hastings steps in the SGMCMC sampler. However, as shown in
Appendices A and B, CPE is not solely an artifact
of poor posterior approximation and can be observed theoretically.

#### Likelihood corruption due to data augmentation.

It has been widely observed that the use of data augmentation often amplifies
CPE (Izmailov et al., ), while being less pronounced when data augmentation
is turned off. It is argued that the augmented data violate the usual
independent and identically distributed (i.i.d.) assumption imposed on the
data. However, CPE has been found to persist even after accounting for this
assumption violation (Nabarro et al., ). Therefore, CPE is unlikely to be a
mere artifact of data augmentation. In a separate analysis,
Kapoor et al. () argues that the SGMCMC sampler will converge to a
tempered posterior in the presence of data augmentation. They concluded that the
likelihood is implicitly raised to a power equal to the number of augmentations,
and raising the likelihood to a power reciprocal of this number should
approximately recover the standard posterior. However, they do not find this
adjustment alone sufficient to remove CPE completely.

#### Model misspecification.

It has been argued that data augmentation and curation may lead to model
misspecification, in particular overestimating the aleatoric uncertainty in the
data (Aitchison, ; Kapoor et al., ). Therefore, tempering is
proposed to be an effective tool for correction. This finding aligns with
Bachmann et al. (), which showed the optimal temperature dependent on both
the aleatoric uncertainty in the data and the ‘invariance’ of the model to
augmented data.

#### Prior misspecification.

While the exact interpretation of a ‘well-specified prior’ is debatable,
Fortuin et al. () conducted a large-scale experiment to study the effect
of tempering under Gaussian (isotropic or with correlated covariance) and
heavy-tailed priors (Student’s-t or Laplace) on four metrics: test LPD,
accuracy, expected calibration error (Naeini et al., ), and
out-of-distribution detection accuracy. They observed that $\beta=1$ is indeed
optimal for test LPD and accuracy when using a heavy-tailed prior. However, the
standard posteriors derived from heavy-tailed priors also tend to underperform
compared to tempered posteriors derived from Gaussian priors. Notably, according
to their experiments, warmer temperatures tend to perform better for expected
calibration error, and there is no general trend in optimal $\beta$ for
out-of-distribution detection accuracy. Therefore, their work suggests that the
optimal $\beta$ is not only dependent on the truth-model-prior triplet but also
on the evaluation metric.

#### Attempts to ‘fix’ the CPE.

Within much of the CPE literature, tempering is often seen as a ‘hack’ that
strays from the Bayesian principle. This has prompted the development of various
‘fixes’ — models and priors that induce standard posteriors with similar
predictive performance to tempered posteriors (see,
e.g., Fortuin et al., ; Kapoor et al., ; Marek et al., ). However, we argue
that this is unnecessary, as tempered posteriors of arbitrary $\beta>0$ are
special cases of the generalized Bayes posterior (Bissiri et al., ); see
Appendix L for a primer. Therefore, posterior tempering does not
deviate from the Bayesian principle.

## 6 CPE through the lens of Generalized Bayes

In GB, it is a standard practice to tune $\beta$ to improve the utility of the tempered posterior (Zhang, ; Grünwald, ; Bissiri et al., ). Specifically, the likelihood is frequently (but not
always) tempered to a warmer temperature, which contrasts with the cold
temperature used in BDL. In this section, we review some recent developments in
GB to clarify the role of tempering in GB and this apparent contradiction.

#### Ensuring posterior concentration on the KL minimizer.

Posterior
concentration on the KL minimizer
$\theta_{\dagger}=\operatorname*{arg\,min}_{\theta}D_{\mathrm{KL}}(q(y|x)|p(y|x
,\theta))$(^3^33In GB, the minimizer is assumed to be unique, though this is often not true for neural network models.) as $n\to\infty$ is often a desired property for many statistical
applications. Proofs establishing this property
(Barron & Cover, ; Zhang, ; Grünwald, ) typically assume that
the following inequality holds for all $\theta$ in the parameter space:

$$ $\mathbb{E}_{q(x,y)}\left[\left(\frac{p(y|x,\theta)}{p(y|x,\theta_{\dagger})} \right)^{\beta}\right]\leq 1,\quad\text{for all }\theta\in\Theta.$ $$

It is easy to verify that the inequality holds at $\beta=1$ for well-specified
models, where $q(x,y)=p(y|x,\theta_{\dagger})q(x)$. The same
cannot be said for misspecified models in general. However,
Grünwald () shows that this inequality can still hold for many
misspecified models for some warm temperatures
$\beta\leq\beta_{\text{critical}}<1$, and proposed the SafeBayes algorithm
to determine $\beta_{\text{critical}}$. Therefore, SafeBayes is presented as a
tool to achieve posterior concentration on $\theta_{\dagger}$.

However, posterior concentration is not always desirable in the context of
maximizing test LPD under model misspecification. Note that taking
$\beta\leq\beta_{\text{critical}}$ results in the corresponding PPD
concentrating on the ‘plug-in’ predictive density $p(y|x,\theta_{\dagger})$ as $n\to\infty$. This
may be undesirable if one prioritizes the test LPD of the PPD, since the plug-in
$p(y|x,\theta_{\dagger})$ may not exhibit good test LPD. For example, when the
misspecified model is non-convex, the PPD can lie outside the model family and
be closer to the truth than the plug-in $p(y|x,\theta_{{\dagger}})$. Therefore,
avoiding posterior concentration in this situation can actually result in a PPD
with a higher test LPD.

#### Calibrating credible and prediction intervals.

Constructing well-calibrated credible and prediction intervals with the nominal
frequentist coverage probability is known to be challenging, as obtaining the
posterior variance of a misspecified model is difficult (Kleijn & van der Vaart, , Example
2.1). These issues can be mitigated by appropriately
tuning $\beta$, which affects the spread of the posterior. To this end,
Syring & Martin () and Wu & Martin () developed algorithms to
select $\beta$ for calibrating credible and prediction intervals, respectively.

#### Calibrating prior-to-posterior information gain.

In decision-theoretic GB (Bissiri et al., , see Appendix L for a
primer), tempered posteriors (with likelihood tempering only)
are seen as an update rule that combines data information with prior belief.
From this perspective, $\beta$ can be interpreted as a ‘learning rate’ at which
information is ‘transferred’ to the posterior. Hence, it is reasonable to
calibrate the information gain at each update. Temperature selection algorithms
that follow this approach include Holmes & Walker () and
Lyddon et al. ().

#### Improving test LPD.

Concurrent to our work, McLatchie et al. () analyzed the role of
$\beta$(^4^44In McLatchie et al. (), the learning rate and, confusingly, temperature refer to $\beta$ in this work.) in
improving test LPD theoretically. They concluded that, in the moderately large
$n$ regime and assuming posterior concentration at $\theta_{{\dagger}}$, the
test LPD shows little improvement once $\beta$ becomes sufficiently large.
Although their motivation aligns with ours, their theoretical results do not
extend to singular models (e.g., neural networks) or over-parameterized models
(e.g., ResNet20 on CIFAR10).

### 6.1 Limitations of the GB framework for analyzing BDL models

The existing theory in GB, though elegant, is limited to regular and
under-parameterized models and cannot fully account for CPE observed in BDL.
Moreover, the metrics emphasized by GB often differ from those prioritized by
BDL practitioners, such as test LPD. As a result, GB works do not provide a
prescriptive methodology for selecting an appropriate temperature in BDL.
Additionally, many current temperature selection algorithms
(Grünwald, ; Syring & Martin, ; Wu & Martin, ) require repeated
posterior computations, making them impractical for modern deep learning models.
Finally, the GB literature primarily considers posteriors with likelihood
tempering only, i.e., leaving the prior without tempering. However, since many
priors $p(\theta)$ encountered in BDL are proper and bounded, prior tempering
can be regarded as rescaling the prior (Kapoor et al., , Section C.2).
Therefore, many of the arguments in GB still apply.

## 7 Conclusion

In this work, we proposed a data-driven approach for selecting a good $\beta$
for use in either of the PPDs (3) and (8). Our approach
circumvents the costly grid search method, i.e., sampling posterior at each
$\beta$ across the grid, by optimizing a likelihood function (7)
to obtain a good $\beta$. The $\hat{\beta}^{*}$ obtained via our method was shown
to achieve comparable test LPD to that obtained from a grid search, all without
performing any extra posterior sampling.

Additionally, we presented a detailed discussion to address the seeming
contradiction in the optimal $\beta$ recommendations from the BDL and GB
communities. We concluded that the optimal $\beta$ can differ depending on the
specific downstream task.

#### Limitation.

Our method is subject to several limitations that require
future attention. Firstly, a poorly-tuned optimization procedure may result in a
poor estimation of $\hat{\beta}^{*}$. Secondly, our method has only been
empirically tested with the ubiquitous Gaussian priors. Lastly, while our method
is empirically effective for modern neural network models and theoretically
justified for Gaussian linear regression, a formal guarantee for general neural
network models is still lacking.

#### Future directions.

We hope that this work will inspire further research on data-driven approaches
to select $\beta$, aimed not only at minimizing test LPD but also at optimizing
other metrics, such as expected calibration error.

## 8 Reproducibility Statement

The experiments details are summarized in the introduction of
Section 4. Proof of Lemma 3.1 is given
in Appendix E. SGMCMC is described in Appendix H.
Diagnostic check of the SGMCMC is given in Appendix I.
Implementation details of our proposed method is given in
Appendix C. The set of hyperparameters to reproduce the
experiments is given in Appendix J. Computational environment
and packages are given in Section J.1. Evaluation metrics are defined
in Appendix F. Model, prior and datasets are described in
Appendix G. Source code is published on
[https://anonymous.4open.science/r/tempered-posteriors-106E](https://anonymous.4open.science/r/tempered-posteriors-106E).

## 9 Acknowledgment

We would like to thank David Frazier and Dino Sejdinovic for helpful
discussions. This research is supported by Cloud TPUs from Google’s TPU Research
Cloud. KN is supported by the Australian Government’s Research Training Program
and a joint top-up scholarship from the Statistical Society of Australia and the
Australian Bureau of Statistics. LH is supported by the Australian Research
Council through a Discovery Early Career Researcher Award (DE240100144).

## Appendix A Optimal temperature of a Gaussian Toy Model

In this section, we compute the optimal temperature that minimizes the
2-Wasserstein distance and the KL divergence between the truth and PPD of a toy
Gaussian model. We show that: 1) the optimal temperature depends on both the
evaluation metric and the truth-model-prior triplet, and 2) the optimal
temperature that maximizes test LPD is not necessarily 1, even in the case of a
well-specified model.

In our setup, we assume a set of i.i.d. samples
$\mathcal{D}=\{x_{i}\}_{i=1}^{n}$ drawn from a univariate Gaussian truth
$q(x)=\mathcal{N}(x|0,\tau^{2})$ with a known variance $\tau^{2}$. Our model
is given by $p(x|\mu)=\mathcal{N}(x|\mu,\sigma^{2})$ with a fixed variance
$\sigma^{2}$. We also assume a Gaussian prior on $\mu$,
i.e., $p(\mu)=\mathcal{N}(\mu|0,\sigma^{2}_{p})$. While we focus on the
unsupervised setting for the ease of exposition, our argument also applies to
the supervised settings.

The tempered posterior at $\beta=\frac{1}{T}$ is given by:

$$ $\displaystyle p_{\beta}(\mu|\mathcal{D})$ $\displaystyle\propto p(\mu)^{1/T}\prod_{i=1}^{n}\mathcal{N}(x_{i}|\mu,\sigma^{ 2}T)$ $\displaystyle=\mathcal{N}\left(\mu\middle|n\bar{x}\frac{\sigma^{2}_{post}}{ \sigma^{2}T},\sigma^{2}_{post}\right)$ $\displaystyle=\mathcal{N}\left(\mu\middle|\bar{x}\left(\frac{\sigma^{2}T}{n \sigma^{2}_{p}}+1\right)^{-1},\sigma^{2}_{post}\right),$ $$

where the posterior variance is given by
$\sigma^{2}_{post}=T\left(\frac{1}{\sigma^{2}_{p}}+\frac{n}{\sigma^{2}}\right)^
{-1}$.

The PPD as defined in (3) is then:

$$ $p_{\beta}(x|\mathcal{D})=\int\mathcal{N}(x|\mu,\sigma^{2})p_{\beta}(\mu| \mathcal{D})\textrm{d}\mu=\mathcal{N}\left(x\middle|n\bar{x}\frac{\sigma^{2}_{ post}}{\sigma^{2}T},\sigma^{2}_{post}+\sigma^{2}\right)$ $$

### A.1 Optimal Temperature for Minimizing 2-Wasserstein Distance

As both $q(x)$ and $p_{\beta}(x|\mathcal{D})$ are Gaussian, the 2-Wasserstein
distance $W_{2}$ between these two density has a closed-form expression:

$$ $W_{2}=\left(\left(\frac{\sigma^{2}}{n\sigma^{2}_{p}}+1\right)^{-1}\bar{x} \right)^{2}+T\left(\frac{1}{\sigma^{2}_{p}}+\frac{n}{\sigma^{2}}\right)^{-1}+ \sigma^{2}+\tau^{2}-2\tau\left[T\left(\frac{1}{\sigma^{2}_{p}}+\frac{n}{\sigma ^{2}}\right)^{-1}+\sigma^{2}\right]^{1/2}$ $$

Taking the gradient of $W_{2}$ with respect to $T$ and set it to 0, we can
derive the optimal temperature $T^{*}$:

$$ $\displaystyle\partialderivative{W_{2}}{T}=\left(\frac{1}{\sigma^{2}_{p}}+\frac {n}{\sigma^{2}}\right)^{-1}-\tau\left[T\left(\frac{1}{\sigma^{2}_{p}}+\frac{n} {\sigma^{2}}\right)^{-1}+\sigma^{2}\right]^{-1/2}\left(\frac{1}{\sigma^{2}_{p} }+\frac{n}{\sigma^{2}}\right)^{-1}=0$ $\displaystyle\implies T^{*}=(\tau^{2}-\sigma^{2})\left(\frac{1}{\sigma^{2}_{p} }+\frac{n}{\sigma^{2}}\right),\quad T^{*}\in(0,\infty)$ $$

Therefore, we can conclude that the optimal $T^{*}\to 0$ in the well-specified
case ($\tau^{2}=\sigma^{2}$).

### A.2 Optimal Temperature for Minimizing KL Divergence

We can also obtain the closed-form expression KL divergence,
$D_{\mathrm{KL}}(q(x)\|p_{\beta}(x|\mathcal{D}))$:

$$ $\displaystyle D_{\mathrm{KL}}(q(x)\|p_{\beta}(x|\mathcal{D}))$ $\displaystyle=\mathbb{E}_{q(x)}\left[\log\frac{q(x)}{p_{\beta}(x|\mathcal{D})}\right]$ $\displaystyle=\left[\left(\frac{\sigma^{2}}{n\sigma^{2}_{p}}+1\right)^{-1}\bar {x}\right]^{2}\frac{1}{2(\sigma^{2}_{post}+\sigma^{2})}$ $\displaystyle\quad+\frac{1}{2}\left(\frac{\tau^{2}}{\sigma^{2}_{post}+\sigma^{ 2}}-1-\ln\frac{\tau^{2}}{\sigma^{2}_{post}+\sigma^{2}}\right).$ $$

Note that this object directly corresponds to test log predictive density (LPD),
which is the primary object of interest in Bayesian deep learning and our work.
Taking the derivative of the divergence with respect to $T$ and set it to 0, we
can solve for the posterior variance that minimizes the divergence:

$$ $\displaystyle\partialderivative{D_{\mathrm{KL}}}{T}=\left[-\left[\left(\frac{ \sigma^{2}}{n\sigma^{2}_{p}}+1\right)^{-1}\bar{x}\right]^{2}\frac{1}{(\sigma^{ 2}_{post}+\sigma^{2})^{2}}+\left(\frac{-\tau^{2}}{(\sigma^{2}_{post}+\sigma^{2 })^{2}}+\frac{1}{\sigma^{2}_{post}+\sigma^{2}}\right)\right]\frac{1}{2} \partialderivative{\sigma_{post}^{2}}{T}=0$ $\displaystyle\implies\sigma^{2}_{post}=\left[\left(\frac{\sigma^{2}}{n\sigma^{ 2}_{p}}+1\right)^{-1}\bar{x}\right]^{2}+\tau^{2}-\sigma^{2}$ $$

From this posterior variance, we can deduce the optimal $T^{*}$:

$$ $\displaystyle T^{*}$ $\displaystyle=\left[\left[\left(\frac{\sigma^{2}}{n\sigma^{2}_{p}}+1\right)^{- 1}\bar{x}\right]^{2}+\tau^{2}-\sigma^{2}\right]\left(\frac{1}{\sigma^{2}_{p}}+ \frac{n}{\sigma^{2}}\right)$ $\displaystyle=\left(\frac{n}{\sigma^{2}}\bar{x}\right)^{2}\left(\frac{1}{ \sigma^{2}_{p}}+\frac{n}{\sigma^{2}}\right)^{-1}+(\tau^{2}-\sigma^{2})\left( \frac{1}{\sigma^{2}_{p}}+\frac{n}{\sigma^{2}}\right),\quad T^{*}\in(0,\infty).$ $$

We can see that $T^{*}$ is non-trivial, and $T^{*}\neq 1$ in the well-specified
case ($\tau^{2}=\sigma^{2}$). This contradicts with the popular belief that
the test LPD is maximized by the PPD induced from a well-specified standard
posterior (Adlam et al., ).

## Appendix B Optimality Conditions for Test LPD at β = 1 𝛽 1 \beta=1 italic_β = 1

In this section, we follow the work of Zhang et al. () (which assumes
likelihood tempering rather than full posterior tempering) and argue that the
PPD (3) is rarely optimal in terms of test LPD at $\beta=1$. Our
approach is to demonstrate that the gradient
$\nabla_{\beta}\operatorname{LPD}(p_{\beta}(y|x,\mathcal{D}))$ is rarely zero at $\beta=1$.

###### Lemma B.1 .

Let
$p(\mathcal{D},\theta)=p(\theta)\prod_{(x,y)\in\mathcal{D}}p(y|x,\theta)$,
where we suppress the conditional dependency on $x$ in $p(\mathcal{D},\theta)$.
Then, the gradient of $\operatorname{LPD}(p_{\beta}(y|x,\mathcal{D}))$ with respect to $\beta$
is given by

$$ $\nabla_{\beta}\operatorname{LPD}(p_{\beta}(y|x,\mathcal{D}))=\mathbb{E}_{q(x,y )}\left[\mathbb{E}_{p_{\beta}(\theta|\mathcal{D}\cup(x,y))}[\log p(\mathcal{D} ,\theta)]\right]-\mathbb{E}_{p_{\beta}(\theta|\mathcal{D})}[\log p(\mathcal{D} ,\theta)],$ $$

where
$p_{\beta}(\theta|\mathcal{D}\cup(x,y))\propto p(y|x,\theta)p_{\beta}(\theta|
\mathcal{D})$
represents an update to the posterior $p_{\beta}(\theta|\mathcal{D})$ after
observing an extra data point $(x,y)$ from the truth, and the expectation
$\mathbb{E}_{q(x,y)}[\cdot]$ integrates over $(x,y)$ in
$p(\theta|\mathcal{D}\cup(x,y))$.

###### Proof.

Let $\mathbb{E}_{p_{\beta}}[\cdot]\coloneq\mathbb{E}_{p_{\beta}(\theta|\mathcal{D})
}[\cdot]$
represent the expectation with respect to the tempered posterior
$p_{\beta}(\theta|\mathcal{D})$, and let the normalizing constant of the tempered
posterior be
$Z(\mathcal{D},\beta)\coloneq\int p(\mathcal{D},\theta)^{\beta}\mathrm{d}{\theta}$.

We begin by deriving an identity that will be useful for the gradient computation. For any arbitrary $f:\mathbb{R}^{d}\to\mathbb{R}$, the following identity holds:

$$ $\displaystyle\nabla_{\beta}\mathbb{E}_{p_{\beta}}f(\theta)$ $\displaystyle=\int\nabla_{\beta}[f(\theta)p_{\beta}(\theta|\mathcal{D})] \mathrm{d}\theta$ $\displaystyle=\int f(\theta)\nabla_{\beta}\log p_{\beta}(\theta|\mathcal{D})p_ {\beta}(\theta|\mathcal{D})\mathrm{d}\theta$ $\displaystyle=\mathbb{E}_{p_{\beta}}[f(\theta)\nabla_{\beta}\log p_{\beta}( \theta|\mathcal{D})]$ $\displaystyle=\mathbb{E}_{p_{\beta}}[f(\theta)\nabla_{\beta}\log p(\mathcal{D} ,\theta)^{\beta}]-\mathbb{E}_{p_{\beta}}[f(\theta)]\nabla_{\beta}\log Z( \mathcal{D},\beta)$ $\displaystyle=\mathbb{E}_{p_{\beta}}[f(\theta)\log p(\mathcal{D},\theta)]- \mathbb{E}_{p_{\beta}}[f(\theta)]\mathbb{E}_{p_{\beta}}[\log p(\mathcal{D}, \theta)],$ $$

where $\nabla_{\beta}\log Z(\mathcal{D},\beta)$ in the second last line becomes

$$ $\displaystyle\nabla_{\beta}\log Z(\mathcal{D},\beta)=\frac{\int\nabla_{\beta}p (\mathcal{D},\theta)^{\beta}\mathrm{d}\theta}{Z(\mathcal{D},\beta)}=\int\log p (\mathcal{D},\theta)\frac{p(\mathcal{D},\theta)^{\beta}}{Z(\mathcal{D},\beta)} \mathrm{d}\theta=\mathbb{E}_{p_{\beta}}[\log p(\mathcal{D},\theta)].$ $$

The gradient can then be computed as follows:

$$ $\displaystyle\nabla_{\beta}\operatorname{LPD}(p_{\beta}(y|x,\mathcal{D}))$ $\displaystyle=\nabla_{\beta}\mathbb{E}_{q(x,y)}\log\mathbb{E}_{p_{\beta}}p(y|x ,\theta)$ $\displaystyle=\mathbb{E}_{q(x,y)}\frac{\nabla_{\beta}\mathbb{E}_{p_{\beta}}p(y |x,\theta)}{\mathbb{E}_{p_{\beta}}p(y|x,\theta)}$ $\displaystyle=\mathbb{E}_{q(x,y)}\left[\frac{\mathbb{E}_{p_{\beta}}[p(y|x, \theta)\log p(\mathcal{D},\theta)]}{\mathbb{E}_{p_{\beta}}p(y|x,\theta)}\right ]-\mathbb{E}_{p_{\beta}}[\log p(\mathcal{D},\theta)]$ $\displaystyle=\mathbb{E}_{q(x,y)}[\mathbb{E}_{p_{\beta}(\theta|\mathcal{D}\cup (x,y))}\log p(\mathcal{D},\theta)]-\mathbb{E}_{p_{\beta}}[\log p(\mathcal{D}, \theta)],$ $$

where the term inside $\mathbb{E}_{q(x,y)}[\cdot]$ in the third equality simplifies to

$$ $\displaystyle\frac{\mathbb{E}_{p_{\beta}}[p(y|x,\theta)\log p(\mathcal{D}, \theta)]}{\mathbb{E}_{p_{\beta}}p(y|x,\theta)}$ $\displaystyle=\int\log p(\mathcal{D},\theta)\frac{p(y|x,\theta)p_{\beta}( \theta|\mathcal{D})}{\int p(y|x,\theta^{\prime})p_{\beta}(\theta^{\prime}| \mathcal{D})\mathrm{d}\theta^{\prime}}\mathrm{d}\theta$ $\displaystyle=\mathbb{E}_{p_{\beta}(\theta|\mathcal{D}\cup(x,y))}\log p( \mathcal{D},\theta).$ $$

∎

We can view the terms in the RHS of Lemma B.1 as a training loss under
a posterior $\rho$:

$$ $\mathbb{E}_{\rho}L(\mathcal{D},\theta)\coloneq\mathbb{E}_{\rho}[-\log p( \mathcal{D},\theta)]=\mathbb{E}_{\rho}[-\log p(\theta)-\sum_{(x,y)\in\mathcal{ D}}\log p(y|x,\theta)].$ $$

This is similar to the training loss as defined in Zhang et al. (), but with
an additional regularizer $p(\theta)$ included in the loss. Then, we can define
underfitting as
$\mathbb{E}_{q(x,y)}[\mathbb{E}_{p_{\beta}(\theta|\mathcal{D}\cup(x,y))}L(
\mathcal{D},\theta)]<\mathbb{E}_{p_{\beta}(\theta|\mathcal{D})}L(\mathcal{D},\theta)$,
i.e., the posterior $p_{\beta}(\theta|\mathcal{D}\cup(x,y))$ will, on average,
have a lower training loss than the original $p_{\beta}(\theta|\mathcal{D})$ after
receiving an extra observation $(x,y)$ from the truth. Similarly, we can define
the converse as overfitting.

Therefore, it is not difficult to imagine that the posterior is rarely
well-fitted at $\beta=1$, and we should not expect optimality in terms of the
test LPD from the PPD at $\beta=1$.

###### Lemma B.1 .

###### Proof.

## Appendix C Details of the Temperature Selection Algorithm

Our proposed procedure for constructing a PPD at an optimal temperature is
summarized in Algorithm 1 with the following details:

- 1.
We first run SGD to maximize the log-likelihood of the tempered model on
a training set. We track the validation log-likelihood at the end of
every $\lfloor L/20\rfloor$-th epoch, where $L$ is the total number
of epochs, to save computation time. We then select the temperature with
the largest validation log-likelihood as our optimal temperature $\hat{\beta}^{*}$.
- 2.
Subsequently, we run SGMCMC to draw samples from the tempered posterior
at $\hat{\beta}^{*}$. We only start collecting samples after a burn-in
phase. During the burn-in phase, we use a linear ramp function to
control $T=\frac{1}{\beta}$ in (13), i.e. we
first run SGMCMC at $T=0$, then ramp up $T$ from 0 to
$\hat{T}^{*}=\frac{1}{\hat{\beta}^{*}}$ and continue the burn-in at
$\hat{T}^{*}$. This SGMCMC scheme is almost identical to Wenzel et al. (, Appendix
A.1), except that we run extra burn-in epochs at $\hat{T}^{*}$.

Figure: Algorithm 1 Procedure to construct a PPD at the optimal temperature.

## Appendix D Variants of the Temperature Selection Method

In this section, we compare a variant of (7) by maximizing

$$ $\hat{\theta}^{*},\hat{\beta}^{*}\coloneq\operatorname*{arg\,max}_{\theta,\beta }\frac{1}{n}\sum_{(x,y)\in\mathcal{D}}[\log p(y|x,\theta,\beta)+\log p(\theta)].$ (9) $$

We shall call this the maximum-a-posteriori (MAP) method, as $\theta$
is now being constrained by a prior $p(\theta)$. This is in contrast to the
maximum likelihood (MLE) method as introduced in (7).
We compare both MAP and MLE on the CIFAR10 experiments and show the achieved
test LPD and accuracy in Table 2. In general, we find that MAP and
MLE perform similarly well.

Another possible strategy here is the ‘marginal likelihood method’ which selects
$\beta$ that maximizes the marginal likelihood of the tempered posterior.
However, marginal likelihood is generally computationally intractable. Existing
methods that rely on Laplace approximations Immer et al. () are generally
unreliable due to the singularity of neural network models Wei et al. (),
although progress has been made to address this limitation
(Hodgkinson et al., ).

**Table 2: Comparison between strategies to select $\beta$ on the CIFAR10 experiments. Reported values are mean $\pm$ standard error across five repetitions. All values are evaluated on a test set. The accuracy and LPD are all within one standard error of difference between MAP and MLE.**
| Data | Strategy | Accuracy $\uparrow$ | LPD $\uparrow$ | $\hat{\beta}^{*}$ |
| --- | --- | --- | --- | --- |
| CIFAR10 | MAP | $89.89\pm 0.25$ | $-0.321\pm 0.010$ | $10.54\pm 9.60$ |
| MLE | $89.92\pm 0.25$ | $-0.317\pm 0.006$ | $10.44\pm 9.54$ |  |
| CIFAR10 (DA) | MAP | $92.81\pm 0.27$ | $-0.232\pm 0.008$ | $14.28\pm 0.52$ |
| MLE | $92.81\pm 0.20$ | $-0.232\pm 0.007$ | $12.79\pm 3.06$ |  |

## Appendix E Proof of Lemma 3.1

###### Lemma E.1 .

Consider a linear regression model
$p(y|x,\theta)=\mathcal{N}(y|x^{\top}\theta,\sigma^{2})$ with a
$d$-dimensional input $x$ and known variance $\sigma^{2}$, and a prior
$p(\theta)=\mathcal{N}(\theta|0,\sigma^{2}_{p})$ with finite variance
$\sigma^{2}_{p}$. Let
$\bm{X}\coloneq(x_{1},\ldots,x_{n})^{\top}\in\mathbb{R}^{n\times d}$
and
$\bm{\Sigma}\coloneq(\bm{X}^{\top}\bm{X}+\frac{\sigma^{2}}{\sigma^{2}_{p}}\bm{I
})^{-1}$.
The test LPD of the PPD in (8) at a fixed $\beta$ is bounded from
below:

$$ $\operatorname{LPD}(\mathbb{E}_{p_{\beta}(\theta|\mathcal{D})}[p(y|x,\theta, \beta)])>\mathbb{E}_{q(x,y)}\log p(y|x,\hat{\theta}_{\textrm{MAP}},\beta)- \frac{1}{2}\mathbb{E}_{q(x,y)}\log{(1+x^{\top}\bm{\Sigma}x)},$ $$

where $\hat{\theta}_{\textrm{MAP}}\coloneq\bm{\Sigma}\bm{X}^{\top}\bm{y}$
is the *maximum-a-posteriori* solution of the posterior
$p_{\beta}(\theta|\mathcal{D})$ at $\beta=1$ and
$\bm{y}\coloneq(y_{1},\ldots y_{n})^{\top}\in\mathbb{R}^{n}$.

###### Proof.

Let $\sigma^{2}_{\beta}=\frac{\sigma^{2}}{\beta}$. The tempered posterior
can be shown to follow a Gaussian distribution

$$ $p_{\beta}(\theta|\mathcal{D})\propto\mathcal{N}(\theta|0,\sigma^{2}_{p}/\beta) \prod_{x,y\in\mathcal{D}}\mathcal{N}(y|x^{\top}\theta,\sigma^{2}_{\beta}) \propto\mathcal{N}(\theta|\hat{\theta},\sigma^{2}_{\beta}\bm{\Sigma}).$ $$

Therefore, the PPD (8) of linear regression can be derived using
the identities of conditional Gaussian densities (see Bishop, , Section
2.3.3)

$$ $\displaystyle\log\mathbb{E}_{p_{\beta}(\theta|\mathcal{D})}[p(y|x,\theta,\beta)]$ $\displaystyle=\log\int p(y|x,\theta,\beta)p_{\beta}(\theta|\mathcal{D})\textrm {d}\theta$ $\displaystyle=\log\int\mathcal{N}(y|x^{\top}\theta,\sigma^{2}_{\beta})\mathcal {N}(\theta|\hat{\theta},\sigma^{2}_{\beta}\bm{\Sigma})\textrm{d}\theta$ $\displaystyle=\log\mathcal{N}(y|x^{\top}\hat{\theta},\sigma^{2}_{\beta}(1+x^{ \top}\bm{\Sigma}x))$ $\displaystyle=-\frac{1}{2}\log(1+x^{\top}\bm{\Sigma}x)-\frac{1}{2}\log(2\pi \sigma^{2}_{\beta})-\frac{(y-x^{\top}\hat{\theta})^{2}}{2\sigma^{2}_{\beta}(1+ x^{\top}\bm{\Sigma}x)}$ $\displaystyle>-\frac{1}{2}\log(1+x^{\top}\bm{\Sigma}x)-\frac{1}{2}\log(2\pi \sigma^{2}_{\beta})-\frac{(y-x^{\top}\hat{\theta})^{2}}{2\sigma^{2}_{\beta}}$ $\displaystyle=-\frac{1}{2}\log{(1+x^{\top}\bm{\Sigma}x)}+\log\mathcal{N}(y|x^{ \top}\hat{\theta},\sigma^{2}_{\beta})$ $\displaystyle=-\frac{1}{2}\log{(1+x^{\top}\bm{\Sigma}x)}+\log p(y|x,\hat{ \theta},\beta)$ $$

Note that $\bm{\Sigma}$ is positive definite, and thus
$1+x^{\top}\bm{\Sigma}x>1$ and the inequality holds. Taking expectation
with respect to $q(x,y)$ at both sides of the inequality concludes the claim.
∎

###### Lemma E.1 .

###### Proof.

## Appendix F Evaluation Metrics

In addition to the LPD evaluated on a test set $\mathcal{D_{\mathrm{test}}}$,

$$ $\mathbb{E}_{q(x,y)}[\log p_{\beta}(y|x,\mathcal{D})]\approx\dfrac{1}{|\mathcal {D_{\mathrm{test}}}|}\sum_{x,y\in\mathcal{D_{\mathrm{test}}}}\log p_{\beta}(y| x,\mathcal{D}),$ $$

we also report MSE for the regression examples

$$ $\dfrac{1}{|\mathcal{D_{\mathrm{test}}}|}\sum_{x,y\in\mathcal{D_{\mathrm{test}} }}(y-\hat{y})^{2},\quad\hat{y}=\mathbb{E}_{p_{\beta}(y|x,\mathcal{D})}[y]$ $$

and accuracy for the classification examples

$$ $\dfrac{1}{|\mathcal{D_{\mathrm{test}}}|}\sum_{x,y\in\mathcal{D_{\mathrm{test}} }}\mathds{1}[y=\hat{y}],\quad\hat{y}=\operatorname*{arg\,max}_{y}p_{\beta}(y|x ,\mathcal{D})$ $$

where $\mathds{1}$ is an indicator function. Note that the MSE and accuracy are
indifferent to the constructions of PPD (3) and (8), as
they both produce identical point predictions $\hat{y}$ in our setup.

## Appendix G Model and Prior Details

### G.1 Network architecture and datasets

The network architecture and datasets are described in this section. The
dimension of $x$ and $\theta$, and the sizes of the training, validation and
testing sets are presented in Table 3.

#### One-layer ReLU network on UCI datasets.

The scalar mean function $\mu(\cdot;\theta)$ is parameterized with a 64-neuron
hidden layer using ReLU activations. The datasets are Concrete
(Yeh, ), Naval (Coraddu & Figari, ) and Energy
(Tsanas & Xifara, ). They are all provided under the CC BY 4.0 license from
UC Irvine Machine Learning Repository.

#### CNN on MNIST.

We utilize the convolutional neural network (CNN) provided in the MNIST example
within the Flax tutorial. In broad terms, it comprises two convolutional layers
(with 32 and 64 filters, respectively), followed by two fully connected layers
(with 256 and 10 outputs). The convolutional layers employ $3\times 3$
convolutions with ReLU activations and $2\times 2$ average pooling. The MNIST
dataset (LeCun et al., ) is provided under the CC BY-SA 3.0 license. Code
for the CNN model can be found
in [https://github.com/google/flax/blob/main/examples/mnist/train.py](https://github.com/google/flax/blob/main/examples/mnist/train.py) under
the Apache License, Version 2.0.

#### ResNet20 on CIFAR10.

We use the ResNet20 architecture (He et al., ) as implemented and ported from
Wenzel et al. (). We also use the following data augmentation scheme, as in
Wenzel et al. ():

- •
Random left and right flipping, then;
- •
Border-padding 4 zero values in both horizontal and vertical direction,
followed by random cropping of the image to its original size.

The CIFAR10 dataset (Krizhevsky, ) can be found
in [https://www.cs.toronto.edu/~kriz/cifar.html](https://www.cs.toronto.edu/~kriz/cifar.html). Code for the ResNet20
model can be found
in
[https://github.com/google-research/google-research/blob/master/cold_posterior_bnn/models.py](https://github.com/google-research/google-research/blob/master/cold_posterior_bnn/models.py)
under the Apache License, Version 2.0.

### G.2 Error Variance of the Gaussian regression model

As the training data $y$ in our experiment is standardized to unit variance, it
is safe to assume that error variance $\sigma^{2}$ of the trained model will be
less than 1. Therefore, we set $\sigma^{2}=0.1$.

### G.3 Prior variance

We use an isotropic Gaussian $\mathcal{N}(0,\sigma^{2}_{p})$ for the neural
network weights. The variance $\sigma^{2}_{p}$ is specified in
Table 3.

**Table 3: Details of the datasets, size of neural network and prior variance. We use the same prior for both CIFAR10 experiments with or without data augmentation.**
|  | Concrete | Energy | Naval | MNIST | CIFAR10 |
| --- | --- | --- | --- | --- | --- |
| $\operatorname{dim}(x)$ | 8 | 8 | 14 | $28\times 28\times 1$ | $32\times 32\times 3$ |
| $\operatorname{dim}(\theta)$ | 641 | 641 | 1025 | 824458 | 273258 |
| $|\mathcal{D}|$ | 824 | 614 | 9547 | 60000 | 50000 |
| $|\mathcal{D_{\mathrm{valid}}}|$ | 103 | 77 | 1194 | 5000 | 5000 |
| $|\mathcal{D_{\mathrm{test}}}|$ | 103 | 77 | 1193 | 5000 | 5000 |
| Prior variance $\sigma^{2}_{p}$ | 0.1 | 0.1 | 1 | 0.1 | 1 |

## Appendix H SGMCMC

We use the implementation of SGMCMC as presented in Wenzel et al. (). This
corresponds to the stochastic gradient Hamiltonian Monte Carlo
(Chen et al., ) with rescaled hyperparameters and an adaptive scaling
on the Gaussian noise to ensure efficient sampling from a tempered posterior.
The samples of
$p_{\beta}(\theta|\mathcal{D})\propto\exp{-\beta U(\theta)}=\exp{-U(\theta)/T}$,
where $U(\cdot)$ is the posterior energy function defined as

$$ $U(\theta)\coloneq-\log p(\theta)-\sum_{x,y\in\mathcal{D}}\log p(y|x,\theta),$ (10) $$

can be drawn by simulating the following Langevin stochastic difference equation
(SDE) over $\theta\in\mathbb{R}^{d}$ and momentum $m\in\mathbb{R}^{d}$

$$ $\displaystyle\textrm{d}\theta$ $\displaystyle=\mathbf{M}^{-1}\,m\,\textrm{d}t,$ (11) $\displaystyle\textrm{d}m$ $\displaystyle=-\nabla_{\theta}U(\theta)\,\textrm{d}t-\gamma m\,\textrm{d}t+ \sqrt{2\gamma T}\mathbf{M}^{1/2}\textrm{d}\mathbf{W},$ (12) $$

for any friction $\gamma>0$. Here, $\mathbf{W}$ is a Wiener process,
which can be loosely interpreted as a generalized Gaussian distribution
(Leimkuhler & Matthews, ). The mass matrix $\mathbf{M}$ is a
preconditioner that can help in speeding up the convergence to the limiting
distribution of this SDE. We also prefer working with $T$ instead of the
inverse temperature $\beta$ in the sampler to facilitate sampler diagnostics and
temperature ramp-up.

In practice, the gradient of $U(\theta)$ is approximated by a minibatch gradient
estimator

$$ $\nabla_{\theta}\tilde{U}(\theta)\coloneq-\nabla_{\theta}\log p(\theta)-\dfrac{ |\mathcal{D}|}{|\mathcal{B}|}\sum_{x,y\in\mathcal{B}}\nabla_{\theta}\log p(y|x ,\theta),$ $$

where $\mathcal{B}$ denotes a minibatch, and $|\mathcal{B}|$ and $|\mathcal{D}|$ denote
the batch size and number of training samples respectively. The SDEs are then
solved numerically with a first-order symplectic Euler discretization scheme
using this minibatch gradient estimator, resulting in the following update
equations

$$ $\displaystyle m^{(t)}$ $\displaystyle=(1-h\gamma)m^{(t-1)}-h\nabla_{\theta}\tilde{U}(\theta^{(t-1)})+ \sqrt{2\gamma hT}\mathbf{M}^{1/2}\mathbf{R}^{(t)}$ (13) $\displaystyle\theta^{(t)}$ $\displaystyle=\theta^{(t-1)}+h\mathbf{M}^{-1}m^{(t)}$ $$

where $\mathbf{R}^{(t)}\sim\mathcal{N}(0,I_{d})$ is a standard Gaussian
vector. Note that the temperature shows up in the update equations and
effectively scales the random noise. This is helpful for drawing samples from
cold posteriors, which tend to have narrow, high density regions. The scaling
prevents the Markov chain from taking steps that are too large and missing the
high-density region.

The step size $h$ is often also modulated with a scheduler
$C(t):\mathbb{R}^{+}\to[0,1]$

$$ $h=h_{0}C(t),$ (14) $$

where $h_{0}$ is the initial step size. Where appropriate, we also use the
layerwise preconditioner as proposed in Wenzel et al. () to speed up
convergence and reduce approximation error.

## Appendix I SGMCMC Diagnostics

In this work, we use the kinectic temperature diagnostic (Wenzel et al., ) to
assess the quality of SGMCMC samples. This is a departure from the common
practice of using test LPD as a proxy for the posterior approximation error,
since this been shown to be a poor proxy (Deshpande et al., ). In addition, we
also report the rank-normalized split-$\hat{R}$ statistics
(Vehtari et al., ) on the (unnormalized) posterior density,
which are invariant to the permutation of neural network weights.

### I.1 Kinetic temperature

We report the expected kinetic temperature of each SGMCMC chain at different
temperatures, as proposed in Wenzel et al. (, Appendix I). The kinetic
temperature estimator is given by

$$ $\hat{T}(m)=\dfrac{m^{\top}\mathbf{M}^{-1}m}{\operatorname{dim}(m)}$ (15) $$

where $m$ and $\mathbf{M}$ are the (random) momentum and the mass matrix in the
SDE (11)-(12). For a perfect simulation of the SDE,
(15) is an unbiased estimator of the temperature of the system,
i.e., $\mathbb{E}[\hat{T}(m)]=T$ (Leimkuhler & Matthews, , Section
6.1.5).

In Table 4, the expected kinetic temperatures were computed by
averaging over the temperature samples over the whole Markov chain. We use these
estimates to gauge the simulation quality of the SDEs — an estimate closer to
the target temperature indicates a better approximation of the SDEs. We
generally expect the simulation quality to worsen in the lower temperature
regime, due to the errors from both discretization and gradient sub-sampling
becoming more prominent relative to the randomness in $m$. Otherwise, there is
no major concern about the sample quality, as observed in Table 4.

**Table 4: The expected kinetic temperatures of each SGMCMC chain are presented here. In this table, the temperatures are inverted. Kinetic temperatures that are closer to the target indicate better simulation of the SDE. The kinetic temperatures in the $\hat{\beta}^{*}$ row differ across repetitions due to the variety of $\hat{\beta}^{*}$ obtained from SGD in each repetition.**
| Target $\beta$ |  | Concrete | Energy | Naval | MNIST | CIFAR10 | CIFAR10 (DA) |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 0.1 | Rep. 1 | 0.08 | 0.08 | 0.10 | 0.10 | 0.10 | 0.10 |
| Rep. 2 | 0.10 | 0.09 | 0.10 | 0.10 | 0.10 | 0.10 |  |
| Rep. 3 | 0.09 | 0.09 | 0.10 | 0.10 | 0.10 | 0.10 |  |
| Rep. 4 | 0.08 | 0.09 | 0.10 | 0.10 | 0.10 | 0.10 |  |
| Rep. 5 | 0.09 | 0.09 | 0.10 | 0.10 | 0.10 | 0.10 |  |
| 0.3 | Rep. 1 | 0.29 | 0.24 | 0.30 | 0.30 | 0.30 | 0.30 |
| Rep. 2 | 0.30 | 0.29 | 0.30 | 0.30 | 0.30 | 0.30 |  |
| Rep. 3 | 0.30 | 0.30 | 0.30 | 0.30 | 0.30 | 0.30 |  |
| Rep. 4 | 0.30 | 0.26 | 0.30 | 0.30 | 0.30 | 0.30 |  |
| Rep. 5 | 0.30 | 0.25 | 0.30 | 0.30 | 0.30 | 0.30 |  |
| 1 | Rep. 1 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 |
| Rep. 2 | 0.99 | 1.00 | 0.98 | 1.00 | 1.00 | 1.00 |  |
| Rep. 3 | 0.99 | 0.99 | 1.00 | 1.00 | 1.00 | 1.00 |  |
| Rep. 4 | 0.99 | 1.00 | 0.99 | 1.00 | 1.00 | 1.00 |  |
| Rep. 5 | 1.00 | 0.99 | 0.99 | 1.00 | 1.00 | 1.00 |  |
| 3 | Rep. 1 | 2.98 | 2.97 | 2.98 | 3.00 | 2.99 | 2.99 |
| Rep. 2 | 2.95 | 3.00 | 2.95 | 3.00 | 2.99 | 2.99 |  |
| Rep. 3 | 2.98 | 2.99 | 2.98 | 3.00 | 2.99 | 2.99 |  |
| Rep. 4 | 2.98 | 2.99 | 2.95 | 3.00 | 2.99 | 2.99 |  |
| Rep. 5 | 2.96 | 2.98 | 2.96 | 3.00 | 2.99 | 2.99 |  |
| 10 | Rep. 1 | 9.87 | 9.94 | 9.82 | 9.99 | 9.96 | 9.92 |
| Rep. 2 | 9.88 | 10.00 | 9.69 | 9.99 | 9.96 | 9.93 |  |
| Rep. 3 | 9.87 | 9.97 | 9.89 | 9.99 | 9.96 | 9.92 |  |
| Rep. 4 | 9.88 | 9.97 | 9.68 | 9.99 | 9.96 | 9.93 |  |
| Rep. 5 | 9.90 | 9.97 | 9.78 | 9.99 | 9.96 | 9.93 |  |
| 30 | Rep. 1 | 29.61 | 29.80 | 28.48 | 29.98 | 29.77 | 29.50 |
| Rep. 2 | 29.28 | 29.85 | 28.64 | 29.98 | 29.77 | 29.59 |  |
| Rep. 3 | 29.39 | 29.76 | 28.73 | 29.96 | 29.77 | 29.54 |  |
| Rep. 4 | 28.93 | 29.91 | 27.61 | 29.96 | 29.77 | 29.55 |  |
| Rep. 5 | 29.04 | 29.85 | 28.73 | 29.97 | 29.77 | 29.63 |  |
| 100 | Rep. 1 | 95.81 | 99.07 | 84.76 | 99.87 | 98.21 | 96.34 |
| Rep. 2 | 92.08 | 98.88 | 85.24 | 99.88 | 98.20 | 97.12 |  |
| Rep. 3 | 92.25 | 97.94 | 85.66 | 99.83 | 98.19 | 96.58 |  |
| Rep. 4 | 92.20 | 99.26 | 72.74 | 99.81 | 98.15 | 96.69 |  |
| Rep. 5 | 91.28 | 98.51 | 87.84 | 99.86 | 98.21 | 97.34 |  |
| 300 | Rep. 1 | 259.87 | 293.28 | 196.04 | 299.19 | 285.58 | 276.00 |
| Rep. 2 | 248.86 | 293.13 | 207.08 | 299.19 | 285.32 | 280.62 |  |
| Rep. 3 | 228.33 | 287.56 | 187.88 | 299.06 | 285.15 | 276.42 |  |
| Rep. 4 | 221.00 | 289.86 | 133.78 | 298.99 | 285.40 | 276.79 |  |
| Rep. 5 | 229.55 | 293.08 | 213.83 | 299.14 | 285.24 | 281.91 |  |
| 1000 | Rep. 1 | 547.10 | 939.37 | 346.19 | 992.38 | 860.32 | 775.48 |
| Rep. 2 | 649.51 | 943.01 | 399.87 | 992.25 | 857.79 | 821.73 |  |
| Rep. 3 | 503.20 | 875.31 | 286.99 | 991.97 | 853.45 | 792.92 |  |
| Rep. 4 | 436.40 | 880.59 | 202.85 | 991.61 | 853.14 | 798.56 |  |
| Rep. 5 | 446.14 | 958.23 | 432.83 | 991.98 | 855.56 | 844.81 |  |
| $\hat{\beta}^{*}$ | Rep. 1 | 0.16 | 7.44 | 47.53 | 2.65 | 21.04 | 7.28 |
| Rep. 2 | 0.17 | 4.86 | 44.12 | 2.68 | 3.49 | 13.93 |  |
| Rep. 3 | 0.11 | 7.91 | 33.33 | 4.13 | 3.44 | 13.72 |  |
| Rep. 4 | 0.08 | 4.12 | 42.88 | 3.58 | 3.46 | 14.17 |  |
| Rep. 5 | 0.10 | 4.58 | 27.26 | 2.03 | 20.48 | 14.28 |  |

### I.2 Rank-normalized split- R ^ ^ 𝑅 \hat{R} over^ start_ARG italic_R end_ARG

In Table 5, we report the rank-normalized, split-$\hat{R}$
(Vehtari et al., ) on the potential
energy (10), which is invariant to the permutation of
neural network weights. This statistic is a strict improvement over the
traditional potential scale reduction factor $\hat{R}$ (Gelman & Rubin, )
and split-$\hat{R}$ (Gelman et al., ). The $\hat{R}$ statistic compares
between-chain and within-chain variances, thus requiring multiple
independently-initialized Markov chains to compute. It was designed based on the
idea that, for Markov chains that are mixing well, they should converge to the
same limiting distribution regardless of initialization. Therefore, the
between-chain and within-chain variances should roughly be the same, and a
$\hat{R}$ closer to 1 is considered better. The term ‘split’ indicates that the
Markov chains are divided in half, resulting in double the number of Markov
chains with half the original length. The variances are computed across these
doubled number of Markov chains to detect poor convergence within an original
chain. The ‘traditional’ $\hat{R}$ is computed on the original values of the
potential energy, while the ‘rank-normalized’ $\hat{R}$ is computed on the
rank-normalized values of the potential energy and does not require the limiting
distribution to have a finite mean or variance.

In Bayesian deep learning, the posterior distribution is almost always
multi-modal (Izmailov et al., ), and we should not expect the
independently-initialized Markov chains to always converge to the same limiting
distribution. Therefore, we should not be overly alarmed by the relatively
‘large’ readings of $\hat{R}$ in Table 5. These values are large by the
usual standard in Bayesian statistics but are common in Bayesian deep learning
(Izmailov et al., ; Fortuin et al., ).

**Table 5: The split-$\hat{R}$ values for the log-posterior at different (inverse) temperatures. These split-$\hat{R}$ values are computed from 5 SGMCMC chains initialized at different $\theta$. Values closer to 1 indicate better mixing of the Markov chains. $\hat{\beta}^{*}$ is left out as the temperature (and hence the posterior) differs across repetitions.**
| $\beta$ | Concrete | Energy | Naval | MNIST | CIFAR10 | CIFAR10 (DA) |
| --- | --- | --- | --- | --- | --- | --- |
| 0.1 | 1.26 | 1.19 | 1.82 | 1.07 | 1.01 | 1.04 |
| 0.3 | 1.04 | 1.16 | 1.83 | 1.03 | 1.02 | 1.08 |
| 1 | 1.54 | 1.09 | 1.94 | 1.02 | 1.13 | 1.07 |
| 3 | 2.09 | 1.10 | 2.24 | 1.03 | 1.17 | 1.46 |
| 10 | 2.79 | 1.30 | 2.78 | 1.14 | 1.69 | 1.99 |
| 30 | 2.68 | 1.77 | 3.07 | 1.27 | 1.42 | 2.30 |
| 100 | 2.83 | 2.38 | 3.25 | 1.27 | 1.24 | 2.29 |
| 300 | 2.81 | 2.40 | 3.10 | 1.19 | 1.19 | 2.19 |
| 1000 | 3.97 | 2.87 | 3.20 | 1.16 | 1.15 | 2.38 |

## Appendix J Hyperparameters and the Compute Environment

In this section, we provide a detailed explanation of the hyperparameters for
SGD and SGMCMC. The values are summarized in Table 6.

#### Learning rate and scheduler (SGD).

This is the learning rate scheduler. The cosine schedule starts from the
indicated learning rate and gradually decreases to 0 throughout the entire SGD
run. The ‘piecewise’ scheduler is used in Wenzel et al. (), where the initial
learning rate is multiplied by a value at specified epochs. We write it in the
format of (epoch, multiplier): (80, 0.1), (120, 0.01), (160, 0.001), (180,
0.0005).

#### Weight decay.

We subtract a weight decay term from the transformed gradient before scaling the
transformed gradient with a learning rate. This subtraction is performed because
we are maximizing the likelihood rather than minimizing a loss.

#### Gradient clipping.

We clip the gradient norm to a specified threshold, as implemented in Optax
(DeepMind et al., ). This helps stabilize the training procedure.

#### Learning rate and momentum (SGMCMC).

The learning rate and momentum terms will control both the initial step size
$h_{0}$ and the friction $\gamma$. The relationship between them is given in
Wenzel et al. (, Appendix B) and is repeated here:

$$ $h_{0}=\sqrt{\textit{learning rate}/n},\quad\gamma=(1-\textit{momentum})/h_{0},$ $$

where $n$ is the size of the training set.

#### Scheduler and cycle length (SGMCMC).

This is the scheduler $C(t)$ in (14) that modulates
$h$. During the burn-in phase, the scheduler is fixed at 1, i.e. $h=h_{0}$.
Then, during the sampling phase, the step size is modulated with a cyclical
cosine schedule (Zhang et al., ) with a period of the specified epoch.
One sample is collected at the end of each cycle.

#### Ramp start and end.

As we run our SGMCMC algorithm from $T=0$ and only start increasing $T$
after a specified number of epochs, these two hyperparameters indicate the epoch
when $T$ was gradually increased and the epoch when $T$ reaches the target
temperature.

#### Burn-in epochs, total epochs and usable samples (SGMCMC).

The ‘total epochs’ indicate the total number of epochs run by SGMCMC, including
epochs during the burn-in phase. ‘Usable samples’ indicate the number of samples
collected after the burn-in phase.

**Table 6: Hyperparameters for SGD and SGMCMC. We use the same set of hyperparameters for both CIFAR-10 experiments, with or without data augmentation.**
|  |  | Concrete | Energy | Naval | MNIST | CIFAR10 |
| --- | --- | --- | --- | --- | --- | --- |
| SGD | Learning rate | $10^{-6}$ | $10^{-6}$ | $10^{-8}$ | $10^{-6}$ | $10^{-6}$ |
| Scheduler | Cosine | Cosine | Cosine | Cosine | Piecewise |  |
| Momentum | 0.9 | 0.9 | 0.9 | 0.9 | 0.9 |  |
| Nesterov | No | No | No | No | No |  |
| Weight decay | 1 | 1 | 1 | 1 | 500 |  |
| Batch size | Full | Full | 128 | 128 | 128 |  |
| Total epochs | 15000 | 15000 | 10000 | 10 | 200 |  |
| Gradient clipping | $10^{6}$ | $10^{6}$ | $10^{4}$ | $10^{6}$ | $10^{6}$ |  |
| SGMCMC | Learning rate | $10^{-3}$ | $10^{-3}$ | $10^{-4}$ | 0.01 | 0.1 |
| Scheduler | Cyclical | Cyclical | Cyclical | Cyclical | Cyclical |  |
| Cycle length | 200 | 200 | 100 | 10 | 50 |  |
| Preconditioner | None | None | None | None | Layerwise |  |
| Momentum | 0.98 | 0.98 | 0.98 | 0.98 | 0.98 |  |
| Batch size | Full | Full | 128 | 128 | 128 |  |
| Ramp start (epoch) | 4800 | 4800 | 900 | 10 | 100 |  |
| Ramp end (epoch) | 5000 | 5000 | 1000 | 20 | 150 |  |
| Burn-in epochs | 10000 | 10000 | 5000 | 200 | 500 |  |
| Total epochs | 30000 | 30000 | 15000 | 1200 | 2000 |  |
| Usable samples | 100 | 100 | 100 | 100 | 30 |  |
| Gradient clipping | $10^{6}$ | $10^{6}$ | $10^{6}$ | $10^{6}$ | $10^{6}$ |  |

### J.1 Computational environment and resources

The algorithms were implemented in JAX (Bradbury et al., ). We used
Flax (Heek et al., ) and Tensorflow Probabilities
(Dillon et al., ) to implement neural networks and models. Plots were
generated with tidyverse (Wickham et al., ) and ggplot2
(Wickham, ). Rank-normalized split-$\hat{R}$ was computed with the
posterior package (Bürkner et al., ). We also used
xarray (Hoyer & Hamman, ) and arviz (Kumar et al., ) to
conduct explanatory analysis, and wandb (Biewald, ) to
track our experiments. Our code will be released.

The experiment was conducted on Google Cloud Platform utilizing TPU-V3. In each
machine, there are 8 TPU cores with 16GB of TPU memory attached to each core
(i.e., totaling 128GB of memory in each machine). However, we only utilize one
TPU core due to the difficulty in parallelizing the MCMC chains across different
cores, and effectively only have access to 16GB of memory in each machine.

On a CIFAR10 experiment, a SGD run (for our proposed algorithm) takes roughly
0.5 hour to complete. A SGMCMC run at a particular temperature takes roughly 4.5
hours. In our main experiment, as we are computing over a grid of 9 temperature
plus the optimal temperature, one repetition of the experiment takes 45.5 hours.
Multiplying this by five repetitions and we get 227.5 hours to produce one of
the CIFAR10 subplot in Figure 1. The rest of the
subplot are considerably cheaper: 16 hours for MNIST, and less than an hour for
each of Concrete, Naval and Energy.

In total, it takes 474 hours to generate the posterior samples for computing
Table 1 and Figure 1. It takes an
additional 279 hours for the experiment that varies data augmentation strength,
i.e., Figure 2. Hyperparamters tuning takes an additional 200
hours. We have also spent roughly 100 compute-hours to try out various prior for
$\beta$.

### J.2 Standard error calculation

The standard error (SE) of the means are computed with sd routine in
R. More specifically, it computes the square root of an unbiased
estimate of the variance. The upper and lower bound of the shaded areas in all
figures are ‘mean $+$ SE’ and ‘mean $-$ SE’ respectively.

## Appendix K Extra Figures From the Main Experiment

In this section, we provides extra figures to complement the results shown in
the main text.

We first shows the validation LPD plotted against the temperature in
Figure 3, and find that our method can recover temperatures in the
high performing regions. We further provide a better visualization of
Figure 1 by separating the test LPD computed with
SM-PD (3) and TM-PD (8) into Figure 4
and Figure 5 respectively. A SGD reference is also included.

Lastly, we shows the performance of the point predictions across different
temperatures in Figure 6, to complement Table 1. Note
that the point predictions are the same across SM-PD and TM-PD. We observe that
the predictions from tempered posteriors generally outperforms SGD.

Figure: Figure 3: Validation LPD plotted against inverse temperature $\beta$. This is computed with SM-PD as defined in (3). Solid lines and shaded area represent mean $\pm$ standard error across five repetitions. The vertical dotted lines indicate the PPD at $\beta=1$. There are five red dots in each plot, each of them corresponding to a repetition of $\hat{\beta}^{*}$ from our method. Higher LPD indicates better performance.
Refer to caption: x3.png

Figure: Figure 4: Test LPD plotted against inverse temperature $\beta$ with SGMCMC (green, solid). This is computed with SM-PD as defined in (3). The SGD solution (horizontal, black, dashed) is included as a reference. The SGD reference in MNIST and CIFAR10 examples performs considerably worse and is out of range. Lines and shaded area represent mean $\pm$ standard error across five repetitions. The vertical dotted lines indicate the PPD at $\beta=1$. There are five red dots in each plot, each of them corresponding to a repetition of $\hat{\beta}^{*}$ from our method. Higher LPD indicates better performance.
Refer to caption: x4.png

Figure: Figure 5: Test LPD plotted against inverse temperature $\beta$ with SGMCMC (blue, solid). This is computed with TM-PD as defined in (8). The SGD solution (horizontal, black, dashed) is included as a reference. Lines and shaded area represent mean $\pm$ standard error across five repetitions. The vertical dotted lines indicate the PPD at $\beta=1$. There are five red dots in each plot, each of them corresponding to a repetition of $\hat{\beta}^{*}$ from our method. Higher LPD indicates better performance.
Refer to caption: x5.png

Figure: Figure 6: Test MSE (Concrete, Energy, Naval) and accuracy (MNIST, CIFAR10) of the point predictions of the PPDs, plotted against inverse temperature $\beta$. The results for SM-PD (3) and TM-PD (8) are consolidated, as both produce identical point predictions by definition. The SGD solution (horizontal, black, dashed) is included as a reference. Lines and shaded area represent mean $\pm$ standard error across five repetitions. The vertical dotted lines indicate the PPD at $\beta=1$. There are five red dots in each plot, each of them corresponding to a repetition of $\hat{\beta}^{*}$ from our method. Note the MSE has been reversed, and higher value indicates better performance in all plots.
Refer to caption: x6.png

## Appendix L The Decision-Theoretic Generalized Bayesian Framework

The classical Bayesian rule (1) updating prior belief
$p(\theta)$ to the posterior belief $p(\theta|\mathcal{D})$ arises from Bayes’
Theorem. An implicit assumption in this procedure for optimal performance is
that the model is correctly specified: there exists some $\theta_{0}$ such that
$p(y|x,\theta_{0})=q(y|x)$ for all $(x,y)$. The adage that “all models
are wrong” Box () highlights that this assumption should not be
expected to hold. Consequently, it does not make sense to infer on a parameter
$\theta$ that lacks a connection to the true data-generation distribution.

Using a decision-theoretic argument, Bissiri et al. () offers an
alternative justification for using (1), even when the
model is misspecified. Suppose that we are interested in a quantity $\theta_{0}$
that minimizes the expected loss under the truth $q(x,y)$, that is,

$$ $\theta_{0}\coloneq\operatorname*{arg\,min}_{\theta}\int\ell(\theta,x,y)q(x,y) \,\textrm{d}y\,\textrm{d}x$ (16) $$

for some loss function $\ell$. Then, we would like to derive an update rule
$\psi$ that takes an observation $(x_{1},y_{1})$ and updates our prior belief
on $\theta_{0}$ to a posterior belief,

$$ $p(\theta|x_{1},y_{1})=\psi\{\ell(\theta,x_{1},y_{1}),p(\theta)\}.$ $$

Furthermore, the update rule should be *coherent*. That is, suppose that we
would like to update our posterior with two data points
$\{(x_{1},y_{1}),(x_{2},y_{2})\}$, the update rule should satisfy

$$ $\psi\{\ell(\theta,x_{2},y_{2}),\psi\{\ell(\theta,x_{1},y_{1}),p(\theta)\}\} \equiv\psi\{\ell(\theta,x_{2},y_{2})+\ell(\theta,x_{1},y_{1}),p(\theta)\}.$ $$

The coherent property will ensure that the same posterior arises regardless of
the order which the data are processed. Bissiri et al. () shows that all
coherent update rules take the form

$$ $p_{\operatorname{GB}}(\theta|\mathcal{D})\propto p(\theta)\exp\left(-\sum_{(x, y)\in\mathcal{D}}\ell(\theta,x,y)\right),$ (17) $$

which we will refer to as the general Bayesian update, and the
associated posterior $p_{\operatorname{GB}}(\theta|\mathcal{D})$ as the
generalized Bayes posterior, which has been widely analyzed
(Zhang, ; ; Jiang & Tanner, ; Bhattacharya et al., ; Alquier et al., ; Alquier & Ridgway, ). Note
that we must take care to have a finite normalizing constant to
(17),

$$ $0<\int\exp\left(-\sum_{(x,y)\in\mathcal{D}}\ell(\theta,x,y)\right)p(\theta) \textrm{d}\theta<+\infty.$ $$

Under certain regularity conditions, the generalized Bayes posterior will
concentrate on $\theta_{0}$ as $n\to\infty$; a proof can be found
in McLatchie et al. (, Lemma 5). The key advantage of this framework
in lies in its ability to directly infer the quantity of interest
(i.e., $\theta_{0}$) using a loss function, without necessitating any
probabilistic assumptions that relate $\theta_{0}$ to the data. For formal
proofs, we direct interested readers to (Bissiri et al., , Section 1).

### L.1 Calibrating information gain

The tempered posterior, as defined in (2), is a special
parameterisation of the generalized Bayes posterior achieved by setting the loss
to be $-\log p(y|x,\theta)$ and scaling it by a factor of $\beta>0$,
i.e. $\ell(\theta,x,y)=-\beta\log p(y|x,\theta)$, and concentrating the
prior to an equivalent extent. Therefore, the tempered posterior has an
interpretation of a coherent update rule for the KL minimizer
$\theta_{{\dagger}}=\operatorname*{arg\,min}_{\theta}-\beta\log p(y|x,\theta)=
\operatorname*{arg\,min}_{\theta}D_{\mathrm{KL}}(q(y|x)\|p(y|x,\theta))$,
as defined in the main text. From this perspective, the temperature has the role
of controlling the amount of information being ‘added’ to the prior in each
update step. Holmes & Walker () proposed setting $\beta$ such that the
expected information gain of our setup (LHS
of (18)) matches that of a hypothetical
experiment (RHS of (18)),

$$ $\int D(p_{\beta}(\theta|\mathcal{D}),p(\theta))q(x,y)\mathrm{d}y\mathrm{d}x= \int D(p(\theta|\mathcal{D}),p(\theta))p(y|x,\theta_{\dagger})q(x)\mathrm{d}y \mathrm{d}x$ (18) $$

where $D$ is a divergence that measures the information gain from a prior to a
posterior, and $p_{\beta}$ is a tempered posterior (with likelihood tempering
only).

On the LHS of (18), we are measuring the
expected information gain in our existing setup, i.e., $q(x,y)$ is unknown,
from a prior to a tempered posterior $p_{\beta}$. Note that this integral is a
function of $\beta$. Then, Holmes & Walker () proposed picking a $\beta$
that matches up the expected information gain from a well-specified experiment
that targets the same quantity of interest $\theta_{{\dagger}}$, i.e., having the
truth be $p(\cdot|x,\theta_{\dagger})q(x)$. In the well-specified setting, the
optimal $\beta$ is 1, and thus the information gain is computed from the prior
to a posterior at $\beta=1$. This proposed method is based on the principle
that the expected information gain should match across posteriors that are
targeting the same quantity of interest. In their follow up work,
Lyddon et al. () proposed selecting $\beta$ such that the asymptotic
Fisher information number of the generalized Bayes posterior to match to that
derived from the “loss-likelihood” bootstrap method, which is a sampling
method for general Bayes posteriors.

### L.2 Data augmentation in generalized Bayes

Accounting for data augmentation in BDL is non-trivial, as the data are no
longer i.i.d. This has led to the development of several techniques to address
this issue (Nabarro et al., ; Kapoor et al., ). However, data augmentation
can be easily incorporated into decision-theoretic GB (Bissiri et al., ),
as the generalized Bayes posterior is derived from iterative belief updates
rather than relying on the independence assumption. We can then regard the
posterior as reflecting our belief in the KL minimizer from the ‘augmented’
truth to our model.

For example, in the ResNet20-CIFAR10 example with SGMCMC in Wenzel et al. (),
the data was augmented in each epoch, and we effectively have $Mn$
(instead of $n$) number of observations in the training set
$\tilde{\mathcal{D}}=\{(\tilde{x}_{i},y_{i})\}^{Mn}_{i=1}$, where $M$ is the
number of epochs, and $\tilde{x}_{i}$ represents an augmented image. Then,
setting $\ell(x,y,\theta)=-\log p(y|x,\theta)$ and working ‘backward’,
the posterior can be regarded as representing our belief in a new quantity
$\tilde{\theta}_{{\dagger}}$. This quantity is the minimizer of the KL divergence from
an ‘augmented’ truth $\tilde{q}(x,y)=\tilde{q}(y|x)\tilde{q}(x)$, from which
the augmented data have arisen, to the model $p(y|x,\theta)$

$$ $\tilde{\theta}_{{\dagger}}=\operatorname*{arg\,min}_{\theta}\mathbb{E}_{\tilde {q}(x)}D_{\mathrm{KL}}(\tilde{q}(y|x)\|p(y|x,\theta))\,.$ $$

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 4
questions: Firstly I would like to say that although I find the paper nice to read, I am not entirely sure how novel the presented methods are, and I find the experimental section quite lacking. I will elaborate on that here:

1. **Regarding lack of novelty**: The authors claim that the use of a hold-out dataset and performing gradient based optimisation for temperature scaling is novel, yet to the best of my knowledge, this is also what is done in [1] when temperature scaling is done. The process of doing this _during_ training of the model parameters may be new, but [1] is never discussed and contrasted to. Could the authors perhaps give their viewpoint on the novelty of their methods in contrast to [1]? And as a minimum, I would expect the authors to compare their methods to these previously proposed temperature scaling methods.
2. **Regarding overclaims**: Unless I am misunderstanding something, I think the statement "The PPDs at β ̸= 1generally outperform SGD and the PPD at β = 1, and our method can achieve comparable, if not better, performance than the grid search" in Table 1 caption is incorrect. In fact there are no cases where the proposed methods statistically significantly outperform grid-search. It may perform on par, but for the claims on cost of grid-search to be validated, some metrics on fit times for the proposed methods would have to be shown.
3. **Regarding lacking metrics**: Could the authors possibly include LPD or ECE in Table 1? I find it strange to not show the LPD in this setting, when uncertainty quantification/calibration is part of the motivation for these methods, but these metrics are never shown when contrasting to other methods.

Minor comments:
1. In line 114 and 115, the authors write that BDL focuses on PPD, while GB focuses on robustness of the model including misspecification and calibrated uncertainty. I would generally disagree with this, as the BDL community largely cares about exactly uncertainty quantification and model calibration. Could the authors possibly comment on this, or as a minimum include some references?
2. The authors never cite [1].

    

[1]: On Calibration of Modern Neural Networks, Guo et. al, 2017
soundness: 2
strengths: * The paper reads very well.
* The topic of this paper is highly relevant with increasing focus on uncertainty quantification of neural networks being highly popular.
* The discussion on the cold posterior effect is interesting, and provides additional insights and bridges understanding between two communities that often are disjoint.
summary: In this paper, the authors propose a likelihood based method for optimising the temperature of the posterior predictive of Bayesian Deep Learning models. The authors discuss the commonly observed fact that the temperature of the posterior predictives obtained from BNNs are sub-optimal, requiring post-hoc optimisation of the temperature. The authors argue that grid search is too expensive, and instead propose to different ways of optimising the temperature of two posterior-predictive density formulation: one in which the temperature is introduced directly in the posterior to obtain a tempered weight posterior, and another in which the temperature is included in the posterior predictive density yielding what the authors call a tempered model. The authors propose tempering during training, finding both optimal weights and temperature during training. The authors evaluate these methods and compare them on a number of datasets. Moreover, they provide two different viewpoints on cold posterior effect (CPE), both from the generalised Bayes and Bayesian Deep learning perspective.
weaknesses: I will elaborate on these weaknesses in questions.

1. Possible lack of novelty.
2. Overclaims in experimental results.
3. Lacking metrics in experimental results.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 1
questions: - Why did you include the "singular learning theory" result?
- Why did you not include a time comparison with grid search over temp. values?
soundness: 1
strengths: I do think that an approach like the proposed could be interesting. It might a good way to find temperatures for tempered Bayesian models, but there is still active debate whether tempering Bayesian models is actually necessary or whether it is an artifact of poor optimization or data augmentations.
summary: The paper proposes a new way to set the posterior temperature in a more efficient way than simple grid search.
The paper does an MLE on both the weights and the temperature for a tempered prediction distribution to come up with the temperature parameter used for the Bayesian model.
weaknesses: - Written in a misleading way. The authors introduce theory for choosing the inverse temperature $\beta$ that is not part of the standard Bayesian literature in Section 3.0. They do not use it later on, though. This feels like a tactic to make the paper look more complex than it is. What they do in the end is find the temperature using a simple MLE. Just like one would usually train a neural network just with a tempered prediction distribution, which they derive intuitively from a dataset-dependent posterior $\tilde{p}$.
- The justification in 3.2 seems to be much better, but is not presented with enough spotlight and is lacking a lot of cases as it only treats the linear regression case and only does so for the *tempered model* PPD, not the real PPD for which we do the tuning.
- Is this grid search really expensive? I think the actually expensive part should be the SG-MCMC, no? I am missing a time comparison. Thus, the main argument of this paper is not experimentally validated.
- In line 183 you write your new posterior $\tilde{p}$ depends on x, but it depends on the xs in $D$ instead.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: The introduction seems to identify BDL with sampling. Could it be stated clearly which types of approximate bayesian inference this work applies to?

"Our primary interest in this work..." The abstract alludes to other important metrics, such as calibration and robustness. Why is the focus here on LPD?

You state "warmer temperatures are often better ... only for metrics that differ fundamentally from test LPD". However example in appendix seems to suggest otherwise?

The aim of the discussion of Watanabe's lemma 3 seems to be to "have some theoretical grasp o nhow temperature affects the test LPD...". However, it is a bit unclear to me, what conclusions you draw from this discussion?

"As beta can be used to capture aleatoric uncertainty in the data..." I am not sure exactly what is meant here. For example in the regression setting you outline, it is assumed that the variance is known and fixed. This would be the aleatory uncertainty, right? The temperature then simply scales this variance in the likelihood. Is there an assumption of model mismatch?

In the regression setting, does the maximum likelihood estimate in eq. 7 differ from a standard maximum likelihood estimate with a learned homoschedastic noise?

Is it correctly understood, that your proposed method consists of first estimating the temperature using maximum likelihood and then estimating the posterior distribution of the model parameters using SGMCMC?

In the regression examples for beta=1 how is the model variance chosen?

How useful is lemma 3.1 when it only holds for the tempered model?

Is there a typo on page 14 line 748? Should the T be there in the last equation?

Minor comments:

"singular models singular models" repeated words
soundness: 3
strengths: The paper is notationally clear and seems solid.

The topic is important and of interest to the community.

Theories related to the cold posterior effect are discussed in good detail.

Additional material in the appendix is illuminating.

Source code is published for reproducibity
summary: The paper discusses the socalled cold posterior effect in Bayesian deep learning, and presents a data driven method to estimate an optimal temperature. The method is demonstrated on standard regression and classification benchmarks. The includes detailed theoretical discussions surrounding the topic of the cold posterior effect.
weaknesses: My primary critique is that the empirical study is somewhat limited. While the proposed method appears straightforward and effective, its comparative performance with alternative approaches is not entirely clear. For instance, while the paper demonstrates how the tempered posterior can be reframed as an equivalent 'tempered model,' it does not include a comparison with direct Bayesian inference within this model, treating the temperature as a model parameter. Another common technique involves fitting the temperature posthoc on validation data, which is also absent from the analysis. Overall, the paper leaves it unclear how this method distinguishes itself from other approaches. Additionally, given that the proposed method essentially follows a specific two-step procedure for inference in models with certain parameterizations, there are numerous other inference procedures that could be relevant for comparison.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: * Do you think that why the effective learning objective for the optimal temperature parameter is different depending on the task type ? Is there any way to make consistency of the learning objective in regardless of the task type ? 


*  As mentioned above, the advantage of learning the temperature parameter does not seem clear based on the given manuscript. Are there other advantages such as computational efficiency or ease of implementation? If they exist, presenting those advantages would be better to persuade the necessity of learning the temperature parameters.

*  Based on Figure 6 in [1], the CPE still holds for varying batch size. However, the proposed methods have not been demonstrated for varying batch size. In the setting of Table 1, would the SM-PD and TM-PD still be effective in finding the consistent optimal temperature parameter if they are applied over varying batch sizes $\\{32,64,128,256\\}$?

[1] How Good is the Bayes Posterior in Deep Neural Networks Really? - ICML 20
soundness: 3
strengths: * This work tackles an interesting question for CPE that trains the temperature parameter while preserving the CPE.

* Training the temperature parameters seems marginally effective for training with domain augmentation.

* The related works are described well.
summary: This work  introduces a method for training the temperature parameter $\beta$ unlike the existing approach for Cold Posterior Effect (CPE), where the temperature parameter is set as the hyperparameter. To this end, the authors propose optimizing the posterior predictive density $p_{\beta}(y|x,\mathcal{D},\theta)$ (SM-PD) or the log likelihood of the tempered model $p(y|x,\theta,\beta)$ (TM-PD), directly and updating the model parameters $\theta$ and temperature parameter $\beta$ jointly. For regression task, the authors provide a rationale explaining how the proposed objective enables learning of the temperature parameter $\beta$. Empirically, they also demonstrate that the training the temperature parameter can yield the same performance with the performance of the grid search on UCI regression task and CIFAR 10 classification task.
weaknesses: * The proposed methods seems less consistent according to the task type because the TM-PD is effective for regression and the SM-PD  is effective for classification in obtaining the optimal temperature parameter.

* The advantage of learning the temperature parameter are not clear because the performances of the trained temperature in Table 1 seem similar to those of the grid search.

Meta-review: This paper presents a method for data-driven tuning of the temperature when tempering the posterior of a Bayesian deep learning model.  This suggests an interesting approach to addressing the so called "Cold Posterior" problem.  I.e. rather than treat this as a nuisance parameter or as a sign of Bayesian impurity, the authors propose to embrace tempering as a part of the model - i.e. parameterize the temperature and learn it from the data.  The reviewers find the work well motivated, well-written, the approach interesting and the paper clear.  However, the reviewers all voted to reject the paper (5, 3, 5, 3).  The reviewers in general found that the experiments weren't comprehensive enough and that the presented experiments didn't really seem to back up the claims of the authors.   There was some concern that there were confusing or misleading statements in the work.  There were also questions regarding novelty, e.g. compared to Guo et al.  Overall, the consensus sentiment seems to be that the paper isn't really ready from an empirical standpoint - i.e. that the experiments don't comprehensively and convincingly demonstrate that the method works.  

Given that the consensus of the reviewers is to reject the paper, the recommendation is to reject.  The reviews seem to suggest that the paper could be made much stronger with more comprehensive empirical evaluation, clearer discussion of how this relates to Bayesian theory and moving metrics like LPD and ECE into the main paper.  Hopefully, these reviews will be helpful in making the paper stronger for a future submission.

Decision: Reject