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_cc_2024_Conference
Published: 2024-05-15 10:40:59

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 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 semantic regularity and the innovative power of ICL. Additionally, prior work often focuses on simplified, unrealistic scenarios with linear transformers or unrealistic loss functions, and due to their technical limitations, they come up with results exhibiting 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 new unseen tasks encoded with multiple cross-concept semantics. Inspired by empirical studies on the 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.

Main paper: 
 introduction: 
 Recently, a variety of transformer-based large language models (LLMs) have demonstrated remarkable performance across a broad spectrum of machine learning tasks, including natural language understanding [1] , symbolic reasoning [2] , and even heuristics design [3, 4] . One crucial emerging ability of these models is their in-context learning (ICL) capacity [5] , which allows them to learn from a few demonstrations and conduct predictions on new queries without requiring any further 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 LLMs in unseen ICL tasks largely unexplained.In line with traditional topic models [6] , [7, 8] propose that latent concepts / topics underlie natural texts, providing a 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 a result of concept or topic learning [9, 10] , where the representations within-concept have positive inner products while representations cross-concepts exhibit near-orthogonal relationships. This structured semantic geometry has been well-documented in recent research on pre-trained LLMs [11, 12, 10, 13] . However, the connection between this observed multi-concepts latent geometric structure and the LMs' remarkable ICL capabilities remains unclear. Separately, recent theoretical analyses have modeled ICL 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.Additionally, existing theoretical work has been conducted on unrealistic, oversimplified settings, such as linear or ReLU transformers [16, 17, 18, 19] , MLP-free attention-only models [16, 20] , QK-combined softmax attention [21, 20, 19, 22, 23] , unrealistic infinite dimensional assumption [14, 24, 21, 19] and impractical loss functions like square loss [16, 9, 25, 20, 26] and hinge loss [27, 28] . Furthermore, existing works have only been able to derive linear or sub-linear convergence rates for the 0-1 loss.Therefore, there is a need for a more advanced analysis that can bridge the understanding between the multi-concept semantic regularity and the mechanisms underlying transformer-based ICL. This naturally leads to the research question: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 -within-concept positive inner products and cross-concept orthogonal geometric properties -that aligns with the findings in [9, 10, 11] . Our main contributions are highlighted as below.1. First, we provide a comprehensive analysis of the learning dynamics for a two-layer transformer model, comprising one attention layer followed by a 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 techniques, 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 a logarithmic number of iterations.3. We provably show how the multi-concept encoded semantic geometry can enable transformer to efficiently perform certain 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 LLMs can explain their OOD extrapolation abilities. 
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" [31, 32] . This phenomenon has been further explored in more recent studies analyzing the exponential convergence of stochastic gradient descent (SGD) [33, 34, 35, 36, 37] , as well as in 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, examining NN's feature direction reconstruction and noise memorization as a 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 [43, 44] . Our work extends this line-of-research to challenging nonlinear Attention-MLP transformers with non-orthogonal structured data representations.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 conceptencoded representations to transformers' powerful ICL abilities. Another line of research has studied the learning dynamics of ICL, including analyses of linear transformers [17, 19] , 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 it 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 ε -2 , where ε is the test error) and long context lengths, which are not required in our result; and (iii) they utilize an impractical hinge loss and only achieve linear convergence without a relation to ε, whereas we analyze the more practical cross-entropy loss and derive an exponential convergence rate in terms of the test error ε. However, we note that this is only an informal comparison due to the differences in the models and primary findings. A detailed Related Work Section is deferred to Appendix C. In addition, we denote span(v 1 , v 2 , . . . , v k ) as the linear subspace spanned by the vectors v 1 , v 2 , . . . , v k , and conic(v 1 , v 2 , . . . , v k ) denotes the conic hull (the set of all non-negative linear combinations) of the vectors v 1 , v 2 , . . . , v k . 
data distribution: 
 The data distribution employed in this study draws inspiration from a range of empirical and theoretical research works [46, 47, 9, 10, 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 a 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 representation. Definition 1. Polysemous Word Model (D x , D y , D z , D ξx , D ξy ). We assume there exists K 1 taskrelevant concepts, each characterized by two semantically-opposite word's feature vectors µ + k1 and µ - k1 , and their corresponding label's feature vectors q + k1 and q - k1 , ∀k 1 ∈ [K 1 ]. There are also K 2 task-irrelevant concepts denoted by ν k2 , ∀k 2 ∈ [K 2 ]. The word samples x ∈ R d X and their labels y ∈ R d Y are generated from distributions parameterized by a shared latent concept variable z = (z 1 , • • • , z K ) ∈ {0, 1} K (K < d X ) capturing the concept-specific information:z ∼ D z , ξ x ∼ D ξx = N (0, σ 2 ξ I d X ), ξ y ∼ D ξy = N (0, σ 2 ξ I d Y ), x = Mz + ξ x ∼ D x , y = Qz + ξ y ∼ D y ,where the feature dictionary M =[µ + 1 , µ - 1 , µ + 2 , µ - 2 , • • • , µ + K1 , µ - K1 , ν 1 , ν 2 , • • • , ν K2 ] ∈ R d X ×Kexhibits positive inner products within concepts and orthogonality across concepts, and the label dictionary Q = [q + 1 , q - 1 , q + 2 , q - 2 , • • • , q + K1 , q - K1 , 0, • • • 0] ∈ R d Y ×K has similar geometric properties. Specifically, we have ∀k 1 ∈ [K 1 ], k 2 ∈ [K 2 ], ∥µ ± k1 ∥ = ∥ν k2 ∥ = ∥u∥, ∥q ± k1 ∥ = ∥q∥, and there exist constants 0 < κ x , κ y < 1 such that 0 < ⟨µ + k1 , µ - k1 ⟩ ≤ κ x ∥u∥2 and 0 < ⟨q + k1 , q - k1 ⟩ ≤ κ y ∥q∥ 2 . The detailed formal definition can be found in Appendix E. By this definition, a single word or label can possess different features corresponds to different concepts. The illustration of Figure 1 in [12] can be an example, where the "Dog" vector in the representation space of LLM is decomposed to a 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 a 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 z among in-context words, which is a special prompt version of PLSA [49] and LDA [6] . The detailed formal version is available in Appendix E. Definition 2. Concept-specific Contextual Prompt Distribution 2 . During training, each prompt sample S = x 1 , y 1 , • • • , x L , y L , x L+1 would share at least one co-concept, which is drawn from a mixture distribution D S defined as:EQUATIONwhere P ± k,L+1 denotes the k-th concept-specific prompt distribution, and π ± k = (2K 1 ) -1 denotes the equal chance of a sample to belong to P ± k,L+1 . Specifically, a sample S n ∼ P e k,L+1 , e ∈ [±] means that the query's label y n L+1 is q e k , and we denote y Sn := e as the real value label of this prompt. In addition, every demonstration pairs (x n l , y n l ), l ∈ [L] in P e k,L+1 contain either (µ + k , q + k ) or (µ - k , q - k ) with equal chance. Also, every z n l , l ∈ [L + 1] would satisfy P(z n l,¬(2k-1∨2k) = 1) = K -1 , denoting the equal chance to have diverse features other than the current co-concept of the P e k,L+1 .This definition suggests that for prompt S sampling from D S , there exists e ∈[±], k ∈ [K 1 ],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 to be +1 or -1.Following [17, 20, 28] , our embedding E(•) of prompt S is formulated as H:H = E(S) = x 1 x 2 • • • x L x query y 1 y 2 • • • y L 0 := (h 1 , h 2 , • • • , h query ) ∈ R (d X +d Y )×(L+1) ,The learning model is a single-head, one-layer Transformer with one self-attention layer and one two-layer perceptron. Mathematically, it can be expressed as follows:f (H; Ψ) = r ⊤ σ R (W O attn(H; Ψ)) , attn(H; Ψ) = L l=1 W V h l σ S (W K h l ) ⊤ W Q h query , where σ R (•) := Relu(•), σ S (•) := softmax(•), W Q , W K ∈ R m qk ×(d X +d Y ) , W V ∈ R mv×(d X +d Y )are the embedding matrices for queries, keys, and values, respectively, and W O ∈ R m×mv and r ∈ R m are parameters in the MLP layer. Typically, min (m qk , m v ) ≥ d X + d Y . Ψ := {W Q , W K , W V , W O , r} 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 letW Q = W x Q * * * , W K = W x KWe utilize mini-batch with-replacement SGD to train the transformer model. The empirical crossentropy loss for each batch B t is written asL Bt (Ψ) = L Bt (Ψ ′ ) := 1 B n∈Bt ℓ (y Sn • f (H; Ψ)) + λ 2 ∥Ψ ′ ∥ 2 F ,where ℓ(z) = log(1 + exp(-z)), y Sn is the real value label of the prompt defined in Definition 2, and the term∥Ψ ′ ∥ 2 F represents ∥W x Q ∥ 2 F + ∥W x K ∥ 2 F + ∥W y O ∥ 2 F, which is the L 2 regularization term with ∥ • ∥ F denoted as the Frobenius norm. The purpose of the regularization in this paper is to accelerate and stabilize the mini-batch with-replacement SGD. The learning step is set to be η t = 2 λ(γ+t) , where γ is an offset parameter. This decaying schedule is standard and also used in prior work [34, 50, 51] studying convergence of SGD. The whole procedure is in Algorithm 1.O are sampled from a i.i.d. Gaussian distributions with mean 0 and variance σ 2 1 . The initialization of W x Q and W x K are diagonal matrices σ 0 I, which are also adopted in other work that consider training W Q and W K separately [28, 25] .Testing Setting. The model performance is measured by 0-1 test error on a test prompt distribution D * : EQUATION′ (0) . for t = 0, 1, . . . , T -1 do If L 0-1 D * (Ψ (t) ) ≤ ε stop else continue. Randomly sample mini batches B t of size B from D S . Update model parameters: Ψ ′ (t+1) = Ψ ′ (t) -η t ∇ Ψ ′ L Bt (Ψ ′ (t) ). end fort ] 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 Proof. Observe that E[ Ec k,e = c 1 η 0 ∥q∥ 2 2mK 1 , c k,e = c 2 η T * (2σ * S -1) 2 (1-κy )∥q∥ 2 16mK 1 , b k,e = e -EQUATIONEQUATION-EQUATION)EQUATIONEQUATION) As cross-entropy loss is L-smooth with L = 1, one can bound the difference byEQUATIONBy Lemma 23, we see that for initialization, we have| E n∈V e k (ef (E(S); E(Ψ (0) )) -E n∈V -e k (-ef (E(S); E(Ψ (0) )))| ≤ 2 2 log( 5Km δ ) • 3σ1(∥c k ∥ + ζ e k ∥d k ∥) 8≤ κ/8. Now we serve to show that the following expected difference| E n∈V e k (ef (E(S); E(Ψ (t) )) -E n∈V -e k (-ef (E(S); E(Ψ (t) )))|is non-increasing. Intuitively, this observation is due to the inherent nature of cross-entropy loss, which always pays more emphasis (has larger derivative) on those low value. Also, another important factor is the update of those ambiguous neurons' coefficient summation would also prefer the low-value one among E EQUATIONO (i,•) ,k | Ψ (t) ] = (1 -ηtλ)eβ (t) O (i,•) ,k -ηt ∥d k ∥ 2 2K1 e∈[±] rie E n∈V e k [ℓ ′ n (t) 1 n O (i) (t) ( l∈S e n,k (σ (t) S ) n l - l∈S -e n,k (σ (t) S ) n l )].Then we haveEQUATIONEQUATIONS ) n l -1) -ηt (K1 -1) 2K1K e∈[±] rie E n∈V e ¬k [ℓ ′ n (t) 1 n O (i) (t) ∥c k ∥ 2 ].By Lemma 28 and Lemma 29, we see that theEQUATION) where we ignore the impact of cross-concept safely due to the large K = Ω(η0C(K1 -1)∥q∥ 2 /(mK1)), as well as the impact of regularization term since λ = O((C log(Km/δ)∥q∥) -1 ) by Condition 1 in the first stage.Similarly, suggest E[ l∈S e n,k (σ (t+1) S ) n l] is also given when considering the update for t + 1, we see thatEQUATIONe(2EQUATIONEQUATIONS ) n l -1)ℓ ′ n (t) )).(49) Interestingly, by Eq.( 39) we see that (2 l∈S e n,k (σ(t) S ) n l -1) = (2 l∈S -e n,k (σ (t) S )n l -1). Thus we can characterize that the magnitude of gradient update of the term in Eq.( 48) and (49) of the e * t would be larger than those of -e * t due to the non-increasing nature of cross-entropy loss. On the other hand, by Lemma 29 the monotonicity ofEQUATIONO (i,•) ,k ] τ =0 τ =t depend on the signal of E[ E n∈V e k (ℓ ′ n (t) ) -E n∈V -e k (ℓ ′ n (t) )]. Specifically, we see thatEQUATIONEQUATIONO (i,•) ,k | Ψ (t) ] τ =0 τ =t+1 = Θ( i∈(W e k,n (τ )-U e k,n (t))∩U -e k,n (τ ) α (t) O (i,•) ,k τ =0 τ =t + i∈(W e k,n (τ )-U e k,n (t))∩U -e k,n (τ ) ηt∥c k ∥ 2 2mK1 [ E n∈V e k (ℓ ′ n (t) ) -E n∈V -e k (ℓ ′ n (t) )]);(50) where the contribution term is shared by the two sequences. Therefore, by Eq.( 50) and ( 46), the evolution ofEQUATIONEQUATIONO (i,•) ,k ] τ =0 τ =twill prefer to grow in the direction of e * t .We then take a look on the decreasing coefficients based on Lemma 29.EQUATIONEQUATIONe(2EQUATIONEQUATIONEQUATIONAs such, we have all preliminaries to characterize the first result of the lemma. We first utilize the induction to prove the following:| E n∈V e k (ef (E(S); E(Ψ (t) )) -E n∈V -e k (-ef (E(S); E(Ψ (t) )))| ≤ κ/8. ∀e ∈ [±].This apparently hold at initialization. Suggest for any t ≤ t -1 the result holds, then we only need to proveE n∈V e * t-1 k (e * t-1 f (E(S); E(Ψ ( t-1) )) - E n∈V -e * t-1 k (-e * t-1 f (E(S); E(Ψ ( t-1) ))) ≥ E n∈V e * t-1 k (e * t-1 f (E(S); E(Ψ (t 1 ) )) - E n∈V -e * t-1 k (-e * t-1 f (E(S); E(Ψ (t 1 ) ))).By the condition of small ηt in Condition 1, Lemma 31, Eq.(48) (49), ( 50) and ( 51), we see thatE n∈V -e * t-1 k (-e * t-1 f (E(S); E(Ψ ( t-1) ))) - E n∈V e * t-1 k (e * t-1 f (E(S); E(Ψ ( t-1) )) -( E n∈V -e * t-1 k (-e * t-1 f (E(S); E(Ψ ( t) ))) - E n∈V e * t-1 k (e * t-1 f (E(S); E(Ψ ( t) ))) ≤ Θ(E[ E n∈V e * t-1 k ((ℓ ′ n ( t-1) ) - E n∈V -e * t-1 k (ℓ ′ n ( t-1) )) i∈W e * t-1 k,n ( t-1) ηT * ∥c k ∥ 2 2mK1 + (2 l∈S e * t-1 n,k (σ ( t) S ) n l -1) i∈W e * t-1 k,n ( t-1) ηT * ∥d k ∥ 2 2mK1 E n∈V e * t-1 k ((2 l∈S e * t-1 n,k (σ ( t-1) S ) n l -1)) ]) ≤ 0,and thus we have| E n∈V e * t-1 k (e * t-1 f (E(S); E(Ψ ( t) )) - E n∈V -e * t-1 k (-e * t-1 f (E(S); E(Ψ ( t) )))| ≤ | E n∈V e * t-1 k (e * t-1 f (E(S); E(Ψ ( t-1) )) - E n∈V -e * t-1 k (-e * t-1 f (E(S); E(Ψ ( t-1) )))|.Therefore, we complete the induction. Then we have|E[ E n∈V e * t-1 k (ℓ ′ n ( t) ) - E n∈V -e * t-1 k (ℓ ′ n ( t) )]| ≤ | E n∈V e * t-1 k (e * t-1 f (E(S); E(Ψ ( t) )) - E n∈V -e * t-1 k (-e * t-1 f (E(S); E(Ψ ( t) )))| ≤ | E n∈V e * t-1 k (e * t-1 f (E(S); E(Ψ ( t-1) )) - E n∈V -e * t-1 k (-e * t-1 f (E(S); E(Ψ ( t-1) )))| ≤ • • • ≤ | E n∈V e k (ef (E(S); E(Ψ (0) )) -E n∈V -e k (-ef (E(S); E(Ψ (0) )))| ≤ 2 log( 5Km δ ) • 3σ1(∥c k ∥ + ζ e k ∥d k ∥) 4 ≤ κ/8.This completes the proof of the first result.To obtain the continuous ODE upper bound of E[A k,e t ], we first recall the updateEQUATIONEQUATIONS ) n l -1)β (t) O (i,•) ,k ) -ηt 1 2K1 e∈[±] rie E n∈V e k [ℓ ′ n (t) 1 n O (i) (t) ∥c k ∥ 2 + ∥d k ∥ 2 (2 l∈S e n,k (σ (t+1) S ) n l -1)(2 l∈S e n,k (σ (t) S ) n l -1) -ηt (K1 -1) 2K1K e∈[±] rie E n∈V e ¬k [ℓ ′ n (t) 1 n O (i) (t) ∥c k ∥ 2 ].Then, utilizing Lemma 31 and the fact |W e k,n (t)| ≤ m, we have constant c1 > 0 such thatEQUATION)where we also neglect the impact of cross-concept due to the large K = Ω(η0C(K1 -1)∥q∥ 2 /(mK1)) in Condition 1 and appropriately chosen c1.To obtain the lower bound ODE couterpart, we examine the update of the correct contributor neurons, as shown in Eq.( 48), ( 50) and (49) . In terms of the update of E[α(t) O (i,•) ,k ] where i ∈ E[U e k,n (t) ∩ (W -e k,n (t) -U -e k,n (t))], we see that its update is controlled by E[ E n∈V e k (ℓ ′ n (t) ) -E n∈V -e k (ℓ ′ n (t) )]: α (t+1) O (i,•) ,k = Θ(α (t) O (i,•) ,k - ηt∥c k ∥ 2 2mK1 [ E n∈V e k (ℓ ′ n (t) ) -E n∈V -e k (ℓ ′ n (t) )]).Then by the first result in this lemma we knowEQUATION, and thusα (t+1) O (i,•) ,k = Θ(α (t) O (i,•) ,k ± ηtκ∥c k ∥ 2 64mK1).By the condition on the small initialization in Condition 1 such that σ1 = O((2σ * S -1) 2Cm 3/2 ∥q∥ ), due to the large C, we see that the κ = O((2σ EQUATIONEQUATION)where we ignore the impact of regularization term at this stage since λ = O((C log(Km/δ)∥q∥) -1 ) and appropriately chosen c2. The third inequality is due to the definition of d k .Collaborating with Lemma 10, the proofs are completed.For the last results, following the techniques in [43] , first it's easy to check that b ( 2c k,e 3b k,e + 2 3 ) ≤ E[A k,e t ] ≤ log( c k,e b k,e t + 1) + c k,e 1 + b k,e .Proof of Lemma 30. We use induction to prove this lemma. All conclusion holds naturally at t = 0. Suppose there exists T ≤ T * such that the six conditions hold for any 0 ≤ t ≤ T -1, we prove that these conclusions also hold for t = T .We now proveEQUATION)Recall the update ruleβ (t+1) O (i,•) ,k = (1 -ηtλ)β (t) O (i,•) ,k -ηt ∥d k ∥ 2 2K1 e∈[±] ri E n∈V e k [ℓ ′ n (t) 1 n O (i) (t) ( l∈S e n,k (σ (t) S ) n l - l∈S -e n,k (σ (t) S ) n l )].As we ignore the regularization term at the first stage, we can easily seen that E[e • β (t)O (i,•) ,k 1(i ∈ U e k,n (t))] increases with t. Assume t β + ,k as the last time ∃i ∈ E[U e k,n (t)] such that E[e • β (t) O (i,•) ,k 1(i ∈ U e k,n (t))] -e • β (0) O (i,•) ,k ≤ (σ * S ) -1 log(T * ), then for i ∈ E[U e k,n ( t)] we have E[e • β ( t) O (i,•) ,k ] ≤ E[e • β (t β + ,k ) O (i,•) ,k ] -η0 ∥d k ∥ 2 2K1 e∈[±] riE[ℓ ′ n (t) 1 n O (i) (t) ( l∈S e n,k (σ (t) S ) n l - l∈S -e n,k (σ (t) S ) n l )] t=t β + ,k -η0 t β + ,k <t< T ∥d k ∥ 2 2K1 e∈[±] riE[ℓ ′ n (t) ( l∈S e n,k (σ (t) S ) n l - l∈S -e n,k (σ (t) S ) n l )] ≤ e • β (0) O (i,•) ,k + (σ * S ) -1 log(T * ) + η0∥d k ∥ 2 2mK1 - t β + ,k <t< T η0∥d k ∥ 2 2mK1 E[ℓ ′ n (t) ] ≤ e • β (0) O (i,•) ,k + 2(σ * S ) -1 log(T * ) - t β + ,k <t< T η0∥d k ∥ 2 2mK1 E[ℓ ′ n (t) 1 n O (i) (t) ],where the first inequality is by the positive nature of regularization term as well as the contribution of the gradient; second inequality is byE[-ℓ ′ n (t β + ,k ) ] ≤ 1 and E[( l∈S e n,k (σ (t) S ) n l -l∈S -e n,k (σ (t) S ) n l )] ≤ 1; the third inequality is by the condition η0 = O( mK 1∥q∥ 2 ) and thus η 0 ∥d k ∥ 2 2mK 1 ≤ 1 ≤ log(T * ), as well as (σ * S ) -1 ≥ 1. The remaining job is to prove that- t β + ,k <t< T η0∥d k ∥ 2 2mK1 E[ℓ ′ n (t) 1 n O (i) (t) ] ≤ (σ * S ) -1 log(T * ).Observe that|E[ℓ ′ n (t) ]| = E[ 1 1 + exp(yS n • e∈[±] e m i∈{r i = e m } σR(W y O (i,•) (t) l∈[L] (σ (t) S ) n l y n l ) ) ] ≤ E[exp   ( i∈W y Sn k,n (t)-U y Sn k,n (t) - i∈U y Sn k,n (t))(α(t) O (i,•) ,k + (2 l∈S y Sn n,k (σ (t) S ) n l -1)yS n β (t) O (i,•) ,k )   ] ≤ E[exp(κ/2 - 1 m i∈U y Sn k,n (t) (2σ * S -1)yS n β (t) O (i,•) ,k )] ≤ 2 exp(-log(T * )).Here the first inequality is by O (i,•) ,k 1(i ∈ U e k,n (t))] ≥ (σ * S ) -1 log(T * ).Then we have that- t β + ,k <t< T η0∥d k ∥ 2 2mK1 E[ℓ ′ n (t) 1 n O (i) (t) ] ≤ t β + ,k <t< T η0∥d k ∥ 2 2mK1 • 2 exp(-log(T * )) ≤ T η0∥d k ∥ 2 mK1 exp(-log(T * )) ≤ T * η0∥d k ∥ 2 T * mK1 ≤ (σ * S ) -1 log(T * ).We complete the proof that 0 ≤ E[e • β (t)O (i,•) ,k 1(i ∈ U e k,n (t))] -e • β (0) O (i,•) ,k ≤ 3(σ * S ) -1 log(T * ) ≤ (σ * S ) -1 α.We now prove a strong augmented hypothesis that there exist i * ∈ E[U e k,n (t)] for ∀0 ≤ t ≤ T * , we haveE[|α (t) O (i,•) ,k |/(e • β (t)EQUATION)where we set Ĉ = 2C ′ 2 log( 5Km δ ) for some constant C ′ . i * can be any element satisfies |β (0)O (i * ,•) ,k | = σ1/2∥d k ∥,which exists at t = 0 by Lemma 7 as well as the fact that ∥c k ∥ > ∥d k ∥ by their definition in Lemma 25.Suppose Eq.( 55) holds at 0 ≤ t ≤ T -1, recall the update rule and the large K condition, we can have a constant C > 1 such thatE[e • β ( t) O (i,•) ,k 1(i ∈ U e k,n ( T -1))] = E[(1 -η T -1 λ)e • β ( T -1)O (i,•) ,k 1(i ∈ U e k,n ( T -1))] -η T -1 ∥d k ∥ 2 2mK1 E[ℓ ′ n ( T -1) ( l∈S e n,k(σ( T -1) S ) n l - l∈S -e n,k(σ( T -1) S ) n l )], ≥ E[(1 -η T -1 λ)e • β ( T -1)O (i,•) ,k 1(i ∈ U e k,n ( T -1))] + η T -1 σ * S ∥d k ∥ 2 2mK1 E[|ℓ ′ n ( T -1) |], E[|α ( t) O (i,•) ,k |] = E |(1 -η T -1 λ)α ( T -1) O (i,•) ,k -η T -1 ∥c k ∥ 2 2K1 e∈[±] [eri • E n∈V e k (ℓ ′ n ( T -1) 1 n O (i) ( T -1) )] -η T -1 (K1 -1)∥c k ∥ 2 2K1K e∈[±][eriE n∈V e ¬k (ℓ ′ n ( T -1) 1 n O (i) ( T -1) )]| ≤ E[|(1 -η T -1 λ)α ( T -1)O (i,•) ,k |] + Cη T -1 ∥c k ∥ 2 2mK1 E[|ℓ ′ n ( T -1) |].where the first inequality is due to the definition of σ * S and U e k,n ( t); the second inequality is due to the large K = Ω(η0C(K1 -1)∥q∥ 2 /(mK1)). Then we haveE[|α ( t) O (i,•) ,k |] E[e • β ( t) O (i,•) ,k 1(i ∈ U e k,n ( t))] ≤ max{ E[|α ( T -1) O (i,•) ,k |] E[e • β ( T -1) O (i,•) ,k 1(i ∈ U e k,n ( T -1))] , C∥c k ∥ 2 σ * S ∥d k ∥ 2 } ≤ Ĉ ∥c k ∥ 2 σ * S ∥d k ∥ 2 ,where the last inequality is by the induction hypothesis and the C ′ can be taken as C, which completes the induction.We now prove0 ≥ E[e • β (t) O (i,•) ,k 1(i ∈ W e k,n (t) -U e k,n (t))] -e • β (0) O (i,•) ,k ≥ - Ĉ∥c k ∥ 2 σ * S 2 ∥d k ∥ 2 α - σ1(σ * S 2 ∥d k ∥ 2 + Ĉ∥c k ∥ 2 ) σ * S 2 ∥d k ∥ 2 log( 5Km δ ).Recall the update ruleβ (t+1) O (i,•) ,k = (1 -ηtλ)β (t) O (i,•) ,k -ηt ∥d k ∥ 2 2K1 e∈[±] ri E n∈V e k [ℓ ′ n (t) 1 n O (i) (t) ( l∈S e n,k (σ (t) S ) n l - l∈S -e n,k (σ (t) S ) n l )].Easy to see thatE[e • β (τ ) O (i,•) ,k 1(i ∈ W e k,n (τ ) -U e k,n (τ ))] τ =0 τ =t≤ 0 and it's decreasing. As we know that theneuron i ∈ E[W e k,n (t) -U e k,n (t)] would be deactivated at t + 1 once E[α (t+1) O (i,•) ,k + e • ( l∈S e n,k (σ (t+1) S ) n l - l∈S -e n,k (σ (t+1) S ) n l β (t+1) O (i,•) ,k ] ≤ 0. 
theoretical results: 
 In this section, we present our main theoretical results, which is based on the following conditions. We consider the learning iterations 0 ≤ t ≤ T * , whereT * = Ω(m -1 σ -1 0 σ -1 1 mλ -2 K 1 ∥q∥ 2 ((L - 1)∥u∥ 2 + 1) log(ε -1)) denotes the maximum admissible iteration. Condition 1. Suppose that there exists a sufficiently large constant C, such that the following hold:1. d X , d Y ≥ max{C log(KLBT * /δ), K}, d Y ≥ C log(m/δ), m ≥ C log(K/δ). 2. γ ≥ C max{∥q∥ 2 /(mK 1 λ), 10/λ}, λ ≤ min{(C log(Km/δ)∥q∥) -1 , (Cσ 0 /2∥u∥ 2 ) -1 } 3. K ≥ {CK 1 , C∥u∥/(σ ξ √ d X )}. 4. σ ξ ≤ min{λm/(C √ d X ∥u∥∥q∥ 1/2 ), ∥q∥/(C d Y )}. 5. σ 0 ≤ K -1 log( ∥u∥ 2 λK1 log( ∥q∥ 2 mλK1 ))/(C∥u∥), σ 1 ≤ min{(Cσ 0 ∥u∥ 4 ∥q∥ log(5Km/δ)/K 1 ) -1 , w * 2 /(Cm 3/2 ∥q∥)}.Here,w * = 1 -e -σ0 2 (1-κx) 2 ∥u∥ 4 /2 1 + e -σ0 2 (1-κx) 2 ∥u∥ 4 /2 .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 on dimensionality d X , d Y and the network width m ensure the learning problem is in a sufficiently overparameterized setting [41, 42, 52, 43] . The condition on γ 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 to control the impact of cross-concept contribution in the Attention's learning dynamic, which can actually be relaxed at the cost of a denser analysis. The condition on σ ξ is to ensure that the gradient flows be mildly influenced by the noise. Last but not least, the conditions on σ 1 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 is delayed to Appendix H. Theorem 2. Exponential Convergence of 0-1 loss. Under Condition 1, defineν := min{2 √ 2σ 1 /(1 + κ y ), σ 0 (1 -κ x )e -log(5Km/δ) σ 2 1 ∥u∥ 4 (1+e -σ 2 0 ∥u∥ 2 ) (1-e -σ 2 0 ∥u∥ 2 )}. Then, for ∀ε > 0 there exist some positive constants C 1 and C 2 , with probability no less than 1 -δ, for T ≥ T = C 1 σ 1 mλK 1 γ (1 + κ y ) log(5Km/δ)/w * 2 (1 -κ y )∥q∥, we haveL 0-1 D * (Ψ (T ) ) ≤ exp(- C 2 ν 2 mλ 2 (γ + T ) K 1 ∥q∥ 2 ((L -1)∥u∥ 2 + 1)).T ε = K 1 ∥q∥ 2 ((L -1)∥u∥ 2 + 1) C 2 ν 2 mλ 2 log( 1 ε ) iterations, we have L 0-1 D * (Ψ (T ) ) ≤ ε.Note that the bound is valid only when T ≥ T , a common threshold in prior convergence rate analyses [34, 33, 36] . Importantly, the existence of T does not affect the convergence rate as ε → 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, our techniques relax the batch size requirement, enabling more general results. Consequently, the sample complexity for Bayes-optimal test error is N = T ε .Before introducing the next proposition, we highlight a key observation from the semantic geometry in Definition 1. For anyk 1 ∈ [K 1 ], defining a k1 := (µ + k1 + µ - k1 )/2 and b k1 := (µ + k1 -µ - k1 )/2, we find that for k ′ 1 ̸ = k 1 , {a k1 , b k1 } ⊥ {a k ′ 1 , b k ′ 1 } and ⟨a k1 , b k1 ⟩ = 0.This structure is exemplified in Figure 1 (b) of [12] , where "[Bird]" consists of orthogonal steering vectors: "plant ⇒ animal" and "mammal ⇒ bird," corresponding to the concept feature a k and semantic label features b k . Here, the term eb k1 in µ e k1 determines the label assignment. Similarly, defining c k1 := (q + k1 + q - k1 )/2 and d k1 := (q + k1 -q - k1 )/2 yields analogous properties. Detailed definitions are provided in Appendix I. The following proposition explores the model's ability to handle OOD unseen ICL tasks. Proposition 1. Out-of-Distribution-Generalization D * S = K1 k=1 π + k * P + k,L * +1 * + π - k * P - k,L * +1 * .Specifically, the new D * S satisfies the following properties.• The prompt length L * can be any positive integer.z can enjoy arbitrary distribution, satisfying that each prompt in the distribution has at least one co-concept k ∈ [K 1 ], and still each word has equal chance to have positive or negative semantic labels over its concepts4 .x ×D * y can enjoy a great family of data shift.∀k ̸ = k ′ ∈ [K 1 ], k 2 ∈ [K 2 ], we can have new M * and Q * such that µ ± k * = a * k ±b * k , q ± k * = c * k ±d * k , ν k2 = ν * k2 . Here, a * k , b * k , c * k , d * k are any vectors belong to the conic hulls of {a k } K1 k=1 , {b k } K1 k=1 , {c k } K1 k=1 , {d k } K1 k=1 respectively, satisfying ∥b * k ∥ ≥ ∥a * k ∥ = Θ(∥u∥) and ∥d * k ∥ ≥ ∥c * k ∥ = Θ(∥q∥). ν * k2 = Θ(∥u∥)are any vectors from the complement space of span(M).Again, the learned model satisfies L 0-1 D * (Ψ (T * ) ) ≤ ε. This proposition demonstrates the strong Out-of-Distribution Generalization ability of transformer utilizing multi-concept 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 * z denotes that each prompt can enjoy multi-co-concepts and each word-label pair can appear in at least ∥z∥ 0 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 ×D * y 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 ICL classification and regression tasks. Our results differ by focusing on compositional generalization of learned concepts, grounded in the concept-specific linear latent geometry observed in LLMs.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. 
proof idea: 
 In a big picture, we simply extend standard expectation-variance reduction techniques [34] to our setting. Section 5.1 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.Idempotent Operator Trick. Define U := span(M) and its complement space U ⊥ . By definition, we know that dim(U) = K and dim(U⊥ ) = d X -K. Then we can let {{a k1 } K1 k1=1 , {b k1 } K1 k1=1 , {ν k2 } K2 k2=1 , {u w } d X -K w=1 } be the set of standard orthogonal basis for R d X , where u ⊥ 1 , • • • , u ⊥ d X -K are the standard orthogonal basis of U ⊥ .Then we can derive an idempotent decomposition of the identity matrixEQUATIONSimilar techniques are also applied to the label's dictionary: Q := span(Q), where we defineq ⊥ 1 , • • • , q ⊥ d Y -K1as 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(Ψ (t) ) evolves. For E(Ψ (t) ), every gradient descent update by all concept's samples within a soft "weight", and thus the analysis is equivalent to gradient descent with an ideally-balanced prompt set. Leveraging the symmetry of the prompt distribution, as well as the symmetry of W Lemma 1. We can decomposeE[W x Q ], E[W x K ] and the i-th row of E[W y O ] (i ∈ [m]) via the following (scaled) projection matrices and projection directions.E[W x Q (t) ] = K1 s=1 α (t) Q,s • a s a s ⊤ ∥a s ∥ 4 + K1 s=1 β (t) Q,s • b s b s ⊤ ∥b s ∥ 4 + K2 r=1 τ (t) Q,r • ν r ν r ⊤ ∥u∥ 4 + d X -K w=1 ρ (t) Q,w • u ⊥ w u ⊥ w ⊤ , E[W x K (t) ] = K1 s=1 α (t) K,s • a s a s ⊤ ∥a s ∥ 4 + K1 s=1 β (t) K,s • b s b s ⊤ ∥b s ∥ 4 + K2 r=1 τ (t) K,r • ν r ν r ⊤ ∥u∥ 4 + d X -K w=1 ρ (t) K,w • u ⊥ w u ⊥ w ⊤ , E[W y O (i,•) (t) ] = K1 k=1 α (t) O (i,•) ,k • c k ⊤ ∥c k ∥ 2 + K1 k=1 β (t) O (i,•) ,k • d k ⊤ ∥d k ∥ 2 + d Y -K1 w=1 ρ (t) O (i,•) ,w • q ⊥ w ⊤ .Here αEQUATIONK,s and α(t)O (i,•) ,k represent the expected concept learning process, βEQUATIONK,s and βO (i,•) ,k represent the expected concept-specific semantic learning process and τ(t) Q,r , τ (t) K,r , ρ (t) Q,w , ρ (t) K,w and ρ (t)O (i,•) ,w represent the expected memorization of the concept irrelevant noise. It holds thatEQUATIONfor ∀e ∈ [±], i ∈ [m], k ∈ [K 1 ] and for ∀e ′ ∈ [±], s ′ ∈ [K 1 ], r ∈ [K 2 ], w ∈ [d X -K], ∀u ∈ {µ e ′ s ′ , ν r , u ⊥ w }, it holds that E[(W x K (t) u)] ⊤ E[W x Q (t) µ e s ] = 0.Similar conclusions hold when the query vectors are ν r and u⊥ w , ∀r ∈ [K 2 ], w ∈ [d X -K].As such, our remaining task is to scrutinize the coefficients evolution, which would be the key contributors to the expected 0-1 loss convergence. EQUATION( i∈U y Sn k,n (t) - i∈W y Sn k,n (t)-U y Sn k,n (t) )   α (t) O (i,•) ,k + (2 l∈S y Sn n,k (σ (t) S ) n l -1)y Sn β (t) O (i,•) ,k    ≥ 0, (5) holds, we have L 0-1 D * (E(Ψ ′ (t) )) = 0.As such, the following lemmas show the learning outcomes of the E(Ψ (t) ) along the iterations. ). There exist constantC 1 > 0, ∀t ≥ T = C 1 σ 1 mλK 1 γ (1 + κ y ) log(5Km/δ)/w * 2 (1 -κ y )∥q∥, we have L 0-1 D * (E(Ψ ′ (t))) = 0. Lemma 4. (Regularizing the models). Under Condition 1, it holds thatEQUATION).In addition, our analysis provides three asymptotic properties of the coefficients evolution, which are delayed to Appendix I.1.3 and I.2 for room limitation.Proposition 2. ∀t ≥ T , when∥Ψ ′ (t) -E(Ψ ′ (t) )∥ F ≤ ν holds, we have L 0-1 D * (Ψ ′ (t) ) = 0. Here, ∥Ψ ′ ∥ 2 F := ∥W x Q ∥ 2 F + ∥W x K ∥ 2 F + ∥W y O ∥ 2 F .By definition of 0-1 loss, then we only need to prove the 0-1 loss convergence by seeing the speed of Ψ ′ (t) converging to E(Ψ ′ (t) ) with an error of ν in terms of ∥ • ∥ F .Drawing insights from [34] , we see B 0 , • • • , B T -1 as a i.i.d. random variables following the same distribution. Then ∀t ∈ {0, • • • , T }, it holds thatEQUATIONare martingale difference sequences, and for ∀X ∈ {Q, K, O} and its corresponding +1) ]. Then we utilize the following lemma in [34, 55] to give a bound over the variance.W ∈ {W x Q , W x K , W y O }, we have T t=0 D t X = W (T +1) -E[W (TLemma 5. Let D 1 , • • • , D T -1 be a martingale difference sequence. Suppose ∃c T > 0 such that T t=0 ∥D t ∥ 2 ∞ ≤ c 2 T , where ∥ • ∥ ∞ is the essential supremum of ∥ • ∥ F .Then for ∀ϵ > 0, we haveEQUATION).Therefore, we need to see if there exists a decaying positive constant c T (with decaying rate O(1/T q ), q > 0), such that supremum of ∥D t X ∥ F . Subsequently, by controlling the martingale sequence norm tail similarly in [34, 55] , we can obtain an exponential convergence rate after T 1 .T t=0 ∥D t X ∥ 2 ∞ ≤ c T 2 , ∀X ∈ {Q, K, O}, where ∥ • ∥ ∞ is the essentialFor W ∈ {W x Q , W x K , W y O },to check the decaying c T , we adopt the techniques of [34, 33, 36] in the following manner. Let B t ′ be an independent variable from B 0 , • • • , B T and let W t (T +1) be an output of the algorithm depending on(B 0 , • • • , B t-1 , B t ′ , B t+1 , • • • , B T ). Then we have ∥D t X ∥ ∞ ≤ E[∥W (T +1) -W t (T +1) ∥ ∞ | B 0 , • • • , B t ].Therefore, one may estimate c T X 2 by bounding∥W (T +1) -W t (T ) ∥ 2 ∞ uniformly w.r.t. B 0 , • • • , B T -1. Such a bound can be derived utilizing stability property of stochastic gradient descent [34, 56] . For 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-conflict. The norm requirements and constraints on L * 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 4, laying the groundwork for the proof.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes]Justification: We have provided the complete configuration in Section 3 and 6. We have uploaded the code with instructions in the supplementary material.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/ CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source 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.cc/public/guides/ CodeSubmissionPolicy) for more details. • The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc. • The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of 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) is recommended, but including URLs to data and code is permitted. 
experiments: 
 In this section, we demonstrate the validity of our theoretical analysis through simulations of Algorithm 1. We use the following parameter settings in Figure 2 : The parameter settings are: the length L = 4, the number of co-concepts K 1 = 2, dictionary size K = 104, the number of test instancesn test = 5000, dimension d X = d Y = 1000, MLP width m = 50, feature strengths ∥u∥ = ∥q∥ = 10, ∀k ∈ [K 1 ], the cosine ⟨µ + k , µ - k ⟩/∥u∥ 2 = ⟨q + k , q - k ⟩/∥q∥ 2 = 0.5, the initialization parameters σ 0 = 0.1, σ 1 = 0.01, and the noise deviation σ ξ = 0.01. For the optimization, we use λ = 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) L * = 5, (b) L * = 2, (c) a 0.8 fraction for the first concept and a 0.2 fraction for the second concepts, and (d) 2 validates our Theorem 2 and Lemma 4, which showcases the fast convergence rate and the evolution of coefficients. Figure 3 validates Proposition 1, where the learned model permits certain data shifts.µ ± 1 * = a 1 ± b 2 , µ ± 2 * = a 2 ± b 1 . FigureQuestion: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: All the experimental settings can be found in Section 3 and 6.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
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 ICL tasks by leveraging the multi-concept semantic property, highlighting their innovative potential. An important future direction is to extend the analysis to more complex scenarios.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: This paper executes algorithms 10 times and reports the average results to reduce randomness. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to 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.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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 use 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 the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets 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. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [No] Justification: This paper does not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [No] Justification: This paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects 
 
 Appendix: The theoretical analysis provided in this work introduces novel perspectives on optimization and generalization, but 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.We implement our methods using 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.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. In 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, 32] 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, 34, 35, 36, 37] have analyzed the exponential convergence of stochastic gradient descent under various settings. Meanwhile, [35] have investigated hard-margin and exponential rates in the context of structured prediction, which encompasses traditional classification as a special case. Besides, recent work also obtain the exponential rates in generalized settings such as Multi-class classification [38] 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 ReLU-activated MLP over the sparse coding data model, whose surrogate loss function is the cross-entropy loss.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, 41, 57, 53, 45, 58, 42, 52, 43, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70] . Rather than directly examining the evolution of the 0-1 loss, this line of work explicitly studies the process of reconstruction of the data's feature directions and memorization of disrupted noise in the network'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 the case of nonlinear attention-MLP transformers with within-concept positive inner products and cross-concept orthogonal data representations. The key to our analysis is assuming good initialization of attention matrices and a sufficiently low-noise condition in the SGD setting, where noise has only a mild impact on shaping neural network matrices or influencing gradient flow.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, 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 linear-attention transformers [16, 17, 71, 72] , QKcombined attention-only models [21, 73, 20, 26, 74, 75, 76, 77, 78, 54, 79] , ReLU-free MLP [80, 81, 54] or without MLP [17, 25] , impractical squared or hinge loss [25, 26, 27, 28] . 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, 83, 84] . Recent research shows that large-scale generative models, such as diffusion models and transformers, effectively encode hierarchical concepts in their latent spaces [11, 12, 13, 46, 85, 86, 87] . Moreover, [88, 89, 73] show that transformers can capture hierarchical and compositional structures in data. From a Bayesian perspective, [14, 7, 8] 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 different concepts are nearly orthogonal [9, 10, 11, 12] . Building on these insights, we explore in a theoretical context how the compositional nature of concept representations relates to transformers' ability to generalize to OOD tasks through a sparse coding modeling. We believe our OOD results are not only coincides with the transformer'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 1a in [86] ).D.1 Probablistic Lemmas on Concentration Lemma 6. Suppose that δ > 0 and ∀d ∈ {dX , dY } = Ω(log( KN L δ )), where N = BT * . Then with probability at least 1 -δ,σ 2 ξ d 2 ≤ ∥ξi∥ 2 2 ≤ 3 σ 2 ξ d 2 , |⟨ξi, ξ i ′ ⟩| ≤ 2σ 2 ξ • d log 6(N (L + 1)) 2 δ , |⟨ξi, µ⟩| ≤ ∥µ∥2σ ξ • 2 log( 6KN (L + 1) δ )for all ξi, ξ ′ i ∼ D ξx ( or D ξy ), µ ∈ Dx( or Dy), l ∈ {1, 2}.Proof. See Lemma B.4 in [42] for a proof.Lemma 7. Suppose that δ > 0, dY = Ω(log(m/δ)), m = Ω(log(K/(δ))). Then with probability at leastEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONMoreover, for some ζ ∈ (0, 1]EQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONO (i,•) ,k ) -0 ≤ 2 log( 5Km δ ) • 5σ1(∥c k ∥ + ζ∥d k ∥) 16 , i∈{i∈[m]|α (0) O (i,•) ,k +e ′ ζβ (0) O (i,•) ,k >0} ri • β (0) O (i,•) ,k -0 ≤ 2 log( 5Km δ ) • 5σ1∥d k ∥ 16 .(8) In addition, for a sufficient large m = Ω(log(K/(δ))/(1 -ω ζ )) the lower bound inequalities regarding maximum value in Eq.( 7) hold at any above index set of i in Eq. (8) . For example, there existEQUATIONEQUATIONEQUATIONO (i,•) ,k < 0}, such that α (0) O (i,•) ,k ≤ -σ1/2∥c k ∥.Proof. First, notice that W yO (i,•) (0) ∼ N (0, σ1I d Y ), then by Bernstein's inequality as well as dY = Ω(log(m/δ)), with probability at least 1 -δ/(5m), for ∀i ∈ [m]|∥W y O (i,•) (0) ∥ 2 -σ1dY | ≤ O(σ 2 1 • dY log(5m/δ)) ≤ σ 2 1 dY /2.By union bound we can have the first inequality in the lemma hold with probability at least 1 -δ/5.Next, we notice thatEQUATIONO (i,•) ,w = ⟨W y O (i,•) (0) , q ⊥ w ⟩are all Gaussian random variable with mean 0 and variance σ 2 1 . Then by Gaussian tail bound and union bound, with probability at least 1 -δ/10, for all i ∈ [m] and q∈ k,w { c k ∥c k ∥ , d k ∥d k ∥ , q ⊥ w }, it holds that |⟨W y O (i,•) (0) , q⟩| ≤ 2 log(5Km/δ) • σ1. Notice P(σ1/2 > |⟨W y O (i,•) (0), q⟩|) is an positive constant, then following the techniques of Lemma B.5 in [42] and the condition m = Ω(log(K/δ)), we haveP(σ1/2 ≤ |⟨W y O (i,•) (0) , q⟩|) = 1 -P(σ1/2 > max{|⟨W y O (i,•) (0) , q⟩|}), = 1 -P(σ1/2 > |⟨W y O (i,•) (0) , q⟩|) mK ≥ 1 -δ/10,then with probability 1 -δ/5, the second and third inequality hold.For ζ ∈ (0, 1], we see that the variable α (0) EQUATIONO (i,•) ,k ∼ N (0, σ 2 1 (∥c k ∥ 2 + ζ∥d k ∥ 2 )),EQUATIONEQUATIONO (i,•) ,k > 0} as a binomial variable with p = 1/4, n = m, then by the property of binomial tail, condition m = Ω(log(K/(δ))) as well as Hoeffding's inequality, with probability at least 1 -δ/5 we haveEQUATIONEQUATIONO (i,•) ,k > 0}| m - 1 4 | ≤ log(10K1/δ) 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 -δ/5.EQUATIONEQUATIONO (i,•) ,k > 0}, we haveEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONO (i,•) ,k > 0), = 1 2 • 1 + ω ζ 2 ,.where1 + ω ζ 2is the probability of the conditional event {α (0)EQUATIONO (i,•) ,k > 0 | α (0) O (i,•) ,k + e ′ ζβ (0) O (i,•) ,k > 0}, and ω ζ > 0 due to the larger variance of α (0) O (i,•) ,k compared to e ′ ζβ (0)O (i,•) ,k . We denote the probability with ω ζ since the true value is hard to compute. Subsequently, the event {i ∈EQUATIONO (i,•) ,k ± ζβ (0) O (i,•) ,k > 0} can be seen as a binomial variable with p = 1 + ω ζ 8, n = m, then we can have the sixth inequality hold with probability at least 1 -δ/5, utilizing the property of binomial tail, condition m = Ω(log(K/(δ))) as well as Hoeffding's inequality.The seventh inequality is a natural inference of the third and forth inequality, where the m = Ω(log(K1/δ)) ensure m log(10K1/δ)/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.Lemma 8. (1.1.P5 in [90] ) Let A ∈ Mn be idempotent, that is,A 2 = A. Then, each eigenvalue of A equals to the rank of A, which is either 0 or 1. Beside, identity matrix I is the only nonsingular idempotent matrix.Lemma 9. For a matrix A = d i=1 µiPi, where Pi are symmetric idempotent matrices with rank(Pi) = 1, and thus d i=1 µiPi is the idempotent decomposition of matrix A by Pi. Then we see that∥A∥F = tr(A T A) = d i=1 µ 2 i = d i=1 λ 2 i, where λi are eigenvalues of A.A T A = d i=1 µ 2 i P T i Pi = d i=1 µ 2 i PiPi = d i=1 µ 2 i Pi.Then, by Lemma 8 we havetr(A T A) = tr( d i=1 µ 2 i Pi) = d i=1 µ 2 i tr(Pi) = d i=1 µ 2 i rank(Pi) = d i=1 µ 2 i = d i=1 λ 2 i .Lemma 10. (Lemma C.1 in [43] ). Suppose that a sequence at, t ≥ 0 follows the iterative formulaat+1 = at + c 1 + be a t , for some 0 ≤ c ≤ 1 and b ≥ 0. Then it holds that xt ≤ at ≤ c 1 + be a 0 + xt for all t ≥ 0.Here, xt is the unique solution ofdxt dt = c 1 + be x t , x0 = a0 ⇔ xt + be x t = ct + a0 + be a 0 .Lemma 11. (Coupled ODE System 1). Suppose that there are two coupled sequences yt, zt, t ≥ 0 follows the iterative formulayt+1 = yt + aztyt 1 2 + e -2y t 2 + e 2y t 2 , y0 > 0, a > 0, zt+1 = zt + b, z0 < 0, b > 0,for some a, b ≥ 0. Then it holds that y(t) ≤ yt, z(t) = zt, for all t ≥ 0. Here, y(t), z(t) are the unique solutions of the following ODE System respectivelyEQUATIONAs such, for t1 = min{t ∈ Z | zt ≥ 0}, we haveyt 1 ≥ y(0)e -az(0) 2 (1 + e -2y(0) 2 ) 4b(1 -e -2y(0) 2 ),and t1 ≥ -z(0)(1 + e -2y(0) 2 ) b(1 -e -2y(0) 2 ).Proof. From the condition we see that z0 < 0 and zt is an increasing sequence (zt ≥ z0). Besides, as y0 > 0, during the period where zt ≤ 0, we see that yt is monotonically decreasing. Then by (2 + e -2y t2 + e 2y t 2 ) -1 ≤ 1/4 as well as Comparison Theorem, it's obvious that the continuous coupled ODE in Eq.( 9) is the lower bound of yt. Then one can readily obtain the result by solving the ODE.Lemma 12. (Coupled ODE System 2). Suppose that there are two coupled sequences yt, zt, which are the sequences after t1 in Lemma 11, and t ≥ t1 follows the iterative formulayt+1 = yt + aztyt 1 for some a, b ≥ 0, and c ′ ≤ ℓ ′ t ≤ 1. Then it holds that y(t) ≤ yt ≤ y(t), z(t) ≤ zt ≤ z(t),for all t ≥ t1. Here, y(t), y(t), z(t), z(t) are the unique solutions of the following ODE System respectively1 2 (Ei(2y(t) 2 ) + Ei(-2y(t) 2 ) + 4 log(y(t))) = abc ′2 1 -e -2y(t 1 ) 2 1 + e -2y(t 1 ) 2 (t -t1) 2 2 + 1 2 (Ei(2yt 1 2 ) + Ei(-2yt 1 2 )) + 4 log(yt 1 ), z(t) = bc ′ 1 -e -2y(t 1 ) 2 1 + e -2y(t 1 ) 2 (t -t1), 1 2 (Ei(2y(t) 2 ) + Ei(-2y(t) 2 ) + 4 log(y(t))) = ab(t -t1) 2 2 + 1 2 (Ei(2yt 1 2 ) + Ei(-2yt 1 2 )) + 4 log(yt 1 ) z(t) = b(t -t1), where Ei(x) = x -∞ e t t dt = γEuler + ln x + exp(x/2) ∞ n=1 (-1) n-1 x n n!2 n-1 ⌊(n-1)/2⌋ k=0 1 2k + 1 .Proof. We see that as zt ≥ 0, t ≥ t1, the yt is monotonically increasing. As such, by Comparison Theorem we see that the upper and lower bound of the coupled system would depends on 1-e -2y(t) 2 1+e -2y(t) 2 and ℓ ′ t . Easy to see that1 -e -2y(t 1 ) 2 1 + e -2y(t 1 ) 2 ≤1 -e -2y(t) 2 1 + e -2y(t) 2 ≤ 1, and then collaborating with c ′ ≤ ℓ ′ t ≤ 1 we can obtain the result by solving the ODE. Observing thatdy(t) dt = abc ′2 1 -e -2y(t 1 ) 2 1 + e -2y(t 1 ) 2 (t -t1)y(t)dt 1 + e 2y(t) 2 + e -2y(t) 2 ⇔ 1 2 (Ei(2y(t) 2 ) + Ei(-2y(t) 2 ) + 4 log(y(t))) = abc ′2 1 -e -2y(t 1 ) 2 1 + e -2y(t 1 ) 2 (t -t1) 2 2 + const, z(t) = bc ′ 1 -e -2y(t 1 ) 21 + e -2y(t 1 ) 2 (t -t1). Thus by the monotonicity the system is unique, which is also ture for the upper bound ODE. The proof is completed.This section provided the detailed formal definitions of the prompt distribution. Definition 3. (Polysemous Word Model (Dx, Dy, Dz, D ξx , D ξy ) ). We assume there exists K1 concepts of words totally. Specifically, each concept k1 ∈ [K1] is characterized by two semantically-opposite feature vectors separately, denoted as µ + k 1 and µ - k 1 , and the label vectors that describe their semantics under the co-concept are q + k 1 and q - k 1 . Our word samples x ∈ R d X and their corresponding labels y ∈ R d Y are generated i.i.d. from distribution Dx and Dy, which can be written as the following forms via reparameterization:z ∼ Dz, ξx ∼ D ξx = N (0, σ 2 ξ I d X ), ξy ∼ D ξy = N (0, σ 2 ξ I d Y ), x = Mz + ξx ∼ Dx, y = Qz + ξy ∼ Dy,where z ∈ R K (K < dX ). We denote z as the sparse latent signal and ξ as the spurious dense noise, and each x-y pair are reparameterized by one shared z. We have the following assumptions on M, z, ξ respectively:• The sparse latent variable z = (z1, • • • , zK ) ∈ {0, 1} k is sampled from Dz. P (zj = 1) = Θ( log log K K ). • M = [µ + 1 , µ - 1 , µ + 2 , µ - 2 , • • • , µ + K 1 , µ - K 1 , ν1, ν2, • • • , νK 2 ] = [M1, • • • , MK ] ∈ R d X ×K is the feature dictionary matrix, where {µ ± k 1 } K 1 k 1 =1 are concept-relevant features, {ν k 2 } K 2 k 2 =1are conceptirrelevant features, and ∀k ∈ [K], ∥M k ∥ = ∥u∥. We assume that features of the same concept have positive inner product:∃0 < κx < 1, ∀k1 ∈ [K1], 0 < ⟨µ + k 1 , µ - k 1 ⟩ ≤ κx∥u∥ 2 . Meanwhile, we let the features of different concept be orthogonal: ∀e ∈ [±], e ′ ∈ [±], s ′ ∈ [K1], r ̸ = r ′ ∈ [K2], u ∈ {µ e ′s ′ , νr}, we have ⟨µ e s , u⟩ = ⟨νr, ν r ′ ⟩ = 0.• Q = [q + 1 , q - 1 , q + 2 , q - 2 , • • • , q + K 1 , q - K 1 , 0, • • • 0] ∈ R d Y ×Kis the corresponding label dictionary matrix, where ∥q ± k ∥ = ∥q∥, for ∀k ∈ [K1]. Similarly, we let the labels of the same concept to have positive inner product:∃0 < κy < 1, ∀k1 ∈ [K1], 0 < ⟨q + k 1 , q - k 1 ⟩ ≤ κy∥q∥ 2, while the labels of different concept to be orthogonal:⟨q ± k , q ± k ′ ⟩ = 0, ∀k ̸ = k ′ ∈ [K1]. Definition 4. (Concept-specific Contextual Prompt Distribution)We consider the case that each prompt is concept-specific (i.e., 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 -1 ), and the chance for selecting the two semantically-opposite vectors of the same concept is 1 2 . During training, each promptS = {x1, y1, • • • , xL, yL, xL+1}is sampled from the mixture distribution DS defined as below.EQUATIONwhereπ + k = π - k = 1 2K1, and the P + k,L+1 and P - k,L+1 are prompt distributions characterized by the k-th concept, defined asP + k,L+1 = S | x ∼ Dx, y ∼ Dy, P L+1,2k-1 = 1, ∀l ∈ [L + 1], j ̸ = {2k -1, k}, P l,j = 1 K , {z l,2k-1 = 1} ∪ {z l,2k = 1} = Ω, {z l,2k-1 = 1} ∩ {z l,2k = 1} = ∅, ∀l ∈ [L], P l,2k-1 = P l,2k = 1 2 , P - k,L+1 = S | x ∼ Dx, y ∼ Dy, P L+1,2k = 1, ∀l ∈ [L + 1], j ̸ = {2k -1, k}, P l,j = 1 K , {z l,2k-1 = 1} ∪ {z l,2k = 1} = Ω, {z l,2k-1 = 1} ∩ {z l,2k = 1} = ∅, ∀l ∈ [L], P l,2k-1 = P l,2k = 1 2 ,where P l,j := P (z l,j = 1). ∀n ∈ [N ] where N is the training size, if the training prompt Sn is sampled fromP e k,L+1 , e ∈ [±], k ∈ [K1],then by Definition 1, the label vector of the query should contain q e k , and we call yS n = e as the real value label of this k-th concept prompt. Specifically, for ∀k ∈ [K1] we define the index set of training prompts sharing the k-th co-concepts asV k = V + k ∪ V - k , where V + k = n | Sn ∼ P + k,L+1 , V - k = n | Sn ∼ P - k,L+1 . For sample x l where n ∈ V k , k ∈ [K1], l ∈ [L + 1], we define the index set for its non-zero elements of z n l besides z n 2k-1,l and z n 2k,l , namelyM n l := {k ∈ [K] | z n l,k = 1, k / ∈ {2k -1, 2k}}. Also, for each prompt sharing the k-th co-concept, we define the index set of demonstration in the context:S + n,k = {l ∈ [L] | n ∈ V k , z n l,2k-1 = 1}, S - n,k = {l ∈ [L] | n ∈ V k , z n l,2k = 1},In this section, we provide several important definitions and compute the original gradients of attention. Lemma 13. (Contributing and Misleading Neurons)EQUATIONEQUATIONt) are neurons that correctly contribute to the prediction. The following lemma computes the original gradients. Lemma 14. (Gradient Update) DenoteEQUATIONEQUATIONSubsequently, we directly compute the update of the attention matrices along the feature directions as below.Lemma 15. (Concept Learning of Attention) For ∀ k ∈ [K1],we have the single step of learning of the concept part of the features:EQUATIONwhereEQUATIONI (t)Q,a k ,contri are defined as below.EQUATIONriW yEQUATIONSimilarly, IK,a k ,chaos andI (t)K,a k ,contri are defined as below.I (t) K,a k ,chaos = 1 B k̸ = k∈[K 1 ] e∈[±] n∈V e k ∩B t e • ℓ ′ n (t) a ⊤ k W x Q (t) (a k + eb k + ξ n x,L+1 + r∈M n L+1 Mr) i∈W e k,n (t) ri • l,j∈[L] (σ (t) S ) n l (σ (t) S ) n j (W y O (i,•) (t) (q y n l k + s∈M n l QS + ξ n y,l ))a k ⊤ ( s∈M n l Ms + ξ n x,l - s∈M n j Ms -ξ n x,j ) + 1 B ê∈[±] n∈V ê k ∩B t êℓ ′ n (t) a ⊤ k W x Q (t) (a k + eb k + ξ n x,L+1 + r∈M n L+1 Mr) • i∈W ê k,n (t) ri• l,j∈[L](σEQUATIONriW yEQUATION) Lemma 16. (Label Semantic Learning of Attention) Also, for ∀ k ∈ [K1],we have the single step of learning of the concept-specific semantically-opposite part of the features:EQUATIONwhereEQUATIONI (t)Q,b k ,contri are defined as below.EQUATIONriW yEQUATIONSimilarly, IK,b k ,chaos andI (t)K,b k ,contri are defined as below.EQUATIONriW yEQUATIONG Model details: MLP Part Lemma 17. (Tensor Update)EQUATIONEQUATIONEQUATIONEQUATIONwhereEQUATIONI (t)O (i,•) ,c k ,contri are defined asEQUATIONEQUATION)Remark 2. (Informal Discussions). Interestingly, the gradient of MLPs' Concept Learning is very large. We have the following situations.• When the neuron is activated (i.e., {n ∈V ê k ∩ Bt, if (W x K (t) b k ) ⊤ W x Q (t) b k > 0} , and α (t) O (i,•) , k + ê • (2 l∈S ê n, k (σ (t) S ) n l -1)W y O (i,•) (t) d k > 0), the neuron is likely to be activated (i ∈ W ê k,n (t)).1. If (1) ri • ê > 0, i ∈ W ê k,n (t) ⇔ i ∈ W ê k,n (t) ∩ U ê k,n (t), the gradient will advance the W yO (i,•) (t) c k ; 2. if (2) ri • ê < 0, i ∈ W ê k,n (t) ⇔ i ∈ W ê k,n (t) -U ê k,n (t), the gradient will diminish the W yEQUATIONwhereI (t) O (i,•) ,d k ,chaos and I (t)O (i,•) ,d k ,contri are defined asEQUATIONEQUATIONNote 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 on dimensionality dX , dY and the network width m ensure the learning problem is in a sufficiently overparameterized setting where the norm and the inner products of the Gaussian noise and initialized NN can be controlled within a certain range with high probability 1 -δ, which is standard requirements in recent feature learning line-of-research [41, 57, 53, 45, 58, 42, 52, 43] . The weak requirement on network width m allows us to conduct a fine-grained analysis based on the network projection length, which is fundamentally differs from the NTK line of research [91] that requires an infinitely wide network to perform linear regression over a prescribed feature map. The condition on γ ensures the learning step to be small and thus learning process enjoys an approximation to gradient flow rather than the challenging "Oscillation" regime [92] , 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 at the origin point, and ensure that we can analyze the expected learning dynamic with limited impact of the regularization at the initial stage, which is also adopted in [53] . The condition on K is to control the impact of cross-concept contribution in the Attention's learning dynamic, which can actually be relaxed at the cost of a denser analysis. The condition on σ ξ 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 norms, which ensures the gradient flows to be not too noisy and could converge to the expected gradient flow exponentially. Last but not least, the conditions on σ1 guarantee that the initial beliefs of MLP is small and the gradients of SGD can update the model effectively. The condition of σ0 is only used when discussing the OOD scenario.In this section, we assume all the events in the Section D hold, denoted as ΥPre.We examine the evolution ofE(Ψ ′t ) := {E(W x Q (t) ), E(W x K (t) ), E(W (t)O (i,•) )} at the whole iteration 0 ≤ t ≤ t, where the expectation E[•] is taken over the stochastic batches. As such, we can see every stochastic gradient update within each batch as a gradient update upon noise-free and category-balanced concept-specific prompts.Lemma 21. For ∀k1 ∈ [K1], we define a k 1 := µ + k 1 + µ - k 1 2 and b k 1 := µ + k 1 -µ - k 1 2. By definition, we then haveEQUATIONfor ∀k ′ 1 ̸ = k1 ∈ [K1]. Remark 3.We observe that, through this formulation, the shared component a k 1 can be interpreted as the "concept" part of the two features, while the terms ±b k 1 represent their opposing semantic aspects. The relevance of this modeling is exemplified by Figure 1 (b) in [12] , where the concept "[Bird]" is composed of orthogonal steering vectors: "plant ⇒ animal" and "mammal ⇒ bird." These vectors correspond to the concept feature a k and the semantic label features b k , respectively.Idempotent Operator Trick. Define U := span(M) and its complement space U ⊥ . By definition, we know that dim(U) = K and dim(U ⊥ ) = dX -K. Then we can have a set of standard orthogonal basis for R d , defined asβ U⊕U ⊥ = { a1 ∥a1∥ , b1 ∥b1∥ , a2 ∥a2∥ , b2 ∥b2∥ , • • • , aK 1 ∥aK 1 ∥ , bK 1 ∥bK 1 ∥ , ν1 ∥u∥ , ν2 ∥u∥ , • • • , νK 2 ∥u∥ , u ⊥ 1 , • • • , u ⊥ d X -K }, where u ⊥ 1 , • • • , u ⊥ d X -K are the standard orthogornal basis of U ⊥ .Then we can derive thatEQUATION) Lemma 22. (Partial Statement of Lemma 1). E[W x Q ] and E[W x K ]are identical and symmetric during the whole iterations. We can decomposeEQUATIONwhere α(t)Q,s and α(t)K,s represent the concept learning process, βQ,s and β(t)K,s represent the concept-specific semantic learning process and τ(t) Q,r , τ (t) K,r , ρ (t) Q,w , ρ (t)K,w represent the memorization of the concept irrelevant noise.Proof. Apparently they hold at t = 0, suppose it holds at step t, thusE[W x K (t) ] = E[(W x K (t) ) ⊤ ] = E[W x Q (t) ] = E[(W x Q (t) ) ⊤ ], we examine t + 1. It holds that EB t [W (t+1) K | E(Ψ ′(t) )] = E[W x K (t) ] -ηtEB t [∂ W x K (t) LB t (E(Ψ ′(t) ))] EB t [W (t+1) Q | E(E(Ψ ′(t) ))] = E[W x Q (t) ] -ηtEB t [∂ W x Q (t) LB t (E(Ψ ′(t) ))]Here, we see E(Ψ ′(t) ) as fixed matrices and the expectation EB t [•] is taken over the stochastic batch at the time step t. As we are considering expectation over the isotropic prompt distribution, which can be seen as a noiseless distribution with an averaged categories of words and labels, the expected gradient form could be written as symmetric form:EB t [W x K (t+1) | E(Ψ ′(t) )] -E[W x K (t) ] = K 1 s=1 (a (t) K,s asas ⊤ + b (t) K,s bsbs ⊤ ) + λ( K 2 r=1 c (t) Q,r νrνr ⊤ + d X -2K 1 -K 2 w=1 d (t) K,w • uwuw ⊤ )with some coefficients a(t) K,s , b (t) K,s , c (t) Q,r , d (t) K,w , ∀s ∈ [K1], r ∈ [K2], w ∈ [dX -2K1 -K2]. It's direct to check that E[W x Q (t)] also has the exactly same outcome. The proof is completed.EQUATIONWe will also haveEQUATIONfor ∀e ∈ [±] and for ∀eEQUATION)Similar conclusions hold when the query vectors are νr and u ⊥ w , ∀r ∈ [K2], w ∈ [dX -K]. Definition 5. Define Q := span(Q) and its complement space Q ⊥ , we can decompose i-th row of W y O via the following decomposition: (σEQUATION)EQUATIONEQUATIONfor ∀e ∈ [±], i ∈ [m], k ∈ [K1].(0) S ) n l ] -1 = exp(σ 2 0 ∥b k ∥ 2 ) -exp(-σ 2 0 ∥b k ∥ 2 ) exp(σ 2 0 ∥b k ∥ 2 ) + exp(-σ 2 0 ∥b k ∥ 2 ), then we have someω ζ e k ∈ (0, ω ′ ζ e k ) where ω ′ ζ ek < 1, the following will holdEQUATIONEQUATIONEQUATIONEQUATIONEQUATION.The parameter ω ′ ζ ek is determined by σ0, σ1, ∥a k ∥, ∥b k ∥, ∥c k ∥ and ∥d k ∥.Proof. We have that EEQUATION)] ̸ = ∅. By Lemma 7, we see that forζ e k = 2 E n∈V k [ l∈S y Sn n,k (σ (0) S ) n l ] -1 = exp(σ 2 0 ∥b k ∥ 2 ) -exp(-σ 2 0 ∥b k ∥ 2 ) exp(σ 2 0 ∥b k ∥ 2 ) + exp(-σ 2 0 ∥b k ∥ 2 ) , we can have corresponding ω ζ e k ∈ (0, ω ′ ζ e k ) where ω ′ ζ ek < 1 to ensure the conclusion holds.Lemma 24. (Coefficient Update) Denote E(Ψ ′(t) ) := {E(W x Q (t) ), E(W x K (t) ), E(W (t)O (i,•) )}, where the expectation E[•] is taken over the stochastic batches. We haveEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATION. By definition, we then haveEQUATION)for ∀k ′ 1 ̸ = k1 ∈ [K1].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 t > 0, we haveEQUATIONEQUATIONwhere r ∈ [K2], w ∈ [dX -K], ŵ ∈ [dY -K1].Lemma 27. For t > 0, we haveEQUATIONProof. The deduction is direct by the symmetric property of prompt distribution in Lemma 22, and the gradient forms in Lemma 15 and Lemma 16.This lemma reveals that the attention layer mainly serves to learn the different semantic part of each concept, and hardly have interest in learning the shared co-concept part. Also, collaborating with Lemma 22, we see thatβ (t+1) Q,k = β (t+1) K,k , this indicates that the signal of β (t) Q,k • β (t)K,k would remain positive.Also, by the symmetry property of learning progress denoted in Lemma 22, we see that ∀k ∈ [K1], αEQUATIONEQUATIONWe see from Lemma 7 that α (0)EQUATIONObviously, The observation in Eq. ( 39), collaborating with the positiveness of βEQUATION(t) Q,k • β (t)K,k , we see that the inequalityEQUATIONn j ] will remain during whole iteration. Also, by Eq.(39), we know thatEQUATIONThis observation under our expectation scenario greatly facilitate our analysis. Since ℓ ′ n (t) < 0, it's obvious that the signal of riβ(t)O (i,•) ,k will determine whether the neuron i ∈ [m] will serve to increase or decrease the βEQUATIONK,k during the gradient update. We therefore start to analyze the MLP's update below based on Lemma 16.Lemma 28. For t > 0, we haveEQUATIONEQUATIONEQUATIONEQUATIONwherek ∈ [K1].Proof. The proof is direct by the symmetric property of prompt distribution in Lemma 22, and the gradient forms in Lemma 19 and Lemma 20.An interesting fact is that the E(IO (i,•) ,c k ,chaos ) also contributes to the learning of 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 29. For a certain iterations t ∈ (0, T1), for ∀k ∈ [K1], e ∈ [±], we haveEQUATION] is non-increasing, and all of this neuron will get deactivated. Additionally, both E[α EQUATIONEQUATIONO (i,•) ,k would monotonically increase. Besides, when there exists constant C ≥ 1 such thatE n∈V e k (ℓ ′ n (t) )] ≤ C E n∈V -e k (ℓ ′ n (t) )]. the E[α (t) O (i,•),k ] would be contributed to increase, otherwise it will decrease. Also, |α (t) EQUATIONO (i,•) ,k ≈ E[(2 l∈S e n,k (σ (t) S ) n l -1)eβ (t) O (i,•) ,k ].Proof. By Lemma O (i,•) ,k compared to eβ (t) O (i,•) ,k in the neuron sets E[U e k,n (t) -(W -e k,n (t) -U -e k,n (t))],whose learning speed ratio is at least (∥c k ∥/∥d k ∥) 2 . Therefore, the absolute value of α (t)O (i,•) ,k will surpass that of eβ (t)O (i,•) ,k , which indicates the neuron would be activated for opposite labels, then the proof is completed. Given that( l∈S e n,k (σ (t) S ) n l -l∈S -e n,k (σ (t) S )n l ) will remain positive, the discussion over eβ (t) O (i,•) ,k is simple since it will always grow in ri's direction, and thus the third and forth point hold.Considering the growth of α (t)O (i,•) ,k , by π + k = π - k , P l,2k-1 = P n l,2k = 1 2, we knowEQUATION)EQUATIONn l ],EQUATIONS ) n l -1)eβ (t) O (i,•) ,k ] ≥ 0.We see that forE n∈V e k [U e k,n (t) ∩ (W -e k,n (t) -U -e k,n (t))], the E[V e k ]will serve to increase the α (t)O (i,•) ,k , but E[V -e k ]will serve to decrease the αO (i,•) ,k . The contribution will tend to be positive ifE n∈V e k (ℓ ′ n (t) (i ∈ U e k,n (t) ∩ (W -e k,n (t) -U -e k,n (t))))] ≥ E n∈V -e k (ℓ ′ n (t) 1(i ∈ U e k,n (t) ∩ (W -e k,n (t) -U -e k,n (t))))].Then O (i,•) ,k ] is comparatably feeble as it receive the contribution oppositely from E n∈V e k (ℓ ′ n (t) 1(i ∈ E[W e k,n (t) -U e k,n (t) ∩ U -e k,n (t)))] and E n∈V -e k (ℓ ′ n (t) 1(i ∈ E[W e k,n (t) -U e k,n (t) ∩ U -e k,n (t)))].Quantatively this is validated by our later results in Lemma 32 where the En∈V e k (ℓ ′ n (t) ) -E n∈V -e k (ℓ ′ n (t)) would be controlled by the initialization. Interestingly, we see that asEQUATIONS ) n l -1)eβ (t) O (i,•) ,k ] grows up, its scale will surpass those of E[α (t) O (i,•) ,k ]. Under this scenario, E n∈V e k [U e k,n (t) ∩ (W -e k,n (t) -U -e k,n (t))] will turn into E[U e k,n (t) -(W -e k,n (t) -U -e k,n (t))], where E[α (t) O (i,•),k ] again continues to grow. Thus finally we haveα (t) O (i,•) ,k ≈ E[(2 l∈S e n,k (σ (t) S ) n l -1)eβ (t) O (i,•) ,k ].Lemma 4 will show that the growing of E[(2 l∈S e n,k (σ(t) S ) n l -1)eβ (t)O (i,•) ,k ] will stuck, and thus the growing ofEQUATIONO (i,•) ,k ] will also stuck at the changing margin from En∈V e k [U e k,n (t) ∩ (W -e k,n (t) -U -e k,n (t))] into E[U e k,n (t) - (W -e k,n (t) -U -e k,n (t))]. The proof is completed.Proof. Proof of Lemma 2. To examine the 0-1 loss, by definition, we knowL 0-1 D * (E(Ψ t )) = PS n∼D * (yS n • f (E(Sn), E(Ψ t )) ≤ 0), = PS n∼DS (yS n • e∈[±]e mi∈{r i = e m } E Ψ (t) [σR(W y O (i,•) (t) l∈[L] (σ (t) S ) n l y n l )] ≤ 0), = P(E[yS n •   e∈[±]e mi∈{r i = e m } σR(W y O (i,•) (t) l∈[L](σEQUATION-EQUATION)(α(t)O (i,•) ,k + (2 l∈S y Sn n,k (σ (t) S ) n l -1)yS n β (t)O (i,•) ,k ) ≤ 0).Therefore, a sufficient condition for L 0-1EQUATIONEQUATIONEQUATIONEQUATIONEQUATIONfor ∀k ∈ [K1], e ∈ [±].We know ∀i ∈ U 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 is under controlled during several periods, which can be safely done due to small initialization by a large γ, as well as the slow quadratic decaying nature of the derivative of η ′ t . We see that at initialization, by Lemma 7 and Lemma 23, theEQUATIONEQUATIONand our remaining job is to see when willE Sn∼D S [f (E(Sn); E(Ψ (t)))] stay positive for some error tolerance.As such, we need to scrutinize the coefficients that would grow along the iterations. Therefore, we defineEQUATIONW yEQUATION(2l∈S y Sn n,k (σ (τ ) S ) n l -1)yS n W y O (i,•) (τ ) d k ] τ =0 τ =t.We will see that the conditional expectation of this sequence (conditioned on E(Ψ ′(t) ), and the expectation is taken over DS) would grow up to conquer the small initialization and make ESn∼D S[f (E(Sn); E(Ψ ′(t) ))] stay positive. Consider the whole training duration 0 ≤ t ≤ T * , the evolving speed of βEQUATIONEQUATIONEQUATIONO (i,•) ,k |},We will show that σ * S is the lower bound ofEQUATIONS )n j ]} along the whole iteration. By Lemma 7, κ can be upper bounded by 8 2 log(5Km/δ) • σ1( (1 + κy)/2∥q∥), and lower bounded by 2 √ 2σ1∥q∥, which is a negligible term due to the small initialization by Condition 1.Lemma 30. Under Condition 1, for the whole iteration0 ≤ t ≤ T * , for ∀i ∈ [m], e ∈ [±], k ∈ [K1], r ∈ [K2], w ∈ [dX -K],we have thatEQUATIONEQUATIONEQUATIONEQUATIONLemma 31. Suppose Eq. ( 45) holds at iteration t ≤ T2, then we haveE n∈V k [yS n f (E(S); E(Ψ (t) ))] -E[A k,y Sn t+1 ] ≤ κ/2.Proof. By definition, we haveE[yS n f (E(S); Ψ (t) )] = E[yS n • e∈[±]e mi∈{r i = e m } σR(W y O (i,•) (t) l∈[L] (σ (t) S ) n l y n l )] = E 1 m ( i∈U y Sn k,n (t) - i∈W y Sn k,n (t)-U y Sn k,n (t) )   α (t) O (i,•) ,k + (2 l∈S y Sn n,k (σ (t) S ) n l -1)yS n β (t) O (i,•) ,k    .Observe thatEQUATIONO (i,•) ,k + (2 l∈S y Sn n,k (σ (τ ) S ) n l -1)yS n β (τ ) O (i,•) ,k τ =0 τ =t ≤ κ/4.Here the inequality holds due to the fact that E[α (t) EQUATIONEQUATION(2l∈S y Sn n,k (σ (τ ) S ) n l -1)yS n β (τ ) O (i,•) ,k ] τ =0 τ =t ≤ κ/4.Combining the two we can see the result is obtained.We then denote the last time when there still exists E[A k,e t ] ≤ κ as T , formally T is the last time wherek∈[K 1 ],e∈[±] {E[A k,e t ] ≤ κ} ̸ = ∅.Latter we will show in Lemma 33 thatT = C1σ1mλK1γ (1 + κy) log(5Km/δ) (2σ * S -1) 2 (1 -κy)∥q∥ .We then denote the learning step at T as η := η T , and thusη = η T = 2 λ( T + γ).By Lemma 31, actually it would hold thatE Sn∼D S [f (E(Sn); E(Ψ ( T ) ))] ≥ κ/2 ≥ 0.And thus the 0-1 loss converges to 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 bound of the sequence A k,e t . Lemma 32. Under Condition 1, suppose Eq.( 45) holds at any iteration t ≤ T * , then for ∀t ≤ T * , ∀k ∈[K1], e ∈ [±], it holds that 1. The difference | E n∈V e k (ef (E(S); E(Ψ (t) )) -E n∈V -e k (-ef (E(S); E(Ψ (t) )))| is none-increasing.2. The difference of the loss derivative is bounded by O(κ):|E[ E n∈V e k (ℓ ′ n (t) ) -E n∈V -e k (ℓ ′ n (t) )]| ≤ κ 8 .i ∈ E[W e k,n ( t) -U e k,n ( t)], E[α ( t) O (i,•) ,k + e • ( l∈S e n,k (σ ( t) S ) n l - l∈S -e n,k (σ ( t) S ) n l β ( t) O (i,•) ,k ] ≥ 0.Now collaborating with Eq. ( 54) and Eq. ( 55), we now can haveE[e • ( l∈S e n,k (σ ( t) S ) n l - l∈S -e n,k (σ ( t) S ) n l β ( t) O (i,•) ,k ] ≥ -E[α ( t) O (i,•) ,k ] ≥ - Ĉ∥c k ∥ 2 σ * S ∥d k ∥ 2 ((σ * S ) -1 α + e • βO (i,•) ) ≥ - Ĉ∥c k ∥ 2 σ * S ∥d k ∥ 2 ((σ * S ) -1 α- 2 log( 5Km δ )σ1∥d k ∥) ⇒ E[e • β (t) O (i,•) ,k 1(i ∈ W e k,n (t) -U e k,n (t))] -e • β (0) O (i,•) ,k ≥ - Ĉ∥c k ∥ 2 σ * S 2 ∥d k ∥ 2 α - σ1(σ * S 2 ∥d k ∥ 2 + Ĉ∥c k ∥ 2 ) σ * S 2 ∥d k ∥ • 2 log( 5Km δ ).The proof is completed.Lemma 33. Under Condition 1, there exist constant C1 > 0, after at mostT = C1σ1mλK1γ (1 + κy) log(5Km/δ) (2σ * S -1) 2 (1 -κy)∥q∥ . iterations, we have L 0-1 D S (E(Ψ t )) = L 0-1 D * (E(Ψ t )) = 0.Proof. For t ≤ t, recall from Eq.( 53) that for the period t ≤ T , it holds thatE[A k,e t+1 | Ψ (t) ] ≥ A k,e t -( c2η(2σ * S -1) 2 (1 -κy)∥q∥ 2 16mK1 • E n∈V e k [ℓ ′ n (t) ]).Note that by definition A k,e 0 = 0, and we recursively use the equation t timesE[A k,e t ] ≥ t-1 s=0 - c2η(2σ * S -1) 2 (1 -κy)∥q∥ 2 16mK1 • E n∈V e k [ℓ ′ n (s) ].For each k ∈ [K1], e ∈ [±], denote by t k,e the last time in the period [0, T * ] satisfying that E[A k,e t ] ≤ κ. Then by Lemma 31 we see thatE n∈V e k [ef (E(S); E(Ψ (t) ))]] ≤ 3κ/2.Thus there exists a positive constant C such that -En∈V e k [ℓ ′ n (t) ] ≥ C for 0 ≤ t ≤ t k,e . Then we have E[A k,e t ] ≥ Cc2η(2σ * S -1) 2 (1 -κy)∥q∥ 2 t 16mK1 .Therefore we see that for ∀k ∈ [K1], e ∈ [±], E[A k,e t ] will reach κ within 16mK1κCc2η(2σ * S -1) 2 (1 -κy)∥q∥ 2 epochs. Recall that in this first stage the impact of decaying learning rate is under controlled by a large γ in Condition 1 as well as the slow quadratic decaying speed of ηt, under which we have η = Θ(η0). By κ ≤ 8σ1∥q∥ (1 + κy) log(5Km/δ)), we see that there exist a positive constant C1 = Θ(64/( Cc2)), the threshold time can beT = C1σ1mλK1γ (1 + κy) log(5Km/δ) (2σ * S -1) 2 (1 -κy)∥q∥ .Then by definition of 0-1 loss we haveL 0-1 D * (E(Ψ T )) = PS n ∼D * (yS n • f (E(Sn), E(Ψ T )) ≤ 0) ≤ PS n ∼D * (E[A k,e T ] -κ/2 ≤ 0) ≤ PS n ∼D * (E[κ/2 ≤ 0) = 0.The proof is completed.We claim that if i∈[m] riβ (0) O (i,•) ,k > 0 during initialization, the expected attention score will not experience this decreasing period due to 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. Lemma 34. Under Condition 1, for ∀k ∈ [K1], after at most a certain iterationsT1 = C3σ1K1γ 10 log(5Km/δ)(1 + e -2σ 2 0 ∥b k ∥ 2 ) 2C4∥d k ∥(1 -e -2σ 2 0 ∥b k ∥ 2 ), where C3 is a positive constant, we would have the β(t) Q,k = β (t)K,k be monotonically increasing during the remaining iterations T1 ≤ t ≤ T * . Besides, it holds that σ * S is the lower bound of the lowest correct attention assignment along the whole iterations:σ * S ≤ min t∈[T * ],k∈[K 1 ] { E n∈D S [ j∈S y Sn n,k (σ (t) S ) n j ]}.Proof. By Lemma 27, the β(t+1) Q,k = E[β (t+1) K,k | Ψ (t)] will be contributed to increase by{i ∈ E[W ± k,n (t)] | ri • β (t) O (i,•) ,k > 0}and they will be contributed to decrease by{i ∈ E[W ± k,n (t)] | ri • β (t) O (i,•) ,k < 0}By the fifth inequality in Lemma 7, we know that|{i ∈ E[W ± k,n (0)] | ri • β (0) O (i,•) ,k > 0}| - m 4 ≤ m 16 , |{i ∈ E[W ± k,n (0)] | ri • β (0) O (i,•) ,k < 0}| - m 4 ≤ m 16 . As E n∈V e k [ℓ ′ n (t) 1 n O (i) (t) ( j∈S + n,k (σ (t) S ) n j )( j∈S - n,k (σ (t) S ) n j )] is shared by all neurons, thus whether the βEQUATIONβ (t)K,k will be contributed to increase or decrease depends on the signal ofi∈E[W ± k,n (t)] ri • eβ (t)O (i,•) ,k . By the last inequality in in Lemma 7, we see that at initialization,EQUATION)By the expected gradient update in Lemma 28, the eβ(t)O (i,•) ,k will grow in ri's direction along the whole iterations. As such, the values ofE[ri • β (t) O (i,•) ,k 1(i ∈ W ± k,n (t))],∀k ∈ [K1] will grow larger. Therefore, after a limited epochs we can havei∈E[W ± k,n (t)] ri • eβ (t) O (i,•) ,k ≥ 0,where the β (σ(t) S ) n l )],which is clearly to be the first iteration where the negative i∈E[W ± k,n (t)] ri • eβ (t)O (i,•) ,k has grown to surpass the 0. By the symmetry property denoted in Lemma 22 and Eq.( 39) we haveEQUATION)Recall thatEQUATIONEQUATIONEQUATIONEQUATIONS ) n l )],and we also see thatEQUATIONEQUATION) As i∈E[W ± k,n (t)] ri • eβ (t)O (i,•) ,k will grow to surpass 0 in a limited number of iterations, we can claim that there exists a constant C ′ , such that for the limited decreasing period of E[( l∈S e n,k (σEQUATION)|] ≥ m/8 by Lemma 7, as well as the fact that E[W e k,n (t)] will at least preserve the neurons of E[U e k,n (0))] along the iterations, without being deactivated as discussed in Lemma 29. Also, we note that in this hypothesised decreasing period, the absolute value of the initially negativeE[ i∈[m] riβ (t) O (i,•) ,k ] and initially positive β (t)Q,k will all decreasing. Then by Condition 1 we see that the small initialization of MLP as well as the small regularization will make the decreasing order of βEQUATIONEQUATIONHere the second inequality is due to Eq.( 58), (56) and the definition of the b k in Eq.( 27); the third inequality is by the condition λ ≤ (Cσ0/2∥u∥ 2 ) -1 and σ1 ≤ (Cσ0∥u∥ 4 ∥q∥ log(5Km/δ)/K1) -1 . Therefore, it holds that during the decreasing period of βQ,k as well as the period where i∈E[W ± k,n (t)] ri • β (t) O (i,•) ,k remain negative, we have i∈E[W ± k,n (t+1)] ri • β (t+1) O (i,•) ,k ≥ i∈E[W ± k,n (t)] ri • β (t) O (i,•) ,k + C4η0∥d k ∥ 2 K1 (2 1 1 + e -2β (0) Q,k 2 /∥b k ∥ 2 -1),Here, by a appropriate chosen small C4, we again ignore the regularization term at this period due to λ = O((C log(Km/δ)∥q∥) -1 ) for a large C by Condition 1, and the impact of the learning rate is also controlled due to the slow quadratic decaying nature of η ′ t and a small initial η0 ≤ O(0.01C -1 ) by Condition 1, so as the changing amount of 1/(1 + e -2β(0) Q,k 2 /∥b k ∥ 2) by Eq.( 59).Therefore, by Eq.( 58) we haveEQUATIONEQUATIONwhere the inequality is by the negative nature ofi∈E[W ± k,n (t)] ri • β (t) O (i,•) ,k, and the decaying nature of ηt and Eq.( 58). Now we can see that there exists two surrogate sequences β(t) Q,k and i∈E[W ± k,n (t)] ri • eβ (t)O (i,•) ,k as the lower bound sequence of the β(t) Q,k and i∈E[W ± k,n (t)] ri • eβ (t)O (i,•) ,k . These two former sequences's initial values are taken as the lower bounds of the latter two (σ0∥b k ∥ 2 and -2 log(5Km/δ)EQUATION), and their update rule areEQUATION,EQUATIONO (i,•) ,k = i∈E[W ± k,n (t)] ri • eβ (t) O (i,•) ,k + C4η0∥d k ∥ 2 K1 (2 1 1 + e -2β (0) Q,k 2 /∥b k ∥ 2 -1).Then by Lemma 11, let a = C4η0∥b k ∥ 4 K1 , b = C4η0∥d k ∥ 2 K1 (2 1 1 + e -2β (0) Q,k 2 /∥b k ∥ 2-1), we have the maximumiterations T1 = -z(0)(1 + e -2y(0) 2 ) b(1 -e -2y(0) 2 ) = σ1K1γ 10 log(5Km/δ)(1 + e -2σ 2 0 ∥b k ∥ 2 ) 2C4∥d k ∥(1 -e -2σ 2 0 ∥b k ∥ 2 ) , set C3 = √ 10/(2C4)we obtain the T1 in the lemma. The lower bound of βQ,k along the decreasing period asβ Q,k = σ0∥b k ∥ 2 e -log(5Km/δ) 25σ 2 1 ∥b k ∥ 4 (1 + e -2σ 2 0 ∥b k ∥ 2 ) 1024(1 -e -2σ 2 0 ∥b k ∥ 2 ) ,Utilizing the scale bounding property (-κx + 1)/2∥u∥ 2 ≤ ∥b k 1 ∥ 2 < ∥u∥ 2 /2 in Eq. ( 27) and ( 36), we can denoted the lower bound of all β (t)Q,k = β (t)K,k for ∀k ∈ [K1] as β - QK , which can be given asβ - QK = σ0(1 -κx)∥u∥ 2 2 e -log(5Km/δ) σ 2 1 ∥u∥ 4 (1 + e -σ 2 0 ∥u∥ 2 ) (1 -e -σ 2 0 ∥u∥ 2 ) ,Recall σ * S is defined asσ * S := 1 1 + e -2 -1 σ 0 2 (1-κx) 2 ∥u∥ 4 e -2σ 1 2 log(5Km/δ) ∥u∥ 4 (1 + e -σ 2 0 ∥u∥ 2 ) (1 -e -σ 2 0 ∥u∥ 2 ), which is actually can be written asσ * S = 1 1 + e -2β - QK 2 .Therefore, we see that σ * S is the lower bound of min t∈EQUATIONn j ]}.Remark 4. As we see that in Lemma 33, we require that the lower bound given in Eq.( 53) depends that the values ofEQUATIONO (i,•) ,k )] surpasses κ, which naturally says that the value ofEQUATIONEQUATIONEQUATIONO (i,•) ,k ] should surpass 0 at T since κ > 0, which indicates that T > T1. We see that the initial period t ≤ T1 is whereEQUATIONO (i,•) ,k ] grow to surpass the initial scale, whose upper bound is κ/8 by the definition of κ.This period's analysis is based on Period 1 in Section I.1.2, or a good initialization such thati∈[m] riβ (0) O (i,•) ,k > 0.Lemma 35. Under Condition 1, consider the duration after T1 in Lemma 34, then for ∀k ∈ [K1], consider the periodT1 ≤ t ≤ T2 = C5 min{ 1+γ λ , ∥u∥∥q∥ λK 1 √ m }, where C5 is a small constant. Then the following holds that• We have y(t), y(t), z(t), z(t) be the lower and upper bounds of the increasing β(t) Q,k = β (t) K,k and i∈E[W ± k,n (t)] ri • eβ (t)O (i,•) ,k respectively. That is, there exists positive constants c3-6, forEQUATIONO (i,•) ,k ≤ z(t),for all t ≥ T1. Here, yare the unique solutions of the following ODE System respectively1 2 (Ei(2y(t) 2 ) + Ei(-2y(t) 2 ) + 4 log(y(t))) = abc ′2 (2σ * S -1) (t -t1) 2 2 + 1 2 (Ei(log( σ * S 1 -σ * S )) + Ei(log( 1 -σ * S σ * S ))) + 4 log(β - QK ), z(t) = bc ′ (2σ * S -1)(t -T1), 1 2 (Ei(2y(t) 2 ) + Ei(-2y(t) 2 ) + 4 log(y(t))) = abt 2 2 + a κ 8 t + 1 2 (Ei( σ 2 0 ∥u∥ 4 2 ) + Ei(-2 σ 2 0 ∥u∥ 4 2 )) + 4 log(σ0/2∥u∥ 2 ), z(t) = bt + κ 8 ,whereEi(x) = x -∞ e t t dt = γEuler + ln x + exp(x/2) ∞ n=1 (-1) n-1 x n n!2 n-1 ⌊(n-1)/2⌋ k=0 1 2k + 1 . • For some limited constant △ such that ∃△, σ * S < △ ≤ △ < 1. Then the β (t) Q,k = β (t)K,k will grow to make the correct attention scoreEQUATIONn j ] achieve the △ in at least a △(1 -△)scaled Gaussian rate such thatβ (t) Q,k ≥ exp( abc ′ △(1 -△)(2σ * S -1) 2 (t -T1) 2 + log(β - QK )).Proof. By Remark 4, we see that at the initial phase during t ≥ T1, we have i∈[m] riβO (i,•) ,k ≤ κ/8, and thus by Eq.( 55) in Lemma 30 we see that α ))]]| is also Θ(κ). This suggest that there still exists a constant C, during a certain amount of subsequent iterations we would still have thatO (i,•) ,k ≤ Ĉ ∥c k ∥ 2 σ * S ∥d k ∥ 2 . This indicates that E[A k,EQUATION)|] ≥ m/8 by Lemma 7, as well as the fact that E[W e k,n (t)] will at least preserve the neurons of E[U e k,n (0))] along the iterations, without being deactivated as discussed in Lemma 29. In addition, recall that in this first stage we also can control the impact of regularization and decaying learning rate by a small λ and a big γ by the sufficiently large C in Condition 1, which indicates we now haveEQUATIONO (i,•) ,k • ℓ ′ (t) 1 1 + e -2β (t) Q,k 2 /∥b k ∥ 2 1 1 + e 2β (t) Q,k 2 /∥b k ∥ 2 ] ,andEQUATIONO (i,•) ,k ] = i∈E[W ± k,n (t)] ri • β (t) O (i,•) ,k + Θ C4η0∥d k ∥ 2 K1 E[ℓ ′ (t)(2 1 1 + e -2β (t) Q,k 2 /∥b k ∥ 2 -1)] .By Lemma 12, we see that the iteration satisfies the ODE System 2 with a positive initialization, where the parameters in Lemma 12 areℓ ′ t = -E[ℓ ′ (t)], a = Θ( C4η0∥b k ∥ 4 K1 ), b = Θ( C4η0∥d k ∥ 2 K1 ), c ′ = C.Then by solving the coupled ODE systems, collaborating the scale bounding property 27) and ( 36), as well as the Comparison Theorem with some constants c3-6, we can have upper and lower bound of(-κx + 1)/2∥u∥ 2 ≤ ∥b k 1 ∥ 2 < ∥u∥ 2 /2, -κy + 1/2∥q∥ 2 ≤ ∥d k 1 ∥ 2 < ∥q∥ 2 /2 in Eq. (β (t) Q,k and i∈E[W ± k,n (t)], which is the result in this lemma.For the second result, given the △, we can directly have a lower bound ODE y △ (t) to be the lower bound of the β(t) Q,k via Comparison Theorem, where y △ (t) satisfies y ′ (t) ≥ y ′ △ (t) = abc ′ △(1 -△)(2σ * S -1))(t -T1)y △ (t) ⇒ y(t) ≥ y △ (t) = exp( abc ′ △(1 -△)(2σ * S -1) 2 (t -T1) 2 + log(β - QK )),where the inequality holds by the decaying nature of g(x) = x(1 -x) when x > 1/2. The proof is completed.Lemma 36. (Asymptotic Property 1). In the first stage, the growing of β(t) Q,k = β (t)K,k as well as the attention score enjoys the asymptotic property thatEQUATIONS ) n j ]] t 4 1 + t 4 = Θ(1).Proof. By the asymptotic property of Ei(x)lim x→+∞ Ei(x) + Ei(-x) exp(x) x = 1.This suggest that when y(t) ≥ y(t) is close to infinity, the lower bound ODE in Lemma 35 will approximately satisfies the followingexp(2y(t) 2 ) 4y(t) 2 + 2 log(y(t)) ≈ abc ′2 (2σ * S -1) t 2 2 + const,This suggest that roughly lim t→+∞ y(t) 2 / log(t 2 ) = Θ(1).Then we see that as y(t) goes to infinity, we have a lower boundEQUATIONS ) n l ] 1 -E[ l∈S e n,k (σ (t) S ) n l ] )/2 log(t 2 ) = Θ(1) ⇒ lim t→+∞ E[ l∈S e n,k (σ (t) S ) n l ]/ t 4 1 + t 4 = Θ(1).On the other hand, obtaining an upper bound over y(t) is relatively easy. Since we have 2+e -2y(t) 2 +e 2y(t) 2 ≤ 4 and (1 -e -2y(t) 2)/(1 + e -2y(t) 2 ) ≤ 1, which gives the upper bound ODE over attention and MLP considering z(0) > 0 1 2 (Ei(2y(t) 2 ) + Ei(-2y(t) 2 ) + 4 log(y(t))) = abt 2 2 + a κ 8 t + const. z(t) = bt + const,where the term "const" ensure that y(0) = y(0). The asymptotic property of this ODE system is the same as the one of lower bound ODE. Then consider t, y(t) both go to infinity, we have some ĉ such thatEQUATIONEQUATIONS ) n l ]/ t 4 1 + t 4 = 1.As the β(t) Q,k = β (t) K,k and eβ (t)O (i,•) ,k are continually growing up, we see that the decaying -E[ℓ ′ (t)], as well as the decaying attention score products( j∈S + n,k (σ (t) S ) n j )(1 -j∈S + n,k (σ (t) S ) n j) 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.Lemma 37. Under Condition 1, consider e = E[yS n ] for all t ∈ [T2, T * ] it holds that eβ (t) O (i,•) ,k = O( ∥q∥ 2 λmK1),EQUATION)),EQUATION))),EQUATION),EQUATIONwhere eβ(t) O (i,•) ,k represents mriβ (t) O (i,•) ,k . That is, we consider the positive growth of E[β (t) O (i,•) ,k ].Proof. We will prove the desired argument based on the following induction hypothesis:eβ (t) O (i,•) ,k = O( ∥q∥ 2 λmK1),β (t) Q,k = O(∥u∥ log( 2∥u∥ 2 ∥q∥ 2 λ 2 mK1 2 )),We split the situations into two cases:(i). eβ (t) O (i,•) ,k ≤ Θ( σ * S 2 (1-κy ) 2 (1+κy ) 2 log( e -κ ∥q∥ 2 (2σ * S -1) λmK 1)), and βEQUATIONEQUATIONS ) n j )] ≥ 1 1+ λ 2 mK 1 2 6∥u∥ 2 ∥q∥ 2. Easy to note that the scales' orders of the case (i)'s quantities are less than those of case (ii), thus this split is plausible.RecallEQUATIONS ) n l )].Then it's easy to check that for case (i), as by Lemma 30 we see that the magnitude of E[|α (t)O (i,•) ,k |] is controlled by some Ĉσ * S -2 (1 + κy) 2 (1 -κy) -2 β (t) O (i * ,•) ,k ≥ Ĉ. That means that the term E[A k,et ] can be controlled by its contributor Θ(eβ(t) O (i,•) ,k ). Then we have Θ(eβ (t) O (i,•) ,k /E[-ℓ ′ n (t) ]) ≤ Θ(eβ (t) O (i,•) ,k (1 + e κ e σ * S -2 (1+κy ) 2 (1-κy ) -2 eβ (t) O (i,•) ,k )) ≤ Θ(e κ e σ * S -2 (1+κy ) 2 (1-κy ) -2 eβ (t) O (i,•) ,k ) ≤ Θ( ∥q∥ 2 (2σ * S -1) λmK1 ),where the first inequality is by the definition of E[-ℓ ′ n (t) ]) (similar to the techniques in Lemma 32) and the definition of E[A k,e t ]; the second inequality is by g(x) = x < e x -1 as well as σ * S -2 (1 + κy) 2 (1 -κy) -2 > 1; the second inequality is by the case (i) hypothesis. Then we would haveEQUATIONThus the growing of eβ EQUATIONEQUATIONO (i,•) ,k ), which directly suggest E[β (T * ) O (i,•) ,k ] = Θ( σ * S 2 (1-κy ) 2 (1+κy ) 2 log( e -κ ∥q∥ 2 (2σ * S -1) λmK 1)) holds since T * is the maximum admissible iterations.Similarly, for the β(t) Q,k in case (i), first recall that β (t+1) Q,k =(1 -ηtλ)β (t) Q,k - 4ηtβ (t) K,k ∥b k ∥ 4 K1 e∈[±] i∈[m] riβ (t) O (i,•) ,k E n∈V e k [ℓ ′ n (t) 1 n O (i) (t) ( j∈S + n,k (σ (t) S ) n j ) (1 - j∈S + n,k (σ (t) S ) n j )],then, as we see thatE[( j∈S + n,k (σ (t) S ) n j )(1 - j∈S + n,k (σ (t) S ) n j )] = 1 1 + e -2β (t) Q,k 2 /∥b k ∥ 2 1 1 + e 2β (t) Q,k 2 /∥b k ∥ 2 = 1 2 + e 2β (t) Q,k 2 /∥b k ∥ 2 + e -2β (t) Q,k 2 /∥b k ∥ 2 ≤ Θ( 1 2 + ∥u∥ 2 σ * S 2 (1-κy ) 2 λK 1 (1+κy ) 2 log( e -κ ∥q∥ 2 (2σ * S -1) λmK 1 ) ) = Θ( ∥u∥ 2 σ * S 2 (1 -κy) 2 λK1(1 + κy) 2 log( e -κ ∥q∥ 2 (2σ * S -1) λmK1 )) -1 ,where the first inequality is by the definition ofEQUATIONS ) n j )] and the induction hypothesis; the second inequality is by the small λ by Condition 1 with a sufficiently large C. Then we see thatΘ(- 4∥b k ∥ 2 K1 E[ e∈[±] i∈[m] riβ (t) O (i,•) ,k E n∈V e k [ℓ ′ n (t) 1 n O (i) (t) ( j∈S + n,k (σ (t) S ) n j )( j∈S - n,k (σ (t) S ) n j )]]) ≥ Θ(λ).Here the inequality is by the case (i) hypothesis upon eβ (t)EQUATIONEQUATION/m] ≤ 2.Thus we see that the growing of β (t)Q,k would also be non-degenerated:β (t+1) Q,k ≥ Θ(β (t) Q,k). This also directly validates that for the maximum admissible iterations T * , it holds thatEQUATION))),E[( j∈S e n,k(σ(T * ) S ) n j )] = Θ( 1 1 + λK 1 (1+κy ) 2 ∥u∥ 2 σ * S 2 (1-κy ) 2 log -1 ( e -κ ∥q∥ 2 (2σ * S -1) λmK 1 )).For case (ii), we directly check thatλeβ (t) O (i,•) ,k ≥ - ∥d k ∥ 2 2K1 e∈[±] E[ri E n∈V e k [ℓ ′ n (t) 1 n O (i) (t) ( l∈S e n,k (σ (t) S ) n l - l∈S -e n,k (σ (t) S ) n l )]].Here the inequality is by -E[ℓ ′ (t)] ≤ 1 andE[( l∈S e n,k (σ (t) S ) n l - l∈S -e n,k (σ (t) S ) n l )] = [E(2 l∈S e n,k (σ (t) S ) n l -1)] ≤ 1.As a result, by the gradient form we see that β O (i,•) ,k ] = O( ∥q∥ 2λmK 1 ). Similarly, as we now haveE[( j∈S + n,k (σ (t) S ) n j )(1 - j∈S + n,k (σ (t) S ) n j )] = 1 1 + e -2β (t) Q,k 2 /∥b k ∥ 2 1 1 + e 2β (t) Q,k 2 /∥b k ∥ 2 = 1 2 + e 2β (t) Q,k 2 /∥b k ∥ 2 + e -2β (t) Q,k 2 /∥b k ∥ 2 ≥ 1 3 1 1 + e 2β (t) Q,k 2 /∥b k ∥ 2 /3 ≥ 1 3 1 1 + 2∥u∥ 2 ∥q∥ 2 λ 2 mK 1 2 , = Θ( λ 2 mK1 2 2∥u∥ 2 ∥q∥ 2 ),where the first inequality is by e-2β (t) Q,k 2 /∥b k ∥ 2≤ 1; the second inequality is by the induction hypothesis of case (ii); the last equality is by the small λ in Condition 1 for a sufficiently large C. Then we observe thatλ ≥ Θ(- 4∥b k ∥ 2 K1 E[ e∈[±] i∈[m] riβ (t) O (i,•) ,k E n∈V e k [ℓ ′ n (t) 1 n O (i) (t) ( j∈S + n,k (σ (t) S ) n j )( j∈S - n,k (σ (t) S ) n j )]]),where the inequality is byEQUATION/m] ≤ 2, as well as the induction hypothesis upon eβ (t)O (i,•) ,k in case (ii). Thus β (t+1) Q,k ≤ Θ(β (t) Q,k), which support our proving goal in this induction process:E[β (t+1) Q,k ] = O(∥u∥ log( 2∥u∥ 2 ∥q∥ 2 λ 2 mK1 2 )).In addition, we can see that even if we suggest the MLP's eβ .O (i,•) ,k isTherefore, we can have some appropriately small constants C5, and when the iteration number is more thanT2 = C5 min{ 1 + γ λ , ∥u∥∥q∥ λK1 √ m}, we need to consider the impact of regularization.Lemma 38. The scale of the coefficients will finally be stabilized at a considerable level:EQUATION)),β (T * ) Q,k = Θ(∥u∥ log( ∥u∥ 2 λK1 log( ∥q∥ 2 mλK1 ))), E[( j∈S e n,k(σ(T * ) S ) n j )] = Θ( 1 1 + λK 1 ∥u∥ 2 log( mλK 1 ∥q∥ 2 )).where eβ (t) O (i,•) ,k represents mriβ (t) O (i,•) ,k .EQUATIONBesides, we see that the expected attention score will grow up considerably, where we can see that there exist some constant c > 1/2, c < E[(σ EQUATIONO (i,•) ,k ]).Again, similar to Lemma 36, we can have asymptotic property when considering the decaying impact of the learning rate, as well as the cross-entropy loss. We directly provide the following two lemmas. Due to the similarity of the proof procedures of Lemma 35 and Lemma 36, we omit the proofs of the following two lemmas as well as the constant details for simplicity. Lemma 39. (Asymptotic Property 2). If we consider the impact of the decaying learning rate at the second stage and do not consider the decaying of cross-entropy loss, for some constants c, d, c, d regarding K1, γ, ∥u∥, ∥q∥, κx, κy, we will have 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, c, d, j, a, b, c, d, j, we will havey(t) ≤ β (t) Q,k = β (t) K,k ≤ y(t), z(t) ≤ i∈E[W ± k,n (t)] ri • eβ (t) O (i,•) ,k ≤ z(ty(t) ≤ β (t) Q,k = β (t) K,k ≤ y(t), z(t) ≤ i∈E[W ± k,n (t)] ri • eβ (t) O (i,•) ,k ≤ z(t),for all t ≥ T2. Here, y(t), y(t), z(t), z(t) are the unique solutions of the following ODE System respectively 1 2 (Ei(2y(t) 2 ) + Ei(-2y(t) 2 ) + 4 log(y(t))) = a(EQUATION)) + 4 log(σ0/2∥u∥ It's obvious that the decaying impact of the learning rate and cross-entropy loss are at the similar order. Also, if we consider decaying learning rate, the right side of the inequality would be smaller. z(t) would be in a Θ(log(log(t))) order when z(t) get large, which will make the right side of the y(t)'s formula contain an intergral of Θ(log log(t)), which is obviously slower.We continue our proof after Lemma 33. In this section, we assume all the events in the Section D hold, denoted as ΥPre.Lemma 41. The Frobenius norm of W y O and its gradient can be bounded:∥W y O ∥ 2 F = O( K1∥q∥ 2 λ 2 m ), ∥∇ W y O (t) LB t (Ψ (t) )∥ 2 F = O( K1∥q∥ 2 m ).Proof. For ∀i ∈ [m], by the gradient update rule in Eq.( 22), as well as Lemma 4's insight we see that the lengths of the W y O on certain projection direction will continue to grow until being stuck by the regularization, which is a λ-scaled W y O itself. Due to the low-noise condition in Condition 1 with a sufficiently large C as well as the isotropy of noise, the learning progress of features would be the main contributor to the F norm of NN matrices and the noise, validated in Figure 2 (iii-iv). We can consider an extreme case where all the samples in a single batch belongs to some concept k ∈ [K1], which we can have the upper bound of the first term of the right side of the inequality over the k-th concept's corresponding projection direction, and thus we can derive an upper boundEQUATIONwhere the first inequality is by O is controlled by the K1 concepts' corresponding lengths in projection space. Then by the definition of Frobenius norm and Eq.( 22) we haveEQUATION∥W y O ∥ 2 F ≤ Θ( K1∥q∥ 2 λ 2 m ), ∥∇ W y O (t) LB t (Ψ (t) )∥ 2 F ≤ Θ( K1∥q∥ 2 m ). Lemma 42. For ∀ k ∈ [K1], a ⊤ k W x Q (t) a k , a ⊤ k W x K (t) a k and b ⊤ k W x Q (t) b k , b ⊤ k W x K (t) b k satisfy a ⊤ k W x Q (t) a k , a ⊤ k W x K (t) a k = O(σ0∥u∥ 2 ), b ⊤ k W x Q (t) b k , b ⊤ k W x K (t) b k = O(∥u∥ log( (L -1)∥u∥ 2 ∥q∥ 2 λ 2 m )).Moreover, the Frobenius norm of W x Q (t) and W x K (t) and its gradient can be bounded as below∥W x Q (t) ∥ 2 F , ∥W x K (t) ∥ 2 F = O(K1 log( (L -1)∥u∥ 2 ∥q∥ 2 λ 2 m )) ∥∇ W x Q (t) LB t (Ψ (t) )∥ 2 F , ∥∇ W x K (t) LB t (Ψ (t) )∥ 2 F = O( K1(L -1)∥u∥ 2 ∥q∥ 2 m ).Proof. By Eq.( 16) and ( 17), we see thatEQUATIONK,a k ,chaos |, |I(t) K,a k ,contri | ≤ ηtΘ(max{|a ⊤ k W x Q (t) a k |, |a ⊤ k W x K (t) a k |}(∥u∥ σ ξ 2 log( KN δ ) + 1 K ∥u∥ 2 ) 2 ( ∥q∥/λm m )) ≤ O(λ max{|a ⊤ k W x Q (t) a k |, |a ⊤ k W x K (t) a k |}).Here, the first inequality is by the scaled identity initialization of W x Q (0) , W x K (0) ,⊤ k W x Q (t) a k and a ⊤ k W x K (t)a k are dominated by their regularization, and thus the scale can not be better than the initialization. By Lemma 7, the conclusion holds.On the other hand, we see that by the scaled identity initialization of W x Q (0) , W x K (0) and orthogonal relationships of vectors in Lemma 27, the initialization of b⊤ k W x Q (t) b k and b ⊤ k W x K (t)b k are the same, and as the gradient update is nearly symmetry, which can lead to the fact thatΘ(b ⊤ k W x Q (t) b k ) = Θ(b ⊤ k W x K (t) b k ) and b ⊤ k W x Q (t) W x K (t) b k = Θ(b ⊤ k W x Q (t) b k b ⊤ k W x K (t) b k /∥b k ∥ 2 ). By the scaled identity initialization ofW x Q (0) , W x K (0), orthogonal relationships of vectors in Lemma 27, Eq.( 19) and ( 20) we can see that t) b k will continue to grow up except always being stuck by the regularization. To see the upper bound under this situation, we consider an extreme case where all the samples in a single batch belongs to some concept k ∈ [K1], and there is only one demonstrations in each prompt share the semantic with the query. Then by the scaled identity initialization of W x Q (0) , W x K (0) , orthogonal relationships of vectors in Lemma 27 and Eq.( 18), we can see that the growing of bEQUATION⊤ k W x Q (t) b k and b ⊤ k W x K (t)b k would satisfy the following and strive to grow up to make the equality holds, which naturally have an upper bound|I (t) Q,b k ,contri |, |I (t) K,b k ,contri | ≥ λ min{|b ⊤ k W x Q (t) b k |, |b ⊤ k W x K (t) b k |} ⇒ ∥u∥ 2 ∥q∥ 2 λm ( j∈S + n,k (σ (t) S ) n j )( j∈S - n,k (σ (t) S ) n j ) ≥ Θ(λ) ⇒ Θ( ∥u∥ 2 ∥q∥ 2 λ 2 m j∈S + n,k e b ⊤ k W x Q (t) W x K (t) b k j∈S - n,k e -b ⊤ k W x Q (t) W x K (t) b k ( j∈S + n,k e b ⊤ k W x Q (t) W x K (t) b k + j∈S - n,k e -b ⊤ k W x Q (t) W x K (t) b k ) 2 ) ≥ 1 ⇒ Θ( ∥u∥ 2 ∥q∥ 2 λ 2 m L -1 (e b ⊤ k W x Q (t) W x K (t) b k + (L -1)e -b ⊤ k W x Q (t) W x K (t) b k ) 2 ) ≥ 1 ⇒ Θ( (L -1)∥u∥ 2 ∥q∥ 2 λ 2 m e -2b ⊤ k W x Q (t) W x K (t) b k ) ≥ 1 ⇒ b ⊤ k W x Q (t) b k , b ⊤ k W x K (t) b k ≤ Θ(∥u∥ 1 2 log( (L -1)∥u∥ 2 ∥q∥ 2 λ 2 m ))Here, the second arrow is by the definition of j∈SEQUATIONS ) n j and Eq.( 31); the third arrow is by our considered extreme case where there is only one demonstration in each of the prompt sample in this all-the-same-concept batch, which is considered for obtaining the upper bound; the forth arrow is by the small λ by Condition 1,which denotes e -2b ⊤ k W x Q (t) W x K (t) b k should be the key contributor.Similar to the claims in Lemma 41, here we see that as the λ is very small, by the scaled identity initialization ofW x Q (0) , W x K (0), orthogonal relationships of vectors in Lemma 27, as well as the low-noise condition by Condition 1, it's safe to say that as the learning proceed, the scales of b⊤ k W x Q (t) b k , b ⊤ k W x K (t) b k would completely dominate a ⊤ k W x Q (t) a k , a ⊤ k W x K (t) a k , as well as a ⊤ k W x X (t) b k , b ⊤ k W x X (t) a k , ν ⊤ r W x X (t) νr, u ⊥ w ⊤ W x X (t) u ⊥ w , ∀X ∈ {Q, K}, r ∈ [K2], w ∈ [dX -2K1 -K2].Collaborating with Lemma 9, we have∥W x Q (t) ∥ 2 F , ∥W x K (t) ∥ 2 F ≤ K1 ∥u∥ 2 max{(b ⊤ k W x Q (t) b k ) 2 , (b ⊤ k W x K (t) b k ) 2 } = O(K1 log( (L -1)∥u∥ 2 ∥q∥ 2 λ 2 m )).On the other hand, we see that the maximum gradient F norm on a single batch comes from the maximum changes of the b⊤ k W x Q (t) b k 2 (or b ⊤ k W x K (t) b k 2). As we see that the extreme case of the growing is every concept k ∈ [K1] has been fully learned 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 ∥λWx Q (t) ∥F (or ∥λW x K (t) ∥F ). Thus ∥∇ W x Q (t) LB t (Ψ (t) )∥ 2 F , ∥∇ W x K (t) LB t (Ψ (t) )∥ 2 F = O(λ 2 K1 log( (L -1)∥u∥ 2 ∥q∥ 2 λ 2 m )) = O( K1(L -1)∥u∥ 2 ∥q∥ 2 m ),where the equality is by g(x) = log(x) ≤ O(x), x > 1.The proof is completed.Remark 5. Worth noting that this upper bound, as well as the upper bound of ∥W yO (i,•) (t) q ± k /∥q ± k ∥∥ in Lemma 41, are looser in the order of K1 -2 and K1 -1 compared to those of β (t) Q,k = β (t) K,k and eβ (t) O (i,•) ,k in Lemma 4.This suits the intuition and statistics since in the practical training setting, for B ≥ 1 we can see that sometimes the samples of a batch all belong to one concept, or sometimes their are not any particular concept in 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 sample appear in every batch scaled by a "soft weight" in the order of Θ(1/K). As the attention 's gradient contain MLP, its order would be Θ(1/K 2 ). Besides, we see that this lemma's result contains the scale of L -1, which comes from the extreme case discussion where there is only one demonstration in each prompt sample that share the semantic to those of query. In contrast, 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 -ℓ ′ to 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β EQUATIONO (i,•) ,k . This actually indirectly demonstrates the superiority of considering expectations.2) ∀t ≥ T , when ∥Ψ ′(t) -E(Ψ ′(t) )∥F ≤ ν holds, we have L 0-1 D * (Ψ ′(t) ) = L 0-1 D * (E(Ψ ′(t) )). Here, ∥Ψ ′ ∥ 2 F := ∥W x Q ∥ 2 F + ∥W x K ∥ 2 F + ∥W y O ∥ 2 F .Proof. By Lemma 33, we see that our convergence of 0-1 loss is based on the intermediate result thatEQUATIONO (i,•) ,k , where would consequently cause E V e k [A k,e t ] ≤ κ/2 that could have potential to deteriorate the 0-1 loss. Given that κ/2 ≥ √ 2σ1∥q∥ by Lemma 23, the decomposition in Eq.( 33) as well as Lemma 9, we see that for some k ∈ [K1], the minimum admissible disparity can be written asΘ(∥( √ 2σ1∥q∥) 2 ( c k ⊤ ∥c k ∥ 2 ) ⊤ c k ⊤ ∥c k ∥ 2 ∥F ) = Θ( √ 2σ1∥(∥q∥) 2 ( c k ⊤ ∥c k ∥ 2 ) ⊤ c k ⊤ ∥c k ∥ 2 ∥F ) ≥ Θ(2 √ 2/(1 + κy)σ1).Therefore, we see that when conditioned onE[W x Q (t) ], E[W x K (t)], the minimum admissible disparity betweenW y O (t) and E[W y O (t) ] to not worsen the 0-1 loss is Θ(2 √ 2/(1 + κy)σ1).On the other hand, when conditioned on E[W y O (t) ], t ≥ T ′ , we compute the minimum admissible dispar-ity between W x Q (T ′ ) , W x K (T ′ ) and E[W x Q (T ′ ) ] = E[W x K T ].Considering all the activated neurons, wheni∈W e k,n (t) E[ri(2 l∈S e n,k (σ ( T ) S ) n l -1)β ( T ) O (i,•) ,k ] = 0, we should have i∈W e k,n (t) E[riα ( T )O (i,•) ,k ] ≥ 0 otherwise some of the neurons must be deactivated, which is contradicted by the definitions of W e k,n (t). In this case we can magnify the impact of i∈WEQUATION) S ) n l -1)β ( T ) O (i,•) ,k ] by consider- ing i∈W e k,n (t) E[riα ( T ) O (i,•) ,k ] = 0.As such, the minimum admissible disparity would be the case where b . RecallEQUATIONν := min{2 √ 2σ1/(1 + κy), σ0(1 -κx)e -log(5Km/δ) σ 2 1 ∥u∥ 4 (1+e -σ 2 0 ∥u∥ 2 ) (1-e -σ 2 0 ∥u∥ 2 )}, the proof is completed.Lemma 44. For t ∈ {1, • • • , T }, for W ∈ {W x Q , W x K , W y O } and X ∈ {Q, K, O} it follows that 1. ∥W (t+1) -Wt (t+1) ∥F ≤ Θ( K 1/2 1 ∥q∥((L -1) 1/2 ∥u∥ + 1) m 1/2 ηt), 2. ∥W (s+1) -Wt (s+1) ∥F ≤ (1 -ηsλ)∥W (s) -Wt (s) ∥F , ∀s ≥ t + 1, 3. T t=0 ∥D t X ∥ 2 ∞ ≤ Θ( K1∥q∥ 2 ((L -1)∥u∥ 2 + 1) mλ 2 (γ + T ) ).Proof. We provide the proof by extending the techniques in [34, 33, 36] to Hilbert-Schmidt space, whose inner product is defined by trace. First we note that η0 = 2 γ+1 ≤ min{1/(LLogist + λ), 1/2λ}, where LLogist is the L-smooth Lipschitz constant of cross-entropy loss ℓ(•), which is 1. The first statement can be shown as follows.Since by definition we see that W (t) = Wt (t) , we only need to check the maximum disparity of the gradient in a single iteration update, then by Lemma 41 and Lemma 42 we readily obtain the results.For the second statement, following the proof in [93, 34] , we see that the Lipschitz smoothness of cross-entropy loss denotes thatEQUATIONThen we have that for s ≥ t + 1,∥W (s+1) -Wt (s+1) ∥ 2 F = (1 -ηsλ) W (s) -Wt (s) -ηs ∂gl (gs, Zs) -∂gl g t s , Zs 2 F = (1 -ηsλ) 2 W (s) -Wt (s) 2 F -2ηs (1 -ηsλ) • ∇ W (s) LB s (Ψ s ) -∇ W t (s) LB s (Ψ s ), W (s) -Wt (s) + η 2 s ∥∇ W (s) LB s (Ψ s ) -∇ W t (s) LB s (Ψ s )∥ 2 F ≤ (1 -ηsλ) 2 W (s) -Wt (s) 2 F -ηs 1 LLogist -ηs • ∇ W (s) LB s (Ψ s ) -∇ W t (s) LB s (Ψ s ) 2 F ≤ (1 -ηsλ) 2 W (s) -Wt (s) 2 Fwhere we utilize the Eq. ( 63) and conditions on learning rates. Utilizing this statement, the stable property of stochastic gradient descent has been shown. Again following the techniques in [34, 33, 36] , we now obtain the bound: for t ∈ {1, . . . , T },EQUATIONFrom the following inequality,T s=t+1 (1 -ηsλ) = T s=t+1 γ + s -2 γ + s < γ + t γ + Twhere the last inequality hold clearly by expanding the product, the right hand side of the Eq.( 64) is upper bounded as followsΘ( K 1 21 ∥q∥((L -1)1 2 ∥u∥ + 1) m 1 2 )ηt T s=t+1 (1 -ηsλ) ≤ Θ( K 1 21 ∥q∥((L -1)1 2 ∥u∥ + 1) m 1 2 ) ηt(γ + t) γ + T = Θ(2 K 1 21 ∥q∥((L -1)1 2 ∥u∥ + 1) m 1 2 ) λ(γ + T ) .We finally obtain the desired bound:T t=0 ∥Dt∥ 2 ∞ ≤ T t=0 Θ( K1∥q∥ 2 ((L -1)∥u∥ 2 + 1) mλ 2 (γ + T ) 2 ) ≤ Θ( K1∥q∥ 2 ((L -1)∥u∥ 2 + 1) mλ 2 (γ + T ) ).Remark 6. Utilizing this lemma, the exponential convergence over the 0-1 loss is readily obtained. x × Dy * pair can appear in at least ∥z∥0 conceptspecific prompts/tasks. Importantly, all of the tasks in this new distribution D * enjoy Bayes-Optimal test errorL 0-1 D * (Ψ (T * ) ) ≤ ε.This lemma demonstrate the strong OOD Generalization ability of transformer utilizing multi-concept semantics, suggesting the efficiency transformer to conduct unseen ICL tasks just by its learned knowledge on the two non-orthogonal dictionaries. Also, this lemma showcases an intriguing phenomenon since it 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 / task) information. For instance, with the prompt "Japan: Sakura; China:", the LLM may output "Penoey" (national flower) or "Panda" (national symbol), reflecting different conceptual (task) interpretations. Both answers are 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 be "national flower". In our theory, we make an elementary explanation to this flexible, context-sensitive in-context learning (ICL) behavior by attributing it to the transformer's ability to harness multi-concept semantics. = k ′ ∈ [K1], k2 ∈ [K2], we let M * 2k-1 = µ + k * = a * k + b * k , M * 2k = µ - k * = a * k -b * k , Q * 2k-1 = q + k * = c * k + d * k , Q * 2k = q - k * = c * k -d * k , M * k 2 +2K 1 = ν k 2 * , Q * k 2 +2K 1 = 0, where a * k ∈ conic({ µ + k + µ - k 2 } K 1 k=1 ), b * k ∈ conic({ µ + k -µ - k 2 } K 1 k=1 ), c * k ∈ conic({ q + k + q - k 2 } K 1 k=1 ), d * k ∈ conic({ q + k -q - k 2 } K 1 k=1 ), ν k 2 * ∈ (span(µ + 1 , µ - 1 , µ + 2 , µ - 2 , • • • , µ + K 1 , µ - K 1 )) ⊥ , satisfying ∥b * k ∥ ≥ ∥a * k ∥ = Θ(∥u∥), ∥d * k ∥ ≥ ∥c * k ∥ = Θ(∥q∥), ν * k 2 = Θ(∥u∥), and {a * k , b * k } K 1 k=1 , {c * k , d * k } K 1 k=1are two collections of pair wise orthogonal vectors. Then we can have a corresponding new prompt distribution D* S = K 1 k=1 π + k * P + k,L * +1 * + π - k * P - k,L * +1 * . Again, the model enjoys Bayes-Optimal test error L 0-1 D * (Ψ (T * ) ) ≤ ε.This lemma suggest that transformer-mlp structure empower ICL ability in solving task involving semantics ("knowledge") originally from other co-concept prompt'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:Today I designed a machine to capture sunlight; Thomas Edison:" to GPT o1, 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-appear with the concept "Fabricate a story" with high chance, the transformers empower the ICL to perform this interesting Out-of-Distribution task.We believe this can serve as an attempt to explain the innovation power of LLM [30, 94, 95] grounded in the linear geometric property of LLM representation, since most of the innovative outcomes of human being generates from cross-concept "Knowledge Intersection", and as it is not an easy task for human specialist to master cross-domain knowledge, we claim that LLM can help innovation by leveraging cross-domain knowledge when deduction over unseen structured task. Similarly, for multi-model scenarios, [86] have shown that compositing different concepts did enable OOD generalization (e.g. "blue square apples" in the Figure 1a in [86] ).This lemma seeks to elementarily 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 [96, 4] , heuristics [3] , acquisition functions [97] , and solutions to combinatorial optimization problems [98] . 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 of 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 of Proposition 1. By Proposition 2, we only need to check the expected 0-1 loss L 0-1 D * (E[Ψ ′ ]) = 0. Denote E[My Sn ] ⊆ [2K1] 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[My Sn ] will be either a collection of odd (corresponding to positive label) or even (corresponding to negative label) numbers, and all of the combination of the features and labels in one prompt are corresponding to the same real value label without "self-conflict". By Lemma 38, we see that the coefficients are all at a substantial scale at T * . Then by the condition on z and Eq. ( 31), 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 • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or 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.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is 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 in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way 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 a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a 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 a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.• 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, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars 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 or tables in the text.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, we here chose a feasible one.Here we do not consider the shift of D ξx , D ξy for the ease of presentation. However, we assert that this can also be addressed by leveraging high-dimensional statistical analysis over other well-behaved noise distributions.The requirement of D * z could be relax with a stricter requirement on L * and a denser analyses.+ e -2y t 2 + e 2y t 2 ℓ ′ t , yt 1 > 0, a > 0, zt+1 = zt + b 1 -e -2y(t) 2 1 + e -2y(t) 2 ℓ ′ t , zt 1 ≥ 0, b > 0, 


Decision: 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.

Reviews: review: paper_summary: paper_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.
The 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.
strength_weakness: strength_weakness: The paper is sometimes hard to follow. More concise and clearer presentation, especially the theoretical content, would be helpful.
Lack of empirical experiments with more realistic datasets to show relevance in real-world setting.
questions: questions: NA
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 

a) Definition 4 states that 'with prob 1 - $\delta$, $L_{D*}(\Phi) = 0$'. What is the source of randomness here? 

b) 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.

2. 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.
questions: 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: limitations: The limitation is adequately addressed.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: questions: I think the authors should improve the presentation of their theory by giving more insights and intuition + simplified proof sketch.
limitations: limitations: The authors did not mention the limitations of their work.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
- 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.
questions: questions: - How sensitive are the results to the specific data distribution assumptions? Would the exponential convergence hold under more general conditions?
- Have you considered extending the analysis to training all weights in the network? What challenges would this introduce?
limitations: 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.

Comments: title: Grateful Response to Reviewer's Positive Feedback
Post rebuttal response
Thank you for your recognition
Thank you
Thank You for Your Feedback and Our Planned Revisions
Thank you and raise confidence score from 3 to 4
Clarification on Definition 4 and the Source of Randomness
comment: Dear Reviewer L2Bx,

We're delighted that you've found our rebuttals have met your satisfaction, and we're encouraged that you've raised the score. We're happy that you found the attached Figure 1 useful, and we would be sure to include it appropriately in the main body of our manuscript.

Once again, we appreciate your invaluable time and comments.

Warmest regards,

Author of Submission 15661
I thank the reviewers for their rebuttal that has addressed most of my concerns. I appreciate the diagram (Figure 1) that allows to get a better sense of the proof sketch and I hope it will be included to the paper. I increase my score by one point.
Dear Reviewer qjJG,

We're greatly encouraged that our rebuttals have been effective in addressing all your core concerns and that you have raised the scores accordingly. We are honored by your recognition and look forward to incorporating the necessary changes.

Thank you for your invaluable time and comments throughout this process.

Warmest regards,

Authors of Submission 15661
I thank the authors for their response. The authors have edit the manuscript to address my core concerns, I have increased the score accordingly.
Dear Reviewer Bq3t,

Thank you for your positive feedback. We are delighted and encouraged that our rebuttals have addressed all your concerns and that you have raised your confidence score. It is an honor to receive your recognition of the value Figure 1 adds to our presentation. We will appropriately include both the discussion and Figure 1 in the main body based on your constructive suggestion.

Once again, we truly appreciate your support and insightful comments.

Warm regards,

Authors of Submission 15661
Thank you for your response. I have checked all rebuttals, and they have addressed my concerns. Thus, I raised my confidence score from 3 to 4. 
The discussion about the **Linearity of LLM Representations** is interesting. The PDF Figure 1 looks nice. It would be good to include them in the main body.
Dear Reviewer dQGe,

Thank you very much for your follow-up. We would like to provide an **edited** version of our response for further clarification.

In the original Definition 4, $\Phi^*$ represents all the estimators that can achieve zero 0-1 loss with high probability at least $1-\delta$.

For clarity, we propose removing Definition 4 and directly stating in subsequent results (Proposition 2 and Lemma 1-2) that with high probability at least $1-\delta$, certain results will hold. The reason for these high-probability statements is to avoid extreme cases of initialization and noise, under which our subsequent results would not rigorously hold. The successful evolution of the estimators (whether the expected $E[\Psi']$ or the SGD-updated $\Psi'$) relies on a non-extreme, non-zero initialization.

We apologize for any confusion and appreciate your continuous patience and constructive feedback. We will include more clarifications in our revisions to make this clearer. Please let us know if this addresses your concern.

Warmest regards,

Authors of Submission 15661
Thank the authors for the detailed response. Regarding Weakness 1.a, I still find the argument confusing, if I understand correctly $\Psi$'  refers to a specific weight of the model, and $\Phi*$ contains all the weights that can minimize the 0-1 loss. Why does this have anything to do with noise or initialization?