A Solvable Attention for Neural Scaling Laws
Authors: Bochen Lyu
Di Wang
Zhanxing Zhu
Venue: ICLR 2025 Poster (2025)
Keywords: self-attention
scaling laws
solution of learning dynamics
Primary area: learning theory
TL;DR: We provide a solution to self-attention and apply it to investigate neural scaling laws of self-attention.

Abstract: Transformers and many other deep learning models are empirically shown to predictably enhance their performance as a power law in training time, model size, or the number of training data points, which is termed as the neural scaling law. This paper studies this intriguing phenomenon particularly for the transformer architecture in theoretical setups. Specifically, we propose a framework for linear self-attention, the underpinning block of transformer without softmax, to learn in an in-context manner, where the corresponding learning dynamics is modeled as a non-linear ordinary differential equation (ODE) system. Furthermore, we establish a procedure to derive a tractable approximate solution for this ODE system by reformulating it as a *Riccati equation*, which allows us to precisely characterize neural scaling laws for linear self-attention with training time, model size, data size, and the optimal compute. In addition, we reveal that the linear self-attention shares similar neural scaling laws with several other architectures when the context sequence length of the in-context learning is fixed, otherwise it would exhibit a different scaling law of training time.

Published as a conference paper at ICLR 2025
A SOLVABLE ATTENTION FOR NEURAL SCALING LAWS
Bochen Lyu! Di Wang"
Zhanxing Zhu
'University of Southampton
DDatCCanvas 'Independent Researcher
"bochen. Liv@gmaill. com
tdi. wai.wang... 2718 18:18880 gmaal. .cmai..com "z.zhu@soton..a ac.
ABSTRACT
Transformers and many other deep learning models are empirically shown to pre- dictably enhance their performance power law in training time, model size, or the number of training data points, which is termed as the neural scaling law. This paper studies this intriguing phenomenon particularly for the transformer ar- chitecture in theoretical setups. Specifically, we propose a framework for linear self-attention, the underpinning block of transformer without softmax, to learn in-context manner, where the corresponding learning dynamics is modeled as a non-linear ordinary differential equation (ODE) system. Furthermore, we estab- lish procedure to derive tractable approximate solution for this ODE system by reformulating it as a Riccati equation, which allows us to precisely characterize neural scaling laws for linear self-attention with training time, model size, data size, and the optimal compute. In addition, we reveal that the linear self-attention shares similar neural scaling laws with several other architectures when the con- text sequence length of the in-context learning is fixed, otherwise it would exhibit a different scaling law of training time.
INTRODUCTION
Large language models (LLMs) (e.g., GPT (Brown et al. 2020) and Llama (Meta 2024)) have made sigmificant achievements across variety of tasks, ranging from question answering to decision making. Adopting the transformer architecture (Vaswani et al., 2017) these I.LSs are large in the sense of both parameters and training data, e.g., the largest Llama model has 405B parameters and is trained on 15.6T tokens (Meta 2024). One of the most fantastic phenomena of such LLMs is their continuing performance gaining as the model size and training steps are scaled up. More remarkably, their performance can behave predictably
power law in the number of parameters, computation or data size (Kaplan et al. 2020 Hoffmann et al. 2022). This impressive power law behavior is termed neural scaling laws.
In particular, for a model with D trainable parameters, neural scaling laws state that the test loss L(D,t) should obey I(D,t) F. AD- Bi- (Kaplan al. 2020 Hoffmann et al.) 2022) where is the number of optimization steps and captures the loss for a generative process on the data distribution. Holding across wide range of orders of magnitude, these neural scaling laws have led to the fundamental belief that autoregressive transformer language models could successively improve their performance when scaling up. Interestingly, they also allow practitioners determine the trade-off between model size and training time for fixed compute budget (Hoffmann et al. 2022] or design dataset with clever pruning (Sorscher al. 2022
Given the significant role of neural scaling laws, the theoretical understanding of their origin and mechanism such as values of their exponents becomes increasingly important recently. Hutter (2021) designed linear model that can exhibit power laws and showed that not all data distributions lead to power laws; Maloney et al. (2022) applied the random matrix theory to identify necessary properties of scaling laws and proposed a statistical model that captures the neural scaling laws; (Bodello et al. (2024); >; Nam et al. (2024) proposed different solvable models to reveal the existence of scaling laws. Although these initial attempts made simplifications model architectures and data for the purpose of analytical tractability, they largely advanced our understandings of neural scaling laws from the theoretical perspective.

---

Published as a conference paper at ICLR 2025
On the other hand, one important aspect commonly absent in these works is that they did not consider the transformer architecture, the universal architecture of current LLMs, which leads the theoretical understanding for neural scaling laws of modern LLMs to be still underexplored. Transformers are special not only because they employ self-attention as the primary component, but also because the way how they perform prediction- an incredible mechanism called in-context learning (Brown) itt al. 2020 Garg et al. 2023) that can adapt their predictions based on data given in context.
The uniqueness of transformer definitely gives rise many intriguing questions from a theoretical perspective. What are the origins of neural scaling laws of transformer? Does transformer induce different neural scaling laws compared to other models? Will in-context learning (e.g. context sequence length) affect neural scaling laws? Due to the importance of transformer and its neural scaling laws, investigating these questions of great interest and necessary.
Answering these questions from a theoretical perspective requires thorough understanding of ex- plicit forms of model predictions during training. which, however, hard since it typically requires solving non-linear ODEs that usually do not admit closed-form solutions. Towards this direction, Saxe et al. (2014) modelled the learning dynamics of deep linear networks as the logistic differential equation that can be solved exactly, Pinson et al. (2023) solved the dynamics of línear convolution neural networks, and Borddloo et al. (2024) applied a DMFT approach from statistical physics to solve random feature models. For transformers, recently Zhang et al. (2023) Tarzanagh et al. (2024) established the forms of converged parameters in regression and classification settings. However, explicit forms of parameters along the training trajectory are still unclear, leading to a gap when investigating neural scaling laws for transformers.
In this paper, we attempt to provide initial answers for the aforementioned questions to fill the gap in part and take a step towards understanding neural scaling laws of LLMs. To conduct an amenable analysis, we focus on the self-attention which stands the core of the transformer architecture, in the linear case. We note that linear self-attention has been widely adopted in recent works (von Oswald et al. 2023; Li al. 2023b; Zhang et al.) (2023) to study properties of transformers. Despite that feature learning is absent, it has the advantage of providing the possibility for clear theoretical characterization. We discuss more related works on learning dynamics, neural scaling laws, and the
analysis of in-context learning for (linear) self-attention in Appendix B
Our Contributions.
1. We design a multitask sparse feature regression (MSFR) problem for the linear self- attention block learn in an in-context manner. More importantly, derive a tractable solution for linear self-attention by modelling its in-context learning dynamics in the MSFR problem cati equations. This is highly nontrivial since non-linear ODE systems are hard to solve, thus our procedure might be of broad interest.
non-linear ODE system and reformulating the system set of Ric-
This solution captures dynamical behaviors of linear self-attention during training explic- itly. To the best of our knowledge, this is the first closed-form solution of self-attentio along the training trajectory. We highlight that it can be applied as an interesting proxy for investigating properties of self-attention and transformers due to its analytical tractability.
2. Built upon this solution, varying time, the size of model, or the number of training data points when data obeys a power-law, which then gives us the scaling law in the optimal compute budget. In addition, we are able to characterize the role of context sequence length in neural scaling laws, revealing that if it obeys different power-law then the time scaling law will be affected, otherwise linear self-attention would share similar neural scaling laws with other models, which well aligns with empirical observations in Kaplan al. (2020).
characterize neural scaling laws of linear self-attention by
2 SETUP OF FRAMEWORK
Notations. We use (1....,
.....,N N2 } to denote all integers between and N. For two vectors a, b E Rd aj to denote its j-th component, the inner product, and diag(a) a) to denote the d x d matrix with its diagonal elements equal to a. We use à to denote the derivative of a with respect time. We let d,,y' be if s' and 0 otherwise.
@ b to denote the elementwise product, a b to denote
2

---

Published as a conference paper at ICLR 2025
For matrix A, use A., denote its i-th row j-th column component. We use P to denote that a is sampled from distribution P. We use O₂ I2 to denote the vector in R
In Section 2.1 we define the problem setting of MSFR problem, and present the concept of in- context learning and the generation of in-context learning data for it in Section 2.2 Finally, in Section
we describe details of linear self-attention block.
2.1 MULTITASK SPARSE FEATURE REGRESSION PROBLEM
There are Nr different tasks in total. We let S be the random variable of picking specific task among N₂ tasks and assume that follows power law distribution:
Pa(S= Zs
(1)
where Z = DE= is the normalization constant and a Since we focus on linear self- attention, we assume the existence of non-linear sparse feature extractor to perform the feature learning. Specifically, for an input data vector x: E Rdd and a task type (l,....Ns). there exists unique feature extractor
(s) RR Rd + (-1,0,13M, IR
where only the s-th component of ((s,, x) can be nonzero, i.e., (s, x) = Furthermore, given task type s, we let the strength for task be A, E R. The target y E IR is now defined through
N
v(s x) A, Σ (s)
We elaborate two properties of this problem before moving on. (i) The reason why this problem is termed as "multitask" is because we have N, different tasks such that each has its own task strength Ag and feature extractor (fs, x), meaning that the model should learn distinet A, for each task. (ii) If we let A E RRV, be the collection of all task strengths, then the target can be written as y(s, x) = A (68, æ)), which is like a linear regression over the feature (68, .)). Since Φ6s, x) like one-hot vector, the problem is "regression with sparse feature". The subtlety lies in that must rely on all task types to learn the complete when compared to standard linear regression. Therefore, our problem is defined as "multitask sparse feature regression".
2.2 IN-CONTEXT LEARNING
A remarkable ability of LLMs is that they can perform in-context learning to adapt to a specific task given context in the form of instructions (Brown et al.) 2020). More specifically, the goal of in-context learning is to enable learner (e.g., transformer) to use the context data to make prediction for the query data. To incorporate this ability, we focus on in-context learning in this paper, and present its details formally for the MSFR problem (Section 2.T in this section.
Generation of in-context data. Given task type s, we let the corresponding context sequence length v. F(s) E R and task strength A, = ((s) E R, i.c., each task s has constant context sequence length t/s and a constant task strength A, determined by the maps F and G, respectively. The generation is composed of four parts (see Fig. [T: (i) task type E (I...., ,N,) is first sampled from the distribution (S length tha and task strength As; (ii) we sample vi, different input vectors E IR and a query vector E Rd from the input data distribution Px, then these data vectors are organized to form a matrix
s) (Eq. (.), which gives us the corresponding context sequence
<(v+.) a) (v+((v+ (iii) we apply the feature extractor ø to each column (+)) of X to obtain the sparse feature (()) E IR.N and generate the target
(x())
(2)
( : y(s, (f) and ğ == v(s, it) according to Eq. (3). then one in-context data point of the task can now be generated
(1))
(s, (1)) Lb(s, ()
((0,)))0))) B(s, (Wa))
Φis ax)
(4) (iv) repeating the above procedure for N times can give us an in-context dataset with N data points, where the numbers of data points for different tasks obey the power law Eq. ( Finally, given in- context data (s. X) [RV++1)x(ee++)) (Eq. (4) and loss function L, in-context learning aims to
( X) = y
3

---

Published as a conference paper at ICLR 2025
Sequence Length (ba) F(8)
Step (1)
Step (i)
a ~ Pa = s)
R ~ Px(s,
[ar(1) ax
ml d) goxto,
Task Type
Data
Data Matrix
Query Data
A(s) = 9(x)
Step (iv)
Task Strongth
Peature Extraatorr
Repeat
Step (fii)
Øls 2) y(t)
dis 2(a,2(1)) 1/2)
bb 2(10)) y(s))
E
E
0
In-context data
Figure 1: In-context data generation of multitask sparse feature regression (MSFR).
learn model f (+A.+)) + R such that 0* arg ming (()(() 0), Note that MSFR can be seen as limiting case of in-context regression under source/capacity condition (generaliza- tion of setup in Lu et al. (2024) with that Cui et al. (2022). see Appendix H.l
2.3 LINEAR SELF-ATTENTION BLOCK
Self-attention block stands at the core of the transformer architectures VVaswani al., 2017). A single-head self-attention block (without residual connection) f Rdxde H Raxdz parameterized by 0 updates input G E Raxdz to
G f(= = PVG softmax ((WxG) E Bddd
where = VP,V, WK. We). P E Rddd is the projection matrix, V E Rd. is the value matrix, and WK, Wo E Rdde xd are the key matrix and query matrix, respectively. Note that softmax applied column-wise.
In this paper, we study a simplified version of the self-attention by removing the softmax operation and merge the key matrix and query matrix as single matrix WKQ W W Wo Rd which has been a popular choice in recent works, e.g., et al. (2023); von Oswald et al. 2023). due to its analytical tractability as well the ability of capturing properties of the standard sell-attention.
In particular, given the MSFR problem (Section 2.1) and in-context data ((s, X) (Eq. 44 ), we study the in-context learning of linear self-attention block V WKQ E NR (vba+t) where WKQ e
1)N,+))*(X++1) To obtain the scalar prediction of the query data 6 adopt the output of the self-attention as f (); 0) = [V "V W slight change to self-attention, it provides the convenience to investigate its intricate in-context learning dynamics. Furthermore, decomposing WKQ as (W w-1) where
Despite that this formulation makes a
W = (w)
wN,) R(N++1)XN, Vi
N,.. w. E R.
(5)
we can write the output of the linear self-attention block for the task as ( [V " W r). which will be used in the rest part of this paper.
=
3 A TRACTABLE SOLUTION OF LINEAR ATTENTINNN
Section establishee our in-context learning framework of multitask sparse feature regression prob- lem for linear self-attention. In this section, will closely investigate the corresponding learning dynamics by modelling as non-linear ODE systems in Section 3.I and give a tractable solution of it in Section 3.2

---

Published as a conference paper at ICLR 2025
3.1 IN-CONTEXT LEARNING DYNAMICS
Given the in-context dataset generated according to the procedure described in Section 2.2 with N data points {*(s(m) ) we use the mean-square mnea-s error (MSE) loss such that
2N Σ E(s (1(((5m) x(()) -g(n)) 2
1
T((0)
(6)
where İ is the empirical loss. The goal of in-context learning now becomes the traditional empirical loss minimization 0 arg ming i(()), which can be solved by various optimization algorithms, and we focus on the general gradient descent in the continuous time limit, i.e., gradient flow (GF):
V = -VvL(V W): W = - Vw W).
To further investigate the learning dynamics, considering the formulation of the feature extrac- tor Eq. (2) and in-context data Eq. (4) and denoting the standard basis vector in R.V, as e, =
(0 0 0 ...)" E R.V. for [l,.....NG) such that the only nonzero component of
e, its s-th component, we find that H, (R(NG+))X(VN-1))) defined by (Appendix D.1
H. = 2 (s(m) (m) x" 8" (s ((()) = diag ((t's + 1)ex) U2A.ee""
VeA.e, V.
does not change for different n, where vi is the context sequence length and A, is the task strength for the task s, both of which only depend on the task type s. H, is composed of the feature covari- ance and target (Eq. (24)). In addition, if we further decompose V as V (w)
VNC), Vi e ,.....,,... } Vs E N2 and recall the decomposition of W Eq. (5). then we can rewrite the olriginal empirical loss Eq. (6) as (Appendix D.I
N.
# 2 N
Empirical loss function: i
(D" H.W, -.,)²
(8)
where #. denotes the number of in-context data points for the task type s in the dataset ( (#(s("), { i... #. Σn ), Eq. (8) indicates that the dynamics of v, and w, for different are decoupled: the s-th row of V and s-th column of W are responsible for learning and predicting the task strength of the task type s, rendering self-attention adapting itself different tasks according to the in-context data. With this empirical loss function, we can now use set of non-linear ODE systems E ,o,)
In-context learning dynamics: ú, = # N A,) (,) w. =
#o (fa- A.) H.U, (9)
N
to describe the in-context learning dynamics by GF where we denote Sa = D² H&W, for simplicity. We note that fa is sufficient for investigate the dynamical behaviors of the output of self- attention for task type s and the empirical loss. Thus, by abusing of definition, we refer to the solution of fa as the solution of the in-context learning dynamics, which also be applied to give solutions of v. and w.,
We highlight that the ODE systems above non-linear for both us R.A, and w, RR-V, and, ob- viously, are different from the logistic differential equations obtained from the GF dynamics of deep linear networks (Saxe al., 2014 Nam et al.) 2024) and different from the Lotka-Volterra predator- prey model (Volterra) 1928). In this sense, the dynamics of linear self-attention (and transformer) different from that of deep linear networks. Meanwhile, we note that non-lineat ODE systems, including Eq. (9) which are non-lineat ODE systems for vectors, typically do not admit closed-form solutions. Therefore we emphasize that solving Eq. (9). which might be of independent interest, to obtain the explicit dynamical behaviors of linear self-attention is novel as well as intriguing.
3.2 SOLUTION OF IN-CONTEXT LEARNING DYNAMICS
Although is intractable to give the exact closed-for form solution to the in-context learning dynam- ics Eq. (9). in this section, we will provide solution that can be approximately exact under the following condition. We defer technical details of this section to Appendix E
Assumption 3.1. Vs (1,...,N..) the context sequence length vh. > 1.
5

---

Published as a conference paper at ICLR 2025
Procedure sketch. Before diving into detailed procedure for deriving the solution, we first present a rough sketch for it. The first step is to transform ODE systems Eq. ( to a more sym- metrical form Eq. (IO) by changing of variables. Then decompose Eq. (IO) two sets of ODE systems by comparing both sides of Eq. (10) to the zero-th and first orders of c,:= 1/th, since H, can be decomposed as two parts, H and c, H with 6s < 1. We then apply a change of variable again and derive a new set of ODEs Eq. (IT as Riccati equations Eq. (12) which admit closed-form solutions by noticing the existence of an important conserved quantity of the dynamics.
We now discuss the procedure in detail. Our first crucial observation is that the ODE of v. Eq. non-linear with respect to Wg. which makes it hard solve. Therefore, we first convert Eq. ( into a more symmetrical form by changing of variables: let 17, = v, + w. and Pa = wa, then the dynamies of 17, E TR-V++1 and ps E &RN++1 can be obtained according to Eq. (
# y- ht
#
g- hs
n N
A.) Hsss Ps
As H.Ps:
(10)
N 4
where we define g. n H, and h2 - -..H..p.
In this way, by solving Eq. (10). we can find the solution of the self-attention and the empirical loss
function Eq. (8). However, Eq. (10) is still not directly solvable. Fortunately, recalling the definition
of H. in Eq. we can rewrite H, as sum of two matrices H, v/ (HQ + H,H.) where
diag(ee) A,e
Age=]
ddag (es) 0
H
H
0
A
and =/16/ < 1 according to Assumption 3.1,
which allows treat cH, an insignificant perturbation in the dynamics Eq. (10) and solve using the perturbation analysis.
Specifically. suppose that the solutions of Eq. (10) can be written 17. = n/. cem. and p. = pa + ExP;, such that 77, and pa are treated as perturbations to ny and pa respectively, then gs and hg can also be written in a perturbed form gs = 92 + es and h, -1 +e,hh, accordingly (Appendix Ip. Now we can obtain ODEs for 17, Ts P% and P. by comparing terms to the zero-th and first orders of E, in both sides of Eq. ([O). respectively (Appendix E..I. This will finally give us ODEs for oh hg. gh and h, the final ODEs that we aim to solve since the output of self-attention for the task s can be written as fa = (t) (t) [o..- 12) + co h/]]/4.
Our strategy for finding solutions of fa is now composed of two parts using the perturbation analysis: (i) solve the ODEs for of. and h. exactly and (ii) find nº and p.. according to the solved ge. and hD. then put them into ODEs of y and h!g to find their solutions. As mentioned earlier, (L) [L) is far less significant than (t) to the dynamical behaviors of self-attention given Assumption 3.1 thus we defer the discussion of f. (t) to Appendix J and only focus on (t).
We now discuss the first step for (1). Our key observation is that HD is like an idempotent matrix: (H2)² (A3 + H) which gives us the dynamics of 9, and ht as new set of non-linear ODEs:
0- (g h2 4A.) /2. ig = (9 h: 4A.) //
(11)
where we let ag v. + 1)/N for ease of notation. Though Eq. (II) still non-linear ODE system, it is much more tractable than the original in-context learning dynamics Eq. (9. Our following key observation can drastically simplify Eq. (TT) even further: VI 2 0 goh 2C, where C2 is a constant determined by the initialization, i.e., goh2 is conserved for the dynamics, since h()) dt = O. In this way, Eq. [T] becomes the following set of Riccati equations that can be solved (Appendix E.2 to give our main results:
80 Aa, a.()?)²/ 2/22///22222 + a,C., i. - -a, (h,))// C.C..
(12)
Theorem 3.1 (Solution for in-context learning dynamics of linear self-attention: zero-th order). For MSFR problem and the in-context learning dynamics by GF of the linear self-attention block Eq. (9) the solution ft) can be approximately written as an expansion (t) - f.f. (t) at large vb with
X
(2) = Ag +
(13)
2
1 P. exp((addtt) 1+ Q, exp(a.Ast)
where a, = = #.W. + 1)/N, and
A, = V413 + 2C., P. 4J(( + 2C, + A. A. A(T 2C, (() A,) 2A,)
2A, Xa
Q. :2=PP 2A, + X,
6

---

Published as a conference paper at ICLR 2025
determined by the initialization. In addition, when 00) = (0), the constant C. 0 and, denoting A. = (Ag f0 (0), the solution can be simplified to be a standard function
As
(2 1+ De Aa,
when (0) (
(14)
Remark. The closed-form solution (t) obtained in Theorem 3.1 is to the zero- th order of 6, and is a good approxima- tion under Assumption 3.1 Uss. Pa. Qs. and X5 jointly control the learning pro- cess, where Ps. Qs. and X are determined by the initialization and task strength, c.g., when (0) = A, both Ps: Q2 + thus (t) A₂ for t 2 O. ag is determined by the dataset, e.g., sequence length, and task strength. In Fig. 2 we compare the loss that is calculated using (1) with that obtained from direct empirical simulation for different context sequence lengths. can be seen that when the context sequence length to > (blue lines), the theoretical prediction of the test loss matches with the empirical results precisely, which validates the accuracy of (t). In addition, we emphasize that Theorem 3.1 is applicable for any P. (Eq, ( of the task type, F(s) of the sequence length, and (s)) of the task strength for the MSFR problem in Section 2 ie., the derivation of (t) does not require assumptions on forms of P.. F(s), and (ss).
2.00
1.75
1.50
1.25
1.00
0.75
0.50
0.25
0.00
107
10
10'
30
Figure 2: Loss L()) for different context sequence lengths 25 and 100) during training. Solid lines are for theoretical predictions while dashed lines are for empirical simulations.
(t) can explicitly characterize the dynamical behaviors of the self-attention block including the influence of various parameters. suggesting that it could contribute to the understanding of self- attention in variety of aspects. In this paper, our focus will be neural scaling laws. For another ample (t) shows that self-attention learns different tasks in different rates that depend on various parameters such the sequence length, task strength, number of data points, and the initialization. The difference of learning speeds for different tasks might lead to the grokking phenomenon (Power) al. 2022), since the model can quickly fit fraction of tasks while learns the rest extremely slowly.
NEURAL SCALING LAWS FOR LINEAR SELF-ATTENTION
In this section, we examine neural scaling laws for linear self-attention with its special in-context learning dynamics. We note that this focus is different from those architectures considered in pre- vious works (Bordelon et al., 2024 Nam et al., 2024, Hutter] 2021 Michaud et al.] 2023) Mal-] oncy et al.) (022),, providing garding neural scaling laws from a theoretical perspective. In particular, we will investigate the test loss according to the solution obtained in Theorem 3.1 under Assumption 3.1 and the relation ft) = (t), which enables us to write the test loss as
possibility to compare self-attention with other architectures
N =s -A,12 (15)
Test Loss: L(t) LxwPx aN P. (((((((
Concretely, we consider neural scaling laws of linear self-attention with respect to each one of size of the model D, training time t, and the number of training data N when the other two factors are not the bottleneck of training. Finally, we consider scaling laws for the optimal compute C. We list the discussion of these factors as follows and defer details to Appendix
Size of the model D. To quantify the model size D, inspired by Michaud al. (2023); Bordelon & Pehlevan (2022); Nam et al. (2024) which considered different models empir- ically or theoretically, we assume that there is a cutoff D such that the model cannot learn any task strength A, with D. Specifically, we let WKO E DxD V E (Rx=1)xX and apply feature extractor such that ø' (s,8, a)) E R-D1 only extracts the first D
7

---

Published as a conference paper at ICLR 2025
elements of the original 2() When D is the bottleneck of training, let N oo (they are sufficient for the training) to derive scaling laws with respect to D.
Training time t. Training time t is equivalent to the number of optimization steps. To investigate the scaling law with respect to t. we remove the bottleneck caused by the size of model and the number of data points by letting N - oc and D Nss.
Number of training data N. When the training is bottlenecked by N, we let / - oo and the cutoff D = =N, following the above arguments.
Optimal compute C. This is the case when the number of data points is sufficient for the training. while training time t or the size of model D is the bottleneck given the compute budget C = Dt such that either t or D scales differently with C. Specifically, if Ltt, D) a,t at + -apD-an then we can derive the optimal test loss as L x C (-aran/(actan) given C = tD (Appendix(F.)
To inspect the effects of the context sequence length the E R and task strength A, E I on neural scaling laws, we let 3 F(s) x - which is inspired by the underlying power-law correlations in language sequence (Ebeling & Pöschel 1994), and A, ((s) x s7 with B, 7 0. We will investigate two different cases. In the first case (Section 4. I) we assume that 7 = 3 = such that the context sequence length v) and task strength A, do not depend the task type In this setting, we can compare neural scaling laws of linear self-attention with other architectures. In the second case (Section 4.2), we study positive and 7, which are unique self-attention. Details of Section
and 4.2 will be deferred to F.2 and respectively.
c=1a
18
e:
1
E 1o-
10-1
30-1
(a) Model Size Law with N, f 0C
(b) Time Law with N + oc, D N.
10
(000
12.1
50
100
10
(c) Data Size Law with co. D N
(d) Optimal Compute with N - co, C. = D
Figure 3: Neural scaling laws for linear self-attention with different values of a 1.8, 2.1. In cach figure, we use solid lines to represent empirical simulation results and dashed lines for power law curves. We also plot thcoretical predictions of test loss in (b) with dotted lines as comparison. In (d), we set a = 1.8 and use different levels of transparency to reflect different model sizes D within the range 10.1000,///5.
4.1 FIXED CONTEXT SEQUENCE LENGTH AND TASK STRENGTH
For simplicity assume that the model is initialized as C. C' and fo fo. which implies that X = X, P. = P. and Q, = Q for s E [1, do not depend on the task type s. When the context sequence length vh, and task strength A, are the same for different task types s, i.e.,
8

---

Published as a conference paper at ICLR 2025
y 3 = 0, the test loss Eq. (1))
be written
Zx²
1
L(I)
(16)
I PeaLde
1 + Qea.A
according to Theorem 3.1 of L(t)) with respect to each one of t. N, and D when the other two are not the bottleneck of training using Eq. [6). Then
To derive neural scaling laws, plan to find the asymptotic behaviors
use the relation C Dt find the scaling law with respect to the
optimal compute budget. We summarize our results in Table verify them in Fig. 3
Table I: Neural scaling laws for linear self-attention when 3 =
Scaling Law
Condition
Size of model D
D D
t + xc, N + 00
Time t
N + 00, D =Na
Number of data N
N 1=
I + 0o, 1) =Ns
Compute budget C
C 2=1
N + 00, t x Ca D x Co+i
Architecture does not matter for scaling laws when context sequence length is fixed. The results summarized in Table [ reveal that, when data admits similar power-law structure, linear self- attention shares the same neural scaling laws with ReLU MLPs (Michaud et al. (223)) and diagonal
linear networks (Nam et al., 2024) with respect to t, N and D. Linear self-attention also exhibits a similar time scaling law as the linear models considered in Bordelon et al. (2024) and Hutter (202T) These similarities indicate that the architecture of model does not affect exponents of neural scaling laws significantly, which well aligns with the empirical conclusion reached by Kaplan et al.) (2220) where they showed that transformers share similar exponents of neural scaling laws with other models when the power-law structures hold.
4.2
VARIED CONTEXT SEQUENCE LENGTH AND TASK STRENGTH
To further capture how the context sequence length ah and the task strength A, affect neural scaling laws for the in-context learning of self-attention, we let th, F(s) s Ag G(s) > as in Section 2.7
We can write the test loss as
Z N
A exp( -2,
L(t)
(17)
1 A exp( A2,,A
Following a similar procedure as in Section 4.1 we derive neural scaling laws of linear self-attention for MSFR problem in Table 2
which gives us the following insights.
Table Neural scaling laws for linear self-attention when both vhs and A, depend on
Scaling Law
Condition
Size of model D
D D-a-2y+1
-- x0, N co
Time t
848-39
0+22-1
N + xc, D = No
Number of data N
N
2+25-1
t - 00, D No
a+2y-1
a+3y+8 N oo,t x CateitoHi D x C C+t3tt441
Compute budget C
a43y+0-1
Varied context sequence length affects the scaling law of time. Table 2 reveals that varied con- text sequence length makes the learning process slower (Fig. 4ap. According to Table 2 nonzero positive 0 leads to larger exponent of time law, thus the test loss will decrease slower than the case when 3 O. This suggests that is better to balance the context sequence length for different tasks to obtain satisfied test loss given a limitation of optimization steps. This conclusion is special to self-attention compared to other architectures considered in previous theoretical works since they lack the place for the context sequence length. On the other hand, also find that 3 does not appear in scaling laws for the size of model and number data points, indicating that self-attention can still admit similar scaling laws for D and N other architectures when 7 = 0.

---

Published as a conference paper at ICLR 2025
Smaliton
Theary
4x10
C Reg
3x10
2x10
20-1
12"
10
10
10
10
(a) Time Law, 1,7 =
(b) Optimal Compute Law, 3 = 1, = 0
Simuaaion
101
10:4
1003
10
10
y
10
D
(c) Time Law, 3 = 7 = 1.5
(d) Model Size Law, 8 = 1 = 1.5
20a
1(0)
LG
On 175
13-2
Dn 200
1-2
12
12)
18
30a
10
(e) Data Size Law, 1,7 1.5
(l) Optimal Compute Law, 3 = 1,7 1.5
Figure 4: Neural scaling laws for self-attention with varied sequence length vi, - and strength A, x - each figure, we set a = 1.8 and use solid lines represent empirical simulation results while use dashed lines for power law curves. In addition, we plot the power law curves (black dashed lines) with 7, B = 0 from Table as comparison. In (a) and (b), we set 7 = 0 to examine effects of sh on time and optimal compute different levels of transparency to reflect varying model sizes D.
laws only. In (b) and (f),
Varied task strength affects all scaling laws. Table 2 reveals that a varied task strength reduces the requirements of the size of model (Fig. 4d) or the number of data points (Fig. 4ch in our MSFR problem. Specifically, due to the existence of a positive y, exponents scaling laws for both size of model and number of data points become smaller, thus the learning requires fewer number of data points or smaller size of model to achieve a similar test loss when they are the bottleneck.
5 CONCLUSION
In this paper, we target on understanding learning dynamics of self-attention, which stands at the core of the transformer architectures, from a theoretical perspective. For this purpose, we first design a multitask sparse feature regression problem for the self-attention to learn in an in-context manner, whose learning dynamics is then modelled as non-linear ODE systems. We then give tractable solution to the ODE systems, which should be of broad interest since non-linear ODE systems typically do not admit closed-form solutions. We also highlight that this solution can be employed as an interesting proxy for studying variety of properties of self-attention and transformers. Finally, we use the proposed solution to investigate neural scaling laws of alf-aatention with respect to each one of training time, number of data points, and size of the model when the other two not the bottleneck of the learning process, which in turn allows us to establish the neural scaling law with respect the optimal compute budget.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Ben Adlam and Jeffrey Pennington. The neural tangent kernel in high dimensions: Triple descent and a multi-scale theory of generalization, 2020. URL httpppp://aaxii..or//a/aba/
06786
Subutai Ahmad and Gerald Tesauro. Scaling and generalization in neural networks: A case study. In D. Touretzky (ed.), Advances in Neural Information Processing Systems, volume 1. Morgan- Kaufmann, 1988. URL http://cppe/ppoceedinnnnneeeeeee neurips cctppper. fiter_ffee/ppprrr///// 1988/ file//illl//////e a114414444. (tlpe//00002000000000000003333a0
Pappssaaaappppp..ppppp pdf
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement
preconditioned gradient descent for in-context learning, 2023a. URtptts:s//axxi..oo
abs/2306. 00297
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning, 2023b. URL/htttss:/arxxiv.oo
abs/2306. 00297
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning
algorithm is in-context learning? investigations with linear models, 2023. URL https://
arxiv.org/abs/ /211 01221.1166661 15661
Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks kemell learners:
The silent alignment effect, 2021. URL httphts:s////x org/ aass//a////// / 1t1///111...00010
Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. variance-limited behavior for networks in the lazy and rich regimes, 2022.
The onset of
URL https:
12147
Alexander Atanasov, Jacob A. Zavatone- Veth, and Cengiz Pehlevan. Scaling and renormalization in high-dimensional regression, 2024. URL https://xxxx org/ abs/2405. 0C592
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jachoon Lee, and Utkarsh Sharma. Explaining neural
scaling laws. Proceedings of the National Academy of Sciences, 121(27), June 2024. ISSN 1091-6490. doi: (0.1073//nas.231187888 2311878121
http:t:: /xx http:///0/o......///0000011////0
Boaz Barak, Benjamin I.. Edelman, Surbhi Goel, Sham M. Kakade, eran malach, and Cyril Zhang. Hidden progress in deep learning: SGD learns parities near the computational limit. Al- ice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neu- ral Information Processing Systems, 2022. URL http::tp//// htRp::tts:////ottt /ttp::/:///peeevieew net forum?id 8XWP 2ewX im
Carl M. Bender and Steven A. Orszag. Advanced Mathematical Methods for Scientists and Engi- neers. McGraw Hill, 1978.
Blake Bordelon and Cengiz Pehlevan. Learning curves for SGD on structured features. In Interna- tional Conference on Learning Representations, 2022. URL httpas://pperrvview net/
Forum?id=WB 12vbkA130
Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. goeeendumdependeet learning curves in kemnel regression and wide neural networks, 2021. Utl//ttpss///rrxivoorr//bbb//
02561
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. A dynamical model of neural scaling laws. In Forty-first International Conference Machine Learning, 2024. URL https:// // openreview Seenreview.eet/Fooeee www.:e///ooruue ftnt:/forrmmddmbbor d=nboY1 1mttRc)
Lukas Braun, Clémentine Dominé, James Fitzgerald, and Andrew Saxe. Exact learn- ing dynamics of deep linear networks with prior knowledge. In S. Koyejo. S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 6615-6629. Curran Associates, Inc., 2022. URL htt.uattsss/ppppeeaaa.....aaree.
http:////....pppeeeeeeee//////// fi F11229aab23311111977712223232222

---

Published as a conference paper at ICLR 2025
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Bernrr, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL httpp:://prcceeednnntt n newrips cttessccc/paeee cc/paper http:///wwppppp.aaaaeeeeee//eeee fi file/155cc66ffo00077111118222222
fflesspapee/2222/
A. Caponnetto and E. De Vito. Optimal rates for the regularized least-squares algorithm. Foun- dations of Computational Mathematiccs 7(3):31 - 2007. doi: 10.1007/s10208-000 1000000000--16--88
URL /tttps:///oil.rg/1101100000
0208 006 0196-8
Hugo Cui, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Generalization error rates in
kernel regression: the crossover from the noiseless to noisy regime". Journal of Statistical Me-
chamics: Theory and Experiment, 2022(111:: 14004,, November 2022. ISSN 1742-5468, doi:
1088/1742-5468//c22222 URL StLLttt:///ee //dx doi org/
Stéphane d' Ascoli, Maria Refinetti, Giulio Biroli, and Florent Krzakala. Double trouble in double descent: Bias and variance(s) in the lazy regime, 2020. URL
2003. 01054
Karthik Duraisamy. Finite sample analysis and bounds of generalization error of gradient descent in in-context linear regression, 2024. URL https: http://axxvvvvvov wttp:/iwwwrr//b//////0// 02462
W Ebeling and Pöschel. Entropy and long-range correlations in literary english. Europhysics Letters (EPL), 26(4):241 May 1994. ISSN 1286-4854. doi: 10. 10.12090099555555564400000
URL http:::/t..cii.oooo// www.xxxoo.oooooooooo.oo
1209 02009/02955 007522255550000220
Benjamin Edelman, Surhhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and vari- able creation in self-attention mechanisms, 2022. URL
10090
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes, 2023. URL htttps://axxiv..ott wwwpsiiiorrg/b.....
2208.00 01066
Mario Geiger, Arthur Jacot, Stefano Spigler, Franck Gabriel, Levent Sagun, Stéphane d' A'Acolli, Giulio Biroli, Clément Hongler, and Matthicu Wyart. Scaling description of generalization with number of parameters in deep learning. Journal of Statistical Mechanics: Theory and Experi- URL
ment, 20200) February 2020. ISSN 1742-5468 doi:
http://dx. doi. wrw/o0.gg/..... 0080/08//// 1742- 58 5468/ab5333
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of dis- crete gradient dynamics in linear neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances Neu- ral Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL httpss:/pproceedinggt wwwrips.einccmmmm cc/paper. http:/papr_fffeeeeeeeeeeeeee fiwesf F39e99ffaaa1111992232221222212 http://pppr...lleee///////////// ptsss//err//2///llee
Paper. .ddf]
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, 2017. htttss
Jordan Hoffmman,, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen. Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optime large language models, 2022. URL https: // /arxiviv.v. ooo.org//aa. 22003. 15556
12

---

Published as a conference paper at ICLR 2025
Marcus Hutter. Learning curve 2021.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks, 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, 2020. URL http::::: http:://aaaxii..ooraa///0010.1.0
Shuai Li, Zhao Song, Yu Xia, Tong Yu, and Tianyi Zhou. The closeness of in-context learning and weight shifting for softmax regression, 2023a. httptpp:://axx..vrrgraaaaab/
13276
Yingcong Li, M. Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in-context learning, 2023b. Jttppss::///
org h0s/abs/23011 07067
Yue M. Lu, Mary I. Letey, Jacob A. Zavatone-Vette Anindita Maiti, and Cengiz Pehlevan. Asymp- totic theory of in-context learning by linear attention, 2024. URL httpps://arxv..rrr
abs/2405. 11751
Alexander Maloney, Daniel A. Roberts, and James Sully. A solvable model of neural scaling laws, 2022.
16859
Meta. The llama 3 herd of 22
Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quantization model of neural scaling. In seirth- Conference Neural Information Processing Systems, 2023, URL httpss//openrevie..rrt net/ wtt//oorummmm =orum?2dd3:bbl122a333
3tbTw2ga8K
Yoonsoo Nam, Nayara Fonseca, Scok Hycong Lee, Chris Mingard, and Ard A. Louis. An exactly solvable model for emergence and scaling laws, 2024. URL https: wwwxiv.org/aa..... abs/ 2404 17663
Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington. 4+3 phases of compute- optimal neural scaling laws, 2024. / wtt.o/g//bb///ss//ee 15074
Hannah Pinson, Joeri Lenaerts, and Vincent Ginis. Linear cnns discover the statistical structure of the dataset using only the most dominant frequencies, 2023. URL https::/arxiv..or
abs/ fas/2303. 221/2333..022 02034
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gener- alization beyond overfitting on small algorithmic datasets, 2022. URL httpt/:sss//aii....///
abs/ abs/2201. 02177
Jonathan S. Rosenfeld, Amir Rosenfeld, Yonatan Belinkov, and Nir Shavit. A constructive prediction prediccioo of the generalization error across scales, 2019. URL http:///wwxwxww.g/a///o///ss
12673
Jonathan S. Rosenfeld, Jonathan Frankle, Michael Carbin, and Nir Shavit. On the predictability of pruning across scales, 2021.
abs/ hat///aab//0...... 201.:.000...100 10621
Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dy- namics of learning in deep linear neural networks, 2014. WRLLhttss://axxt org/ abs/
1.312.6221
Utkarsh Sharma and Jared Kaplan. Scaling laws from the data manifold dimension. J. Mach. Learn. Res., 23(1), January 2022. ISSN 1532-4435.
James B. Simon, Madeline Dickens, Dhruva Karkada, and Michael R. DeWeese. The eigenlearning framework: A conservation law perspective on kernel regression and wide neural networks, 2023.
URL h /ttps://aaxiie http:///rwwiv..or//////10/.0
03922
13

---

Published as a conference paper at ICLR 2025
Ben Sorscher, Robert Geirhos. Shashank Shekhar, Surya Ganguli, and Ari Morcos. Be- yond neural scaling laws: beating power law scaling via data pruning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems, volume 35, pp. 19523-19536. Curran Associates, Inc., nwwrips cc httppppprrrfflleppppa/e//// Conference. File/ 7 Ftle://7aabbblaaaaaa335eee77o7f0 Paper org.reeeeeeectaaaaaaooooereeeeer
2022. URL
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent on separable data, 2024. URL htt.pps://exxii.orr///e///
1710. 10345
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 doi: 10. 10.008//142--448/
Experiment, 2020(12):124001, December 2020. ISSN 1742-5468.
abc6ld. /tree//////ttt////a.oooooooo///r -ttp:///10.1110//222-55555333330
Ingo Steinwart, Don Hush, and Clint Scovel. Optimal rates for regularized least squares regression.
In Proceedings of the 22nd Annual Conference on Learning Theory, 01 2009.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines, 2024. URL httlhttt///xxxiii org/ abs abwaaaasa55 /2008 16898
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention all you need. Luxburg, S. Bengio, H. Wallach, Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL httpss:/pproeedinns www.:ss...e.urisscocco......e cc/papr files/paper/2017/ File/ Fale/3ffee24333 (el 3al1://5eee244447700000001000003 PaperaaaPppp...... pdf
Guyon, U. Von
Vito Volterra. Variations and Fluctuations of the Number of Individuals in Animal Species living
together. ICES Journal of Marine Science, 3(1):3-51, 04 1928. ISSN 1054-3139, doi: 10.1093/
icesjms/3.1.3. URL wwt.httpss.// //dtip:///oo..... org 10. 1 093/1093330 3
Johannes Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, 2023. URLtttts. http//aaxvvvvv. www.v..ogg/aab.......ao abs 221.22222222 222.200222277
Alexander Wei, Wei Hu, and Jacob Steinhardt. More than toy: Random matrix models predict how real-world neural representations generalize, 2022. URL
wrg.ivvorrraaae abs/
1203.06176
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
hatt::://axxxv.orrr/bb//// 003
Ruiqi Zhang, Spencer Frei, and Peter L.. Bartlet. Trained transformers learn linear models in- context, 2023. URL https: /ttp///aivvvv//////////ee//aaav wwtg/ab///aa 09927
14

---

Published as a conference paper at ICLR 2025
APPENDIX
The structure of the Appendix is as follows.
In Appendix
A
we discuss the in-context learnability of the MSFR setup.
In Appendix
B
introduce additional related works.
In Appendix
present key definitions from the main paper for convenience.
In Appendix
we establish the in-context learning dynamics as non-linear ODE systems.
In Appendix we solve the ODE systems to the zero-th order of under Assumption 3.1
Appendix we focus on neural scaling laws.
In Appendix we list details for our numerical experiments.
In Appendix
further diseuss the generality of the MSFR problem and our solution by connecting the MSFR problem with other types tasks.
In Appendix we present supplement numerical experiments to extend our claims on neu- ral scaling laws for linear self-attention to softmax self-attention.
In Appendix we construct the complete solution of self-attention for the MSFR problem up the first order of
A IN-CONTEXT LEARNABILITY OF LINEAR SELF-ATTENTION IN THE SPARSE FEATURE REGRESSION
We provide simple and straightforward example to explain how the model can adapt itself for different input sequences. For simplicity, we only consider one task 8, then the prediction for given query token 0ss, .:) and training sequence ((8, X) of the self-attention at oc. We further assume common constraint for the linear self-attention as in /Zhnn et al. (2023); Wu et al. (2024); Lu et al. (2024): the first N₂ components of v, and the last component w. i.c., +s are zero for any & 2 0. These simplifications enable us to write the output for the task
Sa U (w:) (xx) (o....
As /s = A, at = oo because the model is perfectly trained, we must have the following relation
Us ( w: (oo) 1/tb..
What will the prediction the self-attention for the query token Ō(s, be given a different input sequence (s, X)? Below we use our example to illustrate that the perfectly trained linear self- attention will adapt itself for (s, X). In particular, let the task strength for the input tokens in (ss X) prior to the query token be A. * A, (the task strength is unseen in the training data). Then the prediction of the self-attention for this input sequence ( X) becomes
o." (o0) w. (.,)))
(18)
(o))
which indicates that linear self-attention successfully predicts the unseen task strength Åo by per- forming in-context learning to explore in-context data to adapt itself to a new task.
In addition, namics, hence our method for solving the dynamics Section 3.2 can still be applied here even without Assumption3..] i.c., the solution obtained here will be exact one. To see this point more clearly, we rewrite the output of the model with the constraints as
show that the above newly imposed constraint does not affect the learning dy-
f. = v. H. w. Us +1 +1th ApWss
where v, Us:V,+1 and w Wss are I-dimensional vectors while H, A, is 1 1 matrix that satisfies the idempotent- condition. Thus the dynamics will be exactly the same as Eq. ! and Theorem 3.1 can still be applied.
15

---

Published as a conference paper at ICLR 2025
Finally, we briefly explain the reason why we impose these constraints. For the task and the input sequence (ess with task strength As, the output of the model is
1 " = t/ Wa: + (M + =1sss + 41:
We note that the first term does not depend on A, explicitly, meaning that it does not explore the relationship between ((8; x:) and y, thus it does not contribute to the in-context learning task. In ad- dition, the last term depends on A, in a non-linear way, leading to an inappropriate dependence on the relationship between æ and y. Due to such terms, without any conditions, the linear self-attention hard perfectly adapt itself for new input sequences ((8, X) with different task strengths. This problem is caused by its structure. A possible way to achieve a perfect in-context adaption is impos- (the constraints used in our example and ZZhng et al. (2023); Wu et al. (2024); Lu et al. 2024 ) to eliminate the aforementioned terms. Therefore, we expect that an optimal in-context learner should set the first NV components as O, since they contribute to many
ing the constraints Vs:s W Nx:N
terms that are not necessary for the model to achieve perfect in-context adaption.
B ADDITIONAL RELATED WORKS
Due to the great success of LLMs, the understanding for the transformer architecture becomes in- creasingly important. As Garg et al. (2023) proposed the in-context learning for learning particular functions, besides Zhang et al. (2023 von Oswald et al. 2023);(Akyine et al. (2023) studied trans- formers for the linear regression problem, while Li et al. (2023a) revealed the similarity between a single-layer self-attention and gradient descent on regression problem with softmax. Edelman
et al. (2022) showed that the self-attention is able to learn sparse functions of the input sequence. For the loss landscape of single-layer transformer without softmax, |Ahn et al.] (2023a) showed the existence of solution of the model parameters that can achieve perfect test loss. Furthermore, the idea of MSFR is also partly inspired by the multitask parity problem proposed by Michaud et al. 2023 Barak et al. (2022) that empirically exhibits neural scaling laws.
For other neural networks, the study of the learning dynamics is always an important topic. There is line of research that investigated properties of the converged parameters in both classification and regression settings, e.g.. [Soudry et al. (2024). To characterize the whole training trajectory, Saxe et al. (2014) built an exact solution of deep linear networks that depends on time explicitly. In addition, various tools borrowed from random matrix theory and statistical physics have also been applied to investigate the learning dynamics of linear models (Spigler et al. 2020 Bordelon & Pehlevan 2022, Simon et al.] 2023: Bahri et al. 2024). There are also a fruitful results (Adlam & Pennington 2020; d' Ascoli et al.) 2020 Geiger al. (020) that analyzed the learning dynamics of neural networks in the limiting case with the neural tangent kernel (NTK) (Jacot et al. 2020).
B. ADDITION RELATED WORKS ON LEARNING DYNAMICS OF DEEP NEURAL NETWORKS
Besides the aforementioned works on learning dynamics of neural networks or random feature mod- els. Pinson et al. (2023) studied the learning dynamics of gradient descent for linear convolution neural networks. Particularly, they discovered an interesting interplay between the data structure and network structure that determines the phases of the network along the training trajectory. The learning dynamics is also analyzed by Gidel et al. (2019) while Braun et al. (2022): Atanasov et al. (021 focused different initialization regimes. Our focus in this paper is the learning dynamics of linear self-attention with the in-context learning.
B.2 ADDITIONAL RELATED WORKS ON NEURAL SCALING LAWS
; Hoffmann et al. (002 there are a number of recent works that ex- plored scaling laws in deep neural networks empirically (Rosenfeld et al.] 2021 Hestness et al.) 2017 Rosenfeld et al.) 2019). The study neural scaling laws can be found in some earlier works (Caponnetto & De Vito. 2007 Steinwart et al. 2009 Ahmad & Tesauro 1988). From the theoretical perspective, various works developed solvable models in the context of random feature models (Bahri et al. 22224] Atanasov et al. 2022 2024 Bordelon et al.] 2024 Paquette et al.) 2024) to study the neural scaling laws in certain limits. In addition, Wei et al. (2022); [ordelon et al.
Besides Kaplan et 2020):
220T): Sharma & Kaplan 2022 Bordelon & Pehlevan (0022) also conducted theoretical analysis
16

---

Published as a conference paper at ICLR 2025
on linear models to investigate the neural scaling laws. These works improve the theoretical under- standing for the neural scaling law to a large extent. As a comparison, our focus in this paper is particularly on the linear self-attention with the in-context learning, which is not widely discussed in previous works.
B.3 ADDITIONAL RELATED WORKS ON ANALYSIS OF IN-CONTEXT LEARNING REGRESSION
The success of the transformer architecture Vaswani al. 2017) has encouraged majority body of works to investigate its theoretical understanding, especially the intriguing in-context learning mechanism. common setup along this direction the study of linear regression using the linear self-attention (Duraisamy 2024 Ahn et al. 2023b; Wu et al.) 2024 Lu et al.
2024, Zhang et al.)
2023). Our work also falls into this category as mentioned in Section 2.3
We present a more
extensive comparison and connection to Zhang et al. (2023)
and Lu et al. 2024
in the following.
Comparison and connection to Zhang et al. 2023). Zhang et al. (023) considered linear regres- sion using linear self-attention in the in-context learning manner. By assuming the infinite training dataset size limit, Zhang et al. (2023) revealed that the converged linear self-attention can achieve competitive performance compared to the best linear predictor over the test data distribution. As a comparison, our setting is also a regression for the linear self-attention to learn in an in-context manner, while our problem is multitask version and the distribution of the task type obeys a power law. In addition, besides the converged solution of the dynamics, we derive its (approximate) form along the whole training trajectory, which in turn makes possible characterize the neural scaling laws with respect to time, data size, model size, and the optimal compute. And we note that the characterization for the time scaling law (and the optimal compute law) cannot be derived solely by the converged solution.
Comparison and connection to Lu et al. (2024). Lu al. (2024) proposed a solvable model of in-context learning for a linear regression task by linear self-attention. Specifically, assuming limit where the input dimension, the context sequence length, the training task diversity, and the data size are all taken to infinity following certain ratios, they revealed double-descent learning curve with respect to the number of examples. As a comparison, our problem is also a multitask regression for the linear self-attention to learn in in-context way. In addition, as will explain in Appendix H. our problem can be seen as a limiting case of the multitask in-context regression under the source-capacity condition, which is also generalized version of the setup considered in Lu et al. (2024). Furthermore, our solution (to the first order of e, when the context length is large) captures the whole training trajectory and study the neural scaling laws with respect to various parameters so that we do not assume that all parameters are taken to infinity together.
C NOTATIONS AND DEFINITIONS
We present useful definitions in the main paper here for convenience.
Feature Extractor Eq.
((8,0) Rd x RRR {-1,0, E ERM,
Target Eq. 33)
v(5,2)) A Σ (s. x))
An in-context data point Eq.
(4)
= T0(8 (1))
((1))
((s,5))
(ss ==
y
y
(bba)
0
Lu(s, x())
y(s ((v)) 0
Empirical loss function
N
Σ
s (0(s()),)())))))))
xx()):0) - =
2N
1=
17

---

Published as a conference paper at ICLR 2025
Notation
Definition
Index
Pa x
distribution of task type data vector
Eq. Q
Page
((s, x))
feature extractor target
Eq.
Eq. (3
Ws
context sequence for task task strength for task in-context data point empirical MSE loss number of data for task
Page
Ag
Page 3
(s, X)
Eq.
L
Eq. (6 Eq.
8
#
Eq.
H.
(s,, (,, x)7
Eq.
7
e,
standard basis vector in IR
Page
1/1/s
Page
fs
output linear self-attention for the task
Page
W
Page
ps
U. w.
Page
9.
ng H,
Eq. d0
hg
p HaPs
Eq. Io
Ho and H
decomposition of H. zero-th and first order solution of f. #. (A3 1)/N
Page
f and f!
Page
ag and he
Eq. (12)
zero-th order terms of and h, first order terms of gs and h, tD, optimal compute
Page
g and h
Page
C
Page
Table 3: Notation and definition with the corresponding index
9, and h
ys 17. -H, h2 P. H.P
In-context learning dynamics Eq. (9)
#e (S. - AA) H,W,, w. N
# (fo - Ag) H.U.
(19)
N
Test loss Eq. (15
1 N
L(t) = Ex P.
Σ
Pa(S= A.J².
(20)
2
To navigate the paper, we also present table for notation and definitions with the corresponding index in Table R
D IN-CONTEXT LEARNING DYNAMICS
In D.I derive simpler form of the empirical loss function Eq. 8 In D.2 we present details of the learning dynamics Eq.
D.1 FORMULATION OF EMPIRICAL LOSS FUNCTION
By definition, the output of self-attention is
f ( X):0) = "V
WKO 1.
where W, V are defined by
W = (w)
WNN, E R(N++I)XN, ViE [I.....N.:: : E R
V (v) U..) WWG vi (1....,A.) : RW+++
18

---

Published as a conference paper at ICLR 2025
Deriving H.. To derive H, (Eq. 77. we decompose the in-context data point ((s, X) (Eq. 44 for the task as
P
$
(ss X) =
(21)
y
0
where
(1)
ys
Dluaa)
Ps Φ
ys
E
(22)
(4%)
Li/e
and Φ. = (s, Then
((s,
P3
$
P
yo
y
P.P.T + 6.
P.y=
(23)
y P
y ys
There are four terms for to compute to get the form of H, in Eq.
(23
(i) the first one is
(1) P.P. = o
1 2 diag(e.) Φ
(24)
(
where the standard basis vector in TR-V is C. = (0 J.....,,,) such that the only nonzero component of e, its s-th component; (ii) the second one is Φ. = diag(ee), which can be easily verified; (iii) the third one
0 . E R for E
(1)
U
P.y. = (600)
(1))
B (1) = A.A.es
Lys
(iv) the final one is v.y. y. = Eg
(y) 21:
A3 Combining these terms gives us the form of H.:
H. = [diag ((1h, + 1)e,)
As a result, the output of self- afffettention for the in-context data point ((((n) x()) will be
f ( x(n))) vz (m)) W 2..0,
which gives us the empirical loss as
2
1
Σ
f (0()((),((1))) x(x))) -s(n)]
2N
n=1
N
Σ (s(()) ((())) 2
No
gln) s(m) W Ag(n)
H
2N
n=1
k=
H.(n) Wg(m) Ac(mi] (s(()( (2(
((()))
2N
( H W060 Ascw]] ]
2N
(v' H.w. A,)2
(25)
2
N
where we use the definition of empirical loss in a, the definition of target y in b, the decomposition
of W Eq. (5) in c, Φb .) _ according definition of & Eq. (2) in and recall that # denotes the number of in-context data points with (x) for (1...., N) in
19

---

Published as a conference paper at ICLR 2025
D.2 IN-CONTEXT LEARNING DYNAMICS AS ODE SYSTEMS
We adopt the continuous time limit of gradient descent, i.e., gradient flow, to perform the empirical
loss minimization
v= VvVVv W),
W V L(V, W).
(26)
Specifically, for empirical loss function Eq. ( we can directly obtain the learning dynamics as non-linear ODE systems: Vs E No.
#s
(fa Ao) H.W.
N
#a
0. N
(Js A₂) H.U.
where recall that we denote fa = v², H.W..
E SOLUTION OF SELF- ATTENTION
In E.1 we derive the ODEs of 9, and hg, while solve the ODEs to the zero-th and first order of 6s under Assumption 3.1 and respectively.
E.1 ODES OF Is AND he
There are four steps to derive ODEs for and ha which will be discussed as follows.
Step I: change of variable. As discussed earlier, the ODE of v, Eq. (9) is non-linear with respect w. thus transform more symmetrical form by the change of variable: let
7 Us Wss P. Us Ws.
then the dynamics of ns NN+11 and Ps E IR can be obtained according Eq. (
# 9+ h. Ag Ho P.
#
ge (ge-h
n1 N
N (44 Ag Hap.
(27)
as a result, the ODE 77 is non-linear with respect to ne while that pa is non con-linear to Ps.
Step I:: deriving ODEs to different orders of ex.
According to the definition of H, Eq.
can rewrite H. sum of two matrices
H. th H ..H))
where
H diag(ee)
A,e.
[dagg(e;) e,)
...
A, A,e"
A2
and 6, 1/4. 1 given Assumption 3.1 Therefore
treat eH!. insignificant pertur-
bation the dynamics Eq. (27) Let the solutions of Eq. (27pbe be
7 + +
= p- +e
such that n. and p are treated as perturbations to n and p.. respectively. Then, according to the definitions of gs and h2, we can also write Is and h, in the perturbed form
Is = 6e + +xal
ns H.n.
v. (n. + + (H. + +eeH.) (n. c.m))
v. n a + the (n H.n + +2n H + +((2)
(28)
h- h2 + calh p. H.P:
v.p2 H P. + UhEa (p) H.P. 2p H.P.) + O
(29)
20

---

Published as a conference paper at ICLR 2025
Putting the above perturbation forms back to Eq.
(27 to the first order of c.. we have
c.
[0_
h1 + Ca (g: h1:) 4A,]
H n + ca H nh + H + O(2))
4N
P. +
[9.
h2 E. (9) hit)) 4A.]
H p² cs (H H + Oe
4N
(30)
Matching both sides of Eq. (30p the zero-th and first order of e, respectively gives us
v. #
t
(ge Nº 4A.) H
p..
(ge 11 4A.) HaD!
(31)
4N
4N
and
vo#t
[992 - - - 4A.) H + HDn!) + (g. H²
4N
(32)
8 th
4N [(g] ho -4.)) ( p% + H.) + (g. -hh) H.) P.]
Step III: deriving ODEs for gs and h, to zero-th and first orders of e.. We can now obtain the ODEs for 93 and h2 to the zero-th order of e, by directly applying the definitions and Eq. (31)
d 41
#o mm H H
5 v. dl
Ho
-b, (g) 1: 4A,)
2N
H no
- (g.
h 4A,)
Zero-th Order:
2
(33)
ag9o
(o" 1o 4A,)
2
aeh2
h (0th hg - 4A.)
2
where we use (H (A3 + H) in and derive the equation for ho in similar way. Suppose that we obtain the solution to the zero-th order by solving the above ODEs, then, according to the definition of g in Eq. (28) and the definition of h! in Eq. 29). only need to find the solutions of
m. = n.. H and n, = v.p! H
(34)
to derive the solution to the first order of c.. since we can obtain 7/ H.m and pe H2PP using the solutions of Eq, (see Appendix J.T). This means that based on Eq. (32) we need to solve the following ODEs:
m.
Ven H + v.n. H
(0 11 1A.) [nd
H H + n H + H)))
4N
First Order:
V.#.
(g) 11 H.
4N
(g ho 4A,) a
2 #
(g] h1) 20gh
H H n3
4N m
4
(35)
and, similarly,
h. = p3: H 00
+ 13 HOPO
322
(0) h 4Ao) [p. H + - H HHH!P!
4N
First Order:
2 #
+ 4N
(gs
N3 HO HD P3
22 4N
agho
= (0 no 4Aa)
+
Ps
H
+ (gg - 1122
2
4
(36)
Then we can write the solution of 9. and 113 as
9 2/. 17. H 17. + 2m,, h = wh w. H. 2n.-
We will solve the zero-th th order ODEs in the next section and discuss the solution of the first order ODEs in D
21

---

Published as a conference paper at ICLR 2025
E.2 SOLUTION OF SELF-ATTENTION ZERO-TH ORDER
Eq. 33 is still hard to solve. Fortunately, we note that
dt g2 = (0. - h2 4AA) as (9 - 12 0,
d
(37)
2
which immediately implies that
4t 2 0: 0. = 2C,
(38)
a constant, i.e., h2hhg is conserved quantity of the learning dynamics. This important quantity allows us rewrite Eq. (33 as
9 20,A.g a. + C.
(39)
ig = 20,A.02 a. (2()))// + a.C.s
which are exactly the Riccati equation. To determine Cs, we note that
-n² (p' H.) (10)) H.. + (10) 40 H, (0), (40)
where we use ||a f a² Ha for positive definite matrix H and vector and we note that C, (0) - w (0). In the following, we adopt a series of change variable to solve Eq. 33). For simplicity, omit the subseript and recover it in the final solution.
Step I. Let
ag
ah
p
9
(41)
then we can transform Eq.
39 to
ag
2aAp + +2² a²C
2
2
(42)
0 ah
2²C
2aAg + q²
Step II.
Let -1/7,9
-0/0, then
17
72
+ p²
(43)
222
?
Putting back to Eq. (42).
obtain the ODE of and 0 follows
8. 2aA (- 7
a²Or
a²0r
2aA2
0,
2
2
(44)
a²Co
Ø+ 2aAo
0.
2
Step III. Eq. (44 are just the second-orde linear ODEs, which can be solved following a standard approach. Specifically, let
7 ad 0 bt
then putting them back into Eq.
44 gives us
a²C
V 44 Ca
2 2Aac
=0 0 = c=
2
2
(45)
aCC
24 + 20²C
8 2Aab
=0
b -aa +
2
For ease notation, we denote
V + +22C
S A f=
04 - + 0_ = 5
(46)
2
22

---

Published as a conference paper at ICLR 2025
then solutions of and are
7 = Ae + Be-
0 E Ee- + Fe -
(47)
where A, B, E, and F are constants that need to determined by the initial condition. Now we are ready to recover the solution to Eq.
33
2p
2 7
20 Ar + -e-
9
=
ay
1e + Bee
(48)
qa
0
20_ o- + 1 Fe-
h
2
of
Ee Fe-
Recall that = [g - h) in Eq.
(7).. we are now able to derive the solution the zero-th order
solution of self-attention:
0 h
4
1
0 Ae + Bo_e 0 B. Ee to,Fe
20
A + Bea-t
E Fe-
1
0+ +o_ + B(o-- -04e-
Flo, 0 --)c² c
2a
At + Be
Eca + Fe
F
B
A
+
(49)
2
Ecast +
Acade + B
where we define
= v 4A + 2C,
the last equality.
Step IV. It is now left for to determine A, B, E and F according to the initial condition, which can be obtained by noting the conserved quantity gh in Eq. 38 and can be satisfied when
0 ++ + _o_e- 0 Ee att + Fo ot
2C
(50)
1e + Be
Ee-o t Pe-oit
aCC
of + BB
(AF + BE) = O. 2
(51)
Noticing that
0 S2 25 + 22
0² = 2SE + 5²
can further simplify Eq.
51 to
AFo = BE._.
(52)
On the other hand, considering the initial condition by letting in Eq. 49
A 2
1
f(0) -
(53)
E
A/B +1
if we denote
E
Q 0_
2A X
D 2(f(0) - A)
P
(54)
F'
0+
2A +
then easily see that A/B F F and Fq.
53
becomes
D (1 + PQ) + P P-PQQ 0
P = [D( + 1) +1 Q1 + / D(Q 1)+1-0+1-0000 [)(++)++++@)) 4D2Q
(55)
2DQ
Recall that g, h according to their definitions (H positive-definite), we only take the minus sign in the above solution, which can be simplified by conducting some tedious algebra as
P 4f(0)A 2C v 2C /C)
(56)
2(ʃ(0) 4A + 2C) 2A
23

---

Published as a conference paper at ICLR 2025
We summarize the solution of self-attention to the zero-th order of cg to prove Theorem 3.1 by recovering the subscript @ in Eq. (49)
with solution of P in Eq. (56):
X (2()) As
1
(57)
1 Ps exp(ag),()
1 Qs exp(a.Xat)
where 02 # + 1)/N, and
A = 443 + 2C,, P. 4f2(0)A, + 2C, + A + +2C, 2( A.)(A, 2A,)
2A, X
Q, = 2A, + A
E.3 ZERO-TH ORDER SOLUTION FOR SPECIAL BALANCED INITIALIZATION
When (0) (0), we have C₂ 0, which implies that F a_ 0 according to Eq. (46) and
Eq. 51
and
AS = 2A.
As result, we can rewrite the solution as
A (2 = A, 1 A/Be
A exp( - -2agdat)
(58)
A. 1.-1"10) (2(0)
1
where we use the initial condition Eq.
53 in the second equality.
F NEURAL SCALING LAWS OF SELF-ATTENTION
We first present the overall procedure for deriving neural scaling laws in F. then discuss them in detail for fixed context sequence length in F2 and for varied context length in F3 We use a to mean that is approximately equal to b (by neglecting irrelevant coefficients and constants) if a, E R.
PROCEDURE
For convenience, we first present the test loss described in the main paper
1 Nº.
L(t)
> Pa(==s) [t,]] 2
(59)
Note that as a, is determined by the dataset and it satisfies the following relation
U t. 1 (S = s)t. (A² + "
ag
(60)
N
as the number of training data points N + co. Furthermore, as long as #. + such that a, * 0, crucial property of (t) is
lim (o) = Ag. t-00
(61)
otherwise it would be
Ag lim - (2() = A.
1
T
(62)
2 1+P.
which means that self atffnnttento cannot learn the task if there no data point for it in the training set and is similar to the one shot learner property of diagonal linear networks in Nam et al. (2024). These properties will repeatedly applied the following sections. And we now discuss procedures for deriving different neural scaling laws.
24

---

Published as a conference paper at ICLR 2025
Size of the model D. To quantify the model size D, assume that there is cutoff D such that the model cannot learn any task strength As with $ 2 D. As a result, according to the solution for the task in Eq. (57), we have
vt 02 (2(t) A, = (2) - A..
(63)
When D is the bottleneck of training, we let t, N (they are sufficient for the training) derive scaling laws respect to D. Thus the over all test loss will be
[f) - A. N, 2 2s
L(D)
(64)
To derive the neural scaling law of model size D, we only need to find the asymptotic behavior of Eq. 664 by letting N₁ and replacing the summation with an integral.
Training time 4. Training time is equivalent to the number of optimization steps. To investigate the scaling law respect t, we remove the bottleneck caused by the size model and the number of data points by letting N + oo (thus we have Eq. (60)) and D = NN (thus Eq. (63) does not satisfy for any D E No,). As a result, the overall test loss will be
Zs (t) Aa]² ds L(t) 2 2==
(65)
where we let N₁ -+ oo and replace the summation with integral. We sill study the asymptotic behavior of Eq. (65) with the Laplace method (Bender & Orszag 1978) to investigate the time scaling law.
Number of training data points N. When the training is bottlenecked by N, we let oo (thus Eq. 61 will be satisfied for all task types s if there exist training data points for them otherwise Eq. would be satisfied) and the cutoff D Ns (thus Eq. (63) does not satisfy for any [I.....NG}. As result, we conclude that the probability of lim (t = (2) is exactly the same as the probability that the training data set ( { does not have any training data point for the task i.e, Vn E (1...., N): s()) * s. Therefore, we can rewrite the test loss as
L(N)
Z (0) - J, - P(S = s))
2 s=1
(66)
Z ()) A.1² (1-2x-9)
6=1
where, again, we let N, + 00 and replace the summation with the integral, and we sill study the asymptotic behavior of Eq. (65) with the Laplace method to investigate the data scaling law.
Optimal compute C. This is the case when the number of data points is sufficient for the training (N - Doo), while training time or the size model D is the bottleneck given the compute budget C Dt such that either or D scales differently with C. Specifically, if
Ttt. D) = agt at aa -a
then we can rewrite the loss
L(D) = =00)=ac-o Dar + aan
(67)
To obtain the optimal loss given D, we let
aD a,0,,
oetoD
(aDDD
agtup
0pL(D) =0 D
Cartan,
Caston,
As result, we can derive the optimal compute budget test loss
LCC) (=aoop/(octap)
(68)
given C tD, where a; & can obtained from the neural scaling laws for time and model size, respectively.
25

---

Published as a conference paper at ICLR 2025
F.2 NEURAL SCALING LAWS WITH FIXED SEQUENCE LENGTH AND STRENGTH
In this case, according to Eq. (6), the test loss can be written as
((t) ~ ZA N
+ Pexp - + 1]t]
1+ Q exp a X + 1)d]
(69)
In the following, we will investigate neural scaling laws using the above test loss and the procedures described in F.
Model Scaling Law. the behavior
According to Eq.
64), the model scaling law can be obtained from studying
Zs-a 12
1
1 } ds D D
L(D)
Zs
(70)
J==D
1+P
1 1+Q
where D is the cuttoff of the task such that our model will only learn the first D - 1 tasks and we let / +
N,, + oc and replace the summation with integral in the test loss, we have (omitting irrelevant coefficients and we denote vú (A? + 1) for ease of notation)
Time scaling law.
L(t) ~ S, ast aett +
1
)
2a " a In ds
(71)
nát Q
Now let
F(s) = s " + Ins
then, applying the Laplace method (Bender & Orszag 1978), for large
L(t) x
F(()
rte
-F(())+""(( ds
e F(ojt lete c (c) do
c-c
V2T
b - F(ojt
(72)
/c)(c)t
where we expand F(s) around its minimal F(c) in and b is simply the Gauss integral. We can solve (c) 0 to determine the value of c:
ac
a
= c= (t
(73)
ct
This further gives
In(
-(e)t 1
F(c) -1
E
(74)
exit
and
(c(c) = a(a + 1) a- a
= a(o + 1))b(1(0)
(a+2)
= a 2 c -11
(75)
Putting F(c) and (c) obtained above back Eq. immediately gives the time scaling law:
V2T 1ta-
T(I)
x e-1
(76)
co
26

---

Published as a conference paper at ICLR 2025
Data Scaling Law. In this case the bottleneck of training is the number of data while t - and D = No. According to Eq. (66). the test loss has the form of
1
]
L(N)
Zs
(1 Zs a)
8
1+P
1+Q
e Maaaan/NNNNii " ds.
(77)
We apply the Laplace method again to study the asymptotic behavior of L(N) to derive the data scaling law, Let
F(s) = a In N
In(1 Zss "))"),
(78)
then can expand F(s) around its minimal F(c) where the value of i determined by F*(c) 0:
(*(s) a a7.
(79)
Ns 1 Z
c= ((N + 1)2)
(80)
The loss function can be written as
2'T
L(N) - F(O)N
(81)
Fc)
where
F(c) In((N + 1)2)
In
N
N +1
(N(NZ) + N-1
(82)
where assume that N > 1 in the second line. It is now left for us find the value of F" (c), which is done as follows.
==)())== -c a(N +1)
ac
a Zac
N
1 - Zc
C c(1 ac-
-2 ( Zc
1 Z
+
1- Zc
(1 - Zc N
a(a
a(a 1)% (at22/aaa
where we use N > again in the last line. Putting F(c) and (c) back to Eq. scaling law with respect to N:
81
gives us the
22T L(N) NZe Z- Vala+ 1)NN ~
(83)
F.3 NEURAL SCALING LAWS WITH VARIED SEQUENCE LENGTH AND STRENGTH
In general, derive neural scaling laws with similar spirit as in the previous section. Ad- ditionally, we assume for simplicity that A? > 1 such that a, # A2/N and the model is
initialized (0) (0) and (0) = O(1) for all As result, the test loss Eq. (17)
N.
cxp( -2, 2
Σ L(f) X 2 8=1
22
(84)
+ A exp( -2a,A.t)
which can be casily derived using Theorem 3.1 and v. oc s - A, x 7 and we also initialize
the model such that the initial prediction of the model is equally away from the strength for different tasks to exclude influence from other aspects, i.e., A,/ (0) are similar for all s. As we assume that A 1 in Section 4.2 we can rewrite a, when N oo
0g -2
(85)
where use 7 to denote irrelevant normalization constants.
27

---

Published as a conference paper at ICLR 2025
Model scaling law. According to Eq. (64). the model scaling law can be obtained from studying the behavior
A
L(D)
Zs
ds D D-a-2y+1
(86)
s=
|+
where D is the cuttoff of the task such that our model will only learn the first D - 1 tasks and we let t + 0o.
Time scaling law. With similar procedure as in previous section, we will apply the Laplace method to derive the time scaling law. Specifically,
exp [- aaAs-a---30 3-0333 + at27 In
L(t)
ds
(87)
(1+ + De 2,,2,
where we use 2 to absorb all irrelevant constants. Now we let
+ 2y In 8,
F(s) --s-------- 8--8-37
(88)
then the asymptotic behaviors of L(t) can be written
2m
L(t) ~ P(e)a
(89)
(c)tt)t
where Fc(c) = as before. Note that the first derivative of F(s) is
F'((s = -Ā((a + 8 3 (a+ +
a a 27
(90)
st
42 0
c:= (M
(91)
+ 27
Therefore, we obtain that at s =
F(c) = a + 27
(1+ (I t-
(92)
+ B + 3y
F(c)t
042457
(93)
(eM
Furthermore, the second derivative of F(s) with respect to 8 is
F* = M(a 27) 3 By + 1)s (a+8+3++))
a+ 27
(94)
821.
which gives us
(c) = M² 04/133 कमा.अअ (a + 2y) (a + 37).
(95)
Putting F(c) and F" (c) back to Eq. 87) gives the time scaling law
L(I) x 31
(96)
Data Scaling Law. Similar to previous section, this case the bottleneck of training N and we let f - oo and D = N₁. According to Eq. (66), the test loss will become
L(N) x
27 Zs a3
c (( In/N - dls.
(97)
Following a similar procedure, we let
F(x) = (a + 27) Ins
In(1 aa
(98)
N
28

---

Published as a conference paper at ICLR 2025
and we let Fc(c) = 0. Then the loss function can be written
22T
L(N) _ P(c)t
(99)
F (c)
To obtain e, we need to compute
a + 27
abs a0--a-1
Z(N +1+ 11222//a)
('(s)
c= 78-
(100)
sN
1 1427/a
Note that we need to make sure
( + + 27/a) > 1 + 27/a.
(101)
which can be easily satisfied, to apply the Laplace method. By conducting some algebra, we obtain the following results:
a + 27 a.VV
Z(N +1+ 23)
1 1+
F(c)
14
27
N N+1+
2y &
o+29
1 1+ I NM
6 F(c)N
where use N >> 1 in the second line to study the asymptotic behavior. Furthermore, the second derivative is
(c) -c- -2a(N +1
ac
c- 2y Z 2n
Zac
1 7c-a
c(1 - 7c-a)2 N
Z
ac 2
+
(102)
1- Zc
(1-Zc ac-
aN
which, noting that
7c 1+ 2x
(103)
N
according to the solution of Eq. (100) gives us
(c) x N
(104)
Combining these results gives the complete data scaling law
L(N) N
(105)
G NUMERICAL. EXPERIMENT DETAILS
For all numerical experiments, we generate the dataset exactly as the process described in Sec- tion 2.2 The model structure is linear self-attention as specified in Section 2.3 If not specified, we set the initialization as
v, (0) = A x 1M++1 w.) = v (0) + 0.1 x A 1
(106)
where = 0.1 is a constant and we use 1 E Rdd to represent a vector with all elements equal to 1. For the discrete GD training, we the learning rate as 10-3 and the number of total optimization steps as 5000. The theoretical prediction using the solution f²(t) is simulated with the forward Euler method such that kn where is the optimization step and n is the learning rate.
Fig.1
N
100 for both 3 and 100.
Fig.1
AC 500. We set the context sequence length 100, and the task strength A 0.5.
29

---

Published as a conference paper at ICLR 2025
Fig. 4
In this case, set N,, 500, and set the initialization
(0) (, A x 1N,+1
(107)
where A 10-5.
H GENERALITY OF THE MSFR PROBLEM
In this section, we demonstrate the generality of the MSFR problem. Specifically, we show that our method can also be applied (or is limiting case of) other generalized types of the MSFR setup considered in Section 2. namely multitask in-context regression under the source-capacity condition (Appendix H..I). MSFR with approximately sparse feature (Appendix [H.2). and tasks
with idempotent-ln H. (Appendix H..) We also highlight that our solution can be applied to
study other properties of attention besides the neural scaling laws considered in Section and
H.1 MULTITASK IN-CONTEXT REGRESSION UNDER THE SOURCE-CAPACITI CONDITION
Interestingly, the MSFR problem can seen as a limiting case of the multitask version of the context regression under the source-capacity condition (Cui et al. 2022), which is defined as follows and can be seen as generalization of the setup in Lu et al. (2024).
Multitask in-context regression under the source-capacity condition. Following the settings of the MSFR problem in Section 2.1 there are NN different tasks in total. We let S be the ran- dom variable of picking a specific task among N,, tasks and assume that S follows the power law will show that each task is constructed as in-context regression, thus we term this setting as multitask in-context regression. We do not use the sparse feature extractor (bss, a) defined in Eq. 2 Instead, following Cui et al. (2022), use the feature extractor (ss,, x) R.V. such that for data æ Px
distribution Eq. ①1. In the following,
Σ. = ExpPx [5(s ( ())()) )
= diag (o3)) = diag (I, (w))
(108)
We satisfies the source/capacity condition
where ~h -
"x. R. and Wi
w2 WNr|" R.V, such that
Wk x k kT
(109)
and P, is a simple rearrangement of elements of w such that
a
0 W2
W) WAG-o+1
(110)
0f = WNN--+2
a WW
i... the s-th eigenvalue of Eg is the largest given task type s. Finally, given task type we let the strength for task be A, N, and the target y E IR is
y(s x) A, Ō(s x).
(111)
The in-context regression data $(8, X) now generated according to the process in Section 2.2
((8 X) = (5s5 (1))
((s,)()))) ((.))
o(s, x)
(112)
(1))
0
while assume that the sequence length t/. is fixed for each task
30

---

Published as a conference paper at ICLR 2025
MSFR is a limit of multitask in-context regression under source-capacity condition. We note that, a large т in Eq. (I09) indicates that the spectrum the covariance matrix Σ, shows a very fast decay. If r is large enough, we can conclude that only the largest eigenvalue of Eg (i.c., w.) is significant, thus
(6(5 ((5, (1),)))((s)))))
lunge ((5,2))(s.x))
Σ3 = Egwpy
(113)
where ((s, x)) is defined in Eq. (2). Since BewPx (6(5 2)(,,)) and ((s, x)" deter- mine the in-context learning dynamics (because they are the main component of H, according to Eq. (23) and H. captures the in-context learning dynamics Eq. ). we conclude that the MSFR problem can be seen as limiting case for the multitask in-context regression under source-capacity condition with large т. Therefore, we expect that the neural scaling laws derived in the MSFR problem can be generalized to the multitask in-context regression under source-capacity condition.
Numerical Experiments. To validate the above claims, we conduct numerical experiments to investigate neural scaling laws of softmax attention and we use the parameterization W&W,, Wq rather than merging them single matrix W WKQ. For the feature (fss, 2),
x(,, (O, E.)
(114)
where (0, Σ.) is a Gaussian distribution with zero mean and covariance E, diag((*) as in Eq. (08). We let = 3 in Eq. (199.. To make Σ;= Φ(s, x()) x(1))"// close enough to E. we let 4/. = 1000 for all The task strength A, NO0, 7). The softmax attention is trained by GD with learning rate 2 10-1 We report the neural scaling laws with respect to time and with respect the optimal compute in Fig. where
clearly see that the softmax self-attention for the multitask in-context regression under the source/capacity condition also displays neural scaling laws similar to those derived in the MSFR problem (Table These numerical experiments support our claims and the generality of the MSFR problem.
18
cl 2.
107
5:5
10
25-3
10'
10
10'
10'
(a) Time Law with N + oo, D = N,
(b) Optimal Compute with N + 00, C = Dt
Figure 5: Neural scaling laws for softmax self-attention in the multitask in-context regression under the source-capacity condition. In each figure, we use solid lines to represent empirical simulation results and dashed lines for power law curves. In (b), we set a = 1.8.
H.2 MSFR WITH APPROXIMATELY SPARSE FEATURE
In the MSFR problem in Section 2.1 we let the feature extractor be
( IR x R + (-1,0, e RR
(115)
such that the feature (ss, x) ) is sparse. In fact, solution to the zero-th order of cg, i.e., (t), can still be exact under Assumption when the above sparsity condition is relaxed for large sequence length th.. We give example follows.
31

---

Published as a conference paper at ICLR 2025
MSFR with approximately sparse feature. For the MSFR problem in Section 2.1 now sider new feature extractor ((s, x) such that
Φ(8. ((s, + ((),
(116)
where ((s, æ)) R.V, can be random noise to the first order of c, (Ç does not need to be sparse). We call this task MSFR with approximately sparse feature. This task will give us the same set of non-linear ODEs to the zero-th order of 6, under Assumption 3.1 as that for the original MSFR problem in Section 2.1 Therefore, Theorem 3.1 can still be applied in this case
Numerical Experiments. In Fig. 6 we let Ç (O 2.T) be Gaussian noise vector for each task s. We compare the loss calculated according to f (t) in Theorem 3.1 with that obtained from empirical simulation. can be seen that our theoretical prediction is still highly exact with the existence of the noise vector Ç when the context sequence length va is large.
0.00
2.00
1.75
0.025
1.25
0.020
1.00
L(e).
0.015
0.010
0.25
0.05
0.00
100
10
202
10
Figure 6: Loss L(t) of MSFR with approximately sparse feature for different context sequence lengths 4111 and 100) during training. Solid lines are for theoretical predictions while dashed lines are for empirical simulations.
H.3 GENERAL. TASKS WITH IDEMPOTENT-LIKE H₂
From mathematical perspective. besides the MSFR problem considered in this paper, strat- egy for solving the ODEs Eq. (9) can be applied to any cases when the matrix H, in Eq. (9) is idempotent-like without Assumption 3.1
H H.
(117)
where 14g E IR is a constant. In such cases, the solution of the model prediction still (t) in Theorem 3.1 except for that we now define a = #. #.H.M.NN and (t) is exact do not need Assumption 3.1 We think it will be an interesting future direction to explore other tasks (for self-attention other machine learning models) where H, has the idempotent-like structure.
Numerical Experiments. To verify the above claim, we consider a simple example where (we omit the subscript and consider the case where only have one type of task)
3
H = Σ Eauuy us Uj = E7
(118)
and we let E, 2 for i 1, 3, which will give μ 2 in Eq. 117). The learning dynamics is Eq. with # N. In Fig. 7 we compare the loss calculated according to the solution in Theorem 3.1 with that obtained from empirical simulation. It can be that theoretical prediction matches with the empirical simulation well because is an exact solution in this case.
We that our characterization of (f) in Appendix iis no longer applicable this case, and we believe the characterization of f! (t) and the generalization of our methods more complicate feature extractor Φ(s. be an interesting future direction.
32

---

Published as a conference paper at ICLR 2025
10-1
L(E)
10a2
Theory Simulation
10a
Figure Loss L(t) for learning dynamics Eq. (9) with idempotent-like H.
I ADDITIONAL NUMERICAL EXPERIMENTS
In this section, we conduct additional experiments to explore the generality of our conclusion for the
neural scaling laws. In particular, in Appendix L. we explore the neural scaling laws of softmax
self-attention for the MSFR problem where we train the model with GD, while in Appendix 1.2 we train the model with AdamW.
1.1 NEURAL SCALING LAWS OF SOFTMAX SELF-ATTENTION FOR MSFR PROBLEM
We replace the linear self-attention with the softmax self-attention in the numerical experiments of Fig. 3 and Fig. to investigate the neural scaling laws for the MSFR problem. For complete- ness, we adopt the W Wo decomposition rather than a single merged WKQ. i.c.. f(G; 0) VG softmax G" W K W
All the other settings are the same as those of Section 4.1 and 4.2
Fixed Context Sequence Length. In Fig. g we report the neural scaling laws when the context sequence length is fixed as in Section 4.1 It can seen that the scaling laws with respect to time t, model size D, data size N, and the optimal compute C' are similar to those reported in Table [
Varied Context Sequence Length. For the varied context sequence length, we let th = F(s) o - 8 as in Section 4.2 while we keep A, fixed. We note that the neural scaling laws with respect to the model size D and data size N are not affected by varied context sequence length as reflected in Table 2 which is due to the fact that GD can still learn the task strength A, for the task s as t ++ oo when the context sequence length is varied. We report the scaling laws with respect to time t in Fig. where we can see that the softmax self-attention still admits similar time scaling law compared to the linear self-attention for varied context sequence length. As result, the optimal compute scaling law of softmax self-attention will also be similar to that of linear self-attention, as it is consequence of the time scaling law and model size scaling law and these laws do not change.
These numerical experiments reveal that our claims regarding neural scaling laws for the linear self-attention be generalized to the softmax self-attention.
I.2 NEURAL SCALING LAWS SOFTMAX SELF-ATTENTIO TRAINED ADAMW
To examine the effects of optimization algorithms on neural scaling laws in the MSFR problem, train softmax self-attention with AdamW and we also use the WT parameterization. We focus on the case where the context sequence length and the task strength are fixed. We present our parameters in the following table.
33

---

Published as a conference paper at ICLR 2025
g-1d g-t
2e18
ger g
107
D
(a) Model Size Law with N, f + 0c
(b) Time Law with N + oc, D = No
Nite cm13
10-
L00
LICO
107
100
107
(c) Data Size Law with L co, D = N₂
(d) Optimal Compute with N - 0o, C = Dt
Figure Neural scaling laws for softmax self-attention trained by GD with different values of 1.8, 2.1 when the context sequence length is fixed. In each figure, use solid lines to represent empirical simulation results and dashed lines for power law curves. In (d), we set a = 1.8.
Simulation iile
t- 100
10a
Li0
6x00--
40-0
3x10
10a
10a
t
Figure 9: Neural scaling laws with respect to time for softmax self-attention trained by GD with 1.8 when the context sequence length is fixed. We let N oc, N,, Solid lines represent empirical simulation results while dashed lines represent power law curves obtained from Table 2
Neural scaling laws with respect to model size D and data size N. We expect that Adam W will show similar neural scaling laws with respect to the model size D and data size N when compared to GD. This is because AdamW W can still learn the task strength A, for the task given sufficient training time (Fig. 10c), which is similar GD. We report the corresponding neural scaling laws
34

---

Published as a conference paper at ICLR 2025
learning rate 1)
10
B
0.9
82
0,999
weight decay
10-5
eps
10-8
Table 4: Parameters for Adam W
in Fig. 10a and 10b where it can be seen that the softmax self-attention trained by Adam W still admits similar neural scaling laws with respect D and N.
Neural scaling law with respect to time t. However, AdamW typically exhibits very different dynamics during training compared to GD, as AdamW has very different learning dynamics (e.g.. it converges faster than GD). Thus we expect that AdamW will lead to a very different time scaling law (Fig. 10c). which will further lead to a different neural scaling law for the optimal compute (Fig. TOd). We additionally note that these observations are similar to the observations in Hoff- mann et al. (2022), where the authors revealed that, when compared to Adam, AdamW shows different test loss behavior against the optimization steps (training time), indicating that the type of optimization algorithm can affect the time scaling laws.
L(N)
12-4
105
13a
102
102
12a
12
D
N
Model Size Law with N,
(b) Data Size Law with t xc, D N.
10-
LGC)
30
131
10'
124
(c) Time Law with N oo, D No
(d) Optimal Compute with N 00, Dt
Figure 10: Neural scaling laws for softmax self eeffaattenionn trained by AdamW with different values of a 1.8, 2.1. each figure, we use solid lines to represent empirical simulation results and dashed lines for power law curves that are obtained Table [ (when the self-attention is trained by GD). In (d), we set a = 2.1.
35

---

Published as a conference paper at ICLR 2025
J COMPLETE SOLUTION OF SELF-ATTENTIA UP TO THE FIRST ORDER
We first discuss how to derive the solution model parameters (1) and w to the zero-th order of exs then present the solution to the first order of 6x. which can give us the complete solution of self-attention up to the first order of eg. In the following sections, we omit the subscript 8 for convenience and recover it in the final solution.
J.1 ZERO-TH ORDER SOLUTION OF MODEL PARAMETERS
E.2 gives us the solution of the model output (2( To obtaín the forms of (1) and ()), we only need to solve 17 and p since
170 + p (t)
170 p
w w((t)
according to their definitions. The ODEs of n7 and pa Eq. 31 can be rewritten using the compo- nents (note that v/ is the s-th component of v7²) as:
#
(6) nt - 44) (1) + A
4N
NN - A #
(g) - - - 4A) (nh + A
4N
(119)
8% V#
(0 1 4A) (0 + =Ap%++1
4N
PN.+1 = A #
0
(g) n 44) (0_ - AAPh+1))
4N
An interesting property of these ODEs is that
dt (Ar - n = = Anp 78 = C, Aph - p& =,
(120)
which gives relation between 1/ and n and similar between pa and P On the other hand, since we already know the solution of gd from E.2 and 9d can be written as
go (v An
(121)
we can solve 17 and p based on these relations, which will also give us v and w
J.2 SOLUTION UP TO THE FIRST ORDER
We now discuss the solution the first-order of cg. According to the definition of of and h! in Eq. 28 and (29). they can be rewritten as
g = 1/2 '1100 + 2m
(122)
h vpp H pd + 2n
where m and n follow the dynamics Eq. 335) and Eq. (36). Therefore, to obtain the formulations of g¹ and h², the ODEs of m and will be the only equations that need to be solved since we can obtain 17" '' directly from In the following, we focus on how to solve m and n. For convenience, we first present ODEs for m and n ( (Eq. (35) and Eq. (36) without the subscript s)
i = - (g
h
4A)
am
#
HH 17
- (g ) ag
2
4N
4
(123)
à = (g)
no 4A)
an
2
#
p
H H pm
+ (g]
1: a
2
4N
4
36

---

Published as a conference paper at ICLR 2025
The above equations too complex, thus attempt to reformulate them to simpler forms: can obtain a new set of ODEs from the above equations
d dt
mhi 0 i vihh"
#V²
H T - (g' -h') ag
-0 - h1" 4A) TO
(124)
4N
4
di n9 nn + igg
#
p H H p'g + (g' - 11') hg -
= = - 1100 -4A
(125)
AN
which implies that
d
#b² H H nh p² H''"'
(mhh + ngg) -(g" 12t 4A) di
(n
(126)
4N
The above equation gives relation between mh! and ng Fortunately, according the defi- nitions gb. h", H and H (Eq. (7), we can expand the terms inside the second bracket of the above equation:
n¹ H' h p² H = o [o + APp - (ni + AN
-2CC(CPA,+H - Cn(s+1)
(127)
where we use Eq. (20) in the second equality. the model is initialized as C = Ö 0, then under this condition, we can immediately conclude that
Vt 2 0: m/h" = -gg + Ō di (mh 9 + ng =
(128)
where Ö' is determined by the initial condition and we let Ô = 0 in the following. Noting that g¹ h := r + 2(m n) = (n 'n110 - pp H p²) + 2(m n), allows to simplify the ODEs for m and h further by interchangeably using mh! and
(129)
Eq. 128
ng
in
am
9
4A
+
ng"
ag r
2
(g)
-
14) H
2
m
4
4N
0g
22 #
-am (g) 2A)
(g) - 11
4A)
n¹ HH H²
(130)
4
4N
ah
ñ on (h"" + 2A)
+
(g
h 4A)p² H H p
(131)
4N
Eq. 130 and 131 exactly solvable since they simply first order linear ODEs. Specifically,
let
e(t) aa (g)
-2A).
((t) rr
vb2 # 4N (9 - h 4A) n. H'''m
(132)
then Eq. (30 can rewritten as
m o(t)m + Q(t).
(133)
The standard procedure for solving this letting á = uu and multiplying u to both sides of Eq. (133), then obtain
di um ud um = / u(t))(t)dt const. = m I (u(t)o(t)dt const.
d
(134)
Similarly, to solve n, we let
<(t) = (he 2AA
(()) ah O
+ 1/2 #
(g
-
11 4 H H po
(135)
4
4N
= zE(t)
37

---

Published as a conference paper at ICLR 2025
then
((t)p((t)dt + const.
n
(136)
2
As a result, the solution of self-attention up to the first order e, for the task type s under Assump- tion 3.1 now becomes
f 90 1
Solution:
S [u(T)((T)dd + const. [=((T)((F))T + const. )
+ ++
(137)
J.3 FORMULATION OF THE FIRST ORDER SOLUTION
We examine each term in Eq. ho can be written explicitly as
[17 in the following. Since we let C = C = 0 in Eq.
120). g and
1o = (AA + 1)) g (12 + 1)?()))?
(138)
Form of r. Note that r is defined as
(/) H n"
+ [(1]]] (D2)a] p². D H²
90 - 11
(139)
(A2 + 1)2
Form of and ((())(FF)dr. We first derive u:
ù uu = u= = exp (-) e(t)dt)
(140)
where
/ e( = / algo - -AAdtt = 22At + f 1
2aatt + 2lm
(141)
where we use Eq. 48 in the second line. Putting the above integral back the expression of obtain
(e 22 = (Ac
+ Be be==))
(142)
where { = 2 + C/2 defined in Eq. (46). We now derive ʃ widt. To start, we examine
each term of first:
ag"
((g0)2 2C'
4(A2 1)2
1/2 #
110 -4A) n H H²m² (g
4(A2 + 1)N 0 110 - 4A) gd
4N
(143)
(2 2C 4Ag9)
4(A2 + 1)2
(())) - 2C 2Ag
2(A² + 1)2
Using this in the integral and considering the form of u in Eq. 142). we have
f wddt
/ (Actt + Be BE [[ 20 2Ag" dt
2(A² + 1)22
2
(AbOEE + Bo_e =51 dt (A2+1) 1)2
aC
/CACS Be st)ddd
a(A² + 1)2
2A
/ (AGSI + Be St) (A + Bo_ (t) dt
+ (A + 1)2
(144)
where we frequently use the solution of 9d in Eq. 33
38

---

Published as a conference paper at ICLR 2025
Form of and f =(T)))())dT By the similar procedure of deriving and f uddt, can also derive z and f z.dl..
exp (ʃ Edit = ()) = (Eeft + Fe)
(145)
where we use Eq. follows:
48 in the second equality. Similar to the derivation of 0 we can derive yp as
r = ((h9)2 2C)
4(A2 + 1)2
2 #
#
(g)
110- - -44) pa H H H'
4(A2 + 1)N (g 110 - 4A) no
4N
(146)
a (() 20 + 4AA")
4(A2 + 1)2
p= ((h)22 2C 2A4
2(A² + 1)2
thus
/ zppit
(EeSt + Fe Fe (t)2 [010]- 20 ] dt
=2(22 + 1)2
2
/ (Eo (Eo_eetFoooee FFo4e dt + (A2+ + 1)2 / ((BOS++ + Fe
aC
a(A3 + 1)²
2A
/ (Be& Et (Eo_ Poye-811 dl,
+ (A² +1
(147)
where we frequently use the solution of h1 in Eq. 33
Results of integrals. It is now left for us to solve all the integrals to obtain the complete solution. We list the results below.
J [ + Be- (o+ Ae + Bo - -
A
B
251
-25t + 2AB
251
2
25
[ + FFF Fole-flddt
E2o
F2o+
25t
-2Et
2Et + Aat
25
2
/ (o_Ac + e -
A202
B20²
2Et
2Et + -- ABI
2
25
| (o_Echt -
E²0²
F2
251
2Et + +20+0_ER
25
2
3.
1 (Ac8++ Be -e
A2
B²
28t
-5t
2ABt
25
2
f (BES Fe 2t
E²
F2
25t
2Et
2EFt.
2
2
Complete solution. With these integrals, we now ready to find the explicit forms of the solu- tion Eq. (17).. In particular, we have
f wudt E(A2 + 1)2a
25
(-02 + 2
a²C
ahoo.)
B 25t
a²C
o2
aAo
E(A2+ 1)2a
2
ABt
24A2a²+ + +2²0 - -4000-
(A2 + 1)2a
A2E281A A
(
Aa
B2E-25LA
Au
4aABl
1
(1
(A2 + 1)2 (A2 +C) (148)
(A² + 1)²
{
(A2 + 1)²
39

---

Published as a conference paper at ICLR 2025
and
E2
-0- + d²C
F2e-2EE 25t
a²C
zipdt
+ aAo_
+
0 2
ado+)
E(A2 + +1)aa
E(A² + 1)2a
EFF
( + C2aC -0-0--
(A² + 1)
E2E2ErA A
(1 (1- 5
Aa
(A2 +13 (1 (1+ Aa
F2
4aEFt
(A2 + 1)2 (A² C).
(149)
(A² + 1)²
These equations are sufficient for to find m and n. For m have
-2
(1+ Ae & + 2 (-1+ An) 1aQt (A2 + C)
m7 I Juddt
(150)
(A² + () te e-t2
where Q has already be determined in Theorem
For we have
22 (1 5
26t 1
4aPt (A2 C)
f zppt
(151)
71
(12 + ()2 + 5-51)2
where has already be determined in Theorem 3.1 Finally, by using the solved m and n above and (Eq. (139) in Eq. [17 and recovering the subscript in all relevant terms, obtain the complete solution of self-attention up to the first order of under Assumption 3.1 Note that as - oc, we can 22A/(A² 1)² and r 4A/(A2 1)?, thus g¹ -h' 0.
easily verify that in Eq. (129)
As result, (1) + e. ()) A, as desired.
40


Reviews:
[{"review_id": "LVkuPxwP7Q", "reviewer": "Reviewer_FJbQ", "summary": "The paper performs a theoretical study on the problem of in-context learning for attention models, where a number of $(x,y)$ pairs compose the input string to the model. Compared to a standard transformer, the paper simplifies the architecture by analyzing a single self-attention layer. Thus, the learnable parameters are a matrix corresponding to the merged queries a keys' parameters, and the value's parameter matrix. The data distribution is composed of a power-law distribution over a number of tasks, a sparse feature extractor, and a \"task strength\" parameter that weights the sparse feature extractor to generate the labels. The sequence length and the task strength also follow a power law distribution. \n\nThe paper shows that the training dynamics of this model can be solved in closed form, assuming the sequence length is sufficiently large for all tasks, and keeping the zeroth order term of the expansion at large sequence length. The authors then use these closed-form solutions to get the test loss and devise neural scaling laws in various scenarios by varying the variables (model size, time, number of data points). They also cover the compute-optimal setting.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The exposition of the theory is overall clear. I particularly appreciated the \"Procedure sketch\", which guides the reader through the proof technique, from the non-tractability of the original problem to perturbation analysis and the several changes of variables involved. Also, I appreciate the structure and attention to detail of the appendix. It can be seen that the authors put considerable effort into making the proofs clear and well-organized.\n\n2. The results in Section 4.1. recover well-known scaling laws previously observed both empirically and theoretically on simplified models. The results in Section 4.2 are more specific to the in-context setting and the toy model proposed and seem entirely new to me. \n\n3. The theory is supported by a sufficient suite of experiments.", "weaknesses": "In many cases, it is unclear which set of assumptions are due to reasonable empirical observations in in-context learning and neural scaling laws, and which ones are mainly there to obtain tractable calculations (and if so, why so). For instance:\n\n1. Sparsity. Why the feature extractor can only be {-1, 0, 1}. No justification is provided as to why this should be realistic. In fact, it seems precisely constructed to have the expansion of $H$ as the sum of two matrices, with the second one having the sequence length as a multiplicative prefactor. \n\n2. Why should the task strength follow a power law distribution?\n\n3. Why the distribution of the tasks should also be a power law?\n\nI would appreciate it if the authors stated the purpose of the assumptions more explicitly. More broadly, this task seems specifically designed to get closed-form solutions and it is thus unclear how they can be re-used in future works.\n\nAnother issue is that the analysis is the *sequential* scaling of large context length limit, and then the other quantities (time, number of samples, model size). Assumption 3.1 on large context length allows the authors to use perturbation analysis to get the zero-th order correction of the expansion at large context length (which results in Theorem 3.1), and essentially drop all the higher order terms. However, in Section 4.2 they study the *joint* behavior of sequence length with the other scaling quantities. How do the authors know that other terms of the expansion would not be relevant if the joint limit were taken? \n\nOther issues:\n1. $f^0_s(t)$, $f^1_s(t)$ is defined in Theorem 3.1 but referred to earlier (line 303). \n2. Why notation changes from d to D in Section 4.\n3. Why approx in Eq. 15. \n4. I feel this paper on the asymptotic theory of in-context learning should be cited [1].\n\n[1] Asymptotic theory of in-context learning by linear attention (https://arxiv.org/abs/2405.11751)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9iZW8FQ3LO", "reviewer": "Reviewer_5pWj", "summary": "This paper presents a theoretical analysis of power laws in in-context learning using Transformers. Specifically, it introduces a task setup called multitask sparse feature regression and approximates the learning dynamics of self-attention trained by gradient descent on this setup through perturbation analysis. The authors verify that the test loss follows scaling laws with respect to training time, model size, the amount of training data, and the optimal compute budget.", "questions": "- $f_s^0$ depends on the training data, so is it ok not to take the expectation over the training data in the definition of Test Loss (eq. (14))?\n- In Section 2.2 on the generation of in-context data, is it correct to assume that the data sequence $\\boldsymbol{x}^{(1)},\\dots,\\boldsymbol{x}^{(\\psi_s)}$ is sampled independently and identically from $\\mathcal{P}_X$?\n- Apologies if mistaken, but in eq. (31), is $\\frac{\\psi_s \\\\#_s }{N}$ missing? Wouldn’t this result in $a_s = \\psi_s^2 \\\\#_s (1+\\Lambda_s^2) / N$?\n- When transforming the summation to an integral in eq. (62), is $f_s^0$ defined for non-integer $s \\in \\mathbb{R} \\setminus \\mathbb{N}$? Similarly, does $\\mathcal{P}(S=s)$ in eq. (64) have the same consideration?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- This paper demonstrates a commendable effort in addressing the challenging theoretical analysis of neural scaling laws in in-context learning, which is both complex and impactful.\n- It provides a visually clear figure explaining the multitask sparse feature regression task, along with an accessible proof outline that enhances reader comprehension.\n- The paper establishes an interesting connection between in-context learning dynamics and the Riccati equation. It is particularly intriguing that $g_s^0h_s^0$ emerges as a conserved quantity.", "weaknesses": "- **The problem setup appears somewhat artificial.** While the gradient descent analysis for multitask sparse feature regression is impressive, the problem setting itself seems contrived and may lack applicability. For instance:\n  + The probability of task selection $\\mathcal{P}(S=s)$, task strength $\\Lambda_s$, and context sequence length $\\psi_s$ are all exponentially proportional to $s$. The assumption that these three factors are coordinated (e.g., larger $\\mathcal{P}(S=s)$ implies larger $\\Lambda_s$ and $\\psi_s$) seems limiting for practical applications.\n  + In this multitask sparse feature regression task, only $\\phi(s,\\hat{\\boldsymbol{x}})$ is necessary to predict $\\hat{y}$, making $\\phi(s,\\boldsymbol{x}^{(1)}),\\dots,\\phi(s,\\boldsymbol{x}^{(\\psi_s)})$ effectively irrelevant. This raises doubt as to whether this setup truly qualifies as in-context learning. A token-wise feed-forward network would likely suffice, giving the impression that self-attention has been unnecessarily added to complicate the problem.\n- **Several typos and ambiguities in presentation.** A few observations are listed below:\n  + In eq. (2), $\\mathbb{R}^d \\times \\mathbb{R}$ should be $\\mathbb{R} \\times \\mathbb{R}^d$.\n  + In Section 2.2, $f:\\mathbb{R}^{(\\mathcal{N}_S+1) \\times (\\psi_s+1)} \\mapsto \\mathbb{R}$,(page 4, line 182), the $s$ is incorrectly capitalized, and $\\mapsto$ should be $\\to$.\n  + In Section 2.3, $\\mathbb{R}^{\\mathcal{N}_s \\times (\\psi_s+1)}$ should be $\\mathbb{R}^{(\\mathcal{N}_s+1) \\times (\\psi_s+1)}$.\n  + Below eq. (7), the phrase “does not depend on $n$, where ...” is misleading, as eq. (7) depends on the task type $s^{(n)}$; rephrasing this may clarify the intent.\n  + In the second line of eq. (26), it seems that $\\mathcal{O}(\\epsilon_s^2)$ is unnecessary, whereas it may be needed in eq. (28).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QEBJPzqMVs", "reviewer": "Reviewer_7YzX", "summary": "The paper establishes a theoretical framework to study the gradient flow learning dynamics of a single linear attention layer on a synthetic multi-task regression task.\nIt provides closed-form solutions for the dominant ODE dynamics based on perturbation theory and uses them to derive neural scaling laws w.r.t. training time, data set size, model size, and compute budget.", "questions": "Please see the questions from W1. Additional minor questions:\n\n4. L387: Why does the compute not scale with $N$, the data set size?\n   I believe that computing a gradient should scale linearly in $N$.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- **S1: Presentation & Contribution** Although the paper's content is relatively dense and introduces a lot of symbols, the presentation is very clear and, in my opinion, strikes a very good balance between mathematical detail and prose explanations.\n  Specifically, I appreciated the author's clear outlines before executing the mathematical steps, which helped me understand many of the details, and convinced me that the contribution is non-trivial.\n\n- **S2: Generality** I believe the synthetic multi-task sparse regression task, as well as the derived solutions of the dominant ODE dynamics can be of interest to study other phenomena in transformers, not only neural scaling laws.\n  Therefore, the presented framework seems useful for future works.", "weaknesses": "The following are mostly minor concerns which, if addressed, would strengthen my confidence.\n\n- **W1 (minor): Gap to 'practical' attention** The paper studies a relatively simplistic attention layer without softmax, skip connections, MLPs, depth, etc.\n  I think this is okay, because the contribution is to derive closed-form expressions for the dynamics, which is likely intractable when including these components.\n  It would be interesting to further investigate how useful the theoretical predictions are in a setting that more closely resembles a practitioner's setup.\n  Specifically:\n  1. Theoretically: How does the parameterization $W_K^\\top W_Q$ (used in practise), rather than $W_{KQ}$, affect the dynamics, i.e. how would the presented results change?\n  2. Empirically: How well does the theoretical prediction (linear attention) match the simulation results in Figs. 3/4 when using softmax attention?\n  3. Empirically: Many attention-based models are trained with AdamW. How well does the theoretical prediction (gradient flow) match the simulation results in Figs. 3/4 when training with AdamW+small learning rate?\n\n- **W2 (minor): presentation/editing suggestions**\n  - It would be great if the authors could provide some practical recommendations based on the theoretical insights into neural scaling laws in Section 4.\n  - It would be great if the authors could highlight more connections and differences with related works throughout the text, e.g. other works studying the learning dynamics of MLPs and CNNs (the author may want to cite https://proceedings.mlr.press/v202/pinson23a/pinson23a.pdf), or contrasting the neural scaling laws in Tables 1,2 with other empirical or theoretical results.\n  - In the paragraph above Theorem 3.1, could you highlight which set of equations have the same form as the Riccati equation?\n    I also believe it would be good to extract some of the currently inlined math into environments to facilitate looking up symbols, specifically on page 6.\n  - Typos:\n    - Line 108: Why not just use $a \\approx b$ to indicate approximately equal?\n      Is this notation used in the main text?\n    - Equation 2: I think it should be $\\mathbb{R} \\times \\mathbb{R}^d$\n    - Equation 7: Can you say what $H_s$ is?\n      It looks like it contains the second moments of the data and labels.\n    - L269: 'can be rather exact' sounds weird; do you mean 'approximate, tractable solution'\n    - L285: Maybe use '^T' instead of '\\dot' for consistency with L253.\n    - L297: $\\epsilon$ should be $\\epsilon_s$.\n    - L331: To me it sounds weird to say the solution is 'considerably exact'.\n      Maybe replace with 'is a good approximation'.\n    - L379/381: 'with respect' is missing 'to'\n    - Figs. 3d/4b/4f: Add a legend for each curve.\n    - L536: 'a variety properties' misses an 'of'", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Lzxt5Ky9XN", "reviewer": "Reviewer_9RQV", "summary": "This manuscript studies the dynamics of learning a particular structured in-context regression task with linear attention. Its main results are approximate exponents for the decay of error with time, data, and model size.\n\n**Edit following author response: I have raised my score from 5 to 8**", "questions": "- The discussion in lines 64-71 is not very accurate, as the cited paper of Bordelon et al studies scaling laws for *random feature models*, not neural networks that can learn features. \n\n- Why do you bother to write (6) in terms of a generic loss $\\ell$ rather than just writing the MSE, which is used throughout? This creates unnecessary overhead for the reader. \n\n- Unless I'm mistaken, all experiments use linear attention. The claims regarding architecture (in)-dependence of certain scalings would be much stronger if you could add experiments with softmax attention.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I found this to be a generally interesting paper, and its topic should be of broad interest. I do have some concerns that preclude my recommending publication in its current form, but I think these should be addressable.", "weaknesses": "Here I list some concerns; their order should not be ascribed any meaning. \n\n- The title is inaccurate and vague, since the authors are not the first to study linear attention, and their solutions are approximate. I would suggest switching to something more descriptive, and would strongly suggest mentioning \"linear attention\" in the title. This critique extends to the abstract, which is similarly vague. \n\n- In that vein, please refer to \"linear attention\" or \"softmax attention\" rather than calling everything \"self-attention\"; this will make the paper less confusing to read. \n\n- The review of models for neural scaling laws in Lines 44-51 is inadequate. The authors should cite the original work on source-capacity conditions in kernel regression by Caponnetto and de Vito (from 2007!), and more of the substantial literature on scaling laws in random feature models that precedes the paper of Maloney et al. Suitable references are reviewed, for instance, in the cited work of Bordelon et al. or in the expository work of Atanasov et al. (https://arxiv.org/abs/2405.00592). The authors should also cite Paquette et al (https://arxiv.org/abs/2405.15074), which is contemporaneous with Bordelon et al 2024. \n\n- The authors don't discuss the growing literature on in-context learning of linear regression, which is closely related to their analysis. First, a more extensive comparison to Zhang et al (which I mention is now published as https://jmlr.org/papers/v25/23-1042.html) is required. Next, it could be interesting to compare the results presented here to those of Lu et al https://arxiv.org/abs/2405.11751, who studied the asymptotic behavior of linear attention trained using ridge regression, i.e., to convergence. \n\n- The paper would be much improved if the authors could relate their data model to other tasks and to natural data. At the end of the day, you're doing gradient flow on the MSE, so at least in certain limits you should be able to replace the specific data model with a Gaussian covariate model of matched moments. Would you get the same scaling laws if you studied in-context regression under source-capacity conditions, i.e., the generalization of the task studied in Lu et al to structured covariates and targets? \n\n- The main analytical results of the paper are approximations, and those approximations aren't clearly organized. For instance, the statement of Theorem 3.1 references Assumption 3.1 rather than directly stating that its result is an expansion at large $\\psi_s$. It would be clearer to state this directly. Also, though the authors provide reasonably compelling numerical evidence that their approximations produce reasonable results, it would be better to make this mathematically precise, i.e., to give explicit error bounds. Otherwise, please state your results as Results rather than Theorems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)