Transfer Learning in Infinite Width Feature Learning Networks

TL;DR

This paper develops a theoretical framework for transfer learning in infinitely wide neural networks, analyzing how pretraining influences generalization in various feature learning regimes and validating findings on multiple datasets.

Contribution

It introduces a comprehensive theory of transfer learning in infinite-width networks, encompassing fine-tuning and joint feature learning, with interpretable performance insights.

Findings

Pretraining improves generalization depending on data and task alignment.

Adaptive kernels depend on source data, labels, and target data after pretraining.

Theoretical predictions are validated on linear, polynomial, and real datasets.

Abstract

We develop a theory of transfer learning in infinitely wide neural networks under gradient flow that quantifies when pretraining on a source task improves generalization on a target task. We analyze both (i) fine-tuning, when the downstream predictor is trained on top of source-induced features and (ii) a jointly rich setting, where both pretraining and downstream tasks can operate in a feature learning regime, but the downstream model is initialized with the features obtained after pre-training. In this setup, the summary statistics of randomly initialized networks after a rich pre-training are adaptive kernels which depend on both source data and labels. For (i), we analyze the performance of a readout for different pretraining data regimes. For (ii), the summary statistics after learning the target task are still adaptive kernels with features from both source and target tasks. We test our theory on linear and polynomial regression tasks as well as real datasets. Our theory allows interpretable conclusions on performance, which depend on the amount of data on both tasks, the alignment between tasks, and the feature learning strength.