Conceptio › Archive › arXiv CS

arXiv CSopen access

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

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

neural-networks

machine learning, deep learning, neural networks

O PTIMAL R EPRESENTATION S IZE : H IGH -D IMENSIONAL A NALYSIS OF P RETRAINING AND L INEAR P ROBING A P REPRINT

arXiv:2605.20105v1 [cs.LG] 19 May 2026

Valentina Njaradi Gatsby Computational Neuroscience Unit University College London [email protected] Rachel Swanson Sainsbury Wellcome Centre University College London

Clémentine Dominé Institute of Science and Technology Austria

Marco Mondelli† Institute of Science and Technology Austria

Andrew Saxe† Gatsby Computational Neuroscience Unit Sainsbury Wellcome Centre University College London

A BSTRACT 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.

Introduction

Large language models, vision transformers, and related architectures have demonstrated language, reasoning and mathematical capabilities widely considered out of reach just a decade ago [1–4]. This progress has been made possible by multi-stage training pipelines capable of leveraging massive amounts of unlabelled data: models are first pretrained on a large corpus of unlabelled data, then fine-tuned on labelled tasks and finally refined, e.g., via reinforcement learning from human feedback [5, 6], low-rank adaptation [7] or test-time training [8–10]. Such a paradigm unlocks remarkable sample efficiency, as large pretrained models generalise on new tasks where they have access only to thousands (or even hundreds) of samples [11, 12]. Co-senior authors marked with † .

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

At their core, these training pipelines rely on a key fact: pretraining on large datasets yields rich representations that effectively transfer to downstream tasks. This idea is not new, as greedy layer-wise training and stacked autoencoders were designed to learn reusable representations in pretraining a decade before modern transformers [13–15]. However, different pretraining objectives emphasize different aspects of the data, and it remains unclear which features of a representation are most useful for downstream learning. Characterizing these representations, how they support computation, transfer to new tasks and enable adaptive behaviour in artificial and biological systems, has been a longstanding focus in machine learning, as well as cognitive science and neuroscience [16–18]. In machine learning, structured representations and prior knowledge can accelerate learning and enhance generalisation, but may also hinder adaptation, as exemplified by catastrophic interference in continual and reversal learning [19–22]. Similar trade-offs arise in biological systems, where representations must support flexible and efficient adaptation [23]. A crucial property underlying these effects is the size (or dimensionality) of representations. Both high and low-dimensional representations have been observed across a range of brain regions and tasks [24–27] as well as in machine learning models [11, 28], yet a principled account of how dimensionality interacts with task and environment structure to shape learning performance and sample efficiency remains missing. Here, we study a two-stage learning setting in which representations are first learned via principal component analysis (PCA) from unlabelled samples drawn from a spiked covariance model, and then reused by linear regression on an independent labelled dataset. The covariance structure plays a central role: it produces both a tractable model for analysis and controlled alignment between the pretraining and supervised training stages, thus enabling us to isolate how pretrained representations influence downstream generalisation. More precisely, our contributions are as follows: • We derive exact asymptotic expressions for estimation, generalisation and training error (Theorem 1, 2 and 3, respectively). From this, we establish the improvement in generalisation coming from tuning the representation size, as well as the improvement due to adding more pretraining and finetuning data. For optimally-tuned errors, we identify regions where unlabelled samples are more valuable than labelled samples (Figure 2). • We characterize when either (i) low-dimensional representations are useful for generalisation on downstream tasks, or (ii) tasks benefit from detailed high-dimensional representations (Figure 3). In particular, we identify phase transitions where compression becomes useful (Corollaries 1, 2), and describe their behaviour for an infinite number of pretraining samples (Corollary 4). • We finally show that the theoretical predictions derived in our idealized setting extend to trained autoencoders when the number of samples is larger than the number of features, and further demonstrate that the insights unveiled by our theory persist in more realistic settings through linear probing of pretrained transformers (Figure 4).

Related work

High-dimensional regression. High-dimensional regression, in the regime where number of features and samples grow proportionally, has been the subject of intense investigation, leading to the precise characterization of the test error [29–32] and phenomena such as benign overfitting [33] and double descent [34, 35]. This line of work further characterizes the distribution of the empirical risk minimiser [36], the impact of spurious correlations [37], performative learning [38], as well as training on multiple or surrogate data sources [39–42]. Principal component regression (PCR), where regression is performed on principal components (PCs) of the sample covariance, provides a natural way to restrict regression to a low-dimensional subspace. Its performance depends critically on the number of retained components, which governs the bias–variance trade-off [43–46], and has been characterized both non-asymptotically [47] and in proportional asymptotics for population and sample PCs [44, 45, 48]. In contrast, we study pretrained regression (PR), with PCs formed during a pretraining phase and reused for a downstream task, separating the sample sets for PCA and subsequent regression. As a result, both the asymptotic risk and the optimal number of principal components for PR differ fundamentally from the classical PCR setting. Pretraining and transfer learning. More broadly, the reuse of pretrained representations across tasks has been extensively studied in transfer learning, primarily in supervised settings where generalisation depends on task similarity or alignment with pretrained features [49–55]. These approaches typically do not model representations learned from unlabelled data. Complementary work establishes benefits of unsupervised and self-supervised pretraining [56–59]. Closest to our paper, recent work analyses PCA-based pretraining in spiked models [60], but considers fixed representations without addressing the role of representation dimensionality in generalisation. Representation dimensionality. A growing body of work highlights that pretrained representations often lie in low-dimensional subspaces, as evidenced by studies of intrinsic dimensionality [11, 28], low-rank adaptation [7], and linear probing [61]. Recent theoretical work has begun to formalise aspects of this structure [62, 63]. While these papers study the structure of pretrained features and how they support downstream prediction, they do not provide a 2

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

closed-form characterisation of how the optimal retained dimension depends jointly on pretraining, downstream data, and task structure. Addressing this question through an analytically tractable theory is the focus of our work. Spiked covariance models. The spiked covariance model [64] provides a tractable framework for studying lowdimensional structure in high-dimensional data. A key feature is the BBP phase transition [65–67]: above this threshold, empirical principal components align with the population spike, while below it they are asymptotically uninformative [65, 68–70]. This transition makes the spiked models well-suited for analysing how representation quality depends on sample size and signal strength. Accordingly, they have been used to analyse kernel methods, neural networks and feature learning [71–73], as well as the emergence of low-rank structure during training [74–77]. Finally, spiked covariances have been used to analyse PCR, showing how the projection onto empirical PCs has a regularising effect [78]. Our setting differs from these works: the signal subspace is estimated from an independent unlabelled pretraining sample and then used for a separate labelled task, allowing us to characterise the optimal representation dimension.

Setup

We study two-stage learning with unlabelled pretraining followed by supervised learning: 1. In the pretraining stage, we observe inputs Xu ∈ Rnu ×p whose rows are i.i.d. N (0, Σ), forming the unlabelled dataset Du . From this data, the model extracts a representation via PCA: it computes the top m principal components of Xu , stacks them into Um ∈ Rp×m , and defines a projection of retained directions Pm = Um U⊤ m with a discarded complement P⊥ = I − Pm = U⊥ U⊤ . If m > n , the PCA subspace is completed with uniformly u ⊥ random orthonormal vectors. 2. In the downstream stage, we observe a labelled dataset Dl = {Xl , yl }, where yl = Xl w⋆ + ξ ∈ Rnl , the rows of Xl ∈ Rnl ×p are i.i.d. N (0, Σ) and independent of Xu , w⋆ ∈ Rp is the ground-truth signal and ξ ∈ Rnl is Gaussian noise with i.i.d. entries of variance σ 2 . The pretrained regression model uses the representation learned from Du to fit the downstream task by regressing on the projected inputs (see Figure 1(a) for a schematic representation): ŵ = (Xl Pm )† yl = (Xl Pm )† Xl w⋆ + (Xl Pm )† ξ,

(1)

†

where (·) denotes the pseudoinverse. We denote by Π the projection onto the row space of Xl Pm , i.e., the subspace spanned by the projected training inputs. This setup follows [45], with the key difference that the projection subspace is learned from independent unlabelled data rather than the downstream dataset itself. The quantities of interest are estimation, generalisation and training error: 1 E est = Eξ ∥w⋆ − ŵ∥2 , E gen = Ex,y,ξ (y − x⊤ ŵ)2 , E train = Eξ ∥yl − Xl ŵ∥2 nl where the test point (x, y) in E gen is drawn from N (0, Σnew ). We work in the high-dimensional regime where p/nl → γl ∈ (0, ∞), p/nu → γu ∈ (0, ∞), m/p → α ∈ [0, 1] with α the fraction of dimensions to keep, and γu , γl the aspect ratios for the datasets Du , Dl . Due to the projection of the downstream inputs, we define the effective aspect ratio γeff := nml → αγl and assume γeff ̸= 1. The population covariance follows a rank-one spiked model Σ = Ip + (λ − 1)vv⊤ , where λ > 1 and the unit vector v ∈ Rp are fixed. The alignment of w⋆ with the spike eigenvector of Σ is captured by (w⋆ ⊤ v)2 /∥w⋆ ∥2 → η, and the mass spectrum of w⋆ on the eigenbasis of Σ is ηδλ + (1 − η)δ1 . The test population covariance follows a spiked PT ⊤ model with T spike directions {vnew,i }Ti=1 and eigenvalues {νi }Ti=1 , i.e., Σnew = Ip + i=1 (νi − 1)vnew,i vnew,i .

Theoretical analysis of two-stage learning

In this section, we derive asymptotic expressions for the estimation, generalisation, and training errors of the estimator in (1). We then characterise the phase transitions in the optimal representation dimensionality α, specifically, the regimes in which the optimum shifts from retaining all PCs to retaining only a few PCs, as a function of nu and nl . We ∗ 2 further study how nu , nl , λ and the signal-to-noise ratio SNR = ∥wσ2∥ affect both the errors and the optimal α, as well as the limiting behaviour with infinite pretraining data (γu = 0). Proofs of all results are deferred to Appendix A. 4.1

Asymptotic expressions

The following deterministic limits characterise how the learned projection Pm , its complement P⊥ , and the downstream effective row-space projector Π capture fixed signal directions. Under the high-dimensional scaling assumptions of 3

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

(a)

(b) Unlabelled Pretraining (PCA)

(c)

Downstream Supervised task

PCA m PCs

Figure 1: Analytically tractable model of two-stage learning. (a) In the pretraining stage, PCA extracts the top-m principal components from nu unlabelled samples. The downstream task with nl labelled samples is learned via regression on inputs projected through Pm = Um U⊤ m . Theoretically derived generalisation (b) and training errors (c) match numerical simulations. Section 3, for any deterministic vectors a, b with bounded normalized overlaps with the population spike directions, the bilinear projection terms satisfy a⊤ P⊥ b a.s. −−→ p̄⊥,a,b ∥a∥∥b∥

a⊤ Πb a.s. −−→ π̄a,b , ∥a∥∥b∥

(2)

with p̄⊥,a,b and π̄a,b given by the expressions in (34) and (41), involving integrals over limiting spectral distributions ⊤ of X⊤ u Xu /nu and (Xl Pm ) (Xl Pm )/nl (see Lemmas 5 and 7 in Appendix A.6). Theorem 1 (Asymptotic estimation error). As nl , nu , p, m → ∞ with p/nl → γl , p/nu → γu , m/p → α and γeff ̸= 1, the estimation error converges a.s. to    min{γeff , 1} ∗ 2  (λ − 1)2 2 est ⋆ ,w⋆ + + w̄ σ̄eff + σ 2 , E∞ = w̄∗  1 − π̄ p̄⊥,w⋆ ,v π̄v,v  w  | {z } 2 |γeff − 1| {z } | λ̄ {z } | missing signal leaked signal

∗

(3)

variance

⋆ 2

where w̄ is the deterministic limit of ∥w ∥ , λ̄ = 1 + (λ − 1)(1 − p̄⊥,v,v ) is the effective spike eigenvalue, and the 2 p̄2⊥,w⋆ ,v . variance from the unrecoverable part of w⋆ is σ̄eff = p̄⊥,w⋆ ,w⋆ + λ−1 λ̄ The expression in (3) admits a simple geometric interpretation: the estimator observes w⋆ through the subspace induced by Xl Pm , and the error reflects how w⋆ decomposes relative to this subspace. By decomposing into two components, aligned or orthogonal to Um , w⋆ = w⋆ ∥ + w⋆ ⊥ , the error splits into three contributions: (a) the missing signal, (b) the leaked signal, and (c) the variance term. (a) The term 1 − π̄w⋆ ,w⋆ captures the missing signal, i.e. directions that lie in the retained subspace but cannot be reliably estimated from finite samples. (b) The key non-trivial effect is the leak. In an isotropic model, the components Xl Pm and Xl P⊥ are independent, so w⋆ ⊥ is completely invisible to the estimator. Here, the spike in the covariance couples these two parts of the data, so directions outside the retained subspace become correlated with it. As a result, part of w⋆ ⊥ leaks into Π, 2 where it appears as a spurious signal aligned with v, producing the bias term (λ−1) p̄2⊥,w⋆ ,v π̄v,v . For a detailed λ̄2 discussion of parameter effects on the leak and a breakdown of different terms in the estimation error, see Figure 6 in Appendix B. (c) The remaining part of w⋆ ⊥ does not couple back into the learned subspace. Together with label noise, this contributes to the variance, which exhibits a double descent [34] controlled by γeff = αγl . Thus, even without explicit regularisation in the regression step, the model mitigates double descent by controlling the representation size α, see Figure 6 in Appendix B. Theorem 2 (Asymptotic generalisation error). Under the same regime, let ρw⋆ ,vnew,i be the deterministic limit of w⋆ ⊤ vnew,i . The generalisation error converges a.s. to ∥w⋆ ∥

2

 gen est E∞ = E∞ + w̄∗

T X   λ−1 2 (νi − 1)  p̄⊥,w⋆ ,v π̄v,vnew,i  w⋆ ,vnew,i − π̄w⋆ ,vnew,i − ρ  +σ . λ̄ | {z } {z } | i=1 missing signal

leaked signal

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

