PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Blake Bordelon, Mary I. Letey, Cengiz Pehlevan

arXiv:2510.01098v1stat.MLcond-mat.dis-nncs.LG

TL;DR

This paper asks how transformer depth and width shape in-context linear regression, extending scaling-law analysis beyond prior work’s emphasis on width and pretraining resources. It analyzes deep linear attention and finds that whether depth helps at long contexts depends on covariance structure: fixed covariances yield a brittle whitening solution, whereas randomly rotated covariances support a general-purpose in-context gradient-descent algorithm.

  • Problem

    Prior theoretical work on neural scaling laws has emphasized width, pretraining data, or pretraining time, while this paper examines scaling resources for in-context linear regression.

    Source: §1 · p. 1

  • Method

    The paper develops a solvable deep linear attention model and analyzes in-context learning under three data covariance structures.

    Source: §1 · p. 2, §2.1 · p. 3

  • Results

    For fixed covariances, the model learns a whitening preconditioner that is brittle to covariance changes; for randomly rotated covariances, it learns a general-purpose in-context gradient-descent algorithm with width and depth bottlenecks.

    Source: §3.2 · p. 6, §6 · p. 10

  • Takeaways & Limitations

    The results show that pretraining statistics determine the learned solution and whether depth is necessary at large context lengths.

    Source: §6 · pp. 9–10

    In fixed-covariance pretraining, the shallow zero-loss solution at long contexts is specialized to the pretraining covariance and is brittle to changes in it.

    Source: §3.2 · p. 6

Abstract

from arXiv · show

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

The paper asks how transformer width and depth should scale under a compute budget, and whether their effects depend on task statistics rather than only total parameter count. It develops a solvable model of in-context linear regression and analyzes scaling behavior across covariance settings.

1 INTRODUCTION

  • Existing theory does not establish compute-optimal transformer shapes or distinguish width and depth gains at fixed compute, motivating the paper’s scaling-law analysis.The introduction frames this as the question of whether scaling laws depend only on total parameters or treat width and depth differently.Source: §1 · p. 1, §1.2 · p. 2
  • The paper studies how depth, width, context length, and pretraining time affect in-context linear regression, with learned solutions depending on pretraining task statistics.It contrasts in-context learning, which conditions on past inputs without parameter updates, with in-weight learning.Source: §1 · p. 1

1.1 CONTRIBUTIONS

  • Depth benefits ISO and fixed structured covariance tasks only at finite context-to-dimension ratio; randomly rotated structured covariances benefit from depth even as context length grows.The fixed-covariance solution is also brittle to test-time covariance changes.Source: §1.1 · p. 2
  • The model introduces an asymptotic regime where context length, masked points per context, and contexts per training step scale linearly with data dimension.The stated compute and total-data requirement to converge is reduced by a factor of D relative to prior work (Lu et al., 2025).Source: §1.1 · p. 2
  • For power-law data in the randomly rotated structured setting, the pretraining scaling law has separate width and depth terms alongside training steps and context length.The law is L(t, N, L, P) = c_t t^-β_t + c_N N^-β_N + c_L L^-β_L + c_P P^-β_P.Source: §1.1 · p. 2
  • Compute-optimal joint width and depth scaling follows L ∼ N^ν, with ν depending on the structure of the ICL covariates.The paper introduces ISO, fixed structured covariance (FS), and randomly rotated structured covariance (RRS) task settings.Source: §1.1 · p. 2

1.2 RELATED WORKS

  • Prior theoretical ICL work has studied implementable algorithms and task diversity, while this paper develops a solvable model for comparing scaling resources across covariance structures.The related-work discussion also describes prior asymptotic analysis of task diversity in a shallow linear attention model (Lu et al., 2025).Source: §1 · p. 2, §1.2 · p. 3

2 DATA, ARCHITECTURE AND REDUCED MODEL

