Learning Expressive Priors for Generalization and Uncertainty Estimation
  in Neural Networks
In this work, we propose a novel prior learning method for advancing
generalization and uncertainty estimation in deep neural networks. The key idea
is to exploit scalable and structured posteriors of neural networks as
informative priors with generalization guarantees. Our learned priors provide
expressive probabilistic representations at large scale, like Bayesian
counterparts of pre-trained models on ImageNet, and further produce non-vacuous
generalization bounds. We also extend this idea to a continual learning
framework, where the favorable properties of our priors are desirable. Major
enablers are our technical contributions: (1) the sums-of-Kronecker-product
computations, and (2) the derivations and optimizations of tractable objectives
that lead to improved generalization bounds. Empirically, we exhaustively show
the effectiveness of this method for uncertainty estimation and generalization.
Dominik Schnaus, Jongseok Lee, Daniel Cremers, Rudolph Triebel
# Learning Expressive Priors for Generalization and

###### Abstract

In this work, we propose a novel prior learning method for advancing generalization and uncertainty estimation in deep neural networks. The key idea is to exploit scalable and structured posteriors of neural networks as informative priors with generalization guarantees. Our learned priors provide expressive probabilistic representations at large scale, like Bayesian counterparts of pre-trained models on ImageNet, and further produce non-vacuous generalization bounds. We also extend this idea to a continual learning framework, where the favorable properties of our priors are desirable. Major enablers are our technical contributions: (1) the sums-of-Kronecker-product computations, and (2) the derivations and optimizations of tractable objectives that lead to improved generalization bounds. Empirically, we exhaustively show the effectiveness of this method for uncertainty estimation and generalization.

Machine Learning, Uncertainty Estimation, Neural Networks

## 1 Introduction

Within the deep learning approach to real-world AI problems such as autonomous driving, generalization and uncertainty estimation are one of the most important pillars. To achieve this, Bayesian Neural Networks (BNNs) (MacKay, 1992; Hinton & van Camp, 1993; Neal, 1996) leverage the tools of Bayesian statistics in order to improve generalization and uncertainty estimation in deep neural networks. Due to their potential and advancements so far, BNNs have become increasingly popular research topics (Gawlikowski et al., 2021). However, one of the open problems in BNNs is the prior specifications. While it is widely known that prior selection is a crucial step in any Bayesian modeling (Bayes, 1763), the choice of well-specified priors is generally unknown in neural networks. Consequently, many current approaches have resorted to uninformative priors, like isotropic Gaussian, despite reported signs of prior misspecifications (Wenzel et al., 2020; Fortuin, 2022).

To address the problem of prior specification, we propose a prior learning method for BNNs. Our method is grounded in sequential Bayesian inference (Opper, 1999), where the posteriors from the past are used as the prior for the future. We achieve this by relying on Laplace Approximation (LA) (MacKay, 1992) with Kronecker-factorization of the posteriors (Ritter et al., 2018; Daxberger et al., 2021). Within this framework, we devise key technical novelties to learn expressive BNN priors with generalization guarantees. First, we demonstrate the sums-of-Kronecker-product computations so that the posteriors in matrix normal distributions, _i.e_., Gaussian with Kronecker-factorized covariance matrices (Gupta & Nagar, 2000), can be tractably used as expressive BNN priors. Second, to explicitly improve generalization, we derive and optimize over a tractable approximation of PAC-Bayes bounds (McAllester, 1999; Germain et al., 2016) that lead to non-vacuous bounds, _i.e_., smaller than the upper bound of the loss. Finally, as an added benefit of our idea, a Bayesian re-interpretation to a popular continual learning architecture, namely progressive neural networks (Rusu et al., 2016), is provided for uncertainty-awareness, generalization, and resiliency to catastrophic forgetting.

By design, our method has many advantages for generalization and uncertainty estimation in neural networks. The proposed method achieves non-vacuous generalization bounds in deep learning models, while potentially avoiding prior misspecification for uncertainty estimation using BNNs. For these features, we provide computationally tractable techniques in order to learn expressive priors from large amounts of data and deep network architectures. We contend that such probabilistic representation at a large scale, e.g., pre-trained BNNs on ImageNet, can be beneficial in downstream tasks for both transfer and continual learning set-ups. Therefore, we further provide ablation studies and various experiments to show the aforementioned benefits of our method within the small and large-scale transfer and continual learning tasks. In particular, we empirically demonstrate that our priors mitigate cold posterior effects
(Wenzel et al., 2020) - a potential sign of a bad prior - and further produce generalization bounds. Moreover, within a practical scenario of robotic continual learning (Denninger and Triebel, 2018), and standardized benchmarks of few-shot classification (Tran et al., 2022) for the recent uncertainty baselines (Nado et al., 2021), considerable performance improvements over the popular methods are reported.

**Contributions** To summarize, our primary contribution is a novel method for learning scalable and structured informative priors (Section 2), backed up by (a) the sums-of-Kronecker-product computations for computational tractability (Section 2.2), (b) derivation and optimizations over tractable generalization bounds (Section 2.3), (c) a Bayesian re-interpretation of progressive neural networks for continual learning (Section 2.4), (d) exhaustive experiments to show the effectiveness of our method (Section 4).

## 2 Learning Expressive Priors

### Notation and Background

Consider a neural network \(f_{\mathbf{\theta}}\) with layers \(l\in[L]:=\{1,\dots,L\}\) which consists of a parameterized function \(\phi_{l}:\Omega_{l-1}\rightarrow\Omega_{l}\) and a non-linear activation function \(\sigma_{l}:\Omega_{l}\rightarrow\Omega_{l}\). We denote an input for a network as \(\mathbf{x}\in\Omega_{0}=:\mathcal{X}\) and all learnable parameters as a stacked vector \(\mathbf{\theta}\). Then the pre-activation \(\mathbf{s}_{l}\) and activation \(\mathbf{a}_{l}\) are recursively applied with \(\mathbf{a}_{0}=\mathbf{x}\), \(\mathbf{s}_{l}=\phi_{l}(\mathbf{a}_{l-1})\) and \(\mathbf{a}_{l}=\sigma_{l}(\mathbf{s}_{l})\). The output of the neural network \(\mathbf{y}\) is given by the last activation \(\mathbf{y}=\mathbf{a}_{L}=f_{\mathbf{\theta}}(\mathbf{x})\). The structure of the learning process is governed by different architectural choices such as fully connected, recurrent, and convolutional layers (Goodfellow et al., 2016). The parameters \(\mathbf{\theta}\) are typically obtained by maximum likelihood principles, given training data \(\mathcal{D}=\left((\mathbf{x}_{i},\mathbf{y}_{i})\right)_{i=1}^{N}\). Unfortunately, these principles lead to the lack of generalization and calibrated uncertainty in neural networks (Jospin et al., 2020).

