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

Provably Transformers Harness Multi-Concept Word Semantics for Efficient In-Context Learning

Abstract

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

1 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 on transformer 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 [19, 20, 21, 22, 23], unrealistic infinite dimensional assumption [14, 19, 21, 24] and impractical loss functions like square loss [9, 16, 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:

Essential Questions

Whether and how do the geometric regularity of the multi-concept-encoded representation facilitate transformer in conducting efficient ICL?

To answer the above question, following the meaningful data modeling ideas in [9, 29], we conduct theoretical analysis on a concept-specific sparse coding prompt distribution for classification tasks, where the sparse latent variable encodes the information denoting the word's belonging concept. Importantly, the features in both the word's and label's dictionaries exhibit concept-specific geometric properties - 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.

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.

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.

We provably show how the multi-concept encoded linear 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.

2 Related Work

Theory of Exponential Convergence Rate of Stochastic Gradient Descent. Our analysis of the exponential convergence rate for the 0-1 loss builds upon prior work linking the excess risk and essential supremum norm to exponentially fast convergence under the 'hard low-noise condition' [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.

Theory of Transformers and In-Context Learning The literature on Transformers and ICL is wide-ranging, and we will selectively address the most relevant ones. Prior studies have analyzed how transformers learn topic/concept semantics [9], the origins and biases of LLM representations using latent variable models [10], and ICL from a model averaging perspective [14]. However, albeit incorporating concept variables, these works do not connect the geometric properties of 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.

3 Problem Setup

Notations. For l 2 and Frobenius norms we utilize ∥ · ∥ and ∥ · ∥ F to denote their computations. Considering two series a n and b n , we denote a n = O ( b n ) if there exists positive constant C > 0 and N > 0 such that for all n ≥ N , | a n | ≤ C | b n | . Similarly, we denote a n = Ω( b n ) if b n = O ( a n ) holds, and a n = Θ( b n ) if a n = O ( b n ) and a n = Ω( b n ) both hold. Our 1 ( · ) is to denote the indicator variable of an event. 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 .

3.1 Data Distribution

The data distribution employed in this study draws inspiration from a range of empirical and theoretical research works [9, 10, 46, 47, 48]. This distribution captures context-awareness and can be viewed as a specialized prompt version of PLSA [49] and LDA [6]. In this distribution, each word and label has multiple feature embeddings, each embedding corresponding to a different concept. This is achieved through the use of 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 µ + k 1 and

µ -k 1 , and their corresponding label's feature vectors q + k 1 and q -k 1 , ∀ k 1 ∈ [ K 1 ] . There are also K 2 task-irrelevant concepts denoted by ν k 2 , ∀ 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:

where the feature dictionary M = [ µ + 1 , µ -1 , µ + 2 , µ -2 , · · · , µ + K 1 , µ -K 1 , ν 1 , ν 2 , · · · , ν K 2 ] ∈ R d X × K exhibits positive inner products within concepts and orthogonality across concepts, and the label dictionary Q = [ q + 1 , q -1 , q + 2 , q -2 , · · · , q + K 1 , q -K 1 , 0 , · · · 0] ∈ R d Y × K has similar geometric properties. Specifically, we have ∀ k 1 ∈ [ K 1 ] , k 2 ∈ [ K 2 ] , ∥ µ ± k 1 ∥ = ∥ ν k 2 ∥ = ∥ u ∥ , ∥ q ± k 1 ∥ = ∥ q ∥ , and there exist constants 0 < κ x , κ y < 1 such that 0 < ⟨ µ + k 1 , µ -k 1 ⟩ ≤ κ x ∥ u ∥ 2 and 0 < ⟨ q + k 1 , q -k 1 ⟩ ≤ κ 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:

where P ± k,L +1 denotes the k -th concept-specific prompt distribution, and π ± k = (2 K 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 S n := 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, ¬ (2 k -1 ∨ 2 k ) = 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 .

3.2 Transformer Model

Following [17, 20, 28], our embedding E ( · ) of prompt S is formulated as H :

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:

where σ R ( · ) := Relu( · ) , σ S ( · ) := softmax( · ) , W Q , W K ∈ R m qk × ( d X + d Y ) , W V ∈ R m v × ( d X + d Y ) are the embedding matrices for queries, keys, and values, respectively, and W O ∈ R m × m v 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 let

where W x Q , W x K ∈ R d X × d X , W y V ∈ R ( m v -d X ) × d Y , W y O ∈ R m × d Y . Here, we set the elements other than W x Q , W x K , W y V and W y O to be zero. Besides, we fix W y V to be I ( m v -d X ) × d Y . We sample r i from a uniform distribution Unif {-1 , 1 } and fixed during the training process. Based on this setting, the trainable part we need to consider is actually Ψ ′ := { W x Q , W x K , W y O } . This problem remains highly non-convex and challenging.

We utilize mini-batch with-replacement SGD to train the transformer model. The empirical crossentropy loss for each batch B t is written as

where ℓ ( z ) = log(1+exp( -z )) , y S n 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.

Initialization Setting. All initial values of W y 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 [25, 28].

Testing Setting . The model performance is measured by 0-1 test error on a test prompt distribution D ∗ :

Algorithm 1 Training algorithm

Input: Training distribution D S , Test distribution D ∗ , Batch size B , step size η t = 2 λ ( γ + t ) , stopping criterion ε and total epochs T .

Initialize model parameters Ψ ′ (0) .

for

t

= 0

,

1

, . . . , T

-

If L 0 -1 D ∗ (Ψ ( t ) ) ≤ ε stop else continue.

1

do

Randomly sample mini batches B t of size B from D S .

Update model parameters:

end for

4 Theoretical Results

In this section, we present our main theoretical results, which is based on the following conditions. We consider the learning iterations 0 ≤ t ≤ T ∗ , where T ∗ = Ω( 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:

Ψ

′

(

t

+1)

′

= Ψ

(

t

)

-

η

∇

′

L

t

Ψ

B

t

(Ψ

′

(

t

)

)

.

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.

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(5 Km/δ ) /w ∗ 2 (1 -κ y ) ∥ q ∥ , we have

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 any k 1 ∈ [ K 1 ] , defining a k 1 := ( µ + k 1 + µ -k 1 ) / 2 and b k 1 := ( µ + k 1 -µ -k 1 ) / 2 , we find that for k ′ 1 = k 1 , { a k 1 , b k 1 } ⊥ { a k ′ 1 , b k ′ 1 } and ⟨ a k 1 , b k 1 ⟩ = 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 e b k 1 in µ e k 1 determines the label assignment. Similarly, defining c k 1 := ( q + k 1 + q -k 1 ) / 2 and d k 1 := ( q + k 1 -q -k 1 ) / 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 3 . During testing, the learned model admits probability distribution shift on D ∗ z and data shift on D ∗ x ×D ∗ y to generate a new prompt distribution

Thus after

D ∗ S = ∑ K 1 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.

D ∗ z can enjoy arbitrary distribution, satisfying that each prompt has at least one co-concept k ∈ [ K 1 ] , at least one pair shares the query word's co-concept's label, and still each word has equal chance to have positive or negative semantic labels over its concepts 4 .

̸

D ∗ 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 , ν k 2 = ν ∗ k 2 . Here, a ∗ k , b ∗ k , c ∗ k , d ∗ k are any vectors belong to the conic hulls of { a k } K 1 k =1 , { b k } K 1 k =1 , { c k } K 1 k =1 , { d k } K 1 k =1 respectively, satisfying ∥ b ∗ k ∥ ≥ ∥ a ∗ k ∥ = Θ( ∥ u ∥ ) and ∥ d ∗ k ∥ ≥ ∥ c ∗ k ∥ = Θ( ∥ q ∥ ) . ν ∗ k 2 = Θ( ∥ u ∥ ) are any vectors from the complement space of span ( M ) .

Again, the learned model satisfies L 0 -1 D ∗ S (Ψ ( 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.

5 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.

5.1 Idempotent Operator Techniques

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 k 1 } K 1 k 1 =1 , { b k 1 } K 1 k 1 =1 , { ν k 2 } K 2 k 2 =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 matrix

Similar techniques are also applied to the label's dictionary: Q := span ( Q ) , where we define q ⊥ 1 , · · · , q ⊥ d Y -K 1 as the standard orthogonal basis of the complement space Q ⊥ . In our subsequent derivation, the expectation E [ · ] is taken over the stochastic gradient descent. Similar to the idea in [34, 33, 36], we first serve to see how E (Ψ ( 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

Wv : Rmx(dx+dy)

H : R(dx+dy)x(L+1)

WK:Rmx(dx+dy)

hquery : R(dx+d)

Wo:Rmxme

pa

IRdy

rReLU(

★

Wy

0

=1

L

大

Wy

X1

XLXL+1

softmax

x7

y

W?

X1

y10

XL

y

∈R1xdy

Orthogonal

口

Decomposition

Components

WQ : Rmakx(dx+dy)

K

(αO(i)k

(t)

I|c=|12

Ck

T

Ilds:12

dk.T

十··

dk

qk

IdempotentDecomposition

1

口

(t)

Ck

bk

E[W(] =)

8=1

K1

αQ,s

(t)

asas

I|asl14

十

>

8=1

K

Q(t)

Q，s

bsbsT

[6 14

ak

E[(WKμ)]TE[(Wxμs)] = αQ,s ·αK,s ±βQ,s · βK,s

μk

with an ideally-balanced prompt set. Leveraging the symmetry of the prompt distribution, as well as the symmetry of W (0) Q and W (0) K , we introduce the following decompositions.

Lemma 1. We can decompose E [ 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.

Here α ( t ) Q,s , α ( t ) K,s and α ( t ) O ( i, · ) ,k represent the expected concept learning process, β ( t ) Q,s , β ( t ) K,s and β ( t ) 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 that

for ∀ 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.

5.2 Convergence of the Expectation

Denote U y Sn k,n ( t ) and W v k,n ( t ) -U y Sn k,n ( t ) as the activated neuron set for { i ∈ [ m ] | r i y S n > 0 } and { i ∈ [ m ] | r i y S n < 0 } separately, and ∑ l ∈ S y S n n,k ( σ ( t ) S ) n l represents the correct attention weight, where the detailed definitions are delayed in Appendix E. We then introduce the following lemma.

Under Condition 1, when

TrainingandTestLoss

EvolutionofAttentionWeights

LearningProgress ofAttention

LearningProgressofMLP

0.7

1.0

70

0.7

α% (0)

60

α (1)

0.6

0.6

IB% (0)

0.5

0.9

50

αo*αx（0）

0.5

IB%I (1)

LoSS

S

0.4

0.3

Test 0-1 Loss

TrainLossperbatch

0.8

40

30

βo*βk(0)

αo*ak（1)

βo*βx (1)

0.3

0.4

0.2

0.7

train attn weights (0)

20

Noise Product Max

0.2

0.1

train attn weights (1)

test attn weights (0)

10

0.1

0.0

test attn weights (1)

0

0.0

0

25

50

75

100

0102030405060708090

0102030405060708090

0102030405060708090

Epoch

Epoch

Epoch

Epoch

As such, the following lemmas show the learning outcomes of the E (Ψ ( t ) ) along the iterations.

Lemma 3. (Convergence of the Expectation). There exist constant C 1 > 0 , ∀ t ≥ ˆ T C 1 σ 1 mλK 1 γ √ (1 + κ y ) log(5 Km/δ ) /w ∗ 2 (1 -κ y ) ∥ q ∥ , we have L 0 -1 D ∗ ( E (Ψ ′ ( t ) )) = 0 .

Lemma 4. (Regularizing the models). Under Condition 1, it holds that

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.

5.3 Exponential Convergence of 0-1 loss

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 that

are martingale difference sequences, and for ∀ X ∈ { Q,K,O } and its corresponding W ∈ { W x Q , W x K , W y O } , we have ∑ T t =0 D t X = W ( T +1) -E [ W ( T +1) ] . Then we utilize the following lemma in [34, 55] to give a bound over the variance.

Lemma 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 have

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 ∑ T t =0 ∥ D t X ∥ 2 ∞ ≤ c T 2 , ∀ X ∈ { Q,K,O } , where ∥ · ∥ ∞ is the essential

=

20.4

0.7

0.6

0.5

0.2

0.1

0.0

E°0

Training and Test Loss in in OOD1 Evolution of Attention Weights in 0oD1

0

25

Train Loss per batch

Test 0-1 Loss

Epoch

50

75

100

1.0

0.9

0.8

0.7

0 10 20 30 40 50 60 70 80 90

Epoch

terit aftn weights (1)

trgin aftn weights (1)

test attn weignits (0)

60

50

40

30

20

10

Learming Progress of Attention in OOD1

0

06 08 0L 09 0S 09 0E 0Z 01 0

a*acf1l

8g*0z 10)

*%(0)

R (1)

Noise Product Max

Epoch

0.2

0.6

0.4

0.2

0.0

Learming Progress of MLP in OOD1

010 20 30 40 50 60 70 80 90

a, (0)

a (1)

ISI (0)

II(1)

Complement CoefE Max

UP9y Csaiy iSSON

Epoch

0.4

0.7

0.6

0.5

E'0

0.2

0.1

0.0

Training and Test Loss in in OOD1 Evolution of Acttention Weights in OoD1

0

25

Test 0-1 Loss

Train Loss per batch

Epoch

50

75

100

0.9

0.8

1.0

0.7

0.6

010 20 30 40 50 60 70 80 90

Epoch

terit attn weights (1)

tra/n atn weights (0)

tain attn weights (1)

test attn wtights (0)

60

50

40

30

20

10

Learning Progress of Attention in OOD1

D

010 20 30 40 50 60 70 80 90

*az (0)

A*Ar (0)

A*Ar (1)

ao*0x(1)

Complemepf Prgduct Max

Epoch

=0.2

0.6

0.4

0.2

0.0

Learning Progress of MLP in OOD1

06 08 0L 09 0% 0+ 00 02 01 0

e (0)

α (1)

I#6) (0)

#6I (1)

Epoch

0.4

0.6

0.5

E'0

0.2

0.0

0.7

0.1

Training and Test Loss in in OOD2 Evolution of Attention Weights in OOD2

SZ

Train Loss per batch

Test 0-1 L05s

Epoch

50

75

100

1.0

0.9

0.8

0.7

010 20 30 40 50 60 70 80 90

Epoch

train attn weights (0)

[L] suiam upe uetit

seit attn weighits (0)

test attn weighits (3)

60

50

40

30

20

10

Leaming Progress of Actention in OoD2

06 08 0L 09 0S 0+ 0E 0Z 0T 0

*#x (0)

e;*4 (I)

Ao*Rx (0)

Ao*R (1)

Noise Prodoct Maa

Epoch

0.7

0.6

0.5

0.4

E'0

0.2

0.1

0.0-

Learming Progress of MLP in OOD2

0 10 20 30 40 50 60 70 80 90

a, (0)

a (1)

II (0)

IMI(1)

Epoch

0.4

0.6

0.2

0.0

Training and Test Loss in in OOD3 Evolution of Attention Weights in OOD3

0

25

Train Loss per batch

Test 0-1 Loss

Epoch

50

75

100

1.0

0.9

0.8

0.7

0.6

06 08 0 09 0S 0+ 0E 0Z 0T 0

Epoch

test attn weights (1)

train attn weights (0)

train attn weights (1)

test attn weights (0)

09

50

40

30

20

10

Learning Progress of Attention in OoD3

0 10 20 30 40 50 60 70 80 90

A *Ar (0)

A *Ar (1)

Neoise ProduEt Mtagr

*0x (0)

o*(1)

Complemsent Pnoduc1 Max

Epoch

=0.2

0.6

0.4

0.0

0.2

Leaming Progress of MLP in 0003

0 10 20 30 40 50 60 70 80 90

e (0)

α (1)

(邮6) (1)

URatN owRn aijON

邮I (0)

Epoch

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 .

For 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

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 D ∗ z 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.

6 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 instances n 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) µ ± 1 ∗ = a 1 ± b 2 , µ ± 2 ∗ = a 2 ± b 1 . Figure 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.

7 Conclusion

This work provides the first exponential convergence analysis of 0-1 loss for transformers with softmax attention and ReLU-MLP, trained on a non-orthogonal concept-specific prompt distribution by practical cross-entropy loss. Furthermore, the results demonstrate transformers can perform certain OOD ICL tasks by leveraging the multi-concept semantic linearity, highlighting their innovative potential. An important future direction is to extend the analysis to more complex scenarios.

8 Acknowledgment

We thank the anonymous reviewers for their instrumental comments. D.B. and H.W. are supported in part by the Research Grants Council of the Hong Kong Special Administration Region (Project No. CityU 11206622). W.H. is supported in part by JSPS KAKENHI (24K20848). A.N. is supported in part by National Research Foundation, Singapore and Infocomm Media Development Authority under its Trust Tech Funding Initiative, the Centre for Frontier Artificial Intelligence Research, Institute of High Performance Computing, A*Star, and the College of Computing and Data Science at Nanyang Technological University. T.S. is supported in part by JSPS KAKENHI (24K02905) and JST CREST (JPMJCR2115, JPMJCR2015).

References

Li Dong, Nan Yang, Wenhui Wang, Furu Wei, Xiaodong Liu, Yu Wang, Jianfeng Gao, Ming Zhou, and Hsiao-Wuen Hon. Unified language model pre-training for natural language understanding and generation. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'AlchÉ-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed Chi, Quoc V Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems , volume 35, pages 24824-24837, 2022.

Fei Liu, Xialiang Tong, Mingxuan Yuan, Xi Lin, Fu Luo, Zhenkun Wang, Zhichao Lu, and Qingfu Zhang. Evolution of heuristics: Towards efficient automatic algorithm design using large language model. arXiv preprint arXiv:2401.02051 , 2024.

Fei Liu, Yiming Yao, Ping Guo, Zhiyuan Yang, Xi Lin, Xialiang Tong, Mingxuan Yuan, Zhichao Lu, Zhenkun Wang, and Qingfu Zhang. A systematic survey on large language models for algorithm design. arXiv preprint arXiv: 2410.14716 , 2024.

Sheng Lu, Irina Bigoulaeva, Rachneet Sachdeva, Harish Tayyar Madabushi, and Iryna Gurevych. Are emergent abilities in large language models just in-context learning? arXiv preprint arXiv: 2309.01809 , 2023.

David Blei, Andrew Ng, and Michael Jordan. Latent dirichlet allocation. In Advances in Neural Information Processing Systems , 2001.

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference, 2022.

Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning. In Workshop on Efficient Systems for Foundation Models @ ICML2023 , 2023.

Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In Proceedings of the 40th International Conference on Machine Learning , pages 19689-19729, 2023.

Yibo Jiang, Goutham Rajendran, Pradeep Ravikumar, Bryon Aragam, and Victor Veitch. On the origins of linear representations in large language models. arXiv preprint arXiv: 2403.03867 , 2024.

Kiho Park, Yo Joong Choe, and Victor Veitch. The linear representation hypothesis and the geometry of large language models. arXiv preprint arXiv: 2311.03658 , 2023.

Kiho Park, Yo Joong Choe, Yibo Jiang, and Victor Veitch. The geometry of categorical and hierarchical concepts in large language models. arXiv preprint arXiv: 2406.01506 , 2024.

Yibo Jiang, Goutham Rajendran, Pradeep Ravikumar, and Bryon Aragam. Do llms dream of elephants (when told not to)? latent concept association and associative memory in transformers. arXiv preprint arXiv: 2406.18400 , 2024.

Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv: 2305.19420 , 2023.

Fabian Falck, Ziyu Wang, and Christopher C. Holmes. Are large language models bayesian? a martingale perspective on in-context learning. In ICLR 2024 Workshop on Secure and Trustworthy Large Language Models , 2024.

Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference on Machine Learning , volume 202, pages 35151-35174. PMLR, 2023.

Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv: 2306.09927 , 2023.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Advances in Neural Information Processing Systems , volume 36, pages 57125-57211, 2023.

Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv: 2402.01258 , 2024.

Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv: 2310.05249 , 2023.

Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In Advances in Neural Information Processing Systems , volume 36, pages 71911-71947, 2023.

Yingcong Li, Yixiao Huang, Muhammed E. Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with self-attention. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics , volume 238, pages 685-693, 2024.

Chenyu Zheng, Wei Huang, Rongzhen Wang, Guoqiang Wu, Jun Zhu, and Chongxuan Li. On mesa-optimization in autoregressively trained transformers: Emergence and capability. arXiv preprint arXiv:2405.16845 , 2024.

Shokichi Takakura and Taiji Suzuki. Approximation and estimation ability of transformers for sequence-to-sequence functions with infinite dimensional input. arXiv preprint arXiv: 2305.18699 , 2023.

Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multihead softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv: 2402.19442 , 2024.

Yu Huang, Zixin Wen, Yuejie Chi, and Yingbin Liang. Transformers provably learn featureposition correlations in masked image modeling. arXiv preprint arXiv: 2403.02233 , 2024.

Hongkang Li, Meng Wang, Sijia Liu, and Pin yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015 , 2023.

Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. How do nonlinear transformers learn and generalize in in-context learning? arXiv preprint arXiv: 2402.15607 , 2024.

Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In Proceedings of the 38th International Conference on Machine Learning , volume 139, pages 11112-11122. PMLR, 2021.

Patrik Reizinger, Szilvia Ujváry, Anna Mészáros, Anna Kerekes, Wieland Brendel, and Ferenc Huszár. Position: Understanding LLMs requires more than statistical generalization. In Proceedings of the 41st International Conference on Machine Learning , volume 235, pages 42365-42390, 2024.

Enno Mammen and Alexandre B Tsybakov. Smooth discrimination analysis. The Annals of Statistics , 27(6):1808-1829, 1999.

Pascal Massart and Élodie Nédélec. Risk Bounds for Statistical Learning. The Annals of Statistics , 34(5):2326 - 2366, 2006.

Loucas Pillaud-Vivien, Alessandro Rudi, and Francis Bach. Exponential convergence of testing error for stochastic gradient methods. In Proceedings of the 31st Conference On Learning Theory , volume 75, pages 250-296, 2018.

Atsushi Nitanda and Taiji Suzuki. Stochastic gradient descent with exponential convergence rates of expected classification errors. In Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics , volume 89, pages 1417-1426, 2019.

Vivien A Cabannes, Francis Bach, and Alessandro Rudi. Fast rates for structured prediction. In Proceedings of Thirty Fourth Conference on Learning Theory , volume 134 of Proceedings of Machine Learning Research , pages 823-865. PMLR, 15-19 Aug 2021.

Shingo Yashima, Atsushi Nitanda, and Taiji Suzuki. Exponential convergence rates of classification errors on learning with sgd and random features. In Proceedings of The 24th International Conference on Artificial Intelligence and Statistics , volume 130, pages 1954-1962, 2021.

Kazusato Oko, Taiji Suzuki, Atsushi Nitanda, and Denny Wu. Particle stochastic dual coordinate ascent: Exponential convergent algorithm for mean field neural network optimization. In International Conference on Learning Representations , 2022.

Stefano Vigogna, Giacomo Meanti, Ernesto De Vito, and Lorenzo Rosasco. Multiclass learning with Margin: Exponential rates with no bias-variance trade-off. In Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 22260-22269. PMLR, 17-23 Jul 2022.

Vivien Cabannnes and Stefano Vigogna. A case of exponential convergence rates for svm. In Proceedings of The 26th International Conference on Artificial Intelligence and Statistics , volume 206, pages 359-374. PMLR, 25-27 Apr 2023.

Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. In The Eleventh International Conference on Learning Representations , 2023.

Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convolutional neural networks. In Advances in Neural Information Processing Systems , volume 35, pages 25237-25250, 2022.

Yiwen Kou, Zixiang Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting in two-layer reLU convolutional neural networks. In Proceedings of the 40th International Conference on Machine Learning , volume 202, pages 17615-17659, 2023.

Xuran Meng, Difan Zou, and Yuan Cao. Benign overfitting in two-layer relu convolutional neural networks for XOR data. arXiv preprint arXiv: 2310.01975 , 2023.

Zhiwei Xu, Yutong Wang, Spencer Frei, Gal Vardi, and Wei Hu. Benign overfitting and grokking in reLU networks for XOR cluster data. arXiv preprint arXiv: 2310.02541 , 2023.

Wei Huang, Yuan Cao, Haonan Wang, Xin Cao, and Taiji Suzuki. Graph neural networks provably benefit from structural information: A feature learning perspective. arXiv preprint arXiv: 2306.13926 , 2023.

Hiroaki Yamagiwa, Momose Oyama, and Hidetoshi Shimodaira. Discovering universal geometry in embeddings with ica. arXiv preprint arXiv: 2305.13175 , 2023.

Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In Proceedings of the 38th International Conference on Machine Learning , pages 11112-11122, 2021.

Chi Han, Jialiang Xu, Manling Li, Yi Fung, Chenkai Sun, Nan Jiang, Tarek Abdelzaher, and Heng Ji. Word embeddings are steers for language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 16410-16430, 2024.

Thomas Hofmann. Probabilistic latent semantic indexing. In Proceedings of the 22nd Annual International ACM SIGIR Conference on Research and Development in Information Retrieval , page 50-57, 1999.

LÉon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review , 60(2):223-311, 2018.

Atsushi Nitanda and Taiji Suzuki. Optimal rates for averaged stochastic gradient descent under neural tangent kernel regime. arXiv preprint arXiv: 2006.12297 , 2021.

Yiwen Kou, Zixiang Chen, Yuan Cao, and Quanquan Gu. How does semi-supervised learning with pseudo-labelers work? a case study. In The Eleventh International Conference on Learning Representations , 2023.

Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. Understanding the generalization of adam in learning neural networks with proper regularization. In The Eleventh International Conference on Learning Representations , 2023.

Tong Yang, Yu Huang, Yingbin Liang, and Yuejie Chi. In-context learning with representations: Contextual generalization of trained transformers. arXiv preprint arXiv: 2408.10147 , 2024.

Iosif Pinelis. Optimum bounds for the distributions of martingales in banach spaces. The Annals of Probability , 22(4):1679-1706, 1994.

Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In Proceedings of The 33rd International Conference on Machine Learning , volume 48, pages 1225-1234, 2016.

Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. The benefits of mixup for feature learning. In Proceedings of the 40th International Conference on Machine Learning , volume 202, pages 43423-43479, 2023.

Jinghui Chen, Yuan Cao, and Quanquan Gu. Benign overfitting in adversarially robust linear classification. In Robin J. Evans and Ilya Shpitser, editors, Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence , volume 216, pages 313-323, 2023.

Spencer Frei, Niladri S Chatterji, and Peter Bartlett. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Proceedings of Thirty Fifth Conference on Learning Theory , volume 178, pages 2668-2703, 2022.

Spencer Frei, Gal Vardi, Peter Bartlett, and Nathan Srebro. Benign overfitting in linear classifiers and leaky relu networks from kkt conditions for margin maximization. In Proceedings of Thirty Sixth Conference on Learning Theory , volume 195, pages 3173-3228, 2023.

Yiwen Kou, Zixiang Chen, and Quanquan Gu. Implicit Bias of Gradient Descent for Twolayer ReLU and Leaky ReLU Networks on Nearly-orthogonal Data. In Advances in Neural Information Processing Systems , volume 36, pages 30167-30221. Curran Associates, Inc., 2023.

Wei Huang, Ye Shi, Zhongyi Cai, and Taiji Suzuki. Understanding convergence and generalization in federated learning through feature learning theory. In The Twelfth International Conference on Learning Representations , 2024.

Dake Bu, Wei Huang, Taiji Suzuki, Ji Cheng, Qingfu Zhang, Zhiqiang Xu, and Hau-San Wong. Provably neural active learning succeeds via prioritizing perplexing samples. In Proceedings of the 41st International Conference on Machine Learning , volume 235, pages 4642-4695, 2024.

Yiwen Kou, Zixiang Chen, Quanquan Gu, and Sham M. Kakade. Matching the statistical query lower bound for k-sparse parity problems with stochastic gradient descent. arXiv preprint arXiv: 2404.12376 , 2024.

Alexander Tsigler. Benign Overfitting in Linear Regression and Classification . PhD thesis, UC Berkeley, 2024.

Junhyung Park, Patrick Bloebaum, and Shiva Prasad Kasiviswanathan. Benign overfitting for regression with trained two-layer relu networks. arXiv preprint arXiv: 2410.06191 , 2024.

Eshaan Nichani, Alex Damian, and Jason D. Lee. Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks. In Advances in Neural Information Processing Systems , volume 36, pages 10828-10875, 2023.

Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns lowdimensional polynomials with sgd near the information-theoretic limit. arXiv preprint arXiv:2406.01581 , 2024.

Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Learning sum of diverse features: computational hardness and efficient gradient-based training for ridge combinations. arXiv preprint arXiv:2406.11828 , 2024.

Yunwei Ren and Jason D. Lee. Learning orthogonal multi-index models: A fine-grained information exponent analysis. arXiv preprint arXiv:2410.09678 , 2024.

Spencer Frei and Gal Vardi. Trained transformer classifiers generalize and exhibit benign overfitting in-context. arXiv preprint arXiv:2410.01774 , 2024.

Wei Shen, Ruida Zhou, Jing Yang, and Cong Shen. On the training convergence of transformers for in-context classification. arXiv preprint arXiv:2410.11778 , 2024.

Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du. JoMA: Demystifying multilayer transformers via joint dynamics of MLP and attention. In The Twelfth International Conference on Learning Representations , 2024.

Yu Huang, Zixin Wen, Yuejie Chi, and Yingbin Liang. How transformers learn diverse attention correlations in masked vision pretraining. arXiv preprint arXiv: 2403.02233 , 2024.

Masahiro Sakamoto and Hitomi Sato. Benign or not-benign overfitting in token selection of attention mechanism. arXiv preprint arXiv:2409.17625 , 2024.

Roey Magen, Shuning Shang, Zhiwei Xu, Spencer Frei, Wei Hu, and Gal Vardi. Benign overfitting in single-head attention. arXiv preprint arXiv:2410.07746 , 2024.

Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv: 2402.14735 , 2024.

Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. arXiv preprint arXiv: 2409.10559 , 2024.

Hongru Yang, Bhavya Kailkhura, Zhangyang Wang, and Yingbin Liang. Training dynamics of transformers to recognize word co-occurrence via gradient flow analysis. arXiv preprint arXiv:2410.09605 , 2024.

Jiarui Jiang, Wei Huang, Miao Zhang, Taiji Suzuki, and Liqiang Nie. Unveil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization. arXiv preprint arXiv:2409.19345 , 2024.

Bingrui Li, Wei Huang, Andi Han, Zhanpeng Zhou, Taiji Suzuki, Jun Zhu, and Jianfei Chen. On the optimization and generalization of two-layer transformers with sign gradient descent. arXiv preprint arXiv:2410.04870 , 2024.

Yoshua Bengio. Learning Deep Architectures for AI . Now Publishers Inc, 2009.

Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning . MIT Press, 2016.

Zeyuan Allen-Zhu and Yuanzhi Li. Backward Feature Correction: How Deep Learning Performs Deep Learning. In Conference on Learning Theory , COLT '23, 2023. Full version available at http://arxiv.org/abs/2001.04413.

Core Francisco Park, Maya Okawa, Andrew Lee, Ekdeep Singh Lubana, and Hidenori Tanaka. Emergence of hidden capabilities: Exploring learning dynamics in concept space. arXiv preprint arXiv:2406.19370 , 2024.

Yongyi Yang, Core Francisco Park, Ekdeep Singh Lubana, Maya Okawa, Wei Hu, and Hidenori Tanaka. Dynamics of concept learning and compositional generalization. arXiv preprint arXiv:2410.08309 , 2024.

Lingjing Kong, Guangyi Chen, Biwei Huang, Eric P. Xing, Yuejie Chi, and Kun Zhang. Learning discrete concepts in latent hierarchical models. arXiv preprint arXiv: 2406.00519 , 2024.

Zeyuan Allen-Zhu and Yuanzhi Li. Physics of Language Models: Part 1, Learning Hierarchical Language Structures. ArXiv e-prints , abs/2305.13673, May 2023. Full version available at http://arxiv.org/abs/2305.13673.

Michael Hahn and Navin Goyal. A theory of emergent in-context learning as implicit structure induction. arXiv preprint arXiv: 2303.07971 , 2023.

Emanuele Marconato, Sébastien Lachapelle, Sebastian Weichwald, and Luigi Gresele. All or none: Identifiable linear properties of next-token predictors in language modeling. arXiv preprint arXiv:2410.23501 , 2024.

Roger A. Horn and Charles R. Johnson. Matrix Analysis . Cambridge university press, 2012.

Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Miao Lu, Beining Wu, Xiaodong Yang, and Difan Zou. Benign oscillation of stochastic gradient descent with large learning rate. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning , 2023.

Yurii Nesterov. Introductory Lectures on Convex Optimization: A Basic Course , volume 87. Springer Science & Business Media, 2013.

Karan Girotra, Lennart Meincke, Christian Terwiesch, and Karl T. Ulrich. Ideas are dimes a dozen: Large language models for idea generation in innovation. SSRN , 2023.

Anil Rajnikant Doshi and Oliver Hauser. Generative artificial intelligence enhances creativity but reduces the diversity of novel content. SSRN , 2023.

Fei Liu, Xialiang Tong, Mingxuan Yuan, and Qingfu Zhang. Algorithm evolution using large language model. arXiv preprint arXiv: 2311.15249 , 2023.

Yiming Yao, Fei Liu, Ji Cheng, and Qingfu Zhang. Evolve cost-aware acquisition functions using large language models. arXiv preprint arXiv: 2404.16906 , 2024.

Fei Liu, Xialiang Tong, Mingxuan Yuan, Xi Lin, Fu Luo, Zhenkun Wang, Zhichao Lu, and Qingfu Zhang. An example of evolutionary computation + large language model beating human: Design of efficient guided local search. arXiv preprint arXiv: 2401.02051 , 2024.

A Limitation and Broader Impact

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.

B Additional Experiment Details

Weimplement 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.

C Additional Related Work

Theory of Convergence Rate of Stochastic Gradient Descent . Our analysis of the exponential convergence rate for the 0-1 loss builds upon a rich body of prior work. 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, 42, 43, 45, 52, 53, 57, 58, 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 nonlinear attention-MLP transformers with within-concept positive inner products and cross-concept orthogonal data representations. The key to our analysis is the assumption of good initialization of attention matrices and a sufficiently low-noise condition, which is reasonable for modeling language rather than images. In this setting, SGD allows noise to have only 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 [20, 21, 26, 54, 73, 74, 75, 76, 77, 78, 79], ReLU-free MLP [54, 80, 81] 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, [73, 88, 89] show that transformers can capture hierarchical and compositional structures in data. From a Bayesian perspective, [7, 8, 14] 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 statisticallyindependent concepts are nearly orthogonal [9, 10, 11, 12, 90]. Interestingly, aligning with the findings in [46, 90], Independent Component Analysis (ICA) is naturally more suitable than Principal Component Analysis (PCA) for obtaining meaningful feature or label vectors in our prompt modeling. This is because the features or labels are nearly statistically independent and of equal strength, especially with a large K , while the noise is

feeble in our modeling. 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 Preliminary Lemmas

D.1 Probablistic Lemmas on Concentration

Lemma 6. Suppose that δ > 0 and ∀ d ∈ { d X , d Y } = Ω(log( KNL δ )) , where N = BT ∗ . Then with probability at least 1 -δ ,

σ

2

ξ

2

|⟨

ξ

|⟨

d

i

,

ξ

i

,

≤ ∥

ξ

i

∥

ξ

i

2

2

′

⟩| ≤

µ

2

σ

⟩| ≤ ∥

µ

2

ξ

∥

σ

2

ξ

d

2

√

·

2

σ

ξ

·

,

6(

2 log(

for all ξ i , ξ ′ i ∼ D ξ x ( or D ξ y ) , µ ∈ D x ( or D y ) , l ∈ { 1 , 2 } .

Proof. See Lemma B.4 in [42] for a proof.

Lemma 7. Suppose that δ > 0 , d Y = Ω(log( m/δ )) , m = Ω(log( K/ ( δ ))) . Then with probability at least 1 -δ , for ∀ i ∈ [ m ] , k ∈ [ K 1 ] , w ∈ [ d Y -K 1 ] ,

̸

Moreover, for some ζ ∈ (0 , 1] for ∀ e = e ′ , ∈ [ ± ] , ∃ ω ζ ∈ (0 , ω ′ ζ ) where ω ′ ζ < 1 ,

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 exist i ∈ { i ∈ [ m ] | r i = e m , α (0) O ( i, · ) ,k + ζβ (0) O ( i, · ) ,k > 0 , α (0) O ( i, · ) ,k -ζβ (0) O ( i, · ) ,k < 0 } , such that α (0) O ( i, · ) ,k ≤ -σ 1 / 2 ∥ c k ∥ .

d

log

√

3

(

N

(

L

+1))

δ

6

KN

(

L

δ

2

)

+1)

)

,

Proof. First, notice that W y O ( i, · ) (0) ∼ N ( 0 , σ 1 I d Y ) , then by Bernstein's inequality as well as d Y = Ω(log( m/δ )) , with probability at least 1 -δ/ (5 m ) , for ∀ i ∈ [ m ]

By union bound we can have the first inequality in the lemma hold with probability at least 1 -δ/ 5 . Next, we notice that

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

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 have

then with probability 1 -δ/ 5 , the second and third inequality hold.

For ζ ∈ (0 , 1] , we see that the variable α (0) O ( i, · ) ,k + e ′ ζβ (0) O ( i, · ) ,k ∼ N (0 , σ 2 1 ( ∥ c k ∥ 2 + ζ ∥ d k ∥ 2 )) , and it's independent to the event { r i = e m } , ∀ e ∈ [ ± ] . Therefore, we can see the count of { i ∈ [ m ] | r i = e m , α (0) O ( i, · ) ,k + e ′ ζβ (0) O ( i, · ) ,k > 0 , e ′ ζβ (0) O ( 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 have

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 .

For the event { i ∈ [ m ] | r i = e m , α (0) O ( i, · ) ,k ± ζβ (0) O ( i, · ) ,k > 0 } , we have

where 1 + ω ζ 2 is the probability of the conditional event { α (0) O ( i, · ) ,k -e ′ ζβ (0) O ( 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 ∈ [ m ] | r i = e m , α (0) O ( 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( K 1 /δ )) ensure √ m log(10 K 1 /δ ) / 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.

D.2 Matrix Theories

Lemma 8. (1.1.P5 in [91]) Let A ∈ M n 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 µ i P i , where P i are symmetric idempotent matrices with rank( P i ) = 1 , and thus ∑ d i =1 µ i P i is the idempotent decomposition of matrix A by P i . 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 .

Proof. By definition,

Then, by Lemma 8 we have

D.3 ODE Systems

Lemma 10. (Lemma C.1 in [43]). Suppose that a sequence a t , t ≥ 0 follows the iterative formula

for some 0 ≤ c ≤ 1 and b ≥ 0 . Then it holds that

for all t ≥ 0 . Here, x t is the unique solution of

Lemma 11. (Coupled ODE System 1). Suppose that there are two coupled sequences y t , z t , t ≥ 0 follows the iterative formula

for some a, b ≥ 0 . Then it holds that

Proof. From the condition we see that z 0 < 0 and z t is an increasing sequence ( z t ≥ z 0 ). Besides, as y 0 > 0 , during the period where z t ≤ 0 , we see that y t is monotonically decreasing. Then by (2 + e -2 y t 2 + e 2 y 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 y t . Then one can readily obtain the result by solving the ODE.

for all t ≥ 0 . Here, y ( t ) , z ( t ) are the unique solutions of the following ODE System respectively

As such, for t 1 = min { t ∈ Z | z t ≥ 0 } , we have

Lemma 12. (Coupled ODE System 2). Suppose that there are two coupled sequences y t , z t , which are the sequences after t 1 in Lemma 11, and t ≥ t 1 follows the iterative formula

for some a, b ≥ 0 , and c ′ ≤ ℓ ′ t ≤ 1 . Then it holds that

for all t ≥ t 1 . Here, y ( t ) , y ( t ) , z ( t ) , z ( t ) are the unique solutions of the following ODE System respectively

where

Proof. We see that as z t ≥ 0 , t ≥ t 1 , the y t 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 -2 y ( t ) 2 1+ e -2 y ( t ) 2 and ℓ ′ t . Easy to see that

and then collaborating with c ′ ≤ ℓ ′ t ≤ 1 we can obtain the result by solving the ODE. Observing that

Thus by the monotonicity the system is unique, which is also ture for the upper bound ODE. The proof is completed.

E Data Distribution

This section provided the detailed formal definitions of the prompt distribution.

Definition 3. ( Polysemous Word Model ( D x , D y , D z , D ξ x , D ξ y ) ). We assume there exists K 1 concepts of words totally. Specifically, each concept k 1 ∈ [ K 1 ] 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 D x and D y , which can be written as the following forms via reparameterization:

where z ∈ R K ( K < d X ) . 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:

where

The sparse latent variable z = ( z 1 , · · · , z K ) ∈ { 0 , 1 } k is sampled from D z . P ( z j = 1) = Θ( log log K K ) .

̸

M = [ µ + 1 , µ -1 , µ + 2 , µ -2 , · · · , µ + K 1 , µ -K 1 , ν 1 , ν 2 , · · · , ν K 2 ] = [ M 1 , · · · , M K ] ∈ 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 =1 are conceptirrelevant features, and ∀ k ∈ [ K ] , ∥ M k ∥ = ∥ u ∥ . We assume that features of the same concept have positive inner product: ∃ 0 < κ x < 1 , ∀ k 1 ∈ [ K 1 ] , 0 < ⟨ µ + k 1 , µ -k 1 ⟩ ≤ κ x ∥ u ∥ 2 . Meanwhile, we let the features of different concept be orthogonal: ∀ e ∈ [ ± ] , e ′ ∈ [ ± ] , s ′ ∈ [ K 1 ] , r = r ′ ∈ [ K 2 ] , 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 × K is the corresponding label dictionary matrix, where ∥ q ± k ∥ = ∥ q ∥ , for ∀ k ∈ [ K 1 ] . Similarly, we let the labels of the same concept to have positive inner product: ∃ 0 < κ y < 1 , ∀ k 1 ∈ [ K 1 ] , 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 ′ ∈ [ K 1 ] .

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 Θ( K 1 -1 ) , and the chance for selecting the two semantically-opposite vectors of the same concept is 1 2 . During training, each prompt S = { x 1 , y 1 , · · · , x L , y L , x L +1 } is sampled from the mixture distribution D S defined as below.

where π + k = π -k = 1 2 K 1 , and the P + k,L +1 and P -k,L +1 are prompt distributions characterized by the k -th concept, defined as

̸

̸

where P l,j := P ( z l,j = 1) . ∀ n ∈ [ N ] where N is the training size, if the training prompt S n is sampled from P e k,L +1 , e ∈ [ ± ] , k ∈ [ K 1 ] , then by Definition 1, the label vector of the query should contain q e k , and we call y S n = e as the real value label of this k -th concept prompt. Specifically, for ∀ k ∈ [ K 1 ] we define the index set of training prompts sharing the k -th co-concepts as

For sample x l where n ∈ V k , k ∈ [ K 1 ] , l ∈ [ L +1] , we define the index set for its non-zero elements of z n l besides z n 2 k -1 ,l and z n 2 k,l , namely M n l := { k ∈ [ K ] | z n l,k = 1 , k / ∈ { 2 k -1 , 2 k }} . Also, for each prompt sharing the k -th co-concept, we define the index set of demonstration in the context:

F Model details: Attention Part

In this section, we provide several important definitions and compute the original gradients of attention.

Lemma 13. (Contributing and Misleading Neurons)

W k,n ( t ) := W + k,n ( t ) ∪ W -k,n ( t ) are neurons that can be activated, among which U k,n ( t ) := U + k,n ( t ) ∪ U -k,n ( t ) are neurons that correctly contribute to the prediction. The following lemma computes the original gradients.

Lemma 14. (Gradient Update) Denote

∇ W x Q ( t ) L B t (Ψ ( t ) ) ∈ R d X × d X can be derived as

K

Subsequently, we directly compute the update of the attention matrices along the feature directions as below.

Lemma 15. (Concept Learning of Attention) For ∀ ˆ k ∈ [ K 1 ] , we have the single step of learning of the concept part of the features:

where I ( t ) Q, a ˆ k , chaos and I ( t ) Q, a ˆ k , contri are defined as below.

̸

Similarly, I ( t ) K, a ˆ k , chaos and I ( t ) K, a ˆ k , contri are defined as below.

̸

Lemma 16. (Label Semantic Learning of Attention) Also, for ∀ ˆ k ∈ [ K 1 ] , we have the single step of learning of the concept-specific semantically-opposite part of the features:

where I ( t ) Q, b ˆ k , chaos and I ( t ) Q, b ˆ k , contri are defined as below.

̸

Similarly, I ( t ) K, b ˆ k , chaos and I ( t ) K, b ˆ k , contri are defined as below.

̸

G Model details: MLP Part

Lemma 17. (Tensor Update)

Lemma 18. (Gradient Update) ∇ W y O ( i, · ) ( t ) L B t (Ψ ( t ) ) ∈ R 1 × ( d X + d Y ) can be derived as

̸

Lemma 19. (Concept Learning of MLP) For ∀ i ∈ [ m ] , ˆ k ∈ [ K 1 ] ,

̸

When the neuron is activated (i.e., { n ∈ V ˆ e k ∩ B t , if ( W x K ( t ) b k ) ⊤ W x Q ( t ) b k > 0 } , and α ( t ) O ( i, · ) , ˆ k + ˆ e · (2 ∑ l ∈ S ˆ e n, ˆ k ( σ ( t ) S ) n l -1) W y O ( i, · ) ( t ) d ˆ k > 0 ), the neuron is likely to be activated ( i ∈ W ˆ e ˆ k,n ( t ) ).

If (1) r i · ˆ e > 0 , i ∈ W ˆ e ˆ k,n ( t ) ⇔ i ∈ W ˆ e ˆ k,n ( t ) ∩ U ˆ e ˆ k,n ( t ) , the gradient will advance the W y O ( i, · ) ( t ) c ˆ k ;

if (2) r i · ˆ e < 0 , i ∈ W ˆ e ˆ k,n ( t ) ⇔ i ∈ W ˆ e ˆ k,n ( t ) - U ˆ e ˆ k,n ( t ) , the gradient will diminish the W y O ( i, · ) ( t ) c ˆ k , thus help deactivate this neuron.

Lemma 20. (Label Semantic Learning of MLP) For ∀ i ∈ [ m ] , ˆ k ∈ [ K 1 ] ,

where I ( t ) O ( i, · ) , d ˆ k , chaos and I ( t ) O ( i, · ) , d ˆ k , contri are defined as

H Discussions over Parameter Settings

Note that we do not have any requirement upon demonstration length L and batch size B for training, thus the training can be really flexible compared with the strict requirement in [28]. The condition on dimensionality d X , d Y 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 [92] 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 [93], 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.

I Convergence of Expectation

In this section, we assume all the events in the Section D hold, denoted as Υ Pre .

Weexamine the evolution of E (Ψ ′ 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 ∀ k 1 ∈ [ K 1 ] , we define a k 1 := µ + k 1 + µ -k 1 2 and b k 1 := µ + k 1 -µ -k 1 2 . By definition, we then have

̸

for ∀ k ′ 1 = k 1 ∈ [ K 1 ] .

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 ⊥ ) = d X -K . Then we can have a set of standard orthogonal basis for R d , defined as

where u ⊥ 1 , · · · , u ⊥ d X -K are the standard orthogornal basis of U ⊥ . Then we can derive that

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 decompose E [ W x Q ( t ) ] and E [ W x K ( t ) ] by (scaled) idempotent matrices.

where α ( t ) Q,s and α ( t ) K,s represent the concept learning process, β ( t ) 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 , thus

we examine t +1 . It holds that

Here, we see E (Ψ ′ ( t ) ) as fixed matrices and the expectation E B 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:

with some coefficients a ( t ) K,s , b ( t ) K,s , c ( t ) Q,r , d ( t ) K,w , ∀ s ∈ [ K 1 ] , r ∈ [ K 2 ] , w ∈ [ d X -2 K 1 -K 2 ] . It's direct to check that E [ W x Q ( t ) ] also has the exactly same outcome. The proof is completed.

Worth noting that

We will also have

for ∀ e ∈ [ ± ] and for ∀ e ′ ∈ [ ± ] , s ′ ∈ [ K 1 ] , r ∈ [ K 2 ] , w ∈ [ d X -K ] , ∀ u ∈ { µ e ′ s ′ , ν r , u ⊥ w } ,

Similar conclusions hold when the query vectors are ν r and u ⊥ w , ∀ r ∈ [ K 2 ] , w ∈ [ d X -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:

where q ⊥ 1 , · · · , q ⊥ d Y -K 1 are the standard orthogonal basis of the complement space Q ⊥ . Then we have

for ∀ e ∈ [ ± ] , i ∈ [ m ] , k ∈ [ K 1 ] .

Lemma 23. At initialization, for some e ∈ [ ± ] and ∀ k ∈ [ K 1 ] , define

then we have some ω ζ e k ∈ (0 , ω ′ ζ e k ) where ω ′ ζ e k < 1 , the following will hold

The parameter ω ′ ζ e k is determined by σ 0 , σ 1 , ∥ a k ∥ , ∥ b k ∥ , ∥ c k ∥ and ∥ d k ∥ .

̸

Proof. We have that E n ∈V e k [ U e k,n (0) ∩ ( W -e k,n (0) -U -e k,n (0))] = ∅ . By Lemma 7, we see that for

we can have corresponding ω ζ e k ∈ (0 , ω ′ ζ e k ) where ω ′ ζ e k < 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 have

where e ∈ [ ± ] , s ∈ [ K 1 ] , r ∈ [ K 2 ] , w ∈ [ d X -K ] .

Lemma 25. For ∀ k 1 ∈ [ K 1 ] , we define c k 1 := q + k 1 + q -k 1 2 and d k 1 := q + k 1 -q -k 1 2 . By definition, we then have

̸

for ∀ k ′ 1 = k 1 ∈ [ K 1 ] .

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 have

where r ∈ [ K 2 ] , w ∈ [ d X -K ] , ˆ w ∈ [ d Y -K 1 ] .

Lemma 27. For t > 0 , we have

Proof. 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 ∈ [ K 1 ] , α ( t ) Q,k = α ( t ) K,k , β ( t ) Q,k = β ( t ) K,k . Observe that for ∀ k ∈ [ K 1 ] ,

We see from Lemma 7 that α (0) Q,k = α (0) K,k = σ 0 ∥ a k ∥ 2 , β (0) Q,k = β (0) K,k = σ 0 ∥ b k ∥ 2 . Therefore, for t = 0 , ∀ k ∈ [ K 1 ] , we have

Obviously, E n ∈V y S n k [ ∑ j ∈ S y S n n,k ( σ (0) S ) n j ] > 0 . 5 > E n ∈V y S n k [ ∑ j ∈ S -y S n n,k ( σ (0) S ) n j ] . Meanwhile we see that E n ∈V y S n k [ ∑ j ∈ S y S n n,k ( σ (0) S ) n j ] ≈ 0 . 5 due to the small σ 0 = O ( ∥ u ∥ -2 ) by Condition 1.

The observation in Eq. (39), collaborating with the positiveness of β ( t ) Q,k · β ( t ) K,k , we see that the inequality E n ∈V y S n k [ ∑ j ∈ S y S n n,k ( σ ( t ) S ) n j ] > E n ∈V y S n k [ ∑ j ∈ S -y S n n,k ( σ ( t ) S ) n j ] will remain during whole iteration. Also, by Eq. (39), we know that

This observation under our expectation scenario greatly facilitate our analysis. Since ℓ ′ n ( t ) < 0 , it's obvious that the signal of r i β ( t ) O ( i, · ) ,k will determine whether the neuron i ∈ [ m ] will serve to increase or decrease the β ( t ) Q,k and β ( t ) K,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 have

where k ∈ [ K 1 ] .

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 ( I ( t ) O ( 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 , T 1 ) , for ∀ k ∈ [ K 1 ] , e ∈ [ ± ] , we have

The neuron set E [( W e k,n ( t ) -U e k,n ( t )) -U -e k,n ( t )] is non-increasing, and all of this neuron will get deactivated. Additionally, both E [ α ( t ) O ( i, · ) ,k ] and e · β ( t ) O ( i, · ) ,k would monotonically decrease. Also, it holds that e · β ( t ) O ( i, · ) ,k > 0 and | α ( t ) O ( i, · ) ,k | ≤ e · β ( t ) O ( i, · ) ,k ;

The neuron set E [ U e k,n ( t ) -( W -e k,n ( t ) -U -e k,n ( t ))] is non-increasing, and all neurons in it will turn into E n ∈V e k [ U e k,n ( t ) ∩ ( W -e k,n ( t ) -U -e k,n ( t ))] . Additionally, both E [ α ( t ) O ( i, · ) ,k ] and e · β ( t ) O ( i, · ) ,k would monotonically increase. Also, it holds that e · β ( t ) O ( i, · ) ,k > 0 and | α ( t ) O ( i, · ) ,k | ≤ e · β ( t ) O ( i, · ) ,k ;

For E n ∈V e k [ U e k,n ( t ) ∩ ( W -e k,n ( t ) -U -e k,n ( t ))] , the e · β ( t ) O ( i, · ) ,k would monotonically increase. Besides,

when there exists constant C ≥ 1 such that

the E [ α ( t ) O ( i, · ) ,k ] would be contributed to increase, otherwise it will decrease. Also, | α ( t ) O ( i, · ) ,k | ≥ E [ | e (2 ∑ l ∈ S e n,k ( σ ( t ) S ) n l -1) β ( t ) O ( i, · ) ,k | ] and E [ α ( t ) O ( i, · ) ,k ] > 0 ;

All the neurons in E n ∈V e k [ U e k,n ( t ) ∩ ( W -e k,n ( t ) -U -e k,n ( t ))] will ultimately either have its coefficient update stuck due to regularization, or grow into a changing margin into E [ U e k,n ( t ) -( W -e k,n ( t ) -U -e k,n ( t ))] where

Proof. By Lemma 28, we see that ∀ i ∈ E [ U e k,n ( t )] , α ( t ) O ( i, · ) ,k and eβ ( t ) O ( i, · ) ,k would be contributed by V e k to increase, and also ∀ i ∈ E [ W e k,n ( t ) - U e k,n ( t )] , α ( t ) O ( i, · ) ,k and eβ ( t ) O ( i, · ) ,k would be contributed by V e k to decrease. As such, the first and second point hold naturally by definition. The ultimate transformation of E [ U e k,n ( t ) -( W -e k,n ( t ) -U -e k,n ( t ))] into E n ∈V e k [ U e k,n ( t ) ∩ ( W -e k,n ( t ) -U -e k,n ( t ))] attributes to the faster changing

speed of α ( t ) 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 r i 's direction, and thus the third and forth point hold.

Considering the growth of α ( t ) O ( i, · ) ,k , by π + k = π -k , P l, 2 k -1 = P n l, 2 k = 1 2 , we know

hence if i ∈ E n ∈V e k [ U e k,n ( t ) ∩ ( W -e k,n ( t ) -U -e k,n ( t ))] , it indicates that

We see that for E 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 α ( t ) O ( i, · ) ,k . The contribution will tend to be positive if

Then, as E [(2 ∑ l ∈ S e n,k ( σ ( t ) S ) n l -1) eβ ( t ) O ( i, · ) ,k ] of the neurons in and E n ∈V e k [ U e k,n ( t ) ∩ ( W -e k,n ( t ) - U -e k,n ( t ))] will continue to grow, and finally it will be comparable to the E [ α ( t ) O ( i, · ) ,k ] . Otherwise it will continue to grow while the evolving speed of E [ α ( t ) 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 E n ∈V e k ( ℓ ′ n ( t ) ) -E n ∈V -e k ( ℓ ′ n ( t ) ) would be controlled by the initialization. Interestingly, we see that as E [(2 ∑ l ∈ S e n,k ( σ ( t ) S ) 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

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 of E [ α ( t ) O ( i, · ) ,k ] will also stuck at the changing margin from E n ∈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 know

Therefore, a sufficient condition for L 0 -1 D ∗ ( E (Ψ t )) = 0 is

We know ∀ i ∈ U e k,n ( t ) , E [ e · β ( t ) O ( i, · ) ,k ] in the left side of the inequality is increasing, and ∀ i ∈ E [ W e k,n ( t ) -U e k,n ( t )] , the E [ e · β ( t ) O ( i, · ) ,k ] in the right side of the inequality is decreasing, which is a good news since we want the left side exceed the right side. By Lemma 29, we see that all the neurons in E [( W e k,n ( t ) -U e k,n ( t )) -U -e k,n ( t )] will be deactivated, and all the neurons in E [ U e k,n ( t ) -( W -e k,n ( t ) -U -e k,n ( t ))] will turn into E n ∈V e k [ U e k,n ( t ) ∩ ( W -e k,n ( t ) -U -e k,n ( t ))] .

I.1 First Stage: Growing of Coefficient

In this stage, the coefficient update dynamic is continually changing without being much influenced by the comparably feeble regularization. Also, the impact of the decaying learning step η t 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, the E S n ∼D S [ f ( E ( S n ); Ψ (0) )] satisfies

and our remaining job is to see when will E S n ∼D S [ f ( E ( S n ); E (Ψ ( t ) ))] stay positive for some error tolerance. As such, we need to scrutinize the coefficients that would grow along the iterations. Therefore, we define

We will see that the conditional expectation of this sequence (conditioned on E (Ψ ′ ( t ) ) , and the expectation is taken over D S ) would grow up to conquer the small initialization and make E S n ∼D S [ f ( E ( S n ); E (Ψ ′ ( t ) ))] stay positive. Consider the whole training duration 0 ≤ t ≤ T ∗ , the evolving speed of β ( t +1) Q,k , β ( t +1) K,k , α ( t +1) O ( i, · ) ,k and β ( t +1) O ( i, · ) ,k depends on E [ ℓ ′ n ( t ) ] , E [ 1 n O ( i ) ( t ) ] and E [ ∑ l ∈ S e n,k ( σ ( t ) S ) n l ] . Denote

We will show that σ ∗ S is the lower bound of min t ∈ [ T ∗ ] ,k ∈ [ K 1 ] { E n ∈D S [ ∑ j ∈ S y S n n,k ( σ ( t ) S ) n j ] } along the whole iteration. By Lemma 7, κ can be upper bounded by 8 √ 2 log(5 Km/δ ) · σ 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 iteration 0 ≤ t ≤ T ∗ , for ∀ i ∈ [ m ] , e ∈ [ ± ] , k ∈ [ K 1 ] , r ∈ [ K 2 ] , w ∈ [ d X -K ] , we have that

Lemma 31. Suppose Eq. (45) holds at iteration t ≤ T 2 , then we have

Proof. By definition, we have

Observe that

Here the inequality holds due to the fact that E [ α ( t ) O ( i, · ) ,k 1( i ∈ ( W e k,n ( t ) -U e k,n ( t )) -U -e k,n ( t ))] is decreasing the initial value α (0) O ( i, · ) ,k 1( i ∈ ( W e k,n (0) -U e k,n (0)) -U -e k,n (0)) , and it's absolute value will not surpass that of E [ e (2 ∑ l ∈ S y S n n,k ( σ ( t ) S ) n l -1) β ( t ) O ( i, · ) ,k ] ≤ κ/ 8 , which is positive (by definition) and also decreasing by Lemma

(45)

29. On the other hand,

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 where

Latter we will show in Lemma 33 that

We then denote the learning step at ˆ T as η := η ˆ T , and thus

By Lemma 31, actually it would hold that

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 ∈ [ K 1 ] , e ∈ [ ± ] , it holds that

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.

The difference of the loss derivative is bounded by O ( κ ) :

E [ A k,e t ] is non-decreasing. The lower and upper bounds of the gradient update have continuous ODE counterpart. Specifically, there exist positive constant c 1 , c 2 , we can define c k,e = c 1 η 0 ∥ q ∥ 2 2 mK 1 , c k,e = c 2 η T ∗ (2 σ ∗ S -1) 2 (1 -κ y ) ∥ q ∥ 2 16 mK 1 , b k,e = e -κ/ 2 , b k,e = e κ/ 2 . Let x k,e t , x k,e t be the unique solutions of

then it holds that

Specifically, we have

As cross-entropy loss is L -smooth with L = 1 , one can bound the difference by

By Lemma 23, we see that for initialization, we have

Now we serve to show that the following expected difference

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 n ∈V e k ( ef ( E ( S ); E (Ψ ( t ) )) , ∀ e ∈ [ m ] . To better present this observation, we define

which further means that e ∗ t satisfies E [ E n ∈V e ∗ t k ( ℓ ′ n ( t ) ) -E n ∈V -e ∗ t k ( ℓ ′ n ( t ) )] < 0 due to the non-positive and non-increasing property of cross-entropy loss.

Recall the update rule, we have

By Lemma 28 and Lemma 29, we see that the E [ ∑ i ∈U e k,n ( τ ) -( W -e k,n ( τ ) -U -e k,n ( τ )) α ( τ ) O ( i, · ) ,k ] ∣ ∣ ∣ τ =0 τ = t , ∀ e ∈ [ ± ] is increasing such that

where we ignore the impact of cross-concept safely due to the large K = Ω( η 0 C ( K 1 -1) ∥ q ∥ 2 / ( mK 1 )) , 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 that

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 of

depend on the signal of E [ E n ∈V e k ( ℓ ′ n ( t ) ) -E n ∈V -e k ( ℓ ′ n ( t ) )] . Specifically, we see that

where the contribution term is shared by the two sequences. Therefore, by Eq.(50) and (46), the evolution of

will prefer to grow in the direction of e ∗ t .

We then take a look on the decreasing coefficients based on Lemma 29.

As such, we have all preliminaries to characterize the first result of the lemma. We first utilize the induction to prove the following:

This apparently hold at initialization. Suggest for any t ≤ ˜ t -1 the result holds, then we only need to prove

By the condition of small η t in Condition 1, Lemma 31, Eq.(48) (49), (50) and (51), we see that

and thus we have

Therefore, we complete the induction. Then we have

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 update

Then, utilizing Lemma 31 and the fact |W e k,n ( t ) | ≤ m , we have constant c 1 > 0 such that

where we also neglect the impact of cross-concept due to the large K = Ω( η 0 C ( K 1 -1) ∥ q ∥ 2 / ( mK 1 )) in Condition 1 and appropriately chosen c 1 .

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 ) )] :

Then by the first result in this lemma we know | E [ E n ∈V e k ( ℓ ′ n ( t ) ) -E n ∈V -e k ( ℓ ′ n ( t ) )] | ≤ β 4 ≤ κ 32 , and thus

By the condition on the small initialization in Condition 1 such that σ 1 = O ( (2 σ ∗ S -1) 2 Cm 3 / 2 ∥ q ∥ ) , due to the large C , we see that the κ = O ((2 σ ∗ S -1) 2 /m ) is far more feeble. Thus the gradient contributions made by neuron set E [ α ( t ) O ( i, · ) ,k ] where i ∈ E [ U e k,n ( t ) ∩ ( W -e k,n ( t ) -U -e k,n ( t ))] can be neglected compared to the increasing update of E [ α ( t ) O ( i, · ) ,k 1 ( i ∈ U e k,n ( t ) -( W -e k,n ( t ) -U -e k,n ( t )))] and E [ e (2 ∑ l ∈ S e n,k ( σ ( t ) S ) n l -1) β ( t ) O ( i, · ) ,k )] . Besides, we

see that E [ W e k,n ( t )] will at least preserve the neurons of E [ U e k,n (0))] , which will not be deactivated by Lemma 29.

Then there exists c 2 > 0 , recall σ ∗ S is defined in Lemma 34 as the lower bound of min t,k { E n ∈D S [ ∑ j ∈ S y S n n,k ( σ ( t ) S ) n j ] } and E [ |U e k,n (0)) | ] ≥ m/ 8 , it holds that

where we ignore the impact of regularization term at this stage since λ = O (( C log( Km/δ ) ∥ q ∥ ) -1 ) and appropriately chosen c 2 . 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

thus

Thus

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 prove

Recall the update rule

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

where the first inequality is by the positive nature of regularization term as well as the contribution of the gradient; second inequality is by E [ -ℓ ′ 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 2 mK 1 ≤ 1 ≤ log( T ∗ ) , as well as ( σ ∗ S ) -1 ≥ 1 . The remaining job is to prove that

Observe that

Here the first inequality is by 1 / (1 + exp( z )) ≤ exp( -z ) ; the second inequality is by Lemma 31; the last inequality is by the feeble κ/ 2 and E [ e · β ( t ) O ( i, · ) ,k 1 ( i ∈ U e k,n ( t ))] ≥ ( σ ∗ S ) -1 log( T ∗ ) . Then we have that

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 have

where we set ˆ C = 2 C ′ √ 2 log( 5 Km δ ) 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 that

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 = Ω( η 0 C ( K 1 -1) ∥ q ∥ 2 / ( mK 1 )) . Then we have

where the last inequality is by the induction hypothesis and the C ′ can be taken as C , which completes the induction.

We now prove

Recall the update rule

Easy to see that E [ 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 the neuron i ∈ E [ W e k,n ( t ) -U e k,n ( t )] would be deactivated at t +1 once

This indicates that for the neuron i ∈ E [ W e k,n ( ˜ t ) -U e k,n ( ˜ t )] ,

Now collaborating with Eq. (54) and Eq. (55), we now can have

The proof is completed.

I.1.1 Expected 0-1 loss Convergence

Lemma 33. Under Condition 1, there exist constant C 1 > 0 , after at most

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 that

Note that by definition A k,e 0 = 0 , and we recursively use the equation t times

For each k ∈ [ K 1 ] , 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 that

Thus there exists a positive constant ˜ C such that -E n ∈V e k [ ℓ ′ n ( t ) ] ≥ ˜ C for 0 ≤ t ≤ ˜ t k,e . Then we have

Therefore we see that for ∀ k ∈ [ K 1 ] , e ∈ [ ± ] , E [ A k,e t ] will reach κ within 16 mK 1 κ ˜ Cc 2 η (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(5 Km/δ )) , we see that there exist a positive constant C 1 = Θ(64 / ( ˜ Cc 2 )) , the threshold time can be

Then by definition of 0-1 loss we have

The proof is completed.

I.1.2 Period 1: Decreasing Period of Correct Attention Score

We claim that if ∑ i ∈ [ m ] r i β (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 ∈ [ K 1 ] , after at most a certain iterations

where C 3 is a positive constant, we would have the β ( t ) Q,k = β ( t ) K,k be monotonically increasing during the remaining iterations T 1 ≤ t ≤ T ∗ . Besides, it holds that σ ∗ S is the lower bound of the lowest correct attention assignment along the whole iterations:

Proof. By Lemma 27, the β ( t +1) Q,k = E [ β ( t +1) K,k | Ψ ( t ) ] will be contributed to increase by

and they will be contributed to decrease by

By the fifth inequality in Lemma 7, we know that

and β ( t ) K,k will be contributed to increase or decrease depends on the signal of ∑ i ∈ E [ W ± k,n ( t )] r i · eβ ( t ) O ( i, · ) ,k . By the last inequality in in Lemma 7, we see that at initialization,

By the expected gradient update in Lemma 28, the eβ ( t ) O ( i, · ) ,k will grow in r i 's direction along the whole iterations. As such, the values of E [ r i · β ( t ) O ( i, · ) ,k 1 ( i ∈ W ± k,n ( t ))] , ∀ k ∈ [ K 1 ] will grow larger. Therefore, after a limited epochs we can have

where the β ( t ) Q,k and β ( t ) K,k would be contributed positively and monotonically increase.

Now we serve to find the lower bound of the evolution of E [( ∑ l ∈ S e n,k ( σ ( t ) S ) n l )] , which is clearly to be the first iteration where the negative ∑ i ∈ E [ W ± k,n ( t )] r i · eβ ( t ) O ( i, · ) ,k has grown to surpass the 0 . By the symmetry property denoted in Lemma 22 and Eq.(39) we have

Recall that

and we also see that

As ∑ i ∈ E [ W ± k,n ( t )] r i · 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 ( σ ( t ) S ) n l )] , we have 1 ≥ -E ( ℓ ′ n ( t ) ) ≥ ˜ C . Also, m ≥ E [ |U e k,n (0)) | ] ≥ 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 negative E [ ∑ i ∈ [ m ] r i β ( 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 β ( t ) Q,k negligible, as

Here 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(5 Km/δ ) /K 1 ) -1 . Therefore, it holds that during the decreasing period of β ( t ) Q,k as well as the period where ∑ i ∈ E [ W ± k,n ( t )] r i · β ( t ) O ( i, · ) ,k remain negative, we have

Here, by a appropriate chosen small C 4 , 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 . 01 C -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 have

where the inequality is by the negative nature of ∑ i ∈ E [ W ± k,n ( t )] r i · β ( 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 )] r i · eβ ( t ) O ( i, · ) ,k as the lower bound sequence of the β ( t ) Q,k and ∑ i ∈ E [ W ± k,n ( t )] r i · 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(5 Km/δ ) · 5 σ 1 ∥ d k ∥ 16 ), and their update rule are

Then by Lemma 11, let a = C 4 η 0 ∥ b k ∥ 4 K 1 , b = C 4 η 0 ∥ d k ∥ 2 K 1 (2 1 1 + e -2 β (0) Q,k 2 / ∥ b k ∥ 2 -1) , we have the maximum iterations T 1 = -z (0)(1 + e -2 y (0) 2 ) b (1 -e -2 y (0) 2 ) = σ 1 K 1 γ √ 10 log(5 Km/δ )(1 + e -2 σ 2 0 ∥ b k ∥ 2 ) 2 C 4 ∥ d k ∥ (1 -e -2 σ 2 0 ∥ b k ∥ 2 ) , set C 3 = √ 10 / (2 C 4 )

we obtain the T 1 in the lemma. The lower bound of β ( t ) Q,k along the decreasing period as

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 ∈ [ K 1 ] as β -QK , which can be given as

Recall σ ∗ S is defined as

which is actually can be written as

Therefore, we see that σ ∗ S is the lower bound of min t ∈ [ T ∗ ] ,k ∈ [ K 1 ] { E n ∈D S [ ∑ j ∈ S y S n n,k ( σ ( t ) S ) n j ] } .

Remark 4. As we see that in Lemma 33, we require that the lower bound given in Eq.(53) depends that the values of E [ r i (2 ∑ l ∈ S e n,k ( σ ( t ) S ) n l -1) β ( t ) O ( i, · ) ,k )] surpasses κ , which naturally says that the value of E i ∈U e k,n (0)) [ r i β ( t ) O ( i, · ) ,k ] should surpass κ since E [(2 ∑ l ∈ S e n,k ( σ ( t ) S ) n l -1)] ≤ 1 . Therefore E i ∈U e k,n (0)) [ r i β ( t ) O ( i, · ) ,k ] should surpass 0 at ˆ T since κ > 0 , which indicates that ˆ T > T 1 . We see that the initial period t ≤ T 1 is where E i ∈U e k,n (0)) [ r i β ( t ) O ( i, · ) ,k ] grow to surpass the initial scale, whose upper bound is κ/ 8 by the definition of κ .

I.1.3 Period 2: Increasing Priod of Correct Attention Score

This period's analysis is based on Period 1 in Section I.1.2, or a good initialization such that

Lemma 35. Under Condition 1, consider the duration after T 1 in Lemma 34, then for ∀ k ∈ [ K 1 ] , consider the period T 1 ≤ t ≤ T 2 = C 5 min { 1+ γ λ , ∥ u ∥∥ q ∥ λK 1 √ m } , where C 5 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 ± ( t )] r i · eβ ( t ) O ,k respectively. That is, there exists positive constants c 3 -6 , for

k,n ( i, · ) a = c 3 (1 -κ x ) ∥ u ∥ 4 λγK 1 , a = c 4 ∥ u ∥ 4 λγK 1 , b = c 5 (1 -κ y ) ∥ q ∥ 2 λγK 1 ) , b = c 6 ∥ q ∥ 2 λγK 1 ) , c ′ = ˜ C , it holds that

for all t ≥ T 1 . Here, y ( t ) , y ( t ) , z ( t ) , z ( t ) are the unique solutions of the following ODE System respectively

where

For some limited constant △ such that ∃△ , σ ∗ S < △≤△ < 1 . Then the β ( t ) Q,k = β ( t ) K,k will grow to make the correct attention score E n ∈V y S n k [ ∑ j ∈ S y S n n,k ( σ ( t ) S ) n j ] achieve the △ in at least a △ (1 -△ )

scaled Gaussian rate such that

Proof. By Remark 4, we see that at the initial phase during t ≥ T 1 , we have ∑ i ∈ [ m ] r i β (0) O ( i, · ) ,k ≤ κ/ 8 , and thus by Eq.(55) in Lemma 30 we see that α (0) O ( i, · ) ,k ≤ ˆ C ∥ c k ∥ 2 σ ∗ S ∥ d k ∥ 2 . This indicates that E [ A k,e t ] ≤ Θ( α ) and thus by Lemmar 31 we see that the scale of | E n ∈V e k [ ef ( E ( S ); E (Ψ ( t ) ))]] | is also Θ( κ ) . This suggest that there still exists a constant ˜ C , during a certain amount of subsequent iterations we would still have that ˜ C ≤ -E [ ℓ ′ ( t )] ≤ 1 . Also, m ≥ E [ |U e k,n (0)) | ] ≥ 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 have