Figure 2: Benefits of optimal representation size. (a) Minimum achievable generalisation error with optimal α⋆ . (b) Gain in generalisation error from using α⋆ instead of all PCs (α = 1). (c) Same gain relative to α ≈ 0. (d) Marginal gen gen /∂nu ) / (∂E∞ /∂nl ); values above 1 (red) indicate rate of substitution between unlabelled and labelled data, (∂E∞ that unlabelled data reduces the error more than labelled data, and vice-versa below 1 (blue). (e) Generalisation-optimal representation size α⋆ . Simulation details are given in Appendix D. The calculation of the generalisation error builds on that of the estimation error, with an additional contribution measuring how the prediction error w⋆ − ŵ projects onto the spikes {vnew,i }Ti=1 in the test population covariance. Each direction contributes proportionally to its strength (νi − 1). For each vnew,i , the error depends on the mismatch between the true alignment ρw⋆ ,vnew,i and the alignment recovered by the estimator. As in the estimation error, this mismatch decomposes into a missing signal term and a leaked signal term, arising from finite-sample limitations and the coupling induced by the spike. The final σ 2 term reflects the irreducible error due to label noise. For additional discussions and a breakdown of different terms in the generalisation error, see Figure 7 in Appendix B. Theorem 3 (Asymptotic training error). Under the same regime, the training error converges a.s. to train 2 E∞ = w̄∗ σ̄eff + σ 2 max(0, 1 − γeff ) . This expression has the standard interpolation form. When γeff > 1, the model is overparameterised and perfectly fits the training data, yielding zero error. When γeff < 1, the model cannot interpolate, and the error decreases linearly with 2 γeff , with slope set by the total residual variance w̄∗ σ̄eff + σ2 . Altogether, these results provide exact expressions for generalisation, estimation, and training errors, predicting several key trends. In Figure 1(b-c), we illustrate these errors as functions of the compression ratio α under varying data structures λ. The plots demonstrate an excellent agreement between theoretical predictions and experimental results for moderate values of nu , nl , p. Increasing λ reduces the bias from the missing signal, the leak and also the variance. However, it does not reduce the errors across α uniformly, as errors with low α benefit more from stronger structure than errors with high α. We validate our theory across a range of parameter settings, with experiments exploring variations in the signal-to-noise ratio, the alignment between the spike v and the signal w⋆ , and the number of unlabelled and labelled samples (see Figure 5 in Appendix B). Importantly, optimising the representation dimensionality reduces generalisation error across a broad region of parameter space, both relative to low-dimensional representations with a single PC (α ≈ 0, Figure 2(c)) and to full-rank representations (α = 1, Figure 2(b)), with the optimum shown in Figure 2(a). These gains are particularly pronounced in regimes with substantial pretraining data and limited downstream samples, where standard PCR or regular regression fail to achieve comparable performance even in misaligned task settings (see Figure 12 in Appendix C.1 for a comparison). Next, we study the trade-off between pretraining and supervision at the optimal α⋆ by quantifying how much error reduction is obtained from additional unlabelled versus labelled samples (Figure 2(d)). The resulting phase map reveals a nontrivial dependence on both data sources. In particular, there exists a broad interior region of the phase diagram—spanning moderate to large nu and increasing nl —in which additional pretraining data yields larger reductions in generalisation error than labelled samples. Interestingly, this regime is closely linked to the optimal representation size α⋆ that minimises generalisation error (see Figure 2(e)), which is discussed below. 4.2

Optimal representation size

We now study the optimal representation size quantified by the fraction of retained principal components, α = m/p, that minimises generalisation and training error. We analyse how the resulting optimum depends on the interplay between unlabelled and labelled samples nu and nl , signal-to-noise ratio (SNR), spike strength λ, and task alignment η. If m does not grow with p (e.g., m = 1), then m/p → α = 0. Thus, in simulations, α ≈ 0 refers to taking m not growing with p (e.g., retaining only the top principal component m = 1).

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

Figure 3: When is compression useful. (a) Heatmap of the generalisation-optimal α⋆ reveals distinct phases; white curves mark the theoretical phase-transition boundaries where low α is optimal. Phase boundaries in the (nu , nl ) plane for varying SNR (b) varying λ (c) and varying η (d). Non-varying parameters are λ = 5, SNR = 9 and η = 1. Simulation details are given in Appendix D. Generalisation error. We find that the optimal level of compression minimising generalisation error adapts to the availability of unlabelled and labelled data. For fixed SNR, λ and an aligned task w⋆ ∥ v, Figure 3(a) reveals the following three regimes. (i) Scarce nl and nu : not enough pretraining samples to reliably estimate PCs, thus downstream tasks should not be compressed (α⋆ ≈ 1). (ii) Large nu , low nl : maintain a compact representation subspace (α⋆ < 1). (iii) High nl : little compression needed (α⋆ ≈ 1), as the number of task samples is enough to learn. We argue that regime (ii) is the most representative of modern machine-learning pipelines, where pretraining data is abundant but downstream task-specific data are limited. In practice, modern models are typically pretrained on large-scale datasets [2, 79] and then fine-tuned on substantially smaller task-specific datasets [3]. Altogether, we observe two phase transitions: one between regimes (i) and (ii), and another between regimes (ii) and (iii). We discuss them below in detail. Phase transitions for generalisation error. We derive approximate analytical phase transition lines as a function of signal-to-noise ratio (SNR), spike strength (λ), and sample allocation between pretraining and downstream tasks (nu and nl ). These approximations characterise the boundaries separating regimes in which different representation dimensions (α⋆ ) are optimal, and they provide a tractable description of the underlying transitions observed in simulations. Corollary 1 describes the transition between region (i) and (ii) when γl > 1. In this regime, the generalisation error exhibits a double descent, with two local minima: one at α = 1, and one at small α (see Figure 2(b)). Although the exact location of the left minimum can in principle be obtained from the closed-form error expressions, the resulting formulas are not analytically transparent. Instead, Corollary A.7 provides an interpretable approximation by comparing the error at the two extremes α = 1 and α = 0, corresponding to retaining all components and only a fixed number of PCs. ∗ Corollary 1 (Phase transition for γl > 1). Let γl > 1 and assume Σnew = Σ. Let SNR → S = w̄ σ 2 and define the overlap between the population and sample covariance spike eigenvector c(γu ) = 1γu <(λ−1)2

γu 1 − (λ−1) 2 γu 1 + λ−1

gen Then, the transition between the small-α and large-α minima is approximated by the endpoint condition E∞ (0) = gen E∞ (1), which yields a parametric curve

1 + η(λ − 1) − η

λ2 c(γu ) γl − 1 λγl (γl − 1) 1 = (1 − η) +η + . 1 + (λ − 1)c(γu ) γl (γl − 1 + λ)2 S(γl − 1)

(4)

Along the curve (4), the two local minima of the generalisation error approximately coincide; crossing the curve then induces an abrupt transition in α⋆ . The derivation of this result is in Appendix A.7. Corollary 2 characterises the transition between regions (ii) and (iii) in the underparametrised regime γl < 1. In this regime, the generalisation error exhibits a single minimum (see e.g. Figure 5(e,g,h) in Appendix B.1). The point α = 1 is optimal as long as decreasing α increases the generalisation error, which is determined by the condition gen ∂E∞ = 0. ∂α α=1−

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

∗

Corollary 2 (Phase transition for γl < 1). Let γl < 1, SNR → S = w̄ σ 2 and denote by Dv (γu ) the α-derivative of p̄v,v at α = 1. Then, the solution α = 1 is locally stable if and only if γl ≤ γlcrit :  λγu   , γu < 1, √  S [(1 − η) + ηDv (γu )] (λ − 1 + γu )2 γlcrit = where Dv (γu ) = γu  1 + S [(1 − η) + ηDv (γu )]  , γu > 1.  γu − 1 + λ Depending on the regime, the resulting boundary induces either a linear or nonlinear scaling of the critical number of labelled samples nl with the unlabelled sample size nu . Below this threshold, the optimal representation shifts to α < 1, and compression improves downstream generalisation. The derivation of this result is in Appendix A.8. Of particular interest is the limit γu → 0, corresponding to infinite unlabelled data—a regime reminiscent of modern pretraining datasets [1, 2, 80]. Here, errors and phase boundaries are independent of γu and have simple dependencies on η, SNR, λ, and γl . Precise expressions are in Appendix A.9 (Corollary 3), where we also consider the limit of infinite labelled data, γl → 0 (Corollary 4). In this last regime, the generalisation error simplifies (see Eq. (43)), the double-descent behaviour disappears, and the effective projector (XPm )† XPm reduces to the population projector onto Pm . The full phase transition curve, derived in Corollaries 1 and 2, matches the phase transition observed in the heat map of Figure 3(a) (shown in white; see also additional heatmaps in Figures 8 and 10 in Appendix B.2). Increasing the SNR progressively shrinks the region where compression is beneficial (Figure 3(b)), making full-rank representations (α = 1) optimal over a wider range of sample sizes. The converse is true when increasing the spike strength λ or the task-spike alignment η (Figure 3(c,d)). These trends are evident from Corollary 2: when γl < 1, the stability boundary depends on S(1 − η), so increasing SNR or decreasing η directly reduces γlcrit . Additional heatmaps for α in Appendix B.2 (Figures 8 and 10) further validate the joint dependence on SNR, λ and η, as well as the match with the theoretically derived transitions. In particular, increasing SNR, and decreasing λ and η all act to suppress the benefits of compression: for γl < 1 the phase boundary flattens, while for γl > 1 it shifts along the nl axis, reducing the region where α < 1 is optimal. Intuitively, a weaker spike is harder to estimate, lower alignment reduces the relevance of prior structure, and higher SNR makes the downstream task easy to learn from scratch. These effects diminish the value of compression, even when labelled data is abundant. Training error. Our analysis further characterises the optimal value of α for minimising the training error. In the oversampled regime (nl > p, corresponding to γl < 1), the optimal representation is full-rank, i.e., α = 1. Conversely, in the undersampled regime (nl < p, corresponding to γl > 1), the training error is equal to zero. In this case, the optimal dimensionality matches the sample size, i.e., α = nl /p, since increasing the dimension beyond this point provides no additional information. This is illustrated in Figures 9 and 11 in Appendix B.2. Altogether, our exact asymptotic analysis provides a theoretical framework to understand when low- or high-dimensional representations are optimal for downstream generalisation. We next validate such framework in practical settings.

Validation in autoencoders and pretrained transformers

In this section, we evaluate the extent to which theoretical predictions are valid beyond the idealised setting employed in their derivation. We first consider a controlled synthetic experiment in which the representation learned by an autoencoder approximates PCA. We then relax the assumptions of linearity and spiked isotropic covariance, allowing for a more general covariance structure. Finally, we show that our predictions extend to representations learned by modern neural networks. 5.1

Autoencoders as learned representations

We train the autoencoder on a pretraining task and use the learned representations for regression on downstream tasks. This approach is motivated by the fact that, in several regimes, autoencoders with a bottleneck dimension m recover the subspace spanned by the top-m principal components of the pretraining data [81–84]. We discuss below the regimes under which this correspondence holds. Spiked isotropic covariance. We first train autoencoders on the spiked identity covariance structures considered in the theory. For linear autoencoders, the optimal bottleneck size, defined as the value minimising the generalisation error on a downstream task, closely follows our theoretical results (compare Figure 4(a) with Figure 3(a)). We highlight that the boundaries of the phase transitions in the heatmaps precisely match the predictions from Section 4.2. Nonlinear autoencoders exhibit the same qualitative behaviour, except when the number of pretraining samples is small, see 7

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

(e)

Nonlinear AE

400

4,000

300

1,000

200

250

100

253

568

884

Unlabelled nu

1200

(d)

253

568

884

Unlabelled nu

1200

0.6

Optimal

(c)

0.75

143,000

10 2

253 1

10 3 1

253

568

884

Unlabelled nu

1200

253

568

884

Unlabelled nu

1200

20,000

143,000

(g)

10 1

0.0

568

2,000

Pretraining step Step

0.2

884

128

(f)

0.4

Density

Labelled nl

20,000

0.8

253

1200

Spiked general cov.

500

Optimal m

1.0

568

Labelled nl

884

Pretrained LLM

67,349

Labelled nl

Spiked identity cov.

(b)

Linear AE

Val accuracy

(a) 1200

100

Eigenvalue

125

150

0.70 0.65 0.60 0.55 0.50

Step

2 8 32 128 512 2,000 5,000 20,000 80,000 143,000

100

PCA components m

Figure 4: Optimal representation in autoencoders and pretrained LLMs. Optimal representation size as a function of unlabelled (nu ) and labelled (nl ) sample sizes. Linear autoencoders are considered in panels (a,c) and nonlinear ones in panels (b,d). Autoencoders are trained on Gaussian data with either spiked identity covariance Σ = Ip + λvv⊤ (panels (a,b)) or a spiked Toeplitz covariance Σ = H + λvv⊤ , where Hi,j = ρ|i−j| and ρ = 0.5 (panels (b,d)). (e) In a language transformer (PYTHIA models), PCA on representations also improves downstream linear probing on the sst2 task from the GLUE benchmark. (f) Eigenvalue spectrum of last-token representations from the final hidden layer, showing a bulk with a few spikes. (g) Validation accuracy versus representation size m for a small downstream dataset (nl = 2000); colors denote different checkpoints and shaded regions indicate standard deviation across seeds. Simulation details are given in Appendix D. Figure 4(b). This discrepancy is explained by a data-dependent transition: both linear and nonlinear autoencoders recover PCA only once sufficient pretraining data is available, see Figure 15 in Appendix C.2. In the low-sample regime, autoencoders instead tend to memorise training samples [85–87]. Despite this, when nu < p, higher-dimensional representations still yield lower errors on downstream tasks. General covariance structure. For a spiked general covariance, the same qualitative behaviour persists. In fact, while phase boundaries shift, the overall structure is preserved: no compression when pretraining data is limited; small bottlenecks for large pretraining and small downstream datasets; and intermediate-sized bottlenecks for large pretraining datasets. This is illustrated in Figure 4(c,d), and it is consistent with PCA-based regression under general covariance, see Figure 13 in Appendix C.2. Furthermore, when optimising for training error, the results closely match the spiked identity case (Figure 13 in Appendix C.2). At lower signal-to-noise ratios, the regime where compression is optimal expands (Figure 14 in Appendix C.2), in agreement with the theoretical analysis. Nonlinear autoencoders again exhibit deviations from PCA at low pretraining sample sizes. Overall, our empirical results show that the theoretical predictions derived in the idealised setting extend to trained linear and nonlinear autoencoders when enough pretraining data is available, and that the qualitative phase structure persists under general covariance spectra.

5.2

Linear probing transformers

We next examine whether representations learned by modern architectures exhibit similar dimensionality effects under linear probing. We use models from the Pythia family [79], which provide access to intermediate checkpoints throughout pretraining. As these models are trained for ∼1.5 epochs on the Pile dataset [80, 88], the training step serves as a proxy for the number of pretraining samples nu . The learned representations exhibit a spiked covariance structure, with a clear bulk and a small number of outlying eigenvalues (Figure 4(f)), and the spikes become more pronounced as training progresses, see Figure 18 in Appendix C.3. For each checkpoint, we project representations onto the top-m principal components and train linear probes on a downstream sentiment classification task (SST-2 from GLUE [3]), selecting m to optimise validation or training accuracy. 8

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

The dependence on dataset size is consistent with theoretical predictions, see Figure 4(e). For small downstream datasets, models that are pretrained longer benefit from stronger compression, see Figure 4(g). For larger downstream datasets, performance improves monotonically with m, indicating that little or no compression is optimal (see the right-most plot of Figure 17 in Appendix C.3). For extremely small downstream datasets (∼100 samples), the optimal dimensionality is high again (see the left-most plot in Figure 17 in Appendix C.3), consistent with the weak alignment between the task signal and dominant pretraining directions. Finally, tuning the representation dimensionality via PCA consistently matches or outperforms using the full feature space, see Figure 19 in Appendix C.3. In summary, these results show that real-world representations exhibit a theoretically predictable qualitative phase structure when optimising dimensionality for downstream generalisation.

Discussion

We present a theoretical model quantifying how representation dimensionality affects training and generalisation error on downstream tasks. We demonstrate that (i) tuning this dimensionality significantly improves generalisation compared to full representations or naive fixed dimensionality reduction, and that (ii) the optimal dimensionality is determined by the amount of pretraining and downstream data, the task signal-to-noise ratio, the strength of structure in the inputs, and the alignment between task-relevant signals and dominant covariance directions. In particular, when pretraining data is abundant and downstream data is limited, the optimal representations are maximally compressed. Furthermore, we unveil a trade-off between pretraining and supervision, identifying regimes where unlabelled data yields larger gains than labelled samples. These results provide a principled explanation for why large models trained on extensive data and exhibiting low intrinsic dimensionality [11], are especially effective for fine-tuning and linear probing. The question of optimal representations is central to machine learning as well as neuroscience, where internal representations support categorisation, reasoning, planning, navigation, and decision-making [89–91]. A key property of these representations is their dimensionality. High-dimensional representations may support flexible readout across tasks [92–95], whereas low-dimensional representations are associated with abstraction, efficient coding, or structured prior knowledge [27, 96]. Our results provide a quantitative characterization of how task structure and data availability determine the optimal representation tradeoff in a simplified setting, offering predictions for how biological systems might organize representations for effective learning. Several limitations motivate future work: relaxing the shared-spike assumption to capture data mismatch, extending the analysis to ridge-regularised learning, and considering richer covariance models. Another important direction is to replace the closed-form regression estimator with gradient-based training, enabling the study of optimisation-induced bias and its dependence on initialisation from pretraining. Finally, extending the framework beyond PCA-based projections to learned feature representations would allow for a closer connection to empirical observations in deep networks.

Contributions A.S. and R.S. initiated and conceptualized the project. V.N. and C.D. developed the theoretical analysis and derived the main results, with V.N. leading the derivations. V.N. and C.D. designed and performed the experiments. M.M. co-supervised the project and contributed to the technical development of the work, including proofs, experimental design, and interpretation of results. A.S. co-supervised the project and contributed to the conceptual framing of the project, interpretation of the theoretical results, experimental design and overall scientific direction. V.N. and C.D. wrote the manuscript. M.M. and A.S. reviewed and edited the paper. All authors discussed the results and commented on the manuscript.

Acknowledgements This research was funded in whole or in part by the Austrian Science Fund (FWF) (10.55776/COE12), Gatsby Charitable Foundation (GAT3850 and GAT4058), Sainsbury Wellcome Centre Core Grant from Wellcome (219627/Z/19/Z). For the purpose of open access, the authors have applied a CC BY public copyright license to any Author Accepted Manuscript version arising from this submission.

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

References [1] OpenAI et al. GPT-4 Technical Report. arXiv preprint arXiv:2303.08774, 2024. [2] Aaron Grattafiori et al. The Llama 3 Herd of Models. arXiv preprint arXiv:2407.21783, 2024. [3] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, 2018. [4] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. In Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks, volume 1, 2021. [5] Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In Advances in Neural Information Processing Systems, volume 30, 2017. [6] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems, volume 35, 2022. [7] Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. [8] Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International conference on machine learning. PMLR, 2020. [9] Yuejiang Liu, Parth Kothari, Bastien Van Delft, Baptiste Bellot-Gurlet, Taylor Mordan, and Alexandre Alahi. Ttt++: When does self-supervised test-time training fail or thrive? In Advances in Neural Information Processing Systems, volume 34, 2021. [10] Ekin Akyürek, Mehul Damani, Adam Zweiger, Linlu Qiu, Han Guo, Jyothish Pari, Yoon Kim, and Jacob Andreas. The surprising effectiveness of test-time training for few-shot learning. In International Conference on Machine Learning. PMLR, 2025. [11] Armen Aghajanyan, Sonal Gupta, and Luke Zettlemoyer. Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), 2021. [12] Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), 2018. [13] Yoshua Bengio, Pascal Lamblin, Dan Popovici, and Hugo Larochelle. Greedy layer-wise training of deep networks. In Advances in Neural Information Processing Systems, volume 19, 2006. [14] Dumitru Erhan, Yoshua Bengio, Aaron Courville, Pierre-Antoine Manzagol, Pascal Vincent, and Samy Bengio. Why does unsupervised pre-training help deep learning? Journal of Machine Learning Research, 11, 2010. [15] Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Yoshua Bengio, and Pierre-Antoine Manzagol. Stacked denoising autoencoders: Learning useful representations in a deep network with a local denoising criterion. Journal of Machine Learning Research, 11, 2010. [16] Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE Trans. Pattern Anal. Mach. Intell., 2013. [17] Alex Martin. The representation of object concepts in the brain. Annu. Rev. Psychol., 58, 2007. [18] Stefano Fusi, Earl K Miller, and Mattia Rigotti. Why neurons mix: high dimensionality for higher cognition. Current Opinion in Neurobiology, 37, 2016. 10

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