To address these, BNNs provide a compelling alternative by applying Bayesian principles to neural networks. In BNNs, the first step is to specify a prior distribution \(\pi(\mathbf{\theta})\). Then the Bayes theorem is used to compute the posterior distribution over the parameters, given the training data: \(p(\mathbf{\theta}|\mathcal{D})\propto\pi(\mathbf{\theta})\prod_{i=1}^{N}p(\mathbf{y} _{i}|\mathbf{x}_{i},f_{\mathbf{\theta}})\). This means that each parameter is not a single value, but a probability distribution, representing the uncertainty of the model (Gawlikowski et al., 2021). The posterior distribution provides a set of plausible model parameters, which can be marginalized to compute the predictive distributions of a new sample \((\mathbf{x}^{*},\mathbf{y}^{*})\notin\mathcal{D}\) for BNNs: \(p(\mathbf{y}^{*}|\mathbf{x}^{*},\mathcal{D})=\int p(\mathbf{y}^{*}|\mathbf{ x}^{*},f_{\mathbf{\theta}})\rho(\mathbf{\theta})d\mathbf{\theta}\). In the case of neural networks, the posteriors cannot be obtained in closed form. Hence, many of the current research efforts are on accurately approximating the posteriors of BNNs.

For this, LA (MacKay, 1992) approximates the BNN posteriors with a Gaussian distributions around a local mode. Here, a prior is first specified as an isotropic Gaussian. With this prior, the maximum-a-posteriori (MAP) estimates of the parameters \(\mathbf{\hat{\theta}}\) are obtained by training the neural networks. Then, the Hessian \(\mathbf{H}=\frac{d^{2}}{d\mathbf{\theta}^{2}}\ln p(\mathbf{\theta}|\mathcal{D})= \mathbf{H}_{likelihood}+\mathbf{H}_{prior}\) is computed to obtain the BNN posteriors. In practice, the covariance matrix is usually further scaled which more generally corresponds to temperature scaling. In summary, the priors-posteriors pairs of LA are:

\[\text{Prior:}\quad\pi(\mathbf{\theta})=\mathcal{N}(\mathbf{0}, \gamma\mathbf{I}), \tag{1}\] \[\text{Posterior:}\quad p(\mathbf{\theta}|\mathcal{D})\approx\mathcal{N}( \mathbf{\hat{\theta}},(\mathbf{H}_{likelihood}+\mathbf{H}_{prior})^{-1}).\]

As the Hessian is computationally intractable for modern architectures, the Kronecker-Factored Approximate Curvature (KFAC) method is widely adopted (Martens and Grosse, 2015). KFAC approximates the true Hessian by a layer-wise block-diagonal matrix, where each diagonal block is the Kronecker-factored Fisher matrix \(\mathbf{F}_{l}\). Therefore, defining the layer-wise activation \(\mathbf{\bar{u}}_{l}\) and loss gradient w.r.t pre-activation \(\mathcal{D}\mathbf{s}_{l}\), the inverse covariance matrix of the posteriors is (Ritter et al., 2018):

\[\mathbf{H}\approx\mathbf{F}=\operatorname{diag}(\mathbf{F}_{1}, \mathbf{F}_{2},\cdots,\mathbf{F}_{L})+\gamma\mathbf{I}, \tag{2}\] \[\text{where}\ \mathbf{F}_{l}\approx\mathbb{E}[\mathcal{D} \mathbf{s}_{l}(\mathcal{D}\mathbf{s}_{l})^{T}]\otimes\mathbb{E}[\mathbf{\bar{u }}_{l}(\mathbf{\bar{u}}_{l})^{T}]=\mathbf{L}_{l}\otimes\mathbf{R}_{l}.\]

In this way, using LA, BNN posteriors can be obtained by approximating the Hessian of neural networks. We provide more details on LA and KFAC in Appendix A.1.

We note several ramifications of this formulation for learning BNN priors from data. First, the BNN posteriors can be obtained by approximating the Hessian, which can be tractably computed from large amounts of data and deeper neural network architectures (Lee et al., 2020; Ba et al., 2017). Second, the resulting BNN posteriors can capture the correlations between the parameters within the same layer (Ritter et al., 2018). Moreover, several open-source libraries exist to ease the implementations (Daxberger et al., 2021; Humt et al., 2020). All these points result in easily-obtainable BNN posteriors with expressive probabilistic representation from large amounts of data, deep architectures, and parameter correlations. As opposed to isotropic Gaussian, we next demonstrate that these BNN posteriors can be used to learn expressive prior distributions in order to advance generalization and uncertainty estimation within the transfer and continual learning set-ups.

### Empirical Prior Learning with Sums-of-Kronecker-Product Computations

Intuitively, the idea is to repeat the LA with the prior chosen as the posterior from Equation (1), similar to Bayesian filtering. Since we use the LA with a Kronecker-factored covariance matrix in both the prior and the posterior, we
want to approximate the resulting sum of Kronecker products with a single Kronecker product. We denote the task to compute the prior as source task \(\mathfrak{T}_{0}\) and the task to compute the posterior as target task \(\mathfrak{T}_{1}\), both consisting of a data-set \(\mathcal{D}_{t}\) from a distribution \(P_{t}\), \(t\in\{0,1\}\).

Applying LA on \(\mathcal{D}_{0}\), we obtain the model parameters \(\hat{\mathbf{\theta}}^{(0)}\) and the posteriors \(p(\mathbf{\theta}|\mathcal{D}_{0})\) as before (equations 1 and 2). Then, for the target task, we specify the prior using the posteriors from \(\mathcal{D}_{0}\): \(\pi(\mathbf{\theta})=p(\mathbf{\theta}|\mathcal{D}_{0})\). The ultimate goal is to compute the new posteriors on \(\mathcal{D}_{1}\), _i.e._, \(p(\mathbf{\theta}|\mathcal{D}_{1})\propto p(\mathcal{D}_{1}|f_{\mathbf{\theta}})\pi( \mathbf{\theta})\). To achieve this, we train the neural networks on \(\mathcal{D}_{1}\) by optimizing:

\[\begin{split}\hat{\mathbf{\theta}}^{(1)}&\in\operatorname* {arg\,max}_{\mathbf{\theta}}\{p(\mathbf{\theta}|\mathcal{D}_{1})\}\\ &=\operatorname*{arg\,min}_{\mathbf{\theta}}\{-\ln\left(p(\mathcal{D }_{1}|f_{\mathbf{\theta}})\right)\\ &\qquad\qquad+\frac{1}{2}(\mathbf{\theta}-\hat{\mathbf{\theta}}^{(0)})^{ T}(\mathbf{F}^{(0)}+\gamma\mathbf{I})(\mathbf{\theta}-\hat{\mathbf{\theta}}^{(0)})\},\end{split} \tag{3}\]

where \(\hat{\mathbf{\theta}}^{(1)}\) is the MAP estimate of the model parameters for task \(\mathfrak{T}_{1}\) and \(\mathbf{F}^{(0)}\) is the block-diagonal Kronecker-factored Fisher matrix from task \(\mathfrak{T}_{0}\). The final step is to approximate new Hessian on \(\mathcal{D}_{1}\) using the KFAC method. This results in

\[\text{Prior:}\quad\pi=\mathcal{N}(\hat{\mathbf{\theta}}^{(0)},( \mathbf{F}^{(0)}+\gamma\mathbf{I})^{-1}) \tag{4}\] \[\text{Posterior:}\quad\rho=\mathcal{N}(\hat{\mathbf{\theta}}^{(1)}, \tau(\mathbf{F}^{(1)}+\mathbf{F}^{(0)}+\gamma\mathbf{I})^{-1}),\]

where \(\tau\) is the temperature scaling (Wenzel et al., 2020; Daxberger et al., 2021). Here, the precision matrix of the new BNN posteriors is represented by the sum-of-Kronecker-products, _i.e._, \(\bar{\mathbf{F}}^{(1)}=\mathbf{F}^{(1)}+\mathbf{F}^{(0)}+\gamma\mathbf{I}= \mathbf{L}^{(0)}\otimes\mathbf{R}^{(0)}+\mathbf{L}^{(1)}\otimes\mathbf{R}^{(1) }+\gamma\mathbf{I}\otimes\mathbf{I}\).

```
1:Input: left matrices \((\mathbf{L}^{k})_{k\in[K]}\), right matrices \((\mathbf{R}^{k})_{k\in[K]}\), number of steps \(n^{max}=100\), and stopping precision \(\delta=10^{-5}\).
2:Output: The solutions \(\hat{\mathbf{L}}\) and \(\hat{\mathbf{R}}\).
3:\(\operatorname{vec}(\bar{\mathbf{L}}^{(0)})\leftarrow\mathcal{N}(\mathbf{0}, \mathbf{I})\) {initialization of \(\bar{\mathbf{L}}^{(0)}\)}
4:\(\mathbf{L}^{(0)}\leftarrow\frac{\bar{\mathbf{L}}^{(0)}}{\|\bar{\mathbf{L}}^{( 0)}\|_{F}}\) {normalize \(\bar{\mathbf{L}}^{(0)}\)}
5:for\(n=1\) to \(n^{max}\)do
6:\(\bar{\mathbf{R}}^{(n)}\leftarrow\sum_{k=1}^{K}(\mathbf{L}^{k},\mathbf{L}^{(n- 1)})_{F}\mathbf{R}^{k}\) {first power step}
7:\(\mathbf{R}^{(n)}\leftarrow\frac{\mathbf{R}^{(n)}}{\|\bar{\mathbf{R}}^{(n)}\|_ {F}}\) {normalize \(\bar{\mathbf{R}}^{(n)}\)}
8:\(\bar{\mathbf{L}}^{(n)}\leftarrow\sum_{k=1}^{K}(\mathbf{R}^{k},\mathbf{R}^{(n )})_{F}\mathbf{L}^{k}\) {second power step}
9:\(\mathbf{L}^{(n)}\leftarrow\frac{\bar{\mathbf{L}}^{(n)}}{\|\bar{\mathbf{L}}^{( n)}\|_{F}}\) {normalize \(\bar{\mathbf{L}}^{(n)}\)}
10:if\(\|\mathbf{L}^{(n)}-\mathbf{L}^{(n-1)}\|_{F}<\delta\)then
11:break {stopping criterion}
12:endif
13:endfor
14:\(\hat{\mathbf{L}}\leftarrow\mathbf{L}^{(n)}\)
15:\(\hat{\mathbf{R}}\leftarrow\sum_{k=1}^{K}(\mathbf{L}^{k},\mathbf{L}^{(n)})_{F} \mathbf{R}^{k}\) {first power step}
```

**Algorithm 1** Power method for sums of Kronecker products

Unfortunately, the sum-of-Kronecker-products is not a Kronecker-factored matrix. As a result, the above formulation loses all the essence of Kronecker factorization for modeling high-dimensional probability distributions. For example, we can no longer exploit the rule: \((\mathbf{L}\otimes\mathbf{R})^{-1}=(\mathbf{L}^{-1}\otimes\mathbf{R}^{-1})\) where \(\mathbf{L}\) and \(\mathbf{R}\) are smaller matrices to store and invert when compared to \((\mathbf{L}\otimes\mathbf{R})\). Even more, the storage and the inverse of \(\bar{\mathbf{F}}^{(1)}\) may not be computationally tractable for modern architectures.

To this end, we devise the sum-of-Kronecker-products computations. Concretely, we approximate the sum-of-Kronecker-products as Kronecker-factored matrices by an optimization formulation:12

Footnote 1: For the similar causes to maintain the Kronecker factorization, Ritter et al. (2018b;a) assume \((\mathbf{L}\otimes\mathbf{R}+\gamma\mathbf{I})^{-1}=(\mathbf{L}+\gamma\mathbf{ I})^{-1}\otimes(\mathbf{R}+\gamma\mathbf{I})^{-1}\) which does not hold in general (see Section 4.1).

Footnote 2: Our formulation for one single Kronecker-factored matrix is similar to Kao et al. (2021).

\[\hat{\mathbf{L}},\hat{\mathbf{R}}\in\operatorname*{arg\,min}_{ \begin{subarray}{c}\mathbf{L}\in\mathbb{R}^{M\times M}\\ \mathbf{R}\in\mathbb{R}^{N\times N}\end{subarray}}\|\sum_{k=1}^{K}\mathbf{L}^{k} \otimes\mathbf{R}^{k}-\mathbf{L}\otimes\mathbf{R}\|_{F}. \tag{5}\]

