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 [41, 23, 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 [41].
Thereare also specialized priors, including heavy-tailed priors [15, 23], noise-contrastive priors for highuncertainty under distribution shift [16], and input-dependent priors for domain generalization [22].
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, 39, 4, 19, 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 [18], 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 [27].
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 [27].
In our experiments, we find that MAP estimation, diagonalcovariance, and supervised pre-training are all suboptimal in the transfer learning setting.
Leveraging Auxiliary Data in Bayesian Modeling.
A number of works outside of deep learninghave considered knowledge transfer in Bayesian modeling, especially in settings such as domainadaptation or homogeneous transfer learning in which source and target tasks contain similar featureand label spaces but differ in their marginal distributions.
For example, Xuan et al.
[42] considersmethods for probabilistic graphical models, and Karbalayghareh et al.
[26] develop a theoreticalframework for understanding optimal Bayes classifiers given prior knowledge from other domains.
Other works develop Bayesian tools 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[32, 13, 25], or develop kernels based on neural networks trained on previous tasks for Gaussianprocess inference [30, 34].
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 [37], by perturbing weights and using consistencyregularization [11], or via semi-supervised deep kernel learning [24].
Deep kernel learning has alsobeen used for Bayesian meta-learning in few-shot classification and regression [36].
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 [31] (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 [40, 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 ofthese 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, 17], and we choose CIFAR-10 and CIFAR-100 fordownstream tasks [28].
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) [31] 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 [31].
The SWAG approximate posterior distribution is givenby N ( w, 1 2 Σ diag + 1 2 Σ low-rank ) , where w = 1 M (cid:80) M t
=1 w t is the SWA [21] 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 [29, 20]), 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 with an 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) [27] 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 the 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. [Figure]
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 constrainourselves 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)
[40] 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.
Experiments
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.
[Figure]
4.1 Experimental Setting
We adopt the ResNet-50 and ResNet-101 architectures [17], 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 [28], Oxford Flowers-102 [33], and Oxford-IIIT Pets [35].
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 amount of 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.
Intransfer learning, deep ensemble members are initialized with the same value from pre-training on thesource task, and fine-tuning tends to stay in the same basin, such that the ensemble components arerelatively homogenous in their predictions.
Recall that SGD with our learned priors incurs essentiallythe same computational 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 GeneralizationCIFAR
-10.1 [38] 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. [Figure]
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 assign a zero-mean isotropic Gaussian prior to the decoder and atrousconvolution parameters of DeepLabv3+ which are not pre-trained.
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 learnedpriors boost the performance of MAP estimates as well, and thus are preferable to standard transferlearning, even for practitioners who do not intend to perform Bayesian inference.
[Table]
) and a non-learned prior over ResNet-50 and ResNet101 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 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.
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.
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]
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] We include additional 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 they were 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 experi- ments 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., type of GPUs, internal cluster, or cloud provider)?
[Yes] We estimate compute in section A in 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, ifapplicable?
[N A]
) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
[-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 
-]
[-BNN Learned Prior 98 .84 .
-]
[-BNN Non-learned Prior 93 .93 .
-]
[-SGD Ensemble 96 .86 .72 .
-]
[-SGD Transfer Init 
-]