[19] German I Parisi, Ronald Kemker, Jose L Part, Christopher Kanan, and Stefan Wermter. Continual lifelong learning with neural networks: A review. Neural Networks, 113, 2019. [20] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114, 2017. [21] Friedemann Zenke, Ben Poole, and Surya Ganguli. Continual learning through synaptic intelligence. In International Conference on Machine Learning. PMLR, 2017. [22] Lukas Braun, Clémentine Carla Juliette Dominé, James E Fitzgerald, and Andrew M Saxe. Exact learning dynamics of deep linear networks with prior knowledge. In Advances in Neural Information Processing Systems, 2022. [23] David Badre, Apoorva Bhandari, Haley Keglovits, and Atsushi Kikumoto. The dimensionality of neural representations for control. Current Opinion in Behavioral Sciences, 38, 2021. [24] Lara M. Boyle, Lorenzo Posani, Sarah Irfan, Steven A. Siegelbaum, and Stefano Fusi. Tuned geometries of hippocampal representations meet the computational demands of social memory. Neuron, 112, 2024. [25] Hristos S. Courellis, Juri Minxha, Araceli R. Cardenas, Daniel L. Kimmel, Chrystal M. Reed, Taufik A. Valiante, C. Daniel Salzman, Adam N. Mamelak, Stefano Fusi, and Ueli Rutishauser. Abstract representations emerge in human hippocampal neurons during inference. Nature, 632, 2024. [26] Karyna Mishchanchuk, Gabrielle Gregoriou, Albert Qü, Alizée Kastler, Quentin J. M. Huys, Linda Wilbrecht, and Andrew F. MacAskill. Hidden state inference requires abstract contextual representations in the ventral hippocampus. Science, 386, 2024. [27] Ramon Nogueira, Chris C. Rodgers, Randy M. Bruno, and Stefano Fusi. The geometry of cortical representations of touch in rodents. Nature Neuroscience, 26, 2023. [28] Lucrezia Valeriani, Diego Doimo, Francesca Cuturello, Alessandro Laio, Alessio Ansuini, and Alberto Cazzaniga. The geometry of hidden representations of large transformer models. In Advances in Neural Information Processing Systems, volume 36, 2023. [29] Edgar Dobriban and Stefan Wager. High-dimensional asymptotics of prediction: Ridge regression and classification. The Annals of Statistics, 2018. [30] Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan J. Tibshirani. Surprises in high-dimensional ridgeless least squares interpolation. The Annals of Statistics, 50, 2022. [31] Dominic Richards, Jaouad Mourtada, and Lorenzo Rosasco. Asymptotics of ridge(less) regression under general source condition. In International Conference on Artificial Intelligence and Statistics, 2020. [32] Denny Wu and Ji Xu. On the optimal weighted \ell_2 regularization in overparameterized linear regression. In Advances in Neural Information Processing Systems, volume 33, 2020. [33] Peter L. Bartlett, Philip M. Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117, 2020. [34] Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine-learning practice and the classical bias–variance trade-off. Proceedings of the National Academy of Sciences, 116, 2019. [35] Madhu S. Advani, Andrew M. Saxe, and Haim Sompolinsky. High-dimensional dynamics of generalization error in neural networks. Neural Networks, 132, 2020. [36] Qiyang Han and Xiaocong Xu. The distribution of ridgeless least squares interpolators. Journal of Machine Learning Research, 27, 2026. [37] Simone Bombari and Marco Mondelli. Spurious correlations in high dimensional regression: The roles of regularization, simplicity bias and over-parameterization. In International Conference on Machine Learning, 2025. [38] Edwige Cyffers, Alireza Mirrokni, and Marco Mondelli. Optimal regularization for performative learning. arXiv preprint arXiv:2510.12249, 2025. 11

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

[39] Parham Rezaei, Filip Kovacevic, Francesco Locatello, and Marco Mondelli. High-dimensional analysis of synthetic data selection. In International Conference on Learning Representations, 2026. [40] M. Emrullah Ildiz, Halil Alperen Gozeten, Ege Onur Taga, Marco Mondelli, and Samet Oymak. High-dimensional analysis of knowledge distillation: Weak-to-strong generalization and scaling laws. In International Conference on Learning Representations, 2025. [41] Germain Kolossov, Andrea Montanari, and Pulkit Tandon. Towards a statistical theory of data selection under weak supervision. In International Conference on Learning Representations, 2024. [42] Ayush Jain, Andrea Montanari, and Eren Sasoglu. Scaling laws for learning with real and surrogate data. Advances in Neural Information Processing Systems, 37, 2024. [43] Ian T. Jolliffe. A Note on the Use of Principal Components in Regression. Applied Statistics, 31, 1982. [44] Ji Xu and Daniel J Hsu. On the number of variables to use in principal component regression. In Advances in Neural Information Processing Systems, volume 32, 2019. [45] Alden Green and Elad Romanov. The high-dimensional asymptotics of principal component regression. The Annals of Statistics, 53, 2025. [46] William F. Massy. Principal components regression in exploratory statistical research. Journal of the American Statistical Association, 60, 1965. [47] Paramveer S. Dhillon, Dean P. Foster, Sham M. Kakade, and Lyle H. Ungar. A risk comparison of ordinary least squares vs ridge regression. Journal of Machine Learning Research, 14, 2013. [48] Jianfeng Yao, Zhidong Bai, and Shui-Rong Zheng. Large sample covariance matrices and high-dimensional data analysis. Cambridge series in statistical and probabilistic mathematics. Cambridge university press, 2015. [49] Sinno Jialin Pan and Qiang Yang. A survey on transfer learning. IEEE Transactions on Knowledge and Data Engineering, 22, 2010. [50] Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. A comprehensive survey on transfer learning. Proceedings of the IEEE, PP, 2020. [51] Federica Gerace, Luca Saglietti, Stefano Sarao Mannelli, Andrew Saxe, and Lenka Zdeborová. Probing transfer learning with a model of synthetic correlated datasets. Machine Learning: Science and Technology, 2022. [52] Andrew K. Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. In International Conference on Learning Representations, 2019. [53] Javan Tahir, Surya Ganguli, and Grant M. Rotskoff. Features are fate: a theory of transfer learning in highdimensional regression. In Forty-second International Conference on Machine Learning, 2025. [54] Fan Yang, Hongyang R Zhang, Sen Wu, Christopher Re, and Weijie J Su. Precise high-dimensional asymptotics for quantifying heterogeneous transfers. Journal of Machine Learning Research, 26, 2025. [55] Yanke Song, Sohom Bhattacharya, and Pragya Sur. Generalization error of min-norm interpolators in transfer learning. arXiv preprint arXiv:2406.13944, 2024. [56] Jason D. Lee, Qi Lei, Nikunj Saunshi, and Jiacheng Zhuo. Predicting what you already know helps: Provable self-supervised learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. [57] Eyar Azar and Boaz Nadler. Semi-supervised sparse gaussian classification: Provable benefits of unlabeled data. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. [58] Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv:1902.09229, 2019. [59] Tianyi Zhang and Tatsunori B. Hashimoto. On the inductive bias of masked language modeling: From statistical to syntactic dependencies. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, 2021. 12

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

[60] Taj Jones-Mccormick, Aukosh Jagannath, and Subhabrata Sen. Provable benefits of unsupervised pre-training and transfer learning via single-index models. In Proceedings of the 42nd International Conference on Machine Learning, volume 267. PMLR, 2025. [61] Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.01644, 2018. [62] Bin Wang, W. Jeffrey Johnston, and Stefano Fusi. A mathematical theory for understanding when abstract representations emerge in neural networks. arXiv preprint arXiv:2510.09816, 2026. [63] Nicolas Anguita, Francesco Locatello, Andrew M. Saxe, Marco Mondelli, Flavia Mancini, Samuel Lippl, and Clementine Domine. A theory of how pretraining shapes inductive bias in fine-tuning. arXiv preprint arXiv:2602.20062, 2026. [64] Iain M. Johnstone. On the distribution of the largest eigenvalue in principal components analysis. The Annals of Statistics, 29(2), 2001. [65] Jinho Baik, Gérard Ben Arous, and Sandrine Péché. Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices. The Annals of Probability, 33, 2005. [66] V. A. Marčenko and L. A. Pastur. Distribution of eigenvalues for some sets of random matrices. Mathematics of the USSR-Sbornik, 1, 1967. [67] Zhidong Bai and Jack W. Silverstein. Spectral Analysis of Large Dimensional Random Matrices. Springer Series in Statistics. Springer New York, 2010. [68] Jinho Baik and Jack W Silverstein. Eigenvalues of large sample covariance matrices of spiked population models. Journal of multivariate analysis, 97, 2006. [69] Debashis Paul. Asymptotics of sample eigenstructure for a large dimensional spiked covariance model. Statistica Sinica, 2007. [70] Zhidong Bai and Jianfeng Yao. On sample eigenvalues in a generalized spiked population model. Journal of Multivariate Analysis, 106, 2012. [71] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu. Learning in the presence of lowdimensional structure: a spiked random matrix perspective. In Advances in Neural Information Processing Systems, volume 36, 2023. [72] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? In Advances in Neural Information Processing Systems, volume 33, 2020. [73] Alireza Mousavi-Hosseini, Denny Wu, Taiji Suzuki, and Murat A Erdogdu. Gradient-based feature learning under structured data. In Advances in Neural Information Processing Systems, volume 36, 2023. [74] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. In Advances in Neural Information Processing Systems, volume 35, 2022. [75] Hugo Cui, Luca Pesce, Yatin Dandi, Florent Krzakala, Yue M Lu, Lenka Zdeborová, and Bruno Loureiro. Asymptotics of feature learning in two-layer networks after one gradient-step. In International Conference on Machine Learning, 2024. [76] Yatin Dandi, Luca Pesce, Hugo Cui, Florent Krzakala, Yue Lu, and Bruno Loureiro. A random matrix theory perspective on the spectrum of learned features and asymptotic generalization capabilities. In International Conference on Artificial Intelligence and Statistics, pages 2224–2232. PMLR, 2025. [77] Rishi Sonthalia, Michael Murray, and Guido F. Montúfar. Low rank gradients and where to find them. In Advances in Neural Information Processing Systems, 2025. [78] Daniel Gedon, Antônio H. Ribeiro, and Thomas B. Schön. No double descent in principal component regression: A high-dimensional analysis. In International Conference on Machine Learning, 2024. 13

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

[79] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O’Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, 2023. [80] Stella Biderman, Kieran Bicheno, and Leo Gao. Datasheet for the pile. arXiv preprint arXiv:2201.07311, 2022. [81] Maria Refinetti and Sebastian Goldt. The dynamics of representation learning in shallow, non-linear autoencoders. Journal of Statistical Mechanics: Theory and Experiment, 2023, 2023. [82] Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural Networks, 2, 1989. [83] Elad Plaut. From Principal Subspaces to Principal Components with Linear Autoencoders. arXiv preprint arXiv:1804.10253, 2018. [84] H. Bourlard and Y. Kamp. Auto-association by multilayer perceptrons and singular value decomposition. Biological Cybernetics, 59, 1988. [85] Marco Fumero, Luca Moschella, Emanuele Rodolà, and Francesco Locatello. Navigating the latent space dynamics of neural models. In The Fourteenth International Conference on Learning Representations, 2026. [86] Yibo Jiang and Cengiz Pehlevan. Associative memory in iterated overparameterized sigmoid autoencoders. In Proceedings of the 37th International Conference on Machine Learning, 2020. [87] Adityanarayanan Radhakrishnan, Mikhail Belkin, and Caroline Uhler. Overparameterized Neural Networks Implement Associative Memory. Proceedings of the National Academy of Sciences, 117, 2020. [88] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. [89] Mark G Stokes, Makoto Kusunoki, Natasha Sigala, Hamed Nili, David Gaffan, and John Duncan. Dynamic coding for cognitive control in prefrontal cortex. Neuron, 78, 2013. [90] Aldo Genovesio, Peter J Brasted, Andrew R Mitz, and Steven P Wise. Prefrontal cortex activity related to abstract response strategies. Neuron, 47, 2005. [91] Keisetsu Shima and Jun Tanji. Neuronal activity in the supplementary and presupplementary motor areas for temporal organization of multiple movements. Journal of neurophysiology, 84, 2000. [92] David Raposo, Matthew T Kaufman, and Anne K Churchland. A category-free neural population supports evolving demands during decision-making. Nature neuroscience, 17, 2014. [93] Evelyn Tang, Marcelo G Mattar, Chad Giusti, David M Lydon-Staley, Sharon L Thompson-Schill, and Danielle S Bassett. Effective learning is accompanied by high-dimensional and efficient representations of neural activity. Nature neuroscience, 22, 2019. [94] Mattia Rigotti, Omri Barak, Melissa R Warden, Xiao-Jing Wang, Nathaniel D Daw, Earl K Miller, and Stefano Fusi. The importance of mixed selectivity in complex cognitive tasks. Nature, 497, 2013. [95] Silvia Bernardi, Marcus K Benna, Mattia Rigotti, Jérôme Munuera, Stefano Fusi, and C Daniel Salzman. The geometry of abstraction in the hippocampus and prefrontal cortex. Cell, 183, 2020. [96] Kay Tye, Earl Miller, Felix Taschbach, Marcus Benna, Mattia Rigotti, and Stefano Fusi. Mixed selectivity: Cellular computations for complexity. Neuron, 112, 2024. [97] Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2026. [98] Francisco Rubio and Xavier Mestre. Spectral convergence for a general class of random matrices. Statistics & Probability Letters, 81, 2011.

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

Appendix A A.1

Theory Assumptions

Assumption 1 (High-dimensional asymptotic regime). Let the number of features be p, and nl , nu , m be such that p p m → γl ∈ (0, ∞), → γu ∈ (0, ∞), → α ∈ [0, 1]. nl nu p Furthermore, we define γeff := nml → αγl ∈ [0, ∞). In words, α is the fraction of dimensions to keep, γu and γl are the aspect ratios the prior and downstream datasets, and γeff is the effective aspect ratio of the compressed downstream dataset. Assumption 2 (Spiked covariance model). For each p, the population covariance matrix is Σ = Ip + (λ − 1)vv⊤ , where λ > 1 and the unit vector v ∈ Rp are fixed. Note that the eigendecomposition of Σ is Σ = Vdiag(λ, 1, . . . , 1)V⊤ , with V = [v1 , v2 , . . . , vp ] ∈ Rp×p . As we often refer directly to the spike eigenvector, we keep the simpler notation v1 = v. Assumption 3 (Data generation). The input matrices Xl ∈ Rnl ×p and Xu ∈ Rnu ×p have i.i.d. rows distributed as x ∼ N (0, Σ). Equivalently, there exists a matrix Z ∈ Rnl ×p with i.i.d entries distributed as N (0, 1), such that Xl = ZΣ1/2 . Assumption 4 (Task weight vector). Let the alignment between the true signal w⋆ and the data covariance structure Σ be given by the spectral measure p

Ĝw⋆ (τ ) =

1 X ⋆⊤ 2 (w vi ) 1{λi (Σ) ≤ τ }. w̄∗ i=1

Then, there exists a limiting spectrum Gw⋆ such that Ĝw⋆ ⇒ Gw⋆ almost surely. Sample covariance eigenspectrum. We further note that sample covariance matrices Su = n1u X⊤ u Xu and Sl = 1 ⊤ X X have empirical spectral measures l nl l p

F̂ Su (θ) =

1X 1{λi (Su ) ≤ θ}, p i=1

F̂ Sl (θ) =

1X 1{λi (Sl ) ≤ θ}, p i=1

which converge weakly almost surely to Marcenko-Pastur laws [66, 67] with aspect ratios γu and γl respectively: F̂ Su ⇒ Fγu ,

F̂ Sl ⇒ Fγl .

For completeness, we state the Marcenko-Pastur density with an aspect ratio γ: p (λ+ − λ)(λ − λ− ) 1 + 1γ>1 1 − δ(λ), fγ (λ) = 2πλγ γ √ λ+ = (1 + γ)2 , √ λ− = (1 − γ)2 , with λ+ , λ− denoting the edges of the support of the distribution. 15

(5)

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

A.2

Main ideas of the proofs

The key idea of each proof is to decompose the model’s estimator ŵ into a signal term, a bias term arising from the component of w⋆ outside the relevant subspace, and two variance terms arising from noise and model mismatch. The Gaussian structure of the data allows us to express the component of Xl orthogonal to the prior subspace Um as a linear function of Xl Um plus independent noise E. This reduces the analysis to understanding various projections onto the row space of spiked Gaussian matrices. Conditional Gaussian decomposition under prior subspace projections. Let Xl ∈ Rnl ×p be a random matrix whose rows xi ∼ N (0, Σ) are i.i.d., and let Um ∈ Rp×m be a random matrix with orthonormal columns. Let U⊥ ∈ Rp×(p−m) be such that [Um U⊥ ] is an orthogonal matrix. Define Xm = Xl Um ∈ Rnl ×m ,