A solution to this problem is not unique, e.g., one could scale \(\hat{\mathbf{L}}\) by \(\alpha\neq 0\) and \(\hat{\mathbf{R}}\) by \(\frac{1}{\alpha}\). Hence, we assume that \(\hat{\mathbf{L}}\) is normalized, \(\|\hat{\mathbf{L}}\|_{F}=1\) (or alternatively, we can also assume \(\|\hat{\mathbf{R}}\|_{F}=1\)). For the solution, we show the equivalence to the well-known best rank-one approximation problem.

**Lemma 2.1**.: _Let \(M,N,K\in\mathbb{N}\), \(\mathbf{L}^{k}\in\mathbb{R}^{M\times M}\) and \(\mathbf{R}^{k}\in\mathbb{R}^{N\times N}\) for \(k\in[K]\). Then_

\[\|\sum_{k=1}^{K}\mathbf{L}^{k}\otimes\mathbf{R}^{k}-\mathbf{L} \otimes\mathbf{R}\|_{F} \tag{6}\] \[=\|\sum_{k=1}^{K}\operatorname{vec}(\mathbf{L}^{k})\operatorname{ vec}(\mathbf{R}^{k})^{T}-\operatorname{vec}(\mathbf{L})\operatorname{vec}(\mathbf{R})^{T}\|_{F}.\]

Proof.: The proof can be found in Appendix C.1. 

This result indicates that the solution to equation 5 can be obtained using the power method, which iterates:

\[\mathbf{R}^{(n)}\leftarrow\frac{\sum_{k=1}^{K}(\mathbf{L}^{k}, \mathbf{L}^{(n-1)})_{F}\mathbf{R}^{k}}{\|\sum_{k=1}^{K}(\mathbf{L}^{k}, \mathbf{L}^{(n-1)})_{F}\mathbf{R}^{k}\|_{F}}\quad\text{and}\] \[\mathbf{L}^{(n)}\leftarrow\frac{\sum_{k=1}^{K}(\mathbf{R}^{k}, \mathbf{R})_{F}\mathbf{L}^{k}}{\|\sum_{k=1}^{K}(\mathbf{R}^{k},\mathbf{R})_{F} \mathbf{L}^{k}\|_{F}}.\]

Given randomly initialized matrices, the power method iterates until the stop criterion is reached. The full procedure is presented in Algorithm 1, whereas in Appendix B, we discuss the computational complexity of the algorithm.
Importantly, we can now obtain a single Kronecker factorization from the sum of Kronecker products. Such Kronecker-factored approximations have advantages on tractable memory consumption for the storage and the inverse computations, which enables sampling from the resulting matrix normal distributions (Martens and Grosse, 2015). Therefore, such computations allow us to learn expressive BNN priors from the posterior distributions of previously encountered tasks. Finally, we can also prove the convergence of the power method to an optimal rank-one solution.

**Lemma 2.2**.: _Let \(\mathbf{A}=\sum_{k=1}^{K}\operatorname{vec}(\mathbf{L}^{k})\operatorname{vec}( \mathbf{R}^{k})^{T}\) and \(\mathbf{A}=\sum_{i=1}^{r}\sigma_{i}\mathbf{u}_{i}\mathbf{v}_{i}^{T}\) be its singular value decomposition with \(\sigma_{1}\geq\sigma_{2}\geq\cdots\geq\sigma_{r}>0\) and \(\mathbf{u}_{i}^{T}\mathbf{u}_{j}=\mathbf{v}_{i}^{T}\mathbf{v}_{j}=\mathbb{1}[i =j]\). Then there is a solution of Equation 5 with \(\operatorname{vec}(\tilde{\mathbf{L}})=\mathbf{u}_{1}\text{vec}(\tilde{ \mathbf{R}})=\sigma_{1}\mathbf{v}_{1}\). If \(\sigma_{1}>\sigma_{2}\), the solution is unique up to changing the sign of both factors, and our power method converges almost surely to this solution._

Proof.: The proof can be found in Appendix C.2. 

### Derivations and Optimizations over Tractable PAC-Bayes Bounds

So far, we have presented a method for learning a scalable and structured prior. In this section, we show how generalization bounds for the LA can be adapted to explicitly minimize the generalization bounds with the learned prior. In particular, the proposed approach allows tuning the hyper-parameters of the LA on the training set without a costly grid search. This is achieved by optimizing a differentiable cost function that approximates generalization bounds for the LA. We choose to optimize generalization bounds to trade off the broadness of the distribution and the performance on the training data to improve generalization.

A common method in BNNs, and especially in LA, is to scale the covariance matrix by a positive scalar called the temperature \(\tau>0\)(Wenzel et al., 2020; Daxberger et al., 2021). In addition, often either the Hessian (or Fisher matrix) of the likelihood (Ritter et al., 2018; Lee et al., 2020) or the prior (Gawlikowski et al., 2021) is scaled. Including all these terms, the posterior has the following form:

\[\rho=\mathcal{N}(\boldsymbol{\hat{\theta}}^{(1)},\tau(\beta\mathbf{F}^{(1)}+ \alpha\tilde{\mathbf{F}}^{(0)})^{-1}), \tag{7}\]

where \(\tilde{\mathbf{F}}^{(0)}=\mathbf{F}^{(0)}+\gamma\mathbf{I}\) is the precision matrix of the prior. This reweighting of the three scalars allows us to cope with misspecified priors (Wilson and Izmailov, 2020), approximation errors in the Fisher matrices, and to improve the broadness of the distribution. While these values are usually hyperparameters that have to be manually scaled by hand, we argue that these values should be application-dependent. Therefore, we propose to find all three scales - \(\alpha\), \(\beta\), and \(\tau\) - by minimizing generalization bounds.

In the following, we will first introduce our two approaches and then explain how the optimization of PAC-Bayes bounds can be made tractable for the LA.

**Method 1: Curvature Scaling**_Previous approaches typically scale only one or two of the scales (Daxberger et al., 2021; Ritter et al., 2018; Lee et al., 2020). With our automatic scaling, we can optimize all three scales for each individual layer, allowing the model to decide the weighting in a more fine-grained way. We can use this method not only to scale the posterior, but also to compute the prior. Thus, the prior and posterior are defined as_

