A Theory of How Pretraining Shapes Inductive Bias in Fine-Tuning

Pretraining and fine-tuning are central stages in modern machine learning systems. In practice, feature learning plays an important role across both stages: deep neural networks learn a broad range of useful features during pretraining and further refine those features during fine-tuning. However, an end-to-end theoretical understanding of how choices of initialization impact the ability to reuse and refine features during fine-tuning has remained elusive. Here we develop an analytical theory of the pretraining–fine-tuning pipeline in diagonal linear networks, deriving exact expressions for the generalization error as a function of initialization parameters and task statistics. We find that different initialization choices place the network into four distinct fine-tuning regimes that are distinguished by their ability to support feature learning and reuse—and therefore by the task statistics for which they are beneficial. In particular, a smaller initialization scale in earlier layers enables the network to both reuse and refine its features, leading to superior generalization on fine-tuning tasks that rely on a subset of pretraining features. We demonstrate empirically that the same initialization parameters impact generalization in ResNets trained on CIFAR-100 and SVHN as well as Transformers trained on modular arithmetic tasks. Overall, our results demonstrate analytically how data and network initialization interact to shape fine-tuning generalization, highlighting an important role for the relative scale of initialization across different layers in enabling continued feature learning during fine-tuning.

Machine learning algorithms are commonly trained on large general-purpose datasets (for example images paired with categorical labels). This “pretraining” step imbues these models with useful knowledge that allows them to learn related tasks from fewer examples (this step is called “fine-tuning”). This approach is particularly useful in cases where data on the fine-tuning task is limited and it is therefore important to learn from it as efficiently as possible. We aimed to understand how different ways of sampling the models’ initial weights before pretraining affects how well they can learn the fine-tuning task. By considering a simplified neural network model and a simplified set of tasks, we were able to mathematically characterize the model’s performance during fine-tuning. We found that depending on the relationship between the pretraining and fine-tuning task, different kinds of initialization are optimal. Moreover, our theory clarifies how to choose between these different initializations depending on this relationship. Building such a theoretical foundation can help us improve pretraining–fine-tuning pipelines in practice.