A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints

TL;DR

This paper proposes a Bayesian model selection criterion called downstream free energy to identify pretraining checkpoints that are most adaptable for downstream tasks, without requiring access to downstream data.

Contribution

It introduces a novel Bayesian criterion for selecting pretraining checkpoints that correlates with downstream performance and does not need downstream data or prior task knowledge.

Findings

The criterion reliably predicts better finetuning outcomes.

It correlates well with downstream task adaptability.

The method is effective across different models and tasks.

Abstract

Recent advances in artificial intelligence have been fueled by the development of foundation models such as BERT, GPT, T5, and Vision Transformers. These models are first pretrained on vast and diverse datasets and then adapted to specific downstream tasks, often with significantly less data. However, the mechanisms behind the success of this ubiquitous pretrain-then-adapt paradigm remain underexplored, particularly the characteristics of pretraining checkpoints that enhance downstream adaptation. We introduce a Bayesian model selection criterion, called the downstream free energy, which quantifies a checkpoint's adaptability by measuring the concentration of nearby favorable parameters for the downstream task. We demonstrate that this Bayesian model selection criterion can be effectively implemented without access to the downstream data or prior knowledge of the downstream task. Furthermore, we provide empirical evidence that the criterion reliably correlates with improved finetuning performance, offering a principled approach to predicting model adaptability.