X⊥ = Xl U⊥ ∈ Rnl ×(p−m) .

Then, conditioning on Um , the rows of (Xm , X⊥ ) are jointly Gaussian with covariance ⊤ Σmm Σm⊥ Um ΣUm U⊤ ΣU⊥ m = . Σ⊥m Σ⊥⊥ U⊤ U⊤ ⊥ ΣUm ⊥ ΣU⊥ Moreover, due to the conditional distribution, X⊥ given Xm satisfies X⊥ = Xm Σ−1 mm Σm⊥ + E,

(6) where E is independent of Xm and its rows are i.i.d. with distribution ei ∼ N 0, Σ̄ , with Σ̄ := Σ⊥⊥ − Σ⊥m Σ−1 mm Σm⊥ . Under the spiked covariance model of Assumption 2, the population covariance Σmm of Xm equals Σmm = Im + (λ̃ − 1)ṽṽ⊤ , (7) U⊤ v

m where λ̃ = 1 + (λ − 1)∥Pm v∥2 and ṽ = ∥Pm v∥ . Furthermore, after some algebraic manipulations using the Sherman-Morrison formula, it follows that

λ̃ − 1 ⊤ ṽṽ , λ̃ (λ − 1)∥Pm v∥ ⊤ ṽv U⊥ , Σ−1 mm Σm⊥ = λ̃ λ−1 ⊤ ⊤ Σ̄ = Ip−m + U⊥ vv U⊥ . λ̃ Σ−1 mm = Im −

Decomposing the model’s weight estimate. Equation (6), we can write:

(8) (9) (10)

From the definition of ŵ in Equation (1) and the decomposition in

ŵ = (Xl Pm )† Xl Pm w⋆ + (Xl Pm )† Xl P⊥ w⋆ + (Xl Pm )† ξ † ⊤ ⋆ ⊤ † −1 ⊤ ⋆ ⊤ † ⊤ ⋆ ⊤ † = (Xm U⊤ m ) Xm Um w + (Xm Um ) Xm Σmm Σm⊥ U⊥ w + (Xm Um ) EU⊥ w + (Xm Um ) ξ.

Then, using Equation (9) leads to (λ − 1)(w⋆ ⊤ P⊥ v) ⊤ Um v λ̃ | {z }

⊤ ⋆ Σ−1 mm Σm⊥ U⊥ w =

(11)

and by defining Πm := X†m Xm ∈ Rm×m , we obtain ŵ =

⋆ Πw | {z }

p×p Π := (Xl Pm )† Xl Pm = Um Πm U⊤ m ∈R

A | ·{zΠv}

signal from Pm subspace

(Xl Pm )† Ew⋆ ⊥ | {z }

signal from P⊥ subspace that leaked into Pm

variance from model mismatch

(Xl Pm )† ξ | {z }

(12)

variance from label noise

⋆ ⋆ where w⋆ ⊥ := U⊤ ⊥ w . Note that the component of w not involving label noise is

(Xl Pm )† Xl w⋆ = Πw⋆ + AΠv + (Xl Pm )† Ew⋆ ⊥ . The model’s weight estimate decomposes into the four parts discussed below. 16

(13)

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

• Πw⋆ : part of w⋆ that lies in the subspace formed by Pm and is identifiable from the available data. It is equal to Pm w⋆ if m < nl (i.e., γeff < 1), since in that case Πm = Im (see Lemma 6); if m > nl (i.e., γeff > 1), this term causes bias due to the finite sample size. • AΠv: a part of w⋆ that lies outside the subspace Pm can leak into Pm , causing the model to reconstruct it (in the Pm subspace). The reason is the existence of the population covariance spike which makes the rows of Xl P⊥ and Xl Pm correlated, so the part outside the subspace leaks back into the subspace through this covariance structure. • (Xl Pm )† Ew⋆ ⊥ : the part of the w⋆ that lies outside the subspace Pm and does not leak into Pm behaves like noise, eventually causing a double-descent-like curve that scales with ∥w⋆ ∥2 . • (Xl Pm )† ξ: label noise variance, causing another double-descent-like curve that scales with σ 2 . A.3

Proof for the estimation error

Proof of Theorem 1. The proof proceeds in three stages, exploiting the conditional independence structure of the data. First, we use the joint Gaussianity of Xm and X⊥ to rewrite X⊥ in terms of Xm and an independent noise matrix E ∈ Rnl ×(p−m) (see Section A.2). We then average over E conditioned on Xm and Um , then average over Xm conditioned on Um , and finally over Um . Step 1: Decomposition of the error. The estimation error splits into bias and variance: 2 Eξ ∥w⋆ − ŵ∥2 = I − (Xl Pm )† Xl w⋆ + Eξ ξ ⊤ (Xl Pm )†⊤ (Xl Pm )† ξ − 2Eξ ξ ⊤ (Xl Pm )†⊤ I − (Xl Pm )† Xl w⋆ 2 = I − (Xl Pm )† Xl w⋆ + Eξ Tr ξξ ⊤ (Xl Pm )†⊤ (Xl Pm )† 2 † = I − (Xl Pm )† Xl w⋆ + σ 2 Tr (Pm X⊤ l X l Pm ) . Substituting the result in (13) further decomposes the bias: 1 ∥w⋆ ∥2

1 ⋆ 2 ∥w⋆ − Πw⋆ − AΠv − (Xl Pm )† EU⊤ ⊥w ∥ ∥w⋆ ∥2 1 1 ⋆ 2 = ∥w⋆ − Πw⋆ − AΠv∥2 + ∗ (Xl Pm )† EU⊤ ⊥w ⋆ 2 ∥w ∥ w̄ 1 ⋆ (w⋆ − Πw⋆ − AΠv)⊤ (Xl Pm )† EU⊤ −2 ⊥w . ∥w⋆ ∥2 {z } |

2 I − (Xl Pm )† Xl w⋆ =

cross term T

Step 2: E-independent term. The first term ∥w⋆ − Πw⋆ − AΠv∥2 does not depend on E, so its conditional expectation on E is itself. We expand using the property of the projection matrix Π that Π2 = Π: 1 ∥Πw⋆ ∥2 A2 ⋆ ⋆ 2 ∥w − Πw − AΠv∥ = 1 − + ∥Πv∥2 , ∥w⋆ ∥2 ∥w⋆ ∥2 ∥w⋆ ∥2

(14)

where the cross term −2Aw⋆ ⊤ Πv + 2Aw⋆ ⊤ ΠΠv = 0. Substituting A = λ−1 (v⊤ P⊥ w⋆ ) gives λ̃ 1−

∥Πw⋆ ∥2 (λ − 1)2 (v⊤ P⊥ w⋆ )2 (λ − 1)2 2 p̄⊥,w⋆ ,v π̄v,v almost surely, + ∥Πv∥2 ⇒ 1 − π̄w⋆ ,w⋆ + ⋆ 2 ⋆ 2 ∥w ∥ ∥w ∥ λ̄2 λ̃2

(15)

⋆ 2

∥ 2 which matches the first two terms of the bias. The projections ∥Πw ∥w⋆ ∥2 and ∥Πv∥ converge almost surely to ⊤

⋆

P⊥ w π̄w⋆ ,w⋆ , π̄v,v in Equation (39) according to Lemma 6. The quantity v ∥w converges almost surely to p̄⊥,w⋆ ,v in ⋆∥

Equation (34) according to Lemma 5, and λ̃ = 1 + (λ − 1)∥Pm v∥2 converges to λ̄ almost surely by the continuous mapping theorem and the almost sure convergence of ∥Pm v∥2 → p̄v,v from Equation (30) (see Lemma 4). By the continuous mapping theorem, the combined expression in (15) converges almost surely. Step 3: Cross term linear in E.

We write the cross term as T = c⊤ Ed, 17

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

with 1 (Xl Pm )†⊤ (w⋆ − Πw⋆ − AΠv) ∈ Rnl , ∥w⋆ ∥ 1 U⊤ w⋆ ∈ Rp−m . d := ∥w⋆ ∥ ⊥ c :=

Conditioned on (Xm , Um ), the vectors c and d are deterministic and independent of E. Since the rows of E are i.i.d. Gaussian with covariance Σ̄, it follows that T | (Xm , Um ) ∼ N 0, (d⊤ Σ̄d) ∥c∥2 . We now bound the variance. First, d⊤ Σ̄d =

∥P⊥ w⋆ ∥2 λ − 1 (w⋆ ⊤ P⊥ v)2 + , ∥w⋆ ∥2 ∥w⋆ ∥2 λ̃

which by the almost sure limits of ∥P⊥ w⋆ ∥2 , λ̃, and w⋆ ⊤ P⊥ v, converges almost surely to λ−1 2 p̄ . ⋆ λ̄ ⊥,w ,v In particular, there exists a deterministic constant C1 such that for all sufficiently large p, 2 σ̄eff = p̄⊥,w⋆ ,w⋆ +

d⊤ Σ̄d ≤ C1

almost surely.

It therefore remains to control ∥c∥2 . We have ∥c∥2 ≤

1 ∥w⋆ − Πw⋆ − AΠv∥2 . ∥w⋆ ∥2 σmin (Xm )2

By Step 2, the numerator is bounded almost surely for all sufficiently large p. For the denominator, writing Xm = 1/2 GΣmm (see Lemma 1) gives σmin (Xm ) ≥ σmin (G)σmin (Σ1/2 mm ). p 1/2 Note that σmin (Σmm ) = min{1, λ̃} for any p. By the singular value inequality for Gaussian random matrices (see e.g. Exercise 7.13 in [97]), for all t > 0, √ 2 √ P σmin (G) ≤ | nl − m| − t ≤ 2e−t /2 . Since γeff = m/nl ̸= 1, we have √ √ √ √ |1 − γeff | √ | nl − m| = p = p cγ . √ γl √ Choosing t = 21 cγ p and setting c = c2γ /8, we obtain ! c2γ 1 √ P σmin (G) ≤ cγ p ≤ 2 exp − p . 2 8 Noting that λ̃ ≤ λ, and setting aγ := 14 c2γ min{1, λ}, bγ := 18 c2γ , we have P σmin (Xm )2 ≤ aγ p ≤ 2e−bγ p . Therefore, there exists C2 > 0 such that C2 2 P ∥c∥ ≤ ≥ 1 − 2e−bγ p . p Define the event

Ep :=

Then

C2 ∥c∥ ≤ p 2

P(Epc ) ≤ 2e−bγ p . 18

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

For any δ > 0,

P(|T | ≥ δ) = P(|T | ≥ δ, Epc ) + P(|T | ≥ δ, Ep ) ≤ P(Epc ) + P(|T | ≥ δ, Ep ).

Noting that there exists C1 s.t. d⊤ Σ̄d ≤ C1 for any large p, on Ep , Var(T | Xm , Um ) ≤

C1 C2 . p

Thus, on Ep , we have P(|T | ≥ δ | Xm , Um ) ≤ 2 exp(−cδ p), and since this holds for any Xm , Um , by the law of total probability, we have P(|T | ≥ δ, Ep ) ≤ 2e−cδ p . Therefore, Since

−cp

p Ce

P(|T | ≥ δ) ≤ 2e−bγ p + 2e−cδ p = Ce−cp . < ∞, the Borel–Cantelli lemma implies a.s.

T −−→ 0. Step 4: Average of the trace terms.

The variance term contains † ⊤ † Tr[(Pm X⊤ l Xl Pm ) ] = Tr[(Xm Xm ) ].

† ⊤ ⊤ We note that the term Tr[(Xm X⊤ m ) ] is equivalent to the above, since the non-zero eigenvalues of Xm Xm and Xm Xm min{1,γeff } are the same. By Lemma 2, the trace term has a limiting value |γeff −1| almost surely.

Step 5: Quadratic term in E.

Define † ⊤ ⋆ Q := w⋆ ⊤ U⊥ E⊤ (Xm X⊤ m ) EU⊥ w .

Let

⋆ † a := U⊤ A := (Xm X⊤ ⊥w , m) . Conditioned on (Xm , Um ), the vector a and the matrix A are fixed and independent of E. Since the rows of E are i.i.d. Gaussian with covariance Σ̄, the vector c := Ea is Gaussian with c | (Xm , Um ) ∼ N 0, (a⊤ Σ̄a)Inl . Therefore Q = c⊤ Ac.

First, the conditional expectation is EE [Q | Xm , Um ] = (a⊤ Σ̄a)Tr(A). Here

⋆ ⋆ 2 2 a⊤ Σ̄a = w⋆ ⊤ U⊥ Σ̄U⊤ ⊥ w = ∥w ∥ σeff . By the deterministic equivalents used above, a.s.

2 2 ∥w⋆ ∥2 σeff −−→ w̄∗ σ̄eff .

Moreover, by the pseudo-inverse trace computation in Step 4, † a.s. min{1, γeff } . Tr (Xm X⊤ −−→ m) |γeff − 1| Thus a.s.

2 EE [Q | Xm , Um ] −−→ w̄∗ σ̄eff

min{1, γeff } . |γeff − 1|

It remains to show concentration around this conditional expectation. Write p c = a⊤ Σ̄a z, z ∼ N (0, Inl ), 19

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

δ δ2 , ⊤ 2 ⊤ 2 (a Σ̄a) ∥A∥F (a Σ̄a)∥A∥

We now bound the quantities appearing in this inequality. As above, 2 a⊤ Σ̄a = ∥w⋆ ∥2 σeff

is bounded for all sufficiently large p almost surely. Furthermore, since γeff ̸= 1, the nonzero spectrum of n1l X⊤ m Xm is √ bounded away from zero almost surely. Equivalently, the nonzero singular values of Xm are of order p (see argument in Step 3). Hence 1 † ∥A∥ = ∥(Xm X⊤ ) ∥ = O , m p and X 1 1 2 ∥A∥F = =O . σi (Xm )4 p i:σi (Xm )>0

Consequently, for all sufficiently large p, PE (|Q − EE [Q | Xm , Um ]| ≥ δ | Xm , Um ) ≤ 2e−cδ p . The right-hand side is summable in p, so Borel–Cantelli gives a.s.

Q − EE [Q | Xm , Um ] −−→ 0. Combining this concentration statement with the limit of the conditional expectation yields a.s.

2 Q −−→ w̄∗ σ̄eff

min{1, γeff } . |γeff − 1|

Combining. Summing the contributions from Steps 1–5 and applying the continuity mapping theorem concludes the proof.

A.4

Proof for the generalisation error

Proof of Theorem 2. We decompose the error and take averages over E, Xm , Um . Step 1: Decomposition of the generalisation error. We start from the definition of the generalisation error E gen = Ex,y,ξ ∥y − x⊤ ŵ∥2 | w⋆ , Xl , Xu . Using the data model y = x⊤ w⋆ + ξ and independence of the test noise, this becomes E gen = Ex,ξ (w⋆ − ŵ)⊤ xx⊤ (w⋆ − ŵ) | w⋆ , Xl , Xu + σ 2 = Eξ (w⋆ − ŵ)⊤ Σ(w⋆ − ŵ) + σ 2 . P ⊤ Using the spiked covariance model Σnew = Ip + i (νi − 1)vnew,i vnew,i , we obtain h X 2 i E gen = E est + (νi − 1) Eξ (w⋆ − ŵ)⊤ vnew,i + σ2 .

(16)

(17)

(18)

Substituting ŵ = (Xl Pm )† Xl w⋆ + (Xl Pm )† ξ and expanding the quadratic form, the cross term vanishes in ξ expectation and we obtain a bias contribution 2 (w⋆ − (Xl Pm )† Xl w⋆ )⊤ vnew,i (19) 20

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

and a variance contribution proportional to ⊤ † vnew,i (Pm X⊤ l Xl Pm ) vnew,i .

(20)

Using the same decomposition and conditional Gaussian argument as in Equation (13), (Xl Pm )† Xl w⋆ = Πw⋆ + AΠv + (Xl Pm )† Ew⋆ ⊥ , we further decompose the bias ⊤ ⊤ ⊤ (w⋆ ⊤ − (Xl Pm )† Xl w⋆ )⊤ vnew,i = w⋆⊤ vnew,i − vnew,i Πw⋆ − Avnew,i Πv − vnew,i (Xl Pm )† Ew⋆ ⊥ ,

(21)

with the squared term being: λ − 1 ⋆⊤ ⊤ ⊤ ⊤ (w⋆ ⊤ vnew,i − vnew,i Πw⋆ − (w P⊥ v)(vnew,i Πv))2 + (vnew,i (Xl Pm )† Ew⋆ ⊥ )2 | {z } λ̃ | {z } quadratic in E

E–independent

⊤ − 2 (w⋆ ⊤ vnew,i − vnew,i Πw⋆ − |

λ − 1 ⋆⊤ ⊤ ⊤ (w P⊥ v)(vnew,i Πv))(vnew,i (Xl Pm )† Ew⋆ ⊥ ) . λ̃ {z } cross term

Step 2: E-independent term.

The term linear in E is made up of terms that have almost sure limits. In particular: a.s. √ ⊤ vnew,i Πw⋆ −−→ w̄∗ π̄w⋆ ,vnew,i by Lemma 7 a.s.

⊤ vnew,i Πv −−→ π̄v,vnew,i a.s. √ ⋆⊤ w P⊥ v −−→ w̄∗ p̄⊥,w⋆ ,v a.s.

⋆⊤