The paper uses a deep residual linear-attention model to map input contexts to predictions, with masked evaluation points for regression. It then studies a recurrent reduced model whose predictor is parameterized by a single matrix Γ.

2.1 DEEP LINEAR ATTENTION ARCHITECTURE

  • The depth-L residual linear-attention model maps input contexts to predictions using context pairs and additional evaluation points whose targets are masked.The evaluation tokens are prevented from updating the model by positional masking; the full population loss averages over context matrices.Source: §2.1 · p. 3
  • Linear attention is used rather than soft-max attention, and the paper states that this model suffices for its regression tasks while aiding theoretical tractability.The minimal reduced model encodes x information in a subspace orthogonal to y and ties the attention-block weights as a looped or universal transformer.Source: §2.1 · p. 3, §2.2 · pp. 3–4

2.2 RECURRENT REDUCED-Γ MODEL

  • The recurrent reduced model expresses the test-point predictor through a single matrix Γ, with the empirical covariance defined from the context inputs.The analysis first performs gradient descent directly on Γ, then considers learning the decoupled query, key, and value parameters.Source: §2.2 · p. 4

3 LEARNING CURVES IN REDUCED LINEAR ATTENTION MODEL

For isotropic and fixed structured covariances, depth improves in-context regression chiefly when context is limited; with long contexts, a single layer can attain the minimum loss. Fixed-covariance pretraining is specialized to its training distribution, while random rotations across contexts prevent that covariance-specific solution.

3.1 ISOTROPIC COVARIATES AND TASKS (ISO)

  • In the proportional asymptotic regime, successful pretraining requires Θ(D) contexts, each containing Θ(D) data points.The scaling takes P, K, B, and D to infinity with P/D, K/D, and B/D held fixed.Source: §3.1 · p. 4
  • For L = 1, the loss saturates at (1 + α)^−2, while for L→∞ it is [1−α]+; this illustrates the performance gap between a shallow model and a large-depth model.As α→∞, L = 1 achieves the minimal ICL loss among models of any depth.Source: §3.1 · p. 5

3.2 FIXED STRUCTURED COVARIANCE (FS) →PRECONDITIONED GRADIENT DESCENT

  • With fixed covariance, no depth beyond L = 1 improves the large-context limit, where Γ = LΣ^-1 yields zero ICL loss.Pretraining can memorize the input covariance and precondition by its inverse, enabling zero loss after one in-context GD step.Source: §3.2 · p. 5
  • For codiagonalizable Ω and Σ, Γ's eigenvalues follow decoupled dynamics, and the large-depth loss scales as a power law in pretraining time.The dynamics assume zero initialization; the large-depth solution is given for each eigenmode.Source: §3.2 · p. 6
  • Fixed-covariance pretraining yields a solution specialized to its training covariance, and OOD loss rises as the evaluation covariance becomes more dissimilar.The covariance shift is parameterized by rotating Σ with a random skew-symmetric matrix.Source: §3.2 · p. 6
  • Random rotations across contexts prevent Γ from encoding a whitening transform, precluding a zero-loss shallow solution and yielding unconditioned in-context GD.The method uses randomly rotated data and task covariances across contexts.Source: §3.2 · pp. 6–7

4 MODEL OF COMPUTE OPTIMAL NEURAL SCALING LAWS

The section develops a two-point deterministic equivalent for the randomly rotated setting and uses it to obtain asymptotic loss and compute-scaling results. For powerlaw data, its figures show that width and depth must be scaled together for monotonic improvement with compute.

Two-point deterministic equivalent

  • The two-point deterministic equivalent gives the full asymptotics and supports Chinchilla scaling laws for powerlaw data.The section presents the equivalent as a deterministic function obtained by averaging over the random orthogonal matrix.Source: §4 · p. 7
  • The deterministic-equivalent calculation uses two-point resolvent correlations, which capture more than the single-resolvent spectrum.The functions depend on the spectra of the structured covariance and A⊤A.Source: §4 · p. 7