\[\pi=\mathcal{N}(\boldsymbol{\hat{\theta}}^{(0)},(\tilde{\mathbf{F}}^{(0)})^{- 1}),\quad\rho=\mathcal{N}(\boldsymbol{\hat{\theta}}^{(1)},(\tilde{\mathbf{F}}^ {(1)})^{-1}), \tag{8}\]

_where the diagonal blocks of the precision matrix corresponding to each layer are defined as \(\tilde{\mathbf{F}}^{(0)}_{l}=\frac{1}{\tau_{l}^{(0)}}(\beta_{l}^{(0)}\mathbf{F }_{l}^{(0)}+\alpha_{l}^{(0)}\gamma\mathbf{I})\) and \(\tilde{\mathbf{F}}_{l}^{(1)}=\frac{1}{\tau_{l}^{(0)}}(\beta_{l}^{(1)}\mathbf{ F}_{l}^{(1)}+\alpha_{l}^{(1)}\tilde{\mathbf{F}}_{l}^{(0)})\), respectively._

_Remark 2.3_.: Although the temperature scaling could be captured within the other curvature scales, we find it easier to optimize the bounds with all three parameters per layer.

**Method 2: Frequentist Projection**_Going one step further, we propose to scale not only the curvature but also the network parameters using PAC-Bayes bounds. Since the Fisher matrix is known only after training, we assume the same covariance for the posterior as for the prior when optimizing the network parameters. Since this method does not optimize for the MAP parameters, but only for minimizing the PAC-Bayes bounds, we call it frequentist projection._

Both methods aim at improving the generalization of our model. To do this, we use PAC-Bayes bounds (Germain et al., 2016; Guedj, 2019) to derive a tractable objective that can be optimized without going through the data-set. The main idea of PAC-Bayes theory is to upper bound the expected loss on the true data distribution \(P^{N}\) with high probability. The bounds typically depend on the expected empirical loss on the training data and the KL-divergence between a data-independent prior and the posterior distribution. For \(\varepsilon>0\), that is

\[P_{\mathcal{D}\sim P^{N}}( \forall\rho\ll\pi:\ \mathbb{E}_{\boldsymbol{\theta}\sim\rho}[ \mathcal{L}_{P}^{l}(f_{\boldsymbol{\theta}})] \tag{9}\] \[\leq\delta(\mathbb{E}_{\boldsymbol{\theta}\sim\rho}[\hat{ \mathcal{L}}_{\mathcal{D}}^{l}(f_{\boldsymbol{\theta}})],\mathbb{KL}(\rho|\pi),N,\varepsilon))\geq 1-\varepsilon,\]

where the loss on the true data distribution is denoted as \(\mathcal{L}_{P}^{l}(f_{\boldsymbol{\theta}})=\mathbb{E}_{(\mathbf{x}, \mathbf{y})\sim P}[l(f_{\boldsymbol{\theta}}(\mathbf{x}),\mathbf{y})]\) and the empirical loss on the training data is \(\hat{\mathcal{L}}_{\mathcal{D}}^{l}(f_{\boldsymbol{\theta}})=\frac{1}{N}\sum_{ i=1}^{N}l(f_{\boldsymbol{\theta}}(\mathbf{x}),\mathbf{y})\). In general, the bounds balance the fit to the available data (empirical risk term) and the similarity of the posterior to the prior (the KL-divergence term), and the bounds hold with a probability greater than \(1-\varepsilon\). Various forms of the \(\delta\) bound can be found in the literature, with varying degrees of tightness.
For classification, we rely on the McAllester (McAllester, 1999b) and Catoni (Catoni, 2007) bounds. 3

Footnote 3: This work focuses on classification tasks as PAC-Bayes framework is usually for bounded losses. Yet, using recent theories of PAC-Bayes, we also comment on regression in Appendix D.

Since these bounds depend on the expected empirical loss, we have to iterate over the entire training data and sample from the posterior multiple times at each step in order to evaluate the bound once. This makes optimization intractable. Using the error function as a loss, we propose to compute an approximate upper bound by only using quantities that were already computed during the LA, _i.e_. the Fisher matrix \(\operatorname{diag}(\{\mathbf{F}_{l}\}_{l\in[L]})\) and the network parameters \(\tilde{\mathbf{\theta}}\):

\[\mathbb{E}_{\mathbf{\theta}\sim\rho}\left[\frac{1}{N}\sum_{i=1}^{N} \mathbbm{1}[\operatorname*{arg\,max}_{\mathbf{y}^{\prime}}p(\mathbf{y}^{\prime}| \mathbf{x}_{i},f_{\mathbf{\theta}})\neq\mathbf{y}_{i}]\right]\] \[\leq\mathbb{E}_{\mathbf{\theta}\sim\rho}\left[\frac{1}{N}\sum_{i=1}^ {N}-\frac{\ln p(\mathbf{y}_{i}|\mathbf{x}_{i},f_{\mathbf{\theta}})}{\ln 2}\right] \tag{10}\] \[\approx\frac{-\ln p(\mathcal{D}|f_{\tilde{\mathbf{\theta}}})+\frac{1 }{2}\sum_{l\in[L]}\tau_{l}\operatorname{tr}\Big{(}\mathbf{F}_{l}(\beta_{l} \mathbf{F}_{l}+\alpha_{l}\tilde{\mathbf{F}}_{l})^{-1}\Big{)}}{N\ln 2},\]

where \(\tilde{\mathbf{F}}_{l}\) is from the precision matrix of the prior. Here, we first use a second-order Taylor approximation of around \(\tilde{\mathbf{\theta}}\). Furthermore, the expectation can be converted to a trace by using the cyclic property of the trace together with the fact that the posterior is a multivariate Gaussian. The negative data log-likelihood of the optimal parameters can be computed jointly with the Fisher matrix during the LA. We denote this approximation as \(\operatorname{aer}(\alpha,\beta,\tau)\). On the other hand, the KL-divergence can also be computed in closed form for our prior-posterior pair, since both distributions are multivariate Gaussians. Thus, we can plug both terms into the McAllester bound (Guedj, 2019) to obtain the objective

\[ma(\alpha,\beta,\tau)=\operatorname{aer}(\alpha,\beta,\tau)+ \sqrt{\frac{\operatorname{kl}(\alpha,\beta,\tau)+\ln\frac{2\sqrt{N}}{\varepsilon }}{2N}} \tag{11}\]