λ̃ −−→ λ̄ √ ∗ vnew,i −−→ w̄ ρw⋆ ,vnew,i a.s.

by Lemma 7 by Lemma 5 by Lemma 4 by model assumptions

Thus, by the continuous mapping theorem, we have the almost sure convergence of the entire E-independent term to 2 λ−1 w̄∗ ρw⋆ ,vnew,i − π̄w⋆ ,vnew,i − p̄⊥,w⋆ ,v π̄v,vnew,i . λ̄ Step 3: Cross term linear in E.

The linear E term has the form T = c⊤ Ed,

where

c = C (Xl Pm )†⊤ vnew,i ,

d = w⋆ ⊥ ,

and C is the deterministic prefactor defined in Step 2, which converges almost surely to an O(1) limit. This term has the same structure as the term linear in E in Theorem 1, Step 3, with the vector w⋆ − Πw⋆ − AΠv replaced by the unit vector vnew,i and multiplied by the scalar C = O(1). The same arguments apply and we conclude that a.s. T −−→ 0. Step 4: Quadratic term in E.

Define ⊤ †⊤ ⊤ † ⋆ Q := w⋆ ⊤ ⊥ E (Xl Pm ) vnew,i vnew,i (Xl Pm ) Ew ⊥ .

Let a := w⋆ ⊥ ,

⊤ A := (Xl Pm )†⊤ vnew,i vnew,i (Xl Pm )† .

Conditioned on (Xm , Um ), the vector a and matrix A are independent of E, so Q = a⊤ E⊤ AEa is a quadratic Gaussian form. This is exactly the same object as in Theorem 1, Step 5, with a different choice of A. The concentration argument therefore carries over verbatim once we verify that the same norm bounds hold. 21

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

First, as in Theorem 1, Step 3, a⊤ Σ̄a = O(1) Next, since A is rank one,

almost surely.

∥A∥ = ∥(Xl Pm )† vnew,i ∥2 ,

∥A∥2F = ∥A∥2 .

Using the bounds on ∥(Xl Pm )† ∥op established in Theorem 1, Step 3, together with ∥vnew,i ∥ = 1, we obtain 1 1 2 almost surely. ∥A∥ = O , ∥A∥F = O 2 , p p The conditions required for the Hanson–Wright concentration used in Theorem 1, Step 5 are satisfied. Therefore, by the same argument as in Theorem 1, Step 5, a.s.

Q − EE [Q | Xm , Um ] −−→ 0. Finally, the conditional expectation is 1 a.s. EE [Q | Xm , Um ] = (a Σ̄a)Tr(A) = O −−→ 0 p ⊤

Step 5: Variance contribution. Finally, we consider the term from Equation (20) and note it is the same as the norm of the matrix A from Step 4: 1 ⊤ † † 2 vnew,i (Pm X⊤ X P ) v = ∥(X P ) v ∥ = ∥A∥ = O . l m new,i l m new,i l p Thus, this term converges to 0: a.s.

⊤ † vnew,i (Pm X⊤ l Xl Pm ) vnew,i −−→ 0.

Combining. Thus, in the high-dimensional limit the only additional contribution relative to the estimation error is the spike bias term, and the generalisation error converges to 2 X λ−1 gen est ∗ E∞ = E∞ + w̄ (νi − 1) ρw⋆ ,vnew,i − π̄w⋆ ,vnew,i − p̄⊥,w⋆ ,v π̄v,vnew,i + σ2 . λ̄ i

A.5

Proof for the training error

Proof of Theorem 3. Let H = (Xl Pm )(Xl Pm )† ⋆

and recall that the target vector is yd = Xl w + ξ. Then, the residual vector is: yd − Xl ŵ = yd − (Xl Pm )(Xl Pm )† yd − (Xl P⊥ )(Xl Pm )† yd = (Inl − H)yd = (Inl − H)(Xl w⋆ + ξ),

(22)

where we note that the term multiplying two projections is 0: (Xl P⊥ )(Xl Pm )† = Xl U⊥ U⊤ Um (Xl Um )† = 0. | ⊥{z } 0

The training error splits into two terms since the noise ξ is independent and zero-mean: h i 1 2 E train = Eξ ∥(Inl − H)Xl w⋆ + (Inl − H)ξ∥ nl h i 1 1 2 2 = ∥(Inl − H)Xl w⋆ ∥ + Eξ ∥(Inl − H)ξ∥ . n n {z } |l {z } |l Bias2

Variance

(23)

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

We start by computing the bias. Using the decomposition Xl = Xl Pm + Xl P⊥ and the fact that (I − H)Xl Pm = 0, the bias simplifies to: 1 2 Bias2 = ∥(Inl − H)Xl P⊥ w⋆ ∥ . (24) nl Substituting the conditional Gaussian decomposition, we have ⊤ ⊤ Xl P⊥ = Xm Σ−1 mm Σm⊥ U⊥ + EU⊥ ,

where E is independent of Xm . Thus, we obtain ⊤ ⋆ ⊤ ⋆ (Inl − H)Xl P⊥ w⋆ = (Inl − H)Xm Σ−1 mm Σm⊥ U⊥ w + (Inl − H)EU⊥ w = (Inl − H)Ew⋆ ⊥ ,

(25)

⋆ ⋆ ⋆ 2 2 where we used (Inl −H)Xm = 0 and defined w⋆ ⊥ = U⊤ ⊥ w . Note that the entries of Ew ⊥ are i.i.d. N (0, ∥w ∥ σeff ), with λ − 1 (v⊤ P⊥ w⋆ )2 ∥P⊥ w⋆ ∥2 2 + . (26) σeff = ⋆ 2 ∥w ∥ ∥w⋆ ∥2 λ̃

Conditioned on Xm , H is deterministic, and the term ∥(Inl − H)Ew⋆ ⊥ ∥2 is a quadratic form z ⊤ (Inl − H)z where 2 z ∼ N (0, ∥w⋆ ∥2 σeff Inl ). Hence, we have 2 min(m, nl ) 2 ⋆ 2 σeff ⋆ 2 2 EE Bias | Xm = ∥w ∥ Tr(Inl − H) = ∥w ∥ σeff 1 − , (27) nl nl where the last equality holds almost surely (w.r.t. the randomness in Xm ). Next, to compute the variance, recall that the label noise ξ is isotropic N (0, σ 2 Inl ). Thus, we have ⊤ σ2 min(m, nl ) 1 2 Eξ ξ (Inl − H)ξ = Tr(Inl − H) = σ 1 − , nl nl nl

(28)

where the last equality again holds almost surely (w.r.t. the randomness in Xm ). Note that 1−

min(m, nl ) = max(0, 1 − γeff ). nl a.s.

2 2 To conclude, it suffices to recall from the proof of Theorem 1 that σeff −−→ σ̄eff , and the argument is complete.

A.6

Auxiliary lemmas

Lemma 1 (Limiting spectral distribution of Xm ). Consider the setup of Section 3, and define Xm := Xl Um . Then the empirical spectral measure of n1l X⊤ m Xm converges weakly almost surely to the Marchenko–Pastur law with aspect ratio γeff , i.e., 1 X⊤ X F̂ nl m m ⇒ Fγeff almost surely. Proof. Fix a realization of the full sequence (Um )p≥1 . Conditional on this realization, the rows of Xm = Xl Um are i.i.d. Gaussian in Rm with covariance Σmm := U⊤ m ΣUm . Using the form of Σ from Assumption 2, we obtain ⊤ ⊤ ⊤ Σmm = U⊤ I + (λ − 1)vv Um = Im + (λ − 1)(U⊤ p m m v)(Um v) . Hence, Σmm − Im has rank at most one. Since Σmm ≻ 0, define G ∈ Rnl ×m and W ∈ Rm×m by G := Xm Σ−1/2 mm

W :=

1 ⊤ G G. nl

Conditional on Um , G has i.i.d. N (0, 1) entries. The sample covariance can be rewritten: 1 ⊤ 1 ⊤ 1/2 1/2 Xm Xm = Σ1/2 G G Σ1/2 mm mm = Σmm WΣmm . nl nl 23

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

1/2

Since Σmm = Im + R with rank(R) ≤ 1, the spectral decomposition shows that Σmm − Im also has rank at most one. Consequently, 1/2 1/2 1/2 1/2 Σ1/2 mm WΣmm − W = (Σmm − Im )WΣmm + W(Σmm − Im ), and thus 1/2 rank Σ1/2 mm WΣmm − W ≤ 2. Let F̂ A denote the empirical spectral distribution function of a real symmetric matrix A. By the rank inequality (e.g., see Theorem A.43 in [67]), 1 2 X⊤ X sup F̂ nl m m (θ) − F̂ W (θ) ≤ → 0. p→∞ m θ∈R Since W is a standard Wishart matrix of size m × m with aspect ratio nml → γeff , by the Marchenko–Pastur theorem [66], F̂ W ⇒ Fγeff almost surely, conditional on (Um )p≥1 . n 1 ⊤ o X X Let A := F nl m m ⇒ Fγeff . Taken together, the arguments above shows that: P(A | (Um )p≥1 ) = 1. We now remove the conditioning on (Um )p≥1 . By the law of total probability, P(A) = E[P(A | (Um )p≥1 )] = E[1] = 1. This proves that 1

F̂ nl

X⊤ m Xm

⇒ Fγeff

almost surely.

Lemma 2 (Deterministic equivalent for the pseudoinverse trace). Consider the setup of Section 3, and define Xm := Xl Um . Then

 γeff  1 − γ , ⊤ a.s. eff Tr (Xm Xm )† −−→ 1   , γeff − 1

γeff < 1, γeff > 1.

Proof. We have m 1 1 1 1 † † † = · Tr ( X⊤ Tr (X⊤ = Tr ( X⊤ m Xm ) m Xm ) . m Xm ) nl nl nl m nl

(29)

By Lemma 1, the empirical spectral distribution of n1l X⊤ m Xm converges weakly almost surely to the Marcenko–Pastur law with aspect ratio γeff . If γeff < 1, its support is bounded away from zero, and the claim follows by continuity of x 7→ 1/x. If γeff > 1, there is an additional atom at zero; however, the pseudo-inverse removes this null space contribution, and the nonzero spectrum remains bounded away from zero, so the same argument applies on the effective support. Therefore, Z Z 1 ⊤ 1 1 1 1 a.s. ⊤ † n l Xm Xm Tr (X Xm ) = dF (x) −−→ dFγeff (x). m nl m x x Recall that  1  Z , γ < 1,  1 1−γ dFγ = 1  x  , γ > 1. γ(γ − 1) Substituting into (29) and using γ = γeff gives  γeff   1 − γ , γeff < 1, ⊤ eff † a.s. Tr (Xm Xm ) −−→  1 , γ > 1,  eff γeff − 1 which concludes the proof. 24

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

Lemma 3 (Spectral measure of the overlap). Consider the setup of Section 3, and let (ui )pi=1 be the eigenvectors of p Su = n1u X⊤ u Xu . Let a ∈ R be a fixed vector with population mass distribution Ga on the eigenvectors of Σ, i.e., ⊤ 2 ⊤ 2 Ga = (a∥a∥v)2 δλ + 1 − (a∥a∥v)2 δ1 . Define the empirical CDF Ĝsample (θ) = a

p 1 X ⊤ 2 (a ui ) 1{θ ≤ λi (Su )}. ∥a∥2 i=1

Then, Ĝsample ⇒ Gsample almost surely, where Gsample is such that (i) its support satisfies supp(Gsample ) ⊆ a a a a Sγu ∪ {θ∗ } ∪ {0} with Sγu being the support of the Marcenko-Pastur density with aspect ratio γu ; (ii) it has a continuous density gasample (θ) on the interior of Sγu Z τ γu sample dGa (τ ) fγu (θ); ga (θ) = τ 2 − τ (1 − γu + θ) + θ (iii) it has an atom at {θ∗ } where θ∗ =

λ(λ − 1 + γu ) is the location of the spike λ in the sample covariance spectrum λ−1 Gsample ({θ∗ }) = a

γu 1 − (λ−1) 2 γu 1 + (λ−1)

Ga ({λ});

and (iv) it may have an atom at 0 of weight  G γu < 1 Z a ({0}) γu − 1 Gsample ({0}) = a dGa (τ ), γu > 1.  γu − 1 + τ Proof. The result follows by specializing Lemma 3 in [45] to the spiked covariance model of Assumption 2. Lemma 4 (Deterministic limit of ∥Pm a∥2 , ∥P⊥ a∥2 ). Consider the setup of Section 3, and let a ∈ Rp be a fixed vector with population mass distribution Ga on the eigenvectors of Σ. Let Pm be the projection onto the top m eigenvectors of the prior sample covariance Su , with m/p → α. Let Gsample denote the deterministic limiting weighted spectral a measure defined in Lemma 3. Let λt be the deterministic bulk threshold corresponding to retaining an asymptotic fraction α of principal components Fγu (λt ) = 1 − α. Then, almost surely,

∥P⊥ a∥2 ∥a∥2

where λ− = (1 − and

√

Z λ t    gasample (λ), γu < 1,    λ −   Z λt  a.s. −−→ Gsample ({0}) + gasample (λ), γu > 1, α < γ1u , a   λ −    1 − α sample    ({0}), γu > 1, α ≥ γ1u ,  1 − 1 Ga γu

(30)

γu )2 is the infimum of the support of the Marcenko-Pastur distribution as defined in Equation (5) ∥Pm a∥2 = ∥a∥2 − ∥P⊥ a∥2 .

(31)

Proof. Let Ĝsample (θ) = a

p 1 X ⊤ 2 (a ui ) 1{θ ≤ λi (Su )} ∥a∥2 i=1

be the weighted empirical spectral measure of Su . By Lemma 3, Ĝsample converges almost surely to the deterministic a measure Gsample . a Since ∥P⊥ a∥2 /∥a∥2 is exactly the mass of Ĝsample on the discarded eigenspaces, its limit is obtained by integrating a Gsample over the discarded part of the spectrum. a 25

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

If γu < 1, there is no nullspace, so only bulk directions are discarded. This gives Z λt ∥P⊥ a∥2 a.s. −−→ gasample (λ). ∥a∥2 λ− If γu > 1 and α < 1/γu , fewer than all positive-eigenvalue directions are retained, so the entire nullspace is discarded as well. Hence Z λt ∥P⊥ a∥2 a.s. sample −−→ Ga ({0}) + gasample (λ). ∥a∥2 λ− Finally, if γu > 1 and α ≥ 1/γu , all positive-eigenvalue directions are retained, and the retained subspace is completed by adding m − rank(Su ) = m − nu orthonormal directions chosen uniformly at random in the nullspace, independently of a. Therefore only the fraction p − m a.s. 1 − α −−→ p − nu 1 − γ1u of the nullspace remains discarded. Since the total asymptotic mass of a in the nullspace is Gsample ({0}), we obtain a ∥P⊥ a∥2 a.s. 1 − α sample −−→ G ({0}). ∥a∥2 1 − γ1u a The identity ∥Pm a∥2 = ∥a∥2 − ∥P⊥ a∥2 is immediate. Lemma 5 (Deterministic limit of a⊤ Pm b, a⊤ P⊥ b). Consider the setup of Section 3, let a, b ∈ Rp be fixed vectors and let Pm denote the projection onto the top m principal components of Su , with m/p → α. Define the signed population spectral measure X 1 Ga,b (τ ) = (a⊤ vj )(b⊤ vj ) 1{λj ≤ τ }, (32) ∥a∥∥b∥ j where (vj , λj ) are the population eigenpairs of Σ. Let Gsample denote the corresponding deterministic limiting signed a,b sample spectral measure onto the eigenpairs of Su . Then, a⊤ Pm b = a⊤ b − a⊤ P⊥ b,

(33)

Z λt  sample  ga,b (λ), γu < 1,    λ−   Z λt a⊤ P⊥ b a.s.  sample sample −−→ Ga,b ({0}) + ga,b (λ), γu > 1, α < 1/γu ,  ∥a∥∥b∥  λ −    1 − α sample   G ({0}), γu > 1, α ≥ 1/γu .  1 − γ1u a,b

(34)

and, almost surely,

sample Moreover, Gsample ({0}) and the bulk density of ga,b are given by the formulas of Lemma 3 with dGa replaced by a,b the signed measure dGa,b .

Proof. Define the signed mass density onto the sample eigenvectors of Su p

X 1 (a⊤ ui )(b⊤ ui ) 1{θ ≤ λi (Su )}, ∥a∥∥b∥ i=1

(35)

1 dĜsample (λ) = a⊤ (Su − zI)−1 b = Tr ab⊤ (Su − zI)−1 . a,b λ−z

(36)

Ĝsample (θ) := a,b so that its Stieltjes transform is Z mĜsample (z) = a,b

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

By Rubio and Mestre [98, Theorem 1], for any deterministic matrix Θ with bounded trace norm (i.e., Tr (Θ⊤ Θ)1/2 = ∥Θ∥tr ), Tr Θ(Su − zI)−1 − δΘ (z) → 0 a.s. where δΘ (z) is a deterministic function of z. Applying this with Θ = ab⊤ and repeating the argument from [45] (in their paper, see Lemma 1 proven in Appendix A) yields the following almost sure weak convergence Ĝsample ⇒ Gsample , a,b a,b where the Stieltjes transform of Gsample is a,b mGsample (z) = − a,b

1 z