Compute-optimal scaling

  • For powerlaw data, reduced-Γ training loss decays as t^-β/(2+β) before reaching an asymptote that scales as L^-β.The loss landscape has minima near γ ≈ L.Source: Figure 4 · p. 8
  • For powerlaw data with random covariance across contexts, scaling N and L together yields monotonic improvement with compute, whereas scaling either alone creates a bottleneck.Width-only scaling is depth-bottlenecked; depth-only scaling is width-bottlenecked.Source: Figure 5 · p. 8

5 MORE REALISTIC SELF ATTENTION MODELS

The section tests whether the reduced-Γ dynamics extend to less constrained attention models. Untied Γ layers preserve recurrent-model dynamics from balanced initialization, while full linear and softmax attention show related depth-dependent behavior.

5.1 GRADIENT FLOW WITH UNTIED Γ MATRICES

  • When untied Γ layers use a learning rate scaled by depth, their dynamics are equivalent to the recurrent model in the RRS noise-free setting.Balanced initialization maintains equal layer scales, making the loss dynamics identical to the recurrent model.Source: §5.1 · p. 9

5.2 GRADIENT FLOW FOR FULL LINEAR ATTENTION

  • With separate attention weights, fixing read-in and readout weights to unit norm reduces the dynamics to a one-dimensional ODE like the reduced-Γ model.Under the stated setup, the full module is interpreted through the reparameterization γ(t) → w(t)^5; its learning curves closely match the predicted dynamics.Source: §5.2 · p. 9
  • The separate-weight full-attention dynamics are tractable as low-dimensional ODEs but include transient blowup and recovery in the read-in and readout scales.Fixing those scales to unit norm removes this complication in the reduced dynamics.Source: §5.2 · p. 9

5.3 NONLINEAR (SOFTMAX) ATTENTION

  • In softmax attention trained with Adam, depth remains beneficial to ICL performance, consistent with the linear-attention model.The section also reports experiments with non-recurrent softmax attention and Adam.Source: §5.3 · p. 9

6 DISCUSSION

The discussion links the learned ICL solution and the need for depth to the covariance structure of pretraining tasks. It also describes a two-point deterministic-equivalent method for the randomly rotated setting and notes the analysis’s scope.

6 DISCUSSION

  • With fixed covariance, the model learns a whitening preconditioner and does not require large depth at infinite context, but this solution is brittle to covariance changes.With randomly rotated covariances across contexts, it instead learns a general-purpose in-context gradient descent algorithm, with width and depth both able to bottleneck performance.Source: §6 · p. 10
  • The analysis is limited to linear regression with linear attention; nonlinear function approximation and nonlinear attention are identified as future directions.The discussion also frames the work as an advance in solvable models of neural scaling laws and ICL structure.Source: §6 · pp. 9–10

A TWO POINT DETERMINISTIC EQUIVALENT FOR GENERAL FREE PRODUCT FROM DYNAMICAL MEAN FIELD THEORY

  • The two-point deterministic equivalent supplies an exact loss-landscape result for randomly rotated ICL and supports approximations used to derive scaling laws.The method averages over a random orthogonal matrix and expresses the outer product of the v variables through a trace-equivalent function.Source: A TWO POINT DETERMINISTIC EQUIVALENT… · pp. 16, 18

B.1 LINEAR ATTENTION AND POSITIONAL MASKING FOR OUR TASK

The positional mask defines different residual-stream updates for labeled training examples and evaluation points. This keeps test points out of the updates and assigns opposite signs to training and test updates.

B.1 LINEAR ATTENTION AND POSITIONAL MASKING FOR OUR TASK

  • Only the P labeled examples contribute to residual-stream updates; evaluation points are excluded by the positional mask.The target values are provided only for the P labeled examples in each context.Source: §B.1 · p. 18
  • The mask assigns opposite signs to training-point and test-point updates, a design that enables an in-context gradient descent rule.The residual variables can encode test-point residual errors yµ − fℓµ.Source: §B.1 · p. 18
  • Using the same signs for training and test points is also possible if the model output is negated at the end.The alternative is stated as f → −f.Source: §B.1 · p. 18