where we write \(\operatorname{kl}(\alpha,\beta,\tau)=\mathbb{KL}(\rho\|\pi)\) for the KL-divergence to emphasize the dependence on the scales. Similarly for the Catoni bound (Catoni, 2007), we obtain

\[ca(\alpha,\beta,\tau)=\inf_{c>0}\frac{1-\exp(-c\operatorname{aer}(\alpha, \beta,\tau)-\frac{\operatorname{kl}(\alpha,\beta,\tau)-\ln\varepsilon}{N})}{1 -\exp(-c)}. \tag{12}\]

Overall, these objectives can be evaluated and minimized without using any data samples. As opposed to the cross-validated marginal likelihood or other forms of grid searching, we (a) obtain generalization bounds and analysis, (b) do not need a separate validation set, (c) can find multiple hyperparameters, i.e., scale better with the dimensionality of the considered hyperparameters due to the differentiability of the proposed objectives, and (d) the complexity of the optimization is independent of the data set size and the complexity of the forward pass. The last point is because we only use the precomputed terms from the LA. A full derivation of our objective is given in Appendix C.

### Bayesian Progressive Neural Networks

Having the essentials of learning BNN priors with generalization guarantees, we now present our extension to continual learning, which shows the versatility of our prior learning method. Here, we use so-called Progressive Neural Networks (PNNs) (Rusu et al., 2016), where a new neural network (column) is added for each new incoming task, and features from previous columns are also added for positive transfer between each task. Thus, PNNs are immune to catastrophic forgetting at the cost of an increase in memory, while being also applicable in transfer learning more generally (Wang et al., 2017; Rusu et al., 2017). As such, we extend the set-up in Section 2.2 by sequentially considering T tasks: \(\mathfrak{T}_{1},\ldots,\mathfrak{T}_{T}\). Moreover, to keep the generality of our method, an additional task \(\mathfrak{T}_{0}\) is defined for the priors.

The idea behind our Bayesian re-interpretation of PNNs, dubbed as BPNNs for Bayesian PNNs, is as follows. First, the BNN posteriors are learned at task \(\mathfrak{T}_{0}\) (depicted in Figure 1). Then, for an incoming sequence of tasks, the BNN priors are specified from the BNN posteriors from \(\mathfrak{T}_{0}\). The proposed methods of the sums-of-Kronecker-products and PAC-Bayes bounds are used. For the lateral connections, the BNN posterior from which the lateral connection originates is used. This ensures that the weight distribution from the prior already produces reasonable features given the activations from the previous column. Altogether, the resulting architecture accounts for model uncertainty, generalization,

Figure 1: Illustration of our continual learning architecture. BPNNs use the empirically learned prior from task \(\mathfrak{T}_{0}\) (blue) in all columns. The prior for the lateral connections is the posterior from the lateral connection (orange and green). Best viewed in color.
and resiliency to catastrophic forgetting. All these properties are desirable to have within one unified framework. We note that BPNN is in line with our PAC-Bayes theory, i.e., we increase the complexity without increasing the PAC-Bayes bounds. This is because we can reuse features from previous columns even though they do not contribute to the bounds due to their a-priori fixed weight distribution. In Appendices B and C, we provide more details such as its full derivations, and its training and testing procedures.

## 3 Related Work

There are different work streams related to this paper. The primary area is on BNN priors while we also contribute to Bayesian continual learning and PAC-Bayes theory.

**Bayesian Neural Networks Priors** A prior specification is a prerequisite in Bayesian modeling. However, for neural network parameters, the choice of the prior is generally unknown. So far, uninformative priors such as isotropic Gaussian have been the de-facto-standard for BNNs (Fortuin, 2022). Such uninformative priors may cause undesirable effects such as cold posteriors and worse predictive performance than standard neural networks (Wenzel et al., 2020; Fortuin et al., 2021). To this end, recent research efforts have focused on exploiting function-space (Sun et al., 2019), sparsity-inducing weight-space (Carvalho et al., 2009; Ghosh et al., 2018), structured (Louizos and Welling Max, 2016), and learning-based priors (Immer et al., 2021; Fortuin et al., 2021; Wu et al., 2019). Amongst these, we build upon learning-based priors. Learning-based priors can be an alternative to uninformative prior when no useful prior knowledge is available to encode, or when there exists relevant data and tasks to learn from (Fortuin, 2022).

The idea of learning a prior distribution on a similar task is certainly not new. In this domain, Bayesian meta-learning (Thrun and Pratt, 1998) and their modern extensions (Rothfuss et al., 2021; Finn et al., 2018) can be viewed as another form of learning the prior from data, although their focus is typically not on advancing BNNs. Empirical prior learning is closely related to our work (Robbins, 1992). For BNNs, Fortuin et al. (2021) learns the empirical weight distributions during stochastic gradient descent. Wu et al. (2019) used a moment-matching approach, while Immer et al. (2021) exploited the so-called Laplace-Generalized-Gauss-Newton method. In (Krishnan et al., 2020), the mean of a Gaussian prior distribution is learned from a relevant data-set. The concurrent approaches of Shwartz-Ziv et al. (2022); Tran et al. (2022) share a similar spirit of learning expressive priors from large-scale data-set and architectures. The former utilizes so-called the SWAG (Maddox et al., 2019) framework with an inspiring idea of combining self-supervised learning, while the latter

Figure 2: The results of the ablation studies. The results show that the combination of our approaches can lead to non-vacuous bounds (a) and that the learned prior is more data-efficient (b) and needs less temperature scaling (c) than isotropic priors.
builds upon Frequentist batch ensembles (Wen et al., 2019). All these works show the strong relevance of learning-based priors for the current BNNs. Inspired by the aforementioned works, our main idea is to learn scalable and structured posterior as expressive informative priors, like Bayesian foundation models from broader data and models.

**Bayesian Continual Learning** Continual learning (Thrun Mitchell, 1995) methods can be broadly divided into replay-based, regularization-based, and parameter-isolation methods (De Lange et al., 2022). We build on parameter-isolation methods, e.g. Rusu et al. (2016), where different model parameters are dedicated to new tasks. Within this branch, Bayesian methods have been investigated (Ebrahimi et al., 2020; Ardywibowo et al., 2022; Kumar et al., 2021), which could benefit from our work on advancing the BNN priors. Rudner et al. (2022) shares a similar spirit of bringing the state-of-the-art BNN priors to Bayesian continual learning by relying on the function space priors (Sun et al., 2019). On the other hand, using learning-based expressive priors, we design a single unified framework of continual learning for generalization, uncertainty-awareness, and resiliency to catastrophic forgetting.

