🤖 machine learning

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

This paper presents an analytical high-dimensional model of pretraining and linear probing that derives exact generalization error expressions to identify optimal representation sizes, revealing that maximally compressed representations are best when pretraining data is abundant but downstream data is scarce, while higher-dimensional representations perform better with limited pretraining data.

Original authors: Valentina Njaradi, Clémentine Dominé, Rachel Swanson, Marco Mondelli, Andrew Saxe

Published 2026-05-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Valentina Njaradi, Clémentine Dominé, Rachel Swanson, Marco Mondelli, Andrew Saxe

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to learn a new skill, like playing a specific song on the piano. You have two types of resources:

A massive library of sheet music (unlabelled data) that you can read through to understand how music generally works, but you don't know which notes make up your specific song yet.
A few practice sessions (labelled data) where a teacher tells you exactly which notes to play for your song.

This paper explores a modern strategy used by both computers and, potentially, brains: Pretraining followed by Fine-tuning. First, you study the library to build a general understanding of music (Pretraining). Then, you use that understanding to quickly learn your specific song with very few practice sessions (Fine-tuning).

The big question the authors ask is: How much of that general library should you actually keep in your head?

Should you try to memorize every single note and rule from the library (a huge, complex memory)? Or should you boil it down to just the most essential chords and rhythms (a compressed, simple memory)?

The Core Discovery: It Depends on How Much Practice You Have

The authors built a mathematical model to answer this. They found that the "perfect" amount of memory to keep depends entirely on the balance between your library size and your practice time.

Scenario A: You have a giant library, but very little practice time.

The Result: You should compress your memory.
The Analogy: Imagine you have read 10,000 books, but you only have 10 minutes to prepare for a quiz. If you try to remember every detail from every book, your brain will get confused and mix things up. Instead, it's better to distill those 10,000 books down to the "Top 10 Rules" that apply to almost everything. This "compressed" version is robust and prevents you from making mistakes when you don't have time to think deeply.
The Paper's Claim: When pretraining data is abundant but downstream (task) data is scarce, maximally compressed representations are optimal.

Scenario B: You have a small library, but lots of practice time.

The Result: You should keep more details (higher dimensions).
The Analogy: Imagine you only read 5 books, but you have 50 hours to practice. If you try to boil those 5 books down to just "Top 10 Rules," you might throw away the specific details you actually need. Since you have plenty of time to practice, you can afford to keep a more detailed, high-dimensional memory of those 5 books to get the nuances right.
The Paper's Claim: When pretraining data is limited, higher-dimensional representations generalize better.

The "Leak" Problem: Why Compression Helps

The paper explains a tricky phenomenon called "leakage."

Imagine your library has a hidden pattern (a "spike")—maybe 90% of the books are about jazz, and only 10% are about rock.

If you keep everything (no compression), your brain tries to learn the jazz and the rock. But because you have so much jazz data, your brain gets "confused" and thinks the jazz rules apply to your rock song too. This is a "leak" of information that causes errors.
If you compress (keep only the top patterns), you force your brain to focus on the strongest, most reliable patterns (the jazz) and ignore the noise. This actually helps you perform better on the new task, even if the task isn't perfectly jazz-related, because it stops the confusion.

The "Currency" of Data

The authors also calculated a fascinating trade-off: How much unlabelled data is worth one labelled sample?

They found that in certain situations (lots of library, little practice), adding more pages to your library is actually more valuable than getting a single extra minute of practice with a teacher. The "unlabelled" data acts as a powerful substitute for "labelled" data, but only if you know how to compress it correctly.

Real-World Checks

The authors didn't just stop at math. They tested these ideas on:

Autoencoders: Simple neural networks designed to compress data. They found that when these networks had enough data to learn, they naturally started acting like the "compression" the math predicted.
Large Language Models (LLMs): They looked at real AI models (like Pythia). When they took the AI's "brain" (its internal representations) and tried to use it for a new task (sentiment analysis), they found that pruning (removing) some of the AI's internal dimensions actually made it better at the new task, but only when the new task had very little data.

Summary in One Sentence

If you have a massive amount of general knowledge but very little specific practice, the smartest thing to do is to simplify and compress your knowledge to avoid confusion; but if you have limited general knowledge and lots of practice time, you should keep the details to make the most of your training.

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

Problem Statement
Modern machine learning pipelines increasingly rely on a two-stage paradigm: pretraining on abundant unlabelled data to extract reusable structure, followed by fine-tuning or linear probing on limited labelled data for specific tasks. While this approach has yielded significant sample efficiency, the theoretical understanding of how the dimensionality of the pretrained representation influences downstream generalisation remains incomplete. Specifically, it is unclear whether optimal performance requires maximally compressed (low-dimensional) representations or high-dimensional ones, and how this choice depends on the interplay between the volume of pretraining data, the scarcity of downstream labels, task alignment, and signal-to-noise ratios.