B.2 DERIVATION OF REDUCED GAMMA MODEL

The section derives a reduced-Γ recurrent model from aligned attention weights and relates it to preconditioned in-context gradient descent. It then sets up isotropic covariates and tasks for asymptotic analysis.

B.2 DERIVATION OF REDUCED GAMMA MODEL

  • Tying Γ across layers gives recurrent residual-stream dynamics, with training-point variables following a closed recursion.The test-point recursion differs because test points receive attention only from the unmasked training tokens.Source: §B.2 · p. 19
  • The aligned attention-weight assumptions yield a reduced-Γ recurrent model equivalent to preconditioned in-context gradient descent.The input weights project data into an orthogonal subspace, while target read-in and readout weights align.Source: §B.2 · pp. 18–19

C ISOTROPIC (ISO) SETTING THEORY

  • The ISO setting takes both input covariance and task-vector covariance to be identity, with context length and dimension growing at fixed ratio.The empirical covariance spectrum follows the Marchenko–Pastur law, whose eigenvalue density depends on the aspect ratio.Source: C ISOTROPIC (ISO) SETTING THEORY · pp. 20–21

C.1 SHALLOW SGD THEORY

The shallow analysis characterizes how SGD approaches gradient flow and how batch size, context length, and masked evaluation points affect its dynamics. It also describes the isotropic scalar gradient-flow dynamics at greater depth.

C.1 SHALLOW SGD THEORY

  • Figure 7 reports monotonically improving loss as τ, α, and κ increase, with the asymptotic theory predicting the learning curves.The panels vary τ, α, or κ while holding the other listed parameters at one.Source: Figure 7 · p. 22
  • O(D^2) total tokens suffice for convergence in this ICL regime, compared with O(D^3) in prior work.Shallow-model deviations from gradient flow scale as η/τ.Source: §C.1 · p. 23
  • The number of masked evaluation points per batch affects only SGD noise terms, and its effect vanishes in the gradient-flow limit.In online SGD, label noise also regularizes the final value of γ to reduce overfitting.Source: §C.1 · p. 23
  • Controlling η/τ accesses the gradient-flow limit when batch size scales linearly with dimension.The analyzed joint limit scales context length, batch size, masked evaluation points, and dimension proportionally.Source: §C.1 · pp. 21, 23

C.2 GRADIENT FLOW IN DEEP MODELS

  • At greater depth, gradient flow preserves an isotropic Γ(t)=γ(t)I and reduces the dynamics to a one-dimensional loss landscape.Rotational invariance keeps the gradient proportional to the identity.Source: §C.2 · pp. 23–24

D FIXED AND STRUCTURED COVARIANCE (FS) SETTING

For fixed structured covariances, gradient flow aligns Γ with the shared covariance eigenbasis and yields power-law loss under source/capacity assumptions. The included scaling analysis also relates loss decay to growth in γ and notes finite-resource saturation.

D FIXED AND STRUCTURED COVARIANCE (FS) SETTING

  • When the task and data covariances are simultaneously diagonalizable, Γ shares their eigenbasis, and the large-depth dynamics can be solved exactly.The resulting loss scales as a power law under source/capacity assumptions.Source: D FIXED AND STRUCTURED COVARIANCE (FS)… · p. 24
  • The power-law approximation describes large but finite depths as well as the infinite-depth result.The section points to Figure 3 for the finite-depth comparison.Source: D FIXED AND STRUCTURED COVARIANCE (FS)… · p. 24

