Fine-Tuning with Uncertainty-Aware Priors Makes Vision and Language Foundation Models More Reliable

Fine-tuning off-the-shelf pre-trained neural networks has become the default starting point for a wide range of challenging prediction tasks---especially in computer vision and natural language processing, where pre-trained models trained on millions or even billions of data points are publicly available and can be fine-tuned with a moderate compute budget. However, while fine-tuned models have been shown to significantly improve predictive performance in several respects compared to models trained from scratch, they can exhibit poor calibration and fail to reliably identify challenging distribution shifts. In this paper, we improve uncertainty quantification in fine-tuned models by constructing an uncertainty-aware fine-tuning prior and deriving a tractable variational objective for inference. The prior assigns high probability density to parameters that induce predictive functions with high uncertainty on data points that are meaningfully different from the data used for fine-tuning. We evaluate models trained with this prior on different transfer learning tasks and show that fine-tuning with uncertainty-aware priors significantly improves calibration, selective prediction, and semantic shift detection on computer vision and natural language classification tasks.