1 dGa,b (τ ). 1 + τ mγu (z)

(37)

Here, the companion Stieltjes transform is defined as mγu (z) = γu mγu (z) − (1 − γu ) z1 , and mγu is a Stieltjes transform of a Marcenko-Pastur density fγu . This is exactly the same expression as in Lemma 1 in [45] with the sample population measure Ga replaced by the signed measure Ga,b . Hence, the atom at 0 and the bulk density of ga,b ({0}) are obtained from the formulas of Lemma 3 by the same substitution. Now the mass of Ĝsample on the discarded a,b eigenspaces is p X a ⊤ P⊥ b = (a⊤ ui )(b⊤ ui ) (38) i=m+1

noting that the eigenvectors ui are sorted according to their eigenvalues, such that the set {m + 1, . . . , p} contains the indices of the p − m smallest eigenvalues. Therefore, the same geometric argument as in Lemma 4 gives the three regimes above: discarded bulk only when γu < 1; bulk plus the full atom at 0 when γu > 1 and α < 1/γu ; and, when γu > 1 and α ≥ 1/γu , only the fraction 1−α 1 − γ1u of the nullspace mass remains discarded, by the same random nullspace-completion argument. Finally, a⊤ Pm b = a⊤ b − a⊤ P⊥ b, which proves the result. Lemma 6 (Deterministic limit of the effective projection ∥Πa∥2 ). Consider the setup of Section 3, let Xm = Xl Um ∈ Rnl ×m and define p×p Πm := X†m Xm ∈ Rm×m , Π := Um Πm U⊤ . m ∈R Let the eigendecomposition of Σmm be Σmm := U⊤ m ΣUm with (ṽi , λ̃i ) denoting the eigenpairs. For any fixed vector a ∈ Rp with Pm a ̸= 0, define U⊤ ma ã := ∈ Rm , ∥Pm a∥ and let G̃ã be the mass distribution of ã onto the eigenvectors of Σmm : X G̃ã (τ̃ ) = (ã⊤ ṽi )2 1{λ̃i ≤ τ̃ }. i

Then, it holds almost surely G̃ã (τ ) ⇒

p̄2a,v 1− p̄a,a p̄v,v

! δ1 +

p̄2a,v δ , p̄a,a p̄v,v λ̄

and ⊤

  1,

a Πa a.s. Z −−→ p̄a,a ·  ∥a∥2 1 −

γeff < 1, γeff − 1 dG̃ã (τ̃ ), γeff − 1 + τ̃

with p̄a,a , p̄v,v , p̄a,v defined via Equations (30) and (34). 27

(39) γeff > 1.

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

Proof. We start by computing explicitly G̃ã . Recall that ⊤ ⊤ ⊤ Σmm = U⊤ m ΣUm = Im + (λ − 1)(Um v)(Um v) .

Then, Σmm has m − 1 eigenvalues equal to 1 and the remaining one is λ̃ = 1 + (λ − 1)∥Pm v∥2

⇒

λ̄ = 1 + (λ − 1)p̄v,v , U⊤ v

U⊤ v

m where the convergence occurs almost surely. The corresponding unit eigenvector is ṽ = ∥Um ⊤ v∥ = ∥P v∥ . Therefore m m

the spectral measure G̃ã of ã =

U⊤ ma ∥Pm a∥

with respect to the eigenbasis of Σmm is supported on {1, λ̃}, and hence G̃ã = 1 − (ã⊤ ṽ)2 δ1 + (ã⊤ ṽ)2 δλ̃ , with (ã⊤ ṽ)2 =

⊤ ⊤ 2 2 (a⊤ Um U⊤ (a⊤ Pm v)2 ((U⊤ m v) m a) (Um v)) = = . 2 2 2 2 ∥Pm a∥ ∥Pm v∥ ∥Pm a∥ ∥Pm v∥ ∥Pm a∥2 ∥Pm v∥2

Thus, G̃ã is a two-point measure, with the original atom at λ replaced by an atom at λ̃, and with weight determined by the projected overlap. Almost surely, it holds that: ! p̄2a,v p̄2a,v G̃ã (τ ) ⇒ 1 − δ1 + δ . p̄a,a p̄v,v p̄a,a p̄v,v λ̄ Case m < nl . In this regime, Xm ∈ Rnl ×m has full column rank almost surely, so Πm = Im . Hence ã⊤ Πm ã = 1, giving a⊤ Πa a.s. −−→ p̄a,a , ∥a∥2 2

a.s.

m a∥ as ∥P∥a∥ −−→ p̄a,a from Lemma 4. 2

Case m > nl .

In this regime, Xm has a nullspace of dimension m − nl , and ã⊤ Πm ã = 1 − ã⊤ Πm,null ã.

Conditioned on Um , the rows of Xm are Gaussian with covariance Σmm = U⊤ m ΣUm . Therefore, by the nullspace mass formula from Lemma 3 applied to the m-dimensional model with aspect ratio γeff = m/nl > 1, Z γeff − 1 a.s. ⊤ ã Πm,null ã −−→ dG̃ã (τ̃ ). (40) γeff − 1 + τ̃ 2

a.s.

m a∥ Finally, noting that ∥P∥a∥ −−→ p̄a,a according to Lemma 4 and combining the two terms via the continuous mapping 2 theorem gives the required result.

Lemma 7 (Deterministic limit of the bilinear effective projection a⊤ Πb). Consider the setup of Section 3, let a, b ∈ Rp , with Pm a ̸= 0, Pm b ̸= 0, and define ã :=

U⊤ ma ∈ Rm , ∥Pm a∥

b̃ :=

U⊤ mb ∈ Rm . ∥Pm b∥

Let G̃ã,b̃ be the signed joint mass distribution of ã, b̃ onto the eigenvectors ṽi of Σmm : X G̃ã,b̃ (τ̃ ) = (ã⊤ ṽi )(b̃⊤ ṽi ) 1{λ̃i ≤ τ̃ }. i

Then, it holds almost surely G̃ã,b̃ (τ ) ⇒

p̄a,b 1/2 1/2 p̄a,a p̄b,b

−

p̄a,v p̄b,v 1/2 1/2 p̄a,a p̄b,b p̄v,v

! δ1 +

p̄a,v p̄b,v δλ̄ , 1/2 1/2 p̄a,a p̄b,b p̄v,v

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

and

  p̄a,b ,

γeff < 1,

a⊤ Πb a.s. Z −−→ 1/2 1/2  ∥a∥∥b∥ p̄a,b − p̄a,a p̄

γeff − 1 dG̃ã,b̃ (τ̃ ), γeff − 1 + τ̃ with p̄a,a , p̄b,b , p̄v,v , p̄a,v , p̄b,v defined via Eq. (30) and (34). b,b

(41) γeff > 1.

Proof. The proof follows the steps of the proof of Lemma 6. The mass distribution G̃ã,b̃ can be written as G̃ã,b̃ = ã⊤ b̃ − (ã⊤ ṽ)(b̃⊤ ṽ) δ1 + (ã⊤ ṽ)(b̃⊤ ṽ)δλ̃ , with (ã⊤ ṽ)(b̃⊤ ṽ) = which converges almost surely to

⊤ ⊤ ⊤ ⊤ ⊤ (a⊤ Pm v)(b⊤ Pm v) ((U⊤ m a) (Um v))((Um b) (Um v)) = , 2 ∥Pm a∥ ∥Pm b∥ ∥Pm v∥ ∥Pm a∥ ∥Pm b∥ ∥Pm v∥2 p̄a,v p̄b,v 1/2 1/2

p̄a,a p̄b,b p̄v,v

ã⊤ b̃ =

and ⊤ ⊤ (U⊤ p̄a,b m a) (Um b) a.s. −−→ 1/2 1/2 ∥Pm a∥∥Pm b∥ p̄a,a p̄

b,b

This readily gives the almost sure limit of G̃ã,b̃ (τ ). Next, observe 1 ∥Pm a∥ ∥Pm b∥ ⊤ a⊤ Πb = (U⊤ a)⊤ Πm (U⊤ ã Πm b̃. m b) = ∥a∥∥b∥ ∥a∥∥b∥ m ∥a∥ ∥b∥

(42)

Combining this with the almost sure limits from Lemma 4, with the same steps as in the proof of Lemma 6 and with the mass distribution of G̃ã,b̃ proves the lemma. A.7

Phase transition for γl > 1

Here, we derive parametric equations for the two phase transitions we observe in the phase diagrams for α in Figure 3. For small downstream tasks, γl > 1, there is a sharp phase transition due to competing local minima at different sides of the double descent (local minimum at large vs small α), which we tackle first. We find a proxy of the parametric curve of this phase transition by comparing the values of the error at the endpoints in α, and check when gen gen E∞ (α → 0) = E∞ (α = 1). Note that the left endpoint is not exactly 0, as we always retain the first principal component of the data, but in the high-dimensional regime p → ∞, α approaches 0. We define signed normalized overlaps as ρa,b :=

a⊤ b , ∥a∥ ∥b∥

ηa,b := ρ2a,b .