E.1 EXPLOITING SYMMETRY IN THE GRADIENT UPDATES

  • In the RRS analysis, rotational invariance preserves isotropic gradients at isotropic Γ and reduces γ dynamics to a one-dimensional loss landscape.The loss expression is invariant under orthogonal transformations of Γ.Source: §E.1 · p. 25

F.1 DYNAMICS AT INFINITE N, L, P

  • In the infinite N, L, P limit, loss follows a power law whose exponent is set by the growth rate of γ; finite resources cap γ.The stated dynamics are expected to hold in the joint infinite-resource limit.Source: §F.1 · p. 25

F.2 DEPTH SCALING LAW AT INFINITE TIME, WIDTH, AND CONTEXT LENGTH

The section develops depth, width, and context scaling analyses and extends the model to distinct attention-layer parameters. It identifies stability and resource bottlenecks, while showing that balanced layer scales recover recurrent dynamics in the noise-free RRS setting.

F.2 DEPTH SCALING LAW AT INFINITE TIME, WIDTH, AND CONTEXT LENGTH

  • At finite depth, stability along the top eigendirection imposes a maximum γ, and the error at γ⋆ depends on the top eigenvalue.The depth-scaling analysis considers infinite time, width, and context length.Source: §F.2 · pp. 25–26

F.3 MAPPING THE LOSS FUNCTION TO A TWO POINT DETERMINISTIC EQUIVALENT

  • Dynamical mean field theory supplies an effective loss as a function of width N, depth L, and γ.The derivation maps the loss to a deterministic equivalent using resolvents.Source: §F.3 · p. 26

F.4 ORTHOGONAL PROJECTION AND A IS STRUCTURED WISHART

  • The width-bottleneck analysis studies loss scaling with N, while context-limited scaling follows from rank constraints on learnable eigendirections.Rank-N or rank-P matrices allow only the top N or P eigendirections to be learned.Source: §F.4 · p. 27

G.1 DIFFERENT PARAMETERS FOR EACH LAYER

  • With distinct layer parameters in large-width, large-context RRS, the noise-free loss dynamics retain the same power laws.The analysis also considers target noise separately.Source: §G.1 · p. 27, Figure 8 · p. 28

G.1.1 NOISE FREE LOSS FUNCTION AT LARGE WIDTH AND CONTEXT LENGTH

  • Permutation symmetry keeps the layer scales balanced, making their dynamics match the recurrent model when its learning rate is scaled by L.The balance is illustrated in Figure 8.Source: §G.1.1 · p. 28, Figure 8 · p. 28

G.1.2 NOISY TARGET FUNCTION

  • For noisy targets and decoupled layers, the predictor is expressed through a recursion for the weight discrepancy.The discrepancy is defined between the target weights and the learned weights at each layer.Source: §G.1.2 · p. 28

G.2 DECOUPLING WEIGHTS IN ATTENTION LAYERS

  • In decoupled attention layers, isotropic gradient updates reduce the dynamics to scalar weight-scale ODEs under small-initialization assumptions.With specified output weights, the loss simplifies further; the input weight updates twice as quickly as the other weights.Source: §G.2 · p. 29

G.3 SOFTMAX ATTENTION

The section reports numerical experiments with non-recurrent softmax attention and describes a corresponding alternating softmax-attention/MLP architecture. The figure caption states that adding MLP blocks on the residual stream also benefits from increasing depth.

G.3 SOFTMAX ATTENTION

  • Numerical experiments use non-recurrent softmax attention, decoupled layers, and Adam, with CompleteP learning-rate scaling ΘL(1) (Dey et al., 2025).Their results are provided in Figure 6(c).Source: §G.3 · p. 29
  • Figure 9 gives an equivalent figure for a non-recurrent architecture alternating softmax attention with GELU-activated MLP blocks.Source: §G.3 · p. 29
  • The softmax attention model with MLP blocks on the residual stream also benefits from increasing depth.Figure 9 presents this result.Source: Figure 9 · p. 29
Loading 2510.01098v1…