Optimal Representation Size: High-Dimensional Analysis of Pretraining and Linear Probing

TL;DR

This paper develops an analytical model for understanding how the size of learned representations affects generalisation in pretraining and linear probing, revealing optimal strategies depending on data availability.

Contribution

It provides a high-dimensional analytical framework linking representation size, data quantities, and task alignment, guiding optimal pretraining and probing strategies.

Findings

Maximally compressed representations are optimal with abundant pretraining data and scarce downstream data.

Higher-dimensional representations generalise better when pretraining data is limited.

An exact trade-off quantifies unlabelled data needed to replace labelled samples.

Abstract

Learning to generalise from limited data is a fundamental challenge for both artificial and biological systems. A common strategy is to extract reusable structure from abundant unlabelled data, enabling efficient adaptation to new tasks from limited labelled data. This two-stage paradigm is now standard in modern training pipelines, where pretraining is followed by fine-tuning or linear probing. We provide an analytical model of this process: structure extraction is formalized as principal component analysis on unlabelled data, and downstream learning as linear regression on a separate labelled dataset. In the high-dimensional regime, we derive exact expressions for training and generalisation error showcasing their dependence on representation dimensionality, unlabelled and labelled sample sizes, and task alignment. Our results show that pretrained representations strongly influence downstream generalisation, and we characterize the optimal representation size as a function of task parameters: with abundant pretraining data but scarce downstream data, maximally compressed representations are optimal, whereas with limited pretraining data, higher-dimensional representations generalise better. Furthermore, we establish an exact trade-off between pretraining and supervision, quantifying how much unlabelled data is required to replace a single labelled sample. Beyond our idealised model, we observe similar phenomenology in autoencoders and pretrained LLMs. Altogether, we highlight that optimising representation size is critical, giving conditions for when compression during pretraining improves generalisation.