Pre-Train Your Loss: Easy Bayesian Transfer Learning with Informative Priors
Anonymous Author(s)
AffiliationAddress email
Abstract
Deep learning is increasingly moving towards a transfer learning paradigm wherebylarge “foundation models” are fine-tuned on downstream tasks, starting from aninitialization learned on the source task.
But an initialization contains relativelylittle information about the source task.
Instead, we show that we can learn highlyinformative posteriors from the source task, through supervised or self-supervisedapproaches, which then serve as the basis for priors that modify the whole losssurface on the downstream task.
This simple modular approach enables significantperformance gains and more data-efficient learning on a variety of downstreamclassification and segmentation tasks, serving as a drop-in replacement for standardpre-training strategies.
These highly informative priors also can be saved for futureuse, similar to pre-trained weights, and stand in contrast to the zero-mean isotropicuninformative priors that are typically used in Bayesian deep learning
Introduction
The ability to transfer what is learned from one task to another — learning to ride a bicycle to thenride 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 learnedparameter vector is used as an initialization for training in a downstream task.
While this approachhas had great empirical success, reliance on an initialization is a very limited way to perform transferlearning.
Indeed, if we are doing a good job of optimization, then our final solution should beindependent of initialization, barring local minima with identical training loss.
We propose to instead use a re-scaled Bayesian parameter posterior from the first task as a pre-trainedprior for the downstream task.
Since the negative log posterior on the downstream task is our lossfunction, this procedure has the effect of reshaping our training objective on the downstream task,to reflect more nuanced information that we have learned from the source task, as in Figure 1.
Theposterior on the source task is re-scaled before use as a prior on the downstream task to reflect thebelief that the source and downstream tasks are drawn from different distributions.
With our now highly informative prior for the downstream task, we can then proceed with optimization,or perform full Bayesian model averaging with the posterior on the downstream task.
We find bothprocedures profoundly improve upon standard deep transfer learning, making use of a simple pipelinethat involves easy-to-use pre-existing components.
Indeed, the simplicity of the approach, combinedwith its promising empirical performance, is one of its greatest features — enabling its use as adrop-in replacement for standard approaches for deep transfer learning.
The pre-training posteriorcan be saved as a prior for a wide array of downstream tasks, similar to pre-trained weights.
[Figure]
Despite the simplicity and effectiveness of this approach, Bayesian neural networks are almost alwaysused with uninformative isotropic zero-mean priors [46, 25, 15], and have not harnessed recentadvances in self-supervised learning.
Furthermore, as we will show, the details of how we perform
Bayesian transfer learning with modern neural networks are crucial to 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 we expect to provide good solutions on the sourcetask, which significantly outperforms isotropic or diagonal counterparts.
We also provide severalkey conceptual findings — for example, that the informative priors provide an inductive bias thatenables more data efficient fine-tuning than a pre-trained initialization, and that pre-trained priorsbased on self-supervised learning are more transferable than supervised pre-trained priors.
Finally, weobserve through extensive experiments across both image classification and semantic segmentation onmultiple neural architectures, and on both supervised and self-supervised (SimCLR [6]) pre-trainingloss functions, that Bayesian inference is particularly advantageous in the transfer learning setting.
We have included PyTorch pre-trained priors, and code to reproduce our experiments for learningand using priors for downstream inference.
Background
Bayesian Neural Networks.
Consider a neural network f with weights w and a training dataset
D = { ( x ( i ) , y ( i ) ) } of independently drawn samples.
In standard non-Bayesian neural networktraining, we minimize the negative log posterior, − log p ( w |D , f ) ∝ − log p ( D| w, f ) − log p ( w | f ) ,where the likelihood p ( D| w, f ) denotes the probability that a model with weights w would generatethe observed labels { y ( i ) } (e.g., cross-entropy).
The log prior log p ( w | f ) often takes the formof weight decay, corresponding to a zero-mean isotropic Gaussian prior.
In Bayesian modeling,we instead make predictions with a Bayesian Model Average (BMA) of all models weighted bytheir posterior probabilities: p ( y | x, D , f )
= (cid:82) p ( y | x, w, f ) p ( w |D , f ) dw .
There are many ways toapproximate this integral, such as MCMC, or variational approaches, which take a finite numberof approximate samples w j from the parameter posterior to form the simple Monte Carlo estimate:p ( y | x, D , f )
≈ 1 J (cid:80) j p ( y | x, w j , f ) .
Almost always, one uses zero-mean isotropic priors [46].
Thereare also specialized priors, including heavy-tailed priors [15, 25], noise-contrastive priors for highuncertainty under distribution shift [17], and input-dependent priors for domain generalization [24].
Transfer Learning.
In transfer learning , we wish to recycle the representation learned on one taskto improve performance on another.
Transfer learning is now widely applied in deep learning, andforms the basis for foundation models [1, 12, 8, 42, 4, 20, 10], which are exceptionally large neuralnetworks pre-trained on massive volumes of data, and then fine-tuned on a downstream task.
Recentwork has found self-supervised pre-training can transfer better than supervised pre-training [19], inline with our finding in Section 4 that self-supervised pre-trained priors transfer better.
In continual learning , we wish to learn tasks sequentially without forgetting what has been learnedpreviously.
Elastic Weight Consolidation (EWC) prevents forgetting by imposing a penalty, thediagonal of a Fisher information matrix computed on previous tasks, for adapting to a new task [29].
We can contextualize the EWC penalty within Bayesian transfer learning as a negative log Gaussianprior with diagonal covariance used for maximum a posteriori (MAP) estimation on sequential tasks,such as digit classification or RL [29].
In our experiments, we find that MAP estimation, diagonalcovariance, and supervised pre-training are all suboptimal in the transfer learning setting.
Leveraging Auxiliary Data and Knowledge Transfer in Bayesian Modeling.
A number of worksoutside of deep learning have considered knowledge transfer in Bayesian modeling, especially insettings such as domain adaptation or homogeneous transfer learning in which source and targettasks contain similar feature and label spaces but differ in their marginal distributions.
For example,
Xuan et al.
[47] considers Bayesian transfer learning methods for probabilistic graphical models,and Karbalayghareh et al.
[28] develop a theoretical framework for understanding optimal Bayesclassifiers given prior knowledge from other domains.
Other works on Bayesian transfer learninglearn a dirichlet prior over naive Bayes classifiers [43].
Bayesian optimization methods can alsofind transferable hyperparameter settings [39] or select data which will yield transferable models onshifted domains [44].
Bayesian tools have additionally been used for leveraging auxiliary data or multiple domains in deeplearning.
Chandra and Kapoor [2] learn from multiple domains simultaneously using a round-robintask sampling procedure and a single-layer neural network.
Bayesian methods for continual learningupdate the posterior to accommodate new tasks without forgetting how to perform previous ones[34, 13, 27], or develop kernels based on neural networks trained on previous tasks for Gaussianprocess inference [32, 36].
Semi-supervised algorithms can incorporate unlabeled data into thetraining pipelines of BNNs using biologically plausible Bayesian Confidence Propagation Neural
Networks (BCPNN) which model the cortex [40], by perturbing weights and using consistencyregularization [11], or via semi-supervised deep kernel learning [26].
Gao et al.
[16] also show howto harness unlabeled data to learn reference priors, uninformative priors which depend on the amountof training data and allow labels to most efficiently inform inference.
Deep kernel learning has alsobeen used for Bayesian meta-learning in few-shot classification and regression [38].
In contrast to these works, we do not perform multi-task learning nor is our goal to harness auxiliaryunlabeled downstream data; we instead approach transfer learning in which we leverage pre-trainingdata for expressive priors that maximize performance on a single downstream task.
Bayesian Transfer Learning with Pre-Trained Priors
In order to transfer knowledge acquired through pre-training to downstream tasks, we adopt a pipelinewith three simple components, composed of easy-to-use existing tools:1.
First, we fit a probability distribution with closed-form density to the posterior over featureextractor parameters using a pre-trained checkpoint and SWAG [33] (Section 3.1).2.
Second, we re-scale this distribution, viewed as a prior for new tasks, to reflect the mismatchbetween the pre-training and downstream tasks and to add coverage to parameter settingswhich might be consistent with the downstream, but not pre-training, task.
To this end, wetune a single scalar coefficient on held out validation data (Section 3.2).3.
Finally, we plug the re-scaled prior into a Bayesian inference algorithm, along with a zeromean isotropic Gaussian prior over added parameters (e.g. classification head) to solve thedownstream task (Section 3.3).
We use SGLD and SGHMC samplers [45, 5].
We illustrate this pipeline in Figure 5 in the Appendix.
The simplicity and modularity of thisframework are key strengths: by carefully combining easy-to-use existing components, we will seein Section 4 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 standardprocedures 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 thesignificant developments in self-supervised 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 computationalconsiderations in Appendix A.
For experiments in this section, we use a prior learned over the parameters of an ImageNet pre-trained
SimCLR ResNet-50 feature extractor [9, 6, 18], and we choose CIFAR-10 and CIFAR-100 fordownstream tasks [30].
We perform Bayesian inference with stochastic gradient Hamiltonian Monte
Carlo (SGHMC) [5].
The experiments in this section are primarily intended to gain conceptualinsights 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 whichrepresents knowledge we acquire from pre-training.
To this end, we fit the distribution to the Bayesianposterior, or regularized loss function, on the pre-training task.
This pre-training posterior willbecome the prior for downstream tasks and can be saved, or publicly released, for future use, similarto pre-trained weights.
This stage requires two major design choices: a pre-training task and analgorithm for constructing the probability distribution.
Our downstream inference procedures require that our prior is represented as a closed-form densityfunction.
Thus, we must use a method that can provide a closed-form posterior approximationfor the source task, which can be then re-scaled and used as a prior in the downstream task.
Weopt for SWA-Gaussian (SWAG) [33] due to its simplicity, scalability, popularity, and non-diagonalcovariance.
Other procedures that provide closed-form posterior approximations, such as the Laplaceapproximation or variational methods, could also be applied, though as we will see, a non-diagonalcovariance is particularly important to the success of this approach.
SWAG starts from a pre-trained model, and runs a small number M of fine-tuning epochs witha modified learning rate schedule [33].
The SWAG approximate posterior distribution is givenby N ( w, 12 Σ diag + 12 Σ low-rank ) , where w = 1 M (cid:80)
Mt
=1 w t is the SWA [23] solution, Σ diag =
L − 1 (cid:80) Mt
=
M − L +1 diag ( w t − w ) , and Σ low-rank =
1 L − 1 (cid:80) Mt =
M − L +1 ( w t − w ) ( w t − w ) ⊤ , where L isa hyperparameter controlling the rank of the low-rank component of the covariance matrix.
Afterobtaining a closed-form distribution using SWAG, we remove the head from on top of the feature extractor, for example, a linear classification module, and we consider only the distribution’s restrictionto the parameters of the feature extractor.
New layers that are added for downstream tasks receive anon-learned prior over their parameters.
To highlight the versatility of Bayesian transfer learning, we will focus both on supervised image classification and also on self-supervised learning pre-training tasks and leverage existing torchvisionand SimCLR [6] checkpoints.
For both torchvision and SimCLR models, we learn the priorusing the associated loss function — cross-entropy and InfoNCE, respectively, regularized by weightdecay.
In both settings, the regularized loss function can be represented as the sum of a negativelog-likelihood and the negative log Gaussian density (weight decay).
While standard SGD-based transfer learning uses a learned parameter vector, our Bayesian transfer learning approach uses a covariance matrix over feature extractor parameters which containsinformation about the pre-training loss surface geometry.
Can we really “learn” a prior?
A data-dependent prior may sound odd, because a prior reflects ourbeliefs “before we see the data”.
However, it is entirely principled to learn a prior, as long as it is notlearned using exactly the same labeled data we use in our downstream likelihood for the downstreamtask.
Indeed, any informative prior is based on data that has shaped our beliefs.
Alignment of Pre-Training and Downstream Loss Geometry.
If the covariance matrix of ourlearned prior is to be more effective than an isotropic one, then it must not favour poorly generalizingdirections for the downstream task.
Thus, we compare the alignment of the learned covariance matrixwith a downstream task’s test loss.
We begin by computing the top 5 leading singular vectors of a SWAG learned covariance matrix overparameters of the SimCLR ImageNet-trained ResNet-50 feature extractor.
We then train a linearclassifier head on CIFAR-10 training data on top of the fixed pre-trained feature extractor.
Startingat the learned parameter vector, we perturb the feature extractor parameters in the direction of thesingular 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 perturbationdistances are filter normalized (as in [31, 21]), to account for invariance with respect to filter-wiseparameter rescaling.
In Figure 2a, we see the CIFAR-10 test loss is far flatter in the directions ofleading eigenvectors of the pre-trained covariance than in a random direction, indicating the learnedprior indeed promotes directions consistent with the downstream task.
Learned Covariance Outperforms Only a Learned Mean.
After verifying that our pre-trainedpriors do in fact identify flat directions of the pre-training loss which are aligned with the downstreamloss, we directly compare the performance benefits of our learned covariance over an isotropiccovariance with a learned mean.
To this end, we swap out our learned prior’s covariance withan isotropic version αI , where α is tuned on a held out validation set.
Figure 2c shows that alearned covariance consistently outperforms its isotropic counterpart, indicating that the shape ofthe pre-training loss surface’s basin is informative for downstream tasks.
The x -axis here denotesthe number of training samples used for fine-tuning on the downstream task.
We also see Bayesianmodel averaging provides further performance gains, which we discuss further in Section 3.3.
Wenow dissect just how important the low-rank component is.
How is the Prior Best Structured?
In 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 a lowrank component for capturing particularly important directions in the pre-training loss surface fortransferring to downstream tasks.
Omitting the low-rank component slightly reduces the memoryfootprint, but loses potentially important information about the shape of the pre-training posteriormode.
In practice, the rank of the matrix is determined by the number of samples collected whenrunning SWAG.
In Figure 2b, we present the performance of a prior learned via SWAG on a SimCLR
ResNet-50 feature extractor for transfer learning to CIFAR-10 and CIFAR-100.
As rank of thelow-rank component increases from zero (diagonal covariance), we see the performance improvesdramatically until it saturates, indicating that only a small number of dominant directions are importantfor 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 has10 × fewer classes.
3.2 Rescaling the Prior
In incremental learning , it is common to acquire some data, form a posterior, and then use thisposterior as our new prior in acquiring future data.
This procedure is equivalent to forming theposterior from all of the data at once, assuming all data are drawn from the same distribution with thesame likelihood.
However, in transfer learning, we assume the data from the source and downstreamtask are drawn from different but related distributions.
Thus we do not want to directly re-use aposterior from the source as a prior for the downstream task, without modification.
For example, aswe acquire more data from the source task, our posterior will become increasingly concentrated.
Thisconcentration is an issue, as certainty with respect to the ideal parameters for the pre-training taskdoes not imply certainty on the ideal parameters for the downstream task.
To address this consideration, we rescale the learned Gaussian prior by multiplying its covariancematrix by a scalar value.
We select the highest performing scalar value across a grid on a holdoutset from the downstream task.
If we do not scale the covariance enough, our prior will becomeconcentrated around parameters which are inconsistent with the downstream task, and if we scalethe covariance too much, our prior will be diffuse and assign too much mass to regions of parameterspace which are again inconsistent with the downstream task.
We now put this intuition to the test.
We fit the SimCLR pre-training loss using a Gaussian prior withmean µ and covariance matrix Σ .
Given our uncertainty regarding the strength of the relationshipbetween the pre-training and downstream tasks, it is unclear if our learned prior would lead toover-confidence on the downstream task.
To rectify this problem, we instead assign prior covariancematrix λ Σ with λ ≥ 1 .
Figure 2d shows the accuracy of our method on CIFAR-10 as a function of λ .
As expected, we see that we need to make the prior more diffuse if we are to optimize performance.
In fact, prior rescaling can be the difference between strong and poor performance.
We also seethere is an optimal scaling factor where the prior is neither overly diffuse nor concentrated around asolution which is inconsistent with the downstream task.
Additionally, small scaling factors constrain
[Figure]
task, since data from each task is drawn from a different distribution.
If λ is too small, then we assume too much similarity between tasks.
If λ is too large, we lose useful information from the source task.
Downstream inference is performed on CIFAR-10.
ourselves to poor parameters, hurting performance much more than large scaling factors, whichsimply inducing a near-uniform prior, negating the benefits of transfer learning.
3.3 Bayesian Inference
After learning a prior and re-scaling it, we finally must draw samples from the downstream task’sposterior over the parameters of the entire model, including additional modules, such as classificationheads, which were added after pre-training specifically for a particular downstream task.
We use azero-mean isotropic Gaussian prior over these additional parameters, with a scaling that is again tunedon held-out training data from the downstream task.
Since we obtain a closed-form prior, the re-scaleddistributions we learn are compatible with a wide variety of Bayesian inference algorithms.
In ourexperiments, we choose stochastic gradient Hamiltonian Monte Carlo (SGHMC) [5] and Stochastic
Gradient Langevin Dynamics (SGLD)
[45] since these samplers simultaneously provide strongperformance and tractable computation costs.
Once we have obtained samples from the downstreamposterior, we use these samples to form a 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 log-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 thewhole posterior landscape — and the distinctive feature of a Bayesian approach is that it marginalizesover the whole posterior, rather than simply using an optimum as with standard training.
[Figure]
We now conduct a thorough empirical evaluation of our transfer learning pipeline on image classification and semantic segmentation.
We generally consider five approaches, which we find have thefollowing performance ordering: (1) Bayesian inference with learned priors,
(2) SGD with learnedpriors, (3) SGD with standard pre-training, (4) Bayesian inference with non-learned zero-mean priors,(5) SGD with non-learned 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
4.1 Experimental Setting
We adopt the ResNet-50 and ResNet-101 architectures [18], and we scale the input image to 224 × 224pixels to accommodate these feature extractors designed for ImageNet data.
We use a SimCLR (SSL)
ResNet-50 checkpoint [6] pre-trained on the ImageNet 1k dataset [9] and fit our prior to the SimCLRloss function.
For the supervised setting, we use pre-trained torchvision ResNet-50 and ResNet101 backbones.
We perform image classification experiments on four downstream tasks: CIFAR-10,
CIFAR-100 [30], Oxford Flowers-102 [35], and Oxford-IIIT Pets [37].
On semantic segmentation,we use a DeepLabv3+ system [4] with ResNet-50 and ResNet-101 backbone architectures, andwe evaluate performance on the Pascal VOC 2012 [14] and Cityscapes [7] datasets.
All error barsrepresent 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 ofhyperparameters in Appendix C.1.
4.2 Performance Comparison
In Figure 3, we compare the five approaches described above across various dataset sizes.
Weobserve the following: (i) learned priors consistently outperform SGD transfer learning which in turnoutperforms non-learned priors; (ii) conditioned on using the same prior, Bayesian inference oftenoutperforms SGD training; (iii) Bayesian inference adds more value when used with informativepriors; (iv) learned priors are relatively most valuable on intermediate data sizes for the downstreamtask.
Point (iv) is particularly interesting: unlike standard SGD transfer learning, which involves aninitialization from the pretraining task, the learned priors provide an explicit inductive bias, enablingthe resulting model to learn more efficiently from downstream data.
Once there is a sufficient amountof downstream data, the source pre-training becomes less important, and the methods become moresimilar in performance, although there is still a significant gap.
Computation and comparison to deep ensembles.
Our Bayesian model average (BMA) contains10 samples from the posterior, and thus incurs a higher test-time cost than a single SGD-trainedmodel.
We therefore additionally compare to an equally sized ensemble of SGD-trained transferlearning 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 to stay in the same basin, such that the ensemble components are relativelyhomogenous in their predictions.
Recall that SGD with our learned priors incurs essentially the samecomputational costs as standard transfer learning, but has much better performance.
Comparison between self-supervised and supervised pre-training.
In Appendix B.2 we provideadditional evaluations with priors generated using torchvision ResNet-50 and ResNet-101 checkpoints.
These priors differ from the SimCLR prior in that they are learned with labeled data ratherthan a self-supervision task.
We see here also that learned supervised priors paired with Bayesianinference consistently outperforms all baselines.
We also see, notably, that the SSL priors are moretransferable and outperform their supervised counterparts.
Evaluating uncertainty.
We also measure predictive uncertainty via negative log test-likelihood(NLL) in Figure 4b, with additional results in Appendix B.4.
We find that Bayesian transfer learningoutperforms all other methods.
However, we note that even though Bayesian inference with anon-learned prior has inferior accuracy to SGD with pre-training, it has superior test-likelihood— indicating that the confidence of SGD-based transfer learners is significantly miscalibrated, aslikelihood accounts for both accuracy and uncertainty.
We also evaluate the calibration of uncertaintyusing reliability diagrams (Figure 9 in the Appendix).
These plots demonstrate that Bayesian inferencewith learned priors is the best calibrated among methods we consider. .
4.3 Out-of-Distribution Generalization
CIFAR-10.1 [41] is a test set of 2000 natural images modeled after CIFAR-10 with the sameclasses and image dimensions, following a similar dataset creation process to that described inthe original CIFAR-10 paper.
Models trained on CIFAR-10 consistently perform much worse on
CIFAR-10.1 despite the images being similarly easy for humans to classify.
Figure 4a indicates thatour method achieves superior performance to SGD-based transfer learning and Bayesian inferencewith non-learned priors on CIFAR-10.1 across training set sizes, where training sets are sampledfrom CIFAR-10 training data.
4.4 Semantic Segmentation with ImageNet Priors
Popular segmentation methods use ImageNet pre-trained weights as an initialization for backboneparameters [4].
We observe that semantic segmentation models, which contain numerous parametersaside from those of the backbone, still benefit immensely from a learned prior over backboneparameters.
Placing a strong prior, even if over only the backbone component of the segmentationsystem, enables us to more fully harness pre-training and boost performance across benchmarkdatasets PASCAL VOC 2012 and Cityscapes.
When constructing a prior, we apply SWAG to torchvision (“Supervised”) and SimCLR (“SSL”)
ResNet-50 models and a torchvision ResNet-101 model (the authors of SimCLR did not releasea ResNet-101 model).
We
[Figure]
In Table 1, we see that Bayesian inference with learned priors achieves superior Mean-IoU to bothbaselines (SGD and SGLD) without learned priors.
Priors learned on the SimCLR objective againoutperform those learned on labeled data.
We also present experiments with MAP estimates (SGD)for the loss functions induced by both learned priors in Appendix C.2.
We again find that our learned priors boost the performance of MAP estimates as well, and thus are preferable to standard transfer learning, even for practitioners who do not intend to perform Bayesian inference.
[Table]
Discussion
Our work reveals several new key insights about deep transfer learning:
• Modifying the loss surface on the downstream task through informative priors leads tosignificant performance gains, with and without Bayesian inference .
• Bayesian inference provides a particular performance boost with the informative priors .•
Informative priors lead to more data-efficient performance on the downstream task .•
The success of this approach depends on capturing key directions in the loss surface of thesource task, which we represent through a low-rank plus diagonal covariance matrix .•
Standard transfer learning can be significantly miscalibrated, even providing worse likelihood than Bayesian methods from scratch on the downstream task .
• Pre-training with self-supervised learning transfers better than supervised learning .
In short, pre-training your loss with care provides an easy drop-in replacement for conventionaltransfer learning that relies on initialization.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Appendix E.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] We includeadditional details about broader impact in the supplementary material.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main exper-
imental results (either in the supplemental material or as a URL)?
[Yes] We provide
link to our code to replicate all experiments
(b) Did you specify all the training details (e
g., data splits, hyperparameters, how theywere chosen)?
[Yes] See sections B and A in the supp.
material.
(c) Did you report error bars (e
g., with respect to the random seed after running experiments multiple times)?
[Yes] We include error estimates in all figures.
(d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] We estimate compute in section Ain the supp.
material.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
[Figure] [Figure]
A Computational Considerations
[Figure]
Our method involves three steps: (i) inferring posteriors on the source task; (ii) re-scaling theposteriors to form informative priors for the downstream task; and (iii) utilizing the informative priorson the downstream task.
We now evaluate the costs of each component: (i) The posterior on thesource task is learned in roughly a quarter of an epoch of ImageNet training.
We apply SWAG toa pre-trained checkpoint [33].
For burn-in, we run 50 iterations, then 10 cycles of 200 iterationseach for a total of 2050 iterations with a batch size of 128 .
The computations in this step can beavoided entirely by users releasing pre-trained SWAG posteriors, as we do for all of the tasks inour paper.
The only additional cost of re-scaling the posterior comes from tuning a single scalingparameter.
Using a grid search between 10 4 − 10 8 with six samples, we tune this parameter similarlyto other hyperparameters in standard transfer learning on a held-out validation set after we found theother hyper-parameters.
Aside from this hyperparameter, we have seven more for SGD-based finetuning and eight more for Baysian inference-based fine-tuning.
In other words, tuning the additionalprior re-scaling hyparameter adds 1 / 7 of the total runtime of hyperparameter tuning.
For example,hyperparameter tuning (including the re-scaling coefficient) and running MAP on the full Oxford
Flowers-102 dataset takes 3840 minutes, whereas hyperparameter tuning and running SGD-basedtransfer learning takes 3360 minutes on a computer with a single NVIDIA GeForce RTX
Ti
GPU.
There are two options for using the informative prior in the downstream task , both of whichhave the same computational cost as their corresponding baseline competitors.
MAP optimizationwith informative priors using SGD, which performs better than standard initialization-based transferlearning, incurs the same cost.
For Bayesian inference with SGLD using our learned prior, a singlechain of SGLD incurs the same cost as SGD, since the gradient of the log prior density can becomputed without automatic differentiation in the same fashion as weight-decay, which we arereplacing.
When multiple chains are used, we compare our method with deep ensembles, whichrequire the same overhead but result in poorer performance.
As we have described, the first step ofour method has a minor cost (e.g., 1/4 of an epoch on ImageNet) that can be avoided by releasingpre-trained posteriors.
The second step requires tuning a single additional hyperparameter in the samefashion as other hyperparameters, does not require specialized expertise, and adds 1/7 the runtime ofstandard fine-tuning.
For the third step, MAP optimization bears no additional cost, and multiplechains of SGLD bears a cost comparable to deep ensembles.
B ClassificationB.1 Implementation Details
Based on the evaluation protocols in the papers introducing the datasets, we report the top-oneaccuracy for CIFAR-10 and CIFAR-100 and mean per-class accuracy for Oxford-IIIT Pets and
Oxford 102 Flowers.
For Oxford 102 Flowers, we select hyperparameters based on the validationsets specified by the dataset creators.
While tuning hyperparameters, we hold out a subset of thetraining set for validation on the other datasets.
After selecting the optimal hyperparameters from thevalidation set, we retrain the model using the selected parameters on both the training and validationimages.
Test results are reported.
We use ResNet-50 and ResNet-101 architectures [18] with supervised and SSL checkpoints [6]pre-trained on the ImageNet 1k dataset [9].
We use the same hyperparameters as in Chen et al.
Learning the prior.
In order to learn the prior, we use the SWAG algorithm [33] with cycliclearning rate presented in Zhang et al.
We use SGD with a Nesterov momentum parameterof 0 .
9 for optimization.
We select our initial learning rate from a logarithmic grid between 0 . 005and 0 .
5 and evaluate 10000 iterations for each cycle.
Other hyperparameters are from Chen et al.[6], including data augmentation which comprises color augmentation, blurring, random crops, andhorizontal flips.
Downstream classification tasks.
We train for 30 , 000 steps with a batch size of 128 on CIFAR-10and CIFAR-100, 16 for Oxford Flowers-102, and 32 for Oxford-IIIT Pets.
SGD and SGHMC areused with momentum parameter of 0 .
During fine-tuning, we perform random crops with resizingto 224 × 224 and horizontal flips.
At test time, we resize the images to 256 pixels along the shorterside and produce a 224 × 224 center crop.
In our study, we select the learning rate and weight decayfrom a grid of 7 logarithmically spaced learning rates between 0 . 0001 and 0 . 1 , and 7 logarithmicallyspaced weight decay values between 1 e − 6 and 1 e − 2 , as well as without weight decay.
These weightdecay values are divided by the learning rate.
For the SGHMC optimizer, the temperature is selectedfrom a logarithmic grid of 8 values between 1 e 0 and 1 e − 8 , and we use predictions from 5 differentchains for the final evaluation.
B.2 Supervised Pre-Training
As mentioned above, we conduct additional experiments with priors learned on labeled data.
Figures7 and 6 contain downstream task performance comparisons for Resnet-50 and Resnet-101 modelswith priors learned starting with torchvision checkpoints pre-trained on ImageNet 1k.
Thesefigures compare our method to SGD transfer learning with pre-trained initializations and Bayesianinference with non-learned Gaussian prior.
Across both backbones, our method outperforms allbaselines for nearly every number of train samples and for all four datasets, where the Bayesianinference model with a non-learned prior is always worse.
(For full results see tables 2, 3, 4 and 5. [Table] [Table] [Table]
[Figure]
Model/ # samples
[Figure]
BNN Learned {+Prior 98 .84 .
+}
{+BNN Non-learned Prior 93 .93 .
+}
{+SGD Ensemble 96 .86 .72 .
+}
{+SGD Transfer Init 
+}
B.3 Self-Supervised Pre-Training
In this section, we present additional data to Figures 3 and 4 found in the main body.
The tables 10,11, 12 and 13 are corresponding to Figure 3, while Table 14 corresponding to Figure 4.
B.4 Evaluating uncertainty
In Figure 8, we present the negative log test-likelihood (NLL) for both the Bayesian methods and the
SGD-based methods for 4 datasets (CIFAR-10, CIFAR-100, Oxford Flowers-102 and Oxford-IIIT
Pets).
In general, we can see a similar trend for all the datasets where the Bayesian transfer learningoutperforms all other methods.
In Figure 9 we present the reliability diagrams for CIFAR-10 and
CIFAR-100 for ResNet-18 network.
A perfectly calibrated network has no difference betweenaccuracy and confidence, represented by a dashed black line.
Under-confident predictions are thosebelow this line, whereas overconfident predictions are those above.
We can see that Bayesianinference with learned priors are the best calibrated among the methods we consider.
The Error barsare computed on 5 runs, and the radius is one standard error.
[Figure]
Model
[Table]
[Table]
C Semantic Segmentation
C.1 Implementation Details
For our semantic segmentation evaluations, we use the DeepLabv3+ [4] model with ResNet-50 and
ResNet-101 backbone architectures.
We train our models using Pascal VOC 2012 and Cityscapestrainsets and evaluate on their val sets.
We utilize the same hyperparameters employed in Chen et al.[4] with a few modifications: our learning rate schedule is the poly policy [3] with initial value 0 .
01for Pascal VOC 2012 and 0 .
1 for Cityscapes, our crop size is 513 × 513 for Pascal VOC 2012 and768 × 768 for Cityscapes.
Our output stride is 16 for both training and evaluating, and we performrandom scale data augmentation during training.
In our SGD and SGLD optimization, we employ
Nesterov momentum, and we set the momentum coefficient to 0 .
We use batches of 16 images andweight decay with a coefficient of 10 − 4 .
In addition, for the SGLD optimizer, the temperature is setto 2 × 10 − 8 .
We train all models for 30k iterations.
C.2 MAP Estimation
We also compare MAP estimates (SGD) on the loss induced by our learned priors to the baseline withonly weight decay.
Table 15 and 16 contain the results using ResNet-50 and ResNet-101 backbones, [Figure] [Table] [Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
respectively.
On both architectures and both datasets, learned priors boost performance on theseexperiments without using Bayesian inference at all.
D Alternative Methods for Re-Scaling the Prior
As discussed previously, we tune a coefficient to re-scale our learned prior by maximizing thevalidation log-likelihood over a grid of coefficient values.
Our objective in this section is to examinethree alternative methods for tuning coefficients to re-scale the prior.
First, we consider maximizingthe Laplace approximation to the marginal likelihood, which does not require holding out validationdata [22].
Marginal likelihood maximization can be performed online together with MAP estimationin deep learning.
Specifically, we try tuning a single coefficient and
also per-layer coefficients forre-scaling the prior’s covariance matrix in this fashion.
A third alternative is to instead tune twocoefficients: one for the diagonal component of the covariance matrix, and one for the low-rankcomponent.
We use a grid-search in this case to tune the two coefficients.
Table 17 contains theresults for these three different methods as well as the method advocated in the main body of our
[Table]
[Table] [Table] [Figure]
draft using the SimCLR Resnet-50 backbone and SGHMC for downstream fine-tuning.
We observethat on most datasets, tuning a single covariance re-scaling hyperparameter using marginal likelihooddoes not improve the results reported in the paper.
Using different coefficients for each layer resultsin slight improvements on some datasets.
On CIFAR-10, Oxford Flowers-102, and Oxford-IIIT
Pets, test errors decreased by 0 .
03% − 0 .
Finally, on most datasets, re-scaling the diagonal andlow-rank components of the covariance matrix independently results in slightly better performance.
For example, the test errors on CIFAR-100, Oxford Flowers-102, and Oxford-IIIT Pets decreaseby 0 .
03% − 0 .
We note that these alternative methods bear an additional computational andimplementational cost.
E Limitations
Limitations.
Bayesian model averaging provides particularly compelling results, especially in termsof calibration and data efficiency, but does incur some additional computation.
On the other hand, our
SGLD based sampling procedures outperform deep ensembles, which have comparable computationalcosts.
Moreover, our informative priors with optimization still provide a clear performance boostover standard transfer learning, with negligible overhead.
In terms of applications, we have limitedour considerations to vision settings, for a focused exposition.
Directions for future research.
Pre-training has been particularly transformative in natural languageprocessing, where informative priors could be used to great effect.
More broadly studying loss-surfacealignment between tasks could also be informative for understanding how to build models that providegood domain generalization and robustness to spurious correlations.
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]