Research map / Architectures and efficiency / RNNs and normalization

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

NeurIPS 2026 · Architectures and efficiency › RNNs and normalization · approach: Layer-wise analysis of language models · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Analytical high-dimensional model of two-stage learning (PCA on unlabelled spiked-covariance data, then linear regression on projected labelled data) giving exact error asymptotics, optimal representation size phase diagrams, and a pretraining-vs-supervision trade-off, validated on autoencoders and Pythia linear probing.

Problem

How the dimensionality of a representation learned by unsupervised pretraining (PCA) affects downstream linear regression generalisation, and how unlabelled and labelled sample sizes trade off, in the high-dimensional regime.

Approach

Separates the PCA sample from the regression sample in a spiked covariance model, derives exact deterministic limits for estimation, generalisation and training error as functions of representation fraction alpha, and identifies phase transitions where compression becomes optimal and the marginal rate of substitution between unlabelled and labelled data.

Key results

• Exact asymptotic estimation, generalisation and training error formulas derived for PCA-then-regression with spiked covariance, matching simulations.
• Phase transition curves for optimal representation size alpha* are derived for gamma_l>1 and gamma_l<1, matching heatmaps of simulated optimum.
• In Pythia linear probing on SST-2, longer-pretrained checkpoints benefit from stronger compression on small downstream datasets.

Surprising findings

• For small downstream datasets (around 100 samples) the optimal dimensionality is high again, consistent with weak task-pretraining alignment.

Limitations

• Relies on shared spike between pretraining and downstream covariance and Gaussian data.
• Analysis restricted to ridge-free least squares regression and PCA-based projections.
• Autoencoders recover PCA only with sufficient pretraining data; nonlinear autoencoders deviate at small n_u.

Experimental setup

Datasets: synthetic spiked Gaussian data The Pile SST-2 (GLUE)

Benchmarks: SST-2

Metrics: estimation error generalisation error training error validation accuracy

Baselines: principal component regression full-rank linear regression alpha approx 0 single PC

Builds on: principal component regression double descent Marchenko-Pastur law BBP phase transition

Closest papers by meaning

• On the Provable Advantage of Unsupervised Pretraining (ICLR 2024 · similarity 0.87)
• Feature Reuse and Scaling: Understanding Transfer Learning with Protein Language Models (ICML 2024 · similarity 0.86)
• Representations Shape Weak-to-Strong Generalization: Theoretical Insights and Empirical Predictions (ICML 2025 · similarity 0.86)
• Bridging Compute- and Data-Optimal Pretraining (NeurIPS 2026 · similarity 0.86)
• Understanding and Enhancing Mask-Based Pretraining towards Universal Representations (NeurIPS 2025 · similarity 0.86)
• On Linear Representations and Pretraining Data Frequency in Language Models (ICLR 2025 · similarity 0.86)