这篇论文为预训练表征的维度选择提供了理论指导,做迁移学习或模型压缩的研究者可以直接参考其结论来优化训练策略。
该论文通过高维分析框架,研究了预训练表征维度对下游线性探测泛化性能的影响。作者将预训练建模为无标签数据的主成分分析,下游任务建模为有标签数据的线性回归,并推导出训练误差和泛化误差的精确表达式。研究发现,当预训练数据充足但下游数据稀缺时,最大压缩的表征最优;而当预训练数据有限时,高维表征泛化更好。论文还量化了无标签数据与有标签数据之间的替代关系,并在自编码器和预训练大语言模型中观察到类似现象。
Optimal Representation Size: High-Dimensional Analysis of Pretraining and Linear Probing
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.