Learning Expressive Priors for Generalization and
Uncertainty Estimation in Neural Networks

Consider a neural network $f_{\bm{\theta}}$ with layers $l\in[L]:=\{1,\dots,L\}$ which consists of a parameterized function $\phi_{l}:\Omega_{l-1}\to\Omega_{l}$ and a non-linear activation function $\sigma_{l}:\Omega_{l}\to\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 $\bm{\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_{\bm{\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 $\bm{\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(\bm{\theta})$. Then the Bayes theorem is used to compute the posterior distribution over the parameters, given the training data: $p(\bm{\theta}|\mathcal{D})\propto\pi(\bm{\theta})\prod_{i=1}^{N}p(\mathbf{y}_{i}|\mathbf{x}_{i},f_{\bm{\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_{\bm{\theta}})\rho(\bm{\theta})d\bm{\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 $\bm{\hat{\theta}}$ are obtained by training the neural networks. Then, the Hessian $\mathbf{H}=\frac{d^{2}}{{d\bm{\theta}}^{2}}\ln{p(\bm{\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:

As the Hessian is computationally intractable for modern architectures, the Kronecker-Factored Approximate Curvature (KFAC) method is widely adopted (Martens & 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{a}}_{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., 2018b):

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 Section 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., 2018b). 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.

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 $\bm{\hat{\theta}}^{(0)}$ and the posteriors $p(\bm{\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(\bm{\theta})=p(\bm{\theta}|\mathcal{D}_{0})$. The ultimate goal is to compute the new posteriors on $\mathcal{D}_{1}$, i.e., $p(\bm{\theta}|\mathcal{D}_{1})\propto p(\mathcal{D}_{1}|f_{\bm{\theta}})\pi(\bm{\theta})$. To achieve this, we train the neural networks on $\mathcal{D}_{1}$ by optimizing:

where $\bm{\hat{\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

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., $\mathbf{\tilde{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}$.

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 $\mathbf{\tilde{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:¹¹1For 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).²²2Our formulation for one single Kronecker-factored matrix is similar to Kao et al. (2021).

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

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

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

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 hyperparameters 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., 2018b; Lee et al., 2020) or the prior (Gawlikowski et al., 2021) is scaled. Including all these terms, the posterior has the following form:

where $\mathbf{\tilde{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 & 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., 2018b; 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

where the diagonal blocks of the precision matrix corresponding to each layer are defined as $\mathbf{\tilde{F}}^{(0)}_{l}=\frac{1}{\mathbf{\tau}^{(0)}_{l}}(\mathbf{\beta}^{(0)}_{l}\mathbf{F}^{(0)}_{l}+\mathbf{\alpha}^{(0)}_{l}\gamma\mathbf{I})$ and $\mathbf{\tilde{F}}^{(1)}_{l}=\frac{1}{\mathbf{\tau}^{(0)}_{l}}(\mathbf{\beta}^{(1)}_{l}\mathbf{F}^{(1)}_{l}+\mathbf{\alpha}^{(1)}_{l}\mathbf{\tilde{F}}^{(0)}_{l})$, respectively.

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

where the loss on the true data distribution is denoted as $\mathcal{L}^{l}_{P}(f_{\bm{\theta}})=\mathbb{E}_{(\mathbf{x},\mathbf{y})\sim P}[l(f_{\bm{\theta}}(\mathbf{x}),\mathbf{y})]$ and the empirical loss on the training data is $\hat{\mathcal{L}}^{l}_{\mathcal{D}}(f_{\bm{\theta}})=\frac{1}{N}\sum_{i=1}^{N}l(f_{\bm{\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. ³³3This 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 $\bm{\hat{\theta}}$:

where $\mathbf{\tilde{F}}_{l}$ is from the precision matrix of the prior. Here, we first use a second-order Taylor approximation of around $\bm{\hat{\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}(\mathbf{\alpha},\mathbf{\beta},\mathbf{\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

where we write $\operatorname*{kl}(\mathbf{\alpha},\mathbf{\beta},\mathbf{\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

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.

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},\dots,\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, 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.

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.

In Table 2(a), 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 2(b) 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 2(c). 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.

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

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, 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 ⁴⁴4Unlike Tran et al. (2022), we did not use few-shot learning on ImageNet since we obtain our prior using the entire training set.. 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.

For the baselines, we choose MC dropout (Gal & 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.