Theory of Scaling Laws for In-Context Regression: Depth, Width, Context and Time

## Abstract

Abstract We study in-context learning (ICL) of linear regression in a deep linear self-attention model, characterizing how performance depends on various computational and statistical resources (width, depth, number of training steps, batch size and data per context). In a joint limit where data dimension, context length, and residual stream width scale proportionally, we analyze the limiting asymptotics for three ICL settings: (1) isotropic covariates and tasks (ISO), (2) fixed and structured covariance (FS), and (3) where covariances are randomly rotated and structured (RRS). For ISO and FS settings, we find that depth only aids ICL performance if context length is limited. Alternatively, in the RRS setting where covariances change across contexts, increasing the depth leads to significant improvements in ICL, even at infinite context length. This provides a new solvable toy model of neural scaling laws which depends on both width and depth of a transformer and predicts an optimal transformer shape as a function of compute. This toy model enables computation of exact asymptotics for the risk as well as derivation of powerlaws under source/capacity conditions for the ICL tasks.

## 1 Introduction

In recent years, transformer models have become the backbone architecture for modern machine learning systems (Vaswani et al., 2017; Dosovitskiy et al., 2020). The transformer architecture consists of alternating self attention and multi-layer perceptron blocks in a deep residual network. For a variety of tasks, increasing the size of transformers by increasing width and depth of the model leads to empirically predictable improvements in model performance (Hestness et al., 2017; Kaplan et al., 2020; Hoffmann et al., 2022; Achiam et al., 2023). Various protocols for scaling up the width and depth of transformers have been developed that provide stable training limits (Yang et al., 2021; Bordelon et al., b; 2024b; Chizat & Netrapalli, 2024). One common strategy is to scale width $N$ and depth $L$ linearly with fixed aspect ratio $L/N$ (Noci et al., 2023; Dey et al., 2025), however existing theory cannot justify such a practice as compute optimal or distinguish relative performance gain from width and depth from total parameter count. This leads us to our first open question:

In this work, we explore this question for in-context learning (ICL) problems, specifically in-context linear regression problems. In-context learning refers to the ability of models to condition their outputs on the past sequence of inputs provided by the user without updating model parameters (Brown et al., 2020). This contrasts with in-weight learning (IWL) where information about the dataset is encoded in the model parameters during pre-training. Whereas many prior theoretical works on neural scaling laws essentially analyze the role of width, pretraining data, or pretraining time (Bahri et al., 2021; Maloney et al., 2022; Simon et al., 2023; Bordelon et al., 2024a; Paquette et al., 2024; Lin et al., 2024; Bordelon et al., a), our work explores how depth, width, context length, and pretraining time influence the quality of ICL in self-attention models. In our investigations, we find that the architectural requirements to perform ICL regression depend significantly on the statistics of the pretraining ICL tasks. To address the first question, we are forced to investigate:

To address these questions, we develop a solvable model of deep linear attention for three distinct ICL data covariance structures. Our concrete novel contributions and findings are as follows:

### 1.1 Contributions

- 1.
We identify a new asymptotic scaling for in-context regression with linear attention where context length $P$, number of masked points per context $K$, and contexts per step $B$ all scale linearly with dimension $D$. This reduces the amount of compute and total data (by a factor of $D$) to converge than prior works (Lu et al., 2025) and speeds up simulations.
- 2.
We introduce three distinct ICL data models of increasing generality. We study isotropic input data and task vectors (ISO). We then proceed to analyze fixed and structured (FS) covariances for data and task vectors. Lastly, we examine an ICL distribution where data covariances are randomly rotated and structured (RRS) across contexts.
- 3.
We show with our asymptotic theory, that depth $L$ benefits the first two data settings (ISO + FS) only if $\alpha\equiv\text{context length}/\text{data dimension}$ is finite. In the limit of $\alpha\to\infty$, there is no benefit from depth for these ICL tasks. Further, the FS setting is brittle with respect to variations in changes to the data covariance at test time. However, for the RRS setting, increasing depth is beneficial for any distribution and a non-trivial width & depth scaling.
- 4.
We introduce a model width $N$ bottleneck to our ICL model and study the pretraining scaling law for a model trained for $t$ steps on a width $N$ and depth $L$ linear attention model in the RRS setting. For powerlaw data, the scaling law takes the form of a Chinchilla scaling law but with width and depth contributing separate terms
$\displaystyle\mathcal{L}(t,N,L,P)=c_{t}\ t^{-\beta_{t}}+c_{N}\ N^{-\beta_{N}}+c_{L}\ L^{-\beta_{L}}+c_{P}\ P^{-\beta_{P}}.$
(1)
- 5.
From this scaling law, we consider compute optimal joint scaling of width and depth. Depending on the structure of the ICL covariates, we obtain different scalings of $L\sim N^{\nu}$ where $\nu$ depends on properties of the data.

### 1.2 Related Works

##### Infinite Neural Network Width and Depth Limits and Commonly Used Joint Scalings.

The empirical fact that larger networks tend to perform better on natural data (Hestness et al., 2017; Kaplan et al., 2020; Hoffmann et al., 2022) has led to development of scaling procedures to stably increase the size of a model to allow optimization. The mean-field or $\mu$P scaling theory (Geiger et al., 2020; Mei et al., 2018; Chizat & Bach, 2018; Yang & Hu, 2021; Bordelon & Pehlevan, 2022b) allows one to scale up the width of a model in a way that admits a feature learning infinite limit. Further, this scaling protocol provides consistent optimal hyperparameters, while delivering monotonic improvements in performance (Yang et al., 2021). The same program has been carried out for deep residual models, such as transformers (Bordelon et al., b; Yang et al., 2023; Bordelon et al., 2024b; Dey et al., 2025). However, while these infinite width and infinite depth scaling limits have been established to exist and perform better than finite models, no theory currently captures the relative gains in performance from scaling up width or depth at fixed compute. Understanding compute optimal shapes could help guide architectural choices when training large transformer models.

##### Theories of Compute Optimal Neural Scaling Laws.

Following the empirical scaling law results of Kaplan et al. (2020) and Hoffmann et al. (2022), many theoretical works have examined the generalization theory for fully trained kernel methods under power law features (Bordelon et al., 2020; Spigler et al., 2020; Cui et al., 2021; Bahri et al., 2021; Maloney et al., 2022; Atanasov et al., 2023; 2022; Defilippis et al., 2024). More recently, several efforts have begun to incorporate SGD dynamics into these models to gain a notion of compute (and compute optimal tradeoffs between parameters and training time) (Bordelon & Pehlevan, 2022a; Bordelon et al., 2024a; Paquette et al., 2024; Lin et al., 2024; Kunstner & Bach, 2025; Bordelon et al., a). In these works, the model is essentially one or two layers and the notion of finite model size is introduced with a random projection of the features to an $N$ dimensional space. In this way, existing theories more closely resemble scaling laws for width rather than a comparison where depth and width serve different functions. Recent empirical works have pointed out the utility of increasing depth (or virtual depth through looping) for tasks requiring reasoning such as solving Sudoku puzzles (Wang et al., 2025) and solving math problems (Zhu et al., 2025; Geiping et al., 2025). Merrill & Sabharwal (2025) study regular language recognition where a clear computational advantage to scaling depth instead of width is established and experimentally verified.

##### Empirical Studies of ICL.

Brown et al. (2020) demonstrated that large pretrained language models such as GPT-3 exhibit remarkable in-context learning capabilities for natural language tasks. Following this finding, Garg et al. (2022) empirically investigated the ability of transformers to learn simple function classes such as linear regression, sparse regression, and two-layer neural networks. Various works have studied emergence of ICL beyond language tasks, focusing on what data structures (from dynamical systems to classification problems) can be learned in-context (Liu et al., 2024b; Chan et al., 2022; He et al., 2024; McCracken et al., 2025; Park et al., 2024; Chen et al., 2024). Several works have offered experimental evidence that ICL implements Bayesian inference (Xie et al., 2021; Wurgaft et al., 2025; Liu et al., 2024a; Panwar et al., 2024; Zhang et al., 2023b). Further investigations questioned whether attention is even strictly necessary for ICL by studying the performance of MLPs on these tasks (Tong & Pehlevan, 2025; Kratsios & Furuya, 2025).

##### Theoretical Studies of ICL.

Inspired by empirical studies of ICL, many works have theoretically investigated how ICL algorithms such as in-context gradient descent can be implemented in transformers (Zhang et al., 2023a; Ahn et al., 2023; Von Oswald et al., 2023; Li et al., 2023; Kim et al., 2024; Akyürek et al., 2022; Pathak et al., 2023; Mahankali et al., 2023). Some works have demonstrated flexibility of the transformer to adapt to changing statistics (Vladymyrov et al., 2024) or performing model selection (Bai et al., 2023). Studies have pointed out the need for sufficient pretraining task diversity for ICL generalization (Raventós et al., 2023; Wu et al., 2024) which was theoretically analyzed in an asymptotic scaling limit for a shallow linear attention model by Lu et al. (2025). Extensions to structured covariances and distribution shifts revealed the importance of train-test task alignment Letey et al. (2025). Gatmiry et al. (2024) investigate the need for a sufficient number of residual stream steps (either by increasing true depth or loops of attention layers) to solve ICL distributions with multiple condition numbers with larger condition numbers requiring more depth. While many of these works study a final construction of weights, recent theory has described the training dynamics of one layer linear multi-head attention (Zhang et al., 2025).

## 2 Data, Architecture and Reduced Model

### 2.1 Deep Linear Attention Architecture

The most general model we study is a depth $L$, residual linear attention model $f$ that maps inputs contexts to output predictions. The data are formed as $P$ input-output pairs $\{(\bm{x}_{\mu},y_{\mu})\}_{\mu=1}^{P}$ and $K$ evaluation points $\{(\bm{x}^{\star}_{\mu},*)\}_{\mu=P+1}^{P+K}$ (which do not carry target outputs) into a data matrix $\bm{D}$

$$ $\displaystyle\bm{D}=\begin{bmatrix}\bm{x}_{1}&...&\bm{x}_{P}&\bm{x}_{P+1}&...&\bm{x}_{P+K}\\ y_{1}&...&y_{P}&*&...&*\end{bmatrix}$ (2) $$