**Generalization Theory** In supervised learning, we can obtain models with a bound on the true expected loss from its true data-generating process. Typical early works involved Vapnik-Chervonenkis theory and Rademacher complexity (Vapnik and Chervonenkis, 1968; Shalev-Shwartz and Ben-David, 2014). However, for high-dimensional models like neural networks, these methods often provide vacuous bounds. In recent years, PAC-Bayes theory (McAllester, 1999; Germain et al., 2016) has become an alternative method with wide applications. The seminar paper of (Germain et al., 2016) showed the connection to approximate Bayesian inference. Rothfuss et al. (2021b, a) devise compelling meta-learning priors for BNNs with generalization guarantees. These works form our inspiration to explore PAC-Bayes theory for learning-based BNN priors. We note that our goal is not to advance PAC-Bayes theory, but to investigate a method for scaling the BNN priors with an approximate differentiable objective for generalization.

## 4 Results

The goal of the experiments is to investigate whether our approach provides generalization and calibrated uncertainty estimates. To this end, in addition to ablation studies on the presented algorithm, we show its utility for continual learning and uncertainty estimation. Implementation details are presented in Appendix E. The code is released at [https://github.com/DLR-RM/BPNN](https://github.com/DLR-RM/BPNN).

### Ablation Studies

Our method is to improve generalization in the LA and also to provide an informative prior for the posterior inference. Therefore, in the ablation studies, we want to investigate the impact of each of our methods on the generalization bounds. Moreover, we want to study the learned prior in the small-data regime and see if it can mitigate the cold posterior effect, _i.e_. phenomena in BNNs where the performance improves by cooling the posterior with a temperature of less than one (Wenzel et al., 2020). This effect highlights the discrepancy of current BNNs, where no posterior tempering should be theoretically needed. For this, we use a LeNet-5 architecture (LeCun et al., 1989), learn the prior on MNIST (LeCun et al., 1998) and compute the posterior on NotMNIST (Bulatov, 2011). In our experiments, we compare the performance of our learned prior against an isotropic Gaussian prior with multiple weight decays, either with a mean zero or using the pre-trained model as the mean.

Figure 3: The results of the ablation studies. The results show the accuracy of the proposed sums-of-Kronecker product computations.
In Table 2a, one can see the contribution of each method to the final generalization bounds. For this, we report the mean and standard deviation using \(5\) different seeds. We observe that learning the prior improves the generalization bounds, as the posterior can reuse some of the pre-trained weights. The curvature further controls the flexibility of each parameter, leading to a smaller KL-divergence while still providing a small error rate. In addition, we show that scaling the curvature with our approximate bounds reduces the bounds. In particular, the generalization bounds are also better than a thorough grid search. Frequentist projection further leads to a non-vacuous bound of \(0.885\). In the small data regime, we train our method on a random subset of NotMNIST, _i.e_., we vary the number of available samples per class. Figure 2b shows that the learned prior requires a magnitude fewer data to achieve similar accuracy compared to isotropic priors. In Appendix F.2, we further evaluate the impact of using different temperature scalings and weight decays in this experiment. We observe larger improvements when the model is more probabilistic, _i.e_. when the temperature scaling is large, and when the weight decay is small and thus the learned prior is more dominant. Finally, following Wenzel et al. (2020), we examine the cold posterior effect using the learned prior in Figure 2c. Although the optimal value for the temperature scaling is less than one, the temperature scaling can be closer to one compared to the isotropic priors. This suggests that the cold posterior effect is reduced by using a learned prior.

Additionally, we further validate the proposed sums-of-Kronecker product computations (see Figure 3). In Appendix F, further experimental results are provided, including qualitative analysis of our tractable approximate bound used to optimize the curvature scales (Appendix F.1). Furthermore, results for a larger cold posterior experiment using ResNet-50 (He et al., 2016), ImageNet-1K (Deng et al., 2009), and CIFAR-10 (Krizhevsky, 2009) are presented in Appendix F.3. Here, the learned prior improves the accuracy but does not reduce the cold posterior effect for all evaluated weight decays. This suggests the use-case of our method. Overall, our results demonstrate the effectiveness of our method in improving the generalization bounds. Furthermore, we show that the learned prior is particularly useful for BNNs when little data is available.

### Generalization in Bayesian Continual Learning

Another advantage of our prior is its connection to continual learning. In the following experiments, we, therefore, evaluate our method for continual learning tasks. To do so, we closely follow the setup of Denninger and Triebel (2018), who introduces a practical robotics scenario. This allows us to also obtain meaningful results for practitioners. We do not use the typical MNIST setup here because we also want to test the effectiveness of the expressive prior from ImageNet (Deng et al., 2009). Each continual learning task consists of recognizing the specific object instances of the Washington RGB-D data-set (WRGBD) (Lai et al., 2011), while only a subset of all classes is available at a given time. We neglect the depth information and split the data similar to Denninger and Triebel (2018). The prior is learned with the ImageNet-1K (Deng et al., 2009) and the ResNet-50 (He et al., 2016) is used. We specify a total of four lateral connections, each at the downsampling \(1\times 1\) convolution of the first "bottleneck" building block of the ResNet layers. The baselines are PNNs with several weight decays and using MC dropout (Gal and Ghahramani, 2016). Moreover, we compare to both isotropic priors as in the ablations.

Table 1 shows the mean and the standard deviation of the accuracy for each task of the continual learning experiment for \(5\) different random seeds. We report the mean and the standard deviation of two independent runs. Overall, in our experiments, the accuracy averaged over all tasks is improved by \(0.6\) percent points compared to the second best approach PNN with weight decay \(10^{-5}\). In addition, the increase in accuracy is \(1.3\) percent points greater compared to an isotropic prior, and \(2.1\) percent point greater when using zero as the prior mean. Therefore, these experimental results illustrate that with our idea of expressive posterior as BNN prior, we can improve the generalization on the test set for practical tasks of robotic continual learning.

### Few-Shot Generalization and Uncertainty

Finally, we consider the task of few-shot learning. Our key hypothesis is that the choice of prior matters more in the small data regime, _i.e_., the prior may dominate over the likelihood term. To this end, closely following Tran et al. (2022), we use a few-shot learning set-up across several data-sets,

\begin{table}
\begin{tabular}{l c c c c c c c} \hline \hline  & Other & Bananaa & Coffee Mug & Stapler & Flashlight & Apple & Average \\ \hline PNN (weight decay \(10^{-3}\)) & \(95.7\pm 0.4\) & \(98.5\pm 2.3\) & \(\mathbf{100.0\pm 0.0}\) & \(91.7\pm 1.0\) & \(93.8\pm 9.1\) & \(93.3\pm 4.7\) & \(95.5\pm 2.9\) \\ PNN (weight decay \(10^{-5}\)) & \(\mathbf{96.7\pm 0.3}\) & \(99.0\pm 1.2\) & \(\mathbf{100.0\pm 0.0}\) & \(90.7\pm 2.3\) & \(99.7\pm 0.4\) & \(94.1\pm 3.2\) & \(96.7\pm 1.2\) \\ MC Dropout & \(96.3\pm 0.1\) & \(\mathbf{90.3\pm 0.9}\) & \(\mathbf{100.0\pm 0.0}\) & \(92.7\pm 0.8\) & \(99.8\pm 0.4\) & \(90.4\pm 5.1\) & \(96.4\pm 1.2\) \\ zero mean \& isotropic & \(96.3\pm 0.3\) & \(95.5\pm 4.2\) & \(\mathbf{100.0\pm 0.1}\) & \(91.9\pm 1.7\) & \(\mathbf{100.0\pm 0.1}\) & \(87.7\pm 7.8\) & \(95.2\pm 2.4\) \\ isotropic & \(96.1\pm 0.3\) & \(98.7\pm 1.0\) & \(\mathbf{100.0\pm 0.0}\) & \(93.1\pm 0.6\) & \(\mathbf{100.0\pm 0.0}\) & \(87.8\pm 6.6\) & \(96.0\pm 1.4\) \\ learned & \(96.2\pm 0.2\) & \(98.4\pm 1.6\) & \(\mathbf{100.0\pm 0.0}\) & \(\mathbf{93.9\pm 0.9}\) & \(\mathbf{100.0\pm 0.0}\) & \(\mathbf{95.1\pm 4.7}\) & \(\mathbf{97.3\pm 1.2}\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Continual learning experiment. BPNN with a learned prior (our method) improves the averaged accuracy over all tasks compared to using an isotropic prior around zero or around a learned mean. Our method also improves over PNN with and without MC Dropout. Following Denninger and Triebel (2018), each column represents incoming continual learning tasks, where we receive new objects.
namely CIFAR-100, UC Merced, Caltech 101, Oxford-IIIT Pets, DTD, Colorectal Histology, Caltech-UCSD Birds 200, and Cars196. CIFAR-10 is used as an out-of-distribution (OOD) data-set 4. Standardized metrics such as accuracy, ECE, AUROC, and OOD AUROC are examined. The implementations are based on uncertainty baselines (Nado et al., 2021). We use ResNet-50 for the architecture.

Footnote 4: Unlike Tran et al. (2022), we did not use few-shot learning on ImageNet since we obtain our prior using the entire training set.

For the baselines, we choose MC dropout (Gal and Ghahramani, 2016) (MCD) and additionally include deep ensemble (Lakshminarayanan et al., 2017) (Ens) and a standard network (NN). These uncertainty estimation techniques are often used methods in practice (Gustafsson et al., 2020). We do not use the plex (Tran et al., 2022) due to the lack of industry-level resources to comparably train our Bayesian prior. To increase their competitiveness, we uniformly searched over ten weight decays in the range from \(10^{-1}\) to \(10^{-10}\). We also tried fixed representation or only last-layer fine-tuning. Additionally, we add LA which represents the use of isotropic prior. To facilitate the fair comparison between isotropic and learned prior, we carefully searched the following combinations: (a) curvature scaling without PAC-Bayes, with McAllester and Cantoni, (b) ten temperature scaling ranging from \(10^{-10}\) to \(10^{-28}\), and (c) three weight decays \(10^{-6}\) to \(10^{-8}\) and \(10^{-10}\). For all hyperparameters, we selected the best model using validation accuracy.

The results are reported in Figure 4, where we averaged across eight data-sets, closely following Tran et al. (2022). The results per data-set are in Appendix F.4 whereas in Appendix F.1, we also analyze the influence of these weight decays and temperature scales with varying numbers of available samples per class against accuracy. For shots up to 25 per class, we observe that our method often outperforms the baselines in terms of generalization and uncertainty calibration metrics. In particular, in the experiments, our learned prior from ImageNet significantly outperforms the isotropic prior within directly comparable set-ups. We interpret that the prior learning method is effective in this small data regime. This motivates the key idea of expressive posteriors as informative priors for BNNs. Moreover, for isotropic and learned prior, the McAllester bound often resulted in the best model. This also motivates the use of explicit curvature scaling for the generalization bounds.

## 5 Conclusion

This paper presents a prior learning method for BNNs. Our prior can be learned from large-scale data-set and architecture, resulting in expressive probabilistic representations with generalization guarantees. Empirically, we demonstrated how we mitigate cold posterior effects and further obtain non-vacuous generalization bound as low as 0.885 in neural networks using LA. In our benchmark experiments within continual and few-shot learning tasks, we further showed advancements over the prior arts.

Importantly, we find that the use-case of our prior learning method is more within the small data regime, e.g., when prior may dominate over the likelihood term. Finally, one of the fundamental assumptions of prior learning methods is on the existence of relevant data and tasks to learn the prior from. Moreover, for a valid PAC-Bayes analysis, the data-sets for individual tasks should not share common examples. In the future, we would like to see follow-ups that address this limitation, by either (a) learning prior in larger scale data-set and architectures like foundational models, or (b) combining self-supervised pre-training to quickly fine-tune the prior for the domain of the relevant task (Bommasani et al., 2021). Another direction is to obtain tighter generalization bounds by adapting clever tricks such as optimizing the entire covariance instead of individual scales with our objectives and using a subset of the training data to improve the prior (Perez-Ortiz et al., 2021).

**Acknowledgments** The authors would like to thank the anonymous reviewers and area chairs for valuable feedback. This work is also supported by the Helmholtz Association's Initiative and Networking Fund (INF) under the Helmholtz AI platform grant agreement (ID ZT-I-PF-5-1).

Figure 4: Few-shot learning experiments. Results are averaged over eight data-sets. Higher is better for accuracy and AUC measures, while lower is better for ECE measures. The results show the benefits of our method in uncertainty calibration and generalization.