Pre-Train Your Loss: Easy Bayesian Transfer Learning with Informative Priors
Authors: Ravid Shwartz-Ziv
Micah Goldblum
Hossein Souri
Sanyam Kapoor
Chen Zhu
Yann LeCun
Andrew Gordon Wilson
Venue: NeurIPS 2022 Accept (2022)
Keywords: transfer learning
foundation models
priors
Bayesian neural networks

Abstract: Deep learning is increasingly moving towards a transfer learning paradigm whereby large foundation models are fine-tuned on downstream tasks, starting from an initialization learned on the source task. But an initialization contains relatively little information about the source task, and does not reflect the belief that our knowledge of the source task should affect the locations and shape of optima on the downstream task.
Instead, we show that we can learn highly informative posteriors from the source task, through supervised or self-supervised approaches, which then serve as the basis for priors that modify the whole loss surface on the downstream task. This simple modular approach enables significant performance gains and more data-efficient learning on a variety of downstream classification and segmentation tasks, serving as a drop-in replacement for standard pre-training strategies. These highly informative priors also can be saved for future use, similar to pre-trained weights, and stand in contrast to the zero-mean isotropic uninformative priors that are typically used in Bayesian deep learning. 

Pre-Train Your Loss: Easy Bayesian Transfer
Learning with Informative Priors
Ravid Shwartz-Ziy"
Micah Goldblum" New York University goldblum@nyu edu
New York University
ravid. shwartz. zwz.z..vviii zivenyu..dduuuuuu
Hossein Souri
Sanyam Kapoor New York University
Chen Zhu
Johns Hopkins University
University of Maryland
Yann LeCun New York University
Andrew Gordon Wilson
New York University
Meta AI Research
andreugw@cims nyu. edu
Abstract
Deep learning increasingly moving towards transfer learning paradigm whereby large foundation models are fine-tuned on downstream tasks, starting from an initialization learned on the source task. But an initialization contains relatively little information about the source task, and does not reflect the belief that our knowledge of the source task should affect the locations and shape optima on the downstream task. Instead, we show that we can learn highly informative posteriors from the source task, through supervised or eeff-superviee approaches, which then as the basis for priors that modify the whole loss surface on the downstream task. This simple modular approach enables significant performance gains and more data-efficient learning tasks, serving as drop-in replacement for standard pre-training strategies. These highly informative priors also can be saved for future use, similar to pre-trained weights, and stand contrast to the zero-mean isotropic uninformative priors that are typically used in Bayesian deep learning.
variety of downstream classification and segmentation
Introduction
The ability transfer what is learned from task to another - learning to ride bicycle to then ride a unicycle - has historically set apart biological intelligence from machine learning approaches. However, transfer learning is quickly becoming mainstream practice in deep learning. Typically, large "foundation models" are pre-trained on massive volumes of source data, and then the learned parameter vector is used has had great empirical success, reliance on an initialization is a very limited way to perform transfer learning. If we doing good job of optimization, then final solution should be independent of initialization, barring local minima with identical training loss. Moreover, our knowledge of the source task should affect the locations and shapes of optima the downstream knowledge.
initialization for training in a downstream task. While this approach
We propose to instead use re-scaled Bayesian parameter posterior from the first task as a pre-trained prior for the downstream task. Since the negative log posterior on the downstream task is our loss function, this procedure has the effect of reshaping our training objective on the downstream task,
Authors contributed equally.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Test Loss Train Los
Less Train Loss
Loss
Loss
Random Initialization
Initializstion
Initializztion
0
Gap
(a) Standard Transfer Learning
(b) Transfer I. Learning with Learned Priors
Figure 1: Learned priors reshape the loss surface. Red curves represent training loss, and blue curves represent the true population loss. The typical transfer learning approach avoids poor- generalizing minima by improving initializations (left). Learned priors further improve the alignment between training and test loss, improving accuracy (right). The gap between learned and optimal parameters denoted by a double-sided gold arrow.
to reflect more nuanced information that we have learned from the source task, as in Figure [ The posterior on the source task re-scaled before use as prior on the downstream task to reflect the belief that the source and downstream tasks are drawn from different distributions.
With highly informative prior for the downstream task, or perform full Bayesian model averaging with the posterior on the downstream task. We find both procedures profoundly improve upon standard deep transfer learning, making use of simple pipeline that involves easy-to-use pre-existing components. Indeed, the simplicity of the approach, combined with its promising empirical performance, one of its greatest features enabling its use drop-in replacement for standard approaches for deep transfer learning. The pre-tra pre-training posterior can be saved as prior for wide array of downstream tasks, similar pre-trained weights.
then proceed with optimization,
Despite the simplicity and effectiveness this approach, Bayesian neural networks are almost always used with uninformative isotropic zero-mean priors [51] 25 13), and have not harnessed recent advances in self-supervised learning. Furthermore, as we will show, the details of how we perform Bayesian transfer learning with modern neural networks are crucial practical success. For example, it is important that the prior on the source task be represented with an informative covariance matrix, indicating which directions in parameter space task, which significantly outperforms isotropic or diagonal counterparts, We also provide several key conceptual findings for example, that the informative priors provide an inductive bias that enables more data efficient fine-tuning than a pre-trained initialization. and that pre- pre-trained priors based on self-supervised learning are more transferable than supervised pre-trained priors. Finally, we observe through extensive experiments across both image classification and semantic segmentation on multiple neural architectures, and both supervised and supe-suuer (SimCLR (6) pre-training loss functions, that Bayesian inference is particularly advantageous in the transfer learning setting.
expect provide good solutions the source
We emphasize that the proposed approach outperforms standard transfer learning, without requiring any expert intervention or knowledge of Bayesian methods. One only needs to tune a single additional hyperparameter using standard cross-valices validation,, which incurs minimal computational overhead. We release PyTorch pre-trained priors and code for learning and using priors for downstream inference: http:///gtta...cccoooooo////ttat
Background
Bayesian Neural Networks. Consider a neural network f with weights w and training dataset D [(=((0, () of independently drawn samples. standard non-Bayesian neural network training, we minimize the negative log posterior, - p( f) x - log,,ppp)) (D/w, f) log p(w) )) where the likelihood p(D|w, denotes the probability that a model with weights w would generate the observed labels {)) (e.g., cross-entropy). The log prior log pwwl() often takes the form of weight decay, corresponding to zero-mean isotropic Gaussian prior. In Bayesian modeling,

---

we instead make predictions with a Bayesian Model Average (BMA) of all models weighted by their posterior probabilities: p(y| D. f) = p(y|z, w, p)p[(]) f f)dw. There are many ways to approximate this integral, such as MCMC, or variational approaches, which take a finite number of approximate samples w from the parameter posterior to form the simple Monte Carlo estimate: p(y|.; D. = J E, plylæ, w,, f). Almost always, one uses zero-mean isotropic priors [ST] There are also specialized priors, including heavy-tailed priors [ 25 noise-contrastive priors for high uncertainty under distribution shift [17], and input-dependent priors for domain generalization [24).
Transfer Learning. In transfer learning, we wish recycle the representation learned one task to improve performance on another. Transfer learning is now widely applied in deep learning, and which are exceptionally large neural networks pre-trained on massive volumes of data, and then fine-tuned on a downstream task. Recent work has found sule-sup pre-training can transfer better than supervised pre-training [19], in line with our finding in Section 4 that self-supervised pre-trained priors transfer better.
forms the basis for foundation models ①1 12
In continual learning. wish to learn tasks sequentially without forgetting what has been learned previously. Elastic Weight Consolidation (EWC) prevents forgetting by imposing a penalty, the diagonal of Fisher information matrix computed on previous tasks, for adapting We can contextualize the EWC penalty within Bayesian transfer learning as a negative log Gaussian prior with diagonal covariance used for maximum posteriori (MAP) estimation sequential tasks, such as digit classification or RL [29]. In our experiments, we find that MAP estimation, diagonal covariance, and supervised pre- rre-training are all suboptimal in the transfer learning setting.
new task [29].
Leveraging Auxiliary Data and Knowledge Transfer in Bayesian Modeling. A number of works outside of deep learning have considered knowledge transfer in Bayesian modeling, especially in settings such as domain adaptation or homogeneous transfer learning in which source and target tasks contain similar feature and label spaces but differ in their marginal distributions. For example, Xuan et al. [52) considers Bayesian transfer learning methods for probabilistic graphical models, and Karbalayghareh et al. [28] develop theoretical framework for understanding optimal Bayes classifiers given prior knowledge from other domains. Other works on Bayesian transfer learning learn Dirichlet prior over naive Bayes classifiers [44. Bayesian optimization methods also find transferable hyperparameter settings [40] or select data which will yield transferable models on shifted domains 45 Schnaus [47 uses the Laplace approximation learn posteriors with Kronecker- factored covariance and then adjusts the posteriors by optimizing PAC-Bayes generalization bounds create priors for downstream applications.
Bayesian tools have additionally been used for leveraging auxiliary data or multiple domains in deep learning. Chandra and Kapoor [2] learn from multiple domains simultancously using round-robin task sampling procedure and single-layer neural network. Bayesian methods for continual learning update the posterior to accommodate new tasks without forgetting how perform previous ones 335
46 or develop kernels based on neural networks trained on previous tasks for Gaussian process inference B3 37). Semi-supervised algorithms can incorporate unlabeled data into the training pipelines of BNNs using biologically plausible Bayesian Confidence Propagation Neural Networks (BCPNN) which model the cortex [11), by perturbing weights and using consistency regularization [T], or via semi-supervised deep kernel learning [26]. Gao et al. [66] also show how to harness unlabeled data to learn reference priors, uninformative priors which depend on the amount of training data and allow labels to most efficiently inform inference. Deep kernel learning has also
49
13
27
been used for Bayesian meta-learning in sew-shot classification and regression
33
In contrast to these works, we do not perform multi-task learning nor is our goal to harness auxiliary unlabeled downstream data; we instead approach transfer learning in which we leverage pre-training data for expressive priors that maximize performance on a single downstream task.
3 Bayesian Transfer Learning with Pre-Trained Priors
In order to transfer knowledge acquired through pre-training to downstream tasks, we adopt a pipeline with three simple components, composed of easy-to-use existing tools:
1. First, we a probability distribution with closed-form density to the posterior over feature extractor parameters using pre-trained checkpoint and SWAG (34) (Section 3.1)
3

---

2. Second, we re-scale this distribution, viewed as a prior for new tasks, to reflect the mismatch between the pre- pre-training and downstream tasks and to add coverage parameter settings which might be consistent with the downstream, but not pre-training, task. To this end, we tune a single scalar coefficient on held out validation data (Section B.2p.
Finally, we use this re-scaled prior, along with a zero-medt mern isotropic Gaussian prior over added parameters (e.g. classification head), to form a posterior on the downstream task. We then either optimize the posterior, or use it to perform full Bayesian inference with SGLD and SGHMC samplers [42, 5] (Section 3.3)
3. Finally, we plug the re-scaled prior into Bayesian inference algorithm, along with a zero- mean isotropic Gaussian prior over added parameters (e.g. classification head) form posterior on the downstream task. We then either optimize the posterior, or use it to perform full Bayesian inference with SGLD and SGHMC samplers [50 3 (Section 3.3).
We illustrate this pipeline in Figure in the Appendix. The simplicity and modularity of this framework are key strengths: by carefully combining easy-to-use existing components, we will see in Section that we can straightforwardly improve the default approaches to deep transfer learning. The potential for impact is significant: we can use this pipeline as a drop-in replacement for standard procedures used in deploying foundation models. At the same time, there is significant novelty: while Bayesian neural networks are typically used with simple zero-mean isotropic priors, we leverage the significant developments in self superviedd learning to produce highly informative priors. Moreover, in the following subsections [3.1 3.2 and 3.3 we will gain a nuanced understanding of each of
these three components and the key considerations for practical success. We discuss computational considerations in Section 3.4
For experiments in this section, use prior learned over the parameters of an ImageNet pre-trained SimCLR ResNet-50 feature extractor 9 6. 18], and we choose CIFAR- and CIFAR-100 for downstream tasks 30 We perform Bayesian inference with stochastic gradient Hamiltonian Monte
Carlo (SGHMC) [5]. The experiments in this section are primarily intended to gain conceptual insights into each step of our approach. In Section 4 we present our main empirical evaluations.
3.1 Learning Transferable Priors
We begin by building a probability distribution over the parameters of a feature extractor which represents knowledge we acquire from pre-training. To this end, we fit the distribution to the Bayesian posterior, regularized loss function, the pre-training task. This pre-training posterior will become the prior for downstream tasks and can be saved, or publicly released, for future use, similar to pre-trained weights. This stage requires two major design choices: a pre-training task and an algorithm for constructing the probability distribution.
Our downstream inference procedures require that our prior is represented as closed-form density function. Thus, we must use method that can provide a closed-form posterior approximation for the source task, which can be then re-scaled and used as a prior in the downstream task. We opt for SWA-Gaussian (SWAG) [34] due to its simplicity, scalability, popularity, and non-diagonal covariance. Other procedures that provide flseddform posterior approximations, such the Laplace approximation or variational methods, could also be applied, though as we will see, a non-diagonal covariance is particularly important to the success of this approach.
SWAG starts from a pre-trained model, and runs small number M of fine-tuning epochs with a modified learning rate schedule [34]. The SWAG approximate posterior distribution is given by N(W, 3 + Iooow-rnnk) where wे /H ΣM Wt is the SWA (23) solution, Dding T E M- diag(w; - w), and Soow = T-T M - (w: w)) (w, - w)) where I. is a hyperparameter controlling the rank of the low-rank component of the covariance matrix. After obtaining a closed-form distribution using SWAG, we remove the head from on top of the feature ex- tractor, for example, a linear classification module, and we consider only the distribution's 's restriction to the parameters of the feature extractor. New layers that are added for downstream tasks receive nea-led prior over their parameters.
To highlight the versatility of Bayesian transfer learning, we will focus both on supervised image clas- sification and also on self-supervised learning pre-training tasks and leverage existing torchvision and SimCLR [6] checkpoints. For both torchvision and SimCLR models, we learn the prior

---

using the associated loss function - eross- and InfoNCE, respectively, regularized by weight decay. In both settings, the regularized loss function can be represented as the sum of a negative log-likelihood and the negative log Gaussian density (weight decay).
While standard SGD-based transfer learning uses learned parameter vector, our Bayesian trans- fer learning approach uses a covariance matrix over feature extractor parameters which contains information about the pre-training loss surface geometry.
Can we really "learn" a prior? A data-dependent prior may sound odd, because a prior reflects our beliefs "before we see the data". However, it is entirely principled to learn prior, as long it is not learned using exactly the same labeled data we use in our downstream likelihood for the downstream task. Indeed, any informative prior is based on data that has shaped our beliefs.
Alignment of Pre-Trre Pre-Training and Downstream Loss Geometry. If the covariance matrix of our learned prior is to be more effective than an isotropic one, then it must not favour poorly generalizing directions for the downstream task. Thus, we compare the alignment of the learned covariance matrix with downstream task's s test loss.
We begin by computing the top leading singular vectors of SWAG learned covariance matrix over pooggeNet-trainee ResNet-50 feature extractor. We then train a linear
parameters of the SimCLR classifier head on CIFAR-10 10 training data on top of the fixed pre-trained feature extractor. Starting at the learned parameter vector, we perturb the feature extractor parameters in the direction of the singular vectors (fixing fully-connected classifier parameters), and measure the increase in test loss. We compare to the test loss when instead perturbing by each of 10 random vectors. All perturbation distances are filter normalized (as in [3 21), to account for invariance with respect to filter-wise parameter rescaling. In Figure 2a we see the CIFAR- 10 test loss is far flatter in the directions of leading eigenvectors of the pre-trained covariance than in a random direction, indicating the learned prior indeed promotes directions consistent with the downstream task.
After verifying that our pre-trained priors do in fact identify flat directions of the pre- treiiiiinn loss which are aligned with the downstream loss, we directly compare the performance benefits of our learned covariance over an isotropic covariance with learned mean. To this end, we swap out our learned prior's covariance with an isotropic version aT, where a is tuned on a held out validation set. Figure 2c] shows that a learned covariance consistently outperforms its isotropic counterpart, indicating that the shape of the pre-training loss surface's basin is informative for downstream tasks. The r--aiis here denotes the number of training samples used for fine-tuning the downstream task. We also see Bayesian model averaging provides further performance gains, which we discuss further Section 3.3 We now dissect just how important the low-rank component is.
Learned Covariance Outperforms Only a Learned Mean.
How the Prior Best Structured? the context of continual learning, elastic weight consolidation (EWC) [29] uses purely diagonal covariance priors to help prevent forgetting in continual learning. In this paper, we are interested in transfer learning, and wish to understand the benefits of low- rank component for capturing particularly important directions in the pre-training loss surface for transferring to downstream tasks. Omitting the low-rank component slightly reduces the memory footprint, but loses potentially important information about the shape of the pre-training posterior mode. In practice, the rank the matrix is determined by the number of samples collected when running SWAG. In Figure 2b we present the performance of prior learned via SWAG SimCLR Re feature extractor for transfer learning to CIFAR-10 10 and CIFAR-100. As the rank of the low-rank component increases from zero (diagonal covariance), dramatically until it saturates, indicating that only a small number of dominant directions are important for the transferability of the prior. Note that performance saturation occurs later for CIFAR- 100 than CIFAR- 10 - likely due to the higher complexity of CIFAR-100 compared to CIFAR- 10, which has 10x fewer classes.
the performance improves
3.2 Rescaling the Prior
In incremental learning. it is common to acquire some data, form posterior, and then use this posterior as our new prior acquiring future data. This procedure is equivalent to forming the posterior from all of the data at once, assuming all data are drawn from the same distribution with the same likelihood. However, transfer learning, we assume the data from the source and downstream task are drawn from different but related distributions. Thus do not want to directly re-use
5

---

Random
EMAR-1 CIPAR-100
2.3
17.1
Test Error
Downstream
e Eggenvectors
1.3
Perturbation Size (Fittr Normal ized)
Rank of Rank Component
(a) Pre-training and downstream losses are aligned
(b) Optimal rank of low-rank component
184
Test Error
Test Error
6.0
40
BNN Leamed Frtor
SC0 leamed Prtor
BNN Kon-Lzamed Prerr
4.5
Learned Fitor bobropic
10a
10a
10
10a
5 x 10
10a
10a
10f
10P
Traininng Samples
Covariance Scaling Factor
(c) Importance of a learned covariance
(d) Optimal prior scaling factor
Figure 2: Learning Transferable Priors. Error bars corresponds standard error over trials. All experiments performed with ResNet-50 backbone pre-trained via SimCLR on ImageNet. (a) Test loss (CIFAR-10) is less sensitive to perturbations in the directions of leading singular vectors of learned prior's covariance compared to random directions. Estimated over 10 random directions. (b) A non-diagonal covariance is important, though performance on downstream tasks plateaus at a relatively low rank. CIFAR- 100 benefits from a higher rank covariance than CIFAR-10 10, due to its relative complexity. (C) A non-diagonal covariance outperforms an isotropic covariance with learned mean CIFAR-100. Bayesian inference provides further performance gains. (d) We must re-scale the posterior from the source task, since data from each task drawn from a different distribution. If A is too small, then we assume too much similarity between tasks. If À is too large, lose useful information from the source task. Downstream inference performed on CIFAR-10.
posterior from the source as a prior for the downstream task, without modification. For example, as we acquire more data from the source task, our posterior will become increasingly concentrated. This concentration is issuc, certainty with respect to the ideal parameters for the pre-training task does not imply certainty on the idcal parameters for the downstream task.
To address this consideration, we rescale the learned Gaussian prior by multiplying its covariance matrix by a scalar value. We select the highest performing scalar value across grid on holdout set from the downstream task. If we do not scale the covariance enough, our prior will become concentrated around parameters which are inconsistent with the downstream task, and we scale the covariance too much, our prior will be diffuse and assign too much mass to regions of parameter space which are again inconsistent with the downstream task.
We now put this intuition to the test. We the SimCLR pre-training loss using a Gaussian prior with mean a and covariance matrix Σ. Given our uncertainty regarding the strength of the relationship between the pre- rre-triniig and downstream tasks, it is unclear if our learned prior would lead to over-confidence on the downstream task. To rectify this problem, matrix AΣ with X > 1. Figure 2d shows the accuracy of our method on CIFAR-10 10 as a function of X. As expected, we see that we need make the prior more diffuse we are optimize performance. In fact, prior rescaling can be the difference between strong and poor performance, We also see there is an optimal scaling factor where the prior is neither overly diffuse nor concentrated around solution which is inconsistent with the downstream task. Additionally, small scaling factors constrain ourselves to poor parameters, hurting performance much more than large scaling factors, which simply inducing near-uniform prior, negating the benefits of transfer learning.
instead assign prior covariance

---

3.3 Bayesian Inference
After learning prior and re-scaling it, we finally must draw samples from the downstream task's posterior over the parameters of the entire model, including additional modules, such as classification heads, which were added after pre-training specifically for a particular downstream task. We use a zero-mean isotropic Gaussian prior over these additional parameters, with scaling that is again tuned on held-out training data from the downstream task. Since we obtain a closed-form prior, the re-scaled distributions we learn compatible with a wide variety of Bayesian inference algorithms. In our experiments, we choose stochastic gradient Hamiltonian Monte Carlo (SGHMC) [5] and Stochastic Gradient Langevin Dynamics (SGLD) [50 since these samplers simultaneously provide strong performance and tractable computation costs. Once we have obtained samples from the downstream posterior, use these samples to form Bayesian model average for test-time predictions.
In Figure 2c] we also observe the advantage of Bayesian inference (SGHMC) over MAP estimation (SGD) on the negative rog-posterior yielded by our learned prior. Bayesian inference outperforms MAP estimation across all dataset sizes we consider.
In general, Bayesian inference can benefit most from an expressive prior. Indeed, priors reshape the whole posterior landscape- and the distinctive feature of Bayesian approach that marginalizes over the whole posterior, rather than simply using an optimum as with standard training.
3.4 Practical Considerations
While our method requires learned prior and a re-scaling coefficient, it is very easy to use and has minimal computational costs over standard fine-tuning routines. In particular, no expert intervention or knowledge of Bayesian methods is required, and only a single hyperparameter needs to be tuned. In the Appendix Section A we evaluate the costs of each of the three stages of our pipeline in detail: (i) inferring the posterior on the source task; (ii) re-scaling the posterior to become an informative prior for the downstream task; (iii) using the informative prior in the downstream task.
In short, (i) has a minor cost (e.g.. 1 of an epoch on ImageNet) that is a one-time cost when a pre-trained posterior is publicly released, as we have done; (ii) involves a single hyperparameter that can be tuned the same way as other hyperparameters (simple validation grid search), does not require specialized expertise, and adds 7 to the runtime of fine-tuning; (iii) has no additional cost if we do MAP optimization and has a cost comparable to deep ensembles if we run our implementation of Bayesian inference, both in terms of fine-tuning and test-time, and we show that Bayesian inference significantly outperforms deep ensembles.
4 Experiments
We now conduct a thorough empirical evaluation of our transfer learning pipeline on image classifi- cation and semantic segmentation. We generally consider five approaches, which we find have the following performance ordering: (1) Bayesian inference with Iearned priors, (2) SGD with learned priors, (3) SGD with standard pre-training, (4) Bayesian inference with non-learned zero-mean priors, (5) SGD with non-Iearned zero-mean priors. We note that both (1) and (2) are part of our framework, and that SGD with learned priors (2) significantly outperforms standard transfer learning (3).
4.1 Experimental Setting
We adopt the ResNet-50 and ResNet- 101 architectures [8), and we scale the input image to 224 x 224 pixels to accommodate these feature extractors designed for ImageNet data, We use SimCLR (SSL) ResNet-50 checkpoint (6] pre-trained on the ImageNet Ik dataset [9] and fit our prior to the SimCLR loss function. For the supervised setting, we use pre-trained torchvist on ResNet-50 and ResNet- 101 backbones. We perform image classification experiments on four downstream tasks: CIFAR- 10, CIFAR- 100 30) Oxford Flowers- 102 [36), and Oxford-IIIT Pets [38] On semantic segmentation, we use DeepLabv3+ system 4 with ResNet-50 and 101 backbone architectures, and we evaluate performance on the Pascal VOC 2012 [14] and Cityscapes 7 datasets. All error bars represent one standard error over 5 runs. We evaluate over a variety of downstream training set sizes, as transfer learning often involves limited downstream data, We provide a detailed description of hyperparameters in Appendix C.I
7

---

90
65
Test Error
1xter
40
15
10a
10a
10a
10a
5 10
10a
10a
10a
10a
10*
#Training Samples
#Training Samples
(a) CIFAR-10
(b) CIFAR-I00
90
21
112
Test Error
65
65
40
3MK Leamndd Prior tours)
40
Leamned Prior toursl
Kon-Leamed Prtor
15
Tunsfer
10a
10a # #Training Samples
10a
10a
10a
10a 3 10
#Training Samples
(c) Oxford Flowers- 102
(d) Oxford-IIIT Pets
Figure 3: Performance comparison. Learned priors outperform non-learned priors, and Bayesian inference generally outperforms SGD training, especially with learned priors. Moreover, learned priors tend to be more data-efficient. We use a ResNet-50 and SimCLR for pre-trained priors. The horizontal axis denotes the number of downstream training samples on logarithmic scale. Error bars correspond to a standard eтror over 5 trials.
4.2 Performance Comparison
In Figure 3 we compare the five approaches described above across various dataset sizes. We observe the following: (i) learned priors consistently outperform SGD transfer learning which in turn outperforms non-learned priors; (ii) conditioned on using the same prior, Bayesian inference often outperforms SGD training; (iii) Bayesian inference adds more value when used with informative priors; (iv) learned priors are relatively most valuable on intermediate data sizes for the downstream task. Point (iv) is particularly interesting: unlike standard SGD transfer learning, which involves an initialization from the pretraining task, the learned priors provide an explicit inductive bias, enabling the resulting model to learn more efficiently from downstream data. Once there is a sufficient amount of downstream data, the source pre- -reitraining becomes less important, and the methods become more similar in performance, although there is still significant gap.
Computation and comparison to deep ensembles. Our Bayesian model average (BMA) contains 10 samples from the posterior, and thus incurs higher test-time cost than a single SGD-trained model. We therefore additionally compare to an equally sized ensemble SGD-trained transfer learning models in Appendix B.3 We see that our BMA strongly outperforms the deep ensemble. In transfer learning, deep ensemble members are initialized with the same pre-trained checkpoint, and fine-tuning tends stay in the same basin, such that the ensemble components relatively homogenous in their predictions. Recall that SGD with our learned priors incurs essentially the same computational costs standard transfer learning, but has much better performance.
Comparison between self-supervised and supervised pre-training. In Appendix B.2 we provide additional evaluations with priors generated using torchvisicn ResNet-50 and 101 check- points. These priors differ from SimCLB CLR prior in that they are learned with labeled data rather than self-supervision task. We here also that learned supervised priors paired with Bayesian