In particular, ηa,v denotes the squared population-spike overlap of a fixed vector a with the spike eigenvector v. For bilinear quantities, the signs of the corresponding ρ’s must be retained. Define the limiting value of the population and sample spike eigenvector overlap (v⊤ u1 )2 → c, which depends on whether the spike is below or above the BBP transition [45]: ( 1− γu (λ−1)2 γu < (λ − 1)2 γu 1+ λ−1 . c := 0 γu ≥ (λ − 1)2 Endpoint values for the eigenvalues and projections. At the left endpoint, meaning m = 1 and hence α → 0, the retained direction is the empirical spike direction. Therefore, the limiting value of projections for any fixed vectors a, b are given by p̄⊥,a,a (0) = 1 − ηa,v c, p̄⊥,a,b (0) = ρa,b − ρa,v ρb,v c.

p̄a,a (0) = ηa,v c, p̄a,b (0) = ρa,v ρb,v c,

At the right endpoint, α = 1, the projection is the identity, Pm = Ip . Hence, we have p̄a,a (1) = 1, p̄a,b (1) = ρa,b ,

p̄⊥,a,a (1) = 0 p̄⊥,a,b (1) = 0. 29

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

Next, let us analyse the effective limiting projections π̄. At α = 0, the effective aspect ratio satisfies γeff = 0. Hence the compressed task design is overdetermined, and the task projection is exact inside the retained subspace: π̄a,a (0) = p̄a,a (0) = ηa,v c, π̄a,b (0) = p̄a,b (0) = ρa,v ρb,v c. At α = 1, we have λ̄ = λ and γeff = γl . Therefore   γl < 1, 1, π̄a,a (1) = λ 1  ηa,v + (1 − ηa,v ) , γl > 1. γl − 1 + λ γl   ρ , γl < 1,  a,b π̄a,b (1) = λ 1  ρa,v ρb,v + (ρa,b − ρa,v ρb,v ) , γl > 1. γl − 1 + λ γl Furthermore, note that at α = 1 we have λ̄ = λ, and for α = 0 define: λ̄0 = 1 + (λ − 1)c. At α = 1, we have Pm = Ip , hence

Endpoint values for the generalisation error. p̄⊥,w⋆ ,w⋆ = 0,

p̄⊥,w⋆ ,v = 0,

2 σ̄eff = 0.

Therefore, gen E∞ (1) = w̄∗ (1 − π̄w⋆ ,w⋆ (1)) +

T X 2 σ2 ∗ + w̄ (νi − 1) ρw⋆ ,vnew,i − π̄w⋆ ,vnew,i (1) + σ 2 , |γl − 1| i=1

where for the relevant regime for this phase transition (γl > 1) λ 1 + (1 − ηw⋆ ,v ) , γl − 1 + λ γl 1 λ + ρw⋆ ,vnew,i − ρw⋆ ,v ρvnew,i ,v . π̄w⋆ ,vnew,i (1) = ρw⋆ ,v ρvnew,i ,v γl − 1 + λ γl π̄w⋆ ,w⋆ (1) = ηw⋆ ,v

est At α = 0 ((γeff = 0), the variance term in E∞ vanishes. Moreover,

π̄w⋆ ,w⋆ (0) = ηw⋆ ,v c,

π̄v,v (0) = c,

p̄⊥,w⋆ ,v (0) = ρw⋆ ,v (1 − c).

Thus, we have (λ − 1)2 2 1 − ηw⋆ ,v c + ηw⋆ ,v (1 − c) c λ̄20 2 T X λ−1 ∗ + w̄ (νi − 1) ρw⋆ ,vnew,i − ρw⋆ ,v ρv,vnew,i c − ρw⋆ ,v (1 − c)ρv,vnew,i c + σ 2 . λ̄ 0 i=1

gen E∞ (0) = w̄∗

Noting that λ̄0 = 1 + (λ − 1)c, this simplifies further: " X 2 # T λ λ λ gen ∗ E∞ (0) = w̄ 1 − ηw⋆ ,v c 2− + (νi − 1) ρw⋆ ,vnew,i − ρw⋆ ,v ρv,vnew,i c + σ2 . λ̄0 λ̄0 λ̄ 0 i=1 ∗

Let S := w̄ σ 2 , and define ri := ρw⋆ ,vnew,i ,

qi := ρw⋆ ,v ρv,vnew,i ,

η := ηw⋆ ,v .

Then the normalized α = 0 endpoint is λ B0 (c) := 1 − ηc λ̄0

λ 2− λ̄0

2 T X λc + (νi − 1) ri − qi . λ̄0 i=1

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

Thus

gen 1 E∞ (0) = B0 (c) + . ∗ w̄ S

For α = 1 and γl > 1, define λ , γl − 1 + λ

A(γl ) :=

B(γl ) :=

1 . γl

Then B1 (γl ) := 1 − [ηA(γl ) + (1 − η)B(γl )] +

T X

(νi − 1) [ri − qi A(γl ) − (ri − qi )B(γl )] .

i=1

The phase boundary in the (γu , γl )-plane is therefore B0 (c(γu )) = B1 (γl ) + Specialization to Σnew = Σ.

1 S(γl − 1)

In this case, T = 1, ν1 = λ, and r1 = q1 = ρw⋆ ,v ,

Hence

(γl > 1).

η = ρ2w⋆ ,v .

η cλ2 − 2cλλ̄0 + (λ − 1) λ̄20 − 2λ̄0 λc + λ2 c2 2 λ̄0 ηλ2 c = 1 + η(λ − 1) − , 1 + (λ − 1)c

B0 (c) = 1 +

while B1 = (1 − η)

γl − 1 λγl (γl − 1) . +η γl (γl − 1 + λ)2

Thus, the phase transition is: 1 + η(λ − 1) −

ηλ2 c γl − 1 1 λγl (γl − 1) = (1 − η) +η + 1 + (λ − 1)c γl (γl − 1 + λ)2 S(γl − 1)

Specialization to Σnew = Σ and w⋆ =

√ w̄∗ v.

Now everything simplifies even further. We have

B0 =

λ(1 − c) , 1 + (λ − 1)c

and

λγl (γl − 1) . (γl − 1 + λ)2 Due to the two regimes of c, the proxy transition splits into two components: ( √ λ λ < 1 + γu λγl (γl − 1) 1 √ + = λ λ ≥ 1 + γu (γl − 1 + λ)2 S(γl − 1) (λ−1)2 γu B1 =

A.8

Phase transition for γl < 1

Next, we study the regime of larger downstream tasks γl < 1, and find the phase transition between the phases α = 1 d gen and α < 1 by finding a parametric curve that describes the loss of stability of α = 1, i.e., when dα E∞ |α=1 = 0. Derivatives of the pretraining projection at α = 1.

Let λ± = (1 ±

Ha (θ) := (1 − ηa,v ) + ηa,v

√

λγu . λ2 − λ(1 − γu + θ) + θ

Since gasample (θ) = Ha (θ)fγu (θ), 31

γu )2 , and define

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

the derivative of p̄⊥,a,a at α = 1 is  −Ha (λ− ),   

γu < 1,

d = p̄⊥,a,a (α) Gsample ({0}) a  dα  α=1− , − 1 − γ1u

γu > 1.

Equivalently,  λγu    − (1 − ηa,v ) + ηa,v λ2 − λ(1 − γu + λ− ) + λ− , d = p̄⊥,a,a (α)  dα γu  α=1−  − (1 − ηa,v ) + ηa,v , γu − 1 + λ

γu < 1, γu > 1.

Therefore

d d =− . p̄a,a (α) p̄⊥,a,a (α) dα dα − α=1 α=1− For signed bilinear projections, replace ηa,v by ρa,v ρb,v , and replace the constant bulk contribution 1 − ηa,v by ρa,b − ρa,v ρb,v .

Derivative of the effective projection at α = 1. For γl < 1, the compressed task design remains overdetermined in a neighbourhood of α = 1. Thus π̄a,a (α) = p̄a,a (α), γl < 1, and hence d d = . π̄a,a (α) p̄a,a (α) dα dα α=1− α=1− Stability of α = 1 for γl < 1.

Assume γl < 1 and let α → 1− . Define d Dw⋆ := . p̄w⋆ ,w⋆ (α) dα α=1−

2 At α = 1, p̄⊥,a,b = 0 and σeff = 0, thus simplifying

d gen γl σ 2 w̄∗ = E∞ (α) − Dw⋆ . 2 dα (1 − γl ) 1 − γl α=1− Consequently, α = 1 is locally stable iff d gen ≤ 0, E (α) dα ∞ α=1− or equivalently

SDw⋆ w̄∗ , S := 2 . 1 + SDw⋆ σ The equality gives the local stability boundary of the full representation:  λγu   (1 − η) + η (λ − 1 + √γ )2 , γu < 1, u Dw⋆ =  γ u  (1 − η) + η , γu > 1, γu − 1 + λ γl ≤

γlcrit = where

S [(1 − η) + ηDv ] , 1 + S [(1 − η) + ηDv ]

λγu , √ (λ − 1 + γu )2 Dv =  γu   , γu − 1 + λ    

γu < 1, γu > 1.

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

A.9

Infinite data limits

We consider what happens to the errors and optimal α when γu → 0 or γl → 0. Prior data limit γu → 0. In this limit, the empirical spike aligns with the population spike, and the retained subspace contains the spike direction together with an α-fraction of the isotropic bulk. For any fixed vectors a, b, we have p̄a,a → ηa,v + (1 − ηa,v )α, p̄⊥,a,a → (1 − ηa,v )(1 − α), p̄a,b → ρa,v ρb,v + α (ρa,b − ρa,v ρb,v ) , p̄⊥,a,b → (1 − α) (ρa,b − ρa,v ρb,v ) . In particular, p̄v,v → 1,

p̄⊥,v,v → 0,

λ̄ → λ.

For the effective projections π̄, the behaviour depends on γeff . If γeff < 1, then π̄a,a → p̄a,a ,

π̄a,b → p̄a,b .

If γeff > 1, define Aλ :=

λ , γeff − 1 + λ

A1 :=

1 , γeff

then π̄a,a → ηa,v Aλ + (1 − ηa,v )αA1 , π̄a,b → ρa,v ρb,v Aλ + α (ρa,b − ρa,v ρb,v ) A1 . Errors in the limit γu → 0.

In the infinite-prior-data limit, the empirical spike is recovered exactly. Hence p̄v,v → 1,

p̄⊥,w⋆ ,v → 0,

λ̄ → λ,

and p̄w⋆ ,w⋆ → η + (1 − η)α,

p̄⊥,w⋆ ,w⋆ → (1 − η)(1 − α).

Therefore 2 σ̄eff → (1 − η)(1 − α).

Since the leak term vanishes, est E∞ → w̄∗ (1 − π̄w⋆ ,w⋆ ) +

min{γeff , 1} ∗ w̄ (1 − η)(1 − α) + σ 2 , |γeff − 1|

where   η + (1 − η)α, π̄w⋆ ,w⋆ = λ 1  η + (1 − η) , γeff − 1 + λ γl

γeff < 1, γeff > 1.

Moreover, taking Σnew = Σ, we obtain gen est E∞ → E∞ + w̄∗ (λ − 1)η

γeff − 1 γeff − 1 + λ

+ σ2 .

Finally, the training error satisfies train E∞ → σ 2 + w̄∗ (1 − η)(1 − α) max{0, 1 − γeff }. 33

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

Downstream-data limit γl → 0.

For any fixed α > 0, we have γeff = αγl → 0, hence π̄a,a → p̄a,a , π̄a,b → p̄a,b .

Moreover, the variance factor satisfies αγl min{γeff , 1} = → 0, |γeff − 1| 1 − αγl so all downstream finite-sample variance contributions vanish. The pretraining projections p̄ remain determined by γu . At α = 1, p̄a,a = 1,

p̄⊥,a,a = 0,

π̄a,a → 1,

p̄a,b = ρa,b ,

p̄⊥,a,b = 0,

π̄a,b → ρa,b .

and At the spike-retaining endpoint α → 0, π̄a,a (0) → p̄a,a (0), with p̄ determined by the prior overlap c(γu ). Errors in the limit γl → 0.

π̄a,b (0) → p̄a,b (0),

For fixed α > 0, γeff = αγl → 0. Hence π̄a,a → p̄a,a , π̄a,b → p̄a,b ,

and

min{γeff , 1} → 0. |γeff − 1|

Therefore the estimation error satisfies (λ − 1)2 2 est E∞ → w̄∗ 1 − p̄w⋆ ,w⋆ + p̄ p̄ ⊥,w⋆ ,v v,v . λ̄2 The generalisation error satisfies gen est E∞ → E∞ + w̄∗

2 T X λ−1 ⋆ ⋆ ⋆ (νi − 1) ρw ,vnew,i − p̄w ,vnew,i − p̄⊥,w ,v p̄v,vnew,i + σ2 . λ̄ i=1

(43)

λ−1 2 p̄⊥,w⋆ ,w⋆ + p̄ . ⋆ λ̄ ⊥,w ,v

(44)

The training error satisfies train E∞ → σ 2 + w̄∗

Specialization to Σnew = Σ. We further specialize the result in the limit to the setting where the validation data is the same as pretraining and downstream data. Corollary 3. Assume γu = p/nu → 0, so that the empirical spike direction converges to the population spike v. Let η = ηw⋆ ,v , γeff = αγl , S = w̄∗ /σ 2 . Then the estimation error becomes min{γeff , 1} ∗ est E∞ → w̄∗ (1 − π̄w⋆ ,w⋆ ) + w̄ (1 − η)(1 − α) + σ 2 , |γeff − 1| where   γeff < 1, η + (1 − η)α, π̄w⋆ ,w⋆ = λ α  η + (1 − η) , γeff > 1. γeff − 1 + λ γeff For matched test covariance Σnew = Σ,   0, gen est ∗ E∞ → E∞ + w̄ (λ − 1)η γ − 1 2 eff   , γeff − 1 + λ Finally,

γeff < 1, + σ2 . γeff > 1,

train E∞ → σ 2 + w̄∗ (1 − η)(1 − α) max{0, 1 − γeff }. 34

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

Figure 5: Generalisation error of PCR as a function of the number of retained components α = m p . Solid lines show the theoretical prediction; dashed lines show the empirical mean ± one standard deviation across trials. Each column varies one parameter while holding the others fixed: (a,e) varying the pre-training dataset size nu with nl fixed; (b,f) varying the downstream dataset size nl with nu fixed; (c,g) varying the signal-to-noise ratio SNR; (d,h) varying the spike eigenvalue λ. The two rows correspond to two regimes: smaller (top) and larger (bottom) values of nu and nl . Results are shown for a rank-1 spiked covariance model with perfectly aligned signal (η = 1). Consequently, the phase transitions in this limit can either be derived by taking derivatives and limiting values α, or by substituting the limits directly in Corollaries 1 and 2. In this regime, the phase boundaries become independent of γu and depend only on η, SNR, λ, and γl . Corollary 4 (Phase transitions for γu → 0). In the limit γu → 0, the spike direction is recovered exactly. Hence, c(γu ) → 1 and Dv (γu ) → 0. Thus, the transition curve when γl > 1 simplifies to 1 1 λγl (γl − 1) + . =η 2 γl (γl − 1 + λ) S(γl − 1) Furthermore, when γl < 1, the stability boundary reduces to (1 − η)

γlcrit = B

S(1 − η) . 1 + S(1 − η)

Theory-experiment validation

Below, we present additional plots validating the theoretical predictions through numerical experiments. We introduce an overlap measure for the task-spike alignment θ, which can be defined through the overlap η as (w⋆ ⊤ v)2 = η = (cos θ)2 . ∥w⋆ ∥2 B.1

Error curves

Figure 5 illustrates the range of generalisation error curves across parameter settings. When the downstream sample size is limited (nl < p), the error exhibits a double descent profile, with the peak occurring at α ≈ 1/γl . Increasing the size of the pretraining dataset nu uniformly reduces the overall error scale (panels (a,e)), while varying nl shifts the location of the interpolation threshold (panels (b,f)). Increasing the SNR raises the generalisation error (panels (c,g)), as it amplifies the contribution of signal-dependent bias. This increase is non-uniform across α: smaller values of α are affected more strongly, since a larger fraction of the signal lies outside the retained subspace. A similar non-uniform effect is observed when varying the spike strength λ (panels (d,h)). We further decompose these effects by examining the estimation error components in Figure 6. The bias due to missing signal is determined by the downstream sample regime. In the overparametrized regime (nl < p), this bias admits a 35

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

strictly positive lower bound (panel (a)), reflecting the inability to fully reconstruct the signal. In contrast, when nl > p, this component can vanish (panel (e)). Independently of this constraint, the bias decreases with α, since increasing α enlarges the subspace Pm and thus increases the fraction of w⋆ that can be represented. The leak-induced bias follows a similar decreasing trend in α, as it is controlled by p̄⊥,w⋆ ,v , which shrinks as the orthogonal complement P⊥ becomes smaller. However, its dependence on λ is non-monotonic (panels (b,f)): it is negligible for weak spikes (e.g. λ = 1.1), increases at intermediate values (e.g. λ = 2, 5), and decreases again for strong spikes (e.g. λ = 10). This reflects a trade-off between decreasing leak size p̄⊥,w⋆ ,v and increasing the effective component of the leak vector v in the subspace Xl Pm through π̄v . For the parameter ranges considered, the magnitude of the leak term remains small relative to the other components. The variance behaves as expected: in the regime nl < p, it decreases monotonically with λ, reflecting improved signal-to-noise separation (panels (c,g)). Finally, the generalisation error (Figure 7) includes two additional contributions relative to the estimation error. The first is the additive noise variance σ 2 (panels (c,g)). The second is an additional bias term induced by the spike structure in the test distribution. This term combines a residual missing signal component with a correction from the leak, and it is amplified by λ. Its structure leads to a discontinuity at α = γl (panel (b)), inherited from the behaviour of the “missing signal” bias. Overall, the generalisation error closely tracks the estimation error, with the main differences appearing at small α, where these additional bias contributions are most pronounced (panels (e,f)).

Figure 6: Components of the theoretical estimation error as a function of α for different values of λ. The total error (panels (d, h)) is decomposed into the missing signal term (panels (a, e)), the leak term (panels (b, f)), the variance ⋆ 2 term (panels (c, g)). All plots assume fixed parameters nu = 311, nl = [311, 724], SNR = ∥wσ2∥ , and a target signal misaligned with the spike eigenvector (η = 0.75). The number of features is set to p = 500. B.2

Optimal dimensionality

In Figures 8, 9, 10 and 11, we present various phase portraits in α, reflecting their dependence on the signal-to-noise ratio (SNR), the spike strength (λ), and the task overlap. For the optimal α minimizing the training error (Figures 9 and 11), the SNR, task overlap, and spike strength have little effect. In contrast, for the optimal α minimizing the generalisation error, the phase-transition boundaries and the size of the bulk region are highly sensitive to the task overlap, SNR, and spike strength (λ). Figures 8 and 10 show that the phase-transition boundary flattens into a horizontal line at a fixed nl (i) as the task-spike overlap decreases, (ii) as the SNR increases or (iii) as the spike strength decreases. At the same time, the bulk region expands along the nu axis, thereby reducing the region of parameter space in which compression is used. This behaviour can be explained by the fact that a weaker signal is harder to detect, leading to a less accurate estimate of the prior. Consequently, exploiting this poorly-estimated prior through compression is not advantageous across many regimes, even when the amount of new data is large.

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

Figure 7: Components of the theoretical generalisation error as a function of α for different values of λ. The total error (panels (d, h)) is decomposed into the estimation error (panels (a, e)), the missing signal (panels (b, f)), and the ⋆ 2 leak term (c, g). All plots assume fixed parameters nu = 311, nl = [311, 724], SNR = ∥wσ2∥ , and a target signal misaligned with the spike eigenvector (η = 0.75). The number of features is set to p = 500.

Additional experiments

We provide additional figures from the experiments outlined in the main text in Section 5, as well as an additional experiment directly comparing the effectiveness of pretrained regression. C.1

Comparison with regression and principal component regression

We compare three approaches: standard regression, pretrained regression (PR) where PCA is learned on prior data and transferred to the downstream task, and principal component regression (PCR), where PCA is learned directly on the downstream data. In the fully aligned setting (w⋆ ∥ v, cos θ = 1), Figure 12 (left) shows that incorporating a PCA step consistently improves performance over standard regression. Moreover, when the pretraining dataset is larger than the downstream dataset (nu > nl ), PR outperforms PCR, reflecting the more accurate estimation of the signal subspace from additional data. √ In the misaligned setting (cos θ = 1/ 2, right panel in Figure 12), the advantage of PR is restricted to regimes with large nu and small nl . As nl increases, PCR becomes comparable to PR, and both methods approach the performance of standard regression (within a small error gap). This reflects the diminishing benefit of transferring a pretrained subspace when sufficient downstream data is available to estimate the relevant structure directly. C.2

Additional autoencoder experiments

Figure 13 extends panels (a-d) of Figure 4. The top row shows simulations that directly match our theoretical setting (PCA followed by regression), while the middle and bottom rows replace PCA with unsupervised representation learning using a linear and a nonlinear one-hidden-layer autoencoder, respectively. The linear autoencoder closely reproduces the theoretical predictions in both the spiked identity and general covariance settings, confirming that it effectively recovers the same subspace as PCA in these regimes. We next examine the effect of reducing the SNR in Figure 14. The phase boundaries shift in the same qualitative manner as predicted by the theory, where the optimal α is low across a larger part of the (nu , nl ) space, indicating that the dependence on SNR is captured beyond the exact PCA setting. 37

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

Figure 8: Optimal value of α that minimizes the generalisation error, shown as a function of SNR and spike alignment η. We fix λ = 2. Red lines indicate the approximate phase transitions.

Finally, we analyse deviations from PCA in low-sample regimes in Figure 15. We compare the theoretical overlap Pm v with the empirical overlap obtained by passing v through trained linear and nonlinear autoencoders with varying bottleneck dimension. The linear autoencoder deviates from PCA only when nu < p, while the nonlinear autoencoder exhibits larger discrepancies across a broader range of nu . Despite these deviations, the optimal bottleneck dimension that minimizes both training and generalisation error remains close to α ≈ 1 when nu < p. As a result, the performance of linear autoencoder-based pipelines is still well predicted by the theory across all nu . In contrast, the nonlinear autoencoder exhibits a persistent mismatch for nu > p, which only diminishes for sufficiently large pretraining datasets (empirically around nu ≈ 3p). The magnitude of this mismatch also depends on SNR, as seen by comparing Figures 13 and 14. 38

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

Figure 9: Optimal value of α that minimizes the training error, shown as a function of SNR and spike alignment η. We fix λ = 2.

C.3

Additional transformer experiments

We first consider the setting where PCA is applied to representations extracted from the downstream task (Figure 16), rather than from pretraining data as in Figure 4(e). In this regime, compression remains beneficial for generalisation when the downstream dataset is small, with an optimal intermediate dimension m. The dependence on the pretraining dataset size nu is weaker than in the pretraining-PCA setting, but still observable. This likely reflects the evolution of the representation spectrum during training, with increasing spike size (effective λ) as shown in Figure 18. We next return to applying PCA to pretraining data and report validation accuracy across a broader range of downstream dataset sizes (Figure 17). For small downstream datasets, an intermediate compression level is optimal, particularly for 39

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

Figure 10: Optimal value of α that minimizes the generalisation error, shown as a function of spike strength λ and spike alignment θ with w∗ . We fix SNR = 9. Red lines indicate the approximate phase transitions. models trained for longer durations. In contrast, for large downstream datasets, performance is maximized without compression (m ≈ p). At the smallest dataset sizes, the optimal m is less stable across runs. Neither experimental setup provides a direct match to the theoretical model. In particular, the pretraining and downstream datasets differ not only in semantic structure but also in sequence length, which affects the learned representations and limits the transferability of PCA computed on pretraining data. Despite this mismatch, both PCA-based approaches match or outperform the no-compression baseline (Figure 19). D D.1

Implementation details Figure 1

Panel (a)

Diagram of the model made with figma. ⋆ 2

Model. We define the signal-to-noise ratio as SNR = ∥wσ2∥ and the alignment between the regression target and the spike as η = cos θ. Throughout, we fix p = 500 and σ 2 = 1. Panels (b-c) (generalisation and training error vs α). We fix nu ≈ nl ≈ 300, SNR = 9.0, and η = 1.0. The parameter α is swept over 30 values in [0.01, 1.0] for each of four spike strengths λ ∈ {1.1, 2.0, 5.0, 10.0}. Theoretical predictions (solid curves) are compared against empirical averages over multiple realisations (dashed curves, shading corresponds to 1 standard deviation). D.2

Figure 2 ⋆ 2

Model. We define the signal-to-noise ratio as SNR = ∥wσ2∥ and the alignment between the regression target and the spike as η = cos θ. Throughout, we fix p = 500 and σ 2 = 1. 40

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

Figure 11: Optimal value of α that minimizes the training error, shown as a function of spike strength λ and spike alignment θ with w∗ . We fix SNR = 9.

Regression

2172

Labelled nl

2172

1448

724

Task-misaligned

3000

724

1448

2172 Unlabelled nu

3000

1448

PCR

724

Best model (within 2% tie threshold)

Task-aligned

3000

724

1448

2172

Unlabelled nu

3000

Figure 12: Heatmaps over the (nu , nl ) grid (p = 500, SNR = 1.8, λ = 5) showing which method – standard regression (violet), pretrained regression (PR, blue) or principal component regression (PCR, orange) – achieves lower generalisation error. Forward hatching (//) marks regions where PR and PCR are tied (within 2% error); cross hatching (×) marks regions where all three methods (PR, PCR, and standard regression) are tied. When the task is aligned with the pretraining signal, PR wins for nu > nl and PCR wins otherwise, with the boundary tracking the diagonal. Under misalignment, PCR dominates across most of the grid; PR retains an advantage when nu ≫ nl . Panel (a) (minimum generalisation error). For SNR = 9.0, λ⋆ = 5.0, and η = 1.0, we plot minα Egen (α) over a 100 × 100 grid with nu , nl ∈ [1, 3000]. Dashed white lines mark nu = p and nl = p (p = 500). The colour scale is logarithmic. 41

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

Gen-opt

Train-opt

884

253

568 253

1200

884

Labelled nl

568 253

1200

884

Labelled nl

1200

253

568

884

Unlabelled nu

1200 1

253

568

884

Unlabelled nu

0.4

0.2

568

0.0

253 1

1200

0.6

253 1

253

0.8

568

1.0

568

Gen-opt

Optimal

Labelled nl

1200

Labelled nl

1200

Labelled nl

Nonlinear AE

Linear AE

PCA + Regression

Train-opt

253

568

884

Unlabelled nu

1200 1

253

568

884

Unlabelled nu

1200

Figure 13: Extended version of panels (a-d) in Figure 4: optimal bottleneck size in autoencoders on inputs with different covariance structure. We compare PCA on pretraining data (top row) corresponding to our theoretical setup, a linear autoencoder (middle row) and a non-linear autoencoder (bottom row) trained on Gaussian data, with varying bottleneck size m (α = m/p). The input covariance on the first two columns is Σ = I + λvv⊤ , while in the last two columns the covariance is a Toeplitz matrix H with a spike, i.e., Σ = H + λvv⊤ , where Hi,j = ρ|i−j| with ρ = 0.5. The autoencoder results match those of PCA, and the identity vs general covariance exhibit a similar phenomenology. We fix SNR = 9.0, λ = 5.0.

Panel (b) (gain from tuning α relative to α = 1). Under the same setting, we compute Egen (α = 1) − Egen (α⋆ ) over the same grid. Positive values indicate regions where using the full prior representation is suboptimal and tuning α reduces generalisation error. The colour scale is logarithmic.

Panel (c) (gain from tuning α relative to αmin ). Under the same setting, we compute Egen (α = αmin ) − Egen (α⋆ ), where αmin ≈ 0.01 is the smallest value in the sweep (m = 1 principal component retained). This quantifies the benefit of tuning α upward from the minimal representation. The colour scale is logarithmic.

Panel (d) (marginal substitution rate). Under the same setting, we compute the ratio (∂Egen /∂nu )/(∂Egen /∂nl ) evaluated at the optimal α⋆ , with gradients approximated numerically via central differences on the 100 × 100 grid. Values above 1 (red) indicate that collecting an additional unlabelled sample nu reduces generalisation error more than collecting a labelled sample nl ; values below 1 (blue) indicate the converse. The colour scale is centred at 1. Panel (e) (generalisation-optimal α⋆ ). Under the same setting, we plot the value of α⋆ = arg minα Egen (α) over the same grid, normalised so that α ∈ [0, 1]. 42

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

Gen-opt

Train-opt

884

253

568 253

1200

884

Labelled nl

568 253

1200

884

Labelled nl

1200

253

568

884

1200 1

Unlabelled nu

253

568

884

Unlabelled nu

0.6

253 1

253

0.8

568

1.0

568

Gen-opt

Optimal

Labelled nl

1200

Labelled nl

1200

Labelled nl

Nonlinear AE

Linear AE

PCA + Regression

Train-opt

0.4

0.2

568

0.0

253 1

1200

253

568

884

Unlabelled nu

1200 1

253

568

884

Unlabelled nu

1200

Spike direction overlap

Figure 14: Effect of SNR on optimal bottleneck size in autoencoders on inputs with different covariance structure. We compare PCA on pretraining data (top row) corresponding to our theoretical setup, a linear autoencoder (middle row) and a non-linear autoencoder (bottom row) trained on Gaussian data, with varying bottleneck size m (α = m/p). The input covariance on the first two columns is Σ = I + λvv⊤ , while in the last two columns the covariance is a Toeplitz matrix H with a spike, i.e., Σ = H + λvv⊤ , where Hi,j = ρ|i−j| with ρ = 0.5. The autoencoder results match those of PCA, and the identity vs general covariance exhibit a similar phenomenology. We fix SNR = 1.8, λ = 5.0. latent dim = 5

1.0

latent dim = 25

latent dim = 50

latent dim = 75

latent dim = 100

PCA Theory Linear AE Nonlinear AE

0.8 0.6 0.4

PCA Theory Linear AE Nonlinear AE

0.2 0.0 0

nu/p

PCA Theory Linear AE Nonlinear AE 0

nu/p

Figure 15: Spike direction recovery as a function of sample ratio nu /p for linear and nonlinear autoencoders. Each panel shows the overlap between the singular vectors of the encoder weights and the population spike v as a function of the sample ratio nu /p, for a fixed latent dimension m ∈ {5, 25, 50, 75, 100}. Data are generated from a spiked covariance model in p = 100 dimensions with spike strength λ = 5. Solid lines show empirical means over 5 random seeds; shaded bands show ±1 standard deviation. The dotted line shows the theoretical PCA overlap p̄v,v derived from Lemma 3. Both linear and nonlinear autoencoders show a mismatch in low sample regimes, with a more evident mismatch in nonlinear autoencoders and in autoencoders with larger bottlenecks. For large sample regimes, linear and nonlinear autoencoders match PCA well, and show the phase transitions predicted by our theory (see panels (a-d) in Figure 4 discussed in Section 5.1, and Figures 13-14 discussed in Appendix C.2).

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

(b)

Optimal m for validation accuracy (downstream PCA)

Optimal m for training accuracy (downstream PCA)

67,349

500

20,000

400

20,000

400

4,000

300

4,000

300

1,000

200

1,000

200

250

100 8

128

2,000

20,000

Pretraining step (checkpoint)

Optimal m

500

Labelled samples nl

67,349

Optimal m

Labelled samples nl

(a)

250

100

143,000

128

2,000

20,000

Pretraining step (checkpoint)

143,000

Figure 16: Optimal m using PCA on representations of the downstream task. We use the same layout as Figure 4(e), but the PCA basis is fit on the downstream labelled set rather than on the pretraining corpus. (a) Optimal m for validation accuracy; (b) optimal m for training accuracy. Compared to doing PCA on the pretraining corpus (Figure 4(a)), downstream PCA tends to favor larger m and the dependence on the checkpoint is weaker. Labelled nl = 100

(a)

Labelled nl = 500

(b)

Validation accuracy

(d) Labelled nl = 20, 000

(e) Labelled nl = 67, 349

0.75

0.70

0.65

0.60

0.55

0.50

(c)

0.75

0.70 0.65

Labelled nl = 2, 000

0.70 0.65

0.60

0.65 0.60

0.55

0.60 0.55

0.50

0.45 4

256

PCA components m

256

PCA components m

256

PCA components m

256

0.50

PCA components m

Checkpoint 2 8 32 128 512 2,000 5,000 20,000 80,000 143,000

256

PCA components m

Figure 17: Validation accuracy vs number of PCA components m across five downstream task sizes. Note that the values of m are on a log scale. Each panel shows a different value of nl spanning the full range from nl = 100 to nl = 67,349; curves are coloured by pretraining checkpoint. Shaded bands show the standard deviation over three random seeds. For smaller tasks, the curves have a clear peak and become progressively more monotone as nl grows, reflecting that small tasks benefit from aggressive compression while large tasks can exploit finer-grained directions in representation space. Pretraining step 128

(a)

Pretraining step 2,000

(b)

101

Pretraining step 50,000

(d)0 10

(e) Pretraining step 143,000

100

Density

Pretraining step 10,000

(c)

10 1 10 1

100 10 1

10 1

10 2 10 2

10 1 10 2

10 2 0

Eigenvalue

10 3 0

Eigenvalue

100

Eigenvalue

150

Figure 18: Eigenvalue spectra of Pythia-70M-deduped last-token last-hidden-layer representations at five pretraining checkpoints. Each panel shows the empirical distribution of covariance eigenvalues computed from the representation of the sst2 task. Throughout pretraining, the spectrum shows a clear bulk and a few outlying eigenvalues. As pretraining progresses (steps 128 to 143,000), the outlying eigenvalues become more and more prominent. D.3

Figure 3 ⋆ 2

Model. We define the signal-to-noise ratio as SNR = ∥wσ2∥ and the alignment between the regression target and the spike as η = cos θ2 . Throughout, we fix p = 500 and σ 2 = 1. 44

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

Figure 19: Validation accuracy vs pretraining checkpoint for three probing conditions across five downstream task sizes. Each panel shows a different downstream task size nl ; the x-axis is the pretraining step on a log scale. Three curves are compared: Pretrain PCA (pink, solid) fits the PCA basis on Pile representations from the same checkpoint and selects the best m; Downstream PCA (blue, dashed) fits PCA on the downstream task representations (sst2) and selects the best m; No PCA (grey, dotted) uses the full-dimensional representations (largest available m). For each curve, the displayed value is the mean over three seeds at the per-seed optimal m, with the shaded area corresponding to 1 standard deviation. The gap between PCA and the no-PCA baseline is largest at small and intermediate nl , where dimensionality reduction acts as regularization. Pretrain PCA matches downstream PCA for intermediate and large nl , with the mismatch coming from the difference in the type of data available in the pretraining and downstream datasets. Panel (a) (optimal α⋆ heatmap). For SNR = 9.0, λ = 5.0, and η = 1.0, we compute the theoretical optimal representation size α⋆ that minimises generalisation error over a 100 × 100 grid, where nu and nl each vary from 1 to 3000. Analytical phase-transition boundaries (smooth transition and sharp transition) are overlaid. Panel (b) (phase portrait SNR). At fixed λ = 5.0 and η = 1.0, we plot the smooth and sharp transition boundaries in the (nu , nl ) plane for each of the six SNR values {0.1, 0.5, 1.0, 3.0, 5.0, 9.0}. Increasing SNR progressively reduces the region in which full-rank representations (α ≈ 1) are suboptimal. Panel (c) (phase portrait spike strength). At fixed SNR = 9.0 and η = 1.0, we show the corresponding transition curves for λ ∈ {2, 3, 5, 7, 10}. Stronger spikes shift the phase boundaries. Panel (d) (phase portrait alignment). At fixed SNR = 9.0 and λ = 5.0, we plot the transition curves for η ∈ {0.0, 0.25, 0.5, 0.75, 1.0}. All panels are evaluated over nu , nl ∈ [1, 3000]. D.4

Figure 4

Autoencoders Data generation.

Both covariances use p = 200 features, noise standard deviation σ = 1, SNR = 9, ∥w⋆ ∥ = 3.

• Spiked identity (λ = 5.0): Σ = Ip + (λ − 1)vv⊤ ; true weight vector aligned with the spike w⋆ = ∥w⋆ ∥v, where v is the top eigenvector of Σ. • Spiked Toeplitz (λ = 5.0): Σ is a Toeplitz matrix with ρ = 0.5, with λ added to the top eigenvalue; w⋆ again aligned with the top eigenvector. Each grid point (nu , nl ) averages five independent replications, with nu , nl ∈ [1, 1200] each swept over 20 uniform values. Models. For each (nu , nl ) pair, α ∈ {0.01, 0.05, . . . , 1.0} (20 values) sets the latent dimension m = ⌊αp⌋. The three models below are evaluated. 1. PCA + Regression (PR): Eigenvectors of Su = X⊤ u Xu /nu are computed; Xl is projected onto the top-m eigenvectors and OLS is applied in the latent space. 2. Linear AE: Encoder (Rp → Rm ) and linear decoder (Rm → Rp ) trained by minimizing MSE on Xu ; the frozen encoder maps Xl to features for OLS. 3. Nonlinear AE: Same as the linear AE but with a sigmoid nonlinearity applied at the encoder, z = σ(Wenc x). 45

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

All autoencoders are implemented in JAX/Equinox and trained with Adam (lr = 10−3 , 20,000 epochs, batch size 500). √ Weights are initialized as Wi,j ∼ N (0, 1/ din ). Panel descriptions. Each panel is a 20 × 20 heatmap with nu on the x-axis, nl on the y-axis, and color encoding α⋆ ∈ [0, 1] (the value minimizing mean loss over the five replications). The autoencoder component of the figure is organized as 2 rows × 2 columns:

Spiked identity Spiked Toeplitz

Linear AE

Nonlinear AE

(a) (c)

(b) (d)

Columns compare linear vs nonlinear AEs; rows contrast spiked identity and spiked Toeplitz covariance structures; each panel shows generalisation-optimal α⋆ . A single shared colorbar applies to all panels. LLMs Model and representations. All experiments use Pythia-70M-deduped [79] (70M parameters, trained on the deduplicated Pile [80, 88]). We evaluate 10 checkpoints at steps 2, 8, 32, 128, 512, 2000, 5000, 20000, 80000, and 143,000. Representations are last-token hidden states from the final layer of 64-token truncated inputs, yielding vectors in R512 . Downstream task. Binary sentiment classification on SST-2 from the GLUE benchmark [3]. Fixed 872-sentence validation set; training sizes nl ∈ {100, 250, 500, 1000, 2000, 4000, 10000, 20000, 40000, 67349}, nl training examples are randomly chosen out of the maximum 67349 for each seed. PCA. We sweep m ∈ {2, 3, 4, 7, 11, 16, 25, 39, 60, 92, 142, 218, 334, 512} (approximate log2 grid), and consider the two conditions below. 1. Downstream PCA (Figure 16): basis fit on the nl labelled training representations; 2. Pretrain-PCA (Figure 4(e-g)): basis fit on up to 51,200 Pile sequences (50 batches × 1,024 sequences) from the same checkpoint, using torch.linalg.eigh. For earlier checkpoints, the basis is fit on 1024 × ncheckpoint sequences. Linear probe. Projected representations are fed to a 2-output linear layer trained with cross-entropy loss using Adam (lr = 10−3 ), for up to 10,000 epochs with early stopping. Linear probes were randomly initialized with small weights for each seed. Seeds and aggregation. Each (checkpoint, nl , m) configuration is run with 3 random seeds (affecting probe initialization and downstream data shuffling). Figures show mean ± std; optimal-m heatmaps report the mean per-seed argmax. Panel descriptions. • (e) Heatmap for optimal m under pretrain-PCA, selected by validation accuracy. • (f) Eigenvalue histogram at step 143,000, computed from the 872-sentence SST-2 validation covariance; log-density on the y-axis. The full set of 5 checkpoints is shown in Figure 18. • (g) Validation accuracy vs m (x-axis in log2 scale) under pretrain-PCA for nl = 2000, with curves colored by pretraining progress. D.5

Compute

Theory curves. Theory curves were computed on a single CPU (Intel Xeon E5-2680 v4 @ 2.40 GHz, no GPU required). Each configuration sweeps a 100 × 100 grid of (nu , nl ) values over 30 values of α, amounting to approximately 300,000 evaluations of closed-form expressions. The full sweep covers 24 configurations, λ∗ ∈ {1.1, 2.0, 5.0, 10.0} × SNR ∈ {0.1, 0.5, 1.0, 3.0, 5.0, 9.0}, with alignment η = 1.0, totaling approximately 7.2 million evaluations. When executed in parallel, the full sweep completes in approximately 2h. Peak RAM usage was approximately 2 GB per configuration. 46

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

Simulations. Simulations were distributed over a SLURM cluster. Each job was allocated 4 CPUs and 30 GB of RAM. The full grid 6 SNR × 4 λ1 × 4 λ2 × 30 nb produced 2,880 SLURM jobs, each processing 30 values of nl with 100 independent replications at p = 500. Each replication involves an eigendecomposition and pseudoinverse of a p × p matrix for each of the 30 values of α. When executed in parallel, the full sweep completes in approximately 10h. Autoencoder experiment. Simulations were distributed over a SLURM cluster. Each job was allocated 1 GPU, 8 CPUs, and 16 GB of RAM. The full grid 20 nbase × 20 ntask produced 400 SLURM jobs, each processing 20 values of α with 5 independent replications at p = 200. Each replication trains a linear and a nonlinear autoencoder for 20,000 epochs for each value of α, followed by OLS regression in the learned latent space. The individual jobs completed in 10-50min depending on parameters. LLM probe experiment. Simulations were distributed over a SLURM cluster. Each job was allocated 1 GPU, 16 CPUs, and 32 GB of RAM. Two sets of 10 SLURM jobs were submitted: one per checkpoint of Pythia-70M-deduped for the downstream-PCA condition and one for the pretrain-PCA condition. Each job sweeps over 10 training-set sizes and 3 seeds, fitting a linear probe for up to 10,000 epochs with early stopping across up to 14 PCA dimensionalities m. D.6

LLM Usage

Large language models (LLMs) were used throughout this research project to assist with writing, editing, and formatting the manuscript, including improving clarity and refining LATEX expressions. LLMs were also used as an aid for understanding and discussing technical concepts, as well as for debugging and drafting portions of the experimental and simulation code. All scientific results, analyses, and conclusions were verified and validated by the authors.

Record · ID 204779 · SHA-256 941ae86ae1697530
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.