Provably Transformers Harness Multi-Concept Word Semantics for Efficient In-Context Learning
Authors: Dake Bu
Wei Huang
Andi Han
Atsushi Nitanda
Taiji Suzuki
Qingfu Zhang
Hau-San Wong
Venue: NeurIPS 2024 poster (2024)
Keywords: In-Context Learning; Learning Theory
Primary area: learning_theory

Abstract: Transformer-based large language models (LLMs) have displayed remarkable creative prowess and emergence capabilities. Existing empirical studies have revealed a strong connection between these LLMs' impressive emergence abilities and their in-context learning (ICL) capacity, allowing them to solve new tasks using only task-specific prompts without further fine-tuning. On the other hand, existing empirical and theoretical studies also show that there is a linear regularity of the multi-concept encoded semantic representation behind transformer-based LLMs. However, existing theoretical work fail to build up an understanding of the connection between this regularity and the innovative power of ICL. Additionally, prior work often focuses on simplified, unrealistic scenarios involving linear transformers or unrealistic loss functions, and they achieve only linear or sub-linear convergence rates. In contrast, this work provides a fine-grained mathematical analysis to show how transformers leverage the multi-concept semantics of words to enable powerful ICL and excellent out-of-distribution ICL abilities, offering insights into how transformers innovate solutions for certain unseen tasks encoded with multiple cross-concept semantics. Inspired by empirical studies on the linear latent geometry of LLMs, the analysis is based on a concept-based low-noise sparse coding prompt model. Leveraging advanced techniques, this work showcases the exponential 0-1 loss convergence over the highly non-convex training dynamics, which pioneeringly incorporates the challenges of softmax self-attention, ReLU-activated MLPs, and cross-entropy loss. Empirical simulations corroborate the theoretical findings.

Provably Transformers Harness Multi-Concept Word Semantics for Efficient In-Context Learning
Dake Bu', Wei Huang", Andi Han?, Atsushi Nitanda 3.4, Taiji Suzuki Qingfu Zhang', Hau-San Wong
Department of Computer Science, City University of Hong Kong. Hong Kong SAR Center for Advanced Intelligence Project, RIKEN, Japan 3 CFAR and IHPC, Agency for Science, Technology and Research (A*STAR), Singapore College of Computing and Data Science, Nanyang Technological University, Singapore *Department of Mathematical Informatics, the University of Tokyo, Japan dakebu2-c@my httpe/eccyyyiityy... cityu. eecicyy..uuu.u......ooooom.eeeee edu. (wei. www.guuanggg.... VI, andi.har feanknr ip, wtt.iitt...ttooo......
atsushi aitandaecfar. a-star. edu. sg. (qingfu. zhang, cshswong @cityu.eau...................aae
Abstract
Transformer-base large language models (LLMs) have displayed remarkable cre- ative prowess and emergence capabilities. Existing empirical studies have revealed strong connection between these LLMs' impressive emergence abilities and their in-context learning (ICL. capacity, allowing them solve new tasks using only task-specific prompts without further fine-tuning. On the other hand, existing empirical and theoretical studies also show that there is a linear regularity the multi-concept encoded semantic representation behind transformer-based LLMs. However, existing theoretical work fail to build up an understanding of the con- nection between this regularity and the innovative power of ICL. Additionally, prior work often focuses on simplified, unrealistic scenarios involving linear trans- formers unrealistic loss functions, and they achieve only linear sub-lineat convergence rates. In contrast, this work provides fine-grained mathematical analysis show how transformers leverage the multi-concept semanties of words to enable powerful ICI. and excellent out-of-distribution ICL. abilities, offering insights into how transformers innovate solutions for certain unseen tasks encoded with multiple cross-concepy semantics. Inspired by empirical studies on the linear latent geometry of LLMs, the analysis based coding prompt model. Leveraging advanced techniques, this work showcases the exponential 0-1 loss convergence over the highly non-convex training dynamics, which pioneeringly incorporates the challenges of softmax self attention ReLU- activated MLPs, and cross-entropy loss. Empirical simulations corroborate the theoretical findings.
concept-based low-noise sparse
1 Introduction
Recently, a variety of transformer-based large language models (LLMs)) have demonstrated remark- able performance across broad spectrum of machine learning tasks, including natural language understanding [1. symbolic reasoning (2), and even heuristics design B 4. One crucial emerging ability of these models their in-context learning (ICL) capacity from few demonstrations and conduct predictions on new queries without requiring any further
(5), which allows them learn
"Coresponding authors
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

fine-tuning. However, the current theoretical understanding of the mechanisms underlying this ICL capability remains limited, leaving the reasons for the remarkable emergence and generalization power of transformer-based L.LMs in unseen ICI. tasks largely unexplained.
In line with traditional topic models 66,
propose that latent concepts / topics underlie natural texts, providing Bayesian inference framework to elucidate the ICL mechanism via Bayesian Model Averaging (BMA) approach. On the other hand, theoretical and empirical studies have shown that transformer-based models exhibit linear geometric regularities in their latent representations as result of concept or topic learning
where the representations within-con conhiip have positive inner products while representations cross-concepts exhibit near- near-orthyr orthogonal relationships. This structured semantic geometry has been well-documented in recent research on pre-trained LLMs TO 13]. However, the connection between this observed multi-concepts latent geometric structure and the LMSs' remarkable ICI. capabilities remains unclear. Separately, recent theoretical analyses have modeled ICI. as a martingale process driven by latent "concept" variables [14] 15]. Yet, these studies have not incorporated the observed multi-concept semantic regularity into their analyses, nor have they discussed the strong out-of-distribution (OOD) ICL abilities exhibited by transformers.
g 10,
[T] 12
Additionally, existing theoretical work on transformer has been conducted on unrealistic, oversimpli- fied settings, such as linear or ReLU transformers [66
17
18 19), MLP-free attention-only models
[6 20). Co-mbiidd softmax attention [9 20
22 23 unrealistic infinite dimensional as-
sumption [4 19 21 24] and impractical loss functions like square loss (9] [6 25 200 (2) and
hinge loss 27 28 Furthermore, existing works have only been able derive linear or sub-linear
convergence rates for the 0-1 loss.
Therefore, there is need for more advanced analysis that can bridge the understanding between the multi-concept semantic regularity and the mechanisms underlying transformer-based ICL. This naturally leads the research question:
Essential Questions
Whether and how do the geometric regularity of the multi-concept-encoded representation facilitate transformer in conducting efficient ICL?
To answer the above question, following the meaningful data modeling ideas in 9 29], we conduct theoretical analysis on a concept-specific sparse coding prompt distribution for classification tasks, where the sparse latent variable encodes the information denoting the word's belonging concept. Importantly, the features in both the word's and label's dictionaries exhibit concept-specific geometric properties witthnn-concee positive inner products and cross cross-concept orthogonal geometric properties that aligns with the findings in 9 10
Our main contributions highlighted as below.
1. First, we provide a comprehensive analysis of the learning dynamics for a two-layer trans- former model, comprising one attention layer followed by ReLU-activated feed-forward network, which is trained using the cross-entropy loss via stochastic gradient descent over a concept-specific sparse coding prompt distribution. Leveraging advanced analytical tech- niques, we showcase the asymptotic properties governing the coupled learning dynamics of the attention and MLP layers.
2. To the best of our knowledge, we are the first to prove an exponential convergence of the 0-1 loss over this challenging setting. Despite the highly non-convex optimization landscape, we demonstrate that the transformer can achieve Bayes optimal test error with just logarithmic number of iterations.
3. We provably show how the multi-concept encoded linear semantic geometry can enable transformer to efficiently perform certain out-of out-of-distribution ICL tasks. This offers an intuitive explanation for why transformer-based LLMs are able to successfully leverage the polysemous nature of words to tackle diverse, unseen concept-specific tasks, aligning well with users' practical experiences. Furthermore, our analysis takes a step forward in providing a potential theoretical underpinning for the innovative capabilities of LLMs, encompassing their ability to achieve cross-concept knowledge intersection. We believe our findings provide an initial positive response to Question 5.1.4 in the ICML 2024 position paper (30,, which asks whether the observed latent geometry of L.Mss can explain their OOD extrapolation abilities.

---

2 Related Work
Theory of Exponential Convergence Rate of Stochastic Gradient Descent. Our analysis of the exponential convergence rate for the 0-1 loss builds upon prior work linking the excess risk and essential supremum norm to exponentially fast convergence under the "hard low-noise condition" This phenomenon has been further explored in more recent studies analyzing the exponential
32).
convergence of stochastic gradient descent (SGD) 33
34
35
36 37, as well as more generalized
settings such as multiclass classification 38 and support vector machines
(39
Feature Learning in Learning Theory. Recent works in learning theory have extensively studied structured data from a feature learning perspective, examinnin NN's feature direction reconstruction
and noise memorization as proxy for training or 0-1 loss convergence [40] 41 42). While prior studies often assumed orthogonal features, recent efforts have analyzed non-orthogonal scenarios to challenging nonlinear Attention-MA. MLLP transformers with non-orthogonal structured data representations.
43 44
Our work extends this
Theory of Transformers and In-Context Learning The literature on Transformers and ICL is wide-ranging, and we will selectively address the most relevant ones. Prior studies have analyzed how transformers learn topic/concept semantics (9) the origins and biases of LLM representations using latent variable models [10]. and ICL. from a model averaging perspective 14] However, albeit incorporating concept variables. these works do not connect the geometric properties of concept- encoded representations to transformers' powerful ICL abilities. Another line of research has studied the learning dynamics of ICL, including analyses of linear transformers 17
QK-combined attention-only models [45]. and multi-head softmax attention over linear regression without MLP (25). Though relevant, these works rely on simplifications and do not notice the connection between semantic regularity and powerful ICL. While (28] also analyzes the learning dynamics of transformers with softmax attention and ReLU MLPs for in-context classification tasks, making the most relevant prior work, our analysis differs in several key aspects. Specifically, (i) they consider orthogonal dictionary learning with a single label vector, in contrast to our non-orthogonal concept-encoded dictionaries for both words and labels; (ii) their technique requires a large batch size (at least & e where e is the test error) and long context lengths, which are not required in our result; and (iii) they utilize impractical hinge loss and only achieve linear convergence without relation to E, whereas we analyze the more practical cross-entropy loss and derive an exponential convergence rate in terms of the test error However, note that this is only informal comparison due the differences in the models and primary findings. A detailed Related Work Section is deferred to Appendix C
3 Problem Setup
Notations. For la and Frobenius norms we utilize || and I| denote their computations.
Considering two series a, and bn we denote a, 0 (b..) if there exists positive constant C' > 0
and N > 0 such that for all n 2 N. Janl S C [b,). Similarly, we denote a, 2 (b.,) if b,, = 0 (a,) holds, and an = θ (b,) if an = O (b.. and an Ω (b.)) both hold. Our 1-) is to denote the indicator variable of an event. In addition, we denote span(U1, 02.... ........ U&) as the linear subspace
spanned by the vectors U1 1/2 ... Uk. and conic("), ....... ux) denotes the conic hull (the set of all non-negative linear combinations) of the vectors U1 ..,..., Ek-
3.1 Data Distribution
The data distribution employed in this study draws inspiration from a range of empirical and theoretical research works p T0
46
47 48) This distribution captures context-awareness and can be viewed as
a specialized prompt version of PLSA [49] and LDA (6). In this distribution, each word and label has multiple feature embeddings, each embedding corresponding to a different concept. This is achieved through the use of sparse latent concept/topic variable, which happened to be particularly adept at representing language polysemy 47. Adhering to the LLM representation explored in [9] 10, the features in both the word and label dictionaries maintain orthogonality across concepts and positive inner products within concepts. Additionally, the distribution incorporates Gaussian noise accounting for linguistic ambiguity or the imperfection of the LLM's 's representation.
Definition 1. Polysemous Word Model (Da: Dy, Dz, De.: DEy)) We assume there exists Ki task- relevant concepts, each characterized by two semantically- oppoosiee word's 's feature vectors H and
3

---

Mk. and their corresponding label's feature vectors 9k, and 9k. Tki E [K]]. There are also K2 task-irrelevant concepts denoted by Vky, Vk2 E [K2]. The word samples a E g and their labels y E R are generated from distributions parameterized by shared latent concept variable
x (21 ... 2K) E {0,1}* (K < dx)) capturing the concept concept-specific information:
2 Dz, En De -O Lg Ey DDD - o0, Layy)
x M= + En Day y Q= + Ey DDy
where the feature dictionary M lut Mi μo M2 , HK, M1, V,, MKK, E R exhibits positive inner products within concepts and orthogonality across concepts, and the label dic- tionary Q = [qr 91 9g 92 :.. 4k. 9K, 00,- (0,0..0). Bdy has similar geometric properties. Specifically, we have Vky E [Ki], 22 E [K2], |
1x |ull, |1 = ||ql]. and there exist constants 0 < Kiz Kry 1 such that O < (1 Pky Ka ||ul| and 0 < (qi, Ix, S kl
The detailed formal definition can be found in Appendix E By this definition, single word or label can possess different features corresponds different concepts. The illustration of Figure in [1] can be an example, where the "Dog" vector in the representation space of LLM ís decomposed to direct sum of orthogonal vectors: "[Animal] [Mammal] + +...., and we can see "[Animal]" belongs to the concept "Organism's Category" categorized into labels "[Animal]" and "[Plant]", and "[Mammal]" belongs to the concept of "Animal's Category" characterized by labels "[Mammal]", "[Fish]", "[Bird]", "[Reptile]" Besides, Figure 1 in [46] can also be good support for our modeling, where "Ferrari" vector consists of "[Cars] + [Italian] + .....
The following definition models the contextual prompts via specifying the statistical property of among in content words, which a special prompt version of PLSA 49] and LDA ( The detailed formal version is available in Appendix E
Definition 2. Concept-specifie Contextual Prompt Distribution During training, each prompt sample S = æ1, "1.... L. YL. I++1 would share at least one co-concept, which is drawn from a mixture distribution Dg defined as:
K1 Ds
=
T PR + T TT.LL.1)
(I)
k=1
where P denotes the k-th concept-specifyy prompt distribution, and TT (2K1) denotes the equal chance of sample to belong to Pk.L+11 Specifically, a sample Sn ~ k E [+] means that the query's 's label yi+1 is qi and we denote ys.. = e as the real value label of this prompt. In addition, every demonstration pairs (xT), y)), E [4] in Pt,
contain either (1 q or (1T , IT with equal chance. Also, every 1 [T++1] 1] would satisfy P(=),-(2k-- =1 K-1 denoting the equal chance to have diverse features other than the current co-concept of the Pr k
This definition suggests that for prompt S sampling from Ds, there exists E [=], k [K,]. such that all the word-label pairs in this prompt share the k-th concept as their co-concept, and the corresponding real value label of the query in this prompt is e. Besides, the real value label of each word-label pair in the demonstration would have equal chance be +1 or -.
3.2 Transformer Model
Following
17 20
28), our embedding E() of prompt formulated as H:
H
E(S)
D2
Fr..
Tqquery
(h,, hg..... hquery ) IR ((L++))))
y2
yL
The learning model is a single-head, lne-l Transformer with one self-attentio layer and one two-layer perceptron. Mathematically, it can be expressed as follows:
[(H; v) = = Fo (Wo (H; V)),
altn altn(H; v) Σ Webios ((Wxhh)) Wohquery)
*Our theory allows for a broader range of the probability settings stated in the training prompt distribution. but for the sake of simplicity in presentation, here chose feasible one.

---

where πn(.) == Relu(-), (el(( == softmax()) WQ. WK E ( Wi E (mmx((rrdy))) are the embedding matrices for queries, keys, and values, respectively, and Wo E Rmxm. and r E IRm are parameters in the MLP layer. Typically, min (Mak: m.) 2 dx + dy. V (We, WK. Wv, Wo, denotes the set of all model weights.
Training Setting, We fix one layer in both the attention and MLP layers to scrutinize the training dynamics more rigorously. Specifically, we let
W ( W
* * *) WK ( W * * * )). Wv
( (v=(= WV ) Wo Wo= = (* W
where W Q: W E Rdxxxx W V. E (( W E Rmxdy, Here, we set the elements other than W Wi., W. and W to be zero. Besides, we fix W to be [(m, (d)=dx)xdy> We sample from a uniform distribution Unif(-1, } and fixed during the training process. Based on this setting, the trainable part we need to consider is actually Tp' = ( W Wy). This problem remains highly non-convex and challenging.
utilize mini-batch with-replacemmnt SGD to train the transformer model. The empirical cross- entropy loss for each batch B, is written as
LB, (V) = LB,(W) (V)) = B
Σ nEB,
(vs.. ((H; + W
where ((2) log(1 exp( -2)). ys. is the real value label of the prompt defined in Definition and the term | represents | Q 1/ + I || which is the Lg regularization term with II- Il denoted as the Frobenius norm. The purpose of the regularization in this paper to accelerate and stabilize the mini-bater with-replacement SGD. The learning step is set to be 7) A(--1)) where y is an offset parameter. This decaying schedule is standard and also used in prior work [3 50 51 studying convergence of SGD. The whole procedure is in Algorithm [|
Initialization Setting. All initial values of W O are sampled from i.i.d. Gaussian distributions with K are diagonal matrices oul. which are
mean O and variance oj. The initialization of W and W
also adopted in other work that consider training W and W k separately [25 28)2..
Testing Setting. The model performance is measured by 0-1 test error test prompt distribution D*:
PS-D- ((ys ((E(S); v) < O].
Algorithm Training algorithm
Input:
Training distribution Ds, Test distribution D Batch size B, step size 1): x(++2)*
stopping criterion ε and total epochs T.
Initialize model parameters (1)
fort 0,. 17-110 do
If L -L- < & stop else continue. Randomly sample mini batches B, of size B from Ds.
Update model parameters: (r(t+1) = (t()) T) V w LB, (vr((t)). end for
4 Theoretical Results
In this section, we present our main theoretical results, which is based on the following conditions.
We consider the learning iterations 0 < S T*, where T" 2mmm- -1 - -m | 1)]lu|| + log(e-1)) denotes the maximum admissible iteration.
Condition 1. Suppose that there exists sufficiently large constant C, such that the following hold:
1. dx, dy 2 max/Clog( (K 10), K). dy 2 C' Cllo(m/6), m Clog(K/6)).
5

---

sufficiently 2 Cmax|lal"/mKK31)).00 X min((Clog(<m//6601l1] (Co0lo/((/////)ll)000//////////0
(Coo/2(lul))
3. 22 (CK,CClul//o;xdxx))
4. S min(Am/(Cvdz dz / (all/CCC
5. 00 $ VK-1 log( AKı log 11
mAKi /)
01 S min((Cooo|||] (5 k w²² (
((Cm3/21lll)))
1- e 00 (1-xx)))))))))) // Here, w 1
Note that we do have any requirement upon demonstration length L and batch size B for training, thus the training can be really flexible compared with the strict requirement in [28]. The condition on dimensionality dx. dy and the network width m ensure the learning problem
The condition on Y ensures the learning step to be small and thus learning process enjoys an approximation to gradient flow. The condition on the small is to ensure the model's sufficient learning before being stuck by regularization (53) The condition on K is control the impact of cross-concept contribution the Attention's i learning dynamic, which can actually be relaxed at the cost of denser analysis. The condition on de is to ensure that the gradient flows be mildly influenced by the noise. Last but not least, the conditions on σ guarantee that the initial beliefs of MLP is small and the gradients of SGD can update the model effectively. A more detailed discussion over the parameter settings delayed to Appendix
overparameterized setting [1 42 52 43
Theorem 2. Exponential Convergence of 0-1 loss. Under Condition 7 define
logg5kkm//b of (1 28 101?
V == min minn (ky///(11Knr Kny),o(( kz))
lul.
Then, for Ve > there exist some positive constants Ci and Cz. with probability no less than for T Co.mAKK kyy log(5K (05(5mm/60 (om0///ww/ (1 kyy Iall, we have
C2t 2m +
(1:(()) I <expl- K1|lg/ (| 2 +1)
Thus after
T. = KgllalI((L - 1)|lu||2 log(
C222mX²
iterations, we have T. (I((T) S.
Note that the bound is valid only when T' 2 T, common threshold in prior convergence rate analyses 34 33 36.. Importantly, the existence of T does not affect the convergence rate + 0, since T is independent of Our novel analysis generalizes these prior results to our realistic settings handling the challenges of self-attention, ReLU-MLP, and cross-entropy loss simultaneously. By considering extreme cases, techniques relax the batch size requirement, enabling more general results. Consequently, the sample complexity for Bayes-optimal test error is N = Te-
Before introducing the next proposition, we highlight key observation from the semantic geometry in Definition] 1 For any k E [K1]. defining ag == (14 + M )/2 and by, == (M Mx, 1/2, we find that for k df k1, {ak, bx, 1 {ax; by: } and (ax, bx ) = O. This structure is exemplified in Figure ((b) of [12), where "[Bird]" consists of orthogonal steering vectors: "plant + animal" and "mammal + bird," corresponding to the concept feature ag and semantic label features by. Here, the term eb,, in #h, determines the label assignment. Similarly. defining Ck2 = (q + 9x )/2 and de, = (q. 4x,) 3/2 yields analogous properties. Detailed definitions are provided in Appendix The following proposition explores the model's ability to handle OOD unseen ICL tasks.
During testing, the learned model admits probability distribution shift on D and data shift on D x Dī, to generate new prompt distribution
Proposition 1. Saaa---Dsssibuttonneeeereeee
Herr we do not consider the shift of Dz., DE, for the case of presentation. However, we assert that this can also be addressed by leveraging high- -ige-imension statistical analysis other bell-behaved noise distributions.

---

D's = E T S +k + F " "P
Specifically, the new D's satisfies the following
properties.
The prompt length Lt can be any positive integer.
D can enjoy arbitrary distribution, satisfying that each prompt has at least co-concept k E [K]] at least one pair shares the query word's 's co-concept's 's label, and still each word has equal chance have positive or negative semantic labels over its concepts
D Dÿ can enjoy a great greatfamily of data shift. Vk: # W E [K1),k2 E [K2).
we can have new M and Q" such that jog. =b 1- =c Vkz V/y Here, ai big G dj. are any vectors belong to the conic hulls of {au Jk-11 ( {ck]k- j-1k {dx respectively, satisfying ||ba]| || 2 || = ([] and ||djll 2
l = (ll)). Vig = O([[u])) are any
vectors from the complement space of span
Again, the learned model satisfies (.p((()( S.
This proposition demonstrates the strong cuttof-DDstributionn Generalization ability of transformer utilizing
semantics, suggesting the efficiency transformer to conduct unseen ICL tasks just by its learned "Knowledge" on the high-level concept and low-level label semantic information from the two non-orthogonal dictionaries. The admit of shift for D denotes that each prompt can enjoy multi-co-concepts and each word-label pair can appear in at least ||=||o concept-specific prompts/tasks' distribution, which aligns the real-world cases. On the other hand, we also believe the admit of shift for D x DD is inspiring, suggesting that transformer can conduct specific cross-concept semantic "Knowledge Intersection". As such, this lemma suggest that the transformer can master the regularity of unseen ICL tasks' "structure" in the presence the multi-concept encoded representation.
Remark 1. Comparison with Related Work. Theorem 3.4 in [28] and Theorem 2 in [54) address the transformer's OOD capability in specific structured ICI. classification and regression tasks. Our results differ by focusing on compositional generalization of learned concepts, grounded in the concept-specifer linear latent geometry observed in LLMs.
5
Proof Idea
In a big picture. we simply extend standard expectation-variane reduction techniques setting. Section [. defines coefficients to examine NN's expected projection along feature directions. Section 5.2 provides the convergence of the expected estimator through the lens of coefficient evolution; Section 5.3 showcase the exponential convergence by treating the conditional expectations of the NNs as Doob martingales and exploiting the property of the tails under low-noise conditions.
+4 to our
5.1 Idempotent Operator Techniques
Idempotent Operator Trick. Define U span(M) and its complement space U-_ By definition, we know that dim(U) (aa,, =1 [bk (vK, Ka = {u.) )==1"== ddx-K) } be the set of standard orthogonal basis for Rdx where ut " ... ... _ax- K are the standard orthogonal basis U+
K and dim(U-) = dx K. Then we can let
Then we can derive an idempotent decomposition of the identity matrix
K1
a.a,
K
b.b.
T
K
K
T
2
u
u
Ldxxxx-
w=1
Similar techniques are also applied to the label's dictionary: Q = span Q). where we define ....... ddy-K, as the standard orthogonal basis of the complement space Q In our subsequent derivation, the expectation E|.] is taken over the stochastic gradient descent. Similar to the idea in [34 33 (36), we first serve to see how E( (()(()) evolves. For IE ((()()), every gradient descent update by all concept's samples within soft "weight", and thus the analysis is equivalent to gradient descent
'The requirement of 22 could be relax with stricter requirement 7 and denser analyses.
7

---

H R
W, granlleee)
Depcy RM
Wo: gensn
R
R
ReLU
W
Xi
softmma
yj
Wx"
W
I
X;
Yi
Prowe of wh Whit gine
Orthugonal Components
I
Decompusition
Wi= : daeusd
|
B/ "1
Tex
M.
t....
d
Idempotent Decomposition
Pb
B.A. [6.]1
B B
B|
of.
B W
WW
B5
(e5)
aga
aus
Po.
Figure 1: Illustration of our Idempotent Operator Techniques. This allows us to focus on analyzing the evolving coefficients, which are key to the expected ()-1 loss convergence.
with an ddeally-balancee prompt set. Leveraging the symmetry of the prompt distribution, as well as the symmetry of w and w introduce the following decompositions.
Lemma We can decompose E B[ E. El and the i-th row of IE [ (i E [ml) via the following (scaled) projection matrices and projection directions.
W (4)
(t) a,a,
la.1" 4
T
((t) Q
b,b,
b
Qr
(t)
V
u
4
dx K
Σ
PQ.
(8)
ti u
K₁
K2
b.b." T
K₂
dx
E[ [9]
(0)
a,a,
((0)
(1)
V,V,
K,a
||,||
+
K,s
1
K.
||u]|
PK,00
u u
s=1
r=1
W O(L)
(z.)
Ki
(
|
T
+
K
(1)
O
| 1
d2
dy
Ki
(0) Po
qw
k=
k=
represent the expected concept learning process, 8(1) Here Q K. and O
Q.a" K. and
elt)
Pu.,, represent the expected concept-specific semantic learning process and TQ. Kr PQ K
(1)
(e) (
(8)
and P represent the expected memorization of the concept irrelevant noise. It holds that
B[(WT (0)
4 "+= E (1)
(t) [4]] = aQ. K.
1tts//la./ /la..ll.22 + 260 //..//11b1111 lbbl?
Q
(4)
E(WY O
(1) t a0 0. + (00)
(8)
(0,)
for Ve E [+),i E [m]. k [K]] and for Ve' E [+] E [Ki,,r E [K2]. w [dx - K]. Vu
(uo v,, ). holds that B[( (9 [ Q p = O. Similar conclusions hold when the
query vectors are and w Vr E [K2], [xx K]. As such, remaining task is to scrutinize the coefficients evolution, which would be the key contributors to the expected 0-1 loss convergence.
5.2 Convergence of the Expectation
WF, (1) 45.55 (1) as the activated neuron set for {i [m] rryss. > 0} and Denote 2 (1) and We
{i [m] r.ys,, 0} separately, and Es (o (t), represents the correct attention weight,
m.k
where the detailed definitions are delayed in Appendix E We then introduce the following lemma. Lemma 2. Under Condition 7 when
Σ
Σ iEW"Sa USa (1)-21
Oounss +(
Σ
(0 n 1)ys, 800) Po 0, (5)
" USn (0)
S n.
holds, we have Lo (B() 0.

---

Training and Test Loss
Evelution of Attention Weights
Leariing Progress of Attention
Leaming Progress of MLP
70
60-
40
6/"
0.3
Test 0-1 Loss
30
20
Hotrs Froeutt Pax
6.7
weggtt
and (31
75
100
Egoch
Figure 2: Learning dynamics: (i) training and test loss; (ii) correct attention weight; (iii) maximum values of dxg,a <a Sq.a* 3K maximum values of complement products TQpr TK,r or pq.2 PK,2, and maximum values of produut-witt-noise (.) W&q (iv) maximum values of (. v and (80 a. maximum
values of the complement coefficients pou.) w and maximum values of prodhct-with-noise W
As such, the following lemmas show the learning outcomes of the ((D(*)) along the iterations.
Lemma 3. (Convergence of the Expectation). There exist constant C1 0, vt 2 f C1ammXK1 Kyy) (o555k//6)/u"" Lemma 4. (Regularizing the models). Under Condition it holds that
klll we have L. ( (E((T)))
@[[uu]] log( ||u||
log( ||al? mXK,
Q = K.K = =[EE
(T")
(T*)
(.
Q.
(T))
((T))
XK1
})),
(T") a
S 13000 O(log(
|ql 2
BI(
Σ
(os
(0 )
O
mAK
1
AR I 2
log(
mAki
jES
In addition, our analysis provides three asymptotic properties of the coefficients evolution, which are delayed Appendix I.1.3 and 1.2 for room limitation.
5.3 Exponential Convergence of 0-1 loss
Proposition 2. At 2 T, when [((( E(V'")) holds, we have L (p=((t)/t)) 0. Here, |/ = W | + W
By definition of 0-1 loss, then we only need to prove the 0-1 loss convergence by seeing the speed of (() converging to (())(()) with an error of in terms of || ||..
Drawing insights from 34 we see By.... Br-1 as i.id random variables following the same distribution. Then vt E
{0. ... T), it holds that
D E W (T+1) Bu, ...
B.] E W (T+1)
Bo ... BB-11h
D E W
(T-1)
Bo, B.]] W (T+1) Bp ,B
(6)
K
D E (T+1) Bu, B.] E(wy (T+1)
Bo.... ,B.1),
are martingale difference sequences, and for VX E {Q, K, and its corresponding W E (W&, W Wy ]. we have I D3 W(++1) E E[W(T+ Then utilize the follow- ing lemma in 34 55 to give bound over the variance.
Lemma 5. Let DI. DT-1 be a martingale difference sequence. Suppose Eer such that Ete Drll < 27 where ao the essential supremum of ||r. Then for > 0, we have
2
P
sup SE
Σ Dill 2 5] < 2 2exp
20
1=0
Therefore, we need to see if there exists decaying positive constant er (with decaying rate ((1/T9), 0), such that I lI&
v vX (Q, K, O}, where || |loe is the essential

---

43
(a) OOD Scenario 1()): L' during testing.
(b) OOD Scenario 1(i): I," during testing.
(c) OOD Scenario 2: 0.8 fraction for concept 0 and 0.2 (d) OOD Scenario 3: Shift the data as μ fraction for concept during testing.
+ b2
and 2 02 = b, during testing.
Learning dynamic three OOD scenarios. The training settings and plotting methods are identical to those used Figure and the testing settings are: (a-b) utilizes different prompt lengths; (c) adopts skewed distribution over 2; (d) switches the concept sppecffc semantic features.
Figure 3:
supremum of || Subsequently, by controlling the martingale sequence norm tail similarly in we can obtain an exponential convergence rate after Ti.
34 55,,
For W ( W W J. to check the decaying eT, adopt the techniques of [34] 33 36 in the following manner. Let B,' be an independent variable from Bo, ....... By and let W (T+1) be an output of the algorithm depending on (Bu. BB',,,B'' BB+1, ,Br). Then we have
| S [[][ W. (7,0+)1110 1|x0 |Bo, , B.).
Therefore, one may estimate ox² by bounding [ W, (2) ( Il uniformly w.r.t. By.... BB--1 Such bound can be derived utilizing stability property of stochastic gradient descent [34]5]]. the OOD scenario, since we require the data shift to be via conic combination, the new words and labels in each prompt will share the positive/negative real-valued label without any self-conffict. The norm requirements and constraints on D would ensure the Gaussian noise, concepts other than the co-concepts, and probability shifts have limited influence on the prediction compared with the considerable scale of coefficients by Lemma laying the groundwork for the proof.
56]. For
Experiments
In this section, we demonstrate the validity of our theoretical analysis through simulations of Algorithm _
We use the following parameter settings in Figure 2 The parameter settings are:
the length L= 4. the number of co-concepts K₁ 2, dictionary size K 104, the number of test instances Nsst = 5000, dimension dxx = dy = 1000, MLP width m = 50, feature strengths ||||| = |/ 10, vk [K]], the cosine (HE f4 ) ///22 2 (9 9x ) |||2 0.5, the initialization
parameters Tg = 0.1, Ji = 0.01, and the noise deviation dg = 0.01. For the optimization, we use X 0.002, B 16,
10000, and the total training epochs is 100. Figure 3 (a-d) uses the same training settings, but during testing, it applies different configurations: (a) T.* = 5, (b)
I 2, (c) 0.8 fraction for the first concept and 0.2 fraction for the second concepts, and (d) μi ai + b2, 1 a2 + bj- Figure validates our Theorem and Lemma which showcases the fast convergence rate and the evolution of coefficients. Figure validates Proposition T where the learned model permits certain data shifts.
7 Conclusion
This work provides the first exponential convergence analysis of 0-1 loss for transformers with softmax attention and ReLU-MLP, trained on a non-orthogonal concept-specific prompt distribution by practical cross-entropy loss. Furthermore, the results demonstrate transformers can perform certain OOD ICI. tasks by leveraging the multi-concept semantic linearity, highlighting their innovative potential, An important future direction is extend the analysis to more complex scenarios.
10

---

8 Acknowledgment
We thank the anonymous reviewers for their instrumental comments. D.B. and H.W. are supported in part by the Research Grants Council of the Hong Kong Special Administration Region (Project No. CityU 1 11206622). W.H. is supported in part by JSPS KAKENHI (24K20848). A.N. is supported in part by National Research Foundation, Singapore and Infocomm Media Development Authority under its Trust Tech Funding Initiative, the Centre for Frontier Artificial Intelligence Research, Institute of High Performance Computing, Starar and the College of Computing and Data Science at Nanyang Technological University. T.S. supported in part by JSPS KAKENHI (24K02905) and JST CREST (JPMJCR2115, JPMJCR2015).
References
[1] Li Dong, Nan Yang, Wenhui Wang, Furu Wei, Xiaodong Liu, Yu Wang, Jianfeng Gao, Ming Zhou, and Hsiao- Hon. Unified language model pre-training for natural language un- derstanding and generation. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'AlchÊ-Buc. E. Fox, and R. Gamett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[2] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, volume 35, pages 24824-24837, 2022.
[3] Fei Liu, Xialiang Tong, Mingxuan Yuan, Xi Lin, Fu Luo, Zhenkun Wang. Zhichao Lu, and Qingfu Zhang. Evolution of heuristics: Towards efficient automatic algorithm design using large language model. arXiv preprint arXiv:2401.00 000051, 2024.
[4] Fei Liu, Yiming Yao, Ping Guo, Zhiyuan Yang. Xi Lin, Xialiang Tong, Mingxuan Yuan, Zhichao Lu, Zhenkun Wang, and Qingfu Zhang. A systematic survey on large language models for algorithm design. arXiv preprint arXiv: 2410..4716, 2024.
[5] Sheng Lu, Irina Bigoulaeva, Rachneet Sachdeva, Harish Tayyar Madabushi, and Iryna Gurevych. Are emergent abilities in large language models just in-context learning? arXiv preprint arXiv: 2309.01800 2023.
[6] David Blei, Andrew Ng. and Michael Jordan. Latent dirichlet allocation. In Advances in Neural Information Processing Systems, 2001.
[7] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference, 2022.
[8] Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning. In Workshop on Efficient Systems for Foundation Models @ ICML2023, 2023.
[9] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In Proceedings of the 40th International Conference on Machine Learning, pages 19689-19729, 2023.
[10] Yibo Jiang, Goutham Rajendran, Pradeep Ravikumar, Bryon Aragam, and Victor Veitch. On the origins of linear representations in large language models. arXiv preprint arXiv: 2403.03867, 2024.
[11] Kiho Park, Yo Joong Choe, and Victor Veitch. The linear representation hypothesis and the geometry of large language models. arXiv preprint arXiv: 10368,,55 2023.
[12] Kiho Park, Yo Joong Choe, Yibo Jiang, and Victor Veitch. The geometry of categorical and hierarchical concepts in large language models. arXiv preprint arXiv: 2406.01506, 2024.
[13] Yibo Jiang. Goutham Rajendran, Pradeep Ravikumar, and Bryon Aragam. Do 1lms dream of elephants (when told not to)? latent concept association and associative memory in transformers. arXiv preprint arXiv: 2406.18400 2024.
[14] Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv: 2305.19420 2023.
11

---

[15] Fabian Falck, Ziyu Wang. and Christopher C. Holmes. Are large language models bayesian? a martingale perspective on in-context learning. In ICLR 2024 Workshop on Secure and Trustworhy Large Language Models, 2024.
[16] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference on Machine Learning, volume 202, pages 35151-35174. PMLR, 2023.
[17] Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv: 2306.09927, 2023.
[18] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Advances in Neural Information Processing Systems, volume 36, pages 57125-57211, 2023.
[19] Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv: 2402.01258, 2024.
[20] Yu Huang. Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv: 2310.05249, 2023.
[21] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In Advances in Neural Information Processing Systems, volume 36, pages 71911-71944 2023.
[22] Yingcong Li, Yixiao Huang, Muhammed E. Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with self-attention. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238, pages 685-693, 2024.
[23] Chenyu Zheng, Wei Huang, Rongzhen Wang, Guoqiang Wu, Jun Zhu, and Chongxuan Li. On mesa-optimization in autoregressively trained transformers: Emergence and capability. arXiv preprint arXiv:2405. 16845, 2024.
[24] Shokichi Takakura and Taiji Suzuki. Approximation and estimation ability of transformers for sequence-to-sen io-sequence functions with infinite dimensional input. arXiv preprint arXiv: 18699, 2023.
[25] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi- head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv: 2402. 19442, 2024.
[26] Yu Huang. Zixin Wen, Yuejie Chi, and Yingbin Liang. Transformers provably learn feature- position correlations in masked image modeling. arXiv preprint arXiv: 2403.02233, 2024.
[27] Hongkang Li, Meng Wang, Sijia Liu, and Pin yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity, arXiv preprint arXiv:2302. 06015, 2023.
[28] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. How do nonlinear transformers learn and generalize in in-context learning? arXiv preprint arXiv: 2402. 15607, 2024.
[29] Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self supersupe contrastive learning. In Proceedings of the 38th International Conference on Machine Learning, volume 139, pages 11112-11122. PMLR, 2021.
[30] Patrik Reizinger, Szilvia Ujváry, Anna Mészáros, Anna Kerekes, Wieland Brendel, and Ferenc Huszár. Position: Understanding LLMs requires more than statistical generalization. In P Proceedings of the s International Conference on Machine Learning, volume 235, pages
42365-42390, 2024.
[31] Enno Mammen and Alexandre B Tsybakov. Smooth discrimination analysis. The Annals of Statistics, 27(6): 1111888881,,,,11111111111111 1999.
[32] Pascal Massart and Élodie Nédélec. Risk Bounds for Statistical Learning. The Annals of Statistics, 34(5):2326 2366, 2006.
[33] Loucas Pillaud-Vivie Alessandro Rudi, and Francis Bach. Exponential convergence of testing error for stochastic gradient methods. Proceedings of the 31st Conference On Learning Theory, volume 75, pages 250-296, 2018
12

---

[34] Atsushi Nitanda and Taiji Suzuki. Stochastic gradient descent with exponential convergence rates of expected classification errors. In Proceedings of the Twenty-Second International Conference Artificial Intelligence and Statistics, volume 89, pages 1417-1426, 2019.
[35] Vivien A Cabannes, Francis Bach, and Alessandro Rudi. Fast rates for structured prediction. In Proceedings of Thirtyy Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages 823-865, PMLR, 15-19 Aug 2021.
[36] Shingo Yashima, Atsushi Nitanda, and Taiji Suzuki. Exponential convergence rates of classifica- tion errors on learning with sgd and random features. In Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130, pages 1954-1962, 2021.
[37] Kazusato Oko, Taiji Suzuki, Atsushi Nitanda, and Denny Wu. Particle stochastic dual coordinate ascent: Exponential convergent algorithm for mean field neural network optimization. In International Conference on Learning Representations, 2022.
[38] Stefano Vigogna, Giacomo Meanti, Ernesto De Vito, and Lorenzo Rosasco. Multiclass learning with Margin: Exponential rates with no bias- -tarinariance trade-off. In Proceedings of the 39th In- terattional Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pages 12260-222661 PMLR, 17-23 Jul 2022.
[39] Vivien Cabannnes and Stefano Vigogna. A case of exponential convergence rates for svm. In Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206, pages 359-374. PMLR, 25-27 Apr 2023.
[40] Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. In The Eleventh International Conference on Learning Representations, 2023.
[41] Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer con- volutional neural networks. In Advances in Neural Information Processing Systems, volume 35, pages 25237-25250 2022.
[42] Yiwen Kou, Zixiang Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting in two-layer reI.U convolutional neural networks. In Proceedings of the 40th International Conference on Machine Learning, volume 202, pages 17615- 117599, 2023.
[43] Xuran Meng, Difan Zou, and Yuan Cao. Benign overfitting in two-layer relu convolutional neural networks for XOR data. arXiv preprint arXiv: 23/0.01975, 2023.
[44] Zhiwei Xu, Yutong Wang. Spencer Frei, Gal Vardi, and Wei Hu. Benign overfitting and grokking in reLU networks for XOR cluster data, arXiv preprint arxiv: 2310.02541, 2023.
[45] Wei Huang, Yuan Cao, Haonan Wang, Xin Cao, and Taiji Suzuki. Graph neural networks provably benefit from structural information: A feature learning perspective. arXiv preprini arXiv: 2306. 13926, 2023.
[46] Hiroaki Yamagiwa, Momose Oyama, and Hidetoshi Shimodaira. Discovering universal geome- try in embeddings with ica. arXiv preprint arXiv: 2305,13175, 2023.
[47] Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In Proceedings of the 38th International Conference on Machine Learning. pages 11112-11 122-1122 1122, 2021.
[48] Chi Han, Jialiang Xu, Manling Li, Yi Fung, Chenkai Sun, Nan Jiang, Tarek Abdelzaher, and Heng Ji. Word embeddings are steers for language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 16410-16430, 2024.
[49] Thomas Hofmann. Probabilistic latent semantic indexing. In Proceedings of the 22nd Annual International ACM SIGIR Conference on Research and Development in Information Retrieval, page 50-57, 1999.
[50] Lřón Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223-31, 2018.
[51] Atsushi Nitanda and Taiji Suzuki. Optimal rates for averaged stochastic gradient descent under neural tangent kernel regime. arXiv preprint arXiv: 2006.12297 2021.
13

---

[52] Yiwen Kou, Zixiang Chen, Yuan Cao, and Quanquan Gu. How does semi-supervise learning with pseudo-labelers work? a case study. In The Eleventh International Conference on Learning Representations, 2023.
[53] Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. Understanding the generalization of adam in learning neural networks with proper regularization. In The Eleventh International Conference on Learning Representations, 2023.
[54] Tong Yang. Yu Huang, Yingbin Liang, and Yuejie Chi. In-context learning with representations: Contextual generalization of trained transformers. arXiv preprint arXiv: 2408.10147, 2024.
[55] Iosif Pinelis. Optimum bounds for the distributions of martingales in banach spaces. The Annals of Probability, 2274):167- 1994.
[56] Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: stochastic gradient descent. In Proceedings of The 33rd International Conference on Machine Learning, volume 48, pages 1225-1234, 2016.
Stability of
[57] Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. The benefits of mixup for feature learning. In Proceedings of the 40th International Conference on Machine Learning, volume 202, pages 43423-43477 2023.
(58) Jinghui Chen, Yuan Cao, and Quanquan Gu. Benign overfitting in adversarially robust linear classification. In Robin Evans and Ilya Shpitser, editors, Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence, volume 216, pages 313-323, 2023.
[59] Spencer Frei, Niladri Chatterji, and Peter Bartlett. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Proceedings of Thirty Fifth Conference on Learning Theory, volume 178, pages 2668-2703, 2022.
[60] Spencer Frei, Gal Vardi, Peter Bartlett, and Nathan Srebro. Benign overfitting in linear classifiers and leaky relu networks from kkt conditions for margin maximization. In Proceedings of Thirty Sixth Conference on Learning Theory, volume 195, pages 1177-2228, 2023.
[61] Yiwen Kou, Zixiang Chen, and Quanquan Gu. Implicit Bias of Gradient Descent for Two- layer ReLU and Leaky ReLU Networks on Nearoy-orthogonal Data. In Advances in Neural Information Processing Systems, volume 36, pages 30167-30221. Curran Associates, Inc., 2023.
[62] Wei Huang, Ye Shi, Zhongyi Cai, and Taiji Suzuki. Understanding convergence and gener- alization in federated learning through feature learning theory. In The Twelfth International Conference on Learning Representations, 2024.
[63] Dake Bu, Wei Huang, Taiji Suzuki, Ji Cheng, Qingfu Zhang, Zhiqiang Xu, and Hau-San Wong. Provably neural active learning succeeds via prioritizing perplexing samples. In Proceedings of the 41st International Conference on Machine Learning, volume 235, pages 4642-4695, 2024.
[64] Yiwen Kou, Zixiang Chen, Quanquan Gu, and Sham M. Kakade. Matching the statistical query lower bound for k-sparse parity problems with stochastic gradient descent. arXiv preprint arXiv: 2404. /2376, 2024.
[65] Alexander Tsigler. Benign Overfitting in Linear Regression and Classification. PhD thesis, UC Berkeley, 2024.
[66] Junhyung Park, Patrick Bloebaum, and Shiva Prasad Kasiviswanathan. Benign overfitting for regression with trained two-layer relu networks. arXiv preprint arXiv: 2410.06191, 2024.
[67] Eshaan Nichani, Alex Damian, and Jason D. Lee. Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks. In Advances in Neural Information Processing Systems, volume 36, pages 10828-10875, 2023.
[68] Jason D. Lee, Kazusato Oko. Taiji Suzuki, and Denny Wu. Neural network lears low- dimensional polynomials with sgd near the information -heeorii limit. arXiv preprint arX15.2406.01581, 2024.
[69] Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Learning sum of diverse features: computational hardness and efficient gradient-based training for ridge combinations. arXiv preprins arXiv:2406. / 1828, 2024.
[70] Yunwei Ren and Jason D. Lee. Learning orthogonal multi-index models: A fine-grained information exponent analysis. arXiv preprint arXiv:2410.09673 2024.
14

---

[71] Spencer Frei and Gal Vardi. Trained transformer classifiers generalize and exhibit benign overfitting in-context. arXiv preprint arXiv:2410.01774, 2024.
[72] Wei Shen, Ruida Zhou, Jing Yang, and Cong Shen. On the training convergence of transformers for in-context classification. arXiv preprint arXiv:2410. 1778, 2024.
[73] Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du. JoMA: Demystifying multilayer transformers via joint dynamics of MLP and attention. In The Twelfth International Conference on Learning Representations, 2024.
[74] Yu Huang, Zixin Wen, Yuejie Chi, and Yingbin Liang. How transformers learn diverse attention correlations in masked vision pretraining. arXiv preprint arXiv: 2403.02233 2024.
[75] Masahiro Sakamoto and Hitomi Sato. Benign or not-benign overfitting in token selection of attention mechanism. arXiv preprint arXiv:2409.11 17625, 2024.
[76] Roey Magen, Shuning Shang, Zhiwei Xu, Spencer Frei, Wei Hu, and Gal Vardi. Benign overfitting in single-head attention. arXiv preprint (rxXiv:44000077444 2024.
[77] Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. arXiv preprini arXiv: 14735, 2024.
[78] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Prov- able training dynamics and feature learning in transformers. arXiv preprint arXiv: 2409.10559 2024.
[79] Hongru Yang, Bhavya Kailkhura, Zhangyang Wang. and Yingbin Liang. Training dynamics of transformers to recognize word co-occurrence via gradient flow analysis. arXiv preprint arXiv:2410.09605, 2024.
[80] Jiarui Jiang, Wei Huang, Miao Zhang, Taiji Suzuki, and Liqiang Nie. Unveil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization. arXiv preprint arXiv:2409. 19345, 2024.
[81] Bingrui Li, Wei Huang, Andi Han, Zhanpeng Zhou, Taiji Suzuki, Jun Zhu, and Jianfei Chen. On the optimization and generalization of two-layer transformers with sign gradient descent. arXiv preprint 21Xi::2410.04870, 2024.
[82]
Yoshua Bengio. Learning Deep Architectures for AI. Now Publishers Inc, 2009.
[83]
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016.
[84]
Zeyuan Allen-Zhu and Yuanzhi Li. Backward Feature Correction: How Deep Learning Performs
Deep Learning. In Conference on Learning Theory, COLT '23, 2023, Full version available at http://arxiv.or//abs/2000
10444
[85] Core Francisco Park, Maya Okawa, Andrew Lee, Ekdeep Singh Lubana, and Hidenori Tanaka. Emergence of hidden capabilities: Exploring learning dynamics in concept space. arXiv preprint arXiv:2406. 19370, 2024.
[86] Yongyi Yang, Core Francisco Park, Ekdeep Singh Lubana, Maya Okawa, Wei Hu, and Hidenori Tanaka. Dynamics of concept learning and compositional generalization. arXiv preprint 24X1v.2410008309. 2024.
[87] Lingjing Kong, Guangyi Chen, Biwei Huang, Eric P. Xing, Yuejie Chi, and Kun Zhang. Learning discrete concepts latent hierarchical models. arXiv preprint arXiv: 2406.005/9 2024.
[88) Zeyuan Allen-Zhu and Yuanzhi Li. Physics of Language Models: Part Learning Hierarchical
I Language Structures. ArXiv e-prints, abs/2305. 13673, May 2023. Full version available at http://arxiv.oog/abs/233...3335
[89] Michael Hahn and Navin Goyal. A theory of emergent in-context learning implicit structure induction. arXiv preprint arXiv: 2303.07971, 2023.
[90] Emanuele Marconato, Sébastien Lachapelle, Sebastian Weichwald, and Luigi Gresele. All or none: Identifiable linear properties of next-token predictors in language modeling. arXiv preprint 1rXXi:2440.23500) 2024.
[91] Roger A. Horn and Charles R. Johnson. Matrix Analysis. Cambridge university press, 2012.
[92] Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
15

---

[93] Miao Lu, Beining Wu, Xiaodong Yang, and Difan Zou. Benign oscillation of stochastic gradient descent with large learning rate. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
[94] Yurii Nesterov. Introductory Lectures on Convex Optimization: A Basic Course, volume 87. Springer Science & Business Media, 2013.
[95] Karan Girotra, Lennart Meincke, Christian Terwiesch, and Karl T. Ulrich. Ideas dimes dozen: Large language models for idea generation in innovation. SSRN, 2023.
[96] Anil Rajnikant Doshi and Oliver Hauser. Generative artificial intelligence enhances creativity but reduces the diversity of novel content. SSRN, 2023.
[97] Fei Liu, Xialiang Tong, Mingxuan Yuan, and Qingfu Zhang. Algorithm evolution using large language model. arXiv preprint arXiv: 23/1.15249, 2023.
[98] Yiming Yao, Fei Liu, Ji Cheng, and Qingfu Zhang. Evolve cost-aware acquisition functions using large language models. arXiv preprint arXiv: 2404.16906 2024.
[99] Fei Liu, Xialiang Tong. Mingxuan Yuan, Xi Lin, Fu Luo, Zhenkun Wang, Zhichao Lu, and Qingfu Zhang. An example of evolutionary computation + large language model beating human: Design of efficient guided local search. arXiv preprint arXiv: 2401.02051, 2024.
16

---

A Limitation and Broader Impact
The theoretical analysis provided in this work introduces novel perspectives optimization and generalization. hut the data model employed may require additional refinements to better align with practical scenarios, such as adding more layers of attention. The techniques and findings can inform future empirical and theoretical explorations of transformer architectures, though we do not foresee a direct social impact arising from the theoretical advancements presented.
B Additional Experiment Details
We implement our methods using Py PyTorch, ensuring consistent software and hardware environments. Specifically, the experiments are run on Linux servers with NVIDIA A100 graphics cards and CUDA 11.2, and can be completed within one hour.
C Additional Related Work
Theory of Convergence Rate of Stochastic Gradient Descent. Our analysis of the exponential convergence rate for the 0-1 loss builds upon a rich body of prior work. the context of classification, the faster convergence rate mostly based on the excess of risk with some power of the essential supremum norm. Specifically, [31] [2] introduce the Hard low-noise condition over the margin. When there is a hard margin separating the classes, the test error can exhibit exponentially fast convergence as the number of training samples increases, even when the surrogate loss error only decreases polynomially. This phenomenon has been further explored in more recent studies. [33 35| [3) 37 have analyzed the exponential convergence stochastic gradient descent under various settings. Meanwhile, (3S have investigated hard-margin and exponential rates in the context of structured prediction, which encompasses traditional classification as special case. Besides, recent work also obtain the exponential rates in generalized settings such Multi-class classification (88 and SVM (39) Building upon this rich theoretical foundation, our work derives the first exponential convergence analysis for the 0-1 loss in the specific setting of transformer models with softmax attention and the sparse coding data model, whose surrogate loss function is the cross-entropy loss.
MLP over
Theory of Feature Learning of GD-updated Neural Network. A rich body of recent learning theory research has focused on the feature direction' recovery view of neural network representations [40]
42
43] 45
52)
133 57
68 6 [t. Rather than directly examining the evolution of the 0-1 loss, this line of work explicitly studies the process of reconstruction of data's 's feature directions and memorization of disrupted noise in the network's 's latent space as surrogate metrics. While most studies in this area have assumed (near) orthogonal data, recent efforts by [43] and [44] have made initial attempts to analyze non-orthogonal data scenarios. Building upon this foundation, our study extends this line of research to nonlincar attention-ML. P transformers with within-concept positive inner products and cross-concept orthogonal data representations. The key to our analysis is the assumption of good initialization of attention matrices and a sufficiently low-noise condition, which is reasonable for modeling language rather than images. In this setting, SGD allows noise to have only mild impact on shaping neural network matrices or influencing gradient flow.
58
50 60 61 62 [3 64 65
66
67
Theory of Transformers and In- Contente Learning. The literature on Transformers and ICL is wide-ranging, and will selectively address the most relevant ones. Prior studies have analyzed how transformers learn topic/coneept semantics [ the origins and biases LLM representations using latent variable models [10) and ICL from a model averaging perspective [4. However, these works do not connect the geometric properties of concept-encoded representations to transformers' powerful ICL abilities. Another line of research has studied
the learning dynamics of transformer. including analyses of finear-attention transformers [66] [7 [2), QK- combined attention-only models 20
N 79 ReI.U-free MI.P 54
80 81
without MLP 17 25), impractical squared or hinge loss
26 27
Though relevant, these works rely on
simplifications or do not connect the observed linear semantic representation of large model to the transformer's excelling OOD capability.
Concept Learning in Deep Learning. Hierarchical learning has long been regarded as a key factor behind the success of deep learning [82 [4] Recent research shows that large-scale generative models, such diffusion
models and transformers, effectively encode hicrarchical concepts in their latent spaces [2] 46 86 87). Morcover,
73 88 89
show that transformers can capture hicrarchical and compositional structures in data.
From a Bayesian perspective, [ [ [4] interpret ICL as LLMs predicting outputs based on latent (concept) variable inference. Furthermore, studies reveal a linear structure in LLMs' latent space over independent interpretable concepts: representations of the same concept exhibit positive inner products, while statistically- independent concepts are nearly orthogonal [9] 10
[T [2] 19). Interestingly, aligning with the findings in [46] (90) Independent Component Analysis (ICA) is naturally more suitable than Principal Component Analysis (PCA) for obtaining meaningful feature or label vectors in our prompt modeling. This is because the features or labels are nearly statistically independent and of equal strength, especially with large K, while the noise
17

---

feeble in our modeling. Building on these insights. explore in a theoretical context how the compositional nature of concept representations relates transformers' ability generalize OOD tasks through sparse coding modeling.
believe our OOD results are not only coincides with the transformer's 's compositional generalization ability on language tasks [89], but also consistent with other concept learning outcomes of diffusion and multi model model: [87) shows that adjusting the length of semantic representations can directly
affect image generation behaviors (see Figure 5), while [86] reveals that compositing different concepts enables OOD generalization (e.g. "blue square apples" in the Figure la in (86)
D Preliminary Lemmas
D.1 Probablistic Lemmas on Concentration
Lemma 6. Suppose that > and Vd E {dx 2( KNI, )). where BT.
Then with
probability at least 1 8.
odd
|1::12 S o²d
6(N[L+ 1))2
(E:, Er)I 20ट ddlog
5
(E:, uلمl $ 1|1112000 / 2
6KN +1)
for all Ex, E: ~ Dea( or D&y): u D Dy), E (1,
Proof. See Lemma B.4
42 for proof.
Lemma 7. Suppose that 5 > O, dy = (llog(m/8)), m= ((log((/(6)) Then with probability at least
1 8, for Vi E [m), k E [K],, [dy K1,,
oidy
10)1 3 gidy
W
2
(.,)
2
|a
5Km,
||c]||
|
IPo..)). mll S V
2
8 00,
(7)
/22 < max
ic(m
|c
|dall
(0) IPol..,.oolB
I /2
5Km
8
01
Moreover, for some E (0, for Ve # e',E [+). Bus (0, wtt) where w/ <1
[i E [m] r.
e
e'
(( (..g.. 0H} m
m log (10K1/0)
m
2
[[i E [m]]]
(0)
> 0, (00)
PO > 031
m
m log (00K//8)
4
Ki E [m]
r.
10)
B..3. & = =5,30)
07|
(1
we )m
log 10K
(w/ - w
(...).t
2
8
Hi E [m] r,
(
B..I.E & (3)
0,00..... ...gii
(0)
0}
(1 we wemm
m mlog 10K)
2
Σ
ri r-. (o & + e'
+
(
-0 S
2 2log
5Km, 8
50 501(
+ <l(cbel))
16
(e(fE(m))r i o to (0)
(t.)
Σ
rs r
((0))
0 5 2log(
5Km,
5 |d
16
(0) ( (0) So..)
(8)
In addition, for a sufficient large m (log(K/(6))(11 wc]) the lower bound inequalities regarding
maximam value in Eq. 7 hold at any above index set of in (8) For example, there exist E {i E [m] (3)0)
(0) r; m G + (3) g
> O,
0 (0)
(0)
0 * < 0}. such that a o S -
(0)
18

---

o10,)) then by Bernstein's inequality as well as dy (llog(mm66), with probability least 1 6/(m)) for Vi E [m]]
Proef. First, notice that W (G.))
(0)
W ( (6,)
(0) 102 oıdyl < O dy log(5m/6)) S soiyy/2
By union bound can have the first inequality in the lemma hold with probability at least 1 8/5.
Next, we notice that
8
d2
cx
(0)
|xx
|
(W8
(0)
|| ||
Po m. (Wo
(g.))
(0)
gw
are all Gaussian random variable with mean 0 and variance of. Then by Gaussian tail bound and union bound,
with probability at least - 6/10, for all i E [mm] and q E Uk (
Ck
de
4u }. holds that
|eal
[d
kw (0) a S (2og((5mm/)) OT-
Notice P(a1/2 > (
() is an positive constant, then following the techniques of Lemma B.5 in and the condition m Aloog K/))), we have
(c.))
P(or/2 ((W (c.))
qil) 1 P(01/2 > max
(..)) ("), ",(H),
= 1 - P(o/1/2>> > kW o (0) .g1l)
mk
2 1 5/10,
then with probability 8/5, the second and third inequality hold.
(0) For (0, 1]. we that the variable (........ e'Soo NOO of + [[]dd]13)). and it's
independent to the event {r, ). Ve E [=]. Therefore, we can see the count of {i E [m] r. o..).. c,
> 0} as a binomial variable with 1/4, = m, then by the property of binomial tail, condition 22llog( (6)))) as well Hoeffding's incquality, with probability least - 8/5 we have
_,a(0) ) + t > 0, m
((i E [m] |r;
2. m78
e
). k e' c(33)) P
031
1og(10K1/0))
m
2m
which completes the proof of the forth inequality. Similarly, for the fifth inequality we can utilize the same techniques to derive that it holds with probability at least 1 1 0/5.
For the event {i [m]
a0 (..... + =c30 0}, we have
(0)
P(r:
(0) ( &
(1)0))
On.g.h 0)
P(a + e' E.(,80)
> 0)
-F(aO e' (6.))
c
> 0 a 0
j + e
0c..,.k > 0)
P(aou.,.. e' el
0 O
+ e((
0(.) 0);
Ww
where 1 W/ is the probability of the conditional event (a
0
d +
e' E.co 0}. and we due to the larger variance aoo..... compared e' d'(o)))., We de- note the probability with we since the true value is hard to compute. Subscquently, the event {i E [m] r.
(0)
1 wi n m, then we can have the sixth inequality hold with probability at least - 8/5, 5, utilizing the property of binomial tail, condition well as Hoeffding's 's inequality.
m 0o = < > 0} can be seen as binomial variable with p
e
((lloo((/(O)))
The seventh inequality is natural inference of the third and forth inequality, where the m 2(log (8) ensure / m log( 10K, /0)/2 < m/16. and the last inequality is then also a natural inference of the third and fifth inequality.
Therefore, by union bound, the proof is completed.
19

---

D.2 Matrix Theories
Lemma 8. (1.1.5 P5 in 1977) Let A Mr be idempotent, that is. A. Then, each eigenvalue of A equals Beside, identity matrix the only nonsingular idempotent matrix.
the rank of A, which is either
Lemma For a matrix = Σ H.P., where are symmetric idempotent matrices with rank( P.) and thus Σ PrP. is the idempotent decomposition of matrix by Pt. Then we that || vtr(A" A)
Ei= = =VV X2 where x eigenvalues of A.
Proof. By definition,
A Σ PT P
Σ 14 P.P. = Σ I...
Then, by Lemma 8 we have
d
d a tr().) Σ
A) = tr Σ
PP) Σ
Mi rank( (.)
Σ 1 Σ x
1=1
1=1
1=1
D.3
ODE Systems
Lemma 10. (Lemma C.I in 43
Suppose that a sequence az, 2 follows the iterative formula
at+1 = az
+ beas
for same 0 and b Then it holds that
c
st S a: S
+ æt.
+ beao
for all t 2 O. Here, I, is the unique solution of
dæ:
Ji + be
ct + 0e
dt 1 be
Lemma 11. (Coupled ODE System 1). Suppose that there two coupled sequences yo, 2t, t 0 follows the iterative formula
Ve+1 yo + azcVle 2 -22=? + e?
30 > 0,
0,
E++1 Ze + b.
< 0,
> 0,
for some o, b Then it holds that
v(t) S ye:
2(t) = 2t,
for all t Here, y(t), =(t) the unique solutions the following ODE System respectively
y/ (t) 4 =())((t), y(0) yo,
(9)
(t)
z(0) 20-
As such, for 11 = minft E 2 0}. we have
-a= + -2(0))
1b(1 20(1)²
Vb, 2 r(0)e
and 2 -2(0)(1 + 2y(0)) 6(1 --
Proof. From the condition we see that 20 < and z is an increasing sequence (2, 2 z0). Besides, as yo > 0.
during the period where 2, & O, we see that yi is monotonically decreasing. Then by (2 + e -2y + 2ve= 2 <
1/4 as well as Comparison Theorem, it's obvious that the continuous coupled ODE in Eq. is the lower bound of ye. Then readily obtain the result by solving the ODE.
20

---

Lemma 12. (Coupled ODE System 2). Suppose that there are two coupled sequences which are the sequences after Ln in Lemma 17 and 2 L follows the iterative formula
1
M1+1 = + aZtzi
2+ 201 + e2vi
Mt > 0,
a a> O.
R++1 2t b e 2y(t)²
1 e Zy(1j² C..
Rt, 2 0,
6> O,
for some a, b 2 and é' S I 1. Then holds that
y() S V/ S v(t). 2(1) S 2 S z(f),
for all t t1. Here, M(t). g(t), z(t), z(t) the unique solutions of the following ODE System respectively
(E)())(1))) Ei( (y()))(())) log(y(!))) abe! - 2u(te)?
(t -t- 21)
(Ei(2n) 2) + Ei(-2y:, ?))
1 1 e2(t2)
+ 1log((11);
=(1) be - - 2u((t))
(t (t-11), t1),
1 -22(4))
ab(t - 412
(Bi(2((t))) + Fi(-2y(t)) + log(ū(t)))
(Fi(2 3) 1-tFi(22p)) 2yc 2)) + log(ye,)
2 x(t) b(t t2),
2
where
Ei(x) /
a c¹ di "aaio + In exp(x/2)
Σ
(-1)a-1
1(n ((n-1/21 13/21
22
2k +
Proof. We see that as 2, 2 0, 2 (1, the is monotonically increasing. As such, by Comparison Theorem we =2y02)
see that upper and lower bound the coupled system would depends on 140-2 and &. Easy to see that
1 20(t1)²
1 2y(e)²
S + e 20(0)3 S I,
1 + e 2x(G))
and then collaborating with d' < 4 S we can obtain the result by solving ODE. Observing that
dy(() 12 22(t1) (t (t))
abc"
dt 1 2x(t))
1 1 (2(t)= + 22(()²
2v(tr)
(L 11)²
+ (Fi(2(()))+
log(y(t))) abc"
const,
2x(t))
z(t) be!
1 2uttr (t t.).
Thus by the monotonicity the system is unique, which is also ture for the upper bound ODE. The proof is completed.
E
Data Distribution
This section provided the detailed formal definitions of the prompt distribution.
Definition 3. (Polysemous Word Model (Da, Dy, Dz, De Dey) We assume there exists Kı concepts of words totally. Specifically, each concept E [K1] is characterized by two ammanticall-ooposiit feature vectors separately, denoted as ph and peg and the label vectors that describe their semantics under the co-concept qi, and qx.- Our word samples from distribution Dz and Dy, which can be written as the following forms via reparameterization:
Rax and their corresponding labels E R are generated i.i.d.
2 ~ Day Eu ~ Den = =N(0, TTxx): Eu ~ DE, ~N
x Mz + Ea Day y Q= + Ey Dy:
where z E 1 (K < dx). We denote the sparse latent signal and as the spurious dense noise, and each æ-y pair reparameterized by shared We have the following assumptions M, z, respectively:
21

---

The spurse latent variable (2), ;ax) E (0, is sampled from D.. P(zj = 1) log log K
K
M (at Me 162 12."" MK, k Megy M1,,2,".. MKg] [M MKk] R is the
feature dictionary matrix, where { 11-11 are concept-relevant features, {Vkz (uka)k==1 are concept- irrelevant features, and Ak E [K], [|]]] = ||u]|. We assume that features of the same concept have posirive inner product: 30 < K2 < 1, vk1 E [K1], < (WE MEg ) Ri Kallu/ |l. Meanwhile, we let the features of different concept orthogonal: Ve E [+],e' E [ [Ki].r r' E [K2]], t vrb. we have (pes u) (U,, (.,)=
: 9x, -x,, 00. E g the corresponding label dictionary matrix, where 1 || = ||q|]. for Vk E [Ki]. Similarly, we the labels of the same concept to have positive inner product: 30 Ki < 1, 4k1 [K1]. 0 < (g 9h, Kyllqll? while the labels of different concept be orthogonal: (g) 92) O, vk + k' E [Ki].
Q -qq qi 4.. .... 9k.
Definition 4. (Concept-specific Contextual Prompt Distribution) We consider the case that each prompt is specifft (i.., the multi-concept words in one prompt would at least share one co-concept). Specifically; the chance for selecting each concept as the co-concept of one particular prompt is θ(K1 ²), and the chance for selecting the two semantically-opposie vectors of the same concept is 1 2 During training, each prompt fxr, y1. 21 yr., #14+1 } sampled from mixture distribution Ds defined as below.
Dg
[ P1 + Th Ph)44))
(10)
where Ti. = =
and the P Pk,441 and Pu,1.1 are prompt distributions characterized by the k-th
2K,
concept, defined as
p k1L+1 {st Dz, D PL.1,2 1, 11,1 (2k 1, (2k-)))) K'
13
( 1) 12122 1} 02 (21,2k n [,.. 1) 0, v/ [L], 21 Pl.22 2
Ph,+11 { 251 Da, y~ D0, PL41.24 = 1, L + 1),1 + 12k 1,k}, Pt.s K'
(-122-1 1 (21,2k 1) 02 (2122-1 [..2 1) vI [L]. P1,21-1 Pi.t 11
2
where Pt., = P === 1). vn E [N] where the training size, if the training prompt S., sampled from Pi.L,,, E [ [2], I then by Definition 1 the label vector of the query should contain Th and we call ys., = the real value label of this k-th concept prompt. Specifically, for vk E [K1] we define the index set of training prompts sharing the k-th co-concepts
V4 = =V UVT
where
2t {n Sm. Pt.
VI = (m S. PT.141.
For sample æi where n E V E [Ki], E [L 1], we define the index set for non-zero elements of z!
besides 2ik- and zīe,t, namely MF == [k [K] zī,, k & (2k 1, 1.2k)). Also, for each prompt sharing the k-th co-coeep,, we define the index set of demonstration in the context:
S = Ele [L] I In Va, 2 1/. ST,x {1 E [4] n E Vk, 2,,2,.2 = 1h
F Model details: Attention Part
In this section, we provide several important definitions and compute the original gradients of attention. Lemma 13. (Contributing and Misleading Neurons)
wt, (1) = (i E [m]] n E V 18,, (9) >0, 2t, (1) = (i E [m]] I Vr. 18 (0) > 0}.
(11)
(T. (i e [m] VT 18,0 0), U (t) = [m] n VT.r. 13,0 < 0).
22

---

(1) = w U UW, (1) are neurons that can be activated, among which (4,,, == U UZ,, neurons that correctly contribute the prediction. The following lemma computes the original gradients.
Lemma (Gradient Update) Denote
rj r(t].
E, (1) E' (ys.. f(H"; w())),
), softmax (wWhr) w w
(12)
1 (4] 1( 06 altn(H"; w(1) 0).
LB, (v(1)) (B
E Rdx xdx can be derived as
1
(0) (6)
Σ
(o (w3)) oo. w('h") wx (xf-- x3)x0++1
+AW3 (0)
Usa
B
WE(L
(13)
Similarly, V wR (9) Lo. (wo(t) E xedxxxx can be derived as
1
Σ NEB,
En1 Σ
(o (wo).. (a.) w(' (f - =²3) +XW* (4)
B
(EE(M)
(14)
Subsequently, we directly compute update of the attention matrices along the feature directions as below.
Lemma 15. (Concept Learning of Attention) For vk E [K1], have the single step of learning of the concept part of the features:
ag wa(++)) a ag W a = -1 ag T V Vwzy (1) Ls, (I )az
=-n ag 112) 1120), ndaa wa ai.
(15)
ag We (1+1) (++1)
ag 186-0 ag W ag T7 ag
V Vwz (1) LBs (VD(1) )a
(1 a [ - mdaa T W (1)
110
ag.
where 1 (t)
Qq:0q whans uhans and 1 (t)
a woan are defined as below.
(1)
1
Σ
ela (1) af (EH +
Σ M,) M)
Σ
Σ
(o 10. (o)))
chooor
k cE(+)
PEM
(!)
[JE(L]
nEVE
(Wy (O
(4) up
(ak
+ Σ
Qs + (a W "(v" bf + Σ
M. + En
Σ
M. 2)
SEMP
sEMP
SEMY
Tee, (2) (1az)) + ag (C0,++1 Σ M.) Σ
rl Σ
(o
B
EE[+] nevEnB,
rEME
sEW (8)
[JE(L]]
(W
(Wo (m))
(4)
Qs (a w " ((v) b; + C (.
SEMP
(()
(1)) (11a + (6 Σ M.)) Σ
r w: (6.)
QQag..com carm
RE[+]
rEME 2+1
seWe (8) km
de Σ
( 10 )r
Σ
(o j w (4) (é(vi "b + Ch,1
Σ
(a ), cm,
S
1ES-A m.k
SE(L.
(16)
23

---

Similarly, K.ag cheon and I (1) K somi are defined as below.
K iiiin
B
Σ
ag We (ax + cbp + 62,441
Σ
M.) Σ
r
Σ
(o
KARE|KI cE[+] (og
rEM? 1+1
ieW
NEVENB
(o "h 1, (Wo
(4)
+
Σ
Qs + +g. {
M2
En
Σ SEMy
M,-
E2.))
sEMP
Ce, (0) ag wa (ag + eb + C2 +
Σ M.)
Σ W
B
6E[+ REVENg,
r 11
iEW, (9) R,n
Σ ( (o ( (Wb ( " Σ
Qs + (yd))a (Ca, - Eas)
(SE(A]
sEMP
[())
e (0) a W (4) (ag ebg E0,
Σ M,) E
Wy (9)
K.ag.comme contri
B
(6))
=
rEMe
iEWE, (4)
BE
dg 1( Σ (of))
Σ (aa))))aa iEST!
Tea,
"a. -
(o a (B,))]
LESt
SELL]
(17)
Lemma 16. (Label Semantic Learning of Attention) Also, for wh E [Ka], we have the single step of learning of the concept-specifte catanniicallyoopositt part of the features:
bg w
aa (44
by
1// wa (0 = - by "V (w LB, ( (4
no (6, whoos + Q. comme)) - md Wa ('
y(t)
=
(18)
b) W (1
b% W (0)
= =Ti by 'V wg (0) LB, (T)(t)) ybg
= -n -e (1)
K.b. choor + 1(0) K. comm)))) - 9:2b W (1)
ag:
ag:
where T(t)
and (1)
Q comtri are defined below.
Q. choos
[t)
Lea (0) bi T "(E
Σ M,)
Σ
Σ [EELL
Q. chhoos
B
MEME
(e
=E[+]]
nEVENB.
(Wo
(1) (q
+ Σ Qs + (g,)) W (3) ((vi' - yb + Σ
M. + + - Σ M. - -
SEMP
SEMP
SEMS
(1b2)) + +éb (E +
M.))
Σ rp.
RE[ nB
rEM
ieW! (t)
Σ Σ
(a (o (w (Wo...)
(9) Σ
Qs + 16 w (26; + by (C - ()
lESt JEST
MMMP
n,k
Σ
Σ
(a n (o))))))))
(Wy..)
(0) Σ
Qs + Eg bb W (0 (- 2b + (E%. -n,
LES JES! n n
MEM
(10)
1
Σ
2e, (1bbB1 + +
Σ M.)) Σ
rw dd
Q.b contri
B
RET==
NE
rEME+1
lEWE (8)
b.. W (6) (2(
Σ
( Σ (of3))01 + Σ e Σ
(of
Σ
(f)
JEST
JEST
eE|#]
JES- m.k
LES R
m.k
(19)
24

---

Similarly, K. bR
and I (C)
conn are defined as below.
(1)
1
Σ
e.
1 b
W
(ax eby Et S0,++1 +
E
M,)
Σ
r Σ
(
B
k
rEMi+1
(, (0)
(JE[
cE
nEVinE;
(W (Wo
(4)
E
Qs
+ (b T( (E M. + St
E
M,
SMMs
SEM?
.EM
Σ
Σ
& b & W (1) (ca; + b + E
W2
Σ M.)) Σ re.
B
2C nEV
rEME
ieW (8)
Σ Σ
(0s (
(1)
(Wo
(W
( Σ
Qs (2,)) 11 bj "< (2
LES n. SES n
WEMP
Σ
Σ
(os (1 (O..)
( Σ
Qs + (-2bbb1 + by T (E 1)}}
IEST JEST n,k mik
MEMP
(t)
1
Σ
Σ
2en (4)
bg W (t (éag
- by + E +
M.))
Σ
r: (..) dy
K.bbg.comrr contr
B
rEMi+
SEWA 1
(2(
Σ
)
Σ
(o (0b")))0b,++++
Σ e Σ
(
Σ
(o bb
[
JEST K
JES n
cE[+]
JEST
IES
n.a
(00)
G Model details: MLP Part
Lemma 17. (Tensor Update)
T
Wy (G,)
cj =
(D)
1/
V
Wx
(0) Ls, (V(t))) Jck
1=0
(1,))
(21)
T
W ( d;
(4.)
Po
((0)
1/6
Σ
V
wy
((.))
(c) (1) )d..
1=
Lemma 18. (Gradient Update)
(1) LB, (B,1t)) Rxx(dxtdy))
can be derived
(c.))
1 B
E (9) 100,0 [12 Σ
(o (1) -1ddk te Σ (o (0 (ce Σ Qs+ +xwg oa.
(9)
A [1] eE(d:] ne nEVEnI.
lest
]E[L]
SMMP
(22)
Lemma 19. (Concept Learning of MLP) For E [m],k [Ki].
W ((+1) W
(6)
c -p: V wy (b.))
(o LB, (v)(1)) )c
(6.)
(23)
=- O chane 01t.)) Ted WW Ols
(1)
(0)
ck.
where (0s,) .eg choos chhos and (o.,.) Sg conori are defined as
10 o chaoe
(t)
Σ
Σ
Σ
e eo (1)
r.. 18
(1)
Σ
(o (o()))"( Σ
Qs TT ces].
B
[E|4]
*EMP
(t) (. Sqrcummi (E(+)nEVENB,,
Σ
(e.en (t)
18 (11c
(24)
Remark 2. (Informal Discussions). Interestingly, the gradient of MLPs' Concept Learning is very large. We have the following situations.
25

---

When the neuron is activated (ie., {n E V n Ba, if (WW b.) W. > 0}. and a (8)
S, +
é. Eies" n.h
( 1)W8,,, (0) dg > 0), neuron is likely to be activated (i E W (t))
(1) o i W k,n (1) E (1), the gradient will advance the 1. (f( R . ê > 0, E W Kn
Wo (a.)) cz:
2. (2) r < 0, ((,()) + i WE,,(t) (t) (g,(() (t), the gradient will diminish the (9) ct. thus help deactivate this newron.
W (
Lemma 20. (Label Semantic Learning of MLP) For E [mn] & [Ka].
wy (1+1)
W ((.))
de
W W O(..))
dg = 7/ Vwy
(e,)) (4) Lis (ess(VVV (w( dd
(25)
(t) T/e (L0 ) whare + O wome mnd (c.)
(4) de.
where r(() (..))) de ahove and z(t) ( dg pomer are defined as
16
(1.) di chaas Σ
Σ
Σ
e Eo
19
re 18w (0) Σ
(as
(0) Σ
Qs + Eg.) de..
kE|K,] cE(1]
NEVENB.
LE(L)
SEMP
(1)) (...).de.comm
Σ
c (r,, 18
( Σ
(as
Σ
(o I|
(26)
B
RE( nEVinB,
lESé m,k
lES
H
Discussions over Parameter Settings
Note that we do not have any requirement upon demonstration length L and batch size B for training, thus the training can be really flexible compared with the strict requirement in (28). The condition dimensionality dx, dy and the network width m ensure the learning problem in sufficiently overparameterized setting where the norm and the inner products the Gaussian noise and initialized NN can controlled within certain range with high probability
5. which is standard requirements in recent feceuffeature learning
45 58 42 [2] 43 The weak requirement on network width m allows us to conduct a fine- mraiigr analysis hased on the network projection length, which is fundamentally differs from the NTK line of research [2 that requires an infinitely wide network perform linear regression over preseribed feature map. The condition on y ensures the learning step to small and thus learning process enjoys an approximation to gradient flow rather than the challenging "Oscillation" regime [9). which is analyzable but not necessary in presenting our theory. The condition on the small À is to ensure that the learning dynamic of Attention and MLP would not stuck the origin point, and ensure that can analyze the expected learning dynamic with limited impact the regularization the initial stage, which is also adopted in (53) The condition on K control the impact of cross-come concept contribution in the Attention's is learning dynamic, which can actually be relaxed at the cost of a denser analysis. The condition on de is to ensure that the impact of the norms and inner-products involving the Gaussian Noise on the gradient cannot surpass those in the order of feature's 's norms, which ensures the gradient flows to be not noisy and could converge the expected gradient flow exponentially. Last but not least, the conditions on σi guarantee that the initial beliefs MLP is small and the gradients of SGD can update the model effectively. The condition of ao is only used when discussing the OOD scenario.
[41 57
53
I Convergence of Expectation
In this section, we assume all the events in the Section D Dhold, denoted Tre-
We examine the evolution of E(w"") = (E( B(W2 (")), B(W(0 ( where the expectation [[[] is taken over the stochastic batches. As such, we can see every stochastic gradient
og.., )} at the whole itcration 0 &t5 t.
update within each batch as gradient update upon noise-free free and ditegory-balancee concept-specific prompts.
26

---

Mix + Mai
Pix
Mei
Lemma 21. For vkg [Ki], we define Oki
and bk:
By definition, we then
have
Mxx = ac, + bkg Mk =
bkI
(ax1 bxr) 0, {ax) bbr3
{ar; bug
(e ME,) = Ila,; 1| - llbb 11:
| 112
= | 112
+bbb bell]
112 ||||
(27)
||ul| < <|aak, 1
|ull
|ull
bk 11 1
2
2
2
for Vki # kI E [Ki].
Remark 3. We observe that, through this formulation, the shared component ak, can be interpreted as the "concept" part of the two features, while the terms =bx, represent their opposing semantic aspects. The relevance of this modeling exemplifted by Figure 1(b) in [[2], where the concept "[Bird]" is composed of orthogonal steering vectors: "plant + animal" and "mammal bird." These vectors correspond the concept feature ag and the semantic label features b,, respectively.
Idempotent Operator Trick. Define == span(M) and its complement space U+ By definition, we know that dim(T) K and dim(U*) dx K. Then we have set of standard orthogonal basis for B defined as
Puau
||all| Ill |
a1
||22||
a2 b2
||bal|
||ax,
[|b,,
bx
[|l]]
V1
[|ll
22
VK2 |||||
ut
ugx
R
where ut..... .... uTx- are the standard orthogornal basis of U+. Then we can derive that
Ki a
K1
b.b."
K2
D,
dx
a.
1
U Uw
Layxxxx xdx.
(28)
Lemma 22. (Partial Statement of Lemmal whole iterations. We can decompose F. W.
E W and E[ (0) by (scaled) idempotent matrices.
Kl are identical and symmetric during the
and R R
E
W3
11)
Σ
K
Qi
(0)
a,a,
|a,1|
K
b.b T
bb.14
K2
Qir
V,
|
dx K
Σ
PQ...
The th
(29)
B
(1)
K.a
a,a
|a.11
((0) K,
b.b."
||b.
Kr
(e)
V,
||ll
dx
Σ
PK,w
(e)
Uhe
where 2 and O'K, represent the concept learning process, (0) and (
represent the concept-specific semantic learning process and TQ KKrr aa Pk. represent the memorization of the concept irrelevant noise.
(1) (0) (1) (0)
Proof. Apparently they hold suppose it holds at step thus
E[ [1] ( ( Elw5 "1 = E[( J
we examine + 1. holds that
Es, (W(+)) (++1)
B(V
E E [9] Br, 1o (es Li, (())))
Es, (w ())
((E(W(
wa (+) Es, 10 we (e) Lz, (E (1 (E(
Here, (()) fixed matrices and the expectation Es, || is taken over the stochastic batch the time step As we are considering expectation over the isotropic prompt distribution, which can seen noiseless distribution with an averaged categorics of words and labels, the expected gradient form could written as symmetric form:
Eg, W (t+11 E(DD(0)) F. (1) = Σ (a
K. a,a
+ +800 b.b.
dx-2Ki
U,U,
dx
Uhy Uew
with some coefficients d..' K. & Vs [Ki), E [K2), Idx - 2K: K2). It's direct to check
(8)
(0) (t)
that L (0] also has the exactly outcome. The proof is completed.
L
27

---

Worth noting that
14 E. W 11)
14 = dQ2 Q..
(0)
2(0)
m E
A 0
1 K K,
(1) g00
14 W (1)
1u = (1)
sQp -QQ.
U, E| W (5) I1; = K.5 KK.6'
(1)
(30)
v.
RRw n. Qr
(0)
Ur ' (4)
U, =KK
(1)
We will also have
(E 1 In5) B5 W μr (1) a ax.//aa, 12 + Q 2 b.ll.
(31)
(E|W (1) |4 B |n a a ||all 2 1 1 /b.11l
for Ve
[+:] and for Ve' E [+), e [Kil,r E [K2], [dx KJ, Vu E Ees uwb
(E W u) BlWa k =0
(32)
Similar conclusions hold when the query vectors are u, and u. E [Ka], w E [dx - K].
Definition Define Q = span(p)) Q) and complement space Q we can decompose i-th row of W via the
following decomposition:
W
(s.)) W
K2 == (
(1)
Ck
K1
de
dy Rp-hi K1 Σ
T
T
Rixdy (33)
c
|
2
k=1
k=
where qा. qdy-Ki K2 the standard orthogonal basis of the complement space Q+. Then we have
E[ (..) (9) = auu + e. 1(0)
(34)
for Ve [+), E [m]. E (K1).
Lemma 23, At initialization, for some e E [+] and Vk: E [K1]. define
exp(an [bal 12²) exp(-o 10x 13) exp(o 12 () + (xp(-an 1b 112)
Sh := E WEVE IESS
(os 0 10
then we have some wcg E (0,w6z) where 4l4 < 1, the following will hold
810) Q.k
( K.
Q.4k
JQ.
K.r
(0)
(0)
la
1
|a
b
b
1
Tu
1 2
PQ..0
PKK,0
ao.
m log( 10K:
E [ (0]]) = =|(t [m]
(0)
ECt
-:00 -
03} 4
WEVE
2
m log( 1
neve
WE (0) 4E, (0)] (i E [m]]
r;
m
é
(0)
e
D
k
03
S
m
4
(1 + weg )
10K m log(
(1 wkp )m
[ (0) n (W) (0))1 < HEVE
E
A
8
2
8
(1 wey)m m
log( 10K
(1 wkp))) )m
F. [ neVe
8
2
The parameter wleg is determined by oo, σi, ||ax|l Il ||bx I1. ||cu|] and ||dal]-
Proof. We have that neve
E
UE.... (0) n (WT, (4,;(0))) * 2. By Lemma
7 we see that for
exp ball 12)) exp( ob bu 12)
2 R NE
(o
(0)
1 "]
exp o3 by [2) + exp( of b P2)
1ES
we can have corresponding wig E (0, w/zg) where w/p
to ensure the conclusion holds.
28

---

Lemma 24, (Coefficient Update) Denote B(y == (E(W3 ) ), E(W (") B(W oc,,,
where the
expectation E|] taken over the stochastic batches. We have
(T)
a2 V aw Ba, (I,, (V'1)) TB ((V()--)))
0
= Q
(T) KK
V wy Co LL8 (w)(()) B( (1- Ja..
K,
a,
T
8 Q.a
q1 Σ
b.
T Be LLB, (V)l))) (((V)(-1)))0.
LL,
1=1
K.
K.
11 Σ
V w EB, Ee, B(y 11).
1=1
(T)
EB, LLs (V)l))) B(y D))
Qr
Qr
71
Vr
KKr
(T) (0)
K,
V4
V V (w Ba, Le (w)))) B B(w
(35)
(T) (0)
V Vig ELs, (v)l))) BB( (( Jut.
Ju..
PQ.N
PQ..
71
[=1
(T) (0)
7/t Σ
w te) Es,, (1, (1)))) () yut
PK. PK.
(T)
Y/h
V (o
(1) ER Lss, (w"))) BBW 1) Jck.
((7)
Po .k
1 Σ
V (o
(1) Ba, [LB, (w' B( "j.e.
(8.)
P
(2)
Po
V Wor.)
En, (v B(y 1a.
where e [+), [K1),r [K2), w E (dx K].
ar.
Th. 9ss By definition, we then have + Tns and da,
Lemma 25, For thi [K1], we define Cki =
2
qi, = Cxi + day; qh, = Cb1 dx1.
(ct, , d ) 0, {cr, d -(x deg J.
( 9m,)) ||cm 112 - ||da 12, || 11a ||cx; 112 + ||dR, 1|a |||l]"
(36)
||al| < ||er, 212 My + ||ll?
Ni + ||ll S da, 112 ||1l
for VM # ki E [K]].
Based on Lemma 22 and Lemma 24 the following two lemmas compute the update of attention's expected projection along non-feature and feature directions.
Lemma 26. For 1 > 0, we have
Qr = - (1-7:A)-Qr 72
(1-1)
TR = (1 - (2 (1) ((41)
(t+1)
(4+1) (1 p
PK. (1 (1
(1)
(37)
(e+1) = (1 - pp) P.
(e)
Po
where [K2], w E [dx
K], á E (dy- KIJ.
29

---

Lemma 27. For L > O, we have
((+1) = (1-18 (1-n2)aaq.r (t)
(++1)
= = -1- (1-nidaakk
9q.k
aKk
g(++1) (1 (1-r)-h) -1)
10 4rp.
||bell
Ki
Σ
E
10
(0)
( Σ
(ogg)
Σ
(obg))))
K1
(
nEVE
cE( mm(m)
JE
jES
g(++1) PK,
= (1 - 1)
|bx
Σ ΣΣ
E nEVE
12,100
10 10 Σ
(0)))))
Σ
(o (o)")
K₁
<E[+]
JES
JES
(38)
Proof. The deduction is direct by the symmetric property prompt distribution in Lemma 22 and the gradient forms in Lemma [S and Lemma 16
This lemma reveals that the attention layer mainly serves learn the different semantic part of each concept, and hardly have interest learning the shared co-concept part. Also, collaborating with Lemma PQ.
we see that
this indicates that the signal of (84) (44 would remain positive.
Also, by the symmetry property of learning progress denoted in Lemma 22 we see that Vk E [Ki], (0)
Bk Observe that for vk: E [Ki],
K Q.
exp (()) 2
E MSa
Σ
(1) (ds
Q.k
K.
NEV
RS
exp(8
/ bb + exp(-A0). /R///bbl1? bubl))
(39)
exp( 8(0) K.k/ (lbb)111))
(o )
100 n
exp( Q.
(00 K. bb (1=) + exp( --()) B /bbb))
JES,,K
U8
We see from Lemma [K]], we have
that a Q = K,
(0)
(0)
rol ?
: = PKK = σo ball Therefore, for t 0, Vk
((0)
(0)
exp(ot |x 1
(0) =
nEV
explo by 12 + exp( क.. b 12)
(40)
exp(-oãl1a bu 12)
E
Σ
(a
(0) )
nev M
exp(o b. (2) + exp( ob by 12
JES n,k
ven
Obviously,
E.
(0) Σ " (o ),
0.5
F
Σ
vea (a
Mcanwhile we see that
NEV NS.
NEV
jes,
T Σ
E (ag (0) )"1 R 0.5 due to the small du O[[|u|| -2) by Condition
neW 2
The observation in Eq.
39
collaborating with the positiveness of 3/1 Bk4) we see that the incquality
E
SSe "Sn mik
(a
ΣΣ jes usa (tt)) ) will remain during whole iteration. Also, by Eq.
nEV
we know that
[ Σ (a
(
Σ
(a (o())))) = (exp(8f)e / + exp(-86)) Bl 222l/bbx13))
-2 (41)
sest
JEST
This observation under our expectation scenario greatly facilitate our analysis. Since e, (0) 0. it's obvious that the signal of r, r.. will determine whether the neuron i E [m] will serve to increase decrease the 3(()
and 3(t) during the gradient update. We therefore start analyze the MLP's update below based on Lemma 16
30

---

Lemma 28. For L > O, we have
d0 * (1 - -mpd)a
(e+1)
|eel
cer
E neve
(Ln (0)
18 ())
2k,
cE(+]
) sommm)
f/t
(Ki c)llc
Σ
ler
(E 18 (9)
2K1K
nE
cE
(42)
E(1
P -1.d)..
(a.).
Σ
E
(1)
10(0)
(
Σ
2K
neW
cE((]]
LESt
Σ
(ag
S
where k [K1].
Proof. The proof is direct by the symmetric property of prompt distribution in Lemma 22 and the gradient
forms in Lemma and Lemia 20
An interesting fact is that the B(TO (s..) eq chnce) also contributes to the learning k-th concept. This actually suits our intuition that if similar things appear in various fields (concepts), the learning process can help integrate and facilitate the learning. The following lemma demonstrate the lower bound of the attention assignment, which emerge from the good property of our expected attention.
Lemma For certain iterations
(0, (.) Vk E [Ki]. e E [+b] we have
1. The neuron set E((WE, - MK (t)) (t) is non-increasing, and all this neuron will get deactivated. Additionally, both Ela ). and 18t)) Bo... would monotonically decrease. Also,
holds that e Po and (00 e 1...
(1)
The neuron set ET (WE,., (g)( is non-increasing, and neurons in will turn (t) (W (t)) Additionally, both o.. and -80) - would
into E.
nEVE
monotonically increase. Also. it holds that
300) Du.g. 0 and lab S 2(1) D
3. For E (t) (WE. - (E55(()) the Poa.) would monotonically increase. Besides, nEVE
when there exists constant C 2 1 such that
E
(C, (0 SC E (C, J) nEVE"
REVE
the Blab (osy would be contributed to increase, otherwise it will decrease. Also, la (l(e(2 Elese (o (0)) - 1)88%,. rell and Ela
4. All the newrons
E [4E, (1)(WE (t)) will ultimately either have its coefficient update
(t))
nEVE
stuck due to regularization, or grow into a changing margin a into FRLE,, - (WEi (4 where
(1)
[2
Σ
(
( D o..
iES
Prouf. By Lemma we see that ET ( o and eg(()
G would be contributed by VE a e would be contributed by V to decrease. As such, the first second point hold naturally by definition.
increase, also (I) 24,..())
ultimate transformation of (t)) - (WEi(t) Uz (f)) into E [MG (t) (W (.e(e) (f) Ug (f)) attributes to the faster changing nEve
31

---

speed of (a, y. compared ( > in the neuron sets EQ (1) (WE, (,,,,, whose learning speed ratio is at least (I|e=1//||d [1) Therefore, the absolute value of a (8)
(1)
(s..),. will surpass that of e/ ( which indicates the neuron would be activated for opposite labels, then the proof is completed. Given that Ciesy, (o8 (1)) 1 Eies,, (o)")) ") will remain positive, the discussion over e0 (6,). is simple since will
always grow r, direction, and thus the third forth point hold.
Considering the growth of (t) d. k. by The Th P12k-1 Pl22
1
we know
2
E
R Σ
(
nevE
LES
NEVE
hence if
(1 Ug (1))), it indicates that nEVe
EE
+ (2 Σ (o 1)e8 20.
1E5
We that E [44 (1) -4, (1)), (O)), the E[VE] will serve to increase (t)
but E(V
(o.g. k.
nEVE
will serve to decrease the (..).
(4) The contribution will tend to be positive if
I. nEVE
(C, (
E (2, (I (WE) (())
nEVE
Then, [2 (2EE, (o!"), ( of the neurons in and E. (t) (WE,(() (45((1)
nEVE will continue to grow, and finally be comparable the Ria Sl.gh....k kl Otherwise will continue to grow while evolving speed (l is comparatably feeble as receive the contribution oppositely from
nEVE (C, (+) "( E (()E,,. - 145,,, (1) (NEE) and E
(E,
(9) (1)(i 1(i E BIWE,,. - LLE (4)
nEVE
Quantatively this validated by our later results in Lemma
32 where the
(E, (1)
(Ca would be F.
REVE
nEV.
controlled by the initialization. Interestingly, we see that as (2 Eies,
y. grows up, its [45 (t) n (W (t) U (t)) will turn NEVE
(o (0) O
scale will surpass those of Eloo).., ..] Under this scenario,
into BE (f) (W, (t) - (t))(( where Elo again continues grow. Thus finally we have
(1)
(4)
R [[ Σ
( - 1..
LES
Lemma 4 will show that the growing of [2 >
Stese
1) will stuck, and thus the growing of
Elooo .) will also stuck the changing margin from
(t) (t)) -2, (t)) into E(L
nEve
(WT, (t) 445,5100 (0))-- The proof completed.
32

---

Prouf. Proof of Lemma
To examine the 0-1 loss, by definition, we know
Lg ())"))
PSp (ys.. ((E(S.),
(()))
= 10)
= PS (U5..
Σ
Σ
B
bR WY
(1)
Σ
(o5 (0) ( vi - 0),
m
v(t)
(
eE[1]
iE m
[E[L]
P(Elys..
Σ
Σ
on
W
Σ
(all y1)
] S 0):
(4,)))
nE(=
m
(E(r;==)
[e|4]
P(R
Σ
OR
Ol (( (2 Σ
(1)
(0
1y8n 800
6C (r, Win
LE m
V8n
Σ
OR
(t)
+ (2
Σ
(o
n 1)ysn 0
|=0)
iefr==- "ga
IES
= P Σ
E
(2 Σ
(o5 " )r
1) ) 50).
KEM k.m
MSn
(6)
iEW "4 (1)
158
Therefore, a sufficient condition for L (E( = is
B Σ
Oaagee + (2 Σ
(o 1)e BB (.,),,
NE
E
HEUE
LESE A,
ieWs (--ME
(43)
+ (2 (o
Σ
1)e
(())
1Eg8
for 4k: E [Ki], E [+].
We know Vi E 24,,, (4]. the Ele. p .e in the right side of the inequality is decreasing, which is a good news since we want the left side exceed the right side. By Lemma 29 we see that all the neurons in B(OW (t) UK, U5 (t) will be deactivated, and all the neurons in Blu (t) (WI, (t) 24,5. will turn into E
_ in the left side of the inequality is increasing, and Vi E (I)
U4 (t)
nEVE
(WES (1) (g,5 (f))
I.1
First Stage: Growing of Coefficient
In this stage, the coefficient update dynamic is continually changing without being much influenced by the comparably feeble regularization. Also. the impact of the decaying learning step T/e is under controlled during several periods, which be safely done due small initialization by large x well the slow quadratic decaying nature the derivative of 1/è- We see that at initialization, by Lemma and Lemma 23 the E (EE(s)) y (0)) satisfies
Dg
E
Σ
ri
aou. +
Σ 1es
(0
1)9s, 200)
12 - 2
5Km
S DD
tgW (0)
(44)
50 (1 ea + <i (Elldel))
16
and remaining job is when will E ((E((S)) stay positive for some tolerance. SnnDg
As such, we need to scrutinize the coefficients that would grow along the iterations. Therefore, define
K.US.
1
Σ
Σ
wy (F)
(1.))
iEu
2Sn (7)
VE(W"8. V8n (1) MSn (O))
(1)-
+ Σ
Σ
(2
Σ
(a - 1)ys, W dz) r=1
SEN
VS km
(7)
1E( (v) (r)-M"B (e))
USe
v5 (x)
33

---

We will see that the conditional expectation of this sequence (conditioned on E B(W and the expectation is taken over Dş) would grow up conquer the small initialization and make positive. Consider the whole training duration 0 < t S T",, the evolving speed of QQ. 1. K. G0G and g(++1) (o,,) & depends on E (9), B[10 (*) and E E 1Ess n.
E (F(.)) ( ((D(())) stay
S, S.
((++1) g(++1) (e-1)
(o Denote
3|ul|| (1 + -
la Km, m//)
(1 e- -
1 e-
= 4 log(
8max 8max ( J f,,
1800 Po..
We will show that oss is the lower bound of minugtr"),A(Ki R D nEDy Destse ),)) along the whole
iteration. By Lemma 7 can be upper bounded by 8/22 (m/(5 on (V(1 xy)) /2 [l)). and lower
bounded by 2/ |///, which is negligible term due to the small initialization by Condition
Lemma 30. Under Condition [ for the whole iteration 0 S t S T", for Vi [m], [+]. e [Ki],r
[K2], w E [dx K], we have that
0 & Ble 18 A 1(i E (e e PO S03
-05-ta a,
E l|
0 2 Ele Po kk ( E (E, Uh (t)) -e-3)) G.Bo
03 dg
5km,
01 "ros"I(de deel + E((exll))
2
o3² dx
||c/|2 a
(e) 0 s Bllaba,), wll S E
o (/2l|dk |
(45)
Lemma 31. Suppose Eq.
45 holds at iteration t $ T2, then have
lus. [Bs) - [[A ]
S Sx/2.
nev,
Proof. By definition, we have
E ( (0))
Rlys.
Σ
an
W
(0)
Σ
(a (1)
y(a))) y(")
rE[=] m
(E(N==)
[E
= E m
Σ
Σ
aaa. + (2
Σ (o 1 1) Bo.. ].
iEu (t)
-0-44-0 u8 (t)
Observe that
E
Σ
Co (2
(1)
Σ (0 1)ys, a(0) Bou.s:
B
vE(WE, (0)44)..())
1S U mik
Σ
(F)
(2
(o5")² 1)1s,, (c.) ] S k/44.
(EWSS (T)- (1)))
leS
USn n,k
Here the inequality holds due the fact that B[a 10i (WI, (1) - (1)) is decreasing
the initial value (0)
"ou 1(i (WV.,, (0) (0), 4,,5 (0)). and it's absolute value will not surpass that of
Ele(2 Eiss (1)) 1),8( .l 5 k/8, which positive (by definition) and also decreasing by Lemma n.kk
34

---

29
On the other hand,
Σ
10
(2
Σ (o 1)ys, Poc.
]
|-E
E
Σ
(r)
A 28n
(0)
LES
2Sn b.m
(x)
1
2
(o)")
l)ss. 80 | k/11.
REM (r)
Combining the two we can see the result is obtained.
We then denote the last time when there still exists E[A [. as F, formally F is the last time where
U
[E[A S x) + 0.
Latter will show in Lemma 33 that
CoormAKin/ ky) l0
t =
(203 1) xl)ll
We then denote the learning step at T as = M)j. and thus
2 A(π +
=
By Lemma 31 actually it would hold that
B() 2 [)]]]]]]] 2 nx/2 2 0.
E Se D2
And thus the 0-1 loss converges zero with an error tolerance by definition. Our following job is to find T. The following lemma provides the continuous ODEs as the upper and lower hound of the sequence A R.c.
Lemma 32. Under Condition 7
suppose Eq. holds at any iteration S T", then for Vt S VKE
[K]], e [+]. it holds that
The difference
E (()
E (=
none-increasing.
NEVE
nEVE
2. The difference of the loss derivative is bounded by OK))
E E (E (8)
E
(C,
(4)
)]] g
mEVE
VEV
3. B[AA is non-decreasing. The lower and upper bounds of the gradient update have continuous ODE counterpart. Specifically, there exist positive constant C1, C2 we can define ch..
51 Igll?
2mKi
dh cany= (203 (1-)))= The -0 bhei e²12 Let zh. Et the unique solutions
16mKi
of
Ft
R. Tt 2 kic
It + e hig thee
holds
S.r S
SEATTS
E
(C, (0) <
1 ghe
1+ khe
nEvE
Specifically, we have
log( R.e
2 5 <BIA <
Zhic
ztie
log
31eh
+ ghe
35

---

Proef. Observe that B. E
(ex
(0)
E "EVA"
(C, ()
equals to
nEV.
E
ti (ieeugg Eiewg,. (-- (n (0 a +3E ( Deol (1,))1
14c
1- Lieeugg Piewanm--45 un200
(00 ( 12 Diess (- -1ee
1 1te
-I Encup (1) Eieng ( up (o) (06 (o +22 (o -D Do00
[e( 1 ++e 1- ( co - Eiewy, (9) "g ( +03 Iics (ogg)) " - -100 (2)
(46)
As cross-entropy loss is L-smooth with L = 1, one can bound the difference by
B8
(En (1)
E (C, (1)
E (e)(E(S):
(- (-))
NEVE
NEVE
nEvA
NEVE
(47)
| m
Σ
E
a (o + (2
Σ (o 1)
| 1.
iEMe (t)
(0)
1ESS
By Lemma
23 we see that for initialization, we have
5km,
30 (01c/
|c(]|| + +EEldd
E (e((
E (- S 2
"EV;
BEVE
8
8
S <K/8.
Now we serve to show that the following expected difference
E BEVE
((((
B (- NEV.
is non-incrcasing. Intuitively, this observation is due to the inherent nature of cross-entropy loss, which always pays more emphasis (has larger derivative) those low value. Also, another important factor is the update of those ambiguous neurons' coefficient summation would also prefer the low-value among E (Ef() ((V))))))()))) Ve E [m]. To better present this observation, we define
e arg m
R (e)
)), E (-))( nEVE
REVE
which further means that ez satisfics
B
E
(C2 (1)
(% (1) < 0 due to the non- non-positive and EE o
ncv
non-increasing property of entsopy loss. Recall the update rule, we have
0 (1 - ( -
((+1)
11 2K1
1c
Σ
fer,- B
(L,
13 ("))
neve
eE|t]
(Ki 1)
Σ
(er, F
(% (4) 13,, (')),
m
2K1K
nEVE
rE[+]
(1+1) Oa.ii
(0(1) = (1- ne
||da 11?
1/1 2Ki
Σ
Fec E
[e,
(6) 10
( Σ (o
Σ
(o)
neve
eE[+]
LESS
Then we have (t+1)
Ela o + e(2
( 1)s(et))
(1())
E
(a (o []]]] = (1- (a
occ jut + tel2
169"
LES"
(o - 1) 1) (k.ggik)
Σ
re E. nEVE
18
(0) (Ilca 11 + 1|dx
Σ (o - 1)
2K, cE(
LES0
LE
(2 (o ); - 1)) - 7t6
(Ki
Σ
r4e
R.
F
(4)
18 111c
2K, K
R Ve
LE80
cE|t]
36

---

By Lemma 28
Lemma 29
we see that the E.
E.
Ve
[+] is
==
increasing such that
Σ
w(0) 1=0
= 0
Σ
ao
(T)
==
r=1+1
T=
EE R.n (1) (WE,) (r) (1)
EeUE (r) (W (4
n/lex/1 B (C) )0)),
Σ
2mK nevE
(EMg, (W (1) (1)
(48)
where we ignore the impact of cross- concept safely due to the large K = Q( وو 3 (m/,)) )), as well as the impact of regularization term since À (((Cllo((m//)))lg))
by Condition in the first stage.
Similarly, suggest E Eiess (o)+ is also given when considering the update for t + 1, we see that
E
e2
(o ( 1)00
EE
Σ
(ass
(++1)
T=0
(2
Σ
(o (
1)
m
1E44, (T)
IES
CEST
eistt) C
((2
(a
(o())
1)en )).
iEME 2mKi
nEVE
1E8
(49)
Interestingly, by Eq. we see that (2 (o - =
-2
(o ); 1). Thus we can
(0))))
ies
characterize that the magnitude of gradient update of the term in Eq. those -c due to the non-increasing nature cross-entropy loss.
48 and (49) of the et would be larger than
On other hand, by Lemma the monotonicity of
T
Σ
E
Σ
(T)
r=t²
EE (F) ((()))
i, (T) (
depend on the signal of E E (e, (1)
E
(C, (Cn))) Specifically, we see that
nEVE
neve
R
Σ
0 [((t)
=0(
Σ
(7)
O
1-411
tE (T) (r)- (T))
ieug (r)n(W
K. (T)
ml|cc
Σ
E
(E2
(0
(Co 11) })}
2mKi "EVE
EEM2 i (wn (7) (4) (x))
nevE
B
Σ
D ((t)
=
Σ
(1)
7m0
r=
r=
iE(WE (T) (1)) (r)
ie(WE (r) (r) (6))2
n|llcull 2
Σ
(E, (0)
E
(C
(1))));
2mki nEVE
FEWWE (c) ME... ())))))))))
neV
(50)
where the contribution term shared by two sequences. Therefore, by Eq. (50) and (46). evolution of
E SEMG le ()) (x) MC)
Σ
(r) go
Σ
(r)
1=0
SE(WE (4) (0)
will prefer grow in the direction of eç.
37

---

We then take look on the decreasing coefficients based on Lemma 29
Σ
1000
E
=
ie( (r)-24 (r)
iE(
(
1y nelccal1 11
Σ
(En (4)
2mKi nEV.
SEO hnn (1) 22 (0) (1
E 1
Σ
e(2
Σ
(as (
1),8
(1(1) F
(os
(1+1)
iEW k,n (,)) (7)-ME k.n (t)
1E8
IES M,K
1=t+1
egS(t
(2
(oS
(o
)
1)
E
(6.))
17-0
m
lrme
les n.kk
KEWI, (eW2n(12 (x) (2,-))4)) (r)
T n.lld 1
B ((2
Σ
(o n - 1)ee ())
2mK1
ne
LES
(51)
As such, we have all preliminaries to characterize first result the lemma. We first utilize the induction to prove the following:
E (() NEVE
E (-e s/88. Ve [t].
nevi"
This apparently hold at initialization. Suggest for any / s the result holds, then we only need to prove
E
2
NV 1-1
NEV, -
E
(BS)
E
(-e
((())
nEV.
EEV
By the condition of small n). in Condition
Lemma
Eq. (q.48))
49
50
and
S
we see that
R (-e
B(
(E, E(
9-1
NEV,
-( E
(E(S)
E
(B
"EV, (-1
nEV.
< (
R
((C2 (111)
C (--1)))
Dr- |ex11
2mKi
f-1
nEV.
EVV
i iew km (6-1)
+ (2 Σ
(og
n1 )
1)
Σ
mr- |d 2
((2
Σ
(os
(111)) a
); - 1)) 0.
2mKi
JE8
kEW (-11)))) (T-1)
nev. f-1
1ES,
and thus we have
E (e
(()
E (E()
me ?
<
(c
E
38

---

Therefore, we complete the induction. Then we have
E
(C
E
!]
REV (-1
REVE
< E
(B())
E
(B()
nev.
nev. i
S E
(BS)
(-e
neVe
nev.
S .... $
E (eI(E(S);
E
(- (((S):((V)))
nEVE
nEV
5km, 30 (lIcell + CElldell)
/ 2 log
S n/88
4
This completes the proof of the first result.
To obtain the continuous ODE upper bound of B[A, "I, we first recall the update
Elabt)) & + e(2
Σ
(o )
1)08 (0(()
E
(as
(++1)
) = (1 m)))(abl
1ES
IES
e(2
(o's ); "
1) ).k) 1t 2K1
Σ
re
(C,
(0) 10,0 (0) (Icall (|cx11+111
+ +|ddx 1P
Σ
(ag
(o(tt)))"
nEVE
1E8
sE(1]]
1ES 1,5
- 1)(( Σ
(o!)) 11)
7)t
(Ki -1 1)
Σ
re E
1e
(0)
10
2K, K
neVe.
LEST
eE(+]
Then, utilizing Lemma
31
and the fact (W (f) S m, we have constant e such that
R.c E[A
V
E
Σ
(o(tt))))]
A. ki el( Mellall
E ir, ()),
141
2mnKi
neve
les
S k.
c1( !ll|al?
1
2mK1
1 1+--
(52)
thi
A k,c
+ b Ahe
in To obtain the lower bound ODE couterpart, examine the update of the correct contributor neurons, as shown
where we also neglect the impact of cross-concept due to the large K M(poC(Ki Condition] I and appropriately chosen C1.
Eq. (48 (50) and 49
In terms of the update of El
, where E (t) (WE, (f) 24,5 (t))), we see that its update is
controlled by B[ B "EVE
(C* (0)
E (Ca (1))):
nEV.
((+1)
nr m/
(0)
(C, II).
2mKi
neVe
nEV.
8 S 32
Then by the first result this lemma we know
E
(C, (1)
(e, (1)
S
and thus
NEVE
nevE
+ mollex 11?
(++1) a = -
64mKi
By the condition on the small initialization in Condition such that = O (203- due to the large C. Cm3374pm
we see that the k = O((203 - 1)(/m) is far more feeble. Thus the gradient contributions made by neuron set Elobgaa where BB (t) 0 (WI, (())( can be neglected compared to the increasing update of Ela 1li E 2E, (WW,5 - (4,5(( and B[e(2 Eiege (o))) )r - 1))(().),)))) Besides, we
39

---

see that EEWE, (t) will at least preserve the neurons of EIME, (0))) which will not be deactivated by Lemma
29
Then there exists e2
> 0. recall 0g is defined in Lemma 34 the lower bound of
mine of R.
E Se
(og " and ,. 2 m/8, it holds that
(1)
nEDS
E V BE Σ (os+")))) 2 oz( 72(25' 1)a 1d<l"
1 D
8mKi
"EV.
JES
A + C271 1) 1d
8mK1
+ t
(53)
A. + (oup (20g - (1 |11l
16mKi
1 1 ee
A
1 bi
where we ignore the impact of regularization term at this stage since À (((Clog((m///)]ll)) and appropriately chosen e2. The third inequality is due to the definition of d,.
Collaborating with Lemma T0 the proofs are completed.
For the last results, following the techniques [43. first it's easy to check that
the ct S 5x + e k.
S 1..5 "t gheast S It + ghic et S 11.5the die
thus
log( 2
3he 3) x x/t < log( ghe
t +1,
log( 20
38 + 33 s g
h.
< log(
log( che
+ 1).
ghe
ghe
Thus
log( Re
2 SE(A s log(
the
1)+ the
1 1 Kie"
Proof of Lemma 30
We use induction to prove this lemma. All conclusion holds naturally at t 0. Suppose such that six conditions hold for any S t S T we prove that these conclusions
there exists T 1 also hold for T.
We now prove
0 B(e-B()) ((0) Pou.. 1 Uh (t)) e e0000) p (o5) (o5))
a.
(54)
Recall the update rule
Doa... = (1 - 11))300)
|dzll 12
Σ
rs E
neve 1
(0)
(2)
Σ
(o Σ (o
2K,
ce[t]
LES
LeST
As we ignore the regularization term at the first stage. we can easily seen that Ele 8(0)
(i E ((
increases with L. Assume lot.k & as the last time 3i E E[ (t) such that Ele 2(0) / 1 E (E)(((()) -e-
40

---

8 S S ~003))'oo(((TT "oog(T"), then for i ERGE (1) we have
Ble. 30 p. al S llle. 2
|| 2
- 1lu 2K1
Σ
E, R
18(1)
( Σ
(os
(0
Σ
(og
10) )1
cE(2
IEST
iEST
||
B (1)
(o
Σ
(og ()))))))
2K1
'st, <
cE|t]
1ES
IES
S e p. k + (dis) log(T") +
Mol dell
no
nolldall 112
D[e, (4)
2mKi
2m
f0l a SteT
molldal' 2mK1
(0) Se. So. ). k + +2(00)) log(1"))
Σ
Ele, (0)
1o (')
where the first inequality by the positive nature of regularization term as well as the contribution of the
gradient; second inequality is by El-e,, (1 [') < and ([iEss,, (o(")) - Diss, (o5 S 1; the
third inequality is by the condition njo- O ||al and thus d, 12 S 1 S log(T'), ) as well as (o3) -1 2 1.
2mR1
The remaining job is to prove that
willdall BIC (1 18 " (03) (os))=log(TT))
2mKi
Tht SteF ReeF
Observe that
[ (9
[8[ I|
1 1 exp(ys.
(Ere(t)
e Exe(r= = DR(W% ) (().))
(1) Eieuy (o (0))) y(*))
m
< B[exp
Σ
Habo j +( Σ (o - 1)ys, B (8
kew M., (1)
KE (1)
< Blexp(x/2 m
(20's 1)8s, 80 (01,,
FEM 6.m
USw (0)
< 22exp(-log(T")),
Here the first inequality is by 1/(1 + exp(z)) S exp(-2): the second inequality is by Lemma 31 the last incquality is by the feeble k/2 and Ele. Dod (i UE (t)) 2 (oz) log(T"). Then we have that
vol|de 21 2mK1
no 11l1de112
10
(9] s
2mK: 2exp( - log(T"))
fat & K
(gt a &teT
Two |d 1|
T' n0l del (68) log(T").
mki exp((
T-
We complete the proof that < Ble-
(90) Po 11i E (t., e- 300) Bo S 3(03) -1 log(T") s (as)) la
We now prove strong augmented hypothesis that there exist i" E E/UE, (t)) for 40 S15T". we have
Elloo el((e (...., kll S E ||c/||2
(55)
0 |d 1
where we set C = 2C", 22 22 log( 3Km for some constant C³, can be any element satisfics 1860
a 22],, which exists at t = 0 by Lemma 7 as well as the fact that ||cx || > ||du|| by their definition in Lemma
25
41

---

Suppose Eq. 55) holds at 0 S / 2 T 1. recall the update rule and the large K condition, we can have
constant C > 1 such that
Ele
1))] = E[(1 ( Sou.). d 1(i E Uf -))]
B(--11
1
i (T
-1)]
d 12
ETe, (f-1)
(0g ")
Σ
(o 0T ")
7 2mK
IES m.k
1ES
2 E[(1 - - Aje-
P 12 Ui - 1))] +
2mK
BIe (f-1)
Ellao,,, ell = =B[[11 7/ -/-
1)0
11
-1 2
| 1
Σ
ler,
B (Cn
(2-1)
1010
(f-1)))- - 11--1 2K,
(Ki 1)llce11
Σ
eer
neVE
ec[t]
cE,+]
(M (T 13 [F-1))]]
E nev
Cng-1llca11
< <
1
11 bell
Blle, (F-1))
2mK1
where the first incquality is due to the definition of o and (t) the second incquality is due to the large K = (K(poC(K 1|] (m//) Then we have
Bl B all
o d3 } S & ||c*|| 03 [d
C'Ick 2|
< max
(0).. Ble Po..... & 1 (1 E (P),(())
Ele $ sot
1) Ii ut (T 1))]
where the last incquality is by the induction hypothesis and the C' can be taken C, which completes the induction.
We now prove
02 Ele 360
90.. 1 Wk.n (1) 2k,
Ellcx 1
|d (1
+lllcal1)
/
5Km
>
agdl|del V log(
22og
03 de
Recall the update rule
||d
Bitt) k = (1 121)3000 Op.).
Σ
r, E
18,0 1
Σ
(o
Σ
(a
2K1
neve
eg[t]
LES n.
Fasy to see that Ele- 166) D 1 (T, UE (T))
0 and it's decreasing. As we know that the neuron E (E(Wk.(()) (1 (t)) would be dcactivated at once
Bla +e.
Σ
(o
(e+1))) )r
Σ
(o ((tt)))) ((t-1) )r
LESE
les-
This indicates that for the neuron i E E(WE... (1) Uh (7))
B|u n
Bl tel
(o )
Σ
(o "0) 000 o. 1 0.
1E
LEST
42

---

Now collaborating with Eq. (54) and Eq.
55
we now can have
Ele
Σ (a
300
(G.).
Bllu (2)
LES
M
Cl|c2|?
03 /| ((03) - aate e-Bo
C ((05) a-
331ld 1|
5k
V 2 2oo
Jon
|| Il)
E||c211
Ele. 80
Su,,i 1(i E - (t)]
03 1/21|d111
||d
01 (a3² 2
|| 2 cll)))
og 2 de
2 log( 5km
5
The proof completed.
L1.1 Expected 0-1 loss Convergence
Lemma 33. Under Condition 7 there exist constant C 0. after at most
f CoormAKnny/11 Ky) log(5Km/6)
(20g - 1)² - kl)]ggl Lg- (E(U')) 0.
iterations, we have Lbg (B(V"))
Proof. For t S i, recall from Eq.
53 that for the period < f; holds that
c2n 1)² 1/12
E R.c (41
((()
A
E neve
(0)
1)
16mKi
Note that by definition A kic
0, and we recursively use the equation times
can(2r's 1) kylll
E keg
E nEVE
1C, (5)
16mKi
For each &s [1]] []. denote by 7h the last time in the period [0, T") satisfying that P[AA.) Then by Lemma 31 we see that
E (e()
S 3x//2.
nev
Thus there exists positive constant δ such that nEve [e, (9] 2 C for 515 Pie.. Then we have
E[A Coun(203 - (1- lllll
16mKK
16mK2e
Therefore we see that Vk [K]],e E [+]] E[A["]] will reach x within
Ceon(203 1) Ky ||/l
epochs. Recall that in this first stage the impact of decaying learning rate is under controlled by large Condition I as well the slow quadratic decaying speed of ne. under which have (mr) By
801 |ل//| / (1 sy) log(5Km/8)) see that there exist positive constant C1 ((11/(Ccz)), the threshold time can be
t CiammAin Ky log(5km/5)
(203 - (1)
Then by definition of 0-1 loss we have
Lb-1 (B(V²))
wD (ys., ((E(S.). 0))
P w (E[A kg x/2
P P (E/r/2 = 0.
The proof completed.
43

---

1.1.2 Period 1: Decreasing Period of Correct Attention Score
We claim that if Zie(m] ro j this decreasing period due the expected gradient formula in Lemma 27 Our aim for this period is to examine the lower bound of the attention score during a limited number of iterations.
during initialization, the expected attention score will not experience
Lemma 34, Under Condition 7 for vk E [Ki]] after at most a certain iterations
Ti C3or KITV 10 (00loo55mmmjjj + e 2ng (b.ll
2C4 [04/|dkll/ 1111 2- 20
where Cs is positive constant, we would have the 3 5K monotonically increasing during remaining iterations Ti S S. T Besides, it holds that d'g is the lower bound of the lowest correct attention assignment along the whole iterations:
03
min (AEER])))
IEIS Σ
JE8
25
Proof. By Lemma the g((+1)
Bl0kk
will contributed to increase by
{i EIW,, kin (t) rr Po 0)
(
and they will be contrihuted decrease by
{i E ETW k,n (t) r (60)
Pocc.).. <0}
By the fifth inequality in Lemma
7 we know that
(i E E.. r P
> 001
m
S m
[[ E B[W r, 1 < 03
m
m
4
16'
4
16'
As I. I., 19 nEVE
10
( E)Esst
(og
(o$)) ), " is shared by all neurons, thus whether the 8
and K. will be contributed increase decrease depends the signal of E Ex (1)) re eg(e)
(,,g. By
the last incquality in in Lemma
7 we see that at initialization,
Σ r POG.
log 5km,
50 borl|dall
(56)
16
FERDW (0)
By the expected gradient update in Lemma 28 the e will grow in r. 's's direction along the whole iterations. As such, the values of E(r, g(t) a limited epochs we can have
Gogia lli E W (f))) wk [Ka] will grow larger. Therefore, after
Gg.g 20
WEEE (0)
where the (()) QQ and g K.k would be contributed positively and monotonically increase.
Now serve to find the lower bound of the evolution of E((Eies,, (o("))")) which is clearly to be the first
iteration where the negative Σ DE (.,(()))
(6))
F,
has grown surpass the By the symmetry property (c.).kk
denoted in Lemma and Eq. 39
we have
I
Σ
(57)
2015) /11b2
nev
SE8
1
Recall that
4m0
Q.k bell
g0++1) =((tt()) (1 -
Σ
((
E Ie, (6)
18 12 (0)
Ki
Oaiii rk
neVE
eE(=) [m(m)
(ol)) Σ
(a (o()))))
jESt
jEST,K
20 =(1 12A))(8)..
1/ 2K1
Id
Σ
r
E
1
(4)
18,0
(1)
(0 (1)))
(a )"),
nEvE
<E[+]
8
188
44

---

and we also see that
BIC Σ (og (o
-0(0) (exp(e PK (cxx/lbb/l1)) + exp(- - PK bb> 19) -2 Σ (ob") neS
1
4
(58)
As EieE(wE,. (1) To e0 will grow to surpass 0 in limited number iterations, we can claim that there exists constant C², such that for the limited decreasing period of TB((Eies; (og "")) we have
1 -B(C, ( 2 Č. Also, m BIREE....)) 2 m/8 by Lemma 7 as well
fact that E[WE, (t)]
will least preserve the neurons of -. along the iterations, without heing deactivated as discussed in Lemma 29 Also, we note that in this hypothesised decreasing period, the absolute value of the initially negative
E E o al and initially positive & will all decreasing. Then by Condition we see that the small
initialization of MLP as well as the small regularization will make the decreasing order of ( negligible, as
4.8 ||bball
Il
max
-1 Q.
Σ
Σ
r. r
BIC
Σ JEST
(og
) Σ
(o5)) 13
Ki
FE (0))
JERT
(A
50 |||||
||ll|
2 2log
5km
1)
32K;
5
< O(1/C).
(59)
Here the second inequality is due to 5 (56) and the definition the by in Eq. the condition À I (Coo/2[[u|] and σ (Cool|ull" (ogg(5 K Therefore, it holds that during the decreasing period of (6 well the period where EiCEINE (c) we have
27 the third inequality is by
(60) p remain negative,
Σ
((++)) Googk
Σ
r 3(0) D
Conolddal
(2
2. /1b,112 1),
Kn
SEEIWE, (C+D]
KEBIW E. (0)
1
Here, by appropriate chosen small Ca. we again ignore the regularization term at this period due to X O(CClgg log( Km/6)]lql)) -3) for a large C by Condition 1 and the impact of the learning rate also controlled due the slow quadratic decaying nature of v/. and a small initial no S ((0.01C-)) by Condition so as the 2
changing amount 1/(1 + e
20 /llby by 59.
Therefore, by Eq. 58)) we have
g(e+1) 2 gl (1+
bel
F g
Σ (o
(og))) Σ (o
K1
(60)
WE
(0)
jESt
JEST
where the incquality by negative nature Σ
Eierpa (9) Po g(t)
(u.yy and the decaying nature of n. and
Eq. 58 Now we can see that there exists two surrogate sequences
g(t)
QA and En (0)
ege O & as
the lower bound sequence of the g6.0 and EiERPOtt (1) ri G k These two former sequences's 's initial
50
|d and their
values taken the lower bounds of the latter two (oo boollby 12 and 11l-22222000 /2log(5km/6)
16
update rule are
(1+11 8(0) Q.
+
g(0)
Callllbbll*
Σ
g rs
Σ
(o
Σ ()))
Ki
SE (E.. (4))
JESt n.
JE8
Σ
ri Bbttt
Σ
r Teeett)
kER( k.n [++1)]
LERI (0)
Cimildaa??
1).
K1
20 2 //1ba]1?
1te
Then by Lemma 11 let a Cavolibbll'
b= Cavolddall?
(2
1), have the maximum
Ki
Ki
1 20/00 (0) /1b211?
=0) + e 22(0)
σ.KIn log( (5km) () - 20011211
iterations Ti
set C3 (2C()
b(1 - c -2((0)²)
2C, (ch (1 - 20611bb
45

---

we obtain the Ti in the lemma. The lower bound of Q.k along the decreasing period as
25 b& (1 20281116 (og(5mmm/))
B = oo Trolbbxl 1iee
1024(1 - -
Utilizing the seale bounding property (-rig + 1)/2//uull < ||,, < ||||| / in Eq. (7 and 36), we can
denoted the lower bound of all &Q BK, for V/ E [Ki] as Box- which can be given as
ollllll + -
Ka ///l log(5Km/6))
0u(1
(1 e -
Pak-
2
Recall d' is defined as
03
|u|| + e -oollal"
-2m1 log(5km/5))
(1 - e ---
which is actually can be written as
-2P0k 2 1 e
Therefore, we see that oÿ is the lower bound of "rii
(o (09)" )))
neDs
Remark 4. As we see that in Lemma 33 we require that the lower bound given in 53 depends that the values of (2 (o's 1)88 surpasses K. which naturally says that the
value of ExEug (o (onlr: should surpass since El(2 Ei(2E (o -1)] 1. Therefore E.gug, (o)) (on)oorrre (u.,,, should surpass O at t since >> 0, which indicates that T Tı. We see that the initial period Ti where BuEue (0)) ( . grow surpass the initial scale, whose upper bound k/8 by the definition of k.
1.1.3 Period 2: Increasing Priod of Correct Attention Score
This period's 's analysis is based on Period in Section 1.1.2
good initialization such that
Σ
OO. O.
#E[m)
Lemma 35. Under Condition 1 consider the duration after T1 Lemma 34 then for Vk [K1], consider the period Ti $1 & T2 Cs mind 1- [Hll||g] ). where Ca is small constant. Then the following holds that
We have v(t), T(t), 2(1), Z(t) be the lower and upper bounds of the increasing P 9K. and ExcelwE (0) Fs (00 respectively. That is, there exists positive constants C3-6. for
a cs(1 - Ko) ulll c/|l|| b cs(1 - xl ) b
callgll?
XyKi ) = Č. holds that
AyKi
ArK,
AvKi
((() < ((t) 300 s((), z(1)
Σ
g Oo.g. $z((),
kERpwo (0)
46

---

for all L 2 Tr. Here, (((), y((). æ(I). =(1) are the unique solutions of the following ODE System respectively
Bi(-2y(t) + log(y(())) =bb" (203 (1 t2)² + 2 2 (Bi(log( 03 -
2
)) + Bi(log( 1 1-08 + +lgg(g@k)
2(x) be (203 [) Tı);
abl?
(Bi(2u(t))) + Bi(-2y(t)) + log(u(()))
2 a a
2
1 (Bi
aه ||ll
) Ei( od "'llul",
+ 2
)) + (///0//2|||u])
=(1) bt
8
where
L = dt = "Yeate + In + exp(x/2)
(-1)"=1 x"
Ei(x)
{
2
2k +
For some limited constant such that 3A, oZ < make the correct attention score
< I < 1. Then 8 = 3 will grow to
1 EE
(o achieve the A in at least 4(1 - 4)
nEV v8a
scaled Gaussian rate such that
abc'
A(1 4) - 1), (t - Ti)² + log(Sxx))
2 expl
Proef. By Remark we see that at the initial phase during / 2 Tı, we have Eie(m) 1. S k/8, and thus
||c |l
by 55 Lemma we see that a
/ | This indicates that E[A,"] S ((a) and thus by [e (E(S;:
* C
(L.)..
0
Lemmar 31 we see that the scale of
E
is also O(x). This suggest that there still exists MII
nEVE
a constant C, during certain amount subsequent iterations we would still have that Ō -B(L (t)] 1. Also, m 2 E,I (0))]] 2 m/8 by Lemma 7 as well as the fact that E[WE... will at least preserve the along the iterations, without being deactivated as discussed in Lemma 29 In addition,
neurons of recall that in this first stage also can control the impact of regularization and decaying learning rate by small À and big y by the sufficiently large C' in Condition
which indicates
have
g(t+1) ((t) + + (38) Canoolball Q.
Dankk (((I)
K1
FERD A. (9)
1 -
20(1) /llb
e
Q
/llbe 11bba1"
and
Σ
Σ
r
P
KEE ((+))
((DD)
iE (0)
Canol del
+
+
(()(1))
1 - 22 /11b,112 111
Ki
By Lemma 12 we see that the iteration satisfies the ODE System 2 with positive initialization, where the O Canjol dell ), C. Then by K
[E]" o Caroo bbll'
parameters in Lemma 12
Ki
solving coupled ODE systems, collaborating the scale bounding pooundingppropertt (- + 1)/2 u/l < | 1|2 ||u| /2, Kk + 2 < |dx,, 1| < ||qll in Eq. (27) and (36). as well as the Comparison Theorem with
some constants cs s, we can have upper and lower bound of 3(4) and EiERiwet (0) which the result this lemma.
For the second result, given the A, can directly have lower bound ODE y (f) be the lower bound of the Q. via Comparison Theorem, where y (t) satisfies
(())
y 2 v'a (t) = abc' ( ()0 - 1))(t - Ti)yy (t)
((() 2 y (L) = expl abc' 11 - () 1) (1 2i)² + log(33x))
2
47

---

where the inequality holds by the decaying nature of g(x) = x(1 1) when > 1/2. The proof is completed.
Lemma 36. (Asymptotic Property 1). the first stage, the growing of $
3 as well attention
score enjoys the asymptotic property that
E ((1),
E
REV.
1
Sn
(os))
lim +-000
lim 6I0
(().
log(t)
1
Proof. By the asymptotic property of Ei(x)
Ei(x) Ei( x))
lim 2-
exp(2)
This suggest that when y(t) 2 y(t is close to infinity, the lower bound ODE in Lemma 35 will approximately satisfics the following
exp(2y(1) 4y(1)²
+ 2 log(y(t)) abc"² 1) 12 + const,
This suggest that roughly
lim g(t))/log(()) O(1).
Then we see that as y(() goes infinity, we have a bound
F. ElEiEs (o(")
lim log( 6-0
1/2llo((1)) @(1)
= lim 61000
E
(1
= (1).
B|Eiess (o
I
iESE
On the other hand, obtaining upper bound over yv)r is relatively casy. Since have 2+e 2y(1)² ?y(1)
<
and (1 2y(1)² )/(1+ 20(1)²) I 1, which gives the upper bound ODE over attention and MLP considering z(0)
2 (2 Ei( (i(-22(()² + (llo(((())) = abt
K
+ const.
((() = bt const,
where the term "const" ensure that (00) = v0). The asymptotic property of this ODE system is the same as the of lower bound ODE. Then consider t, v(t both go to infinity, we have some such that
(t + 2/
lim Σ
lim
Σ JESS
=1
1+ (t 60/2)"
1+
JESC n.k
I.2 Second Stage: Regularizing the Model
As the 3/1) ((t) and e0(()
....t. continually growing up, see that the dccaying -E(e" well as
the decaying attention products (Ejest, (o () E,est, (o()))) is becoming feeble and feeble, under which we can no longer ignore the regularization term safely when estimating the coefficient gradient dynamics. However, although the regularization can prevent the coefficients from growing. it will maintain their scales without decreasing them.
48

---

Lemma 37. Under Condition 1 consider = Elys,, f.or all 1 é [Tr,T*] it holds that
eg/
(6.) 0 l11l2
AmKi
e0
(.) O 035 (1 ky)
log( ""alal (203 -1) 1),
(1 + Kiy )2y)
AmKi
|1|1/||a|1
00 = 00 ||||| [
X²mKi ))
||ul|'03 (1 - k1-k)
e- (203 - -1) 1)).
0 = e(1au]|
log(
log(
AK1(1 + kw)
AmKi
E((
(o5"),) 01
1 2 ||
A²
JESC
mikk
E[( (o )1)1 = 0
1 [1 (1-Kk)) log
(1
" ((20' -1)
JES
Am
where e(t)
G..). represents mr: S..).. That is, consider the positive growth Bl B
Proof. We will prove the desired argument based on the following induction hypothesis:
g. p. =0
AmkK
2 u 9|
0 Olllall log(
)),
X²MK.²
We split the situations into two cases:
(i). els() S 0 (1--0))) log
(20
AmRi
and a(t) Q. S @1111]] y log( [11l)~32(1--p)) log(
222
Am
(ii). o 2
1|l12 2XmK,
and (t) 2 |lully/ log( 2 = TE. (
(o (og))))
1+ Xm x Easy to note the scales'
[q]2
orders of the case (i)'s quantities are less than those of case (ii), thus this split is plausible. Recall
Po = (1 (1--1/7 11-)))00000
||dall
Σ
E I. neVe
(0)
1000
(t)
Σ
(6
Σ
2K,
ed[t]
LEST
Then it's easy to check that for case (i), as by Lemma
30
we see the magnitude Ella00,., is controlled
by some Cog -3 + )2 Ky)) -2 P00 o
That means that the term E EEA k. can be controlled by
its contributor (() (n.g.))) R). Then we have
((e
O [2| E (1D s O(eg!)
( t
-3- (11+m)))( ( (..).
s e oB ((tt)))(1
(g0(t) (200)) (..)).))))))
so
XmK1
where the first incquality is by the definition of El- (1I) (similar to the techniques in Lemma 32 and the definition of EEAA f-";; the second inequality by g(x) = e" 1 well as o'S -2 k.y) Ky) the second incquality is by the case (i) hypothesis. Then we would have
60(-
|de
Σ eE[+]
Eir. E [C, (1)
100
00 Σ
Σ
(o
"...
(61)
2K1
nEVE
1E8%
188 188
49

---

Thus the growing of 1. would be non-degenerated: Ele 2 ( which directly suggest
BI = 0 (1- Similarly, for the A()) in case (i), first recall that
log(
""lla "Ila|" (203-1)) )) holds since 7 the maximum admissible iterations.
[[+my]
AmK1
4m.5kA
I1bu
A =(1 - 12:1)
K.
ΣΣ
13 ropo
E 1e
" Σ (og"))
K1
nEVE
FE/m)
JES
(1 Σ (ob)))),,
JES
then, as we see that
B(( Σ (09)))1
Σ (o5"))]]
1
+ 2.0 2
// + 2 /1|bE/1
JEST n.
SEST
+ 28 /1bb/112 - Q, 2
/1/bx11
S 0
2 ||u|| (11 log( 1111
AkI(I+ky]
Amhi
0 ||ul] (1 Riy)
log( ""1laa 2|ll 1) ))
AK1(1 + k2y)
AmKi
where the first incquality is by the definition of EE(E (o5 (a)"))) and the induction hypothesis; the second inequality is by the small À by Condition T with sufficiently large C. Then we that
1||bal O K1
E Σ
r( xE|m)
Σ
E NEVE
1e (4)
10
(9( Σ
(o Σ (0)))
JES mik
JES
Here the incquality is by the case (i) hypothesis upon e
-E(E(t) < 1, and
B
E
nEVE 1
=
/m] 2.
eE(L] (E(m]
<E(=)
IEWS
Thus see that the growing of S would also be non-degenerated: validates that for the maximum admissible iterations T*, it holds that
3 ((+1)
> ((((()) This also directly
| 103 k2)² AK + x2)
log( N 11/12 - 1) 1)),
| log(
AmnKi
E(( Σ JES
(o "
00
14
(1+k)))
log 1( e- |(al| '(20'
|ull- (1-k)))
Am
For case (ii), we directly check that
|d ll 2K,
Σ
Eir 12 (1) 11 18
(0) Σ (o!0))
Σ
neve
1ES
lESt
IES n.k
Here the inequality is by (t( < and
[ Σ (
Σ (o
[E(2 Σ (o - 51
JESS
JESC
As result, by the gradient form we that (t++1) Eleg((+1) = 0(
2 and thus prove the induction proving goal
[1q
Kmk5
50

---

Similarly, as we now have
B(( Σ
(a
1--
Σ
(o (
)
1+ 20 2
/1lbe 1 20 2 /1lbe
JES m.
JES
+ e 28 /11bu11
+ e -20 /11bx11
28 /lbull? /3
1 + e
1 2llu i Iall"
X²mK,²
e 2 2
where the first inequality is by -. 2 /11b,,1? S 1: the second inequality is by the induction hypothesis of case (ii): the last equality by the small À Condition D for sufficiently large C. Then we observe that
2 0
4| K
Σ eE|t) ic/m)
rith Oa. nEVE
10
(
Σ JES m,k
(of Σ (,)))))).
JES
where the inequality is by -B(e (t) 1
R Σ Σ
1 (0]]=
Σ
R
Σ
/m] 2,
nEVE
e kE/m]
cE[#]
iEWE
as well as the induction hypothesis upon Du,)) in case (ii). Thus o < ((3(()) which support our proving goal in this induction process:
Bloh112) = = log( 2| 2|
)). X²mKi²
addition, we can see that even we suggest the MLP's e &
is growing a fastest lincar-level speed,
require at least O + to reach the maximum admissible value 2
Meanwhile, we see that even when
considering the fast speed the increasing attention, by the asymptotic perperty discussed Lemma 36 we
see that we still require @ ||ا|| to reach highest admissible correct attention score
XK
X² X2mK
1+
|||u|||
Therefore, have some appropriately small constants Co. and when the iteration number is than
T- Cs min 1 + ? ||l|||l|l we need to consider the impact of regularization.
X "KKY/m
Lemma The scale of the coefficients will finally be stabilized at considerable level:
||q|| mXK, )),
(T') 00 S <eb0 = O(log(
B = O([|u|] y (log( ||1|22
log( |q/l
))),
AK1
mXK,
E((
Σ
(o (oSTT)) (
O(
1+
AK M
log(
mAK
JES k
where D represents mr mr. Br.s
du.g.t That we consider the positive growth of (.f(t)
Proof. Recall last discussion in Lemma we see that as E[(2 Eiese (a 1) ( getting
larger and larger, it will finally reach the scale of d'o k. which has updated a feeble speed controlled by
initialization when the neuron fell into the neuron set E 14 (1) n (WER(0) - 24,1 After oougi S NVVE
51

---

E[(2 E[(2Eiess, Eies,, (og)) -1) J. the neuron would change into the neuron set
IE
M (1) (1))
nEvi
Ug (4)].. As such, the d'0 )* would again increase normal speed, which is faster than e y due to the update rules and the fact that ||c=|| > || d As such, the neuron set n
R (t)
(f))
would again fell back into the neuron set
[M (t) n (W (t) - (t),, where the update speed again E nEVE
feeble. And it will increase until [[ Eiese, (o De (a. catch up.
Besides, we see that the expected attention score will grow up considerably, where can see that there exist some constant č > 1/2, ō < Bl(os ) S 1. As such, ultimately we have BB(Eiss (o
EiEST (o))) (11), ). S eoog. and E[A,"" @(Ble5
Then following the process in Lemma 37 we can obtain the results. Here we omit this part since the proving procedure is
the same to Lemma
37 despite we see (EiEs, (1) -iess (o{)), = @(1) and B[A
have asymptotic property when considering the decaying impact of the learning rate, as well as cross-entropy loss. We directly provide the following two lemmas. Due to the similarity of the proof procedures Lemma 35 and Lemma (3 we omit the proofs of the following two lemmas as well as the constant details for simplicity.
Again, similar to Lemma 36
Lemma 39. (Asymptotic Property 2). If we consider the impact of the decaying learning rate at the sec- ond stage and do not consider the decaying of cross-entropy loss, for some constants c, d, e.d regarding Ki, 7 I |, |gll, Say Ky, we will have
v(1) S PO
sk S T(),
2(1)
Σ
Te.ePo S E((),
VEEWW E. (0)
for all 2 T2. Here, (()) y()). 2(1), z(l) are the unique solutions of the following ODE System rexpectively
-7 - 7 log(et a) log [c( 1) ττ d
J + á
[Bi())))))) Bi(-2y(t))) log(y(t)))
Liz
03 )) + Ei(gog( 1 03 1)) 4 (8xx) (Ei(log(
- 03
+ 2
I-o3 03
z(1) = be (20 1)( Ti),
d gt
log(et d) log c(7 + t) 57 d
a Liz
7 - g
1 (Ei(
adl
+ Ei(-2 2
اBllاll" )) 4 loggoo//221l]
+ 2
2
E(t) Tt
g
where
S-" In(1 - -) dl.
Lia(2) = =-
Additionally, would have asymptotic property that
E
[ ses (o(")""1
lim v(t)²= = lim O(log(log"(t)) 60000 t00+00
lim t- +00
= O(1).
log (t)
1 log (())()
Lemma 40. (Asymptotic Property 3). If we put our sight on the long period and take the decaying property of the E((((t) into account, for some constants a. b, d. a, b, d. we will have
v(1) < -/ = 7(), 2(( E
Σ
O.. $ $((1),
SERTWEE (0)
52

---

for all 1 2 T2. Here, T((), y(). =(1), x(L) are the unique solutions of the following ODE System respectively
Lig 3(d + σt)
log(at + a) log( b+ ejt + d
b + dj
b d dj
2 + + log(y(())) =ā
2j
03. 1) + Ei(log( 1 03 و )]) + log log( + (Bi(log(
1 1-
2(t) be (203- 1) Ti),
Lie( j(d + gt)
b + cjt + d
loglet log(
log(
- d dj
b+ d di
=a(
Ci
ej
Øôllu]|*
oôl ull"
(Ei(
+ Bi(-2
log(ou/ ou/ /2||u|l
2
2
2(1)
Ee
8"
where
In(1 f) dt.
Li
Additionally, we would have asymptotic property that
E
Σ 5e US (oS")"[1
nEV.
lim v(t)²= lim O(log(log(()))
lim
O())
log (t)
1 log (t)
It's obvious that the decaying impact of the learning rate and cross pross-entropy loss are at the similar order. Also, if we consider decaying learning rate, the right side of the inequality would be smaller. =tt) would be in (log(logg()) order when =(t) get large, which will make the right side of the y(t)ss formula contain an intergral of (log lfg(f)), which obviously slower.
J Exponential Convergence of 0-1 Loss
We continue proof after Lemma
35 In this section, we assume all the events in the Section D hold, denoted
Tre--
Lemma 41. The Frobenius norm of W and its gradient can be bounded:
11 0 Kille Killall? )
LB = 0 o
Kille Killgl?
A2m
m
Proof. For Vi [m], by the gradient update rule in Eq. 22 well Lemma insight of the WV certain projection direction will continue to grow until being stuck by the regularization, which a X-sealed W itself. Due to the low-noise condition in Condition [ with a sufficiently large C' as well as the isotropy of noise. the learning progress of features would be the main contributor to the norm of NN matrices and the noise, validated in Figure (iii-iv). We can consider extreme case where all the samples in a single batch belongs to some concept [.). which the inequality over the k-th concept's 's corresponding projection direction, and thus we can derive an upper bound
that the lengths
have the upper bound of the first term of the right side of
(W
(e)
q
)
S -2
(Ilea dx m
30
O
|| X2m
(62)
where the first inequality is by (2 (iEss
(o 1) < 1, and Lemma 6 the last equality is by the low
noise condition de < || /dy in Condition Then by the low noise condition as well as the data model's definition we see that all the 2-norm of the W is controlled by the K1 concepts' corresponding lengths in
projection space. Then by the definition of Frobenius norm and (2)) we have
Killo K||| ) Wg Le, (1) < e K1l1
will < 5
A3m
53

---

Lemma 42, Vor E [Ki], of We ug. aj W (0) ag and byg wa (o bj W bo satisfy
a we ag az w a O(oollul²).
by W (9 by = O(]iully y log( (L. |)ll )).
bg
X²m
Moreover, the Frobenius norm of We (1) and W (9) and its gradient con be bounded as below
IW@ W (1) f = O(Ki log( (L [ 1)|l///|al??
2mm
IV (V 0 01 Ki(L |)|l |q
V Wy Lo,
Proof. By Eq. [6) and we see that
q.a. chhos 1 counll [10)
(2) Keaggehaosl eehaall Kiag. sonel S T/e (maxxxx( w2
agl |a] w (9)
de 2 log
KNN
|
/gll/Am )) S O(A ma w ail la
W
a
K
Here, the first inequality is by the scaled identity initialization of w (0)
w (0)
orthogonal relationships of
vectors in Lemma 27 Lemma 6 Σ Ei (og (1))
(03 (o)) I 1/4 Eq. 62 in Lemma 41 the second inequality
is by the low noise condition og S Am/(CC//a [[Cxxx/|lul[a]]]]] ||]||||||| 1/2) and the large K Clull/ (or /dx) for large C' in Condition 1 Thus the update of ag Q 04 and ag "W7 (4) and thus the scale can not be better than the initialization. By Lemma the conclusion holds.
ag are dominated by their regularization,
On the other hand, we see that by the scaled identity initialization of W (0)
W (0) and orthogonal rela- tionships of vectors in Lemma 27 the initialization of bj WT by and bj W (4) by are the same, and as the gradient update is nearly symmetry, which can lead to the fact that @(b; ( = @(b; W bg) and bj
K
W
Q W by e(b w (+) b&b W b/ /| (6,11 12) By the scaled identity initialization of
Wa
W K
orthogonal relationships of vectors in Lemma
27 Eq 19
and 20
we can see that
(7())
(())) S 1e0 (1x w bel 1
||lll||l|l
144.b, comal K.K.b.comii s <
11
[ w ( by 16 W (4)
b+
< (a)
Am
SESt
Σ JES
(o (o)))))))
We see that b._ ization. To see the upper bound under this situation, we consider an extreme case where all the samples in single batch belongs to some concept k E [Ki]. and there is only one demonstrations in each prompt share the semantic with the query. Then by the scaled identity initialization of W (0) W (0) orthogonal relationships
and b. W (b;; will continue to grow up except always being stuck by the regular-
of vectors in Lemma 27 and Eq. 18). we can see that the growing of b. w (1) by and by W (0);; would satisfy the following and strive to grow up make the equality holds, which naturally have upper bound
Q
b.b .comil 14 .omml Amint(b w (1)
w (c) bul
|all
Σ
Σ
(o 2 ((A)
Amn
JEST
JEB
Ejest
2% w3
wx W Σ EseS",
c oo w (t) CE
| ||1
0
-b we (e Wba 121
Amm
( Ejest
2 Wa (t (0 + EjegTA
E
||||| |||l2||||2
T-1 (T, 1)e
by wy wgtoww 211
X2m
(6
(L-
1)]lul/ ||ll - wa uw (0b))
2
X2m
1 (L- - 1 u aall ))
6. W ( bg. bj wx b s e([lully /
log(
X2m
54

---

Here, the second arrow is by the definition of Eestt, (o and 31.31 the third anrow is by our considered
31
extreme case where there only one demonstration each of the prompt sample in this all-the-same-conee batch, which is considered for obtaining the upper bound; the forth arrow is by the small À by Condition
which denotes e e W (0
(00gg
should he the key contributor.
Similar to the claims in Lemma 41 here see that the À is very small, by the scaled identity initialization of
Wo (0), W (0), orthogonal relationships of vectors in Lemma 27 as well as the low-noise condition by Condi-
tion it's safe to say that as the learning proceed, the scales of by WQ Wogbbg bg. bj W by), would completely
dominate a W ag. a W (1)
ag. well as dg W (0 bg. bj W (A) ag. u. W Vr, uu WWx Ou..
XX E K,, E [K2], W E [dx 2K K2). Collaborating with Lemma 9
we have
I 1 12: I W (1) 13 I
Ki
max[(b] we (1) be) (b) W (9 b2) = O(K, log(
(L 1)l) ))-
||/|2
X²m
On the other hand, see that the maximum gradient F norm on single batch comes from the maximum changes of the b._ w (y) (or bo. w bogg ). As we see that the extreme case of the growing is every concept & E [K1] has been fully leamned such that even a batch full of the same concept can not let the
corresponding concept's feature grow. In this case, we see that the maximum gradient F norm should be at the order of | [1 (or |AWX ([.). Thus
I WW co Ls, )1 II LB, (B. (1) = =OA2KT K1 log( (L 1) [|
= Ki(L - 1)/|u| ||alll//||all Iall
m where the equality is by g(Ø)- log(ar) S O(x), x > 1. The proof is completed.
Remark S. Worth noting that this upper bound. well as the upper bound of [W (i.))) qe / /l III in Lemma
41 are looser the order of Ki 2 and Ki 11 compared to those of B 3 and , Lemma This suits the intuition and statistics since in the practical training setting, for B 2 1 we can see that sometimes the samples of batch all belong to one concept, or sometimes their are not any particular concept a single batch, especially when B is small. Therefore, unless we have the situation where even when every prompt sample of a batch belong to the same concept the regularization can stuck the growing, there is still chance for that concept's features to be learned. In contrast, the expectation considers every concept's 's sample appear in every batch scaled by "soff weight" in the order of (11 K). As the attention gradient contain MI.,, its order would be (11KK)) Besides, we see that this lemma's 's result contains the scale of I. 1, which comes from the extreme case discussion where there is only one demonstration in each prompt sample that share the semantic to those query. In contrust, when considering the expectation, the number of two opposite semantics is the same, under which the L/2 would be eliminated in the numerator and denominator. Last but not least, when estimating the real cases, we have scaled the derivative of -E' its maximum 1. we do so because in real cases due to the imbalanced prompt samples in a single batch, it would be inconvenient to consider it is contributed by severel elements like e a O6..). This actually indirectly demonstrates the superiority of considering expectations.
Vi 2 f. when 1/0 B(G || S v holds, we have T. T = T. (R (R(()(()) Here, [ [wa li + W 1 + [W/1
Lemma 43. (Restatement of Proposition
Proof. By Lemma 33
we see that our convergence of 0-1 loss is based on the intermediate result that
E[AR1] 2 which will ensure that Elys,, ((E(S.), E(W")) )]] 2 sz/2. Therefore, when conditioned on E[W3 (0) E| (*)] minimum admissible disparity between w (1) and E|wy (0] corresponds the the
minimum admissible disparity between W W
(4) ()c.. w (6) (d and a(t) (o......n.. where would
(..) consequently cause Evg [AA < k/2 that could have potential to deteriorate the 0-1 loss. Given that
(...))
x/2 2 /201llgl1 by Lemma 23 the decomposition in Fq. (33) (3.33)
well as Lemma we that for some
k E [K1), the minimum admissible disparity can be written as
e(I
(Ir) =
(lr) 2 (12
|
|
|c
Therefore, we see that when conditioned on B[
(*], the minimum admissible disparity between wy (1) and k[Wg ($) to not worsen the 0-1 loss is 21 )1).
Q (0) W K
55

---

On the other hand, when conditioned on B (1) ]. 2 T", we compute the minimum admissible dispar- ity between W W (T*) and E W (T') E
K ]. Considering all the activated neurons, when (o)"), 1) (..). O, we should have Exewg, (0) p , 2 0 Eiewg,. (1) (2 . (ag
(2)
otherwise some of the neurons must be deactivated, which contradicted by the definitions of WE, (t).
In this case can magnify the impact of Eiews, (2r:(2 ( 1)88 by consider-
ing Eiewg,, (0) D = O. As such, the minimum admissible disparity would be the case where (2)
b wa and bj WK (1)b bu both differ from PQ. SK. by the amount fΦx Recall the defini- tion of Fax in Lemma 34 and collaborating with Lemma we have the minimum admissible disparity be
(mm fflwraat -00mm00
du(1 - -ne)e Re
Recall
1m Allulitaa+ ob a- ogllul? }.
v =
00(1 rz)e
the proof is completed.
Lemma 44. For E , W E (WQ. W W and X E {Q, K, о} it follows that
K
- 1)2/2 u + 1) 77))
1. [w(41)- - w. (1.4)-W.(+1)) s E
1/2
2. [[W(+1) (0t2)) F S n.)i)l/
1)11|
3. ΣL D 2 mX7(7 +T)
Proof. We provide the proof by extending the techniques in 34 33 36) Hilbert-Schhidt space, whose inner product is defined by trace. First we note that rjo = 7+1 S min 11 + A), 1/2A]. where LLggia is the I-smooth Lipschitz constant cross-entropy loss E(-). which 1. The first statement can shown as follows. Since by definition a single itcration update, then by Lemma 41 and Lemma we readily obtain the results.
that ( W, (t), we only need check the maximum disparity of the gradient in
For the second statement, following the proof loss denotes that
we see that the Lipschitz smoothness of cross-entropy
(V (V), W W
- - Vw
(63)
Liagae
Then we have that for s 1,
w W. (a+1) || = (1 nod) (W - -W, - 77 (3gl (go, ko) 2gl (o', 21
= (1 - -1a) w - -w, w/ - 21, (1 - nad)-
(V Wwi (V¹) Lo, (V²", w(() .,(
V (V0") V Vw, (V=))
$ (1 - nTA)² - - w/
Leggit
|| w LB, (V²) - - LB, (W0)
(1 (1-n.A)² W w/
where we utilize the Eq. (3) and conditions learning rates. Utilizing this statement, the stable property of stochastic gradient descent has been shown. Again following the techniques [34 36), we now obtain the bound: for L E (I....,T}.
(w(rt)-w,
S <e(
Kt
(
ull + 1) m1)
II - -)-A))).
(64)
m/2
56

---

From the following inequality.
II 71-
II
+ ++ 7
r ++
s=++
s=t+1
where the last inequality hold clearly by expanding the product, the right hand side of the Eq. 64 is upper bounded as follows
K [g - 1) ||l|| + 1) 1m II (1 - 1--)))) s O
Kł lal ((T- 1) ||lll + 1), (r + +t)
2
m2
7 +
Ki [lall(a - ///l| +1)
0
m
X(7+
We finally obtain the desired bound:
T
Ki [llla|l((L (LL((())lll 1|| 1) s 0 mX²(7 + T)²
Kallall ((L 1)]|u|| md + T)
Σ
1=0
1=0
Remark 6, Utilizing this lemma, the exponential convergence over the 0-1 loss is readily obtained.
K
Out-of-Distriuution Generalization
Lemma 45. OOD Master of Polysemy of Words. During testing, The prompt length L can any positive integer. The D' can have any new probability distribution that differs from the training distribution, satisfying that each prompt has at least one co-concept lo E [K1]. with equal chance to have positive or negative semantic labels. Additionally, single (x. y) Da x Dy" pair can appear in at least || || concept- p T.g." =0-'(1((")) SE
specific prompts/takks. Importantly all of the tasks this distribution D enjoy Bayes-Optinne test error
This lemma demonstrate the strong OOD Generalization ability of transformer utilizing multi-concepp semantics. suggesting the efficiency transformer conduct unseen ICL tasks just its learned knowledge the two non-orthogonal dictionaries. Also, this lemma showcases an intriguing phenomenon since allows multiple concepts with comparable chance along word-demo pairs even with the same input-output pair and query, the model can produce diverse responses when provided varying contextual (concept ttas) information. For instance, with the prompt "Japan: Sakura; China:", the LLM may output "Penocy" (national flower) or "Panda" (national symbol), reflecting different conceptual (task) interpretations. Both answers right since they are all the co-concept tasks. Interestingly, adding another demonstration like "Japan: Sakura, France: Iris germanica, China:" stabilizes the response to "Penoey", since the only co-concept is left to he "national flower". In our theory, make elementary explanation to this flexible, context-sensitive in-context learning (ICL) behavior by attributing the transformer's 's ability harness multi-concept semantics.
Lemma 46. 0OD Innovation. During testing, the distribution of D', x Dy can enjoy data shift. Specificully, suggest we now have a new M and Q" to define new Di, Di, Specifically, vk & [K1], bog E [K2), we let
M 1// ai bx. Mr. Me u bi.
Q2x-1 4: =ci + di, Qz 4s = di.
Mrst2Ki = Vb2', Qizt2ki 0,
where
a conic(1
11 ME Ki 1k=l),
by E conic(
11 Mb Ki
2 beslh
2
e E conic((
rt + 9
Ki 1K
di E conic((
9 ),,,
2
2
Vk2" E (span(y4
uT ,
17, pk )²
satisfying
|B > || =(|uul= O(Iuu[)) || 2 ||e (([al)), vig [[[uu]]
and {ai, bilk (ci, di) d are two collections of pair wise orthogonal vectors. Then we can have a "
corresponding new prompt distribution D's =>K, ( Pitu + ax "Pi, :). Again, the model
enjoys Bayes- Ootresl test error Lo-1 (v((T)) S.
57

---

This lemma suggest that transformer-min structure empower ICL ability in solving task involving semantics ("knowledge") originally from other co-concept prompt's 's training distribution. This cross-concept semantic "understanding" ability ensure the transformer perform an specific OOD ability.
For example, when we show a prompt "Isaac Newton:Todes Today designed machine to capture sunlight; Thomas Edison:" to GPT ol, we would obtain an answer "Today I invented a lamp that shines without fire." During training, even when the concept "Inventors and Their Inventions" may not co- co-appear with the concept "Fabricate a story" with high chance, the transformers empower the ICL. to perform this interesting Out-of-Distribuur
task.
We believe this can serve attempt explain the innovation power of LLM 30 95 [6] grounded in the linear geometric property of L.MM representation, since most the innovative outcomes of human being generates from cross coss-ce "Knowledge Intersection". and as it is not an easy task for human specialist to master cross- gross-domain knowledge, we claim that LLM can help innovation by leveraging cross coos--dmain knowledge when deduction over unseen structured task. Similarly, for different concepts did enable OOD generalization (e.g. "blue square apples" the Figure in (86] )
scenarios, (86) have shown that compositing
This lemma seeks clementarily explain why LLMs' ICL can excel in complex tasks when using evolutionary strategies, especially when the LLM's latent representation based on language only partially captures the relevant features. Such tasks include algorithm design 97.4 heuristics [ acquisition functions [98], and solutions to combinatorial optimization problems (99] Although the resulting solutions may often seem counterintuitive to human experts, a possible explanation is that transformers can perform ICL in OOD scenarios by leveraging weighted combinations their updated "understanding" (i.e., changing the identified underlying concepts in the evolution process) of new demo-query pairs, such as randomly sampled TSP instances. These understandings are rooted in the latent structures of the problem instances and can be effectively updated by evolutionary strategies that selectively refine and discard certain outcomes.
Proof. Proof Proposition 1 By Proposition 2] we only need check the expected loss L (B(
Denote E[Mys., C [2 as the expected index set denoting the expected shared concept-specific features by the query and one demonstration. By definition in the Lemma, as the semantic combination is conic combination, we see that E[M2s.. will be either collection of odd (corresponding to positive label) or even (corresponding negative label) numbers, and all of the combination of the features and labels in one prompt are corresponding the same real value label without "self-conflict". By Lemma 38] we see that coefficients are at a substantial scale at 7". Then by the condition on z and Eq. (11.. we can readily check that even when the probability of the fraction of demonstrations sharing the co-concept label semantic with query is feeble (but at least one), utilizing the same set of notations, we still have
ENED-
IT' s
153
Sn
m.k
(TT 2
L' 126 k
2 @
9. 20 (
(( gg()))
(65)
L /2(ee
LEER[MUS
16.
+
(K-
9.6 K.
1bg
2
logí
1
2 (
AR]
mAXi
1 XK1
Log'(
Iqll" mAKT
+ (K 1)) !||?
1/2,
where the equality and inequality is by casse-case consideration over D small an and À in Condition
I with a
sufficiently large C, as well as the requirement | bi]l 2
| = @([uu])). Besides, by ||d]| 2 ||c| o ([q1))
Eq. 65 Lemma 4
Eg 5 and Lemma we have that
B Σ
(a (T") Ou...h + ys.. (2
Σ
(0 (T 1 1) o(., 2 (
M5 V5n n.k
15E
Collaborating with Lemma the poof is completed.
58

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide a short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked "[Yes] is generally preferable to "[No] ". is perfectly acceptable to answer "[No] provided a proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset used"). In general, answering "[No] oo "[NA] is not phrased in binary way. acknowledge that the answer write a justification elaborate.
the checklist as of the factors in their evaluation. While
grounds for rejection. While the questions often more nuanced, so please just use your best judgment evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
supporting
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
The contributions and scope of this paper well summarized in the abstract and
introduction.
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 be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can he expected to generalize other setrings.
It is fine to include aspirational goals motivation long it is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We explicitly discuss the limitation in Appendix A
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors encouraged create separate "Limitations" section in their paper.
59

---

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 sell-specification. asympiotic approximations only holding locally). The authors should reflect on how these assumptions might be violated practice and what the implications would be.
The authors should reflect the scope of the claims made, c.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-ter toxttttt system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If l'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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role in developing norms that preserve integrity the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The detailed assumptions and proofs for all theorems and lemmas are given in the corresponding positions.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should he numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by
formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Our algorithm is straightforward and easy to implement, and every detail is given in Section 3.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, 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 a 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 various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is 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 the model. In general. releasing code and data often one good way to 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 a large language model), releasing of model checkpoint, other means that appropriate to the research performed.
60

---

While NeurIPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how reproduce that algorithm.
(h) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results way to reproduce the model (e.g.. with open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may he tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may he that access the model is limited in some way (e.g., registered users), but 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 instructions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We have provided the complete configuration in Section B and We have uploaded the code with instructions the supplementary material.
Guidelines:
The answer means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines https::/miis.cccppbbicceeeee CodeSubmissionPolicy for more details.
While we encourage release of code and data, we understand that might not possible,
so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-sourge benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines https://nips.ccpuul////ee CodeSubmissionPoleyy for more details.
The authors should provide instructions on data access and preparation, including how 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. If only a subset experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to 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 is permitted.
Experimental Setting/Detail
Question: Does the paper specify all the training and test details (e.g.. data splits, hyperparameters, how they were chosen, type of optimizer. etc.) necessary understand the results?
Answer: [Yes]
Justification: All the experimental settings can be found Section and 6
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the of the paper level of detail that necessary to appreciate the results and make sense them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate informa-
tion about the statistical significance of the experiments?
Answer: [No]
This paper executes algorithms 10 times and reports the average results to reduce randomness.
Justification:
61

---

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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example. train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should given (c.g., Normally distributed errors).
should be clcar whether the error bar is the standard deviation the standard error of the mean.
report I-sigma error bars, one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CL, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures
symmetric error hars that would yield results that are out of range (e.g. negative error rates). If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the computer resources (type of compute workers, memory, time execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
We have provided sufficient information about computer resources in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for cach of the individual experimental well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g.. preliminary or failed experiments that didn't make 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
Answer: [Yes]
The research does not involve any human subjects, personal data, or interactions that would raise ethical concerns about consent, privacy, or respect for persons. In conclusion, research aligns with the ethical principles outlined in the NeurIPS Code of Ethics.
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure preserve anonymity (e.g., if there is special consideration due laws 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: [Yes]
Justification:
We have discussed the broader impacts in Section A
Guidelines:
The answer NA means that there is societal impact of the work performed.
62

---

If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.- disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploy- ment technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will he foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 generate deepfakes for disinformation. On the other hand, it not needed 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 is 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 to attacks, mechanisms for monitor- ing misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place responsible release of data or models that have a high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [No]
Justification:
This paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should he released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should deseribe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors 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: [No]
Justification: This paper does not use existing assets.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include a URL.
The name of license (e.g., c. CC-BY 4.0) should be included for each asset. For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, paperswithcode coticooaeccm//tttstttissss has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
63

---

Answer: [No]
Justification:
This paper does not release new assels.
Guidelines:
The answer means that paper not release new assets.
Researchers should communicate details the dataset/code/model as part of their sub- missions 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 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, well as details about compensation (if any)?
Answer: [No]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
answer means that paper does not involve crowdsourcing nor research with human subjects.
Including this information in supplemental material is fine, but main contribution of paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS 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.
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 the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on requirements your country or institution) were obtained?
Answer: [No]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer means that paper does not involve crowdsourcing nor research with human subjects.
Depending on the country which research is conducted, approval (or equivalent) may be required for any human subjects research. If you ohtained 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 adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such the institution conducting the review.
64


Reviews:
[{"review_id": "0lCihQIPhp", "reviewer": "Reviewer_qjJG", "summary": "The paper investigates how transformer-based large language models (LLMs) leverage their in-context learning (ICL) capabilities through the lens of multi-concept semantics. The authors cited limitations in existing theoretical work, which uses oversimplified models and unrealistic loss functions, leading to only linear or sub-linear convergence rates. To address this, the authors present an analysis of a three-layer transformer model, consisting of one attention layer and a ReLU-activated feedforward network, trained with logistic loss. The analysis included an examination of the learning dynamics of the transformer model and proved exponential convergence of the 0-1 loss in this complex setting, demonstrating that the transformer can achieve Bayes optimal test error with a logarithmic number of iterations.\nThe paper also demonstrated how multi-concept encoded semantic geometry enables transformers to perform efficiently on out-of-distribution ICL tasks, explaining their success in leveraging the polysemous nature of words for diverse, unseen tasks. The paper included empirical simulations to support these theoretical findings.", "questions": "NA", "limitations": "Yes", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The paper brings a new theoretical perspective linking multi-concept word semantics with in-context learning", "weaknesses": "The paper is sometimes hard to follow. More concise and clearer presentation, especially the theoretical content, would be helpful.\nLack of empirical experiments with more realistic datasets to show relevance in real-world setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for acknowledging our work for linking the multi-concept word semantics to transformer’s ICL capability.\n\n**Q: More concise and clearer presentation, especially the theoretical content, would be helpful.**\n\n**A**:  We very much appreciate your suggestion. We will provide more explanations and figures to improve the overall readability. The following are some explanations (which will be extended in the final version).\n\n- **Definitions and Technical Challenge**. As mentioned in the paper, we leverage a sparse coding approach suitable for capturing language polysemy [1-2] and incorporate observed latent geometric properties of LLMs [3-5], distinguishing our work from previous studies that assume idealistic, noise-free orthogonal settings [6]. Our learning problem considers practical non-linear transformer architectures and losses (softmax attention, ReLU MLP, and cross-entropy loss), contrasting with the oversimplified settings (linear/QK-combined/ReLU/MLP-free/infinite-dimensional transformers with square or hinge loss) in prior work [7-8].\n- **Theoretical Achievement**. Our theorem demonstrates an **exponential** convergence rate for our non-trivial setup, going beyond the linear or sublinear rates achievable in previous simplified setups [6-7] due to their technical limitation. Intuitively, our theory suggests the transformer's learned knowledge after training, with certain lengths and cross-concept orthogonality, enables its capability for certain OOD extrapolation, such that the test prompts can enjoy different lengths, various distributions of latent concepts, and even shifts in word semantics.\n- **Proof Sketch**. We prepare Figure 1 to visualize our Idempotent Operator Techniques in the new PDF file. This technique enables us to rigorously analyze the learning dynamics by conducting scale analyses on different orthogonal components of the data, whose practical meanings are partially validated in [9]. In addition, in Sections 5.2-3, we extend standard expectation-variance reduction techniques [10] to our setting. We treat the conditional expectations of the NN matrices as Doob martingales, and by constructing martingale difference sequences and exploiting the exponential convergence property of the tails, we derive the exponential convergence rate under low-noise conditions.\n\n**Q: Alignment to Real-world.**\n\n**A**:  Thank you for your feedback. We would like to address your concerns as follow.\n\n1. **Empirical Relevance**. Our setting of non-orthogonal dictionaries is inspired by practical observations [4-5, 9], where the LLM's latent geometry exhibits within-concept positive inner products and cross-concept orthogonality, which is also theoretically validated by [3]. Furthermore, according to [1-2], the sparse coding approach is a suitable setup for modelling the polysemy of language.\n2. **Data Modelling in Feature Learning Theory**. Mathematical theories often deal with simplified modelling, aiming to reveal certain **intrinsic capabilities** of models. There is a rich literature that adopts modelling **similar to ours**, such as studying self-supervised contrastive learning [2], or classification upon orthogonal dictionaries [6].\n3. **Alignment to the Real-World Setting**. In contrast to prior theories that considered oversimplified settings like noise-free orthogonal feature dictionaries [6] or linear/QK-combined/ReLU/MLP-free/infinite-dimensional transformers [7-8], our setup with non-orthogonal dictionaries, non-linear transformers, and cross-entropy loss is more aligned with practical scenarios. Our goal is to demonstrate the capability of transformers to utilize latent linear geometry for certain OOD tasks. The practical meanings of the OOD samples are partially validated in [9] and Theorem 8 in [11], where different settings show that the combination of concepts forms new meaningful semantics. We believe this provides an **initial positive response** to Question 5.1.4 in the ICML 2024 position paper [12], which asks whether the observed latent geometry of LLMs can explain their OOD extrapolation abilities.\n\nIn summary, grounded in practical observations, we study this non-trivial learning problem with **comparatively realistic settings**, and is an important step forward to reveal the ICL emergent capability of transformer [13].  Looking ahead, one of our important future directions include exploring more realistic model settings, such as incorporating additional attention layers, as well as considering the hierarchy of language in our theoretical analysis.\n\n**Reference**\n\n[1] Arora et. al. Linear algebraic structure of word senses, with applications to polysemy. TACL 2018\n\n[2] Wen and Li. Toward understanding the feature learning process of self-supervised contrastive learning. ICML 2021\n\n[3] Li et al. How do transformers learn topic structure: towards a mechanistic understanding. ICML 2023\n\n[4] Park et al. 2023: The linear representation hypothesis and the geometry of large language models. ICML 2024\n\n[5] Jiang et al. On the origins of linear representations in large language models. ICML 2024\n\n[6] Li et al. How do nonlinear transformers learn and generalize in In-Context Learning? ICML 2024\n\n[7] Chen et al. Training dynamics of multi-head softmax attention for In-Context Learning: emergence, convergence, and optimality. COLT 2024\n\n[8] Zhang et al. Trained transformers learn linear models In-Context. JMLR 2024\n\n[9] Yamagiwa et al. Discovering universal geometry in embeddings with ICA. EMNLP 2023\n\n[10] Nitanda and Suzuki. Stochastic gradient descent with exponential convergence rates of expected classification errors. AISTATS 2019\n\n[11] Park et al. The Geometry of Categorical and Hierarchical Concepts in Large Language Models. ICML Workshop MI 2024\n\n[12] Reizinger et al. Position: Understanding LLMs Requires More Than Statistical Generalization. ICML 2024\n\n[13] Lu et al. Are emergent abilities in large language models just In-Context Learning? ACL 2024", "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": "70IEdh8PlN", "reviewer": "Reviewer_dQGe", "summary": "This paper performs an in-depth analysis of the optimization dynamics of a simplified Transformer on a sparse coding prompt model. It manages to show the exponential 0-1 loss convergence on this non-convex loss. Experiments on synthetic data verify the convergence.", "questions": "1. It seems that the proof considering mini-batch SGD is showing a concentration to the population gradient descent by Lemma 3. As mentioned in the intro, one of the key properties of the theoretical claim here is that it doesn't require a large batch size. Why is the technique here so robust to batch size?", "limitations": "The limitation is adequately addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. This paper is very original and investigates the interesting question of how a Transformer learns shared concepts within context.\n2. This paper shows the clear OOD generalization of the Transformer model when trained on the sparse coding distribution.", "weaknesses": "1. The notation in the proof sketch is very heavy and not very easy to track. It would improve the paper if more explanation is given. Also, some of the statements seem counterintuitive. For example: \n\na) Definition 4 states that 'with prob 1 - $\\delta$, $L_{D*}(\\Phi) = 0$'. What is the source of randomness here? \n\nb) Proposition 4 states minimization of the loss to 0, which seems highly surprising at first pass. This seems like a specialty due to 0-1 test loss considered here and should be highlighted.\n\n2. Some of the legends in Figure 1 is left unexplained and in general the graph may need more explanations on why it matches Lemma 1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are immensely grateful for your thoughtful and insightful feedback on our work. Your recognition of the originality, soundness, and excellent contributions of our paper is deeply encouraging. We found your comments to be highly professional and constructive, and we would like to respond as follows.\n\n**Q: more explanation for the notation of the proof sketch.**\n\n**A**: Thanks for your suggestions. We would be happy to add more insights and illustrations (Figure 1 in the PDF file) to interpret the notations.\n\nThe key definitions are:\n\n- $a_k$ is the mean vector of $\\mu_k^+$ and $\\mu_k^-$, which is denoted as the \"concept vector\"; $b_k$ is denoted as the \"task-specific vector\", where $b_k=\\mu_k^+ - a_k$. $c_k$ and $d_k$ are defined similarly with regard to $q_k^+$ and $q_k^-$.\n- The coefficients $\\alpha_{Q, k}$ and $\\beta_{Q, k}$ are defined through the idempotent decomposition of $W_Q$. This allows us to show that ${(W_K\\mu_k^{\\pm e})}^{T}W_Q\\mu_k^e = \\alpha_{Q,k}\\cdot\\alpha_{K,k}\\pm\\beta_{Q,k}\\cdot\\beta_{K,k}$ by simple calculations, which pave the way for us to examine the evolution of attention weights by tracking the dynamics of the coefficients. $\\alpha_O$ is defined similarly. Our theory then builds up the relationship between the test error of the expected model and the expected coefficients. This would facilitate us to establish the 0-1 loss exponential convergence by closely examining the evolution of these coefficients.\n\n**Q: source of randomness.**\n\n**A**: The primary source of randomness here stems from two factors: the uncertainty in the model's initialization, and the tail behavior of the Gaussian noise, as detailed in Appendix B.1. Our overparameterization conditions ensure that we can control, with high probability 1-δ, the volumes and directional lengths of the matrix parameters, as well as the influence of the noise. Without careful control of these preliminary properties, our theoretical results would not be able to hold rigorously. We appreciate you noting this important detail.\n\n**Q: test error considered in Proposition 4 should be highlighted.**\n\n**A**: We appreciate your intuition about our results. You are correct that, similar to [1], the consideration of the 0-1 loss and treating cross-entropy loss as a surrogate is crucial to establishing the convergence in our Proposition 4. We will certainly highlight this aspect in our revised manuscript.\n\nAnother key point is that our proposition examines the expected model behavior at each time step, where the randomness here stems from the stochastic batch sampling in SGD. Given the isotropic noise and balanced positive/negative labels in our setting, the expected model can be viewed as being trained by noise-free gradient descent on a balanced dataset. From this perspective, the expected model's test error can be shown to converge to zero very rapidly, as long as the lengths of the neural network matrices grow sufficiently along the critical directions.\n\nWe appreciate your reminder and would highlight the points with interpretations in our revised manuscript.\n\n**Q: more explanation for Figure 1.**\n\n**A**: Thank you for the helpful suggestion. We will replace the existing Figure 1 with a new one (Figure 2 in the supplementary pdf), and ensure the caption provides detailed descriptions.\n\nSpecifically, the new Figure 2 includes:\n\n1. A plot demonstrating the exponential convergence of the test error, validating our theory.\n2. A plot showing that the correct attention weights for both concepts' classification tasks converge to 1, aligning with the last conclusion of Lemma 1.\n3. A plot verifying our Lemma 1, as it shows the products $\\alpha_{Q, s}^{(t)}\\cdot\\alpha_{K, s}^{(t)}$ are all non-increasing, while the values of $\\beta_{Q, s}^{(t)}\\cdot\\beta_{K, s}^{(t)}$ can grow sufficiently.\n4. A plot supporting the claims in Lemma 1 regarding the sufficient growth of the$\\lvert\\beta_{O_{(i, \\cdot)},k}^{(t)}\\rvert$ and $\\alpha_{O_{(i, \\cdot)}, k}^{(t)}$, as well as their relationships.\n\nAs our analysis demonstrates the convergence between the expected model and the SGD-updated model, the new Figure 2 clearly illustrates how the empirical results align with and validate the theoretical claims our theories. Additionally, we include supplementary Figure 3 of OOD scenarios to further collaborate our theories. In the final manuscript, we will provide extended experiments and descriptions in the appendix.\n\n**Q: robustness to batch size.**\n\n**A**: The key reason that our technique is robust to batch size, compared to prior work [2], is that we introduce standard techniques from the theoretical literature on exponential convergence [1] to our settings. Specifically, rather than directly using the batch size (at least $\\varepsilon^{-2}$, where $\\varepsilon$ is the test error) to bound the gradient difference between the expected model and the SGD-updated model via Hoeffding's inequality, as done in [2], our analysis first considers the test error convergence of the expected model. We then examine the rate of convergence for the test error difference between the expected model and the SGD-updated model at each iteration, which could be exponential due to the hard low-noise condition [1]. Crucially, our analysis considers the extreme case where the batch size can be as small as 1, as in [1], to provide an upper bound on the test error. This allows our technique to be robust to the batch size, in contrast with the larger batch size requirements in prior work.\n\n**Reference**\n\n[1] Nitanda and Suzuki. Stochastic gradient descent with exponential convergence rates of expected classification errors. AISTATS 2019\n\n[2] Li et al. How do nonlinear transformers learn and generalize in In-Context Learning? ICML 2024", "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": "6cZTCv4b9j", "reviewer": "Reviewer_L2Bx", "summary": "This paper tries to understand the mechanisms that explain the emergence of in-context learning.The starting point of their approach is the observation that the embeddings have geometric properties to encode within-concept and cross-concept relationships. Their goal is to connect this geometric properties with the ability to conduct efficient in-context learning. They first describe a data distribution that is sparse coding prompt distribution and then describe their model which is a 3-layer transformer. They show that such a model trained with SGD converge in an exponential speed to the 0-1 Bayes optimal test error. Lastly, they show that their learned model can utilize the multi-concept semantics to conduct unseen ICL tasks.", "questions": "I think the authors should improve the presentation of their theory by giving more insights and intuition + simplified proof sketch.", "limitations": "The authors did not mention the limitations of their work.", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I think that the results are pretty strong and I liked the proposition 1 which shows how the transformer can use its multi-concept semantics knowledge to perform in-context learning on unseen tasks.", "weaknesses": "I do not have any knowledge on the topic and so this is why my evaluation is going to be very shallow. Even though the theorems look sound to me, I think the authors could improve the presentation of their results by giving some proof sketch (+drawing) in the introduction.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback and recognition of the strength of our theoretical results, particularly Proposition 1 which demonstrates how transformers can leverage their multi-concept semantic knowledge to perform effective unseen ICL tasks. We appreciate your insightful assessment of the significance of our work.\n\n**Q: knowledge on the topic & more insights and intuitions for presentation.**\n\n**A**: We sincerely appreciate your feedback and consider it an honor to share more background on our topic with you. We also value your advice and aim to provide additional insights and intuitions as follows:\n\n- **Understanding LLM**. In recent years, the remarkable power of LLMs has prompted both practitioners and theoreticians to explore their underlying mechanisms. Practitioners validate their claims through large-scale experiments [1-2], while theoreticians use simplified models for analysis [3-6]. A natural practice is that the theoreticians employ mathematics to explain the phenomena observed by practitioners.\n- **Research Gap and Our Contributions**. On the one hand, empirical findings suggest that the emergent power of LLMs mainly stems from their ICL ability [1], which is partially attributed to data properties and transformer structures [2]. On the other hand, existing theories are often based on unrealistic settings (e.g., noise-free orthogonal data, simplified transformer architectures like linear/QK-combined/ReLU/MLP-free/infinite-dimensional models, and square or hinge loss) [3-6]. This leaves room for providing explanations for emergent ICL capabilities in more realistic settings. To advance this understanding, we theoretically demonstrate that transformers leverage latent linear geometry for certain OOD tasks. We model the data based on empirical observations of LLM latent structures [7-9] and consider realistic non-linear transformer architectures and losses (softmax attention, ReLU MLP, and cross-entropy (logistic) loss). Using advanced analytical techniques, we showcase an **exponential** convergence rate for our non-trivial setup, going beyond the linear or sublinear rates achievable in previous simplified setups [3-5] due to their technical limitation.\n\nInterestingly, our work addresses a research question raised in Question 5.1.4 of the ICML 2024 position paper [10], which inquires whether the observed linear latent geometry of LLMs can explain their OOD extrapolation capabilities. We believe our results provide an **initial positive response** to this question.\n\n**Q: provide (simplified) proof sketch(+ drawing) in introduction.**\n\n**A**: Thank you for the valuable feedback. We appreciate your suggestion to enhance the accessibility of our proof sketch with visual aids. In response, we have prepared Figure 1 for inclusion in the introduction, which would be collaborated with a simplified proof sketch in our revised manuscript.\n\n**Simplified Poof Sketch**. Please refer to our illustration (Figure 1) in the new PDF. The Idempotent Operator Techniques depicted allow us to rigorously analyze the model's learning dynamics by examining matrix lengths across orthogonal components, whose practical meanings are partially validated in [11]. In Sections 5.2-3, we extend expectation-variance reduction techniques [12-13]. By treating conditional expectations as Doob martingales, we exploit exponential convergence properties, deriving the rate under low-noise conditions. Our theory suggests that the transformer's learned knowledge, characterized by specific lengths and cross-concept orthogonality, is essential for OOD extrapolation, particularly in prioritizing words that share components with queries.\n\n**Q: The authors did not mention the limitations of their work.**\n\n**A**: We would like to clarify that we have a Limitations Section in Appendix J. There is potential to enhance our modellings to better align with real-world scenarios, such as increasing attention layers or considering the Markovianity of the prompt distribution. We will highlight these limitations in the main body and move Appendix J to Appendix A in our revisions.\n\n**Summary**\n\nWe are sincerely grateful for the reviewer’s constructive feedback. Should the reviewer have any further suggestions or wish to discuss any points in more detail, we would be delighted to continue this productive exchange. Once again, we deeply appreciate the reviewer’s time and valuable comments.\n\n**Reference**\n\n[1] Lu et al. Are emergent abilities in large language models just In-Context Learning? ACL 2024\n\n[2] Chan et al. Data Distributional Properties Drive Emergent In-Context Learning in Transformers. NeurIPS 2022\n\n[3] Zhang et. al. Trained transformers learn linear models In-Context. JMLR 2024\n\n[4] Kim and Suzuki. Transformers learn nonlinear features In-Context: nonconvex mean-field dynamics on the attention landscape. ICML 2024\n\n[5] Li et al. How do nonlinear transformers learn and generalize in In-Context Learning? ICML 2024\n\n[6] Chen et al. Training dynamics of multi-head softmax attention for In-Context Learning: emergence, convergence, and optimality. COLT 2024\n\n[7] Park et al. 2023: The linear representation hypothesis and the geometry of large language models. ICML 2024\n\n[8] Jiang et. al. On the origins of linear representations in large language models. ICML 2024\n\n[9] Li et. al. How do transformers learn topic structure: towards a mechanistic understanding. ICML 2023\n\n[10] Reizinger et al. Position: Understanding LLMs Requires More Than Statistical Generalization. ICML 2024\n\n[11] Yamagiwa et al. Discovering universal geometry in embeddings with ICA. EMNLP 2023\n\n[12] Nitanda and Suzuki. Stochastic gradient descent with exponential convergence rates of expected classification errors. AISTATS 2019\n\n[13] Yashima et. al. Exponential convergence rates of classification errors on learning with SGD and random feature. AISTATS 2021", "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": "lEWOO4qBJh", "reviewer": "Reviewer_Bq3t", "summary": "This paper provides a theoretical analysis of in-context learning (ICL) in transformer models. The authors prove exponential convergence of the 0-1 loss for a three-layer transformer (attention + MLP) trained on a concept-specific prompt distribution. They also demonstrate how the model can leverage multi-concept semantics for out-of-distribution generalization. The analysis connects the geometric properties of concept-encoded representations to ICL capabilities.", "questions": "- How sensitive are the results to the specific data distribution assumptions? Would the exponential convergence hold under more general conditions?\n- Have you considered extending the analysis to training all weights in the network? What challenges would this introduce?", "limitations": "In Appendix J, the authors acknowledge that the data model may need refinement to better align with practical scenarios. They also note that adding more attention layers could make the model more realistic. These are fair limitations to point out.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- First work to show exponential convergence of 0-1 loss for a realistic transformer model: The authors prove exponential convergence for a three-layer model with **softmax** attention and **ReLU** MLP, which is more practical than previous simplified models. This is a significant theoretical advancement.\n- Solid mathematical analysis: The proof techniques seem rigorous and well-developed. The authors leverage advanced techniques to handle challenges like softmax attention and logistic loss, which have been difficult to analyze theoretically.\n- Experiments support theoretical findings: The simulations in Section 6 validate key aspects of the theory, like the evolution of attention weights and MLP coefficients. \n- Good motivation and connections to empirical observations: The authors motivate their work well by connecting to empirical findings on LLM representations (e.g., within-concept positive inner products). This grounds the theory in practical observations.", "weaknesses": "- Limited training setup: The paper only trains $W_Q, W_K, W_O$, keeping other weights fixed. This is quite restrictive compared to full transformer training and may limit the applicability of the results.\n- Overly specific data model: The concept-specific prompt distribution seems quite structured and simple. It's not clear how well this captures the complexity of real language data. A linear hyperplane can solve all tasks, including OOD ones, in this setup.\n- Lack of non-linear task consideration: The analysis focuses on essentially linear classification tasks. It would be more compelling to see how the model handles tasks requiring non-linear decision boundaries or composition of multiple concepts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful review. We greatly appreciate your acknowledgment of our theoretical advancements, rigorous analysis, well-grounded connections to empirical observations, and fair assessment of the limitations we address.\n\n**Q: Limited training setup & challenges of training all weights.**\n\n**A**: Thank you for raising this point. The reason we left some matrices fixed (e.g., zero entries, $W_V$ and $a$) was to simplify the coupled gradient dynamics and enable rigorous theoretical analysis, which is a common practice in deep learning theory [1-2]. We remark that this restriction does not fundamentally limit the functionality of the transformer with sufficient sample complexity in our learning problem (e.g. one $W_V$ here already can take the full job of $W_O W_V$), which is validated by experiments.\n\nWe appreciate your encouragement to consider training all weights. Indeed, it’s known that adding more trainable layers could **lower the sample complexity**, especially for harder learning targets that exhibit hierarchy [3-4]. Extending the analysis to training all weights would certainly pose additional challenges since there’re many gradient updates of various matrices simultaneously. This would require simplification techniques such as layer-wise training and reparameterization [3], as well as special structural assumptions on the concept classes (e.g., information gap) and algorithm configurations (e.g., different regularizers at different phases) [4].\n\n**Q: Alignment of the data model to real-world & applicability to general conditions.**\n\n**A**: Thanks for raising the points. We would like to address your concerns as follows.\n\n- **Empirical Relevance**. Mathematical theories often simplify models to reveal intrinsic capabilities. To better align with the real world, we leverage a sparse coding approach, proposed as a suitable theoretical setup for capturing language polysemy [5], and in our case, seen as a special prompt version of the recognized LDA [6]. Our purpose is to understand how transformers leverage latent linear geometry for certain OOD tasks, modeled after empirical observations on LLM latent structure [7-9]. The practical meanings of the OOD samples are partially validated in [9-10], where different settings show that the combination of concepts forms new meaningful semantics. While there is room for improvement, such as considering hierarchy, we believe our theory stands out and provides an initial positive response to the research question in [11], which asks whether the observed LLM latent geometry can explain their OOD extrapolation abilities.\n- **Applicability**. The exponential convergence result hinges on the hard low-noise condition [12]. Intuitively it suggests that when the target pattern is sufficiently “low-noise”, the exponential convergence with a suitable learning rate can be achieved. We note that strict assumptions are common in theoretical studies to enable rigorous analysis. While our current theory holds under Condition 1, we believe there is scope to relax these assumptions and develop theories for more general data distributions. Looking ahead, a key future direction is to explore analytical techniques that can handle more general conditions.\n\n**Q: Lack of non-linear task & composition of multiple concepts**\n\n**A**: Thank you for this insightful feedback. We would like to address your concerns as follows:\n\n- **Linearity of LLM Representations**. Recent work has shown that high-level concepts in language models tend to be encoded linearly [7-8]. Our theory aims to connect this observed linear structure to the transformer's capability on certain OOD tasks.\n- **OOD Tasks with Multiple Concepts**. We note that our proposed OOD tasks allow for the composition of multiple co-concepts, as stated in the second point of Proposition 1.\n- **Handling Non-linear Tasks**. Given that our transformer model includes non-linear MLPs, which can theoretically handle non-linear tasks [13], it is feasible to consider non-linear tasks. In this context, we expect that the attention mechanisms would still assign more weight to words in demonstrations that share the most similar \"non-linear patterns\" with the query to complete ICL tasks.\n\nWe believe our focus on empirically grounded linear modeling has allowed us to provide the first theories linking multi-concept geometry to the OOD capabilities of transformers, marking an important step forward. However, we agree that exploring whether and how transformers can excel in certain non-linear tasks is another crucial direction for our future endeavors. We appreciate you highlighting this insightful point.\n\n**Reference**\n\n[1] Jelassi et al. Vision Transformers provably learn spatial structure. NeurIPS 2022\n\n[2] Huang et al. In-Context Convergence of Transformers. ICML 2024\n\n[3] Tian et al. JoMA: Demystifying Multilayer Transformers via Joint Dynamics of MLP and Attention. ICLR 2024\n\n[4] Allen-zhu and Li. Backward feature correction: How deep learning performs deep (hierarchical) learning. COLT 2023\n\n[5] Arora et al. Linear algebraic structure of word senses, with applications to polysemy. TACL 2018\n\n[6] Blei et al. Latent Dirichlet Allocation. NIPS 2001\n\n[7] Park et al. 2023: The linear representation hypothesis and the geometry of large language models. ICML 2024\n\n[8] Jiang et al. On the origins of linear representations in large language models. ICML 2024\n\n[9] Yamagiwa et al. Discovering universal geometry in embeddings with ICA. EMNLP 2023\n\n[10] Park et al. The Geometry of Categorical and Hierarchical Concepts in Large Language Models. ICML Workshop MI 2024\n\n[11] Reizinger et al. Position: Understanding LLMs Requires More Than Statistical Generalization. ICML 2024\n\n[12] Massart and Nedelec. Risk Bounds for Statistical Learning. AISTATS 2006\n\n[13] Kou et al. Matching the Statistical Query Lower Bound for k-sparse Parity Problems with Stochastic Gradient Descent. arXiv:2404.12376", "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:
Dear ACs and Reviewers,

Thank you again for all of your positive and constructive feedback! We are truly encouraged to see so many well-recognized comments on our work, such as the **innovative angle** (Reviewer qjJG, Reviewer dQGe), **advanced techniques** (Reviewer dQGe, Reviewer Bq3t), **significant theoretical achievement** (Reviewer L2Bx, Reviewer Bq3t), **practical modelings** (Reviewer Bq3t), **sound analysis** (Reviewer L2Bx, Reviewer Bq3t), and **empirically-grounded motivation** (Reviewer Bq3t).

Your insights have greatly helped us strengthen our manuscript. Alongside addressing your comments point-by-point, we have attached a new PDF featuring illustrations of our proving technique (Figure 1) and additional experiments with detailed descriptions (Figures 2-3). We particularly thank Reviewer dQGe for the reminder to provide more interpretation for the notations and Reviewer L2Bx for suggesting visual aids.

Your feedback assures us that we have successfully achieved our main goal: understanding how transformers leverage latent linear geometry for certain out-of-distribution tasks, modeled after the empirical observations on LLM latent structure [1-3]. Interestingly, the research gap we’re addressing aligns with a research question raised in the ICML 2024 position paper [4], regarding whether the observed linear latent geometry of LLMs can be leveraged to explain their OOD extrapolation abilities. By emphasizing your points in our revisions, we believe we can excel as an important step forward and continue contributing to the community in future directions inspired by your valuable comments.

Should the reviewer have any further suggestions or wish to discuss any points in more detail, we would be more than delighted to continue our productive exchange. Once again, we deeply appreciate the reviewer’s time and valuable comments. 

Warmest regards,

Authors of Submission 15661

**Reference**

[1] Yamagiwa et al. Discovering universal geometry in embeddings with ICA. EMNLP 2023

[2] Park et al. 2023: The linear representation hypothesis and the geometry of large language models. ICML 2024

[3] Jiang et. al. On the origins of linear representations in large language models. ICML 2024

[4] Reizinger et al. Position: Understanding LLMs Requires More Than Statistical Generalization. ICML 2024

Meta-review: This paper adds to the growing body of working trying to formalize in-context learning of large language models. Despite lower-than-average confidence scores, the reviewers unanimously praised the paper for its originality, theoretical treatment of ICL, and its empirical results. The main concerns pertain to the clarity and presentation of the paper and to some of the assumptions required by its theoretical formulation.

Overall, I agree with the reviewers that this is paper provides an interesting and compelling formalism of LLMs through the lens of mutli-concept word semantics, that leads to a novel and sound interpretation of ICL. Despite its limited evaluation and simplified setting, I believe its strengths outweigh its weaknesses. However, the authors should definitely take into account the reviewer's comments regarding clarity and presentation, since the paper in its current form is unnecessarily opaque.
Decision: Accept (poster)