---

inference consistently outperforms all baselines. We also see, notably, that the SSL priors are more transferable and outperform their supervised counterparts.
Evaluating uncertainty. We also measure predictive uncertainty via negative log test-likelihodd (NLL) in Figure/4h) with additional results in Appendix [B.4] We find that Bayesian transfer learning outperforms all other methods. However, we note that even though Bayesian inference with a non-learned prior has inferior accuracy to SGD with pre-training, it has superior test-likelihood iindicating that the confidence of SGD-based transfer learners is significantly miscalibrated, as likelihood accounts for both accuracy and uncertainty. We also evaluate the calibration of uncertainty using reliability diagrams (Figure the Appendix). These plots demonstrate that Bayesian inference with learned priors the best calibrated among methods we consider.
4.3 Out-of-Distribution Generalization
is test set of 2000 natural images modeled after CIFAR-10 10 with the same classes and image dimensions, following a similar dataset creation process to that described in the original CIFAR-10 10 paper. Models trained on CIFAR- 10 consistently perform much worse CIFAR-10. 1 despite the images being similarly easy for humans to classify. Figure 4alee4aja indicates that method achieves superior performance to SGD-bbaee transfer learning and Bayesian inference with non-learned priors on CIFAR-10 10.1 across training set sizes, where training sets are sampled from CIFAR-10 training data.
CIFAR-10.1 42
65
Test Eroo
40
BNN Leamed Prior leurs)
SGD Loamodd Prior lours)
24N kund eerned Prior
15
5G0 Trnnfer Intt
SGD Nan-Leamed Prior
5
10a
10a
10a
5xx
#Training Samples
#Training Samples
(a) Distributional shift on CIFAR-10.1.
(b) Test Negative LikeLikelihood on CIFAR-10.
Figure 4: Predictive likelihood and distribution shift. The horizontal axis denotes the number of downstream CIFAR-10 10 training samples on a logarithmic scale. Error bars correspond to standard over trials. All experiments performed with ResNet-50 backbone pre-trained via SimCLR on ImageNet. (a) BNNs with our learned priors similarly achieve superior accuracy on OOD testing. (b) Bayesian inference with learned prior achieves superior test NLL., but BNNs with non-learned prior outperforms all other methods despite often performing worst in terms of accuracy.
4.4 Semantic Segmentation with ImageNet Priors
Popular segmentation methods use ImageNet pre-trained weights initialization for backbone parameters [4]. We observe that semantic segmentation models, which contain numerous parameters aside from those of the backbone, still benefit immensely from learned prior over backbone parameters. Placing a strong prior, even if over only the backbone component of the segmentation system, enables us to more fully harness pre-training and boost performance across benchmark datasets PASCAL VOC 2012 and Cityscapes.
When constructing prior, apply SWAG to torchvision ("Supervised") and SimCLR ("SSL") ResNet-50 models and a torchvision ResNet- 101 model (the authors of SimCLR did not release ResNet-101 101 model). We assign zero-mean isotropic Gaussian prior to the decoder and atrous convolution parameters of DeepLabv3+ which are not pre-trdined.
In Table [ we see that Bayesian inference with learned priors achieves superior Mean-loU to both baselines (SGD and SGLD) without learned priors. Priors learned on the SimCLR objective again outperform those leamned on labeled data. We also present experiments with MAP estimates (SGD)

---

for the loss functions induced by both learned priors in Appendix We again find that our learned priors boost the performance of MAP estimates as well, and thus are preferable standard transfer learning, even for practitioners who do not intend to perform Bayesian inference.
O
SGD Transfer
Non-Learned Prior
Learned Prior
Learned Prior SSL
Dataset
Backbone
Supervised
ResNet-50
73.17
73.16
73.72
74.15
PASCAL VOC 2012
ResNet-101
75.53
75.79
76.27
ResNet-50
73.17
75.52
76.09
76.52
Cityscapes
ResNet-101
77.04
77.02
77.63
Table 1: Bayesian inference for semantic segmentation. The performance (Mean-lol -o of SGLD with learned priors (both supervised and SSL) and a non-learned prior over ResNet-50 and ResNet- 101 parameters for segmentation tasks. "SimCLR ResNet-101 checkpoints have not been released.
Discussion
Our work reveals several new key insights about deep transfer learning:
Modifying the loss surface on the downstream task through informative priors leads to significant performance gains, with and without Bayesian inference.
Bayesian inference provides particular performance boost with the informative priors.
Informative priors lead to more dula-efficien performunce on the downstream task.
The success of this approach depends on capturing key directions in the loss surface of the source task, which we represent through low-rank plus diagonal covariance matrix.
Standard transfer learning can be significantly miscalibrated, even providing worse likeli- hood than Bayesian methods from scratch on the downstream task.
Priors learned via self-supervised pre-training transfer better than those learning via supervised learning.
In short, pre- traininig your loss with care provides an easy drop-in replacement for conventional transfer learning that relies initialization.
10


Reviews:
[{"review_id": "AxTwMa9-Wt6", "reviewer": "Reviewer_3R6h", "summary": "This paper presents a Bayesian learning method for transfer learning on target task. Specifically, a posterior approximation method SWAG is used to estimate the posterior distribution of a supervised or self-supervised pre-trained model. This distribution is used as the prior for downstream learning. Experiments and ablation study demonstrate the effectiveness of learning a prior for transfer.", "questions": "Q1: Most of the evaluations in experiments focus on the performance of the learned prior when transferred to a new target task. I'm not sure if it is possible to use Bayesian model average can on source task itself, as it is not clear how good is the approximated posterior distribution. This study on the posterior approximation on the source task can probably show the effectiveness of estimation with close-form density as well as Bayesian learning.\n\nQ2: The rescaling method proposed in section 3.2 is not well developed. Are there any related work for the proposed scaling method? Practically, from Fig. 2d, it seems that the scaling factor is very big. It seems that it is necessary to tune this scaling factor for every new target dataset. Is it possible to estimate this parameter without using holdout validation set?\n\nQ3: Transfer learning results on target classification tasks are reported mainly using test error. A more common comparison is classification accuracies as error depends on the loss function. As seen from the tables the test error of several baseline methods are quite close. What are the classification accuracies for these results?\n\n", "limitations": "See weakness and questions.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n1) The problem of learning a prior for transfer with Bayesian learning is well motivated.\n\n2) The proposed method with three steps is reasonable.\n\n3) Experimental results of the proposed method on semantic segmentation are quite strong.\n\nWeaknesses\n\n1) This paper employs methods of posterior approximation and Bayesian learning and incorporates them into Bayesian transfer by learning a prior from source tasks. These existing methods have been well developed so the overall technical contribution is incremental upon these.\n\n2) The method introduced in section 3.2 looks a bit arbitrary. In order to slightly smooth the prior distribution, this paper proposes to multiply a single scaling factor to the covariance matrix. This does not seem to be a principled way to adapt a prior distribution to a target task. Are there any other possible methods of adjusting the covariance matrix for transfer? The holdout validation method used to select this scaling factor also neglects other data information that is potentially useful for determining the scaling factor.\n\n3) The background review on Bayesian transfer learning seems not very thorough, which makes the overall contribution of the proposed transfer learning with Bayesian learning unclear. The only relevant review on this direction is about continual learning.\n\n4) The proposed method will introduce additional computational burden for transfer learning.\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "9BXXsH9Nszr", "reviewer": "Reviewer_hPGA", "summary": "This paper proposes a Bayesian perspective of transfer learning in deep neural networks, whereby a re-scaled Bayesian parameter posterior from a source task is used as a pre-trained prior for a target task. The proposed procedure effectively reshapes the training objective of the target task to more faithfully reflect knowledge learnt from the source task. The authors find that modifying the loss surface of target tasks through informative priors significantly improves performance and calibration, especially with Bayesian inference. ", "questions": "Since transformers (visual or otherwise) are the main subjects of transfer learning nowadays, it would make a lot of sense to evaluate them too. It is also known that transformers incur different loss landscapes than typical CNNs such as ResNets, it would be interesting to see how performance, robustness and calibration are effected by the proposed transfer learning policy. ", "limitations": "See above.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n\n- Highly important and relevant topic for the research community;\n- Clear and organized exposition;\n- Well motivated arguments and experiments.\n\n**Weaknesses**:\n- Proposed approach can incur significant additional computational costs;\n- May require expert knowledge to use effectively.\n\nOverall I like this paper. I find the arguments and experiments convincing enough and believe it would be valuable for the research community. One concern is whether the proposed Bayesian inference based transfer learning pipeline would be of limited utility to regular practitioners, as expert knowledge may be required to get these systems working stably as intended. There is also the incurred computational cost which may deter uptake if the performance gain does not justify it. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "0zW5Eevj3uR", "reviewer": "Reviewer_X6YH", "summary": "This submission shows how an approximate posterior distribution on DNN weights learned on a source task can be beneficial in transfer-learning to downstream tasks. \n\nA 3-step method is proposed: 1) Fit posterior over weights on source task using SWAG. 2) Re-scale the learned posterior using a single scalar coefficient. 3) Plug the re-scaled prior into a Beysian inference algorithm like SGLD or SGHMC.\n\nExperimental results are provided to support the following claims:\n\n- In terms of performance, learned priors > SGD transfer learning > non-learned priors.\n- SGD with learned priors significantly outperforms standard transfer learning (i.e., learned priors lead to performance gains with or w/o Bayesian inference).\n- Learned priors lead to more data-efficient performance on the downstream task.\n- When using the same prior Bayesian inference often outperforms SGD training -- though Bayesian inference with a non-learned prior has inferior accuracy to SGD with pre-training.\n", "questions": "**Q1.** What would be expected increase in computation (tuning included) w.r.t. conventional transfer learning for applying the proposed method to common tasks?\n\n**Q2.** The method is compared to deep ensembles with regards to accuracy. How does the method compare with regards to overconfidence?\n", "limitations": "The two limitations discussed are the additional computation and the inclusion of only vision applications. ", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\n**S1.** The method is well-motivated, presented with clarity and shown to lead to performance gains.\n\n**S2.** The insights provided prompt further exploration of loss-surface alignment in transfer learning and other settings (e.g., few-shot learning).\n\n**S3.** The method builds on prior art and adds a single hyperparameter. For this reason it should be easy for others to apply.\n\n**S4.** The experimental validation includes classification tasks of different complexities and semantic segmentation. Thus, the method is shown to be beneficial in interesting problems.\n\n# Weaknesses\n\n**W1.** The proposed method claims to be a drop-in replacement for standard transfer learning.Because the additional computation is not expected to be trivial, it would have been beneficial to provide a proper study of the expected increase in computation w.r.t. standard transfer learning.\n\nThe submission includes some discussion about the additional training and inference cost but no discussion about the cost (in practice) of tuning hyperparameters. For example, there are a number of hyperparameters for the methods involved in the different steps of the proposed method. Setting or tuning these hyperparameters for a given transfer would be part of the cost of the method, in particular for steps 2 (scaling $\\lambda$ of prior) and 3 (parameters of SGLD and SGHMC).\n\n**W2.** The main technical novelty in the method is the scaling of the prior (steps 1 and 3 are applications of prior art) but this was not given much attention in the experiments (only Fig. 2c). What is the recommended procedure and expected cost for tuning the scaling parameter $\\lambda$?\n\n**W3.** Plots, legends and captions are not synchronized. For example, in Fig. 2c there is a purple plot but the legend does not include purple; in Fig. 7 the caption speaks of a red plot but there is no red in the plots.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept