In-Context Learning with Representations: Contextual Generalization of Trained Transformers
Authors: Tong Yang
Yu Huang
Yingbin Liang
Yuejie Chi
Venue: NeurIPS 2024 poster (2024)
Keywords: training dynamics
optimization
transformer
multi-head softmax attention
in-context learning
Primary area: optimization_for_deep_networks

Abstract: In-context learning (ICL) refers to a remarkable capability of pretrained large language models, which can learn a new task given a few examples during inference. However, theoretical understanding of ICL is largely under-explored, particularly whether transformers can be trained to generalize to unseen examples in a prompt, which will require the model to acquire contextual knowledge of the prompt for generalization. This paper investigates the training dynamics of  transformers by gradient descent through the lens of non-linear regression tasks. The contextual generalization here can be attained via  learning  the template function for each task in-context, where all template functions lie in a linear space with $m$ basis functions. We analyze the training dynamics of one-layer multi-head transformers to {in-contextly} predict unlabeled inputs given partially labeled prompts, where the labels contain Gaussian noise and the number of examples in each prompt  are not sufficient to determine the template. Under mild assumptions, we show that the training loss for a one-layer multi-head transformer converges linearly to a global minimum. Moreover, the transformer effectively learns to perform ridge regression over the basis functions. To our knowledge, this study is the first provable demonstration that transformers can learn contextual (i.e., template) information to generalize to both unseen examples and tasks when prompts contain only a small number of query-answer pairs.

In-Context Learning with Representations Contextual Generalization of Trained Transformers
Tong Yang* CMU
Yu Huang UPenn
Yingbin Liang! OSU
Yuejie Chil CMU
Abstract
In-context learning (ICL) refers to a remarkable capability of pretrained large language models, which can learn However, theoretical understanding of ICL is largely under-explored, particularly whether transformers can be trained to generalize to unseen examples in prompt, which will require the model to acquire contextual knowledge of the prompt for generalization. This paper investigates the training dynamics of transformers by gradient descent through the lens of non-linear regression tasks. The contextual generalization here be attained via learning the template function for each task in-context, where all template functions lie in a linear space with m basis functions. We analyze the training dynamics of one-layer multi-head transformers to in-contextly predict unlabeled inputs given partially labeled prompts, where the labels contain Gaussian noise and the number of examples in each prompt are not sufficient to determine the template. Under mild assumptions, we show that the training loss for a one-layer multi-head transformer converges linearly to global minimum. Moreover, the transformer effectively learns to perform ridge regression over the basis functions. To our knowledge, this study the first provable demonstration that transformers can learn contextual (i.e., template) information to generalize to both unseen examples and tasks when prompts contain only small number of query-answer pairs.
task given few examples during inference.
Introduction
Transformers [Vaswani et al., 2017] have achieved tremendous successes in machine learning, particularly in natural language processing, by introducing self-attention mechanisms that enable models to capture long-range dependencies and contextualized representations. particular, these self-attention mechanisms endow transformers with remarkable in-context learning (ICL) capabilities, allowing them adapt to new tasks or domains by simply being prompted with few examples that demonstrate the desired behavior, without any explicit fine-tuning or updating of the model's s parameters [Brown et al., 2020].
A series of papers have empirically studied the underlying mechanisms behind in-context learning in transformer models [Garg al., 2022, Von Oswald et al., 2023, Wei al., 2023, Olsson et al., 2022, Xie et al., 2021, Chen and Zou, 2024, Agarwal et al., 2024], which have shown that transformers can predict unseen examples after being prompted on a few examples. The pioneering work of
Department of Electrical and Computer Engineering, Carnegie Mellon University; email:
tongyangeandrev. edu.
DDepartment of Statistics and Data Science, Wharton School, University of Pennsylvania; email:
yuh42@wharton. upenn. edu.
'Department Electrical and Computer Engineering, The Ohio State University; email:
1 Baang88890ouu edu.
$ Department of Electrical and Computer Engineering, Carnegie Mellon University; email: yuejiechi@emu.eer edu.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Garg et al. [2022] showed empirically that transformers can be trained from scratch to perform in-context learning of simple function classes, providing a theoretically tractable in-context learning framework. Following this well-established framework, several works have investigated various properties of in-context learning in transformers. For instance, studies have explored generalization and stability [Li et al., 2023], expressive power [Bai et al., 2024, Akyürek et al., 2022, Giannou et al., 2023], causal structures [Nichani et al., 2024, Edelman et al., 2024], statistical properties [Xie et al., 2021, Jeon et al., 2024], to name a few.
In particular, analysis from an optimization perspective can provide valuable insights into how these models acquire and apply knowledge that enable in in-context learning. A few works [Huang et al., 2023, Chen al., 2024, Li al., 2024, Nichani et al., 2024] thus studied the training dynamics of shallow transformers with softmax attention in order to in in-context learn simple tasks such as linear regression [Huang et al., 2023, Chen et al., 2024], binary classification tasks [Li et al., 2024], and causal graphs [Nichani et al., 2024]. Their theoretical analyses illuminated how transformers, given an arbitrary query token, learn to directly apply the answer corresponding to it from the query-answer pairs that appear in each prompt. Therefore, they all require the sequence length of each prompt to be large enough so that all query-answer pairs have been seen in each prompt with sufficiently high probability, whereas practical prompts are often too short to contain many query examples. This suggests that in-context learning can exploit inherent contextual information of the prompt to generalize to unseen examples, which further raise the following intriguing theoretical question:
How do transformers learn contextual information from more general function classes predict unseen examples given prompts that contain only partial examples?
Since our paper studies ICL of non-linear function regression, the function mapping (which we also term as "template") naturally serves as the "contextual information" that can be learned for generalization to unseen examples. When each prompt contains only small number of (noisy) examples, the template that generates the labels may be underdeternined, i.e., multiple templates could generate the same labels in the prompt. Such an issue of underdeterminatiion further raises series of intriguing questions, such as:
When the template that generates a prompt is underdeterminnd, what is the transformer's 's preference for choosing the template and how good is such choice?
1.1 Our contributions
In this paper, we answer the above questions by analyzing the training dynamics of a one-layer transformer with multi-head softmax attention through the lens of non- nin-linear regression tasks. In our setting. the template function for each task lies in the linear space formed by m nearly-arbitrary basis functions that capture representation (i.e., features) of data. Our goal is to provide insights on how transformers trained by gradient descent (GD) acquire template information from more general function classes to generalize to unseen examples and tasks when each prompt contains only a small number of query-answer pairs. We summarize our contributions are as follows.
We first establish the convergence guarantee of a one-layer transformer with multi-head softmax attention trained with gradient descent on general non-linear regression in-context learning tasks. We assume each prompt contains only a few (i.e., partial) examples with their Gaussian noisy labels, which are not sufficient to determine the template. Under mild assumptions, we establish that the training loss of the transformer converges at a linear rate. Moreover, by analyzing the limit point of the transformer parameters, we are able to uncover what information about the basic tasks the transformer extracts and memorizes during training in order to perform in-context prediction.
We then analyze the transformer's 's behavior at inference time after training, and show that the transformer chooses its generating template by performing ridge regression over the basis functions. We also provide the iteration complexity for pretraining the transformer to reach e-precision with respect to its choice of the template given an arbitrary prompt at inference time. We further compare the choice of the transformer and the best possible choice over the template class and characterize how the sequence length of each prompt influences the inference time performance of the model.

---

Under more realistic assumptions, our analysis framework allows us to overcome a handful of assumptions made in previous works such as large prompt length [Huang al., 2023, Chen et al., 2024, Li et al., 2024, Nichani et al., 2024], orthogonality of data [Huang et al., 2023, Chen et al., 2024, conditions [Chen et al., 2024], special structure of the transformer [Nichani et al., 2024], and mean-field models [Kim and Suzuki, 2024]. Further, the function classes we consider are generalization of those considered in most theoretical works [Huang et al., 2023, Chen et al., 2024, Li et al., 2024, Wu et al., 2023, Zhang et al., 2023a]. We also highlight the importance of multi-heat attention mechanism in this process.
et al., 2024, Nichani et al., 2024], restrictive initialization
To our best knowledge, this is the first work that analyzes how transformers learn contextual (i.e., template) information to generalize to unseen examples and tasks when prompts contain only a small number of query-answer pairs. Table provides a detailed comparison with existing theoretical works in terms of settings, training analysis and generalization of in-context learning.
nonlinear
multi head
task shift
GD convergence
noisy data
representation learning
Reference
attention
Wu al. [2023]
x
Zhang et al. [2023a]
X
X
/
EXX
XX
Huang al. [2023]
X
/
X
Li et al. [2024]
X
V
V
Chen al. [2024]
V
X
V
Kim and Suzuki [2024]
X
X
Ours
Table 1: Comparisons with existing theoretical works that study the learning dynamics transformers in ICL. Here, the last column refers to the fact that the response in the regression task is generated by a linearly weighted unknown representation (feature) model.
1.2 Related work
In-context learning. Recent research has investigated the theoretical underpinnings of transformers' ICL capabilities from diverse angles. For example, several works focus on explaining the in-context learning of transformers from a Bayesian perspective [Xie al., 2021, Ahuja et al., 2023, Han et al., 2023, Jiang, 2023, Wang et al., 2023, Wies et al., 2024, Zhang et al., 2023b, Jeon et al., 2024, Hahn and Goyal, 2023]. Li al. [2023] analyzed the generalization and stability of transformers' in-context learning. Focusing on the representation theory, Akyürek et al. [2022]. Bai et al. [2024] studied the expressive power of transformers on the linear regression task. Akyürek al. [2022] showed by construction that transformers can represent GD of ridge regression or the closed-form ridge regression solution. Bai al. [2024] extended Akyürek al. [2022] and showed that transformers can implement broad class of standard machine learning algorithms in-context. Dai et al. [2022], Von Oswald et [2023] showed transformers could in-context learn to perform GD.
More pertinent to our work, Guo et al. [2023] considered an ICL setting very similar to ours, where the label depends on the input through basis of possibly complex but fixed template functions, composed with a linear function that differs in each prompt. By construction, the optimal ICL algorithm first transforms the inputs by the representation function, and then performs linear ICL on top of the transformed dataset. Guo et al. [2023] showed the existence of transformers that approximately implement such algorithms, whereas our work from a different perspective, showing that (pre)training the transformer loss by GD will naturally yield solution with the aforementioned desirable property characterized in Guo al. [2023].
Training dynamies of transformers performing ICL. A line of work initiated by Garg al. [2022] aims to understand the ICL ability of transformers from optimization perspective. [Zhang et al., 2023a, Kim and Suzuki, 2024] analyzed the training dynamics of transformers with linear attention. Huang et al. [2023], Chen et al. [2024], Li et al. [2024] studied the optimization dynamics of lne-layer softmax attention transformers performing simple in-context learning tasks, such as linear regression [Huang et al., 2023, Chen et al., 2024] and binary classification [Li et al., 2024].

---

Among them, Huang et al. [2023] was the first to study the training dynamics of softmax attention, where they gave the convergence results of a one-layer transformer with single-head attention on linear regression tasks, assuming context features come from an orthogonal dictionary and each token in the prompts drawn from a multinomial distribution. In order to leverage the concentration property inherent to multinomial distributions, they require the sequence length to be much larger than the size of dictionary. Their analysis indicates that the prompt tokens that are the same as the query will have dominating attention weights, which allows the transformer to copy-paste the correct answer from those prompt tokens.
Li et al. [2024] studied the training of one-layer single-head transformer in ICL on binary classifica- tion tasks. Same as Huang et al. [2023]. they required the data to be pairwise orthogonal, and shared the same copy-paste mechanism as in Huang et al. [2023]. To be precise, fraction of their context inputs needs to contain the same pattern as the query to guarantee that the total attention weights on contexts matching the query pattern outweigh those on other contexts.
Chen et al. [2024] studied the dynamics of gradient flow for training a one-layer multi-head softmax attention model for ICL of multi-task linear regression, where the coefficient matrix has certain spectral properties. They required the sequence length to be sufficiently large [Chen et al., 2024, Assumption 2. together with restrictive initialization conditions [Chen et al., 2024, Definition 3.1]. While using the copy-paste analysis framework as in Huang et al. [2023], Li et al. [2024], the attention probability vector in their work is delocalized, so that the attention is spread out to capture the information from similar tokens regression tasks. Kim and Suzuki [2024] studied the dynamies of Wasserstein gradient flow for training a one-layer transformer with an mnfiite-dimensiona. fully- connected layer followed by a linear attention layer for ICL. of linear regression, assuming infinite prompt length. Nichani et al. [2024] analyzed the optimization dynamics of simplified two-layer transformer with gradient descent on in-context learning a latent causal graph.
Notation. Boldface small and capital letters denote vectors and matrices, respectively. Sets denoted with curly capital letters, e.g., W. We let (Rd II-II) denote the d-dimensional real coordinate space equipped with norm ||-|]- Ia is the identity matrix of dimension d. The (P-norm of is denoted by I|wllp- where 1 S S 0o. and the spectral norm and the Frobenius norm of a matrix M are denoted by || and || ||M]] respectively. M stands for the Moore-Penrose pseudoinverse of matrix M, and M.,, stands for its i-th column vector. We let [N] denote [1....,N,, and denote 1 to represent the all-one vector of length N. and by 0 vector or matrix consisting of all O's. We allow the application of functions such as exp(-) to vectors or matrices, with the understanding that they are applied in an element-wise manner. We use e, denote the one-hot vector whose i-th entry is I and the other entries are all 0.
2 Problem Setup
In-context learning with representation. We consider ICL of regression with unknown represen- tation, similar to the setup introduced in Guo et al. [2023]. To begin, let f Rd - Rm be fixed representation map that f(x) = (fi (x), ... fm (x)) for any x E Rd. The map f can be quite general, which can be regarded as a feature extractor that will be learned by the transformer. We assume that each ICL task corresponds to a map X f(-) that lies in the linear span of those m basis functions in f()), where À is generated according to the distribution Dx Thus, for each ICL instance, the (noisy) label of an input un (Vk E [K]) is given as
y/ = X (f(x)) + EK),
A~
Ek i.dd
where 7 > 0 is the noise level.
The goal of ICL is to form predictions on query Qquery given in-context labels of the form (1) on a few inputs, known as prompts. In this paper, we use V to denote the dictionary set that contains all K unit-norm distinct tokens, i... V {vi.... dK} c Rd with each token ||0 = We assume that each prompt P = PA provides the first N tokens (with N < K) and their labels, and is embedded in the following matrix
E" =
U11
U2
UN UN
=
V TT
E RRd+IIXN,
(2)
1/1
1/2
where
V = (V
,UN) E R R
(3)

---

E
E
UN UN+1
K2
M/
(WX)
W
(W H
87ha1
1/1
W
3/2
Value
softmax
Key
Query
head
H (heady
Jh=ad h==1
Figure 1: The structure of one-layer transformer with multi-head softmax attention.
is the collection of prompt tokens, and y == (y1, YNN) is the prompt label. Given the prompt as the input, the transformer predicts the labels for all the K tokens y1,... РK in the dictionary set.
Transformer architecture. We adopt one-layer transformer with multi-head softmax attention [Chen et al., 2024] illustrated Figure - to predict the labels of all the tokens in the dictionary V. where !l is the number of heads. Denote the query embedding as
E
UN+1
UN+2
OK
R
(d)))))(K-),,,
(4)
0
and denote the emhedding of both the prompt and the query as E = (EP, EQ) E (d++))x We define the output of each transformer head as
head,, (E) W EP softmax (E (WX) E)
[H],
(5)
where W (4+1), W E (dx((+1),, and WY E Kx++11 the query, key, and value matrices, respectively, and the softmax applied oolumn-wisely,, i... given a vector input æ. the i-th entry of softmax(x) is given by eti Σ, ets. The attention map of the transformer T T(E) is defined as
head; (E)
T(E) = WO
E R K
(6)
head (E)
where W the output matrix. Following recent theoretical literature to streamline analysis [Huang et al., 2023, Nichani et al., 2024, Deora et al., 2023, Chen et al., 2024], we assume that the embedding matrices take the following forms:
0 = (IK, II) E
BR)RReN RKKHK
W = (0, wa) E kxx((+1)
(7a)
(W TW Q
E (((++)))(+1) vh. E [M]
(7b)
where W/ = (wh.. )T.K) E RRKK and Qn E Raxd are trainable parameters for all h E [H].
The prediction of the labels is provided by the diagonal entries of T(E). which we denote by ÿ = (V,,... ,9x) E RR Note that Th takes the following form under our parameter specification:
Ak E [K]:
Th
y Σ Wh.s softmax(V V Qnex))
(8)
1=1
(Qr, wa } denote all trainable parameters of T. Let (c,, .x) mm denote the noise matrix. Given training data ICL instances, the goal of training is to predict labels ya for all vù E V. Specifically, we train the transformer using gradient descent (GD) by optimizing the following mean-squared population loss:
Training via GD. Let 0
(( Ex.
1
Th
8k )2)
K
5

---

We apply different learning rates nQ: nww > 0 for updating
(Qn /==1 and (wnthe) respectively. i...
the t-th (t 2 1) step, we have
Vh E [ 0 Q 12 Vo, ((0(t-1))
w wo (t-1) The Vw, ((0(t-1)), (10)
(1)
where (()) {Q (1) (1) 1 the parameter at the t-th step.
wh Jh=1
At inference time, given prompt P = Px with N examples, where À may not be in the support of the generation distribution Dx, the transformer applies the pretrained parameters and predicts the labels of all K tokens without further parameter updating.
Inference time.
3 Theoretical Analysis
3.1 Training time convergence
In this section, show that the training loss L converges to its minimum value at linear rate during training, i.c. the function gap
A L()()) inf L + - 0o
(11)
at a linear rate, under some appropriate assumptions.
Key assumptions. We first state our technical assumptions. The first assumption on the distribu- tion Dx for generating the coefficient vector À of the representation maps.
Assumption 1 (Assumption on distribution Dx). We assume that in (1) each entry X is drawn independently and satisfies E[A,] = 0 and E { 1 for all i E [m].
To proceed, we introduce the following notation:
Z (f. ... E BMN,
Z (z mTmIN) 1/2 NXN, F max
|
rE[N]
(12)
where is the i-th column vector of Z for i [N]. We further define c (k (KT,LEN,) and (t) as follows:
(
c (t) = softmax V V Uk.... V TT ())))))) vk) E E NxH
B 2C RR
(13)
To guarantee the convergence, we require the initialization of the parameters (0) satisfies the following condition.
Assumption (Assumption on initialization). For all [K], B has full row rank.
Before stating our main theorem, let us examine when the initialization condition in Assumption is met. Fortunately, only require the following mild assumption on V ensure parameter initialization has good properties.
Assumption (Assumption on V). There exists token matrix V (cf. (3)) such that v; + T.j. Vi ( j-
vector (II, ... _N of the prompt
Assumption 3 implies that V has distinct tokens, i.e., Uj + Us when A k. It is worth noting that Assumption 3 is the only assumption we have on the dictionary V. In comparison, all other theoretical works in Table 1 impose somewhat unrealistic assumptions on V. For example, Huang ct al. [2023], Li et al. [2023], Nichani et al. [2024] assume that the tokens are pairwise orthogonal, which is restrictive since it implies that the dictionary size K should be no larger than the token dimension d, whereas in practice it is often the case that K >> [Reid et al., 2024, Touvron et al., 2023]. In addition, Chen et al. [2024], Zhang et al. [2023a], Wu et al. [2023] assume that each token independently sampled from some Gaussian distribution, which also does not align with practical scenarios where tokens are from fixed dictionary and there often exist (strong) correlations between different tokens.
The following proposition states that when the number of heads exceeds the number of prompts, i.e. > N. we can guarantee that Assumption holds with probability by simply initializing H {Qu 1 using Gaussian distribution.

---

Proposition (Initialization of {Qr) Suppose Assumptions 1, 3 hold and H 2 N. For any fixed 5 > 0, ler Q% (i,j) "id. N(0, ;²), then Assumption holds almost surely.
Proof. See Appendix E.I.
Choice of learning rates. Define
Co min kE(K]
Amin
(B()) ((O
)
(14)
where (0) is the initial function gap (c.f. (11)). Assumption 2 indicates that Co > 0. Let y be any positive constant that satisfies
1/2
128v2
2 So
5/4
5/4
|212 VH fmax K 33/2 (0)
(15)
v2 -
We set the learning rates as
T/Q S 1/L
and Tho 72122
(16)
where
L 8/2H KK
2213
(0)
2 2 11
2
||
(x0²
1096 758
K
X2NA
50
mr
+ 2
16384
k
(1) 2
(17)
K
7
Theoretical guarantee. Now we are ready to state our first main result, regarding the training dynamic of the transformer.
Theorem (Training time convergence). Suppose Assumptions 1, 2 hold. We let w and set the learning rates in (16). Then we have
(( inf ((0)
11-
Mo 2K
)
inf -igf))))n()))))))))))))))))))- WEN.
(18)
Proof. See Appendix C.
Theorem I, together with Proposition 1, shows that the training loss converges its minimum value at linear rate, under mild assumptions of the task coefficients and token dictionary. This gives the first convergence result for transformers with multi-heal softmax attention trained using GD to perform ICL tasks (see Table 1). Our convergence guarantee (18) also indicates that the convergence speed decreases the size K of the dictionary the number H of attention heads increases, which is intuitive because training with larger vocabulary size or number of parameters is more challenging. However, small H will limit the expressive power of the model (see Section 3.3 for detailed discussion), and we require H 2 N to guarantee Assumption 2 holds, as stated in Proposition 1.
3.2 Inference time performance
We now move examine the inference time performance, where the coefficient vector À correspond- ing to the inference task may not drawn from Dxx In fact, we only assume that the coefficient vector inference time is bounded as in the following assumption.
Assumption 4 (Boundedness of X at inference time). for some B >
We assume that at inference time ||||| S B
For notational simplicity, let _ E (p denote
o (f(vN+1),... f(UK)) mK-N)
(19)
The following theorem characterizes the performance guarantee of the transformer's output ğ (after sufficient training) at the inference time.
'We leave tighter, but more complicated, expression of L in the appendix (cf. (61)) in the appendix and present simplified form in the main paper for readability.
7

---

Theorem 2 (Inference time performance). Let be the solution to the following ridge regression problem:
= arg min {2N Σ -i=1 (y: X f(u)))² +
21
(20)
Under the assumptions in Theorem 1, for any and E (0, 1), if the number of training iterates satisfies
log
B² (0)
12112t log( (1g
(m77)
T
(21)
log 1/(1
2K
then given any prompt P that satisfies Assumption the inference time, with probability at least 1 - 6, the output of the transformer y satisfies
with " = x2)
y
(22)
2K
Proof. See Appendix
In Theorem 2, (22) shows that after training, the transformer learns to output the given labels of the first N tokens each prompt, and more importantly, predicts the labels of the rest N tokens by implementing the ridge regression given in (20). Note that Akyürek al. [2022] studied the expressive power of transformers on the linear regression task and showed by construction that transformers can represent the closed-form ridge regression solution. Interestingly, here we show from an optimization perspective that transformers can in fact be trained to do so.
Generalization capabilities of the pretrained transformer. Theorem 2 captures generaliza- tion capabilities that the pretrained transformer can have.
Contextual generalization unseen examples: Theorem suggests that the transformer exploits the inherent contextual information (to be further discussed in Section 3.3) of the function template in the given prompt, and can further use such information predict the unseen tokens.
ii) Generalization to unseen tasks: Theorem also suggests that the pretrained transformer can generalize function map corresponding any À E R the inference time (albeit satisfying Assumption 4), which not necessarily sampled from the support of its training distribution Dx.
We note that the contextual generalization that the transformer has here is different in nature from the prediction ability shown in previous works ICL [Huang et al., 2023, Chen et al., 2024, Li et al., 2024, Nichani et al., 2024]. Those work focuses on a setting where each prompt contains a good portion of tokens similar to the query token, allowing the transformer directly corresponding answers from the prompt as the prediction. However, in practical scenarios, prompts often contain only partial information, and our analysis sheds lights on explaining how transformers generalize to unseen examples by leveraging ridge regression to infer the underlying template.
the label of the
How does the representation dimension affect the performance? Beyond the above discovery, several questions are yet be explored. For instance, while we demonstrate that transformers can be trained to implement ridge regression, how good is the performance of the ridge regression itself? What is the best choice of ridge regression could expect? How close is the transformer's choice to the best possible choice? We address these questions as follows.
Given any prompt P at inference time, since there is no label information about the rest K N tokens, the best prediction we could hope for from the transformer shall be
(1Z9)
y
(23)
where 2 is defined in (19), and X satisfies:
d == arg min Ez 2N Σ
(y) AT (f(p,) + +c)))
(24)

---

In other words, we hope the transformer outputs the given N labels as they are. For the rest K - N labels, the best we could hope for is that the transformer estimates the coefficient vector À by solving the above regression problem to obtain X,, and predict the k-th label by X f(vx) for
k == N +1.... K. Note that (24) is equivalent to the following ridge regression problem (see Lemma 4 in the appendix for its derivation):
{ 2N E A f(x.)) + [1A1
= arggin,
(25)
The only difference between the two ridge regression problems (20) and (25) is the coefficient of the regularization term. This indicates that at the training time, the transformer learns to implement ridge regression to predict the labels of the rest K - N tokens, assuming the noise level is given by N This observation also reflects how the sequence length affects the transformer's preference for choosing templates and its performance inference time:
The closer is to N, the closer the transformer's choice of templates is to the best possible choice, and the better the transformer's prediction will be;
When N m, the transformer tends underfit by choosing À with small E--norm;
When N > m, the transformer tends to overfit since it underestimates the noise level and in turn captures noise the prediction.
3.3 Further interpretation
We provide more interpretation on our results, which may lead to useful insights into the ICL ability
of the transformer.
How does the transformer gain ICL ability with representations? Intuitively speaking. our pretrained transformer gains in-context ability by extracting and memorizing some "inherent in- formation" of all basic function maps fi (i [m]) during the training. Such information allows it to infer the coefficient vector À from the provided labels in each prompt and calculate the inner product (X, [(xx) to compute ys given any token Vr E V at inference time. To be more specific, the "inherent information" of all basic tasks could be described by the N-by-K matrix A defined as follows (see also (34)):
A= T ZZ -1 1zT2+ (MTIN. 0)) RR
where 2 = (S(1)),, ... [(xx)) (z, (2, E R During training, the transformer learns to approximate A by D (o..ssfftax(V' Qn,vk) for each k E [K].
To further elaborate, we take closer look the special case when the labels do not contain any noise, i.e., r = 0, and N 2 m. In this case, becomes Z Ō, and given any prompt P = Px the coefficient vector À could be uniquely determined from the provided token-label pairs in the prompt. is straightforward to verify that the label of each token Uy could represented by the inner product of the given label vector and the k-th column of Z Z, i.c.
yk = y z 2
(26)
Comparing the above equation with (8), it can be seen that in order to gain the in-context ability, the transformer needs learn an approximation of Z Z by E Wr, V Qk(Rk for each 1 E [K].
More generally, in the proof of Theorem 2, show that
(y, A.,)),
(27)
comparing which with (8) suggests that a small training error implies that Σ Wi,,ssoftax(WW V QRvR) is close to A. In fact, this the necessary and sufficient condition for the training loss to be small. A rigorous argument is provided in Lemma 5.

---

The necessity and trade-offs of multi-head attention mechanism. Multi-head attention mech- anism is essential in our setting. fact, is generally impossible to train shallow trans- former with only one attention head to succeed in the ICI. task considered in our paper. This because, as we have discussed above, the key for the transformer to approximate A., by Σx=1 Wh. vg. ) for each & E [K] If H = 1, the transformer could not approximate each A by W1, 15oftmaa (V Qi in general since the entries of the latter vector are either all positive or all negative. transformer with a simple initialization method satisfy our desired property that is crucial to guarantee the fast linear convergence. However, (18) implies that we should not set H to be too large, since larger H yields slower convergence rate.
addition, Proposition 1 indicates that when 2 N, the weights of the
4 Conclusion
We analyze the training dynamics of one-layer transformer with multi-head softmax attention trained by gradient descent to solve complex non-linear regression tasks using partially labeled prompts. In this setting, the labels contain Gaussian noise, and each prompt may include only few examples, which are insufficient to determine the underlying template. Our work overcomes several restrictive assumptions made in previous studies and proves that the training loss converges linearly to its minimum value. Furthermore, we analyze the transformer's strategy for addressing the issue of underdeterminatiion during inference and evaluate its performance by comparing it with the best possible strategy. Our study provides the first analysis of how transformers can acquire contextual (template) information generalize to unseen examples when prompts contain limited number of query-answer pairs.
Acknowledgments and Disclosure of Funding
The work T. Yang and Y. Chi is supported in part by the grants NSF CCF-2207911, DMS-2134080 and ONR N00014- 1--444. The work of Liang was supported in part by the U.S. National Science Foundation under the grants ECCS- 21CS-2213860) DMS-2134145 and 211-21
References
R. Agarwal, A. Singh, L. M. Zhang, B. Bohnet, S. Chan, A. Anand, Z. Abbas, A. Nova, J. D. Co-Reyes, E. Chu, et al. Many-shot in-context learning. arXiv preprint arXiv:2404. /1018, 2024.
K. Ahuja, M. Panwar, and N. Goyal. In-context learning through the bayesian prism. arXiv preprint arXiv:2306.04891, 2023.
E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou. What learning algorithm in-context learning? investigations with linear models. arXiv preprint arXiv:2211.1566) 2022.
Y. Bai, F. Chen, H. Wang, C. Xiong, and S. Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901 2020.
S. Chen, H. Sheen, T. Wang, and Z. Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402. 19442, 2024.
X. Chen and D. Zou. What can transformer learn with varying depth? case studies on sequence learning tasks. arXiv preprint arXiv:2404.01600 2024.
D. Dai, Y. Sun, L. Dong, Y. Hao, S. Ma, Z. Sui, and F. Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers, arXiv preprint arXiv:22 /2. 10559, 2022.
10

---

P. Deora, R. Ghaderi, H. Taheri, and C. Thrampoulidis. On the optimization and generalization of multi-head attention. arXiv preprint arXiv:2310.126800 2023.
B. L. Edelman, S. Goel, Kakade, and C. Zhang. Inductive biases and variable creation in self- attention mechanisms. In International Conference Machine Learning, pages 5793-5831. PMLR, 2022.
B. L. Edelman, E. Edelman, S. Goel, E. Malach, and N. Tsilivis. The evolution of statistical induction heads: In-context learning markov chains. arXiv preprint 2Xii:24422110000 2024.
Garg, D. Tsipras, Liang, and G. Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
A. Giannou, S. Rajput, J.-y. Sohn, K. Lee, J. Lee, and D. Papailiopoulos. Looped transformers as programmable computers. In International Conference Machine Learning, pages 11398-1144 11442. PMLR, 2023.
T. Guo, W. Hu, S. Mei, H. Wang, C. Xiong, S. Savarese, and Y. Bai. How do transformers learn in-context beyond simple functions? a case study on learning with representations. arXiv preprint arXiv:23 10. 106/6, 2023.
M. Hahn and N. Goyal. A theory of emergent in-context learning as implicit structure induction. arXiv preprint 2rXiv:2303.07771 07971, 2023.
C. Han, Z. Wang, H. Zhao, and H. Ji. In-context learning of large language models explained as kernel regression. arXiv preprint arXiv:2305.12700 2023.
Y. Huang, Y. Cheng, and Y. Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
H. J. Jeon, J. D. Lee,, Q. Lei, and B. Van Roy. An information-theorrticc analysis of in-context learning. arXiv preprint arXiv:2401. 2024.
H. Jiang. A latent space theory for emergent abilities in large language models. arXiv preprint arXiv:2304.0m0 09960, 2023.
H. Karimi, Nutini, and M. Schmidt. Linear convergence of gradient and mekximal-gradient methods under the Lojaswwwjj condition. In European Conference on Machine Learning and Knowledge Discovery in Databases, pages 795-811, 2016.
J. Kim and T. Suzuki. Transformers learn nonlinear features in context. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024.
B. Laurent and Massart. Adaptive estimation of a quadratic functional by model selection. Annals of statistics, pages 1302-1338, 2000.
H. Li, M. Wang, S. Lu, X. Cui, and P.-Y. Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.1 15607, 2024.
Y. Li, M. E. Ildiz, D. Papailiopoulos, and S. Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pages 15565-19594 PMLR, 2023.
Q. N. Nguyen and M. Mondelli. Global convergence of deep networks with one wide layer followed by pyramidal topology. Advances in Neural Information Processing Systems, 33:11961-1199 2020.
E. Nichani, A. Damian, and J. D. Lee. How transformers learn causal structure with gradient descent. arXiv preprint 172.435 2024.
C. Olsson, N. Elhage, Nanda, Joseph, N. DasSarma, T. Henighan, B. Mann, Askell, Bai,
A. Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209, 11895, 2022.

---

M. Reid, N. Savinov, D. Teplyashin, D. Lepikhin, T. Lillicrap, J.-b. Alayrac, R. Soricut, A. Lazaridou, O. Firat, J. Schrittwieser, al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.055300 2024.
H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint ar arXiv:23377009265 2023.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
J. Von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pages 35151-35174. PMLR, 2023.
X. Wang, W. Zhu, M. Saxon, M. Steyvers, and W. Y. Wang. Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning. Workshop on Efficient Systems for Foundation Models@ @ ICML2023, 2023.
J. Wei, J. Wei, Y. Tay, D. Tran, A. Webson, Y. Lu, X. Chen, H. Liu, D. Huang, D. Zhou, et al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303.03866 2023.
N. Wies, Y. Levine, and A. Shashua. The learnability of in-context learning. Advances in Neural Information Processing Systemx, 36, 2024.
Wu, D. Zou, Z. Chen, V. Braverman, Q. Gu, and P. Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv:2310.0839) 2023.
S. M. Xie, A. Raghunathan, P. Liang, and T. Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.0208, 2021.
R. Zhang, S. Frei, and P. L. Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.0990 2023a.
Y. Zhang, F. Zhang, Z. Yang, and Z. Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305. / /9420, 2023b.
12

---

(ood)
iieeence locd)
38
(a) l-layer transformer
(b) 4-layer transformer
Figure 2: Training and inference losses of (a) I-layer and (b) 4-layer transformers, which validate
Theorem 2, as well as the transformer's contextual generalization to unseen examples and to unseen tasks.
A Experiments
This section aims to provide some empirical validation to our theoretical findings and verify that some of our results could be generalized deeper transformers.
Setup. We conduct experiments on synthetic dataset, where we randomly generate each token v2 and their representation f(v)) from standard Gaussian distribution. We employ both the 1-layer transformer described in Section 2 and a standard 4-layer transformer in Vaswani et al. [2017] with dmodel 256 and dæ 512. We set the training loss to be the population loss defined in (9), and initialize {Q), (0) The и using standard Gaussian and set ( )
]he[ar] to be identical what is specified in Section 3. We generate À from standard Gaussian distribution to create the training set with 10000 samples and in-domain test set with 200 samples; we also create an out-of-domain (ood) test set with 200 samples by sampling À from (11 m. Given X, we generate the label yı, of token vr using (1), for [K]. We train with a batch size 256. All experiments use the Adam optimizer with a learning rate 1 X 10-4
Training and inference performance. We set N = 30, K 200, d 100, 20, and set H to be 64 and for I-layer and 4-layer transformers, respectively. Figure 2 shows the training and inference losses of both I-layer and 4-layer transformers, where we measure the inference loss by K T 1/ to validate (22): after sufficient training, the output of the transformer ÿ converges to y From Figure 2 we can see that for both 1-layer and 4-layer transformers, the three curves have the same descending trend, despite the inference loss on the ood dataset is higher than that on the in-domain dataset. This experiment also shows the transformer's contextual generalization unseen examples and to unseen tasks, validating our claim in Section 3.2.
Figure 3 plots the performance gap R - Thest 2| of the one-layer transformer with respect to different ranging from 50 to 150, when we fix m 100 and T = 0.01. This verifies that the ridge regression implemented by the pretrained transformer has a better performance when m is close to N, again verifying our claim at the end of Section 2.
Impact of the number of attention heads. We now turn to examine the impact of the number of attention heads. In this experiment, we use the population loss (9), and set the other configurations same as those in Figure 2. Figure 4 shows the training loss curves for different H with respect the iteration number, which validates our claims. From Figure 4, we can see that we need to set H large enough to guarantee the convergence of the training loss. However, setting H too large (H 400) leads to instability and divergence of the loss. Recall that in Proposition 1, we require H 2 to guarantee our convergence results hold. Although this condition may not be necessary. Figure 4 shows that when H < N = 30, the loss stopped descending even when it is far from the minimal value. On the other side, the loss keeps descending when H 30 (though slowly).
13

---

Amh Mman
200
1011
103 120
140
UNkmbb euxmp pmompr)
Figure 3: The performance gap K 1"" that the closer N is to m, the better the transformer's prediction is.
g with different N when m 100, which validates
104
10--
10
200
283
4012
Figure 4: Training losses the 1-layer transformer with different number of attention heads H, where H should be large enough to guarantec the convergence of the training loss, but setting H too large leads to instability and slower divergence.
We also explore how H affects the training of the 4-layer transformer, as displayed in Figure 5, where we set = 200 and the configurations other than H are the same as in Figure 3. We fix the wall-clock time to be 100 seconds and plot the training loss curves with different H. Figure 5 (a) shows the final training and inference losses with respect to H. It reflects that the losses converge faster with smaller H (here the final training loss is the smallest when H 4). The training curves in Figure (b) corresponding to different H within 100s may provide some explanation to this phenomenon: (i) transformers with larger H could complete less iterations within fixed amount of time (the curves corresponding to larger H are shorter); (ii) the training loss curves corresponding to large H (H- 32, 64) descend more slowly. This suggests our claim that larger H may yield slower convergence rate is still valid on deeper transformers. Note that unlike the l-layer transformer, deeper transformers don't require large H guarantee convergence. This is because deeper transformers have better expressive power even when H is small.
B Proof Preparation
B.1 Summary of key notation
We summarize the frequently used notation in Table 2 for ease of reference.
B.2 Auxiliary lemmas
We provide some useful facts that will be repeatedly used later on. Let
2 f(v)) (f(())).falv)) fz(uw).
fm
RM,
Wk e [K].
Recalling (12). we can rewrite
Z= (z ,NN) E m
We further define s E R N
as follows:
$ softmax(V V QuvR))
T vk [K], h e [H].
(28)
14

---

Taining
Traiking Last
101
inference nindimahl
Number of Attention Heads (H)
beration
(a) final losses H
(b) training loss curves for different H
Figure 5: Training losses of a 4-layer transformer with different H, fixing wall-clock time to be 100s. This experiment shows that unlike 1-layer transformers, deeper transformers don't require to be large to guarantee convergence of the loss.
notation K N+ d E N+ E N+ H E N
meaning
total number of tokens
token dimension
number basic tasks
number of attention heads
N N+
number of examples each prompt
Uk Rd k E [K]
the k-th token the i-th basic task coefficient vector the k-th label
$ Rd
+ R i E [m] A R
Vh A (f(Us) + Ek), k [K]
Table 2: Notation for key parameters.
Lemma (Softmax gradient). For all j E [N], k E [K] and h E [H),
we have
0sh aQk
N
h jk
s (u, v.juk
(29)
1=1
where si is defined in (28).
Proof. See the proof Lemma A. 1 in Huang et al. [2023].
Lemma 2 (Smoothness of softmax). For vectors ε &2 E Rd, have
|softmax({:) softmax((aa)) -oftmma S 2 [| - E2ll
(30)
Proof. See Corollary A.7 in Edelman et al. [2022].
We also need make use of the following form of Young's inequality. Lemma 3. For any :1.-.. xt1 E R", we have
Σ T
(31)
The following lemma shows the equivalence between (24) and (25). Lemma (Equivalence of the regression problems). Given any prompt Px (V1, 01, we have the following equivalence:
UN.,NN)
N
1
Ex
yy
T +:))2
A [(x:))²
1
(32)
2N
2N
15

---

Proof. See Appendix E.2.
C Proof of Theorem 1
We first outline the proof. To prove Theorem 1, we first remove the expectation in the expression of the loss function L in (9) by reformulating it to deterministic form (see Lemma 5). With this new form, we show by induction that the loss function is smooth (Lemma 10) and satisfies the Lojase--o (PL) condition (c.f. (49)). Provided with both smoothness and PL conditions, are able to give the desired linear convergence rate [Karimi et al., 2016].
We define
H Σ h- & (z Z -mT)) (Z Z + m
(Zk + mTTE),
if k [N],
8%
(33)
E
if E [K] [N].
We also define the following matrices:
A
7 Z + MTTNN
2T (mTIN. E ERN*K
(34)
H
Á(O)= =
Wh,18 81
wh.Ko%
ERNKK
(35)
1=1
where 2 = (Z),-.. 2K) ) E I
We first reformulate the loss function to remove the expectation in the population loss.
Lemma 5 (Reformulation of the loss function). Under Assumption 1, the loss function ((0) could be rewritten into the following equivalent form:
K
(())=
(zT Z + +mTI)
1/2
8% +
L
28 1 + L",
(36)
2K
2K
k=1
where
L
1
Σ
1-(zz 2か + mTTE) T (z + mTl)
(z 2x + mTex) + 1| + mn)
2K
k=1
K
Σ
(- 2x) (z Z + m (z 2τ) (zTzz) + || 112
+
2K
k=N+1
a constant that does not depend on 0, and Z is defined in (12).
Proof. See Appendix E.3.
Lemma 5 indicates that L* is a lower bound of L. We'll later show that L is actually the infimum of C. i.e., info L(O).
Lemma 5 also indicates that, the necessary and sufficient condition for L(O(t) converge L during training is
Vk [K] 0g¹ 0, t 00, co, which folllows immediately that (37) is equivalent to
(37)
A A 0.
- oc.
(38)
To simplify the analysis, we introduce the following reparameterization to unify the learning rates of all parameters, and we'll consider the losses after reparameterization in the subsequent proofs.
Lemma (Reparameterizzatonn). Define
7 == /n/ Oy == w//r Wh E (H),
(39)
16

---

and let
= (Qh: an anl-= ((E) = (()).
(40)
Then (10) is equivalent to
((() g g(-11) n1 Vge(((1-1) AE E (T].
(41)
Proof. See Appendix E.4.
We denote a as a = (an,))ee(H) E #
The following lemma bounds the gradient norms by the loss function, which is crucial to the proof of Theorem 1.
Lemma (Upper bound of the gradient norms). for all h E [H). we have
Suppose Assumption 1 holds and h S Then
(9(
S 22V27a ( ((E(t)) - L²
(42)
0Q% (t)
Proof. See Appendix E.5.
Now we are ready to give the main proof.
Proof of Theorem 1. To prove Theorem 1, it suffices to prove that under our assumptions, have
(Upper bound of the parameters:)
|
(t)
a
(43a)
(B,BB B z So
(Lower bound of eigenvalues:)
Xmin
(43b)
((0(0) - L* $ 5(15 (1-12 )20)²
-(()))))-'))
(Linear decay of the loss:)
-
(43c)
where
S
a= V2K 4 4
750 v [(((0) L*,
(44)
K
and y, on is defined in (39), is defined in (14). We shall prove (43b).(43b) and (43c) by induction.
Base case. It is apparent that (43b),(43b) and (43c) all hold when t = 0.
Induction. We make the following inductive hypothesis, i.e., when E [t 1), (43a), (43b) and (43c) hold. Below prove that (43b),443) and (43c) hold when = f by the following steps.
Step 1: verify (43b) and the loss w.r.t.
condition. condittion
We first compute the gradient of the
al(E) 1 à
1
Vk E [K]:
Z (7Cbar AA)ll
doc
2K
dac
2K
da
Y
(2Ck)
2
7 B 26%
(45)
K
K
where the first equality follows from Lemma 5, Ck. Bg is defined in (13).
Let b) denote the h-th column vector of Bx, [H], ie., B₁ (b), b then for any k E [K] and t E N,, we have
|| (B)) (b)))
< 1 01
(s
(s)
< ||
(st) (0) (s)
< <
(TQ
(0)
<2
max
1. (Q)
]E[N]
<2
1 !)) Q (0)
(46)
17

---

where the third line uses Lemma 2, and that
)E(E(())
vh E [M]
0
Q (0)
7
aQ%
2v 22 max ((E(5)) L²
==0
1-1
2 2//21ros V [(0(0)) -
Σ
1 2
(V
110
8/27afuux [(((0))-4 -.
(47)
0
where the second inequality follows from Lemma 7 (cf. (42)) and the third inequality follows from the inductive hypothesis and the fact that ((5(+)) = L()(*)), Vs. Combining (47) with (46), we have
H
B
2
/ 11
(0)
1==1
2
VH
((010)) L
1/v2
Vo.
(48)
where the last inequality follows from (15). The above inequality (48) Vx E R
indicates that
B
M
x
B
(B k x
T
-B 2 v /2/
(
which gives (43b).
Therefore, we obtain the following PL condition:
K
OE(E)
72
Σ
(26
B 2
Dan,
K2
2
(1)
|2
-²),
(49)
2K
2
where the equality comes from (45), and the last equality follows from (36).
Step verify the smoothness of the loss function. We first give the following lemma that bounds the Lipschitzness of by and og which will be used later on. For notation simplicity, we let B, Q, o denote B)).. @0), (0), respectively, and let B', Q', a' denote B Q(0'), (00), respectively.
Lemma 8 (Lipschitzness of b% and 6.). For all k E [K] and h E [H], and all transformer parameters 0, 0', if max maxx(on,,l, ,0 I a, then we have
(b((0) by (0') ||2 < 2 2
1 QAlIF
(50)
H Σ
822 5% |1 22VVHa Ha,
QQ
Q/1l + HH H ||ax a'll2.
(51)
h=1
Proof. (50) follows from a similar argument in (46). Regarding the Lipschitzness of 8g we have
H
H
6%
5%
=7
Σ
- (((
+
Σ
(an,k - as
2
Aml
S 7
Jonkk
[s)(0))
(0') ||2 + 7
ohn - ah (
h=1
< 2Y
Σ QQu 021l HHH H | -
18

---

where we use (46) again to bound the first term in the second line, and use the fact that (0)) 21 and Cauchy-Sccie -chhaar inequality bound the second term in the second line.
We also need the following lemma which bounds the norm of B and 6%.
Lemma 9 (Upper bounds of by and 8%) For all k E [K] and h E [H], if max(|an. I then we have
10h S a,
| = 2
(52)
16%112 S HHa + ||
(53)
where A is defined in (34).
Proof. (52) follows from
| - || |s2ll2 S 1|
(53) follows from
H 5
|02111
Σ h=1
lonal 1 lbll2 + -|1Aexllas S SHH ||A
As consequence of Lemma 8 and Lemma 9, For all l E [K], and all transformer parameters 0,
if max{ l |0/,11 S a, we have
V.,, (() V.,. ((E))12
(45) 7
(Bx- B², 25% + B¹ T (55 5% >||
K
< K
?
| 212 |B Bill= |102112
Y
|z|| |
0% 5% 112
K
2²
Y
2| +
Σ
1Q, Q
H
| c
(54)
K
K
from which we obtain the smoothness the l w.r.t. a as follows:
(EE)
va((E')1
K
V Van ((E) - Vay ((E')11
H Σ
71 K²
$
2K
2|| (2 Ha + [|A112))
QQ 12 +2
H2 | l|a
a'I
K
1=1
< 2
1
2y
2 2 2 (2>Ha + [Al1)) +
H2 K2
[2112
|E E'112-
(55)
where the first inequality uses Young's inequality (c.f. Lemma 3).
To obtain the smoothness of the loss function w.r.t. Qn. we first note that by (82) we have
aE(E) 7 N aQh K
N
Σ Σ (26%)
h 2, an
Σ
S (v,
(56)
k=
19

---

Therefore, if max{(an,
10/ S a,
we have
(((())
iQ
00(5')
aQn
fmax 27
K
K
2
5%
11.
as 3b
(0)
Σ
N
t/ 0
k=1
NV
N
6% 112112
ahals (0)
(0)
1=1
NV.
|1
11a
8%
(0)) s3((0'))
Σ (0)
1=1
N
N Σ
+
Σ
||
11,
6%
08 jA (0')
1 (0)
s (0')
2=1
27
max
1Z112
K
0% 0%
68%
a
K
1=1
N
N
+
8/'
by (0)
(0)) + 5%
10 (0) - -122100)
2
j=1
i=1
27fmax
12112
K
t 5% 11 a + |008" |1,
lan, - ahh
K
1=
s 02(00) } $ (0)
+ 2
g'
avN
N
(57)
where the third inequality uses Cauchy- Schwarz inequality. Combining the above inequality (57) with Lemma 8 and Lemma 9, we have
OE(E)
06(5')
Fyax
112
H
<
2Ka
Σ
1Q
2/1 + VK l| a'llp)
aQ
aQ
F
K
h=1
+ (ĄH + (A) VK] an.. - ahrill
+ (HHa + 2 2aVN 2K 11Q% - Qullp } }
(58)
where the last line uses (46) to bound ||,](0) - (00) ||2 The above inequality (58) further gives
Σ Q ((E) Vq, ?F')
Mmax || (2Ka)² [107H]] + 4N 4N (07H+11411))
" Σ h=1
88.
QQQ - 2//12
K
+ K [<((0) + (ayH+
[a0 + [|A|11)) |a la-
<8yfmax || max (2VK))))) [a07H]]
+ 4N (a + (1A11)) ||5" El (59)
where the first inequality makes use of Young's inequality (c.f. Lemma 3).
Combining the above two relations (55) and (59), we obtain the smoothness of l w.r.t. & as follows:
Lemma 10 (Smoothness of the loss function), Let ry == /nw/na. For all transformer parameters
E. if max{ l. /o/, I a, then we have
[IVE((E) Vel(E")llz S T.11 - -E'122,
(60)
where
I 2 K (2> 2 | (2,Ho Ho ||A1|2)
(1
1
71 +
x H² |
K²
(61)
+ 8y/maxx || |21|2 max {1,
[1771133 + 4N (a7H + [111))]]
20

---

Step 3: verify (43a), (45) implies
dE(E) 2
(bb) 28%
dank
K
which, combining with (52), gives
dE(E)
72 K2
4k E [K]. h E [H]
S
|
1
12
2
2
Combining this with (36) we obtain
((E)
2
272
I ||
(E(E) - L*
8cch
K
which indicates
0E(E)
2
||
1111
112 7V
(E(E) - L*).
(62)
doc 1
K
Therefore, we have
)
Q
110
2
dan
1=
2
00(5(1)
(0)
Σ
0
11Q
2
dom
1
2y
+ nQ
|| 1212
V
V (({())
L
K
==
272 (E(0(0) -(*)
1-1
(0)
1Q || V
720
2
K
2-2 (L(0())) -C*)
+ no |
K
where the second inequality follows from (62) and the third inequality follows from the induction hypothesis (43c). (43a) follows from plugging defined in (44) into the above inequality and using the initializtion condition that a un O.
verify the linear convergence rate (43c). Combining (43a), (60) and Lemma 4.3 in Nguyen and Mondelli [2020], we have
Step 4:
((1-1)) + 2 | g g(( 1
L- - ((g((-1))- L + (Ve((e((-1)) -
(63)
2
which indicates when 11Q S 1/L, we have
- 2 L
712 | Ve((e(t-1)) 112 (49)
E (1- (1-))2) n220
(e(e(t-1)) -- (64)
2
which, combined with the fact that L(B(+)) = ((g(s)) for all s (see Lemma 6), verifies (43c).
Note that (36) implies that e < L()) holds for all 0. And from (43c) we know that C() + L
as + co. Therefore, it follows that
=- inf (()).
Consequently, (43c) is equivalent to (18).
21

---

D Proof of Theorem 2
By (43c) we know that (()()) - L as t + cc. Thus from (36) we know that (37) and (38) hold.
By Sherman-Morrioon Morrsssooooo- formula, we have
(MTIN + =)
IN 1 mt
z (MTIm +Zz -1 Z.
(65)
me
Thus we have
A [ Z + mTTIN) -1 T 2 + (mT (MTINN,)) 0)
1
IN z (mT ZZ 2) 1 2 + (mT NN, 0))
Z + (mT IN, 0) z (mTIm= =z=) -1 (MTIm + +Zz 2
MT
MTT (mTIn m + 2z -1 mT ' (MTIm + -
=IN, Z (mT I, + ZZ ) (0)
= IN. T (mTIm + zz -1 zo).
(66)
where 0 is defined in (19). On the other hand, it's straightforward to verify that Ā defined in (20) admits the following closed form:
À (m + Zz
Zy.
(67)
Combining the above two equations, we obtain
y A y (02) (mTIm + Zz-)=² Zy = -2 0",
(Z9)
=
where the last equality follows from (22).
Now we give the iteration complexity for the mean-soo mean-squaree -quarre error between the prediction ÿ and the limit point z² to less than Given any prompt Px, where À satisfies Assumption 4, we have
1/1 = x (z; +E,) ~ N(A z,, |All2
Letting 2, T ||
have 24 NO, 1). Define
AAl T - 1) = y 2
Nr |A 211 T
By Laurent and Massart [2000, Lemma 1], have
Vs > 0: P p 2 B(Z22VN TV3 75
exp
By letting s = log(1/5) and using the definition of Z, we have
(IV T 1 2 N7 | + 2 2 N log (1/8) | + + T log(1/6)) s 8.
(68)
Thus with probability at least 8, we have
|0l12 < Z + yy Z A ll
$ Z 2 + ||ll2 V (N + 2 N log(1/)) + 22 2log(1/6) 1/2
B |2 1| VT (N 2/ + +2 2log(1/8))
(69)
22

---

where we use (68) in the second inequality, and the third inequality follows from Assumption 4.
On the other hand, by (36) have
)(0(0)= = 2K
2(A A)
+ + 2
A +
2K
which gives
A
A
S
2K
(C(O(T)) -*) 1
2K (L()
2
L*) (1
7
no 2K
T/2
(70)
2
mi
mt
Thus we know that w.p. at least
we have
2k 1ー
1
4 A)
S
A
all2
11v112 S
2K
2K
where the last relation follows from (69), (70) and (21).
E Proof of Key Lemmas
E.1 Proof of Proposition
For notation simplicity we drop the superscript (0) in the subsequent proof.
Let D (VT Q QION.".. ... VTQHou) NxH Note that
Vi ( E (H). D3 = V = V (q ,qu), where (i.,)) 3id. NOO, 1x
(71)
This suggests the column vectors Dj. are i.i.d. and the density of each column vector positive at any point E R(V). where R(V) c RV is the row space of V.
Since Z has full rank, prove B& has full rank a.s., we only need argue that Ct 1: .N) has full rank w.p. 1. Below we prove this by contradiction (recall that by definition C₁ softmax (DainaxxxD))....... and we assume H 2 N).
Suppose w.p. larger than 0, there exists one of C=(:: 1: N)'s column vector that could be linearly represented by its other N - column vectors. Without loss of generality, we assume this colomn vector is C& (;, 1) = softmax( D& 1)). Let æ = x(q:) == exp( 1)) exp(V q:). Then æ could be linearly represented by exp (;,
2....,N.
Let Á = exp(pp (t, N)), then w.p. larger than 0. C(A), where C(A) is the column vector space of Á. i.e., we have
P(x C(A))A)) > 0. JRNX(--1)
which further indicates that there exists A Nx such that P(x C(A)) Since the dimension of C(A) is at most N 1, there exists y E R.N. y * such that yCA) Therefore, we have
P(y x = > 0.
(72)
By Assumption 3, without loss of generality, we assume that in = (E11, U12, ... UIN) T has different entries. For any vector w = (wi.... wa) E Rd, we let v = (wz, ... wa) E Rd-1 denote the vector formed by deleting the first entry of w. Let qi (q, q For any fixed q Rd-1 the function g( g(-l@i) IR + P defined by
(لقاوو)
Σ
y:e 9m +0,
Σ d o ex, qı)) (y, x(q1))
has finite zero points and thus {g Rlg(gl@n) 0} is meroume set. Therefore, we have
P( x) gd-1
((g(glg)) 0lq (q))))) 0,
which contradicts (72). Therefore, C2 1 N) has full rank with probability 1.
23

---

E.2 Proof of Lemma 4
Lemma 4 can be verified by the following direct computation (recall that the noise in each label
satisfies E2 a.d N(0, m.), Vi [N]):
1
y T (S(v.) + E:))
2N
-=1
N
E
((ys - - s(vi))) - 22 Es(U: 2TTec
A s(s))) (TGGTN
2N
==1
1
(u. s(v))) + -121
2N
1
X X f(xx))++ + |
2N
E.3
Proof of 5 5
We let ef = (EI,.. ... EN) E R=xN, E = (E1,-.., ,EK) E R Recall that
(y1. UNN R. Then we have
y = (Z =T))T,
(73)
and
K
1 K
((
1
Σ
(0)
_x,
(0x- )²
(74)
K
K
k=
k=
K
1 2K
Ex,E yy ā= - -A (Zx + Ex)112
k=1
1
Ex.e AT + Pyā= AT Ex)112
2K
k=1
1
Ex,c [Z + (=6"@a -2x [x)] 2A [(Z + c²)ā= -2x + Ex)
2K
1
Ee [Z + e²)āz -(z + [w)] T [z - =²"@az - (2x + c+)l
2K
k=1
K
Σ
E.
Zāz - 2x12 + 2(ā 2x) (6P as (6)
6x) + | - cx12))
(75)
2K
A=1
where àg denote the k-th column vector of matrix Á(0) defined in (35), and the fifth line uses Assumption 1.
Note that for all l E [K]. we have
(Z - 2k) T - (c) 0,
(76)
and that
Ex 1 an 1aaaa-en Ex mi (Iāxe + 1) 2mmTTu I {k: E [N]]
(77)
where 1 (I& E [N]} is the indicator function that equals 1 if k: E [N] and 0 otherwise, and we have i.i.d. ~0)
made use of the assumption that eş Combining the above two equations with (75), we know that for & E [N], it holds that
(0) = 2 ( = + mr a ex
24

---

Reorganizing the terms in the RHS of the above equation, we obtain that
(x) = - (z z + mfl) 1/2 ax (z T Z + mT1) -1 ( 24 + mTex) )1² 2
2 (78)
2
where Cg (Z" 2 mTee) T (z Z + mTT) -1 (Z 2 + mTee) + 1| + MT.
By similar argument, we can show that for & E [K]] it holds thet
(x(0) | (z + mTI) 1/2 (@x (z Z + m - 2 2/) | + 2 ck:
(79)
where of - (ZTzs) 2k)) (Z Z mT
(zTzx) || l22
(78). (79) together with (33) and the definition of Le give (36).
E.4
Proof of Lemma
First, it holds that
Q =QQ (t-1) 17Q Vo, = QQ (t-1) 77 Qq. (((
(80)
Second, note that
w/ w/ f/ V w/ ((0(1-1))
(t) (t-1)
yah (1-1) 1) 7² ? 112 V an ((((1-1))
= 7 a (1-1) nιa Val (()(t=1)))
Dividing both sides of the above equality by have
a (t-l)
no Va, ((g(t-1))
(81)
Hence, (41) follows from combining (80) and (8 1).
E,5
Proof of Lemma 7
Throughout this proof, we omit the superscript (f) for simplicity. We first compute the gradient of w.r.t. Qn. By (36) we know that
((E) = C())
Σ
11
2K
and thus we have
OE(E)
K
N
0
N
27
06
Σ
8,
2
aQk
K
06
2
iQk
k= J=1
1=1
Σ
(2ox)
3j
866 (1)
(82)
K
1=1
Note that
G Ghijk]
1p S 2as
(83)
25

---

where we use the fact that [|(0; 1 11 S 2 (recall that we assume each Uy has unit norm, k E (K].) Combining (82) and (83), we have the desired result
)E(E))
Zöx 112 ||2
|| Gh
aQn
K
k=1
j=
Zdu Fmaxasjia
==1
24 fm VK
K
Σ
28 1
K
S 22/27/maaa
((E)
L+
(84)
where fmax is defined in (12) and the third line follows from Cauchy-Schwaad inequaaal
26

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: we clearly state in the abstract and introduction the claims we made, including the contributions made in the paper and important assumptions and limitations.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: we clearly state our assumptions.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- lant role in developing norms that preserve the integrity of the community, Reviewers penalize honesty concerning limitations.
will be specifically instructed
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
27

---

Justification:
we provide the full set of assumptions and complete (and correct) proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
28

---

Answer: [No]
Justification: The experiments are very simple and can be easily reproduced by following the instructions in the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. ec/
oolicylyyy)).. for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: http//uuidee//oeeeee
/ //nips. cc/
wwisssolicccyyii for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Experiment details are included in Section A. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[No]
Justification: stochasticity is not critical in our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
29

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [NA]
Justification: the results are irrelevant to the compute resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https:/ /tte::///eu....//eeeee httpp:///arrpps.cppbbbbbbi//////
Answer: [Yes]
Justification: the research conducted in the paper conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: this is theoretical paper and it has no societal impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
30

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: the paper aims to provide better understanding on existing algorithms and thus poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: the paper does not use existing assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
31

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: the paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/mooel as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
32


Reviews:
[{"review_id": "jkU1G7bTab", "reviewer": "Reviewer_dubr", "summary": "This paper presents a rigorous theoretical analysis of the training dynamics and generalization capabilities of a one-layer transformer with multi-head softmax attention for in-context learning (ICL) of non-linear regression tasks. The authors consider a more challenging and realistic setting where prompts contain only a small number of noisy labeled examples, insufficient to fully determine the underlying template function. This addresses limitations in previous theoretical work that often required unrealistically large numbers of examples per prompt.\n\nA key theoretical contribution is proving convergence of the training loss to its global minimum at a linear rate when optimizing with gradient descent. This is the first such convergence result for transformers with multi-head softmax attention on ICL tasks. The analysis reveals that the transformer effectively learns to perform ridge regression during training. The paper provides precise bounds on the mean squared error between the transformer's predictions and those of the optimal ridge regression solution, as a function of the number of training iterations.\n\nThe authors demonstrate that multi-head attention is crucial for succeeding at this ICL task. They provide a lower bound on the number of heads needed (H ≥ N), while also noting that too many heads can slow down convergence. The analysis shows how the multi-head mechanism allows the transformer to approximate the required matrix operations for ridge regression.", "questions": "Have the authors considered adding any toy domain experiments that would illustrate and further validate their claims ?", "limitations": "- Simplified model: The analysis is limited to a one-layer transformer, which is much simpler than state-of-the-art models used in practice.\n- Lack of empirical validation: While the theoretical results are impressive, the paper doesn't include experimental results to validate its predictions.\n- Focus on regression: The paper only considers regression tasks, and it's not immediately clear how these results would generalize to classification or more complex tasks.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Setting : The paper overcomes several limiting assumptions made in previous theoretical work, such as requiring very long prompts, orthogonal input data, restrictive initialization conditions, or unrealistically wide models. Key technical innovations include a novel reformulation of the loss function to remove expectations, proving smoothness and the Polyak-Łojasiewicz condition for the loss landscape, and carefully bounding various matrix norms and eigenvalues throughout training.\n\n- Intuition and connection to ridge regression : The analysis provides insight into how the regularization strength in the implicit ridge regression depends on the ratio of the feature dimension m to the number of examples N in each prompt. Also, the authors show that the transformer acquires two types of generalization capabilities: contextual generalization to unseen examples within a task by learning to infer and apply the underlying template function, and generalization to unseen tasks by learning a general strategy (ridge regression) that works across different λ vectors.\n\nThe paper provides rigorous mathematical analysis in a realistic setting. In doing so it offers novel insights, in particular the notion of implicitly learning to perform ridge regression and acquire generalizable knowledge about the representation function.\n\nAlso, the paper provides the first convergence guarantees for transformers with multi-head softmax attention on ICL tasks, which is a significant theoretical contribution.", "weaknesses": "- Limited model architecture (one-layer transformer) which is understandable\n\n- Lack of empirical validation:  The paper is purely theoretical and does not include any experimental results to validate its predictions.\n\n- Focus on regression : The analysis is limited to regression tasks, which represent only a subset of the problems typically addressed by transformers and in-context learning. Many real-world applications involve classification, sequence generation, or more complex structured prediction tasks. It's not immediately clear how the insights about ridge regression would translate to these other task types.\n\nLack of coomparison to other ICL approaches: The paper doesn't provide a comprehensive comparison to other theoretical approaches to in-context learning. While it does mention some previous work, a more in-depth discussion of how this approach relates to or improves upon other theoretical frameworks for ICL would have provided valuable context.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to Reviewer dubr\n\nThank you for your review and positive comments. We've **included some experiments to verify our theoretical findings in the supplementary pdf and the global response**. Below we address your other points. If our responses resolve your concerns, we'd highly appreciate your consideration of increasing your current score. Certainly, please also let us know if you have further questions.\n\n>**simplified model**\n\n- Our analysis indeed focuses on shallow transformers, but we believe our methodology lays a foundation for extending to deeper architectures. The core of our proof technique, which combines the Polyak-Łojasiewicz (PL) condition with smoothness to demonstrate convergence (as detailed in the proof of Theorem 1), provides a framework that could potentially be adapted for deeper transformers. We anticipate that our key results, such as the limit point of the transformer's output (see (19) in Theorem 2), the explanation of the transformer's ICL ability, and the analysis of its generalization capability, would extend to deeper models, albeit with more complex derivations.\n\n- However, extending the proofs to deeper architectures presents significant challenges, particularly in estimating the PL coefficient $\\sigma$ and smoothness coefficient $L$. As evident from our proofs, even for a one-layer transformer, this computation is highly non-trivial and requires a meticulous block-by-block analysis. For deeper models, this analysis would become substantially more intricate, requiring novel mathematical techniques to handle the increased complexity.\n\n- To the best of our knowledge, there's no existing work analyzing deep transformers' training dynamics from an optimization perspective. Most theoretical works, including those in Table 1, focus on shallow transformers due to these challenges. Developing rigorous analytical methods for deep transformers remains an open and exciting direction for future research in this field.\n\n>**focus on regression**\n\n- Our framework considers non-linear regression tasks, which are significantly more complex than the linear tasks analyzed in previous theoretical works (as highlighted in Table 1 of our paper). We expect that such a contribution within regression framework is valuable to the community. \n\n-   A key aspect of our work is the analysis of how transformers learn representations (see Section 3.3). This is fundamental to many other machine learning tasks, including feature extraction, translation, and generalization. Our insights into how transformers extract and memorize \"inherent information\" of basic function maps during training (lines 312-317) could be valuable for understanding other more complex tasks. Besides, our novel insights into how transformers acquire contextual generalization ability with limited data and handle underdetermined templates is relevant to a broader range of in-context learning tasks beyond regression. We leave the extensions of our settings as important directions for future work.", "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": "QE0HciAxTM", "reviewer": "Reviewer_TyrV", "summary": "The paper investigates the theoretical understanding of in-context learning, focusing on whether transformers can generalize to unseen examples within a prompt by acquiring contextual knowledge. The paper analyzes the training dynamics of transformers using non-linear regression tasks, demonstrating that multi-head transformers can predict unlabeled inputs given partially labeled prompts. The study shows that the training loss for a shallow multi-head transformer converges linearly to a global minimum, effectively learning to perform ridge regression. Main contributions include: (1) The paper establishes the convergence guarantee of a shallow transformer with multi-head softmax attention trained with gradient descent on general non-linear regression in-context learning tasks. (2) The paper analyzes the transformer’s behavior at inference time after training, demonstrating that the transformer decides its generating template by performing ridge regression. (3) The analysis framework allows overcoming several assumptions made in previous works, such as the requirement for large prompt lengths, orthogonality of data, restrictive initialization conditions, special structure of the transformer, and super wide models.", "questions": "1. Do you have any empirical results or planned experiments that demonstrate the practical applicability of your theoretical findings?\n\n2. What are the potential impacts of the remaining assumptions and simplifications (e.g., specific initialization conditions, structure of the transformer) on the generalizability of your results?", "limitations": "Please review the weaknesses.", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The focus on analyzing training dynamics and convergence guarantees for transformers in ICL tasks is relatively novel. Prior works often concentrate on empirical performance without delving deeply into the theoretical underpinnings.\nThe combination of gradient descent analysis with multi-head softmax attention in the context of non-linear regression is a creative fusion of ideas that opens up new avenues for understanding how transformers learn contextually.\nBy overcoming assumptions such as the need for large prompt lengths, orthogonality of data, and super wide models, the paper advances the state of the art and provides a more realistic framework for understanding transformer performance.\n\n2. The paper’s analysis of the convergence dynamics and ridge regression behavior of transformers is mathematically robust, with detailed proofs and clear logical progression. \n\n3. The paper is well-organized, with a logical flow from the problem statement to the theoretical analysis and conclusions, making it easy to follow the argumentation.", "weaknesses": "While the paper excels in theoretical analysis, it lacks empirical validation of the proposed concepts and results. It is suggested to include empirical experiments that demonstrate the practical applicability of the theoretical findings would significantly strengthen the paper. For instance, running experiments on standard benchmarks to show how the proposed theoretical insights translate to improved performance in real-world tasks would validate the results more strongly.\n\nThe paper’s discussion on the practical implications of the theoretical findings is somewhat limited. This can make it challenging for practitioners to understand how to apply these insights to real-world problems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to Reviewer TyrV\n\nThank you for your valuable feedback. We agree that empirical validation is crucial and **have included experimental results in the supplementary PDF and the global response**. Note that doing in-context training on real datasets is highly demanding and usually requires very large models trained on a lot of different datasets [3]. Therefore, it's standard practice for the line of works (e.g., [4,5]) that analyze transformers' in-context ability to experiment on the synthetic dataset.  Below we address your other points. If these clarifications resolve your primary concerns, we would highly appreciate your consideration of increasing your score. Certainly, please don't hesitate to request any further clarification.\n\n\n>**(Weaknesses/Q1) discussion on the practical implications of the theoretical findings**\n\nBesides those discussed in the global response, one interesting implication is that our paper suggests assigning different learning rates for different parameter blocks, which has been shown to work well by some very recent work (e.g. [1,2]), where the authors show sgd (which assigns the same learning rate to all blocks) doesn't work well for transformers, but it's also not necessary to keep different learning rates for all parameters as Adam does.\n\n\n\n>**What are the potential impacts of the remaining assumptions and simplifications (e.g., specific initialization conditions, structure of the transformer) on the generalizability of your results?**\n\nThank you for bringing up this important question. We'll incorporate the following discussion in our paper:\n\n- While our current analysis focuses on shallow transformers, we believe our methodology has the potential to be extended to deeper architectures. The core of our proof technique, which combines the Polyak-Łojasiewicz (PL) condition with smoothness to demonstrate convergence (as detailed in the proof of Theorem 1), could potentially be generalized to deeper transformers. Therefore, we expect that our key results, such as the limit point of the output of the transformer (see (19) in Theorem 2), the explanation of the transformer's ICL ability and the analysis of its generalization ability would also apply to deep Transformers. \n\n-  The main challenge that we need to overcome in extending the proofs to deeper architectures lies in estimating the PL coefficient $\\sigma$ and smoothness coefficient $L$. To this end, it is possible to leverage our techniques for one-layer transformer and carefully conduct a block-by-block analysis potentially via induction arguments.\n\n- As we comment in line 217-218 and Proposition 1 in the paper, our initialization condition can be easily achieved via random sampling from Gaussian distribution.\n\n- It's worth noting that we assume Gaussian noise for simplicity and show that the transformer decides its generating template by performing ridge regression (see Section 1.1 and Theorem 2). Our technique could potentially be extended to analyze other types of noise. For example, we could hypothesize that if the noise is Laplacian, the transformer might learn to perform Lasso regression instead. We leave this exploration as future work.\n\n---\n\n[1] Y Zhang et al. (2024). Adam-mini: Use Fewer Learning Rates To Gain More.\n\n[2] Y Zhang et al. (2024). Why Transformers Need Adam: A Hessian Perspective.\n\n[3] S Min (2022). MetaICL: Learning to Learn In Context.\n\n[4] S Garg et al. (2022). What Can Transformers Learn In-Context? A Case Study of Simple Function Classes.\n\n[5] S Chen et al. (2024). Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality.", "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": "FB5aJePCvq", "reviewer": "Reviewer_MzSs", "summary": "This paper tackles the problem of task learning with representations using transformers. It presents a theoretical proof for convergence of training of a single-layer multi-head softmax attention transformer to the global minimum of the loss. Specifically, given $N$ example tokens and their corresponding task outputs, the transformer learns to predict the task output for unseen tokens $K-N$, where $K$ is the size of the dictionary. The paper shows that when trained for sufficiently long, the transformers can perform of ridge regression on the example tokens to label the unseen tokens.", "questions": "(1) Proposition 1 provides a way to make (almost) sure that Assumption 3 holds, but requiring $H \\ge N$ seems very demanding in terms of compute. Can you comment on how to obtain similar results by simpler assumption?\n\n(2) I am curios about the motivation behind your problem formulation, i.e., task-based labelling of a finite dictionary? Is there a real-world problem where this might be of use?", "limitations": "(1) One minor limitation is that the performance of the trained transformer depends on the relation between the number $N$ of example tokens and the dimension $m$ of the representational space (lines 295-299).", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The paper shows that the parameters of a single-layer multi-head softmax attention transformers can be trained at linear convergence rate (Theorem 1) to the global minimizer of the training loss. This result elucidates that the remarkably efficient training that a transformer architecture offers.\n\n(2) The paper proves that once trained for sufficiently long, the transformer performs ridge regression (Theorem 2) on the example tokens to label the query tokens during inference. This is a very promising step towards understanding the generalization capabilities of transformers, which is quite different from earlier works that require the number of example tokens $N$ to be large.\n\n(3) The paper is written crisp and clean, making it very easy to follow.", "weaknesses": "(1) The optimal transformer configuration $(\\boldsymbol{\\theta})$ is not explicitly stated; only the functional form of the output is given, which makes it difficult to the interpret the result in the parameter $(\\boldsymbol{\\theta})$ space.\n\n(2) The result in this paper is about the wide single-layer transformers (more heads), rather than the deep ones (more layers), and relies on Assumption 3, which holds when the number of heads, $H\\ge N$ (Proposition 2). This is restrictive and there might be other ways in which the assumption can hold even if $N > H$. Some empirical demonstrations with shallow transformers on simple experiments might support this better.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to Reviewer MzSs\n\nThank you for your time in reviewing our paper. Below We address your points. If our responses adequately address your concerns, we would be grateful if you could consider increasing your current score. And we are happy to answer your additional questions.\n\n> **(W1) The optimal transformer configuration $(\\theta)$ is not explicitly stated; only the functional form of the output is given, which makes it difficult to the interpret the result in the parameter space.**\n\n- As is stated in line 316 and line 482-485, $\\theta=${$Q_h,w_h$} is optimal if and only if it satisfies:\n$$\n\\forall k: \\sum_{h=1}^H w_{h,k} softmax(V^\\top Q_h v_k)=A_{:,k},$$\nwhere matrix $A$ is defined in (33). \n\n    Note that due to the non-linear nature of the softmax function and the complex interplay between Q_h and w_h, an explicit closed-form solution in the parameter space is not feasible.\n\n- This characterization, while not explicitly in the parameter space, shows that the transformer gains in-context ability by learning to approximate the matrix $A$, which encodes the \"inherent information\" of the tasks (basic function maps), through a combination of its attention mechanism (represented by the softmax terms) and output weights. It's worth noting that this characterization shows the role of multi-head attention in approximating complex functions (see line 330-339) and connects the learned parameters to the underlying structure of the tasks.\n\n> **(W2/Q2) Assumption 3 and the $H \\geq N$ condition**\n\nWe recognize the importance of potentially weakening this assumption and will explore this in future work. Thank you for raising this. Meanwhile, we also want to make the following clarifications.\n\n- To weaken (or avoid) Assumption 3, other proof frameworks are needed. For example, one possible way is to extend the analysis in (Huang et al. 2023) [9] based on induction hypothesis. However, due to the highly non-convexity of (even shallow) transformers, it's somewhat unavoidable to make other strong assumptions. For examples, similar works such as those listed in Table 1 make stronger assumptions than us despite they consider simpler tasks, see Section 1 for details. \n\n- While the above argument indicates $H\\geq N$ may not be necessary to guarantee the convergence of the transformer, it's intuitive that $H$ may need to be large in our setting. This is because shallow transformers have limited expressive power, necessitating more attention heads to solve complex tasks. In fact, we have shown that it's generally impossible to train a transformer to succeed in the non-linear regression tasks we consider at least when $H=1$, see line 330-339 for details.\n\n- The experiment in Figure 3 in the supplementary pdf shows that when $H<N(=30)$, the loss stopped descending when it's far from the minimal value. And it keeps descending when $H=N$ (though slowly).\n\n> **(Q2) motivation/real-world applications**\n\nThank you for your question. We'll incorporate the following in our updated paper:\n\n- The non-linear regression setting is quite general and can model a wide range of real-world problems across various domains, including classification (by using a softmax function as the output layer [1]) time series forecasting [2], inverse problems[3], financial modeling [4], and physical/mathematical modeling [5].\n    \n- A key aspect of our work is the analysis of how transformers learn representations (see Section 3.3). This is fundamental to many complex machine learning tasks, including feature extraction, translation and generalization. Our insights into how transformers extract and memorize \"inherent information\" of basic function maps during training (lines 312-317) could be valuable for understanding these more complex tasks. Besides, our novel insights into how transformers acquire contextual generalization ability with limited data and handle underdetermined templates is relevant to a broader range of ICL tasks beyond regression.\n\n- While our analysis uses a finite dictionary, this approach allows us to derive rigorous theoretical results while still capturing key aspects of how transformers learn and generalize. The insights gained from this analysis can inform our understanding of transformer behavior in more general settings.\n\n> **(limitation) the performance of the trained transformer depends on the relation between $m$ and $N$**\n\n- You're correct that the inference-time performance depends on the $m$-$N$ relationship. This behavior highlights a fundamental trade-off in ICL between having enough examples to determine the underlying function and avoiding overfitting to noisy labels. Far from being a limitation, this is a key finding of our work and reveals nuanced behavior in ICL that differs from traditional supervised learning. \n\n- We've conducted experiments to verify this finding, please refer to Figure 2 in the supplemental pdf and the corresponding descriptions in our global response.\n\n\n\n\n---\n\n[1] Liu, C., Zhu, L., and Belkin, M. (2020). Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. \n\n[2] Karimi et al. (2016). Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Łojasiewicz Condition.\n\n[3] Loizou et al. (2021). Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence.\n\n[4] Fischer, Manfred M. (2015). Neural networks. A class of flexible non-linear models for regression and classification.\n\n[5] F Stulajter (2002). Predictions in Time Series Using Regression Models.\n\n[6] R Nickl (2021). On some information-theoretic aspects of non-linear statistical inverse problems.\n\n[7] T Amemiya (1983). Non-linear regression models.\n\n[8] M Alizamir et al. (2020). A comparative study of several machine learning based non-linear regression methods in estimating solar radiation: Case studies of the USA and Turkey regions.\n\n[9] Y Huang et al. (2023). In-Context Convergence of Transformers.", "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": "TSnNkMrhGp", "reviewer": "Reviewer_vG3D", "summary": "This paper studies the training dynamics of multi-head Transformers by gradient descent on non-linear regression tasks via ICL. This work shows the linear convergence to the global minimum of the training and the in-context inference. An impressive contribution is that transformers are proven to learn contextual information to generalize to unseen examples when prompts contain a small number of input-label pairs. I am willing to update my review if my concerns are addressed.", "questions": "1. Regarding Weakness 1, why do you have $\\eta_Q$ in Eqn 63? I think by (41), (60), and Lemma 4.3 in [Nguyen and Mondelli, 2020], $\\eta_Q$ should be replaced by $1$ in Eqn 63.\n\n2. What is the requirement for the initialization of $w\\_{h,k}$ defined in (7)?\n\n3. It seems there is no requirement or assumption for the $f$ function to be learned. Is your bound derived for the worst case of all the possible $f$?", "limitations": "There is no potential negative societal impact of this work.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The problem to be solved is significant, interesting, and challenging.\n2. The paper is well-written.\n3. The comparison with the related works makes the motivation clear.", "weaknesses": "1. It seems that the assumption that $H\\geq N$ in Proposition 1 is too strong. The single-head works in Table 1 do not need this condition for few-shot ICL (Obviously, for single-head, $1<N$).  I know the problem to solve is different, but this means that Transformers cannot handle long contexts without a large number of heads. \n\n2. No experiments are conducted to verify the theory. It is important to verify the theory by experiments, even with the most simplified Transformer model. For example, can you show if $H<N$, the learning will fail, while if $H\\geq N$, the learning becomes successful by a one-layer Transformer?\n\n3. Some statements related to the existing works are wrong. For example, in [Li et al., 2024], $N=\\mathcal{O}(\\epsilon^{-2}T)$ is the number of data for the pretraining to enable ICL for the Transformer. It is not the number of examples in each prompt as stated in line 129 of the submitted manuscript. The number of examples in each prompt in [Li et al., 2024] is $l\\_{tr}$ and $l\\_{ts}$ in Eqn 9 of Theorem 3.3.\n\n4. I am unsure whether the theoretical proof is correct. For example, I don't know how you derive linear convergence of the loss via only PL condition and the smoothness of the loss function. Please see Question 1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to Reviewer vG3D\n\nThank you for your insightful review. We've **added experiments to validate our theory in the supplementary pdf and our global response**. Below we address your concerns. If our responses resolve your questions, we'd highly appreciate your consideration in raising the score. Certainly, please don't hesitate to request any further clarification.\n\n> **(W1) Assumption $H\\geq N$ too strong. The single-head works in Table 1 do not need this condition for few-shot ICL.**\n\nWe recognize the importance of weakening this assumption and will explore this in future work. Meanwhile, we want to make the following clarifications.\n\n-  As explained in Section 3.3, it is generally impossible to train a single-head shallow transformer to succeed in the task we consider. Our task involves more complex non-linear regression with representation learning, which fundamentally requires multi-head attention.\n\n- Note that $H\\geq N$ is a necessary and sufficient condition for the initialization condition (Assumption 3) to hold. To see why it is necessary, note that $B_k^{(0)}$ is an $N$-by-$H$ matrix, and has full row rank only when $H\\geq N$. \n\n- Assumption 3 is crucial to our proof: it guarantees that $\\zeta_0$ (defined in (12)) is positive, which in turn ensures $\\sigma$ (the PL parameter in (44)) is positive. \n\n- However, it's important to note that Assumption 3 may not be a necessary condition for the convergence of the transformer. To weaken (or avoid) Assumption 3, other proof frameworks are needed. But due to the highly non-convexity of (even shallow) transformers, it's somewhat unavoidable to make other strong assumptions (see the examples we list in the next bullet point).\n\n- Though the single-head works in Table 1 do not need $H\\geq N$, they often rely on other stronger assumptions even for tasks without representation learning, see Section 1 and 3.1 for details. To summarize:\n    \n    - Huang et al. [2023], Li et al. [2024], Nichani et al. [2024] assume that the tokens are pairwise orthogonal, which implies the the token dimension $d$ is no smaller than the dictionary size $K$, while in practice $K\\gg d$.\n    \n    - Huang et al. [2023] and Chen et al. [2024] require the context length to be sufficiently large, so they are not analyzing *few-shot ICL*.\n    \n    - Li et al. [2024] requires the width of their model to be larger than the square of the number of data.\n\n\n\n> **(W1) This means that Transformers cannot handle long contexts without a large number of heads/(W2) No experiments are conducted to verify the theory...can you show if $H<N$, the learning will fail, while if $H\\geq N$, the learning becomes successful by a one-layer Transformer?**\n\n- Our above response indicates $H\\geq N$ is only a sufficient (but may not be necessary) condition used to guarantee the convergence of the transformer. Therefore, our results don't indicate transformers cannot handle long contexts without a large number of heads in practice. However, shallow transformers have limited expressive power, which necessitates more attention heads to solve complex tasks. To support this, we have put empirical results in the supplementary pdf, please refer to Figure 3 and Figure 5 and the corresponding description in our global response.\n\n\n\n\n> **(W3) Statements on [Li et al., 2024] are wrong.**\n\n- You are correct, and we apologize for the mistake and will fix it in the updated paper. Thank you for your careful review. \n\n\n> **(W4/Q1)typo in (63)**\n\n- Thank you for your sharp observation. You are right -- $\\eta_Q$ should be replaced by 1 in (63). We apologize for the typo and will fix it in the updated paper.\n\n- Let us emphasize that this typo does not affect the subsequent results, and the rest of our proof remain valid:\n\n    The revised (63) reads:\n$$\n\\ell(\\xi^{(t)})-\\mathcal{L}^\\star \\leq \\ell(\\xi^{(t-1)})-\\mathcal{L}^\\star+\\langle \\nabla_\\xi \\ell (\\xi^{(t-1)}), \\xi^{(t)}-\\xi^{(t-1)}\\rangle +\\frac{L}{2}||\\xi^{(t)}-\\xi^{(t-1)}||_2^2.\n$$ \n    Combining this with (38), we have\n$$\n\\ell(\\xi^{(t)})-\\mathcal{L}^\\star\\leq \\ell(\\xi^{(t-1)})-\\mathcal{L}^\\star + \\eta_Q(L\\eta_Q/2-1)||\\nabla \\ell (\\xi^{(t-1)})||_F^2.\n$$\n\nWhen $\\eta_Q\\leq 1/L$, we have $L\\eta_Q/2-1\\leq -1/2$, which gives the first part of (64):\n$$\\ell(\\xi^{(t)})-\\mathcal{L}^\\star\\leq \\ell(\\xi^{(t-1)})-\\mathcal{L}^\\star - \\frac{\\eta_Q}{2}||\\nabla_\\xi \\ell (\\xi^{(t-1)})||_F^2.$$\nNote that this is the only place where (63) is used.\n\n>**(W4) derive linear convergence of the loss via only PL condition and the smoothness of the loss function**\n\n- The combination of PL condition and smoothness is indeed sufficient for linear convergence of gradient descent. See for example Theorem 13.2 in this note: https://www.stat.cmu.edu/~siva/teaching/725/lec13.pdf\n\n\n- This approach is commonly used in deep learning optimization literature, as PL condition has been proven to hold for over-parameterized neural networks. See [1-3] for example.\n\n> **(Q2) initialization of $w_{hk}$**\n\n- $w_{hk}$ is initialized to be 0, as stated in line 226 of our paper.\n\n> **Is your bound derived for the worst case of all the possible $f$?**\n\n- our bounds apply for any possible $f$, which includes the worst case one. The key insight is that our theoretical guarantees depend on the function's outputs on the dictionary $\\mathcal{V}$, i.e., in our analysis, we work with the finite-dimensional matrix $\\widehat Z=(f_i(v_j))_{i\\in[m],j\\in[K]}\\in R^{m\\times K}$ defined in line 485, rather than the full functional form of $f$.\n\n---\n\n[1] Liu, C., Zhu, L., and Belkin, M. (2020). Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. \n\n[2] Karimi et al. (2016). Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Łojasiewicz Condition.\n\n[3] Loizou et al. (2021). Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence.", "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:
We conducted the following experiments to validate our theoretical findings. The attached file includes the experimental plots. We will add these results as an experiment section in our revision.

**Set up.** We conduct experiments using synthetic dataset (which is standard practice of this line of works that analyze transformers' in-context ability, see [1,2] for example), where we randomly generate each token $v_k$ and their representation $f(v_k)$ from standard Gaussian.  We experiment either on the 1-layer transformer described in Section 2 or a standard 4-layer transformer in [3] with $d_{model}=256$ and $d_{ff}=512$. For the 1-layer transformer experiments, we set the training loss to be the population loss defined in (8), and initialize $Q_h^{(0)}$ using standard Gaussian and set $w_h^{(0)}$ to be 0 ($h\in[H]$), identical to what is specifed in Section 3. For the experiments on 4-layer transformers, we generate $\lambda$ from standard Gaussian distribution to create the training set with 10000 samples and in-domain test set with 200 samples; we also create an out-of-domain test set with 200 samples by sampling $\lambda$ from $N(1e,4I)$. Given $\lambda$, we generate the label $y_k$ of token $v_k$ using (1), $k\in[K]$. We train with a batch size 256. All experiments use the Adam optimizer with a learning rate 1e-4.

Figure 1 in the attached file shows the training and inference loss of the 1-layer transformer, where we measure the inference loss by $\frac{1}{K}||\hat y-\hat y^\star||_2^2$ to validate (19): after sufficient training, the output of the transformer $\hat y$ converges to $\hat y^\star$. We set $N=30$, $K=100$, $d=100$, $m=20$, $H=30$. We generate $\lambda$ from $N(0.5e,0.01I)$, which is out of the training distribution (Assumption 1). Thus this experiment also shows the transformer's contextual generalization to unseen examples and generalization to unseen tasks, validating our claim in Section 3.2. We also validate the two types of generalization capabilities on the 4-layer transformer, as shown in Figure 4, where we set $K=200$ and the other configurations same as in Figure 1. From Figure 4 we can see that the three curves have the same descending trend, despite the inference loss on the ood dataset is higher than that on the in-domain dataset.

Figure 2 verifies our claim at the end of Section 2: the closer $m$ is to $N$, the better the transformer's choice is. To be specific, we fix $m=100,\tau=0.01$, and plot $\frac{1}{K}||\hat{y}^\star-\hat{y}^{\text{best}}||_2^2$ corresponding to $N$ from 50 to 150. When $N=m=100$, this square norm becomes exactly 0.

Figure 3 validates our claim on the number of attention head $H$ at the end of Section 3.3 using the 1-layer transformer. In this experiment we use different $H$ to plot the training loss curves, and set the other configurations same as those in Figure 1.  From Figure 3 we can see that we need to set $H$ large enough to guarantee the convergence of the training loss. However, setting $H$ too large ($H=400$) leads to instability and divergence of the loss. Recall that in Proposition 1, we require $H\geq N$ to guarantee our convergence results hold. Although this condition may not be necessary, Figure 3 shows that when $H<N=30$, the loss stopped descending when it's far from the minimal value. On the other side, the loss keeps descending when $H=30$ (though slowly).

We also explore how $H$ affects the training on the 4-layer transformer, as displayed in Figure 5, where we set $K=200$ and the configurations other than $H$ are the same as in Figure 3. We fix the wall-clock time to be 100 seconds and plot the training loss curves with different $H$. The left figure shows the final training and losses change with $H$. It reflects that the losses converge faster with smaller $H$ (here the final training loss is the smallest when $H=4$). The right figure of training curves corresponding to different $H$ within 100s may provide some explanation to this phenomenon: (i) transformers with larger $H$ could complete less iterations within a fixed amount of time (the curves corresponding to larger $H$ are shorter); (ii) the training loss curves corresponding to large $H$ ($H=32,64$) descend more slowly. This indicates our claim that larger $H$ may yield slower convergence rate is still valid on deeper transformers. Note that unlike the 1-layer transformer, deeper transformers don't require a large $H$ to guarantee convergence. This is because deep transformers have great expressive power even when $H$ is small. 

---

[1] S Garg et al. (2022). What Can Transformers Learn In-Context? A Case Study of Simple Function Classes.

[2] S Chen et al. (2024). Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality.

[3] A Vaswani et al. (2017). Attention is all you need.

Meta-review: This paper provides a theoretical account of in-context learning in a single-layer transformer performing nonlinear regression tasks. This moves beyond prior theoretical accounts along some dimensions — though assumptions like a single layer are still quite limiting. The reviewers generally agreed that the paper offered a useful contribution despite these limitations, especially after clarification in the rebuttal period, but that the papers impact might be limited because of both its strong assumptions, and the relative lack of empirical validation or further practical insight provided.
Decision: Accept (poster)