where $*$ indicates masked target values on the $K$ evaluation points which are provided as 0 entries. The evaluation tokens $\mu\in\{P+1,...,P+K\}$ prevented from updating the model with a positional masking matrix $M_{\mu\nu}$ (Appendix [B](https://arxiv.org/html/2510.01098v1#A2)). The model $f_{\mu}$ is computed from

$$ $\displaystyle\bm{h}^{1}_{\mu}=\bm{W}_{x}\bm{x}_{\mu}+\bm{w}_{y}y_{\mu},\quad\quad\bm{h}^{\ell+1}_{\mu}=\bm{h}^{\ell}_{\mu}+\frac{1}{LP}\sum_{\nu=1}^{P}\left(\bm{k}^{\ell}_{\nu}\cdot\bm{q}^{\ell}_{\mu}\right)\bm{v}^{\ell}_{\nu}\ ,\ \ell\in[L]\ ,\ \mu\in[P+K]$ $\displaystyle\bm{q}^{\ell}_{\mu}=\bm{W}_{q}^{\ell}\bm{h}^{\ell}_{\mu}\ ,\ \bm{k}^{\ell}_{\mu}=\bm{W}_{k}^{\ell}\bm{h}^{\ell}_{\mu}\ ,\ \bm{v}^{\ell}_{\mu}=\bm{W}_{v}^{\ell}\bm{h}^{\ell}_{\mu},\quad\quad f_{\mu}=\bm{w}_{o}\cdot\bm{h}^{L}_{\mu}$ (3) $$

The loss function for context $\bm{D}$ is $\mathcal{L}({\bm{D}})=\frac{1}{K}\sum_{\mu=P+1}^{P+K}(f_{\mu}-y_{\mu})^{2}$ and the full population loss is $\mathcal{L}=\mathbb{E}_{\bm{D}}\mathcal{L}(\bm{D})$ where $\bm{D}$ is the distribution of context matrices. We stress that the operation $\bm{q}_{\mu}\cdot\bm{k}_{\mu}$ corresponds to linear attention rather than commonly used in soft-max attention. For the regression tasks we consider, this model is sufficient to solve the ICL task (Ahn et al., 2023; Von Oswald et al., 2023) and aids theoretical tractability (Lu et al., 2025; Zhang et al., 2025).

### 2.2 Recurrent Reduced- Γ \Gamma Model

Following prior works on ICL in linear regression (Zhang et al., 2023a; Lu et al., 2025; Wu et al., 2024; Zhang et al., 2025), we examine the minimal (simplest) reparameterization of the above model which can solve this task where the residual stream encodes $\bm{x}$ information in a subspace orthogonal to $\bm{w}_{y}=\bm{w}_{o}$ so that $\bm{W}_{x}^{\top}\bm{w}_{y}=0$ and $\bm{W}_{v}\propto\bm{w}_{y}\bm{w}_{y}^{\top}$. Further, instead of optimizing separate hidden weight matrices, we consider looped / universal transformers where each attention block is the identical $\bm{W}^{\ell}_{i}=\bm{W}^{\ell^{\prime}}_{i}$ for $i\in\{k,q,v\}$ (Dehghani et al., 2018; Gatmiry et al., 2024), though we relax this assumption in Section [5.1](https://arxiv.org/html/2510.01098v1#S5.SS1) and Appendix [G.1](https://arxiv.org/html/2510.01098v1#A7.SS1). The reduced model then defines the predictor $f(\bm{x}_{\star})$ for test point $\bm{\bm{x}}_{\star}$ in terms of a single matrix $\bm{\Gamma}\in\mathbb{R}^{D\times D}$

$$ $\displaystyle\bm{\Gamma}$ $\displaystyle\equiv\left(\bm{w}_{o}^{\top}\bm{W}_{v}\bm{w}_{y}\right)\bm{W}_{x}^{\top}\bm{W}_{k}^{\top}\bm{W}_{q}\bm{W}_{x}\ ,\ f(\bm{x}_{\star})=\frac{1}{LP}\bm{x}_{\star}^{\top}\bm{\Gamma}\sum_{\ell=0}^{L-1}\left(\bm{I}-L^{-1}\hat{\bm{\Sigma}}\bm{\Gamma}\right)^{\ell}\bm{X}^{\top}\bm{y}$ (4) $$

where $\hat{\bm{\Sigma}}=\frac{1}{P}\bm{X}\bm{X}^{\top}$ is the empirical covariance. Due to the simplicity of this model and loss landscape structure, we will first focus on the gradient dynamics when we directly perform gradient descent on the matrix $\bm{\Gamma}$. Then in Section [5.2](https://arxiv.org/html/2510.01098v1#S5.SS2), we consider learning dynamics on the decoupled collection of parameters $\{\bm{W}_{i}\}_{i\in\{k,q,v\}}$.

## 3 Learning Curves in Reduced Linear Attention Model

### 3.1 Isotropic Covariates and Tasks ( ISO )

To start our investigation, we begin by considering $D$ dimensional isotropic data and isotropic task distribution. For each context $\bm{D}_{c}$ and each data point $\mu\in[P]$, the distribution of $\bm{x}$ and $y$ are

$$ $\displaystyle\bm{x}_{\mu,c}\sim\mathcal{N}\left(0,{\bm{I}}\right)\quad\quad\bm{\beta}_{c}\sim\mathcal{N}(0,{\bm{I}})\ ,\ y_{\mu c}=\frac{1}{\sqrt{D}}\bm{\beta}_{c}\cdot\bm{x}^{\mu}_{c}+\sigma\epsilon^{\mu}_{c}\ ,\ \epsilon^{\mu}_{c}\sim\mathcal{N}(0,1).$ (5) $$

In Appendix [C.1](https://arxiv.org/html/2510.01098v1#A3.SS1), we analyze SGD with batch of $B$ contexts sampled at each step in the proportional asymptotics $P,K,B,D\to\infty\ ,\ P/D=\alpha\ ,\ K/D=\kappa\ ,\ B/D=\tau$, establishing that successful pretraining requires a total of $Bt=\Theta(D)$ contexts, each of size $P=\Theta(D)$.

Figure: (a) $L=16$, Linear Transformer
Refer to caption: x1.png

For $L=1$, the loss saturates to $\mathcal{L}_{\star}=(1+\alpha)^{-2}$ while for $L\to\infty$, the $\mathcal{L}_{\star}=[1-\alpha]_{+}$, which illustrates the gap in performance between a shallow model and a large depth model. We illustrate loss curves of varying depth $L$ compared to pretrained linear transformers (dots) in Figure [1](https://arxiv.org/html/2510.01098v1#S3.F1) (c). We extend to the case with label noise $\sigma^{2}>0$ in Appendix [C.2](https://arxiv.org/html/2510.01098v1#A3.SS2). In this case optimal $\gamma$ is smaller since early stopping acts as an effective regularization (Advani et al., 2020; Sonthalia et al., 2024).

Figure: (a) Loss Landscape $\alpha=0.5$
Refer to caption: x4.png

### 3.2 Fixed Structured Covariance ( FS ) → \to Preconditioned Gradient Descent

Next, we let the population covariance be arbitrary $\left<\bm{x}\bm{x}^{\top}\right>=\bm{\Sigma}$ across all contexts and task correlations be given by the matrix $\left<\bm{\beta}\bm{\beta}^{\top}\right>=\bm{\Omega}$. The ICL population loss takes the form

$$ $\displaystyle\mathcal{L}=\text{tr}\ \bm{\Omega}\left<\left[\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{L}\right]^{\top}\bm{\Sigma}\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{L}\right>$ (9) $$

We support this finding in Figure [3](https://arxiv.org/html/2510.01098v1#S3.F3) (b) where we show that small depth models are not outperformed by deeper models even after very long training horizons. When the ICL pretraining distribution for $\bm{\Sigma}$ is fixed, the model will memorize statistical information about the covariance of the inputs from the pretraining distribution. By preconditioning with the inverse of the data covariance $\bm{\Sigma}^{-1}$, the model is capable of achieving zero loss after even a single step of in-context GD. The gradient flow dynamics of the $\bm{\Gamma}$ matrix can be decomposed further in the case where $\bm{\Omega}$ and $\bm{\Sigma}$ commute.

Figure: (a) Eigenvalue Evolution
Refer to caption: x7.png

While pretraining the linear transformer in this setting can achieve zero loss for long contexts $P\to\infty$ even in shallow networks, the solution the model finds from gradient descent is brittle to properties of the pretraining data. Rather than learning a generic algorithm that solves ICL regression for any covariance $\bm{\Sigma}$, its solution is specialized to the pretraining covariance.

We illustrate this brittleness in Figure [3](https://arxiv.org/html/2510.01098v1#S3.F3) (c) where we define a family of new covariance matrices $\bm{\Sigma}^{\prime}=\exp(\theta\bm{S})\bm{\Sigma}\exp(-\theta\bm{S})$ where $\bm{S}$ is a random skew-symmetric matrix. As $\theta$ increases and $\bm{\Sigma}^{\prime}$ becomes more dissimilar to $\bm{\Sigma}$, the OOD loss increases for all depths $L$.

### 3.3 Random Rotated and Structured Covariances ⟹ \implies In-Context GD

Next, we attempt to enhance the covariance diversity across contexts. To do so, we now allow that each context $c$ has a random data and task covariances which are randomly rotated

$$ $\displaystyle\bm{x}^{\mu}_{c}\sim\mathcal{N}\left(0,\bm{\Sigma}_{c}\right)\ ,\ \bm{\Sigma}_{c}=\bm{O}_{c}\bm{\Lambda}\bm{O}_{c}^{\top}\ ,\ \bm{\Omega}_{c}=\bm{O}_{c}\bm{\Omega}\bm{O}_{c}^{\top},$ (12) $$

where $\bm{O}_{c}$ is a random $d\times d$ orthogonal matrix sampled from the Haar measure. The idea to pretrain with a diverse set of covariances $\bm{\Sigma}_{c}$ across contexts $c$ is to encourage the model to learn a generic in-context learning algorithm that is not specifically tailored to a particular data covariance $\bm{\Sigma}$. By introducing the random rotation across contexts, the model cannot encode a whitening transform of the data in the matrix $\bm{\Gamma}$ which prevents a zero loss solution in a shallow model with depth $L=1$. Therefore, even the $P/D\to\infty$ limit has the potential to exhibit a nontrivial scaling law in $L$.

This indicates that, provided the covariance is randomly rotated across contexts, the behavior of the learned solution is unconditioned in-context GD (see Appendix [E](https://arxiv.org/html/2510.01098v1#A5)). In the next section, we explore the consequences of this finding for optimal shapes.

## 4 Model of Compute Optimal Neural Scaling Laws

We consider the third setting with power law features and also introduce a notion of width through a projection matrix $\bm{A}\in\mathbb{R}^{N\times D}$. Rather than training on $D$ dimensional inputs $\bm{x}$, the model has access to $N$ dimensional features $\tilde{\bm{x}}=\bm{A}\bm{x}$, which leads to the following condition of $\bm{\Gamma}(t)$

$$ $\displaystyle\tilde{\bm{x}}=\bm{A}\bm{x}\quad{\implies}\quad\bm{\Gamma}(t)=\gamma(t)\left(\bm{A}\bm{A}^{\top}\right)\in\mathbb{R}^{N\times N}$ (14) $$

As before, the loss can again be viewed as a function of the scale parameter $\gamma(t)$. We provide a recipe to compute it below by computing the average over the random orthgonal matrix, resulting in a two-point deterministic equivalent for free products(^1^11Two-point refers to the correlation function of two resolvents evaluated at different arguments, rather than the one point function which is a single resolvent and determines only the spectrum of the random matrix.), in the same spirit of the results of (Bordelon et al., 2024a; Atanasov et al., 2025), using a saddle point method (Appendix [A](https://arxiv.org/html/2510.01098v1#A1)).

Figure: (a) Pretraining time $t$ scaling
Refer to caption: x10.png

In Appendix [F.4](https://arxiv.org/html/2510.01098v1#A6.SS4) we provide formulas for $\Psi_{v\chi}(\omega)$ and $\Psi_{vv}(\omega,\omega^{\prime})$ in the case that $\bm{A}^{\top}\bm{A}$ is a rank $N$ projection (has $N$ eigenvalues equal to $1$) and $\hat{\bm{\Sigma}}=\frac{1}{P}\bm{X}\bm{X}^{\top}$ is a structured Wishart matrix. This formula provides the full asymptotics. We can further extract Chinchilla scaling laws for powerlaw data. We demonstrate the need for both width and depth in Figures [4](https://arxiv.org/html/2510.01098v1#S4.F4) and [5](https://arxiv.org/html/2510.01098v1#S4.F5).

Figure: (a) Fixed Depth $L=1$
Refer to caption: x13.png

## 5 More Realistic Self Attention Models

In this section, we discuss more realistic models which exhibit similar training dynamics and depth dependence. First, we describe the dynamics when the $\bm{\Gamma}^{\ell}$ matrices are not tied. Second, we provide a theory of training when each of the weight matrices is optimized separately. Lastly, we provide experiments with softmax attention models trained with Adam.

### 5.1 Gradient Flow with Untied Γ \Gamma Matrices

In this section, we consider the effect of untied $\bm{\Gamma}^{\ell}$ matrices across layers $\ell\in[L]$. When each of the weights are optimized separately with a learning rate that is upscaled by depth $\eta=\eta_{0}L$, the dynamics are actually equivalent to the recurrent model that we presented previously under the RRS noise-free setting. In this setting, the matrices remain isotropic $\bm{\Gamma}^{\ell}=\gamma^{\ell}\bm{I}$ under gradient flow and leading to permutation symmetric loss function

$$ $\displaystyle\mathcal{L}(\{\gamma^{\ell}\})=\left<\left|\bm{\Lambda}^{1/2}\prod_{\ell=1}^{L}\left[\bm{I}-L^{-1}\gamma^{\ell}\bm{O}\left(\bm{A}^{\top}\bm{A}\right)^{2}\bm{O}^{\top}\hat{\bm{\Sigma}}\right]\bar{\bm{\beta}}\right|^{2}\right>$ (19) $$

where the average is over $\bm{O}$ and $\bar{\bm{\beta}}$. The evolution of $\gamma^{\ell}(t)$ under gradient flow will maintain balance $\gamma^{\ell}(t)=\gamma^{k}(t)$ if they start in a balanced initial condition such as $\gamma^{\ell}(0)=0$ for all $\ell$. As a consequence, the loss dynamics will be identical to the recurrent model that we analyzed in previous sections. We provide additional details in Appendix [G.1](https://arxiv.org/html/2510.01098v1#A7.SS1) and Figure [8](https://arxiv.org/html/2510.01098v1#A7.F8).

### 5.2 Gradient Flow for Full Linear Attention

In this section we consider gradient flow on all of the attention weights $\{\bm{W}_{k},\bm{W}_{q},\bm{W}_{v}\}$ separately rather than gradient flow on the $\bm{\Gamma}$ matrix. This corresponds to the dynamical system

$$ $\displaystyle\frac{d}{dt}\bm{W}_{i}=-\frac{\partial}{\partial\bm{W}_{i}}\mathcal{L}\ ,\ i\in\{x,y,k,q,v,o\}$ (20) $$

The dynamics of this model from small initialization are theoretically tractable as a set of low dimensional ODEs (see Appendix [G.2](https://arxiv.org/html/2510.01098v1#A7.SS2)), but suffers some defects due to transient blowup and recovery in the scale of $\bm{w}_{y}$ and $\bm{w}_{o}$. However, if we fix these weights to unit norm, the dynamics of the above model reduces to a one-dimensional ODE much like the reduced-$\bm{\Gamma}$ model.

We show the learning curves for this full attention module for isotropic covariates and powerlaw covariates in Figure [6](https://arxiv.org/html/2510.01098v1#S5.F6). The theoretical dynamics and predicted powerlaw exponent for this new dynamical system closely match the predictions.

### 5.3 Nonlinear (Softmax) Attention

We also provide experiments showing that similar scaling behavior holds in softmax attention models trained with Adam. In this setting, the attention block has the form $\bm{h}^{\ell+1}_{\mu}=\bm{h}^{\ell}_{\mu}+L^{-1}\sum_{\nu=1}^{P}A_{\mu\nu}\bm{W}_{v}\bm{h}^{\ell}_{\nu}$ where $A_{\mu\nu}^{\ell}=\frac{1}{Z_{\mu}}\exp(\bm{h}_{\nu}^{\top}\bm{W}_{k}\bm{W}_{q}\bm{h}^{\ell}_{\mu})$ with $Z_{\mu}=\sum_{\nu}\exp(\bm{h}_{\nu}^{\top}\bm{W}_{k}\bm{W}_{q}\bm{h}^{\ell}_{\mu})$ as the normalization factor. We train with Adam and show depth is still beneficial to ICL performance in Figure [6](https://arxiv.org/html/2510.01098v1#S5.F6) (c) and [9](https://arxiv.org/html/2510.01098v1#A7.F9), consistent with the phenomenology of our linear attention model.

Figure: (a) Loss Isotropic Features $\alpha=1$
Refer to caption: x16.png

## 6 Discussion

##### Conclusion

In this work we analyzed a solvable model of in-context learning with deep linear attention. We showed that the pretraining statistics strongly determine the type of solution the model converges to and under what conditions scaling depth is necessary at large context lengths. When training on fixed covariance structures, large depth is not necessary (at infinite context length) as the model learns a preconditioner that whitens the data. However, this learned solution is brittle to changes in the data covariance. When training on tasks where the data covariance is randomly rotated across contexts, the model learns a general purpose in-context gradient descent algorithm and exhibits a separable Chinchilla neural scaling law where limited width and depth can both bottleneck performance. Lastly, we show these results are robust to reparameterization of the attention blocks.

##### Limitations and Future Directions

While our work presents an advance in the solvable models of neural scaling laws and the structure of ICL in linear attention, there are many current limitations. The primary limitation is that we focus on linear regression tasks solved with linear attention. Characterization of more general ICL problems such as nonlinear function approximation and nonlinear attention could provide more insights into realistic in-context function approximation. Further, our analysis is focused on online learning in the present work. Future work could investigate overfitting effects caused by repeating tasks or context matrices during training, perhaps with dynamical mean field theory techniques (Mignacco et al., 2020; Bordelon et al., 2024a; Montanari & Urbani, 2025). Future work could also examine the role of large learning rate effects during pretraining dynamics. Another direction that could be interesting to explore in future works is scaling up the number of loops dynamically as the network trains to significantly reduce the total compute required.

## Acknowledgements

The authors would like to thank Jacob Zavatone-Veth, Alex Atanasov, Alexandru Meterez, Clarissa Lauditi, William Tong, Jamie Simon, Boris Hanin, Zhouran Yang, and Jascha Sohl-Dickstein for insightful discussions. B.B. acknowledges support from the Center of Mathematical Sciences and Applications
(CMSA) of Harvard University. C.P. is supported by an NSF CAREER Award (IIS-2239780), DARPA grants DIAL-FP-038 and AIQ-HR00112520041, the Simons Collaboration on the Physics of Learning and Neural Computation, and the William F. Milton Fund from Harvard University. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence.

## References

- Achiam et al. (2023)
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al.
Gpt-4 technical report.
*arXiv preprint arXiv:2303.08774*, 2023.
- Advani et al. (2020)
Madhu S Advani, Andrew M Saxe, and Haim Sompolinsky.
High-dimensional dynamics of generalization error in neural networks.
*Neural Networks*, 132:428–446, 2020.
- Ahn et al. (2023)
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra.
Transformers learn to implement preconditioned gradient descent for in-context learning, 2023.
URL [https://arxiv.org/abs/2306.00297](https://arxiv.org/abs/2306.00297).
- Akyürek et al. (2022)
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou.
What learning algorithm is in-context learning? investigations with linear models.
*arXiv preprint arXiv:2211.15661*, 2022.
- Atanasov et al. (2022)
Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan.
Neural networks as kernel learners: The silent alignment effect.
In *International Conference on Learning Representations*, 2022.
URL [https://openreview.net/forum?id=1NvflqAdoom](https://openreview.net/forum?id=1NvflqAdoom).
- Atanasov et al. (2023)
Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan.
The onset of variance-limited behavior for networks in the lazy and rich regimes.
In *The Eleventh International Conference on Learning Representations*, 2023.
URL [https://openreview.net/forum?id=JLINxPOVTh7](https://openreview.net/forum?id=JLINxPOVTh7).
- Atanasov et al. (2025)
Alexander Atanasov, Blake Bordelon, Jacob A Zavatone-Veth, Courtney Paquette, and Cengiz Pehlevan.
Two-point deterministic equivalence for stochastic gradient dynamics in linear models.
*arXiv preprint arXiv:2502.05074*, 2025.
- Bahri et al. (2021)
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma.
Explaining neural scaling laws.
*arXiv preprint arXiv:2102.06701*, 2021.
- Bai et al. (2023)
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei.
Transformers as statisticians: Provable in-context learning with in-context algorithm selection.
*Advances in neural information processing systems*, 36:57125–57211, 2023.
- Bender & Orszag (2013)
Carl M Bender and Steven A Orszag.
*Advanced mathematical methods for scientists and engineers I: Asymptotic methods and perturbation theory*.
Springer Science & Business Media, 2013.
- Bordelon & Pehlevan (2021)
Blake Bordelon and Cengiz Pehlevan.
Learning curves for sgd on structured features.
*arXiv preprint arXiv:2106.02713*, 2021.
- Bordelon & Pehlevan (2022a)
Blake Bordelon and Cengiz Pehlevan.
Learning curves for SGD on structured features.
In *International Conference on Learning Representations*, 2022a.
URL [https://openreview.net/forum?id=WPI2vbkAl3Q](https://openreview.net/forum?id=WPI2vbkAl3Q).
- Bordelon & Pehlevan (2022b)
Blake Bordelon and Cengiz Pehlevan.
Self-consistent dynamical field theory of kernel evolution in wide neural networks.
*Advances in Neural Information Processing Systems*, 35:32240–32256, 2022b.
- Bordelon et al. (a)
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan.
How feature learning can improve neural scaling laws.
In *The Thirteenth International Conference on Learning Representations*, a.
- Bordelon et al. (b)
Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan.
Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit.
In *The Twelfth International Conference on Learning Representations*, b.
- Bordelon et al. (2020)
Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan.
Spectrum dependent learning curves in kernel regression and wide neural networks.
In *International Conference on Machine Learning*, pp. 1024–1034. PMLR, 2020.
- Bordelon et al. (2024a)
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan.
A dynamical model of neural scaling laws.
In *Proceedings of the 41st International Conference on Machine Learning*, pp. 4345–4382, 2024a.
- Bordelon et al. (2024b)
Blake Bordelon, Hamza Chaudhry, and Cengiz Pehlevan.
Infinite limits of multi-head transformer dynamics.
*Advances in Neural Information Processing Systems*, 37:35824–35878, 2024b.
- Brown et al. (2020)
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.
Language models are few-shot learners.
*Advances in neural information processing systems*, 33:1877–1901, 2020.
- Chan et al. (2022)
Stephanie C. Y. Chan, Adam Santoro, Andrew K. Lampinen, Jane X. Wang, Aaditya Singh, Pierre H. Richemond, Jay McClelland, and Felix Hill.
Data distributional properties drive emergent in-context learning in transformers, 2022.
URL [https://arxiv.org/abs/2205.05055](https://arxiv.org/abs/2205.05055).
- Chen et al. (2024)
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang.
Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality.
*arXiv preprint arXiv:2402.19442*, 2024.
- Chizat & Bach (2018)
Lenaic Chizat and Francis Bach.
On the global convergence of gradient descent for over-parameterized models using optimal transport.
*Advances in neural information processing systems*, 31, 2018.
- Chizat & Netrapalli (2024)
Lénaïc Chizat and Praneeth Netrapalli.
The feature speed formula: a flexible approach to scale hyper-parameters of deep neural networks.
*Advances in Neural Information Processing Systems*, 37:62362–62383, 2024.
- Cui et al. (2021)
Hugo Cui, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová.
Generalization error rates in kernel regression: The crossover from the noiseless to noisy regime.
*Advances in Neural Information Processing Systems*, 34:10131–10143, 2021.
- Defilippis et al. (2024)
Leonardo Defilippis, Bruno Loureiro, and Theodor Misiakiewicz.
Dimension-free deterministic equivalents and scaling laws for random feature regression.
*Advances in Neural Information Processing Systems*, 37:104630–104693, 2024.
- Dehghani et al. (2018)
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser.
Universal transformers.
*arXiv preprint arXiv:1807.03819*, 2018.
- Dey et al. (2025)
Nolan Dey, Bin Claire Zhang, Lorenzo Noci, Mufan Li, Blake Bordelon, Shane Bergsma, Cengiz Pehlevan, Boris Hanin, and Joel Hestness.
Don’t be lazy: Completep enables compute-efficient deep transformers.
*arXiv preprint arXiv:2505.01618*, 2025.
- Dosovitskiy et al. (2020)
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.
An image is worth 16x16 words: Transformers for image recognition at scale.
*arXiv preprint arXiv:2010.11929*, 2020.
- Garg et al. (2022)
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant.
What can transformers learn in-context? a case study of simple function classes.
*Advances in neural information processing systems*, 35:30583–30598, 2022.
- Gatmiry et al. (2024)
Khashayar Gatmiry, Nikunj Saunshi, Sashank J. Reddi, Stefanie Jegelka, and Sanjiv Kumar.
On the role of depth and looping for in-context learning with task diversity, 2024.
URL [https://arxiv.org/abs/2410.21698](https://arxiv.org/abs/2410.21698).
- Geiger et al. (2020)
Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart.
Disentangling feature and lazy training in deep neural networks.
*Journal of Statistical Mechanics: Theory and Experiment*, 2020(11):113301, 2020.
- Geiping et al. (2025)
Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein.
Scaling up test-time compute with latent reasoning: A recurrent depth approach.
*arXiv preprint arXiv:2502.05171*, 2025.
- He et al. (2024)
Tianyu He, Darshil Doshi, Aritra Das, and Andrey Gromov.
Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks.
In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang (eds.), *Advances in Neural Information Processing Systems*, volume 37, pp. 13244–13273. Curran Associates, Inc., 2024.
URL [https://proceedings.neurips.cc/paper_files/paper/2024/file/17d60fef592086d1a5cb136f1946df59-Paper-Conference.pdf](https://proceedings.neurips.cc/paper_files/paper/2024/file/17d60fef592086d1a5cb136f1946df59-Paper-Conference.pdf).
- Hestness et al. (2017)
Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou.
Deep learning scaling is predictable, empirically.
*arXiv preprint arXiv:1712.00409*, 2017.
- Hoffmann et al. (2022)
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al.
Training compute-optimal large language models.
*arXiv preprint arXiv:2203.15556*, 2022.
- Kaplan et al. (2020)
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.
Scaling laws for neural language models.
*arXiv preprint arXiv:2001.08361*, 2020.
- Kim et al. (2024)
Juno Kim, Tai Nakamaki, and Taiji Suzuki.
Transformers are minimax optimal nonparametric in-context learners, 2024.
URL [https://arxiv.org/abs/2408.12186](https://arxiv.org/abs/2408.12186).
- Kratsios & Furuya (2025)
Anastasis Kratsios and Takashi Furuya.
Is in-context universality enough? mlps are also universal in-context, 2025.
URL [https://arxiv.org/abs/2502.03327](https://arxiv.org/abs/2502.03327).
- Kunstner & Bach (2025)
Frederik Kunstner and Francis Bach.
Scaling laws for gradient descent and sign descent for linear bigram models under zipf’s law.
*arXiv preprint arXiv:2505.19227*, 2025.
- Letey et al. (2025)
Mary I. Letey, Jacob A. Zavatone-Veth, Yue M. Lu, and Cengiz Pehlevan.
Pretrain-test task alignment governs generalization in in-context learning, 2025.
URL [https://arxiv.org/abs/2509.26551](https://arxiv.org/abs/2509.26551).
- Li et al. (2023)
Yingcong Li, M. Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak.
Transformers as algorithms: Generalization and stability in in-context learning, 2023.
URL [https://arxiv.org/abs/2301.07067](https://arxiv.org/abs/2301.07067).
- Lin et al. (2024)
Licong Lin, Jingfeng Wu, Sham M Kakade, Peter L Bartlett, and Jason D Lee.
Scaling laws in linear regression: Compute, parameters, and data.
*arXiv preprint arXiv:2406.08466*, 2024.
- Liu et al. (2024a)
Shang Liu, Zhongze Cai, Guanting Chen, and Xiaocheng Li.
Towards better understanding of in-context learning ability from in-context uncertainty quantification, 2024a.
URL [https://arxiv.org/abs/2405.15115](https://arxiv.org/abs/2405.15115).
- Liu et al. (2024b)
Toni JB Liu, Nicolas Boullé, Raphaël Sarfati, and Christopher J Earls.
Llms learn governing principles of dynamical systems, revealing an in-context neural scaling law.
*arXiv preprint arXiv:2402.00795*, 2024b.
- Lu et al. (2025)
Yue M. Lu, Mary Letey, Jacob A. Zavatone-Veth, Anindita Maiti, and Cengiz Pehlevan.
Asymptotic theory of in-context learning by linear attention.
*Proceedings of the National Academy of Sciences*, 122(28):e2502599122, 2025.
doi: 10.1073/pnas.2502599122.
URL [https://www.pnas.org/doi/abs/10.1073/pnas.2502599122](https://www.pnas.org/doi/abs/10.1073/pnas.2502599122).
- Mahankali et al. (2023)
Arvind Mahankali, Tatsunori B. Hashimoto, and Tengyu Ma.
One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention, 2023.
URL [https://arxiv.org/abs/2307.03576](https://arxiv.org/abs/2307.03576).
- Maloney et al. (2022)
Alexander Maloney, Daniel A Roberts, and James Sully.
A solvable model of neural scaling laws.
*arXiv preprint arXiv:2210.16859*, 2022.
- McCracken et al. (2025)
Gavin McCracken, Gabriela Moisescu-Pareja, Vincent Letourneau, Doina Precup, and Jonathan Love.
Uncovering a universal abstract algorithm for modular addition in neural networks, 2025.
URL [https://arxiv.org/abs/2505.18266](https://arxiv.org/abs/2505.18266).
- Mei et al. (2018)
Song Mei, Andrea Montanari, and Phan-Minh Nguyen.
A mean field view of the landscape of two-layer neural networks.
*Proceedings of the National Academy of Sciences*, 115(33):E7665–E7671, 2018.
- Merrill & Sabharwal (2025)
William Merrill and Ashish Sabharwal.
A little depth goes a long way: The expressive power of log-depth transformers.
*arXiv preprint arXiv:2503.03961*, 2025.
- Mignacco et al. (2020)
Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová.
Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification.
*Advances in Neural Information Processing Systems*, 33:9540–9550, 2020.
- Montanari & Urbani (2025)
Andrea Montanari and Pierfrancesco Urbani.
Dynamical decoupling of generalization and overfitting in large two-layer networks.
*arXiv preprint arXiv:2502.21269*, 2025.
- Noci et al. (2023)
Lorenzo Noci, Chuning Li, Mufan Li, Bobby He, Thomas Hofmann, Chris J Maddison, and Dan Roy.
The shaped transformer: Attention models in the infinite depth-and-width limit.
*Advances in Neural Information Processing Systems*, 36:54250–54281, 2023.
- Panwar et al. (2024)
Madhur Panwar, Kabir Ahuja, and Navin Goyal.
In-context learning through the bayesian prism, 2024.
URL [https://arxiv.org/abs/2306.04891](https://arxiv.org/abs/2306.04891).
- Paquette et al. (2024)
Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington.
4+ 3 phases of compute-optimal neural scaling laws.
*arXiv preprint arXiv:2405.15074*, 2024.
- Park et al. (2024)
Core Francisco Park, Andrew Lee, Ekdeep Singh Lubana, Yongyi Yang, Maya Okawa, Kento Nishi, Martin Wattenberg, and Hidenori Tanaka.
Iclr: In-context learning of representations.
*arXiv preprint arXiv:2501.00070*, 2024.
- Pathak et al. (2023)
Reese Pathak, Rajat Sen, Weihao Kong, and Abhimanyu Das.
Transformers can optimally learn regression mixture models, 2023.
URL [https://arxiv.org/abs/2311.08362](https://arxiv.org/abs/2311.08362).
- Potters & Bouchaud (2020)
Marc Potters and Jean-Philippe Bouchaud.
*A first course in random matrix theory: for physicists, engineers and data scientists*.
Cambridge University Press, 2020.
- Raventós et al. (2023)
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli.
Pretraining task diversity and the emergence of non-bayesian in-context learning for regression, 2023.
URL [https://arxiv.org/abs/2306.15063](https://arxiv.org/abs/2306.15063).
- Simon et al. (2023)
James B Simon, Dhruva Karkada, Nikhil Ghosh, and Mikhail Belkin.
More is better in modern machine learning: when infinite overparameterization is optimal and overfitting is obligatory.
*arXiv preprint arXiv:2311.14646*, 2023.
- Sonthalia et al. (2024)
Rishi Sonthalia, Jackie Lok, and Elizaveta Rebrova.
On regularization via early stopping for least squares regression.
*arXiv preprint arXiv:2406.04425*, 2024.
- Spigler et al. (2020)
Stefano Spigler, Mario Geiger, and Matthieu Wyart.
Asymptotic learning curves of kernel methods: empirical data versus teacher–student paradigm.
*Journal of Statistical Mechanics: Theory and Experiment*, 2020(12):124001, 2020.
- Tong & Pehlevan (2025)
William L. Tong and Cengiz Pehlevan.
Mlps learn in-context on regression and classification tasks, 2025.
URL [https://arxiv.org/abs/2405.15618](https://arxiv.org/abs/2405.15618).
- Vaswani et al. (2017)
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.
Attention is all you need.
*Advances in neural information processing systems*, 30, 2017.
- Vladymyrov et al. (2024)
Max Vladymyrov, Johannes von Oswald, Mark Sandler, and Rong Ge.
Linear transformers are versatile in-context learners, 2024.
URL [https://arxiv.org/abs/2402.14180](https://arxiv.org/abs/2402.14180).
- Von Oswald et al. (2023)
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov.
Transformers learn in-context by gradient descent.
In *International Conference on Machine Learning*, pp. 35151–35174. PMLR, 2023.
- Wang et al. (2025)
Guan Wang, Jin Li, Yuhao Sun, Xing Chen, Changling Liu, Yue Wu, Meng Lu, Sen Song, and Yasin Abbasi Yadkori.
Hierarchical reasoning model.
*arXiv preprint arXiv:2506.21734*, 2025.
- Wu et al. (2024)
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L. Bartlett.
How many pretraining tasks are needed for in-context learning of linear regression?, 2024.
URL [https://arxiv.org/abs/2310.08391](https://arxiv.org/abs/2310.08391).
- Wurgaft et al. (2025)
Daniel Wurgaft, Ekdeep Singh Lubana, Core Francisco Park, Hidenori Tanaka, Gautam Reddy, and Noah D Goodman.
In-context learning strategies emerge rationally.
*arXiv preprint arXiv:2506.17859*, 2025.
- Xie et al. (2021)
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma.
An explanation of in-context learning as implicit bayesian inference.
*arXiv preprint arXiv:2111.02080*, 2021.
- Yang & Hu (2021)
Greg Yang and Edward J Hu.
Tensor programs iv: Feature learning in infinite-width neural networks.
In *International Conference on Machine Learning*, pp. 11727–11737. PMLR, 2021.
- Yang et al. (2021)
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao.
Tuning large neural networks via zero-shot hyperparameter transfer.
In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), *Advances in Neural Information Processing Systems*, 2021.
URL [https://openreview.net/forum?id=Bx6qKuBM2AD](https://openreview.net/forum?id=Bx6qKuBM2AD).
- Yang et al. (2023)
Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou.
Tensor programs vi: Feature learning in infinite-depth neural networks.
*arXiv preprint arXiv:2310.02244*, 2023.
- Zhang et al. (2023a)
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett.
Trained transformers learn linear models in-context, 2023a.
URL [https://arxiv.org/abs/2306.09927](https://arxiv.org/abs/2306.09927).
- Zhang et al. (2025)
Yedi Zhang, Aaditya K. Singh, Peter E. Latham, and Andrew Saxe.
Training dynamics of in-context learning in linear attention, 2025.
URL [https://arxiv.org/abs/2501.16265](https://arxiv.org/abs/2501.16265).
- Zhang et al. (2023b)
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang.
What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization.
*arXiv preprint arXiv:2305.19420*, 2023b.
- Zhu et al. (2025)
Ruike Zhu, Hanwen Zhang, Tianyu Shi, Chi Wang, Tianyi Zhou, and Zengyi Qin.
The 4th dimension for scaling model size.
*arXiv preprint arXiv:2506.18233*, 2025.

## Appendix

## Appendix A Two Point Deterministic Equivalent for General Free Product From Dynamical Mean Field Theory

In this section, we describe a method to handle two-point correlation properties of free products of random matrices. Similar results were provided in a simple random feature model of (Bordelon et al., 2024a; Atanasov et al., 2025). This computation provides the necessary technical results to establish the behavior of ICL in the randomly rotated context setting. Note throughout that we will use $\text{tr}\bm{M}$ for normalized trace of the matrix $\bm{M}$ so that for $N\times N$ matrix $\text{tr}\bm{M}=\frac{1}{N}\sum_{k=1}^{N}M_{kk}$.

##### Tracking Linear Dynamics Generated by Free Product

Taking a dynamical approach, we consider the following dynamical system which depends on a random matrix $\bm{M}\in\mathbb{R}^{N\times N}$

$$ $\displaystyle\partial_{t}\bm{v}(t)=-\bm{M}\bm{v}(t)+\delta(t)\bm{v}_{0}\ ,\ \bm{M}=\bm{O}\bm{B}\bm{O}^{\top}\bm{A}$ (23) $$

where $\bm{O}$ is a random $N\times N$ orthogonal matrix drawn from the Haar measure for the orthogonal group. Our starting point is to express the above dynamical system with an integral representation of Dirac-Delta functions $1=\int dz\delta(z)=\int\frac{dzd\chi}{2\pi i}e^{-\chi z}$ where the integration variable $\chi\in(-i\infty,i\infty)$ runs along the imaginary axis. Performing this for each time $t$ of the dynamics yields

$$ $\displaystyle Z=\int\mathcal{D}\bm{v}\mathcal{D}\bm{\chi}\left<\exp\left(-\int dt\bm{\chi}(t)\left[\partial_{t}\bm{v}(t)+\bm{M}\bm{v}(t)-\bm{v}_{0}\delta(t)\right]\right)\right>=1$ (24) $$

where $\mathcal{D}\bm{v}\mathcal{D}\bm{\chi}=\lim_{\delta t\to 0}\prod_{n=-\infty}^{\infty}\frac{d\bm{v}(n\cdot\delta t)d\bm{\chi}(n\cdot\delta t)}{2\pi i}$ is the flat measure over $\bm{v}(t),\bm{\chi}(t)$ in the continuous time limit $\delta t\to 0$. The average $\left<\cdot\right>$ is computed over the random orthogonal matrix $\bm{O}$.

To simplify the calculation, we transform our dynamics into Fourier space

$$ $\displaystyle\bm{v}(t)\equiv\int\frac{d\omega}{2\pi}\exp\left(i\omega t\right)\bm{v}(\omega)\ ,\ \bm{\chi}(t)\equiv\int\frac{d\omega}{2\pi}\exp\left(i\omega t\right)\bm{\chi}(\omega),$ (25) $$

which transforms the original integral over $t$ into an integral over $\omega$ as $\int d\omega\bm{\chi}(\omega)\cdot[(i\omega)\bm{v}(\omega)+\bm{M}\bm{v}(\omega)-\bm{v}_{0}]$.

##### Disorder Average

Averaging the resulting expression over the orthogonal matrix $\bm{O}$, we obtain the following representation of $Z$ as an integral over order parameters $\{\bm{\Sigma},\bm{\Psi}\}$

$$ $\displaystyle Z=\int\mathcal{D}\bm{\Sigma}\mathcal{D}\bm{\Psi}\exp\left(-N\mathcal{S}[\bm{\Sigma},\bm{\Psi}]\right).$ (26) $$

In the above expression, the action $\mathcal{S}$ has the form

$$ $\displaystyle\mathcal{S}[\bm{\Sigma},\bm{\Psi}]=-\text{Tr}\bm{\Sigma}\bm{\Psi}-\frac{1}{N}\ln\mathcal{Z}_{A}(\bm{\Psi})-\text{Tr}\bm{\Sigma}\hat{\bm{\Sigma}}_{\star}-\frac{1}{N}\ln\det\left(\hat{\bm{\Sigma}}_{\star}\otimes\bm{I}+\bm{V}\otimes\bm{B}\right)+\ln\det\bm{\Sigma}$ (27) $$

where $\hat{\bm{\Sigma}}_{\star}$ solves the implicit equation $\frac{\partial\mathcal{S}}{\partial\hat{\bm{\Sigma}}_{\star}}=-\bm{\Sigma}+\text{tr}\left(\bm{\hat{\Sigma}}_{\star}\otimes\bm{I}+\bm{V}\otimes\bm{B}\right)^{-1}=0$ and the single site measure $\mathcal{Z}_{A}$ has the form

$$ $\displaystyle\mathcal{Z}_{A}(\bm{\Psi})$ $\displaystyle=\int\mathcal{D}\bm{v}\mathcal{D}\bm{\chi}\exp\left(-\int d\omega\bm{\chi}(\omega)\cdot\left[(i\omega)\bm{v}(z)-\bm{v}_{0}-\int d\omega^{\prime}\Psi_{v\chi}(\omega,\omega^{\prime})\bm{A}\bm{v}(\omega^{\prime})\right]\right)$ $\displaystyle\exp\left(-\frac{1}{2}\int d\omega d\omega^{\prime}\left[\Psi_{vv}(\omega,\omega^{\prime})\bm{v}(\omega)^{\top}\bm{A}^{2}\bm{v}(\omega^{\prime})+\Psi_{\chi\chi}(\omega,\omega^{\prime})\bm{\chi}(\omega)\cdot\bm{\chi}(\omega^{\prime})\right]\right)$ (28) $$

##### Taking the Large System Size Limit

To study the limit of $N\to\infty$, the saddle point equations over $\bm{\Sigma}$ and $\bm{\Psi}$ give

$$ $\displaystyle\frac{\partial\mathcal{S}}{\partial\bm{\Psi}(\omega,\omega^{\prime})}=-\begin{bmatrix}\Sigma_{vv}(\omega,\omega^{\prime})&\Sigma_{v\chi}(\omega,\omega^{\prime})\\ \Sigma_{\chi v}(\omega,\omega^{\prime})&\Sigma_{\chi\chi}(\omega,\omega^{\prime})\end{bmatrix}+\begin{bmatrix}\frac{1}{N}\left<\bm{v}(\omega)^{\top}\bm{A}^{2}\bm{v}(\omega^{\prime})\right>&\frac{1}{N}\left<\bm{v}(\omega)^{\top}\bm{A}\bm{\chi}(\omega^{\prime})\right>\\ \frac{1}{N}\left<\bm{\chi}(\omega)^{\top}\bm{A}\bm{v}(\omega^{\prime})\right>&\frac{1}{N}\left<\bm{\chi}(\omega)\cdot\bm{\chi}(\omega^{\prime})\right>\end{bmatrix}$ $\displaystyle\frac{\partial\mathcal{S}}{\partial\bm{\Sigma}(\omega,\omega^{\prime})}=-\bm{\Psi}(\omega,\omega^{\prime})-\bm{\hat{\Sigma}}_{\star}(\omega,\omega^{\prime})+[\bm{\Sigma}^{-1}](\omega,\omega^{\prime})=0$ (29) $$

The average over $\left<\cdot\right>$ represents an average over the Gaussian measure defined by $\mathcal{Z}_{A}$. The solution to these equations has the following block structure

$$ $\displaystyle\bm{\Sigma}(\omega,\omega^{\prime})=\begin{bmatrix}\Sigma_{vv}(\omega,\omega^{\prime})&\Sigma_{v\chi}(\omega,\omega^{\prime})\\ \Sigma_{v\chi}(\omega,\omega^{\prime})&0\end{bmatrix}$ (30) $\displaystyle\hat{\bm{\Sigma}}(\omega,\omega^{\prime})=\begin{bmatrix}0&\hat{\Sigma}_{v\chi}(\omega,\omega^{\prime})\\ \hat{\Sigma}_{v\chi}(\omega,\omega^{\prime})&\hat{\Sigma}_{\chi\chi}(\omega,\omega^{\prime})\end{bmatrix}$ (31) $\displaystyle\bm{\Psi}(\omega,\omega^{\prime})=\begin{bmatrix}0&\Psi_{v\chi}(\omega,\omega^{\prime})\\ \Psi_{v\chi}(\omega,\omega^{\prime})&\Psi_{\chi\chi}(\omega,\omega^{\prime})\end{bmatrix}$ (32) $$

##### Off Diagonal Blocks

The off diagonal blocks decouple over different frequencies $\Sigma_{v\chi}(\omega,\omega^{\prime})=\delta(\omega-\omega^{\prime})\Sigma_{v\chi}(\omega)$ and satisfy the following equations

$$ $\displaystyle\Sigma_{v\chi}(\omega)=\text{tr}\bm{A}\left(i\omega+\Psi_{v\chi}(\omega)\bm{A}\right)^{-1}=\text{tr}\left(\hat{\Sigma}_{v\chi}(\omega)+\bm{B}\right)^{-1},$ (33) $\displaystyle\Psi_{v\chi}(\omega)\Sigma_{v\chi}(\omega)=1-\hat{\Sigma}_{v\chi}(\omega)\Sigma_{v\chi}(\omega)$ (34) $$

Introduce the $\tau$-transform $\tau_{M}$ of a matrix $\bm{M}$ as

$$ $\displaystyle\tau_{M}(i\omega)\equiv\text{tr}\bm{M}(i\omega+\bm{M})^{-1}$ (35) $$

as well as its inverse function $i\omega_{M}(\tau)$ (^2^22Up to a change in signs and $i\omega\to-z$ this is equivalent to the $t$-transform of a random matrix (Potters & Bouchaud, 2020).). Then our saddle point equations give us

$$ $\displaystyle\tau=\Sigma_{v\chi}(\omega)\Psi_{v\chi}(\omega)=\tau_{A}(i\omega_{A})=\tau_{B}(i\omega_{B})$ (36) $$

where $i\omega_{A}=\frac{i\omega}{\Psi_{v\chi}(\omega)}$ and $i\omega_{B}=(\tau^{-1}-1)\Psi_{v\chi}(\omega)$. Putting these equations together, we find

$$ $\displaystyle(i\omega_{A}(\tau))(i\omega_{B}(\tau))=\frac{1-\tau}{\tau}(i\omega)$ (37) $$

This equation is to be solved for $\tau(\omega)$. Once this function is determined we can use it to compute the diagonal blocks of $\bm{\Sigma},\bm{\Psi}$ as we describe in the next section.

##### Diagonal Blocks

Now, we can determine the diagonal blocks, which determine the covariance structure of the $\bm{v}(\omega)$ variables

$$ $\displaystyle\Sigma_{vv}(\omega,\omega^{\prime})$ $\displaystyle=\frac{1}{N}\bm{v}_{0}^{\top}\left(i\omega+\Psi_{v\chi}(\omega)\bm{A}\right)^{-1}\bm{A}^{2}\left(i\omega^{\prime}+\Psi_{v\chi}(\omega^{\prime})\bm{A}\right)^{-1}\bm{v}_{0}$ $\displaystyle-\Psi_{\chi\chi}(\omega,\omega^{\prime})\ \underbrace{\text{tr}\bm{A}^{2}\left(i\omega+\Psi_{v\chi}(\omega)\bm{A}\right)^{-1}\left(i\omega^{\prime}+\Psi_{v\chi}(\omega^{\prime})\bm{A}\right)^{-1}}_{\mathcal{A}(\omega,\omega^{\prime})}$ $\displaystyle\Sigma_{vv}(\omega,\omega^{\prime})$ $\displaystyle=-\hat{\Sigma}_{\chi\chi}(\omega,\omega^{\prime})\ \underbrace{\text{tr}\left(\hat{\Sigma}_{v\chi}(\omega)+\bm{B}\right)^{-1}\left(\hat{\Sigma}_{v\chi}(\omega^{\prime})+\bm{B}\right)^{-1}}_{\mathcal{B}(\omega,\omega^{\prime})}$ $\displaystyle\Psi_{\chi\chi}(\omega,\omega^{\prime})$ $\displaystyle=-\hat{\Sigma}_{\chi\chi}(\omega,\omega^{\prime})-\Sigma_{vv}(\omega,\omega^{\prime})\Sigma_{v\chi}(\omega)^{-1}\Sigma_{v\chi}(\omega^{\prime})^{-1}$ (38) $$

where we introduced functions $\mathcal{A}$ and $\mathcal{B}$ which can be determined from the (already obtained) off-diagonal blocks. Combining these equations

$$ $\displaystyle\Sigma_{vv}(\omega,\omega^{\prime})$ $\displaystyle=\frac{\frac{1}{N}\bm{v}_{0}^{\top}\left(i\omega+\Psi_{v\chi}(\omega)\bm{A}\right)^{-1}\bm{A}^{2}\left(i\omega^{\prime}+\Psi_{v\chi}(\omega^{\prime})\bm{A}\right)^{-1}\bm{v}_{0}}{1-\left[\Sigma_{v\chi}(\omega)^{-1}\Sigma_{v\chi}(\omega^{\prime})^{-1}-\mathcal{B}(\omega,\omega^{\prime})^{-1}\right]\mathcal{A}(\omega,\omega^{\prime})}$ (39) $$

From this expression, both $\hat{\Sigma}_{\chi\chi}(\omega,\omega^{\prime})=-\mathcal{B}(\omega,\omega^{\prime})\Sigma_{vv}(\omega,\omega^{\prime})$ and $\Psi_{\chi\chi}(\omega,\omega^{\prime})=-\hat{\Sigma}_{\chi\chi}(\omega,\omega^{\prime})-\Sigma_{vv}(\omega,\omega^{\prime})\Sigma_{v\chi}(\omega)^{-1}\Sigma_{v\chi}(\omega^{\prime})^{-1}$ are determined.

##### Deterministic Equivalent

From the functions $\Psi_{v\chi}(\omega)$ and $\Psi_{\chi\chi}(\omega,\omega^{\prime})$ we obtain the following determinstic equivalent for the outer product of $\bm{v}$ variables

$$ $\displaystyle\bm{v}(\omega)\bm{v}(\omega^{\prime})^{\top}\simeq$ $\displaystyle\left(i\omega+\Psi_{v\chi}(\omega)\bm{A}\right)^{-1}\left[\bm{v}_{0}\bm{v}_{0}-\Psi_{\chi\chi}(\omega,\omega^{\prime})\bm{I}\right]\left(i\omega^{\prime}+\Psi_{v\chi}(\omega^{\prime})\bm{A}\right)^{-1}$ (40) $$

where the $\simeq$ expression holds under trace against a test matrix (i.e. $\bm{M}_{1}\simeq\bm{M}_{2}\implies\text{tr}\bm{C}\bm{M}_{1}=\text{tr}\bm{C}\bm{M}_{2}$) (Potters & Bouchaud, 2020). This function will be used in subsequent expressions to give an exact result for the loss landscape of our randomly rotated ICL loss function. Approximations of this will give rise to our scaling law results as we discuss in Appendix [F.4](https://arxiv.org/html/2510.01098v1#A6.SS4).

## Appendix B Model Definition, Masking Mechanics, Reduced- Γ \Gamma Model

### B.1 Linear Attention and Positional Masking For our Task

In this section, we provide more detail on the structure of the masking used in our task. First, we note that the target values $y_{\mu}$ are only provided for the $P$ labeled examples within each context matrix. Second, we note that the residual attention operations are masked differently for the $P$ labeled examples and the $K$ evaluation points. To define our masking operation precisely, we introduce the notation $\bm{1}_{k}\in\mathbb{R}^{k}$ as a vector consisting of all $k$ entries equal to one. We introduce the masking matrix $\bm{M}\in\mathbb{R}^{(P+K)\times(P+K)}$ for the residual stream which has the following block structure

$$ $\displaystyle\bm{M}=\begin{bmatrix}-\bm{1}_{P}\ \bm{1}_{P}^{\top}&\bm{0}\\ \bm{1}_{K}\ \bm{1}_{P}^{\top}&\bm{0}\end{bmatrix}.$ (41) $$

Now that this positional masking matrix is introduced, we can conveniently express our update rule for the residual stream

$$ $\displaystyle\bm{h}^{\ell+1}_{\mu}$ $\displaystyle=\bm{h}^{\ell}_{\mu}+\frac{1}{PL}\sum_{\nu=1}^{P+K}M_{\mu\nu}\left(\bm{k}^{\ell}_{\nu}\cdot\bm{q}^{\ell}_{\mu}\right)\bm{v}_{\nu}^{\ell}$ $\displaystyle=\bm{h}^{\ell}_{\mu}+\frac{1}{PL}\sum_{\nu=1}^{P}M_{\mu\nu}\left(\bm{k}^{\ell}_{\nu}\cdot\bm{q}^{\ell}_{\mu}\right)\bm{v}_{\nu}^{\ell}\ ,\ \mu\in[P+K]$ (42) $$

where $\bm{k}^{\ell}_{\mu}=\bm{W}_{k}\cdot\bm{h}^{\ell}_{\mu},\bm{q}^{\ell}_{\mu}=\bm{W}_{q}\bm{h}^{\ell}_{\mu},\bm{v}^{\ell}_{\mu}=\bm{W}^{\ell}_{v}\bm{h}^{\ell}_{\mu}$ and we dropped the sum over evaluation points $\{P+1,...,P+K\}$ due to the structure of the positional mask $\bm{M}$. We thus see that the masked update rule has two properties

- 1.
It prevents the test points $\bm{x}_{\mu}$ for $\mu\in\{P+1,...,P+K\}$ from being used in the residual stream updates. Only the first $P$ training points are utilized.
- 2.
It provides an opposite sign for the updates on training points and on testing points. We will see that this will enable the model to implement an in-context gradient descent rule. In such a rule, a subspace of the residual variables will encode residual errors $y_{\mu}-f_{\mu}^{\ell}$ for the training predictions and $+f_{\mu}^{\ell}$ for the test points $\mu\in\{P+1,..,P+K\}$. Instead, one could use the same signs for training points and test points in $\bm{M}$ and simply negate the output of the model at the end $f\to-f$.

### B.2 Derivation of Reduced Gamma Model

In this section, we describe the conditions under which a linear attention transformer model can be reparameterized as the recurrent reduced-$\Gamma$ model we discuss in the main text.

##### Alignment Assumptions

Following Von Oswald et al. (2023); Zhang et al. (2025); Lu et al. (2025), we study configurations of weights that encode information about inputs $\bm{x}$ and targets $y$ in orthogonal subspaces of the residual stream. Concretely, the following assumptions are made on the input weights, which implement in-context preconditioned gradient descent

$$ $\displaystyle\bm{W}_{x}^{\top}\bm{w}_{y}=0\ ,\ \bm{W}_{x}^{\top}\bm{w}_{o}=0\ ,\ \bm{w}_{y}\cdot\bm{w}_{y}=|\bm{w}_{y}||\bm{w}_{y}|.$ (43) $$

Collectively these assumptions imply that read-in weights for the targets $\bm{w}_{y}$ and readout weights $\bm{w}_{o}$ perfectly align and that the read-in weights for $\bm{x}$ $\bm{W}_{x}$ project input data to an orthgonal subspace. Next we study the following set of alignment conditions for the key, query and value matrices

$$ $\displaystyle\bm{W}_{x}^{\top}\left(\bm{W}_{k}^{\ell}\right)^{\top}\bm{W}_{q}^{\ell}\bm{W}_{x}\propto\bm{\Gamma}^{\ell}\in\mathbb{R}^{D\times D}$ (44) $\displaystyle\bm{W}_{v}\bm{W}_{x}=0\ ,\ \bm{W}_{v}\bm{w}_{y}\propto\bm{w}_{y}$ (45) $$

Under these assumptions, we can define a collection of $D\times D$ matrices $\bm{\Gamma}^{\ell}$

$$ $\displaystyle\bm{\Gamma}^{\ell}\equiv\left(\bm{w}_{o}^{\top}\bm{W}_{v}^{\ell}\bm{w}_{y}\right)\bm{W}_{x}^{\top}(\bm{W}_{k}^{\ell})^{\top}(\bm{W}_{q}^{\ell})\bm{W}_{x}\in\mathbb{R}^{D\times D},$ (46) $$

which gives rise to the following residual stream dynamics for $\Delta_{\mu}^{\ell}\equiv\bm{w}_{o}\cdot\bm{h}^{\ell}$

$$ $\displaystyle\Delta^{\ell+1}_{\mu}=\Delta^{\ell}_{\mu}+\frac{1}{LP}\sum_{\nu=1}^{P}M_{\mu\nu}\bm{x}_{\nu}^{\top}\bm{\Gamma}^{\ell}\bm{x}_{\mu}\Delta_{\nu}^{\ell}.$ (47) $$

We note that this can be separated into two distinct dynamical systems, one for the first $P$ training points

$$ $\displaystyle\Delta^{\ell+1}_{\mu}=\Delta^{\ell}_{\mu}-\frac{1}{LP}\sum_{\nu=1}^{P}\bm{x}_{\nu}^{\top}\bm{\Gamma}^{\ell}\bm{x}_{\mu}\Delta_{\nu}^{\ell}\ ,\ \mu\in\{1,...,P\}$ (48) $$

which form a closed dynamical system on the $P$ labeled training points. From these $P$ variables $\{\Delta_{\mu}^{\ell}\}_{\mu\in[P]}$, we can describe how the remaining $K$ points evolve

$$ $\displaystyle\Delta^{\ell+1}_{\mu}=\Delta^{\ell}_{\mu}+\frac{1}{LP}\sum_{\nu=1}^{P}\bm{x}_{\nu}^{\top}\bm{\Gamma}^{\ell}\bm{x}_{\mu}\Delta_{\nu}^{\ell}=\frac{1}{LP}\sum_{\ell=1}^{L}\sum_{\nu=1}^{P}\bm{x}_{\nu}^{\top}\bm{\Gamma}^{\ell}\bm{x}_{\mu}\Delta_{\nu}^{\ell}\ ,\ \mu\in\{P+1,...,P+K\}.$ (49) $$

##### Recurrence

Instead of defining separate $\bm{\Gamma}^{\ell}$ matrices for each layer $\ell$, we instead can examine a recurrent model where the weights are tied $\bm{\Gamma}^{\ell}=\bm{\Gamma}$ for all $\ell\in[L]$. We relax this assumption in Appendix [G.1](https://arxiv.org/html/2510.01098v1#A7.SS1) and in many of the settings we analyze recurrence has no impact compared to decoupling layers. Under this constraint, the residual stream of the model has the following dynamics

$$ $\displaystyle\Delta_{\mu}^{\ell+1}$ $\displaystyle=\Delta^{\ell}_{\mu}-\frac{1}{LP}\sum_{\nu=1}^{P}\bm{x}_{\mu}^{\top}\bm{\Gamma}\bm{x}_{\nu}\ \Delta^{\ell}_{\nu}\ ,\ \mu\in\{1,...,P\}$ $\displaystyle\Delta_{\mu}^{\ell+1}$ $\displaystyle=\Delta_{\mu}^{\ell}+\frac{1}{LP}\sum_{\nu=1}^{P}\bm{x}_{\mu}^{\top}\bm{\Gamma}\bm{x}_{\nu}\ \Delta^{\ell}_{\nu}\ ,\ \mu\in\{P+1,...,P+K\}$ (50) $$

Let $\bm{\Delta}^{\ell}\in\mathbb{R}^{P}$ represent the residual stream variables restricted to the $P$ unmasked training points in the context. This vector has the residual stream dynamics

$$ $\displaystyle\bm{\Delta}^{\ell}=\left(\bm{I}-\frac{1}{LP}\bm{X}^{\top}\bm{\Gamma}\bm{X}\right)^{\ell}\bm{y}.$ (51) $$

However, the recursion is different for the test set since these points only receive attention signals from the $P$ unmasked training tokens. For one of the test points $\bm{x}_{\star}$, the prediction of the model $f_{\star}$ satisfies

$$ $\displaystyle f_{\star}=\frac{1}{LP}\bm{x}_{\star}^{\top}\bm{\Gamma}\bm{X}\sum_{\ell=0}^{L-1}\bm{h}^{\ell}=\frac{1}{LP}\bm{x}_{\star}^{\top}\bm{\Gamma}\bm{X}\sum_{\ell=0}^{L-1}\left(\bm{I}-\frac{1}{LP}\bm{X}^{\top}\bm{\Gamma}\bm{X}\right)^{\ell}\bm{y}$ (52) $$

##### Equivalence to Preconditioned In-Context Gradient Descent

This update rule is equivalent to implementing preconditioned in-context gradient descent steps for a linear regression model $f=\frac{1}{\sqrt{D}}\bm{\beta}\cdot\bm{x}$. To see this, define an in-context training loss $\hat{\mathcal{L}}(\bm{\beta},\bm{D})$ on the $P$ labeled training points for context $\bm{D}$

$$ $\displaystyle\hat{\mathcal{L}}(\bm{D})=\frac{1}{2P}||D^{-1/2}\bm{X}^{\top}\bm{\beta}-\bm{y}||^{2}$ (53) $$

Preconditioned gradient descent with learning rate $D/L$ and a preconditioner matrix $\bm{\Gamma}$ generates the following dynamics on the learned ICL weights $\bm{\beta}$

$$ $\displaystyle\bm{\beta}^{\ell+1}=\bm{\beta}^{\ell}-L^{-1}\bm{\Gamma}\ \nabla\mathcal{L}(\bm{\beta}^{\ell},\bm{D})=\bm{\beta}^{\ell}-\frac{\sqrt{D}}{LP}\bm{\Gamma}\bm{X}\left(D^{-1/2}\bm{X}^{\top}\bm{\beta}^{\ell}-\bm{y}\right)$ (54) $$

Defining a variable $\bm{h}^{\ell}\equiv\bm{y}-D^{-1/2}\bm{X}^{\top}\bm{\beta}^{\ell}$, we note that this variable satisfies the same recursion as the residual stream variables $\bm{h}^{\ell}$ on the training points described above which gives the identical solution $\bm{h}^{\ell}=\left(\bm{I}-\frac{1}{LP}\bm{X}^{\top}\bm{\Gamma}\bm{X}\right)^{\ell}\bm{y}$. The function $f_{\star}$ on a test point $\bm{x}_{\star}$ takes the form

$$ $\displaystyle f_{\star}=\frac{1}{\sqrt{D}}\bm{\beta}^{L}\cdot\bm{x}_{\star}=\frac{1}{LP}\bm{x}_{\star}^{\top}\bm{\Gamma}\bm{X}\sum_{\ell=0}^{L-1}\bm{h}^{\ell}=\frac{1}{LP}\bm{x}_{\star}^{\top}\bm{\Gamma}\bm{X}\sum_{\ell=0}^{L-1}\left(\bm{I}-\frac{1}{LP}\bm{X}^{\top}\bm{\Gamma}\bm{X}\right)^{\ell}\bm{y}$ (55) $$

We will next express the test loss

##### Test Error Formula for Linear Models

We now desire to compute the expected test loss (averaged over the dataset $\bm{X}$ and noise $\bm{\epsilon}$) for (noisy) linear target function $y(\bm{x})$

$$ $\displaystyle y(\bm{x})=\frac{1}{\sqrt{D}}\bm{\beta}_{\star}\cdot\bm{x}+\sigma\epsilon\ ,\ \left<\bm{x}\bm{x}^{\top}\right>=\bm{\Sigma}\ ,\ \left<\epsilon^{2}\right>=1$ (56) $$

where $\bm{\Sigma}$ is the (population) covariance of the inputs. The expected loss under this assumption is

$$ $\displaystyle\mathcal{L}=\frac{1}{D}\left<\left(\bm{\beta}^{L}-\bm{\beta}_{\star}\right)^{\top}\bm{\Sigma}\left(\bm{\beta}^{L}-\bm{\beta}_{\star}\right)\right>_{\bm{X},\bm{\epsilon}}+\sigma^{2}$ (57) $$

Letting $\bm{\epsilon}\in\mathbb{R}^{P}$ represent the noise on the $P$ labeled training points, the target weights have the form

$$ $\displaystyle\bm{\beta}^{L}$ $\displaystyle=\frac{\sqrt{D}}{LP}\bm{\Gamma}\bm{X}\sum_{\ell=0}^{L}\left(\bm{I}-\frac{1}{LP}\bm{X}^{\top}\bm{\Gamma}\bm{X}\right)^{\ell}\left[\frac{1}{\sqrt{D}}\bm{X}^{\top}\bm{\beta}_{\star}+\sigma\bm{\epsilon}\right]$ $\displaystyle=L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\sum_{\ell=0}^{L-1}\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{\ell}\bm{\beta}_{\star}+\frac{\sigma\sqrt{D}}{LP}\ \bm{\Gamma}\sum_{\ell=0}^{L-1}\left(\bm{I}-L^{-1}\hat{\bm{\Sigma}}\bm{\Gamma}\right)^{\ell}\bm{X}\bm{\epsilon}$ $\displaystyle=\bm{\beta}_{\star}-\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{L}\bm{\beta}_{\star}+\frac{\sigma\sqrt{D}}{LP}\bm{\Gamma}\sum_{\ell=0}^{L-1}\left(\bm{I}-L^{-1}\hat{\bm{\Sigma}}\bm{\Gamma}\right)^{\ell}\bm{X}\bm{\epsilon}$ (58) $$

where we defined $\hat{\bm{\Sigma}}=\frac{1}{P}\bm{X}\bm{X}^{\top}\in\mathbb{R}^{D\times D}$. Then we note that for $\alpha\equiv P/D$ that the reducible loss $\mathcal{L}-\sigma^{2}$ has the form

$$ $\displaystyle\mathcal{L}-\sigma^{2}$ $\displaystyle=\frac{1}{D}\left<\left(\bm{\beta}_{\star}-\bm{\beta}^{L}\right)^{\top}\bm{\Sigma}\left(\bm{\beta}_{\star}-\bm{\beta}^{L}\right)\right>$ $\displaystyle=\frac{1}{D}\bm{\beta}_{\star}^{\top}\left<\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{L\top}\bm{\Sigma}\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{L}\right>\bm{\beta}_{\star}$ $\displaystyle+\frac{\sigma^{2}}{\alpha L^{2}}\text{Tr}\ \bm{\Gamma}^{\top}\bm{\Gamma}\sum_{\ell,\ell^{\prime}=0}^{L-1}\left<\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{\ell}\hat{\bm{\Sigma}}\left[\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{\ell^{\prime}}\right]^{\top}\right>$ (59) $$

In the next sections, we will compute this quantity for various distributions for $\bm{\hat{\Sigma}}$ and $\bm{\beta}_{\star}$.

## Appendix C Isotropic ( ISO ) Setting Theory

In this section, we consider isotropic covariates and tasks so that the data covariance $\bm{\Sigma}$ and the task vector covariance $\bm{\Omega}$ are both identity

$$ $\displaystyle\bm{\Sigma}=\left<\bm{x}\bm{x}^{\top}\right>=\bm{I}\ ,\ \bm{\Omega}=\left<\bm{\beta}_{\star}\bm{\beta}_{\star}^{\top}\right>=\bm{I}.$ (60) $$

where the average for target vectors $\bm{\beta}_{\star}$ is over different contexts. We further operate in the high dimensional asymptotic limit where both context length $P$ and input dimension $D$ diverge with fixed ratio

$$ $\displaystyle P,D\to\infty\ ,\ P/D=\alpha.$ (61) $$

In this setting, the spectrum of the empirical covariance matrix $\hat{\bm{\Sigma}}=\frac{1}{P}\bm{X}\bm{X}^{\top}$ follows the well-known Marchenko-Pastur law, where the eigenvalue density $\rho(\lambda)$ depends explicitly on the aspect ratio $\alpha$ (Potters & Bouchaud, 2020; Advani et al., 2020). We will first describe an asymptotic limit of SGD for the shallow $L=1$ case where the dynamics on the matrix $\bm{\Gamma}$ are linear. We will then pursue a description of gradient flow dynamics in an arbitrary depth $L$ model.

### C.1 Shallow SGD Theory

In this section, we consider the SGD dynamics of the shallow architecture $L=1$. This excercise will establish at what rate our dynamics will converge to gradient flow and how finite batchsize $B$, context length $P$ and number of masked evaluation points $K$ impact the SGD dynamics. In this case, the updates have the form

$$ $\displaystyle\bm{\Gamma}(t+1)=\bm{\Gamma}(t)+\eta\ \bm{G}(t)$ $\displaystyle\bm{G}(t)=\frac{1}{B}\sum_{n}\hat{\bm{\Sigma}}_{\star,n}\left[(\bm{I}-\bm{\Gamma}\bm{\Sigma}_{n})\bm{\beta}_{n}-\sigma P^{-1}\sqrt{D}\bm{\Gamma}\bm{X}_{n}\bm{\epsilon}_{n}\right]\left[\bm{\Sigma}_{n}\bm{\beta}_{n}+\sigma P^{-1}\sqrt{D}\bm{X}_{n}\bm{\epsilon}_{n}\right]^{\top}$ $\displaystyle\hat{\bm{\Sigma}}_{n}=\frac{1}{P}\bm{X}_{n}^{\top}\bm{X}_{n}\ ,\ \hat{\bm{\Sigma}}_{\star,n}=\frac{1}{K}\bm{X}_{\star,n}^{\top}\bm{X}_{\star,n}$ (62) $$

where $\bm{X}\in\mathbb{R}^{P\times D}$ is the set of training points and $\bm{X}_{\star}\in\mathbb{R}^{K\times D}$ represents the set of $K$ evaluation points and $\bm{\epsilon}\in\mathbb{R}^{P}$ represents the label noise on the provided targets $\bm{y}$. We remind the reader that we will operate in the joint scaling limit

$$ $\displaystyle P,B,K,D\to\infty\ ,\ P/D=\alpha\ ,\ K/D=\kappa\ ,\ B/D=\tau$ (63) $$

The population ICL loss has the form

$$ $\displaystyle\mathcal{L}(t)=\frac{1}{D}\left<|\bm{\Gamma}(t)-\bm{I}|^{2}\right>,$ (64) $$

where the average is performed over all possible draws of data and tasks. This function admits the recursion

$$ $\displaystyle\mathcal{L}(t+1)=\mathcal{L}(t)-2\eta\text{tr}[\bm{\Gamma}(t)-\bm{I}]^{\top}\left<\bm{G}(t)\right>+\eta^{2}\text{tr}\left<\bm{G}(t)^{\top}\bm{G}(t)\right>$ (65) $$

We have the following mean for the gradient

$$ $\displaystyle\left<\bm{G}(t)\right>=(\bm{I}-(1+\alpha^{-1}+\sigma^{2}\alpha^{-1})\bm{\Gamma})$ (66) $$

Thus if we only updated using mean gradients the model would relax to a fixed point $\bm{\Gamma}_{\star}=\left(1+\alpha^{-1}+\sigma^{2}\alpha^{-1}\right)^{-1}\bm{I}$. However, we will see momentarily that SGD noise in the gradients impact the $\bm{\Gamma}$ that SGD converges to

$$ $\displaystyle\text{tr}\ \left<\bm{G}(t)^{\top}\bm{G}(t)\right>=\frac{1}{B^{2}}\sum_{nm}\text{tr}\ \left<\bm{u}_{n}\bm{v}_{n}^{\top}\bm{v}_{m}\bm{u}_{m}^{\top}\right>$ (67) $$

where $\bm{u}_{n}=\bm{\hat{\Sigma}}_{\star,n}\left[(\bm{I}-\bm{\Gamma}\bm{\Sigma}_{n})\bm{\beta}_{n}-\sigma p^{-1}\sqrt{d}\bm{\Gamma}\bm{X}_{n}\bm{\epsilon}_{n}\right]$ and $\bm{v}_{n}=\bm{\Sigma}_{n}\bm{\beta}_{n}+\sigma p^{-1}\sqrt{d}\bm{X}_{n}\bm{\epsilon}_{n}$. Doing the averages

$$ $\displaystyle\frac{1}{D}\left<\bm{u}_{n}\cdot\bm{u}_{m}\right>=\delta_{mn}\text{tr}\left[\bm{I}-2\bm{\Gamma}+(1+\alpha^{-1})\bm{\Gamma}^{2}+\sigma^{2}\alpha^{-1}\bm{\Gamma}^{2}\right](1+\kappa^{-1})$ (68) $\displaystyle=\delta_{mn}(1+\alpha^{-1}(1+\sigma^{2}))(1+\kappa^{-1})\text{tr}\left[(\bm{\Gamma}-\bm{\Gamma}_{\star})^{2}+\frac{\alpha^{-1}(1+\sigma^{2})}{(1+\alpha^{-1}(1+\sigma^{2}))^{2}}\bm{I}\right]$ (69) $\displaystyle\frac{1}{D}\left<\bm{v}_{n}\cdot\bm{v}_{m}\right>=\delta_{mn}(1+\alpha^{-1}(1+\sigma^{2}))$ (70) $$

From the above results, we get the following second moment structure

$$ $\displaystyle\text{tr}\left<\bm{G}(t)^{\top}\bm{G}(t)\right>$ $\displaystyle=\text{tr}\left<\bm{G}(t)\right>^{\top}\left<\bm{G}(t)\right>$ $\displaystyle+\frac{1}{\tau}(1+\kappa^{-1})(1+\alpha^{-1}(1+\sigma^{2}))^{2}\left[(\bm{\Gamma}-\bm{\Gamma}_{\star})^{2}+\frac{\alpha^{-1}(1+\sigma^{2})}{(1+\alpha^{-1}(1+\sigma^{2}))^{2}}\bm{I}\right].$ (71) $$

Thus we get a recursion

$$ $\displaystyle\bm{\Gamma}(t+1)=\bm{\Gamma}(t)+\eta(1+\alpha^{-1}(1+\sigma^{2}))\left[\bm{\Gamma}_{\star}-\bm{\Gamma}(t)\right]+\eta\ \bm{\Xi}(t)$ (72) $$

where $\bm{\Xi}\in\mathbb{R}^{D\times D}$ is the noise process with variance given above. We define a quantity $C(t)$ which measures the relaxation of $\bm{\Gamma}$ to $\bm{\Gamma}_{\star}$

$$ $\displaystyle C(t)=\frac{1}{D}\left|\bm{\Gamma}(t)-\bm{\Gamma}_{\star}\right|^{2}\ ,\ \bm{\Gamma}_{\star}=\left[1+\alpha^{-1}(1+\sigma^{2})\right]^{-1}\bm{I}$ (73) $$

This function exhibits the following linear dynamics

$$ $\displaystyle C(t+1)$ $\displaystyle=\left[1-\eta\left(1+\alpha^{-1}(1+\sigma^{2})\right)\right]^{2}C(t)$ $\displaystyle+\frac{\eta^{2}}{\tau}(1+\kappa^{-1})(1+\alpha^{-1}(1+\sigma^{2}))^{2}\left[C(t)+\frac{\alpha^{-1}(1+\sigma^{2})}{(1+\alpha^{-1}(1+\sigma^{2}))^{2}}\right]$ (74) $$

This recursion again takes the form

$$ $\displaystyle C(t+1)=a(\eta,\alpha,\kappa,\tau)C(t)+b(\eta,\alpha,\kappa,\tau)$ $\displaystyle a(\eta,\alpha,\kappa,\tau)=\left[1-\eta\left(1+\alpha^{-1}(1+\sigma^{2})\right)\right]^{2}+\frac{\eta^{2}}{\tau}(1+\kappa^{-1})(1+\alpha^{-1}(1+\sigma^{2}))^{2}$ $\displaystyle b(\eta,\alpha,\kappa,\tau)=\frac{\eta^{2}}{\tau}(1+\kappa^{-1})\alpha^{-1}(1+\sigma^{2})$ (75) $$

Now we need to compute the ICL loss

$$ $\displaystyle\mathcal{L}(t)$ $\displaystyle=\text{tr}[\bm{\Gamma}(t)-\bm{\Gamma}_{\star}+\bm{\Gamma}_{\star}-\bm{I}]^{2}=C(t)+2\text{tr}(\bm{\Gamma}_{\star}-\bm{\Gamma}(t))(\bm{I}-\bm{\Gamma}_{\star})+\text{tr}(\bm{\Gamma}_{\star}-\bm{I})^{2}$ $\displaystyle=C(t)+\frac{2\alpha^{-1}(1+\sigma^{2})}{(1+\alpha^{-1}(1+\sigma^{2}))^{2}}\left[1-\eta(1+\alpha^{-1}(1+\sigma^{2}))\right]^{t}+\frac{(1+\sigma^{2})^{2}}{(\alpha+1+\sigma^{2})^{2}}$ (76) $$

Figure: (a) Vary Batchsize $\tau$
Refer to caption: x19.png

Putting this all together, we summarize our findings below.

We verify the validity of these learning curves for shallow pretraining dynamics in Figure [7](https://arxiv.org/html/2510.01098v1#A3.F7).

We thus see that deviations from gradient flow in the $L=1$ model come in at a scale of $\eta/\tau$. The purpose of this result is to stress that only $\mathcal{O}(D^{2})$ total tokens are required (unlike the prior work of Lu et al. (2025) which required $\mathcal{O}(D^{3})$ tokens) for an ICL learner in this regime (^3^33The gap is caused by suboptimal scaling of the number of evaluation points $K$ per context (usually $K=1$ in prior works). Note realistic LLMs get multiple error signals per context (one per token).). We verify these theoretical learning curves in Figure [7](https://arxiv.org/html/2510.01098v1#A3.F7).

We also indicate the following facts about SGD effects

- •
The number of masked evaluation points per batch $\kappa\equiv K/D$ only alters the SGD noise terms (the terms involving $1/\tau$) and its effect vanishes in the gradient flow limit $\eta/\tau\to 0$.
- •
Unlike offline training, in this online SGD setting the model regularizes the final value of $\gamma(t)=$ based on label noise $\sigma^{2}$ to reduce overfitting.
- •
SGD fluctuations can impact the final loss through $\lim_{t\to\infty}C(t)=\frac{b(\eta,\alpha,\tau,\kappa)}{1-a(\eta,\alpha,\tau,\kappa)}$.

One nice consequence of this result is that the gradient flow limit can be accessed simply by controlling the scale of $\eta/\tau$ in a limit where batch size is linear in the dimension $B=\tau D$.

### C.2 Gradient Flow in Deep Models

From the previous section, we saw that in the absence of SGD noise, gradient dynamics generated an isotropic $\bm{\Gamma}$ matrix. This fact is true for larger depth $L\geq 1$ as well. The gradient flow dynamics generate is $\bm{\Gamma}(t)=\gamma(t)\bm{I}$. To see this, we note that the gradient has the form

$$ $\displaystyle\frac{\partial}{\partial\bm{\Gamma}}\mathcal{L}$ $\displaystyle=-\left<\frac{2}{L}\sum_{\ell=1}^{L}\left[\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{\ell}\right]^{\top}\hat{\bm{\Sigma}}\left[\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{L-1-\ell}\right]^{\top}\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{L}\right>$ $\displaystyle+\frac{2\sigma^{2}}{\alpha L^{2}}\bm{\Gamma}\sum_{\ell,\ell^{\prime}}\left<\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{\ell}\hat{\bm{\Sigma}}\left[\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{\ell^{\prime}}\right]^{\top}\right>$ $\displaystyle+\frac{2\sigma^{2}}{\alpha L^{3}}\bm{\Gamma}^{\top}\bm{\Gamma}\sum_{\ell,\ell^{\prime},k}\left<\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{k}\hat{\bm{\Sigma}}\left(\bm{I}-\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{\ell-1-k}\hat{\bm{\Sigma}}\left[\left(\bm{I}-L^{-1}\bm{\Gamma}\hat{\bm{\Sigma}}\right)^{\ell^{\prime}}\right]^{\top}\right>$ (80) $$

Now, evaluating this at $\bm{\Gamma}=\gamma\bm{I}$ we find

$$ $\displaystyle\frac{\partial}{\partial\bm{\Gamma}}\mathcal{L}|_{\bm{\Gamma}=\gamma\bm{I}}$ $\displaystyle=-2\left<\hat{\bm{\Sigma}}\left(\bm{I}-L^{-1}\gamma\hat{\bm{\Sigma}}\right)^{2L-1}\right>+\frac{2\sigma^{2}\gamma}{\alpha L^{2}}\left<\hat{\bm{\Sigma}}\sum_{\ell\ell^{\prime}}(1-L^{-1}\gamma\hat{\bm{\Sigma}})^{\ell+\ell^{\prime}}\right>$ $\displaystyle+\frac{2\sigma^{2}\gamma^{2}}{\alpha L^{2}}\sum_{\ell,\ell^{\prime}}\left<\hat{\bm{\Sigma}}\left(\bm{I}-L^{-1}\gamma\hat{\bm{\Sigma}}\right)^{\ell+\ell^{\prime}-1}\right>\propto\bm{I}$ (81) $$

where we recognize through rotational invariance, that the above averages are proportional to the identity. Thus if $\bm{\Gamma}$ is currently in an isotropic configuration, it will remain in one throughout gradient flow. Further, we can compute the scalar $\gamma(t)=\text{tr}\bm{\Gamma}(t)$ under gradient flow $\partial_{t}\bm{\Gamma}(t)=-\frac{1}{2}\partial_{\bm{\Gamma}}\mathcal{L}(\bm{\Gamma})$.

$$ $\displaystyle\frac{d}{dt}\gamma(t)=-\frac{\partial}{\partial\gamma}\text{tr}\left[\left(\bm{I}-L^{-1}\gamma\hat{\bm{\Sigma}}\right)^{2L}+\frac{\sigma^{2}}{\alpha}\hat{\bm{\Sigma}}^{-1}\left[\bm{I}-\left(\bm{I}-L^{-1}\gamma\hat{\bm{\Sigma}}\right)^{L}\right]^{2}\right]$ (82) $$

Now, we note that this expression can be rewritten in terms of the Marchenko-Pastur eigenvalue density $\rho(\lambda)=\frac{1}{D}\sum_{k=1}^{D}\left<\delta(\lambda-\lambda_{k})\right>$ as

$$ $\displaystyle\frac{d}{dt}\gamma(t)=-\frac{\partial}{\partial\gamma}\int d\lambda\rho(\lambda)\left[\left(1-L^{-1}\gamma\lambda\right)^{2L}+\frac{\sigma^{2}}{\alpha\lambda}\left[1-\left(1-L^{-1}\gamma\lambda\right)^{L}\right]^{2}\right]$ (83) $$

which shows that we can view the dynamics as a gradient flow on a one dimensional loss landscape.

## Appendix D Fixed and Structured Covariance ( FS ) Setting

In this section we discuss the case where the data covariance is structured but potentially structured. We consider the noise free setting and $P/D\to\infty$ for simplicity where the loss has the form

$$ $\displaystyle\mathcal{L}=\text{tr}\ \bm{\Omega}\left<\left[\left(\bm{I}-L^{-1}\bm{\Gamma}{\bm{\Sigma}}\right)^{L}\right]^{\top}\bm{\Sigma}\left(\bm{I}-L^{-1}\bm{\Gamma}{\bm{\Sigma}}\right)^{L}\right>$ (84) $$

In this case, gradient flow will generate dynamics that cause $\bm{\Gamma}$ to pick up anisotropy from the structure of $\bm{\Omega}$ and $\bm{\Sigma}$. We analyze the case where $\bm{\Omega}$ and $\bm{\Sigma}$ are simultaneously diagonalizable with eigenvalues $\omega_{k}$ and $\lambda_{k}$. In this case, $\bm{\Gamma}$ will share the same eigenbasis. Let the eigenvalues be $\gamma_{k}$

$$ $\displaystyle\frac{\partial}{\partial t}\gamma_{k}(t)=2\omega_{k}\lambda_{k}^{2}\left(1-L^{-1}\lambda_{k}\gamma_{k}(t)\right)^{2L-1}.$ (85) $$

The fixed point of the above dynamics is $\gamma_{k}=L\lambda_{k}^{-1}$. In the large depth limit $L\to\infty$, we can solve the dynamics exactly

$$ $\displaystyle\frac{d}{dt}\gamma_{k}(t)$ $\displaystyle=2\omega_{k}\lambda_{k}^{2}e^{-2\lambda_{k}\gamma_{k}(t)}\implies e^{2\lambda_{k}\gamma_{k}}d\gamma_{k}=\omega_{k}\lambda_{k}^{2}dt$ $\displaystyle\implies 2\lambda_{k}\gamma_{k}(t)$ $\displaystyle=\ln(1+4\omega_{k}\lambda_{k}^{3}t)$ (86) $$

Plugging this solution into the loss function we find

$$ $\displaystyle\mathcal{L}(t)=\frac{1}{D}\sum_{k}\omega_{k}\lambda_{k}e^{-2\gamma_{k}\lambda_{k}}=\frac{1}{D}\sum_{k}\omega_{k}\lambda_{k}[1+\omega_{k}\lambda_{k}^{3}t]^{-1}$ (87) $$

Under powerlaw (source/capacity) assumptions on the structure of the data

$$ $\displaystyle\lambda_{k}\sim k^{-\nu}\ ,\ \omega_{k}\lambda_{k}\sim k^{-\nu\beta-1},$ (88) $$

the loss scales in a powerlaw as

$$ $\displaystyle\mathcal{L}(t)=\sum_{k}k^{-\nu\beta-1}[1+k^{-2\nu-\nu\beta}t]^{-1}\sim t^{-\frac{\beta}{\nu+\nu\beta+1}}$ (89) $$

This powerlaw provides a decent approximation to large but finite depth models $L$ as we show in Figure [3](https://arxiv.org/html/2510.01098v1#S3.F3).

## Appendix E Randomly Rotated and Structured Covariance ( RRS ) Setting

### E.1 Exploiting Symmetry in the Gradient Updates

Utilizing the definition of the RRS setting, we can massage the placement of the orthogonal matrices so that the loss function can be expressed as

$$ $\displaystyle\mathcal{L}=\left<|\bm{\Lambda}^{1/2}\left(\bm{I}-L^{-1}\bm{O}^{\top}\bm{\Gamma}\bm{O}\hat{\bm{\Sigma}}\right)^{L}\bm{\beta}_{\star}|^{2}\right>.$ (90) $$

From this expression, it is immediately clear that this function is rotationally invariant with respect to $\bm{\Gamma}$ since the transformation $\bm{\Gamma}\to\bm{V}\bm{\Gamma}\bm{V}^{\top}$ for orthogonal $\bm{V}$ leaves the Haar average unchanged. The form of the loss gradients also reveals a symmetry

$$ $\displaystyle\frac{\partial}{\partial\bm{\Gamma}}\mathcal{L}=-2L^{-1}\sum_{\ell}\left<\bm{O}\hat{\bm{\Sigma}}\bm{M}^{\ell}\bm{\beta}_{\star}\bm{\beta}_{\star}^{\top}\left(\bm{M}^{L}\right)^{\top}\bm{\Lambda}\left(\bm{M}^{L-1-\ell}\right)^{\top}\bm{O}^{\top}\right>\ ,\ \bm{M}=\bm{I}-\bm{O}^{\top}\bm{\Gamma}\bm{O}\hat{\bm{\Sigma}}$ (91) $$

Note that starting from zero initialization, we have $\bm{M}=0$ which implies isotropic gradients $\frac{\partial}{\partial\bm{\Gamma}}\mathcal{L}|_{\bm{\Gamma}=0}\propto\bm{I}$. To see this note that $\left<\bm{O}\bm{C}\bm{O}^{\top}\right>\propto\bm{I}$ for any matrix $\bm{C}$ that is independent of $\bm{O}$. Further, suppose that we evaluated the loss gradient at any isotropic $\bm{\Gamma}=\gamma\bm{I}$, then $\bm{M}=\bm{I}-\gamma\hat{\bm{\Sigma}}$ and the gradients remain isotropic

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial\bm{\Gamma}}$ $\displaystyle=-2L^{-1}\sum_{\ell}\left<\bm{O}\hat{\bm{\Sigma}}\bm{M}^{\ell}\bm{\beta}_{\star}\bm{\beta}_{\star}^{\top}\left(\bm{M}^{L}\right)^{\top}\bm{\Lambda}\left(\bm{M}^{L-1-\ell}\right)^{\top}\bm{O}^{\top}\right>$ (92) $\displaystyle=-2L^{-1}\sum_{\ell}\text{tr}\left<\hat{\bm{\Sigma}}\bm{M}^{\ell}\bm{\beta}_{\star}\bm{\beta}_{\star}^{\top}\left(\bm{M}^{L}\right)^{\top}\bm{\Lambda}\left(\bm{M}^{L-1-\ell}\right)^{\top}\right>\times\bm{I}$ (93) $$

where we explicitly performed the average over the only remaining factors of $\bm{O}$ since $\bm{M}=\bm{I}-\gamma\hat{\bm{\Sigma}}$ is independent of $\bm{O}$. Further, we can express the dynamics of $\gamma(t)$ as a gradient flow on the reduced one-dimensional loss landscape

$$ $\displaystyle\mathcal{L}(\gamma)=\text{tr}\ \bm{\Omega}\left(\bm{I}-\gamma L^{-1}\hat{\bm{\Sigma}}\right)^{L}\bm{\Lambda}\left(\bm{I}-\gamma L^{-1}\hat{\bm{\Sigma}}\right)^{L}$ (94) $$

## Appendix F Scaling Law Theory

In this section, we consider the scaling law theory with a finite projection matrix $\bm{A}$, following Bordelon et al. (2024a). In this case, our RRS setting requires utilizing the free product result in Appendix [A](https://arxiv.org/html/2510.01098v1#A1).

### F.1 Dynamics at Infinite N , L , P N,L,P

The gradient flow in the limit of $N,L,P\to\infty$ has the form

$$ $\displaystyle\frac{d}{dt}\gamma(t)=-\frac{\partial}{\partial\gamma}\sum_{k}\lambda_{k}\omega_{k}e^{-2\gamma(t)\lambda_{k}}\approx-\frac{\partial}{\partial\gamma}\gamma(t)^{-\beta}=\beta\gamma(t)^{-\beta-1}$ (95) $$

This is a separable ODE with solution (under the assumption that $\gamma(0)=0$) up to constants

$$ $\displaystyle\gamma(t)\sim t^{\frac{1}{\beta+2}}.$ (96) $$

This implies a powerlaw loss scaling with exponent set by the rate of growth of $\gamma(t)$

$$ $\displaystyle\mathcal{L}(t)\sim\gamma(t)^{-\beta}\sim t^{-\frac{\beta}{2+\beta}}.$ (97) $$

We expect these dynamics to hold in the limit of $N,L,P\to\infty$. At finite $N,L,P$ the model will saturate at some finite maximal value of $\gamma$.

### F.2 Depth Scaling Law at Infinite Time, Width, and Context Length

We now explore the scaling law in depth $L$. In this case, we can consider taking all other resources to infinity $t,N,P\to\infty$. The value of $\gamma$ will stabilize, resulting in a final loss

$$ $\displaystyle\mathcal{L}(\gamma)=\sum_{k}\lambda_{k}\omega_{k}\left(1-\gamma L^{-1}\lambda_{k}\right)^{2L}$ (98) $$

While $\gamma(t)$ diverges as $t^{\frac{1}{2+\beta}}$ in the infinite depth limit, we see that at large but finite depth, there is a maximal $\gamma$ which enables stability along the top eigendirection. Specifically we need $\gamma<\frac{2L}{\lambda_{1}}$ and approximate the optimal value s

$$ $\displaystyle\gamma_{\star}\approx L/\lambda_{1}.$ (99) $$

where $\lambda_{1}$ is the top (maximal eigenvalue). At $\gamma_{\star}$, the error takes the form

$$ $\displaystyle\mathcal{L}\approx\sum_{k}\lambda_{k}(\beta^{\star}_{k})^{2}\left(1-\lambda_{k}/\lambda_{1}\right)^{L}\approx\int dkk^{-\nu\beta-1}\exp(-Lk^{-\nu})\approx L^{-\beta}.$ (100) $$

which matches the scaling of $L$ steps of gradient descent on a problem with source exponent $\beta$ Bordelon & Pehlevan (2021).

### F.3 Mapping the Loss Function to a Two Point Deterministic Equivalent

We utilize dynamical mean field theory (DMFT) result of Appendix [A](https://arxiv.org/html/2510.01098v1#A1) to compute the effective loss as a function of $N,L$ and $\gamma$. We start with a representation of the relevant polynomial in $\bm{M}$ with the Cauchy integral formula

$$ $\displaystyle\left(\bm{I}-\gamma L^{-1}\bm{M}\right)^{L}=\int_{\mathcal{C}}\frac{d\omega}{2\pi}\left[i\omega+\bm{M}\right]^{-1}\left(1+\gamma L^{-1}i\omega\right)^{L},$ (101) $$

where the contour $\mathcal{C}$ encloses the positive imaginary axis in complex plane (Bender & Orszag, 2013). Next, we define $\bm{v}(\gamma)=\left(\bm{I}-\gamma L^{-1}\bm{M}\right)^{L}\bm{\beta}_{\star}$

$$ $\displaystyle\mathcal{L}=\frac{1}{D}\left<\bm{v}(\gamma)^{\top}\bm{\Lambda}\bm{v}(\gamma)\right>=\text{tr}\ \bm{\Omega}\left[\left(\bm{I}-\gamma L^{-1}\bm{M}\right)^{L}\right]^{\top}\bm{\Lambda}\left(\bm{I}-\gamma L^{-1}\bm{M}\right)^{L}$ (102) $$

The loss can thus be expressed as a double integral involving the resolvents

$$ $\displaystyle\mathcal{L}=\int_{\mathcal{C}\times\mathcal{C}}\frac{d\omega d\omega^{\prime}}{(2\pi)^{2}}(1+\gamma L^{-1}i\omega)^{L}(1+\gamma L^{-1}i\omega^{\prime})^{L}\ \ \text{tr}\ \bm{\Omega}\left[\left[i\omega+\bm{M}\right]^{-1}\right]^{\top}\bm{\Lambda}\left[i\omega^{\prime}+\bm{M}\right]^{-1}$ (103) $$

The main result needed is the following deterministic equivalent from Appendix [A](https://arxiv.org/html/2510.01098v1#A1).

$$ $\displaystyle\text{tr}\ \bm{\Omega}\ \left[\left[i\omega+\bm{M}\right]^{-1}\right]^{\top}\bm{\Lambda}\left[i\omega^{\prime}+\bm{M}\right]^{-1}$ $\displaystyle\simeq\text{tr}\left(i\omega+\Psi_{v\chi}(\omega)\bm{A}\right)^{-1}\left[\bm{\Omega}-\bm{I}\Psi_{\chi\chi}(\omega,\omega^{\prime})\right]\left(i\omega^{\prime}+\Psi_{v\chi}(\omega^{\prime})\bm{A}\right)^{-1}\bm{\Lambda}$ (104) $$

where $\Psi_{v\chi}(\omega)=i\omega/i\omega_{(\bm{A}^{\top}\bm{A})^{2}}$ can be determined from the equation

$$ $\displaystyle i\omega_{(\bm{A}^{\top}\bm{A})^{2}}(\tau)i\omega_{\hat{\bm{\Sigma}}}(\tau)=\frac{1-\tau}{\tau}(i\omega)$ (105) $$

Once $\tau$ is identified as a function of $\omega$, we can compute $\Psi_{v\chi}$ and $\Psi_{\chi\chi}$ from the formulae in [A](https://arxiv.org/html/2510.01098v1#A1).

### F.4 Orthogonal projection and A is Structured Wishart

In this section, we describe what these relations imply for a product matrix that arises in our RRS setting, which is a free product of the width projection matrix $(\bm{A}^{\top}\bm{A})^{2}$ where $\bm{A}\in\mathbb{R}^{N\times D}$ has rank at most $N$ and the empirical covariance for data in the context $\hat{\bm{\Sigma}}=\frac{1}{P}\bm{X}\bm{X}^{\top}\in\mathbb{R}^{D\times D}$,

$$ $\displaystyle\bm{M}=\bm{O}\ \left(\bm{A}^{\top}\bm{A}\right)^{2}\ \bm{O}^{\top}\ \hat{\bm{\Sigma}},$ (106) $$

where $\bm{O}$ is a random orthogonal matrix sampled from the Haar measure. The data $\bm{X}\in\mathbb{R}^{P\times D}$ are comprised of random iid vectors with covariance $\bm{\Lambda}$, which we take to be diagonal without loss of generality. Further, we let $\bm{A}^{\top}\bm{A}$ be a rank $N$ projection matrix

$$ $\displaystyle\bm{A}^{\top}\bm{A}=\sum_{k=1}^{N}\bm{e}_{k}\bm{e}_{k}^{\top}\in\mathbb{R}^{D\times D}$ (107) $$

where $\bm{e}_{k}\in\mathbb{R}^{D}$ are Cartesian (one-hot) unit vectors. To utilize the result in Appendix [A](https://arxiv.org/html/2510.01098v1#A1), we first need to compute the necessary resolvents for these two matrices. Let $\alpha_{N}\equiv N/D$ be the ratio between width $N$ and input dimension $D$ and $\alpha_{P}=\frac{P}{D}$ represent the ratio of context length to input dimension, then

$$ $\displaystyle\tau=\text{tr}(\bm{A}^{\top}\bm{A})^{2}\left(i\omega_{(\bm{A}^{\top}\bm{A})^{2}}+(\bm{A}^{\top}\bm{A})^{2}\right)^{-1}=\frac{\alpha_{N}}{i\omega_{(\bm{A}^{\top}\bm{A})^{2}}+1}\implies i\omega_{(\bm{A}^{\top}\bm{A})^{2}}=-1+\frac{\alpha_{N}}{\tau}$ (108) $$

The other matrix can be easily worked out for a structured Wishart with $P/D=\alpha_{P}$ following the techniques of Bordelon et al. (2024a) which reveals that

$$ $\displaystyle\tau=\text{tr}\ \hat{\bm{\Sigma}}\left(i\omega_{\hat{\bm{\Sigma}}}+\hat{\bm{\Sigma}}\right)^{-1}=\text{tr}\bm{\Lambda}\left(i\omega_{\hat{\bm{\Sigma}}}(1-\alpha_{P}^{-1}\tau)^{-1}+\bm{\Lambda}\right)^{-1}.$ (109) $$

Using our result for products we find

$$ $\displaystyle\left(-1+\frac{\alpha_{N}}{\tau}\right)\left(i\omega_{\hat{\bm{\Sigma}}}(\tau)\right)=\left(-1+\frac{1}{\tau}\right)i\omega$ $\displaystyle\tau=\text{tr}\bm{\Lambda}\left(i\omega_{\hat{\bm{\Sigma}}}(1-\alpha_{P}^{-1}\tau)^{-1}+\bm{\Lambda}\right)^{-1}$ $\displaystyle=\text{tr}\bm{\Lambda}\left(i\omega\ (-1+\tau^{-1})(-1+\alpha_{N}/\tau)^{-1}(1-\tau/\alpha_{P})^{-1}+\bm{\Lambda}\right)^{-1}.$ (110) $$

##### Width Bottleneck

Now we obtain the scaling of the loss with width $N$ in the regime where $t,P,L,D\to\infty$. To do, we examining the structure of the equations at $\alpha_{P}\to\infty$ and $i\omega\to 0$, which correspond to taking gradient flow time $\gamma\to\infty$ (which is the correct solution for noise free problems at infinite depth). In this limit $\tau=\alpha_{N}$ so we have that

$$ $\displaystyle\tau=\frac{N}{D}=\text{tr}\bm{\Lambda}\left(i\omega_{\hat{\bm{\Sigma}}}+\bm{\Lambda}\right)^{-1}=\frac{1}{D}\sum_{k=1}^{D}\frac{\lambda_{k}}{i\omega_{\hat{\bm{\Sigma}}}+\lambda_{k}}.$ (111) $$

For powerlaw features $\lambda_{k}\sim k^{-\nu}$ we find an approximate solution for $i\omega_{\hat{\bm{\Sigma}}}$

$$ $\displaystyle\lambda_{k}\sim k^{-\nu}\implies i\omega_{\hat{\bm{\Sigma}}}\approx N^{-\nu}.$ (112) $$

Thus the loss at width $N$ and context length $P$ can be approximated as

$$ $\displaystyle\lim_{t,L,P\to\infty}\mathcal{L}(t,N,P,L)=\sum_{k}\frac{(i\omega_{\hat{\bm{\Sigma}}})^{2}\lambda_{k}(\beta^{\star}_{k})^{2}}{(i\omega_{\hat{\bm{\Sigma}}}+\lambda_{k})^{2}}\sim\sum_{k}\frac{k^{-\nu\beta-1}}{(1+k^{-\nu}N^{\nu})}\approx N^{-\nu\beta}$ (113) $$

##### Context Bottleneck

Following the same logic, we can investigate the regime where performance is limited by context $P$. To access an approximation for this regime, we take $\alpha_{N}\to 1$ and $i\omega\to 0$. In this case, we have $i\omega=i\omega_{\hat{\bm{\Sigma}}}$. Since $1-\alpha_{P}/\tau\to 0$ as $i\omega\to 0$, we introduce the leading order behavior for small $i\omega$

$$ $\displaystyle 1-\frac{\alpha_{P}}{\tau}\sim r^{-1}i\omega\ ,\ i\omega\to 0$ (114) $$

where $r$ is a (currently) unknown value. Plugging this into the self consistency equation, we identify the value of $r$

$$ $\displaystyle P=\sum_{k}\frac{\lambda_{k}}{r+\lambda_{k}}\approx r^{-\frac{1}{\nu}}$ (115) $$

which gives the scaling $r\approx P^{-\nu}$. We can now use the final value of the resolvent $\left(i\omega+\bm{\Lambda}\right)^{-1}$

$$ $\displaystyle\mathcal{L}=\sum_{k}\frac{r^{2}\lambda_{k}(\beta^{\star}_{k})^{2}}{(r+\lambda_{k})^{2}}\approx\sum_{k}\frac{k^{-\nu\beta-1}}{(1+k^{-\nu}P^{\nu})^{2}}\sim P^{-\nu\beta}.$ (116) $$

##### Rank Deficiency/Null Space Interpretation

These last two scaling laws express the simple fact that rank $N$ or rank $P$ matrices only allow the top $N$ or $P$ eigendirections to be learned (Bordelon et al., 2024a).

## Appendix G Enhancing Realism of the Model

### G.1 Different Parameters for Each Layer

In this section, we analyze the case where the model is allowed $L$ distinct attention layers along the residual stream $\{\bm{\Gamma}^{\ell}\}_{\ell=1}^{L}$ instead of a single attention layer matrix $\bm{\Gamma}$ which is applied $L$ times recurrently. We focus our attention on the RRS setting at large width and large context length where the model still exhibits nontrivial depth and pretraining time scaling laws. We first consider the noise free setting $\sigma^{2}=0$ before considering the role of target noise.

#### G.1.1 Noise Free Loss Function at Large Width and Context Length

First, we note that by the rotational invariance, each of the matrices is isotropic $\bm{\Gamma}^{\ell}=\gamma^{\ell}\bm{I}$ so that

$$ $\displaystyle\mathcal{L}(\{\gamma^{\ell}\})=\left<\left|\bm{\Lambda}^{1/2}\prod_{\ell=1}^{L}\left[\bm{I}-L^{-1}\gamma^{\ell}\bm{M}\right]\bar{\bm{\beta}}\right|^{2}\right>\ ,\ \bm{M}=\bm{O}\left(\bm{A}^{\top}\bm{A}\right)^{2}\bm{O}^{\top}\hat{\bm{\Sigma}}$ (117) $$

We note that, due to commutativity of the matrix products $\prod_{\ell=1}^{L}\left[\bm{I}-L^{-1}\gamma^{\ell}\bm{M}\right]$, which means that this loss function is permutation symmetric in the $\{\gamma^{\ell}\}$ variables. We therefore expect permutation symmetric dynamics and a solution where $\gamma^{\ell}=\gamma$. Indeed, analyzing the gradient flow from $\gamma^{\ell}=0$ for all $\ell$ leads to a balanced solution where all of these are equal since

$$ $\displaystyle\frac{\partial}{\partial\gamma^{\ell}}\mathcal{L}(\{\gamma^{\ell}\})|_{\gamma^{k}=\gamma}=\frac{\partial}{\partial\gamma^{\ell^{\prime}}}\mathcal{L}(\{\gamma^{\ell}\})|_{\gamma^{k}=\gamma}\ ,\ \forall\ell,\ell^{\prime}\in\{1,...,L\}.$ (118) $$

Thus gradient flow will maintain a balance in these parameters. Further, this flow will exactly match the dynamics of the recurrent model if the learning rate is upscaled by $L$. We plot these dynamics and show balancing of the $\gamma^{\ell}(t)$ variables in Figure [8](https://arxiv.org/html/2510.01098v1#A7.F8).

Figure: Figure 8: Dynamics for layer-decoupled reduced $\Gamma$ model. (a) The loss dynamics under powerlaw RRS covariates exhibit the same powerlaws. (b) The different scale factors remain balanced throughout gradient flow $\gamma^{k}(t)=\gamma(t)$ for all $k\in\{1,...,L\}$ due to permutation symmetry.
Refer to caption: x22.png

#### G.1.2 Noisy Target Function

The predictor in the case of a noisy target function with decoupled layers can be expressed in terms of the evolution of the weight discrepancy $\bm{v}^{\ell}\equiv\bar{\bm{\beta}}-\hat{\bm{\beta}}^{\ell}$ which has the recursion

$$ $\displaystyle\bm{v}^{\ell}=\bm{v}^{\ell-1}-\gamma^{\ell}L^{-1}\bm{M}\bm{v}^{\ell-1}-\frac{\sigma\sqrt{D}}{P}\gamma^{\ell}L^{-1}\bm{X}\bm{\epsilon}\ ,\ \ell\in\{1,...,L\}$ (119) $$

Defining the matrices $\bm{A}^{\ell}=\left[\bm{I}-\gamma^{\ell}L^{-1}\bm{M}\right]$ and $\bm{b}=\frac{\sigma\sqrt{D}}{LP}\bm{X}\bm{\epsilon}$ so that $\bm{v}^{\ell}=\bm{A}^{\ell}\bm{v}^{\ell-1}+\gamma^{\ell}\bm{b}$

$$ $\displaystyle\bm{v}^{L}=\prod_{\ell=1}^{L}\bm{A}^{\ell}\bar{\bm{\beta}}+\sum_{\ell=1}^{L}\gamma^{\ell}\left(\prod_{k=\ell+1}^{L}\bm{A}^{k}\right)\bm{b}$ (120) $$

Since the loss function takes the form $\mathcal{L}=\frac{1}{D}\left<\bm{v}^{L}\cdot\bm{\Lambda}\bm{v}^{L}\right>$ we have

$$ $\displaystyle\mathcal{L}=$ $\displaystyle\frac{1}{D}\left<\bar{\bm{\beta}}^{\top}\left(\prod_{\ell}\bm{A}^{\ell}\right)^{\top}\bm{\Lambda}\left(\prod_{\ell}\bm{A}^{\ell}\right)\bar{\bm{\beta}}\right>$ $\displaystyle+\frac{1}{D}\sum_{\ell\ell^{\prime}}\gamma^{\ell}\gamma^{\ell^{\prime}}\left<\bm{b}^{\top}\left(\prod_{k=\ell+1}^{L}\bm{A}^{k}\right)\bm{\Lambda}\left(\prod_{k=\ell^{\prime}+1}^{L}\bm{A}^{k}\right)\bm{b}\right>$ (121) $$

Using the fact that $\left<\bm{b}\bm{b}^{\top}\right>=\sigma^{2}L^{-2}\alpha^{-1}\bm{\Lambda}$ we have the following loss function

$$ $\displaystyle\mathcal{L}=$ $\displaystyle\frac{1}{D}\left<\left|\bm{\Lambda}^{1/2}\prod_{\ell=1}^{L}\bm{A}^{\ell}\bar{\bm{\beta}}\right|^{2}\right>+\frac{\sigma^{2}}{L^{2}\alpha}\sum_{\ell\ell^{\prime}}\gamma^{\ell}\gamma^{\ell^{\prime}}\text{tr}\bm{\Lambda}\left(\prod_{k=\ell+1}^{L}\bm{A}^{k}\right)\bm{\Lambda}\left(\prod_{k=\ell^{\prime}+1}^{L}\bm{A}^{k}\right).$ (122) $$

### G.2 Decoupling Weights in Attention Layers

We note that for the ISO and RRS settings, we can exploit similar symmetry arguments as in [C.2](https://arxiv.org/html/2510.01098v1#A3.SS2) and [E](https://arxiv.org/html/2510.01098v1#A5) to argue that the gradient updates for $\bm{W}_{x},\bm{W}_{k},\bm{W}_{q}$ are isotropic and that $\bm{W}_{v}$ gets an update in the $\bm{w}_{o}\bm{w}_{y}^{\top}$ direction. The gradient flow can be expressed in terms of the original gradients on $\bm{\Gamma}$

$$ $\displaystyle\partial_{t}\bm{W}_{i}=-\frac{\partial\mathcal{L}}{\partial\bm{\Gamma}}\cdot\frac{\partial\bm{\Gamma}}{\partial\bm{W}_{i}}\ ,\ i\in\{x,k,q,v,o,y\}$ (123) $$

As before, under the assumption of small initial conditions, we can reduce the loss to a collection of ODEs on scalars representing the scale of each weight matrix

$$ $\displaystyle\mathcal{L}(w_{o},w_{v},w_{q},w_{k},w_{x},w_{y})$ $\displaystyle=\text{tr}\ \bm{\Omega}\bm{\Lambda}\left[\bm{I}-w_{o}w_{y}\left(\bm{I}-\left[1-L^{-1}w_{v}w_{k}w_{q}(w_{x})^{2}\bm{\Lambda}\right]^{L}\right)\right]^{2}$ (124) $$

Gradient flow dynamics on this loss function can reproduce the dynamics of the decoupled self attention model. Under the further assumption that $w_{o}=w_{y}=1$, we can simplify the loss further to

$$ $\displaystyle\mathcal{L}(w_{v},w_{q},w_{k},w_{x})$ $\displaystyle=\text{tr}\ \bm{\Omega}\bm{\Lambda}\left[\bm{I}-L^{-1}w_{v}w_{k}w_{q}(w_{x})^{2}\bm{\Lambda}\right]^{2L}$ (125) $$

We note that $w_{x}$ will be updated twice as quickly as the other weights under gradient flow. To achieve balance, we can initialize it to have $w_{x}(0)=\sqrt{2}w_{k}(0)$ and have $w_{q}(0)=w_{k}(0)=w_{v}(0)=\sigma$. In this case, the loss can be further reduced to a function of a single variable

$$ $\displaystyle\mathcal{L}(w_{v},w_{q},w_{k},w_{x})$ $\displaystyle=\text{tr}\ \bm{\Omega}\bm{\Lambda}\left[\bm{I}-L^{-1}w(t)^{5}\bm{\Lambda}\right]^{2L}$ (126) $$

Under source and capacity assumptions, this will generate the following dynamics at large depth $L$

$$ $\displaystyle w(t)=t^{\frac{5}{5\beta+2}}\implies\mathcal{L}\sim t^{-\frac{5\beta}{5\beta+2}}.$ (127) $$

Figure: Figure 9: Softmax attention model with MLP Blocks on the residual stream also benefit from increasing the depth.
Refer to caption: x24.png

### G.3 Softmax Attention

We also provide numerical experiments with non-recurrent softmax attention, decoupled layers, and Adam. In this context we use CompleteP scaling for the learning rate $\Theta_{L}(1)$ (Dey et al., 2025). The results are provided in Figure [6](https://arxiv.org/html/2510.01098v1#S5.F6) (c).

Supplementing this figure is an equivalent figure for a non-recurrent architecture with alternating softmax attention and MLP with GELU activation, given by Figure [9](https://arxiv.org/html/2510.01098v1#A7.F9).

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 6
presentation: 4
questions: 1. As I'm mostly concerned about the role of the depth, which plays an important role in the novelty of this work, can the authors justify the validity of assuming equal weight matrices across layers and the corresponding generality?

2. Furthermore, can the authors discuss the difficulty brought by varying weight matrices w.r.t layers and how the current framework can still be applied in that case?

3. Is taking the joint limit of $P, K, B, D, N$ necessary for the results presented in this paper?
soundness: 3
strengths: 1. This paper is well written and well organized. The presentation is very clear and the flow of this paper is consistent, which I think can allow the readers to easily appreciate the core contributions of this work (the summarized theoretical results along with immediate numerical experimental results). 

   The proof is also easy to follow (though I did not check all the details): first transforming the gradient learning dynamics to an equivalent linear model, which has a simpler dynamics, then diving to this new dynamics under different conditions of the covariance. Although this general idea is not new, incorporating the depth to the analysis is novel.

2. The derived theoretical results indeed consider attention in the multi-layer case, which I believe is an improvement over prior works. The authors demonstrate when and why the depth can be necessary. The summarized results are indeed novel and interesting.
summary: This paper studies deep linear self-attention by analyzing the corresponding solvable model. The authors systematically investigate how depth of the model, width, context length, and training steps affect the solution of the model. Specifically, the authors focus on three distinct types of data (termed as ISO, FS, and RRS), where the authors reveal that depth is unnecessary for long contexts on ISO and FS, along with a series of other results that characterize the gradient flow dynamics. Furthermore, the authors derive a separable scaling law for the RSS setting.
weaknesses: While the results are interesting by considering the depth, I think the setting is still not significantly novel compared to prior works. In particular:

 - While equation (3) indicates a dependence on the layer depth $l$, the induced parameter $\Gamma$ in fact does not depend on $l$, because the matrices $W_i$'s are treated equally for different layers given one specific $i$. Instead, this dependence on $l$ is replaced by a simple summation over $l$ in $\Gamma$. As a result, the corresponding analysis in fact does not provide significantly novel analysis compared to prior works in this line of research, i.e., studying the gradient flow dynamics of the induced parameter $\Gamma$ (which is still a linear regression) as a proxy of the true weight matrices of the attention model. It remains unclear whether doing so can really capture the effects of depth.

- The RSS is positioned as the most general case, but the randomly rotated and structured across contexts covariance still cannot effectively capture  the essence of task diversity. In addition, the theoretical framework is built on taking a very specific join limit where $P, K, B, D, N \to \infty$ with fixed ratios. While convenient for analytical tractability, this obscures efects that are relevant at finite scales. 

- Due to the aforementioned limitations, the generality of the derived results remains unclear.


Minor: As this paper considers solvable models and scaling laws of attention, I think the related work [1], which also studies a solvable model for attention and its scaling laws, could be discussed a bit in the related work.

[1]. Lyu et al.  A Solvable Attention for Neural Scaling Laws. ICLR 2025.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. Your theory and experiments focus on linear regression tasks with (mostly) linear attention, Gaussian feature distributions, and controlled covariance structure. How confident should we be that the same depth–width scaling conclusions hold for nonlinear Transformers trained on natural language, vision, or multimodal data? 

2. You mostly analyze single-head linear attention plus residual depth. How do you expect multi-head structure and MLP blocks (i.e., actual transformer blocks) to affect the depth vs. width story? For example, could significantly more attention heads (possibly also scales with D) also benefit the learning process?
soundness: 4
strengths: 1. I appreciate that the paper gives a concrete, theoretically grounded answer to a question that is widely discussed in practice: how should depth vs. width vs. context length scale for in-context learning? Instead of treating “bigger model = better,” it isolates when depth specifically matters, and ties that to properties of the task distribution (shared vs. varying covariance across contexts). Specifically, the model decouples network width $N$ from the problem dimension $D$ when studying the scaling law, and the use of loop transformer ensures that the total number of parameters does not go up with the number of computes---which I believe provides a decoupled and generic test bench that adds to similar work studying scaling with linear model.

2. A demonstration that the usefulness of depth is task-distribution dependent: If all tasks share the same covariance, depth is asymptotically unnecessary (long enough context suffices, as the model weights can "encode" the covariate information). If covariances vary across contexts, depth is fundamentally valuable, even with infinite context (reflecting the philosophy of test-time compute). 

3. A good match between theories and experiments.
summary: The authors tackle the problem of how should we allocate depth, width, context length, and training compute when scaling transformers for in-context learning (ICL) on regression tasks. The paper analyzes a deep loop linear attention transformer trained (via SGD) to do linear regression in context, without finetuning at test time. It studies three task regimes: (i) isotropic data, (ii) fixed but structured covariance, and (iii) randomly rotated structured covariance (task distribution shifts every context), and derives dynamical equations and asymptotic scaling laws linking pretraining time, width, depth, and context length. The contribution lies in a unified, theoretically grounded scaling law that distinguishes the roles of depth, width, context length, and training time in ICL.
weaknesses: 1. The entire analysis is built around linear regression tasks solved via in-context learning with (mostly) loop linear attention. I didn't find much discussion surrounding the use of loop attention block. One benefit I could imagine is the decoupling between total model weights and the depth. However, the use of loop could potentially restrict the model's expressiveness, where model could possibly implemented higher-order optimization algorithm, e.g., Newton's step rather than gradient descent [1], and the model might demonstrate different scaling behavior, especially in depth. 

2. The key conceptual result is that depth becomes essential when task covariances vary across contexts (“randomly rotated structured,” RRS). But the diversity they model is very specific: random orthogonal rotations of a shared spectrum. That’s mathematically nice but arguably still a stylized shift. Real heterogeneity looks more like mixture of domains, sparsity structure, nonstationary label noise, hierarchical latent factors, etc. It’s not obvious that random rotations is the right stand-in for natural distribution shift.

3. The paper argues its results are relevant to large-scale LLM design, but the experiments cap out at synthetic regression and relatively small controlled transformers. There’s no ablation on modern-scale architectures (residual blocks with MLPs, nonlinear attention heads, long-context finetuning) to show even qualitative alignment. So the significance for frontier models is still somewhat speculative.

[1] Giannou, A., Yang, L., Wang, T., Papailiopoulos, D., & Lee, J. D. (2024). How Well Can Transformers Emulate In-context Newton’s Method? arXiv preprint arXiv:2403.03183.

---

confidence: 3
contribution: 2
initial_rating: 2
presentation: 1
questions: The problem considered in this paper is interesting and potentially important. However, recurring issues with presentation and clarity make it difficult to fully assess the contributions. Improving the clarity would make the results more accessible for proper evaluation.
soundness: 2
strengths: The paper studies an interesting and relevant problem of training a deep linear self-attention model on in-context linear regression tasks, considering different in-context task structures. Extensive prior work has examined one-layer linear self-attention trained on in-context linear regression with isotropic task vectors. Analyzing deep models and non-isotropic task vectors represents meaningful progress.
summary: This paper studies deep linear self-attention trained on the in-context linear regression task, characterizing how performance depends on width, depth, number of training steps, batch size and data per context.
weaknesses: - My main concern is the looped transformer assumption, $W_i^l=W_i^{l'}, i\in \{k,q,v\}$, which means that all layers have the same weights. The expressivity and optimization properties of a deep self-attention model can differ significantly with or without this assumption. Why would the scaling limits derived under this constrain reflect those of a real multi-layer attention model, which typically learns different weights across layers?

- In the reduced model in Equation (4), all weight matrices within a layer appear to be merged into a single trainable matrix $\Gamma$, effectively making the self-attention layer "shallow". Since gradient descent dynamics and the loss landscape are sensitive to such reparameterization, the true loss landscape probably differs from those shown in Figure 2 and Figure 4b. If this is indeed the case, it would be helpful to explicitly highlight this distinction.

- I suggest that the authors perform another round of proofreading and polishing. There are presentation issues that make the paper unnecessarily difficult to read smoothly. I list some below.

  It appears that manual vertical spacing commands have been used in several places of the paper. The formatting on page 14 seems irregular.

  The clarity of Equation (3) could be improved by specifying the dimensionality of weight matrices. 

  The authors use inner-product and transpose notations interchangeably. The expectation notation $\mathbb E(\cdot)$ and the angle brackets notation $\langle \cdot \rangle$ are also used interchangeably. Adopting consistent notations throughout would enhance readability.

  The symbols $\boldsymbol X, \boldsymbol y$ in Equation (4) seem to be undefined.

  The symbol $i$ is used inconsistently: sometimes as an index and sometimes as the imaginary number. In particular, it is undefined in Result 7. Clarifying its meaning in each context would avoid confusion.

  In Equations (13), (15), it appears that a scalar is being added to a matrix, e.g., $(1-L^{-1}\gamma\Lambda), (i\omega+\Psi\Lambda)$. Please check these terms for dimensional consistency. Additionally, the trace operator $\text{tr}$ is used without parentheses, which could be ambiguous.

---

confidence: 4
contribution: 4
initial_rating: 6
presentation: 2
questions: - Could the authors elaborate on the necessity of employing DMFT to derive the closed-form loss expression in Result 7? Is this approach required because directly analyzing the gradient flow dynamics in (13) is intractable, perhaps due to the lack of a known closed-form solution for the ODE governing $\gamma(t)$ in the randomly rotated setting with finite width N?
- The current analysis focuses heavily on the proportional asymptotic regime. Are the techniques employed amenable to deriving non-asymptotic results that might provide insights into the behavior of the system with sizes?
- The derivation of the two-point deterministic equivalent using DMFT in Appendix A introduces notation that appears distinct from the parameters used in the main text to describe the Transformer model and its dynamics and it's hard to directly match (22) with the main result. Could the authors provide a clearer mapping between the DMFT variables/order parameters and the model parameters/dynamics described earlier in the paper?
soundness: 4
strengths: - The paper provides a comprehensive theoretical analysis of multi-layer linear self-attention models for in-context linear regression across three distinct covariate settings (ISO, FS, RRS).
- This paper rigorously characterizes the training dynamics using gradient flow analysis, revealing how the model learns under different data structures and providing an interpretation of the learned estimator as implementing multi-step gradient descent with optimal step sizes.
- The derivation of a Chinchilla-like neural scaling law incorporating time, width, depth, and context length for the RRS setting in the context of linear regression with power-law features is a significant theoretical contribution.
- The application of Dynamical Mean Field Theory (DMFT) to derive a two-point deterministic equivalent for the loss landscape under random rotations represents a novel technical approach for analyzing complex learning dynamics in this asymptotic regime.
summary: This paper theoretically investigates in-context learning for linear regression using deep linear self-attention models, analyzing performance based on width (N), depth (L), context length (P), pretraining time (t), and data structure. Examining isotropic (ISO), fixed structured (FS), and randomly rotated structured (RRS) covariance settings, the authors find that depth primarily benefits ICL in the ISO and FS settings only when context length is limited; for infinite context length in these settings, increasing depth beyond L=1 offers no advantage. However, in the more complex RRS setting where covariances vary, depth significantly improves performance even at infinite context length. For this RRS case, the paper derives a Chinchilla-like scaling law and predicts a compute-optimal shape scaling, linking optimal architecture to task data properties.
weaknesses: -The presentation of detailed proofs and derivations within the appendix could be improved for clarity and accessibility, making it challenging to fully verify the technical steps.


Decision: Accept (Poster)