and

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 = Θ( C 4 η 0 ∥ b k ∥ 4 K 1 ) , b = Θ( C 4 η 0 ∥ d k ∥ 2 K 1 ) , c ′ = ˜ C . Then by solving the coupled ODE systems, collaborating the scale bounding property ( -κ 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. (27) and (36), as well as the Comparison Theorem with some constants c 3 -6 , we can have upper and lower bound of β ( 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

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 that

Proof. By the asymptotic property of Ei( x )

This suggest that when y ( t ) ≥ y ( t ) is close to infinity, the lower bound ODE in Lemma 35 will approximately satisfies the following

Then we see that as y ( t ) goes to infinity, we have a lower bound

Onthe other hand, obtaining an upper bound over y ( t ) is relatively easy. Since we have 2+ e -2 y ( t ) 2 + e 2 y ( t ) 2 ≤ 4 and (1 -e -2 y ( t ) 2 ) / (1 + e -2 y ( t ) 2 ) ≤ 1 , which gives the upper bound ODE over attention and MLP considering z (0) > 0

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 ˆ c such that

I.2 Second Stage: Regularizing the Model

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.

This suggest that roughly

Lemma 37. Under Condition 1, consider e = E [ y S n ] for all t ∈ [ T 2 , T ∗ ] it holds that

where eβ ( t ) O ( i, · ) ,k represents m r i β ( 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:

We split the situations into two cases:

eβ

(

t

O

β

)

(

i,

(

t

)

Q,k

·

)

,k

∥

u

∥

2

q

∥

2

λmK

√

1

,

1

2

log(

(ii).

and

6

∥

u

λ

∥

2

∥

q

2

mK

∥

1 2

2

)

E

[(

∑

j

∈

S

+

n,k

(

σ

(

t

S

)

)

n

j

)]

1+

λ

6

∥

1

2

mK

u

∥

2

∥

1 2

q

∥

2

orders of the case (i)'s quantities are less than those of case (ii), thus this split is plausible.

Recall

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 ˆ Cσ ∗ S -2 (1 + κ y ) 2 (1 -κ y ) -2 β ( t ) O ( i ∗ , · ) ,k ≥ ˆ C . That means that the term E [ A k,e t ] can be controlled by its contributor Θ( eβ ( t ) O ( i, · ) ,k ) . Then we have

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 have

. Easy to note that the scales'

Thus the growing of eβ ( t ) O ( i, · ) ,k would be non-degenerated: E [ eβ ( t +1) O ( i, · ) ,k ] ≥ Θ( eβ ( t ) O ( 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

then, as we see that

where the first inequality is by the definition of E [( ∑ j ∈ S ± n,k ( σ ( t ) S ) 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

Here the inequality is by the case (i) hypothesis upon eβ ( t ) O ( i, · ) ,k , -E [ ℓ ′ ( t )] ≤ 1 , and

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 that

For case (ii), we directly check that

Here the inequality is by -E [ ℓ ′ ( t )] ≤ 1 and

As a result, by the gradient form we see that β ( t +1) Q,k ≤ β ( t ) Q,k , and thus we prove the induction proving goal E [ eβ ( t +1) O ( i, · ) ,k ] = O ( ∥ q ∥ 2 λmK 1 ) .

Similarly, as we now have

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

where the inequality is by -E [ ℓ ′ ( t )] ≤ 1 ,

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:

In addition, we can see that even if we suggest the MLP's eβ ( t ) O ( i, · ) ,k is growing in a fastest linear-level speed, it require at least Θ( 1 + γ λ ) to reach the maximum admissible value ∥ q ∥ 2 2 λmK 1 . Meanwhile, we see that even when considering the fast speed of the increasing attention, by the asymptotic perperty 1 discussed in Lemma 36, we see that we still require Θ( ∥ u ∥∥ q ∥ λK 1 √ m ) to reach the highest admissible correct attention score 1 1 + λ 2 mK 1 2 2 ∥ u ∥ 2 ∥ q ∥ 2 .

Therefore, we can have some appropriately small constants C 5 , and when the iteration number is more than T 2 = C 5 min { 1 + γ λ , ∥ u ∥∥ q ∥ λK 1 √ m } , we need to consider the impact of regularization.

Lemma 38. The scale of the coefficients will finally be stabilized at a considerable level:

where eβ ( t ) O ( i, · ) ,k represents m r i β ( t ) O ( i, · ) ,k . That is, we consider the positive growth of | β ( t ) O ( i, · ) ,k | .

Proof. Recall the last discussion in Lemma 29, we see that as E [(2 ∑ l ∈ S e n,k ( σ ( t ) S ) n l -1) eβ ( t ) O ( i, · ) ,k ] getting larger and larger, it will finally reach the scale of α ( t ) O ( i, · ) ,k , which has updated in a feeble speed controlled by initialization when the neuron fell into the neuron set E n ∈V e k [ U e k,n ( t ) ∩ ( W -e k,n ( t ) -U -e k,n ( t ))] . After α ( t ) O ( i, · ) ,k ≤

E [(2 ∑ l ∈ S e n,k ( σ ( t ) S ) n l -1) eβ ( t ) O ( i, · ) ,k ] , the neuron would change into the neuron set E n ∈V e k [ U e k,n ( t ) -( W -e k,n ( t ) -U -e k,n ( t ))] . As such, the α ( t ) O ( i, · ) ,k would again increase at a normal speed, which is even faster than eβ ( t ) O ( i, · ) ,k due to the update rules and the fact that ∥ c k ∥ > ∥ d k ∥ . As such, the neuron set E n ∈V e k [ U e k,n ( t ) -( W -e k,n ( t ) -U -e k,n ( t ))] would again fell back into the neuron set E n ∈V e k [ U e k,n ( t ) ∩ ( W -e k,n ( t ) -U -e k,n ( t ))] , where the update speed is again feeble. And it will increase until E [(2 ∑ l ∈ S e n,k ( σ ( t ) S ) n l -1) eβ ( t ) O ( i, · ) ,k ] catch up.

Besides, 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 [( σ ( t ) S ) n l ] ≤ 1 . As such, ultimately we have E [( ∑ l ∈ S e n,k ( σ ( t ) S ) n l -∑ l ∈ S -e n,k ( σ ( t ) S ) n l )] = Θ(1) , α ( T ∗ ) O ( i, · ) ,k ≤ eβ ( T ∗ ) O ( i, · ) ,k and E [ A k,e t ] = Θ( E [ eβ ( t ) O ( i, · ) ,k ]) . Then following the process in Lemma 37 we can obtain the results. Here we omit this part since the proving procedure is the same to Lemma 37, despite we see E [( ∑ l ∈ S e n,k ( σ ( t ) S ) n l -∑ l ∈ S -e n,k ( σ ( t ) S ) n l )] = Θ(1) and E [ A k,e t ] = Θ( E [ eβ ( t ) O ( 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 K 1 , γ, ∥ u ∥ , ∥ q ∥ , κ x , κ y , we will have

for all t ≥ T 2 . Here, y ( t ) , y ( t ) , z ( t ) , z ( t ) are the unique solutions of the following ODE System respectively

where

Additionally, we would have asymptotic property that

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 have

for all t ≥ T 2 . Here, y ( t ) , y ( t ) , z ( t ) , z ( t ) are the unique solutions of the following ODE System respectively

where

Additionally, we would have asymptotic property that

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.

J Exponential Convergence of 0-1 Loss

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:

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 ∈ [ K 1 ] , 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 bound

where the first inequality is by (2 ∑ l ∈ S e n,k ( σ ( t ) S ) n l -1) ≤ 1 , and Lemma 6; the last equality is by the low noise condition σ ξ ≤ ∥ q ∥ /C √ d Y in Condition 1. Then by the low noise condition as well as the data model's definition we see that all the 2-norm of the W y O is controlled by the K 1 concepts' corresponding lengths in projection space. Then by the definition of Frobenius norm and Eq.(22) we have

Moreover, the Frobenius norm of W x Q ( t ) and W x K ( t ) and its gradient can be bounded as below

Proof. By Eq.(16) and (17), we see that

Here, the first inequality is by the scaled identity initialization of W x Q (0) , W x K (0) , orthogonal relationships of vectors in Lemma 27, Lemma 6, ∑ l,j ∈ [ L ] ( σ ( t ) S ) n l ( σ ( t ) S ) n j ≤ 1 / 4 , Eq.(62) in Lemma 41; the second inequality is by the low noise condition σ ξ ≤ λm/ ( C √ d X ∥ u ∥∥ q ∥ 1 / 2 ) and the large K ≥ C ∥ u ∥ / ( σ ξ √ d X ) for a large C in Condition 1. Thus the update of a ⊤ ˆ 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 of W x Q (0) , W x K (0) , orthogonal relationships of vectors in Lemma 27, Eq.(19) and (20) we can see that

We see that b ⊤ ˆ k W x Q ( t ) b ˆ k and b ⊤ ˆ k W x K ( 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 ∈ [ K 1 ] , 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 b ⊤ ˆ 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

Here, the second arrow is by the definition of ∑ j ∈ S + n,k ( σ ( t ) S ) 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 -2 b ⊤ ˆ 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 of W 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 ∈ [ K 2 ] , w ∈ [ d X -2 K 1 -K 2 ] . Collaborating with Lemma 9, we have

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 ∈ [ K 1 ] 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 ∥ λ W x Q ( t ) ∥ F (or ∥ λ W x K ( t ) ∥ F ). Thus

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 y O ( i, · ) ( t ) q ± k / ∥ q ± k ∥∥ in Lemma 41, are looser in the order of K 1 -2 and K 1 -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β ( t ) O ( i, · ) ,k , α ( t ) O ( i, · ) ,k . This actually indirectly demonstrates the superiority of considering expectations.

Lemma 43. (Restatement of Proposition 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 that E [ A k,e t ] ≥ κ , which will ensure that E [ y S n · f ( E ( S n ) , E (Ψ ˆ T ))] ≥ κ/ 2 . Therefore, when conditioned on E [ W x Q ( t ) ] , E [ W x K ( t ) ] , a minimum admissible disparity between W y O ( t ) and E [ W y O ( t ) ] corresponds the the minimum admissible disparity between W y O ( i, · ) ( t ) c k , W y O ( i, · ) ( t ) d k and α ( t ) O ( i, · ) ,k , β ( t ) O ( 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 ∈ [ K 1 ] , the minimum admissible disparity can be written as

Therefore, we see that when conditioned on E [ W x Q ( t ) ] , E [ W x K ( t ) ] , the minimum admissible disparity between W 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 disparity 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, when ∑ i ∈W e k,n ( t ) E [ r i (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 [ r i α ( ˆ 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 ∈W e k,n ( t ) E [ r i (2 ∑ l ∈ S e n,k ( σ ( ˆ T ) S ) n l -1) β ( ˆ T ) O ( i, · ) ,k ] by considering ∑ i ∈W e k,n ( t ) E [ r i α ( ˆ T ) O ( i, · ) ,k ] = 0 . As such, the minimum admissible disparity would be the case where b ⊤ ˆ k W x Q ( ˆ T ) b k and b ⊤ ˆ k W x K ( ˆ T ) b k both differ from β ( ˆ T ) Q,k = β ( ˆ T ) K,k by the amount of β -QK . Recall the definition of β -QK in Lemma 34, and collaborating with Lemma 9, we have the minimum admissible disparity be

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

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 / ( L Logist + λ ) , 1 / 2 λ } , where L Logist 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 ) = W t ( 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 [94, 34], we see that the Lipschitz smoothness of cross-entropy loss denotes that

Then we have that for s ≥ t +1 ,

where 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 } ,

From the following inequality, where

satisfying

where the last inequality hold clearly by expanding the product, the right hand side of the Eq.(64) is upper bounded as follows

We finally obtain the desired bound:

Remark 6. Utilizing this lemma, the exponential convergence over the 0-1 loss is readily obtained.

K Out-of-Distribution Generalization

Lemma 45. OOD 1: Master of Polysemy of Words . During testing, The prompt length L ∗ can be any positive integer. The D ∗ z can have any new probability distribution that differs from the training distribution, satisfying that each prompt has at least one co-concept k ∈ [ K 1 ] , with equal chance to have positive or negative semantic labels. Additionally, a single ( x , y ) ∼ D ∗ x ×D y ∗ 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 error L 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.

̸

Lemma 46. OOD 2: Innovation. During testing, the distribution of D ∗ x ×D ∗ y can enjoy data shift. Specifically, suggest we now have a new M ∗ and Q ∗ to define new D ∗ x , D ∗ y . Specifically, ∀ k = k ′ ∈ [ K 1 ] , k 2 ∈ [ K 2 ] , we let

and { a ∗ k , b ∗ k } K 1 k =1 , { c ∗ k , d ∗ k } K 1 k =1 are 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. Webelieve this can serve as an attempt to explain the innovation power of LLM [30, 95, 96] 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 [97, 4], heuristics [3], acquisition functions [98], and solutions to combinatorial optimization problems [99]. Although the resulting solutions may often seem counterintuitive to human experts, a possible explanation is that transformers can perform ICL in OOD scenarios by leveraging weighted combinations 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 [ M y Sn ] ⊆ [2 K 1 ] 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 [ M y 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

where the equality and inequality is by worse-case consideration over D ∗ z , a small σ 0 and λ in Condition 1 with a sufficiently large C , as well as the requirement ∥ b ∗ k ∥ ≥ ∥ a ∗ k ∥ = Θ( ∥ u ∥ ) . Besides, by ∥ d ∗ k ∥ ≥ ∥ c ∗ k ∥ = Θ( ∥ q ∥ ) , Eq.(65), Lemma 4, Eq.(5) and Lemma 2, we have that

Collaborating with Lemma 43, the poof is completed.

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.

Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:

You should answer [Yes] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.

IMPORTANT, please:

Delete this instruction block, but keep the section heading 'NeurIPS paper checklist" ,

Keep the checklist subsection headings, questions/answers and guidelines below.

Do not modify the questions and only use the provided macros for your answers .

Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The contributions and scope of this paper are well summarized in the abstract and introduction.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We explicitly discuss the limitation in Appendix A.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The detailed assumptions and proofs for all theorems and lemmas are given in the corresponding positions.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.

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.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Our algorithm is straightforward and easy to implement, and every detail is given in Section 3.

Guidelines:

The answer NA means that the paper does not include experiments.

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.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, 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.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary 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.

7. Experiment Statistical Significance

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.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, 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.

8. 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.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?

Answer: [Yes]

Justification: The research does not involve any human subjects, personal data, or interactions that would raise ethical concerns about consent, privacy, or respect for persons. In conclusion, the research aligns with the ethical principles outlined in the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We have discussed the broader impacts in Section A.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [No]

Justification: This paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have 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.

13. 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.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well 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

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

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.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- First work to show exponential convergence of 0-1 loss for a realistic transformer model: The authors prove exponential convergence for a three-layer model with **softmax** attention and **ReLU** MLP, which is more practical than previous simplified models. This is a significant theoretical advancement.\n- Solid mathematical analysis: The proof techniques seem rigorous and well-developed. The authors leverage advanced techniques to handle challenges like softmax attention and logistic loss, which have been difficult to analyze theoretically.\n- Experiments support theoretical findings: The simulations in Section 6 validate key aspects of the theory, like the evolution of attention weights and MLP coefficients. \n- Good motivation and connections to empirical observations: The authors motivate their work well by connecting to empirical findings on LLM representations (e.g., within-concept positive inner products). This grounds the theory in practical observations.\n\nWEAKNESSES:\n- Limited training setup: The paper only trains $W_Q, W_K, W_O$, keeping other weights fixed. This is quite restrictive compared to full transformer training and may limit the applicability of the results.\n- Overly specific data model: The concept-specific prompt distribution seems quite structured and simple. It's not clear how well this captures the complexity of real language data. A linear hyperplane can solve all tasks, including OOD ones, in this setup.\n- Lack of non-linear task consideration: The analysis focuses on essentially linear classification tasks. It would be more compelling to see how the model handles tasks requiring non-linear decision boundaries or composition of multiple concepts.\n\nQUESTIONS:\n- How sensitive are the results to the specific data distribution assumptions? Would the exponential convergence hold under more general conditions?\n- Have you considered extending the analysis to training all weights in the network? What challenges would this introduce?\n\nLIMITATIONS:\nIn 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

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

STRENGTHS:
I think that the results are pretty strong and I liked the proposition 1 which shows how the transformer can use its multi-concept semantics knowledge to perform in-context learning on unseen tasks.

WEAKNESSES:
I do not have any knowledge on the topic and so this is why my evaluation is going to be very shallow. Even though the theorems look sound to me, I think the authors could improve the presentation of their results by giving some proof sketch (+drawing) in the introduction.

QUESTIONS:
I think the authors should improve the presentation of their theory by giving more insights and intuition + simplified proof sketch.

LIMITATIONS:
The authors did not mention the limitations of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

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

STRENGTHS:
1. This paper is very original and investigates the interesting question of how a Transformer learns shared concepts within context.
2. This paper shows the clear OOD generalization of the Transformer model when trained on the sparse coding distribution.

WEAKNESSES:
1. The notation in the proof sketch is very heavy and not very easy to track. It would improve the paper if more explanation is given. Also, some of the statements seem counterintuitive. For example: 

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:
1. It seems that the proof considering mini-batch SGD is showing a concentration to the population gradient descent by Lemma 3. As mentioned in the intro, one of the key properties of the theoretical claim here is that it doesn't require a large batch size. Why is the technique here so robust to batch size?

LIMITATIONS:
The limitation is adequately addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-14

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

STRENGTHS:
The paper brings a new theoretical perspective linking multi-concept word semantics with in-context learning

WEAKNESSES:
The paper is sometimes hard to follow. More concise and clearer presentation, especially the theoretical content, would be helpful.
Lack of empirical experiments with more realistic datasets to show relevance in real-world setting.

QUESTIONS:
NA

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-14