Methodology
The authors propose an analytically tractable theoretical model to study this two-stage learning process in the high-dimensional regime (where the number of features $p$ , unlabelled samples $n_u$ , and labelled samples $n_l$ grow proportionally).

Pretraining Stage: The model learns representations via Principal Component Analysis (PCA) on an unlabelled dataset $X_u$ drawn from a spiked covariance model $\Sigma = I_p + (\lambda - 1)vv^\top$ . The model retains the top $m$ principal components, defining a projection matrix $P_m$ .
Downstream Stage: A separate labelled dataset $X_l, y_l$ is used for supervised learning. The task is modeled as linear regression $y_l = X_l w^* + \xi$ , where $w^*$ is the ground-truth signal and $\xi$ is Gaussian noise. The regression is performed on the projected inputs $X_l P_m$ .
Analysis: The authors derive exact asymptotic expressions for estimation error, generalisation error, and training error as $p, n_u, n_l, m \to \infty$ . They analyze the dependence of these errors on the representation fraction $\alpha = m/p$ , the aspect ratios $\gamma_u = p/n_u$ and $\gamma_l = p/n_l$ , the spike strength $\lambda$ , and the alignment $\eta$ between the task signal $w^*$ and the population spike $v$ .

Key Contributions

Exact Asymptotic Expressions: The paper derives closed-form expressions for estimation, generalisation, and training errors (Theorems 1, 2, and 3). These expressions decompose error into "missing signal" (unrecoverable due to finite samples), "leaked signal" (bias induced by the coupling of orthogonal subspaces via the population spike), and variance terms.
Characterization of Optimal Representation Size: The authors identify distinct regimes for the optimal dimensionality $\alpha^*$ $α^{*}$ :
- Scarce Pretraining, Scarce Downstream: When pretraining data is insufficient to reliably estimate principal components, high-dimensional representations ( $\alpha \approx 1$ ) are optimal.
- Abundant Pretraining, Scarce Downstream: When pretraining data is abundant but downstream labels are limited, maximally compressed representations ( $\alpha < 1$ ) are optimal. Compression acts as a regularizer, mitigating the variance caused by the double-descent phenomenon in the downstream task.
- Abundant Downstream: When downstream labels are plentiful, little compression is needed ( $\alpha \approx 1$ ).
Phase Transitions: The study identifies precise phase transition boundaries (Corollaries 1 and 2) where the optimal strategy shifts between compression and full-rank representations. These boundaries depend non-trivially on the signal-to-noise ratio (SNR), spike strength $\lambda$ , and task alignment $\eta$ .
Trade-off Quantification: The authors quantify the marginal rate of substitution between unlabelled and labelled data, showing regimes where additional unlabelled data reduces generalisation error more effectively than additional labelled samples.

Results

Theoretical Predictions vs. Simulations: The derived asymptotic formulas show excellent agreement with numerical simulations across various parameter settings, including variations in SNR, spike strength, and sample sizes.
Error Decomposition: The analysis reveals that in the presence of a population spike, components of the signal orthogonal to the retained subspace can "leak" into the learned subspace, creating a bias term. This leak is a non-trivial effect of the spiked covariance structure that distinguishes this pretraining setting from standard Principal Component Regression (PCR).
Validation in Realistic Models:
- Autoencoders: Experiments with linear and nonlinear autoencoders trained on spiked covariance data confirm that the theoretical phase structure persists. Linear autoencoders closely match PCA behavior, while nonlinear autoencoders exhibit deviations only when pretraining data is scarce (where they tend to memorize rather than learn the subspace).
- Pretrained Transformers: Linear probing experiments on Pythia language models demonstrate that representations learned during pretraining exhibit a spiked covariance structure. Tuning the representation dimensionality via PCA on these representations improves downstream performance on the SST-2 sentiment task, particularly for small downstream datasets and models trained for longer durations, aligning with the theoretical predictions.

Significance
The paper provides a principled theoretical framework explaining why and when compression during pretraining improves downstream generalisation. It challenges the intuition that more dimensions are always better, showing instead that optimal dimensionality is a function of data availability and task structure. Specifically, it highlights that in modern pipelines where pretraining data is abundant but downstream data is limited, aggressive compression is theoretically optimal. Furthermore, the work establishes a quantitative trade-off between unsupervised pretraining and supervised learning, offering insights into the sample efficiency of transfer learning. The findings bridge the gap between idealized random matrix theory and empirical observations in deep learning, suggesting that the low intrinsic dimensionality observed in large language models is not merely a byproduct but a feature that can be leveraged for optimal adaptation.

The Core Discovery: It Depends on How Much Practice You Have

The "Leak" Problem: Why Compression Helps

The "Currency" of Data

Real-World Checks

Summary in One Sentence

More like this