Title: How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression
Venue: NeurIPS 2024
Accepted: yes
Keywords: in-context learning

transformers

deep learning theory

learning theory
Abstract: Despite the remarkable success of transformer-based models in various real-world tasks, their underlying mechanisms remain poorly understood. Recent studies have suggested that transformers can implement gradient descent as an in-context learner for linear regression problems and have developed various theoretical analyses accordingly. However, these works mostly focus on the expressive power of transformers by designing specific parameter constructions, lacking a comprehensive understanding of their inherent working mechanisms post-training. In this study, we consider a sparse linear regression problem and investigate how a trained multi-head transformer performs in-context learning. We experimentally discover that the utilization of multi-heads exhibits different patterns across layers: multiple heads are utilized and essential in the first layer, while usually only a single head is sufficient for subsequent layers. We provide a theoretical explanation for this observation: the first layer preprocesses the context data, and the following layers execute simple optimization steps based on the preprocessed context. Moreover, we demonstrate that such a preprocess-then-optimize algorithm can significantly outperform naive gradient descent and ridge regression algorithms. Further experimental results support our explanations. Our findings offer insights into the benefits of multi-head attention and contribute to understanding the more intricate mechanisms hidden within trained transformers.

How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression

Abstract

Despite the remarkable success of transformer-based models in various real-world tasks, their underlying mechanisms remain poorly understood. Recent studies have suggested that transformers can implement gradient descent as an in-context learner for linear regression problems and have developed various theoretical analyses accordingly. However, these works mostly focus on the expressive power of transformers by designing specific parameter constructions, lacking a comprehensive understanding of their inherent working mechanisms post-training. In this study, we consider a sparse linear regression problem and investigate how a trained multi-head transformer performs in-context learning. We experimentally discover that the utilization of multi-heads exhibits different patterns across layers: multiple heads are utilized and essential in the first layer, while usually one single head is dominantly utilized for subsequent layers. We provide a theoretical rationale for this observation: the first layer undertakes data preprocessing on the context examples, and the following layers execute simple optimization steps based on the preprocessed context. Moreover, we prove that such a preprocess-then-optimize algorithm can outperform naive gradient descent and ridge regression algorithms, which is also supported by our further experiments. Our findings offer insights into the benefits of multi-head attention and contribute to understanding the more intricate mechanisms hidden within trained transformers.

1 Introduction

Transformers [45] have emerged as a dominant force in machine learning, particularly in natural language processing. Transformer-based large language models such as Llama [42, 43] and the GPT family [36, 2, 12, 38], equipped with multiple heads and layers, showcasing exceptional learning and reasoning capabilities. One of the fundamental capabilities is in-context learning [12, 52], i.e., transformer can solve new tasks after prompting with a few context examples, without any further parameter training. Understanding their working mechanisms and developing reasonable theoretical explanations for their performance is vital and has gathered considerable research attention.

Numerous studies have been conducted to explore the expressive power of transformers, aiming to showcase their ability to tackle challenging tasks related to memorization [33], reasoning [24, 8,

28, 14], function approximation [26, 39], causal relationship [35], and simulating complex circuits [21, 31]. These endeavors typically aim to enhance our understanding of the capabilities and limitations of transformers when configured with varying numbers of heads[16] and layers. However, it's important to note that the findings regarding expressive power and complexity may not directly translate into explanations or insights into the behavior of trained transformer models in practical applications.

To perform a deeper understanding of the learning ability of transformer, a line of recent studies has been made to study the in-context learning performance of transformer by connecting it to certain iterative optimization algorithms [17, 3, 47]. These investigations have primarily focused on linear regression tasks with a Gaussian prior, demonstrating that a transformer with L layers can mimic L steps of gradient descent on the loss defined by contextual examples both theoretically and empirically[3, 53]. These observations have immediately triggered a series of further theoretical research, revealing that multi-layer and multi-head transformers can emulate a broad range of algorithms, including proximal gradient descent [8], preconditioned gradient descent [3, 50], functional gradient descent [15], Newton methods [22, 19], and ridge regression [8, 4]. However, these theoretical works are mostly built by designing specific parameter constructions, which may not reflect the key mechanism of trained transformers in practice. The precise roles of different transformer modules, especially for the various attention layers and heads, remain largely opaque, even within the context of linear regression tasks.

To this end, we take a deeper exploration regarding the working mechanism of transformer by investigating how transformers utilize multi-heads, at different layers, to perform the in-context learning. In particular, we consider the sparse linear regression task, i.e., the data is generated from a noisy linear model with sparse ground truth w ∗ ∈ R d with ∥ w ∗ ∥ 0 ≤ s ≪ d , and train a transformer model with multiple layers and heads. While a line of works also investigates this problem [20, 8, 1], understanding the key mechanisms behind trained transformers always requires more experimental and theoretical insights. Consequently, we empirically assess the importance of different heads at varying layers by selectively masking individual heads and evaluating the resulting performance degradation. Surprisingly, our observations reveal distinct utilization patterns of multi-head attention across layers of a trained transformer: in the first attention layer, all heads appear to be significant; in the subsequent layers, only one head appears to be significant. This phenomenon suggests that (1) employing multiple heads, particularly in the first layer, plays a crucial role in enhancing in-context learning performance; and (2) the working mechanisms of the transformer may be different for the first and subsequent layers.

Based on the experimental findings, we conjecture that muti-layer transformer may exhibit a preprocess-then-optimize algorithm on the context examples. Specifically, transformers utilize all heads in the initial layer for data preprocessing and subsequently employ a single head in subsequent layers to execute simple iterative optimization algorithms, such as gradient descent, on the preprocessed data. We then develop the theory to demonstrate that such an algorithm can be indeed implemented by a transformer with multiple heads in the first layer and one head in the remaining layers, and can achieve substantially lower excess risk than gradient descent and ridge regression (without data preprocessing). The main contributions of this paper are highlighted as follows:

We empirically investigate the role of different heads within transformers in performing in-context learning. We train a transformer model based on the data points generated by the noisy sparse linear model. Then, we reveal a distinct utilization pattern of multi-head attention across layers: while the first attention layer tended to evenly utilize all heads, subsequent layers predominantly relied on a single head. This observation suggests that the working mechanisms of multi-head transformers may vary between the first and subsequent layers.

Building upon our empirical findings, we proposed a possible working mechanism for multi-head transformers. Specifically, we hypothesized that transformers use the first layer for data preprocessing on in-context examples, followed by subsequent layers performing iterative optimizations on the preprocessed data. To substantiate this hypothesis, we theoretically demonstrated that, by constructing a transformer with mild size, such a preprocess-then-optimize algorithm can be implemented using multiple heads in the first layers and a single head in subsequent layers.

We further validated our proposed mechanism by comparing the performance of the preprocessthen-optimize algorithm with multi-step gradient descent and ridge regression solution, which can be implemented by the single-head transformers. We prove that the preprocess-then-optimize algorithm can achieve lower excess risk compared to these traditional methods, which is also verified by our numerical experiments. This aligns with our empirical findings, which indicated that multi-head transformers outperformed ridge regression in terms of excess risk.

To further validate our theoretical framework, we conducted additional experiments. Specifically, we performed probing on the output of the first layer of the transformer and demonstrated that representations generated by transformers with more heads led to lower excess risk after gradient descent. These experiments provided further support for our explanation on the working mechanism of transformers.

2 Preliminaries

Sparse Linear Regression. We consider sparse linear models where ( x , y ) ∼ P = P lin w ⋆ is sampled as x ∼ N ( 0 , Σ ) , y = ⟨ w ⋆ , x ⟩ + N ( 0 , σ 2 ) , where the Σ is a diagonal matrix and ground truth w ⋆ ∈ R d satisfies ∥ w ⋆ ∥ 0 ≤ s . Then, we define the population risk of a parameter w as follows:

Moreover, we are interested in the excess risk, i.e., the gap between the population risk achieved by w and the optimal one:

Multi-head Transformers. Transformers are a type of neural network with stacked attention and multi-layer perceptron (MLP) blocks. In each layer, the transformer first utilizes multi-head attention Attn to process the input sequence (or hidden states) H = [ h 1 , h 2 , . . . , h m ] ∈ R d hid × m . It computes h different queries, keys, and values, and then concatenates the output of each head:

where V i = W V i H , Q i = W Q i H , V i = V V i H and θ 1 = { W V i , W K i , W Q i ∈ R d hid /h × d hid } h i =1 are learnable parameters. The MLP then applies a nonlinear element-wise operation:

where θ 2 = { W 1 , W 2 } denotes the parameters of MLP. We remark that here some modules, such as layernorm and bias, are ignored for simplicity.

Linear Attention-only Transformers To perform an tractable theoretical investigation on the role of multi-head in the attention layer, we make further simplification on the transformer model by considering linear attention-only transformers. These simplifications are widely adopted in many recent works to study the behavior of transformer models [47, 53, 32, 3]. In particular, the i -th layer TF i performs the following update on the input sequence (or hidden state) H ( i -1) as follows:

where { W V i , W K i , W Q i ∈ R d hid h × d hid } h i =1 and W 1 ∈ R d hid × d hid are learnable parameters, note that as we ignore the ReLU activation in Eq.(2.1), so we merge the parameter W 1 and W 2 into one matrix W 1 . Besides, the mask matrix M is included in the attention to constrain the model focus the first n in-context examples rather than the subsequent m -n queries [3, 32, 54]. To adapt the transformer for solving sparse linear regression problems, we introduce additional linear layers W E ∈ R ( d +1) × d hid and W O ∈ R d hid × 1 for input embedding and output projection, respectively. Mathematically, let E denotes the input sequences with n in-context example followed by q queries,

Multi-Head Attention

Add & Norm

Embedding Projection

FFN

Add & Norm

Output Projection

𝐿 ×

Multi-Head Attention

Input

Input

First layer ( l = 1 )

preprocessing

subsequent layers ( l > 1 )

optimizing

𝑦1 …  𝑦 𝑛

0

⋯  0

𝑦𝑖 =

𝒙𝑖 , 𝒘

+

𝜉𝑖

𝒘 0 ≤ 𝑠

In-Context Sparse Linear regression

𝒙𝑖  ~ 𝑁

0, Σ

𝜉𝑖

~ 𝑁 0, 𝜎 2

(a) Overview of the experiments, including task, data, transformer architecture, and our insights.

5

10

15

20

25

in-context examples

0

1

2

3

4

5

excess risk

n_layer = 4,

= 0.1

ols

ridge

h = 1

h = 2

h = 4

h = 8

1

2

3

4

i-th head

4

3

2

1

j-th layer

h_num = 4, n_layer 4

0.1

0.2

0.3

0.4

0.5

0.6

(c) Heads Assessment

input

1

2

3

4

i-th layer

0

1

2

3

4

5

excess risk

h_num = 4, n_layer 4

= 0.1

pruned h = 4

h = 4

h = 1

(d) Pruning and Probing

Then model processes the input sequence E , resulting in the output ̂ y ∈ R 1 × ( n + q ) :

here, L is the layer number of the transformer, and ̂ y i + n is the prediction value for the query x i + n . During training, we set q > 1 for efficiency, and for inference and theoretical analysis, we set q = 1 and define the in-context learning excess risk E ICL as:

Notations. For two functions f ( x ) ≥ 0 and g ( x ) ≥ 0 defined on the positive real numbers ( x > 0 ), we write f ( x ) ≲ g ( x ) if there exists two constants c, x 0 > 0 such that ∀ x ≥ x 0 , f ( x ) ≤ c · g ( x ) ; we write f ( x ) ≳ g ( x ) if g ( x ) ≲ f ( x ) ; we write f ( x ) ⋍ g ( x ) if f ( x ) ≲ g ( x ) and g ( x ) ≲ f ( x ) . If f ( x ) ≲ g ( x ) , we can write f ( x ) as O ( g ( x )) . We can also write write f ( x ) as ˜ O ( g ( x )) if there exists a constant k > 0 such that f ( x ) ≲ g ( x ) log k ( x ) .

3 Experimental Insights into Multi-head Attention for In-context Learning

While previous work has demonstrated the in-context learning ability for sparse linear regression [20, 8], the hidden mechanism behind the trained transformer for solving this problem remains unclear. To this end, we design a series of experiments, utilizing techniques like probing [5] and pruning [27] to help us gain initial insights into how the trained transformer utilizes multi-head attention for this problem. For all experiments in Sections 3 and 6, we choose an encoder-based architecture as the backbone (see Figure 1a), set the hidden dimension d hid to 256, and use the input sequence format shown in Eq.(2.3), where d = 16 , s = 4 , x ∼ N ( 0 , I ) , with varying noise levels, layers, and heads, Additional experimental details can be found in Appendix B. The experiments we designed are as follows:

ICL with Varying Heads: First, based on the experiment results by [8], we further investigate the performance of transformers in solving the in-context sparse linear regression problem with varying attention heads. An example can be found in Figure 1b, where we display the excess risk for different models when using different numbers of in-context examples. We can observe that given few-shot in-context examples, transformers can outperform OLS and ridge. Moreover, we can also clearly observe the benefit of using multiple heads, which leads to lower excess risk when increasing

the number of heads. This highlights the importance of multi-head attention in transformer to perform in-context learning .

Heads Assessment: Based on Eq.(2.2), we know that the j -th head at the i -th layer corresponds to the subspace of the intermediate output from ( j -1) · d hid /h to j · d hid /h -1 . To assess the importance of each attention head, we can mask the particular head by zeroing out the corresponding output entries, while keeping other dimensions unchanged. Then, let ( i, j ) be the layer and head indices, we evaluate the risk change before and after head masking, denoted by ∆ E ICL ( i,j ) . Then we normalize the risk changes in the same layer to evaluate their relative importance:

An example can be found in Figure 1c. We can observe that in the first layer, no head distinctly outweighs the others, while in the subsequent layers, there always exists a head that exhibits higher importance than others. This gives us insight that in the first attention layer, all heads appear to be significant, while in the subsequent layers, only one head appears to be significant .

Pruning and Probing: To further validate our findings in the previous experiments, we prune the trained model by (1) retaining all heads in the first layer; and (2) only keeping the most important head and zeroing out others for the subsequent layers. Then the pruned model, referred to as the 'pruned transformer', will be fine-tuned with with the same training data. We then use linear probes [6] to evaluate the prediction performance for different layers. An example can be found in Figure 1d, we can find that the 'pruned transformer' and the original model exhibit almost the same performance for each layer. Additionally, compared to the model with single-head attention, we observe that the probing result is largely different between single-head transformers and the 'pruned transformers', the latter has better performances compared to the former. Noting that the main difference between them is the number of heads in the first layer (subsequent layers have the same structure), it can be deduced that the working mechanisms of the multi-head transformer may be different for the first and subsequent layers .

4 Potential Mechanism Behind Trained transformer

Based on the experimental insights from Section 3, we found that all heads in the first layer of the trained transformer are crucial, while in subsequent layers, only one head plays a significant role. Furthermore, by checking the result for probing and pruning, we can find that the working mechanisms of the transformer may be different for the first and subsequent layers. To this end, we hypothesize that the multi-layer transformer may implement a preprocess-then-optimize to perform the in-context learning, i.e., the transformer first performs preprocessing on the in-context examples using the first layer and then implements multi-step iterative optimization algorithms on the preprocessed in-context examples using the subsequent layers.

We note that [24] adapts a similar two-phase idea to explain how transformer learning specific functions in context, in their constructed transformers, the first few layers utilize MLPs to compute an appropriate representation for each entry, while the subsequent layers utilize the attention module to implement gradient descent over the context. We highlight that our algorithm mainly focus on utilizing multihead attention, and it aligns well with the our experimental observation and intuition. The details of our algorithm are as follows:

4.1 Preprocessing on In-context Examples

First, as the multihead attention is designed to facilitate to model to capture features from different representation subspaces [45], we abstract the algorithm implementation by the first layer of the transformers as a preprocessing procedure. In general, for the sparse linear regression, a possible data preprocessing method is to perform reweighting of the data features by emphasizing the features that correspond to the nonzero entries of the ground truth w ∗ and disregard the remaining features. In the idealized case, if we know the nonzero support of w ∗ , we can trivially zero out the date features of x on the complement of the nonzero support, as a data preprocessing procedure, and perform projected gradient descent to obtain the optimal solution.

4: end for

5: Output : Sequence with the preprocessed in-context examples/queries { ( ˜ x i , y i ) } n i =1 , { ( ˜ x i , 0) } n + q i = n +1 .

The preprocessing procedure aligns well with the structure of a multi-head attention layer with linear attention, which motivates our theoretical construction of the desired transformer. In particular, each head of the attention layer can be conceptualized as executing specific operations on a distinct subset of data entries. Then, the linear query-key calculation, represented as ( W K i H ) ⊤ W Q i H , where H = E denotes the input sequence embedding matrix, effectively estimates correlations between the i -th subset of data entries and the corresponding label y i . Here, W K i and W Q i selectively extract entries from the i -th subset of features and the label, respectively, akin to an "entries selection" process. Furthermore, when combined with the value calculation W V i H , each head of the attention layer conducts correlation calculations for the i -th subset of features and subsequently employs them to reweight the original features within the same subset. Consequently, by stacking the outputs of multiple heads, all data features can be reweighted accordingly, which matches the design of the proposed preprocessing procedure in Alg. 1. We formally prove this in the following theorem.

Proposition 4.1 (Single-layer multi-head transformer implements Alg. 1) . There exists a single-layer transformer function TF 1 , with d heads and d hid = 3 d hidden dimension, together with an input embedding layer with weight W E ∈ R d hid × d , that can implement Alg. 1. Let E be the input sequence defined in Eq. (2.3) and ˜ x i = ̂ Rx be the preprocessed features defined in Alg. 1, it holds that

where · · · in third row implies arbitrary values.

4.2 Optimizing Over Preprocessed In-Context Examples

Based on the experimental results, we observe that the subsequent layers of transformers dominantly rely on one single head, suggesting their different but potentially simpler behavior compared to the first one. Motivated by a series of recent work [47, 15, 53, 3] that reveal the connection between gradient descent steps and multi-layer single-head transformer in the in-context learning tasks, we conjecture that the subsequent layers also implement iterative optimization algorithms such as gradient descent on the (preprocessed) in-context examples.

To maintain clarity in our construction and explanation, in each layer, we use a linear projection W ( i ) 1 to rearrange the dimensions of the sequence processed by the multi-head attention, resulting in the hidden state H ( i ) of each layer. We refer to the first d rows of the input data as x , and the ( d +1) -th row as the corresponding y . For example, in Eq.(4.2), we take the first d rows, together with the

̸

In general, the nonzero support of w ∗ is intractable to the learner, so that one cannot perform idealized masking-related data preprocessing. However, one can still perform estimations on the importance of data features by examining their correlation with the target. In particular, note that we have y = ⟨ w ∗ , x ⟩ + ξ i = ∑ d i =1 w ∗ i x i + ξ i , implying that r i := E [ x i y ] = E [ ∑ d i =1 w ∗ i x i · x i ] + E [ ξx i ] = w ∗ i E [ x 2 i ] if considering independent data features. Then it is clear that such a correlation between the feature and label will be nonzero only when | w ∗ i | = 0 . Therefore, instead of knowing the nonzero support of w ∗ , we can instead calculate such a correlation to perform reweighting on the data features. Noting that the transformer is provided with n in-context examples { ( x i , y i ) } n i =1 , such correlations can be estimated accordingly: ̂ r j = 1 n ∑ n i =1 x ij y i , which will be further used to perform the data preprocessing on the in-context examples. We summarize this procedure in Alg. 1.

Algorithm 1 Data preprocessing for in-context examples

1: Input : Sequence with { ( x i , y i ) } n i =1 , { ( x i , 0) } n + q i = n +1 as in-context examples/queries.

2: for k = 1 , . . . , n do

3: Compute ˜ x k by ˜ x k = ̂ Rx k , where ̂ R = diag { ̂ r 1 , ̂ r 2 , . . . , ̂ r d } , where ̂ r j is given by

( d +1) -th row, as the input data entry { ˜ x i , y i } n +1 i =1 . Then, the following proposition shows that the subsequent layers of transformer can implement multi-step gradient descent on the preprocessed in-context examples { ( ˜ x i , y i ) } i =1 ,...,n .

Proposition 4.2 (Subsequent single-head transformer implements multi-step GD) . There exists a transformer with k layers, 1 head, d hid = 3 d , let ̂ y ℓ n +1 be the prediction representation of the ℓ -th layer, then it holds that ̂ y ℓ ( n +1) = ⟨ w ℓ gd , ˜ x n +1 ⟩ , where ˜ x n +1 = ̂ Rx n +1 denotes the preprocessed data feature, w ℓ gd is defined as w 0 gd = 0 and as follows for ℓ = 0 , . . . , k -1 :

The proof of Propositions 4.1 and 4.2 can be found in Appendix C, combining these two propositions , we show that the multi-layer transformer with multiple heads in the first layer and one head in the subsequent layers can implement the proposed preprocess-then-optimization algorithm. In the next section, we will establish theories to demonstrate that such an algorithm can indeed achieve smaller excess risk than standard gradient descent and ridge regression solutions of the sparse linear regression problem.

5 Excess Risk of the Preprocess-then-optimize Algorithm

In this section, we will develop the theory to demonstrate the improved performance of the preprocessthen-optimize algorithm compared to the gradient descent algorithm on the raw inputs. The proof for Theorem 5.1, 5.2, and 5.3 can be found in Appendix D, E, and F, respectively.

We first denote ˜ w t gd as the estimator obtained by t -step GD on { ( ˜ x i , y i ) } n i =1 , which can be viewed as the solution generated by the t +1 -layer transformer based on our discussion in Section 4, and w t gd as the estimator obtained by t -step GD on { ( x i , y i ) } n i =1 . Before presenting our main theorem, we first need to redefine the excess risk of GD on { ( ˜ x i , y i ) } n i =1 . Note that in our algorithm, the learned predictor takes the form x →⟨ ̂ Rx , ˜ w t gd ⟩ . Consequently, the population risk of a parameter ˜ w t gd is naturally defined as ˜ L ( ˜ w t gd ) := 1 2 · E ( x ,y ) ∼ P [ ( ⟨ ̂ Rx , ˜ w t gd ⟩ -y ) 2 ] , and the excess risk is then defined as E ( w ) := ˜ L ( w ) -min w ˜ L ( w ) 2 . Next, we provide the upper bound of the excess risk for E ( ˜ w t gd ) and E ( w t gd ) respectively.

̸

Theorem 5.1. Denote S := { i : w ⋆ i = 0 } and R = diag { r 1 , . . . , r d } , where r j = ∑ d i =1 w ⋆ i Σ ij . Suppose that there exist a β > 0 such that min i ∈S | r i | ≥ β , ∥ R ∥ 2 , ∥ Σ ∥ 2 , ∥ w ⋆ ∥ 2 ≃ O (1) and n ≳ 1 /β 2 · t 2 s · ( Tr 2 / 3 ( Σ ) + Tr( RΣR ) ) · poly(log ( d/δ )) . Then set η ≲ 1 / ∥ RΣR ∥ 2 and

it holds that

with probability at least 1 -δ .

Theorem 5.1 provides an upper bound on the excess risk achieved by the preprocess-then-optimize algorithm, where we tuned learning rate η to balance the bias and variance error. Then, it can be seen that the risk bound is valid if Tr( RΣR ) /n → 0 and Tr( Σ ) s/n 2 → 0 when n →∞ . This can be readily satisfied if we have ∥ w ∗ ∥ 2 and Tr( Σ ) be bounded by some reasonable quantities that are independent of the sample size n , which are the common assumptions made in many prior works [58, 57, 9]. Besides, it can be also seen that the excess risk bound explicitly depends on the sparsity parameter s and lower sparsity implies better performance. This implies the ability of the proposed preprocess-then-optimize for discovering and leveraging the nice sparse structure of the ground truth.

As a comparison, the following theorem states the excess risk bound for the standard gradient descents on the raw features. To make a fair comparison, we consider using the same number of steps but allow the step size to be tuned separately.

Theorem 5.2. Suppose that ∥ Σ ∥ , ∥ w ⋆ ∥ 2 ≃ O (1) and n ≳ t 2 (Tr( Σ ) + log (1 /δ )) . When η ≲ 1 / ∥ Σ ∥ 2 and ηt ≃ ( σ 2 Tr( Σ ) log ( d/δ ) n ) -1 / 2 , it holds that

with probability at least 1 -δ .

We are now able to make a rough comparison between the excess risk bounds in Theorems 5.1 and 5.2. Then, it is clear that E ( ˜ w t gd ) ≲ E ( w t gd ) requires Tr( RΣR ) /β 2 ≲ Tr( Σ ) and s/ ( n 2 β 2 ) ≤ 1 /n . Specifically, we can consider the case that Σ to be a diagonal matrix, assume w ⋆ i ∼ U {-1 / √ s, 1 / √ s } has a restricted uniform prior for i ∈ S and min i ∈S Σ ii ≥ 1 /κ for some constant κ > 1 , we can get β ≥ √ 1 / ( sκ 2 ) , thus Tr( RΣR ) /β 2 ≤ κ 2 ∑ i : w ⋆ i =0 Σ ii and s/ ( n 2 β 2 ) ≤ κ 2 s 2 /n 2 . Note that |S| = s ≪ d , then if the covariance matrix Σ has a flat eigenspectrum such that ∑ i ∈S Σ ii ≪ ∑ i ∈ [ d ] Σ ii = Tr( Σ ) , we have Tr( RΣR ) /β 2 ≤ Tr( Σ ) and s/ ( n 2 β 2 ) ≤ κ 2 s 2 /n if s = o ( min { d, √ n } ) . This suggests that the preprocess-then-optimization algorithm can outperform the standard gradient descent for solving a sparse linear regression problem with s = o ( min { d, √ n } ) .

̸

To make a more rigorous comparison, we next consider the example where x i i.i.d. ∼ N ( 0 , I ) , based on which we can get the upper bound for our algorithm and the lower bound for OLS, ridge regression, and finite-step GD.

Theorem 5.3. Suppose S with |S| = s is selected such that each element is chosen with equal probability from the set { 1 , 2 , . . . , d } and w ⋆ i ∼ U {-1 / √ s, 1 / √ s } has a restricted uniform prior for i ∈ S , ∥ w ⋆ ∥ 2 ≃ Θ(1) and n ≳ t 2 s 3 d 2 / 3 . Then there exists a choice of η and t such that

with probability at least 1 -δ . Besides, let ̂ w λ be the ridge regression estimator with regularized parameter λ , and w ols be the OLS estimator, it holds that

with probability at least 1 -δ , where w ∈ { ̂ w λ , w ols , w t gd } .

It can be seen that for a wide range of under-parameterized and over-parameterized cases, ˜ w t gd has a smaller excess risk than ridge regression, standard gradient descent, and OLS. In particular, consider the setting σ 2 = 1 , in the over-parameterized setting that d ≳ n , the excess risk bound of preprocess-then-optimize is ˜ O ( ds/n 2 ) , which also outperforms the ˜ Ω(1) bound achieved by OLS, ridge regression, and standard gradient descent if the sparsity satisfies s = O ( n 2 /d ) (in fact, this condition can be certainly removed as E ( ˜ w t gd ) also has a naive upper bound ˜ O (1) ). In the underparameterized case that d ≲ n , it can be readily verified that the data preprocessing can lead to a ˜ O ( s/n ) excess risk, which is strictly better than the ˜ Ω( d/n ) risk achieved by OLS, ridge regression, and standard gradient descent. Moreover, it is well known that Lasso can achieve ˜ O ( s/n ) excess risk bound in the setting of Theorem 5.3. Then, by comparing with our results, we can also conclude that the proprocess-then-optimize algorithm can be comparable to Lasso up to logarithmic factors when d ≲ n , while becomes worse when d ≳ n .

6 Experiments

In Section 3, we conduct several experiments, and based on the observations, we propose that a trained transformer can apply a preprocess-then-optimize algorithm: (1) In the first layer, the transformer can apply a preprocessing algorithm (Alg. 1) on the in-context examples utilizing multi-head attention.

𝑇𝐹1

𝑥1

𝑥2

…

𝑥𝑛

…

Embedding Projection

𝑥𝑛+1   …

𝑥𝑛+𝑞

𝑇𝐹2

𝑇𝐹𝐿

𝐸  =

𝐻 1 =

෤

𝑥𝑛+1 … ෤

𝑥𝑛+𝑞-1

෤

𝑥𝑛+𝑞

𝑦𝑛+1  …

𝑦𝑛+𝑞-1

𝑦𝑛+𝑞

K

-step GD

𝑤gd

𝑘

ො

𝑦𝑛+𝑞

𝑘

ℰ 𝑘

=

(

ො

𝑦𝑛+𝑞

𝑘

-

𝑦𝑛+𝑞) 2

-𝜎 2

extract

query

entries

Trained Transformer

Preprocessing Probing

(a) P-probing

0

1

2

3

4

5

6

7

8

9 ×10

steps

2.0

2.5

3.0

3.5

4.0

excess risk

= 0

h = 1

h = 4

h = 8

0

1

2

3

4

5

6

7

8

9 ×10

steps

2.0

2.5

3.0

3.5

4.0

= 0.1

h = 1

h = 4

h = 8

0

1

2

3

4

5

6

7

8

9 ×10

steps

2.0

2.5

3.0

3.5

4.0

= 0.2

h = 1

h = 4

h = 8

0

1

2

3

4

5

6

7

8

9 ×10

steps

2.0

2.5

3.0

3.5

4.0

= 0.4

h = 1

h = 4

h = 8

0

1

2

3

4

5

6

7

8

9 ×10

steps

2.0

2.5

3.0

3.5

4.0

= 0.8

h = 1

h = 4

h = 8

(b) P-Probing result for Transformers trained on varying heads and data settings

0

1

2

3

4

5

6

7

8

9×10 3

steps

1.0

1.5

2.0

2.5

3.0

3.5

4.0

excess risk

= 0

gd

0

1

2

3

4

5

6

7

8

9×10 3

steps

1.0

1.5

2.0

2.5

3.0

3.5

4.0

= 0.1

gd

0

1

2

3

4

5

6

7

8

9×10 3

steps

1.0

1.5

2.0

2.5

3.0

3.5

4.0

= 0.2

gd

0

1

2

3

4

5

6

7

8

9×10 3

steps

1.5

2.0

2.5

3.0

3.5

4.0

= 0.4

gd

0

1

2

3

4

5

6

7

8

9×10 3

steps

2.0

2.5

3.0

3.5

4.0

= 0.8

gd

(c) Directly apply Alg. 1 on input data for gradient descent

(2) In the subsequent layers, the transformer applies a gradient descent algorithm on the preprocessed data utilizing single-head attention. While the second part is supported by extensive theoretical analysis and experimental evidence [47, 15, 53, 3], here we develop a technique called preprocessing probing (P-probing) on the trained transformer to support the first part of our algorithm. We also directly apply Alg. 1 on the in-context examples and then check the excess risk for multiple-step gradient descent to verify the effectiveness of our algorithm and theoretical analysis.

P-probing: To verify the existence of a preprocessing procedure in the trained transformer, we develop a 'preprocessing probing' (P-probing) technique on the trained transformers, as illustrated in Figure 2a. For a trained transformer, we first set the input sequence as in Eq.(2.3), where the first n examples { x } n i =1 have the corresponding labels { y } n i =1 , and the following q query entries only have { x i } n + q i = n +1 in the sequence. Then, we extract the last q vectors in the output hidden state H 1 from the first layer of the transformer and treat these data as processed query entries. Next, we conduct gradient descent on the first q -1 query entries with their corresponding y , computing the excess risk on the last query. Additional experimental details can be found in Appendix B. We adapt this technique based on the intuition that, according to our theoretical analysis, we can extract the preprocessed entry { ˜ x i } n + q i = n +1 from H 1 , besides, the excess risk computed by the preprocessed data has a better upper bound guarantee compared to raw data without preprocessing under the same number of gradient descent steps, so if the trained transformer utilize multihead attention for preprocess, compared with single head attention, the queries entries extract from H 1 by multihead attention can have better gradient descent performance compared with single head attention.

Verifying the benefit of preprocessing: To further support the effectiveness of our algorithm, we directly apply Alg. 1 on the input data { x i , y i } n +1 i =1 , and then implement gradient descent on the example entries { x i , y i } n i =1 and compute the excess risk with the last query { ̂ x n +1 , y n +1 } , we refer this procedure as pre-gd . We compare pre-gd with the excess risk obtained by directly applying gradient descent without preprocessing (referred to as gd ). For all experiments (both P-probing and this), we set w 0 gd = 0 and tune the learning rate η for each model by choosing from [1 , 10 -1 , 10 -2 , 10 -3 , 10 -4 , 10 -5 , 10 -6 ] with the lowest average excess risk.

Based on Figure 2b, we can observe that compared to the transformer with single-head attention ( h = 1 ), the query entries extracted from the transformer with multiple heads ( h = 4 , 8 ) preserve better convergence performance and can dive into a lower risk. This aligns well with our experiment result in Figure 2c, where compared to gd , the data preprocessed by Alg. 1 preserves better convergence performance and can dive into a lower risk space, supporting the existence of the preprocessing procedure in the trained transformer. Moreover, Figure 2c also aligns well with our theoretical analysis, where we provide a better upper bound for convergence guarantee for our algorithm compared to ridge regression and OLS.

7 Conclusions and Limitations

In this paper, we investigate a sparse linear regression problem and explore how a trained transformer leverages multi-head attention for in-context learning. Based on our empirical investigations, we propose a preprocess-then-optimize algorithm, where the trained transformer utilizes multi-head attention in the first layer for data preprocessing, and subsequent layers employ only a single head for optimization. We theoretically prove the effectiveness of our algorithm compared to OLS, ridge regression, and gradient descent, and provide additional experiments to support our findings.

While our findings provide promising insights into the hidden mechanisms of multi-head attention for in-context learning, there is still much to be explored. First, our work focuses on the case of sparse linear regression, and it may be beneficial to implement our experiment for more challenging or even real-world tasks. Additionally, as we adapt attention-only transformers for analysis simplification, the role of other modules, such as MLPs, are neglected. How these modules incorporate in real-world tasks remains unclear. Moreover, our analysis does not consider the training dynamics of transformers, while the theoretical analysis in [13] provides valuable insights into the convergence of single-layer transformers with multi-head attention, the training dynamics for multi-layer transformers remain unclear. How transformers learn to implement these algorithms is worth further investigation.

References

Jacob Abernethy, Alekh Agarwal, Teodor Vanislavov Marinov, and Manfred K. Warmuth. A mechanism for sample-efficient in-context learning for sparse retrieval tasks. In Claire Vernade and Daniel Hsu, editors, Proceedings of The 35th International Conference on Algorithmic Learning Theory , volume 237 of Proceedings of Machine Learning Research , pages 3-46. PMLR, 25-28 Feb 2024. URL https://proceedings.mlr.press/v237/abernethy24a. html .

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems , 36, 2024.

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations , 2022.

Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.01644 , 2016.

Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes, 2017. URL https://openreview.net/forum?id=ryF7rTqgl .

Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.2, knowledge manipulation. arXiv preprint arXiv:2309.14402 , 2023.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection, July 2023.

Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences , 117(48):30063-30070, 2020.

Steven Bills, Nick Cammarata, Dan Mossing, Henk Tillman, Leo Gao, Gabriel Goh, Ilya Sutskever, Jan Leike, Jeff Wu, and William Saunders. Language models can explain neurons in language models. URL https://openaipublic. blob. core. windows. net/neuronexplainer/paper/index. html.(Date accessed: 14.05. 2023) , 2023.

Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, et al. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread , page 2, 2023.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

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.

Xingwu Chen and Difan Zou. What can transformer learn with varying depth? case studies on sequence learning tasks. arXiv preprint arXiv:2404.01601 , 2024.

Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.06528 , 2023.

Yingqian Cui, Jie Ren, Pengfei He, Jiliang Tang, and Yue Xing. Superiority of multi-head attention in in-context linear regression. arXiv preprint arXiv:2401.17426 , 2024.

Nan Ding, Tomer Levinboim, Jialin Wu, Sebastian Goodman, and Radu Soricut. CausalLM is not optimal for in-context learning, September 2023.

Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs. Advances in Neural Information Processing Systems , 36, 2024.

Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers Learn Higher-Order Optimization Methods for In-Context Learning: A Study with Linear Models, October 2023.

Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What Can Transformers Learn In-Context? A Case Study of Simple Function Classes, August 2023.

Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped Transformers as Programmable Computers, January 2023.

Angeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D. Lee. How Well Can Transformers Emulate In-context Newton's Method?, March 2024.

Andrey Gromov, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel A Roberts. The unreasonable ineffectiveness of the deeper layers. arXiv preprint arXiv:2403.17887 , 2024.

Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations, October 2023.

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

Tokio Kajitsuka and Issei Sato. Are Transformers with One Layer Self-Attention Using LowRank Weight Matrices Universal Approximators?, July 2023.

Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient convnets. In International Conference on Learning Representations , 2017. URL https://openreview.net/forum?id=rJqFGTslg .

Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning , pages 19565-19594. PMLR, 2023.

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

David Lindner, János Kramár, Sebastian Farquhar, Matthew Rahtz, Thomas McGrath, and Vladimir Mikulik. Tracr: Compiled Transformers as a Laboratory for Interpretability, November 2023.

Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=De4FYqjFueZ .

Arvind Mahankali, Tatsunori B. Hashimoto, and Tengyu Ma. One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention, July 2023.

Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization Capacity of MultiHead Attention in Transformers, October 2023.

Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. Shortgpt: Layers in large language models are more redundant than you expect. arXiv preprint arXiv:2403.03853 , 2024.

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

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems , 35:27730-27744, 2022.

Onkar Pandit and Yufang Hou. Probing for bridging inference in transformer language models. arXiv preprint arXiv:2104.09400 , 2021.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Shokichi Takakura and Taiji Suzuki. Approximation and Estimation Ability of Transformers for Sequence-to-Sequence Functions with Infinite Dimensional Input, May 2023.

Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as Support Vector Machines, September 2023.

Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer, July 2023.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023.

Alexander Tsigler and Peter L Bartlett. Benign overfitting in ridge regression. Journal of Machine Learning Research , 24(123):1-76, 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need, August 2023.

Roman Vershynin. High-dimensional probability. University of California, Irvine , 10:11, 2020.

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, May 2023.

Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking Like Transformers, July 2021.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771 , 2019.

Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? In The Twelfth International Conference on Learning Representations , 2023.

Zhiyong Wu, Yun Chen, Ben Kao, and Qun Liu. Perturbed masking: Parameter-free probing for analyzing and interpreting bert. arXiv preprint arXiv:2004.14786 , 2020.

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations , 2021.

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

Ruiqi Zhang, Jingfeng Wu, and Peter L Bartlett. In-context learning of a linear transformer block: Benefits of the mlp component and one-step gd initialization. arXiv preprint arXiv:2402.14951 , 2024.

Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Joshua M. Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id=AssIuHnmHX .

Zeyuan Allen Zhu and Yuanzhi Li. Physics of language models: Part 3.1, knowledge storage and extraction. arXiv preprint arXiv:2309.14316 , 2023.

Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, Dean P Foster, and Sham Kakade. The benefits of implicit regularization from sgd in least squares problems. Advances in neural information processing systems , 34:5456-5468, 2021.

Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham Kakade. Risk bounds of multi-pass sgd for least squares in the interpolation regime. Advances in Neural Information Processing Systems , 35:12909-12920, 2022.

A Additional Related Work

In addition to works towards understanding the expressive power of transformers that we introduced before, there is also a body of research on the mechanism interpretation and the training dynamics of transformers:

Mechanism interpretation of trained transformers To understand the mechanisms in trained transformers, researchers have developed various techniques, including interpreting transformers into programming languages [18, 30, 48, 55], probing the behavior of individual layers [37, 51, 11, 7, 56], and incorporating transformers with other large language models to interpret individual neurons [10]. While these techniques provide high-level insights into transformer mechanism understanding, providing a clear algorithms behind the trained transformers is still very challenging.

Training dynamics of transformers In parallel, a body of work has also investigated how transformers learn these algorithms, i.e., the training dynamics of transformers. Tarzanagh et al. [40] shows an equivalence between a single attention layer and a support vector machine. Zhang et al. [53], Ahn et al. [3] analyze the training dynamics of a single-head attention layer for in-context linear regression, where [53] demonstrates that it can converge to implement one-step gradient over in-context examples. Huang et al. [25], Chen et al. [13] extended these findings from linear attention to softmax settings, with [13] revealing that trained transformers with multi-head attention tend to utilize different heads for distinct tasks in various subspaces. Additionally, Tian et al. [41], Li et al. [29] study the convergence of transformers on sequences of discrete tokens. Gromov et al. [23], Men et al. [34] use experiments to show that in large language models, parameters in deeper layers are less critical compared to those in shallower layers. These works provide valuable insights towards the theoretical understanding of the training dynamics of transformers, which offer potential future extension aspects for our work.

B Additional Details for Sections 3 and 6

Architecture and Optimization We conduct extensive experiments on encoder-only transformers with d hid = 256 , varying the number of heads h ∈ { 1 , 2 , 4 , 8 } , layers l ∈ { 3 , 4 , 5 , 6 } , and noise levels σ ∈ { 0 , 0 . 1 , 0 . 2 , 0 . 4 , 0 . 8 } . For the input sequence, we sample x ∼ N ( 0 , I ) . For w , we first sample w ∼ N ( 0 , I ) ∈ R 16 , and randomly choose s = 4 entries, setting the other elements to zero. Note that We don't apply positional encodings in our setting, as no positional information is needed in our input setting. To further support our preprocessing-then-optimize algorithm, we also try a decoder-only architecture(Figure 9), train models on other settings like stardand linear regression task s = d = 16 (Figure 10) and non-orthogonal data distributions (Figure 11) as comparisons in Appendix G. During training, we set n = 12 and q = 4 , with a batch size of 64 . We utilize the Adam optimizer with a learning rate γ = 10 -4 for 320000 updates. Each experiment takes about two hours on a single NVIDIA GeForce RTX 4090 GPU. We fix the random seed such that each model is trained and evaluated with the same training and evaluation dataset. We use HuggingFace [49] library to implement our models.

ICL with Varying Heads We compare the model's performance with ridge regression, OLS, and lasso. For ridge regression and lasso, we tune λ, α ∈ { 1 , 10 -1 , 10 -2 , 10 -3 , 10 -4 } respectively for the lowest risk, as in [20].

From Figure 3, we can find that in most cases, transformers with single head ( h = 1 ) exhibits higher risk compared to models with multiple heads ( h = 4 , 8 ). Note that in thesame subplot, models with different numbers of heads have the same number of parameters. This experiment highlights the importance of multi-head attention for transformers in in-context learning.

Heads Assessment Here, we set n = 10 and q = 1 , with an evaluation data size of 8192 . For a model with h heads and l layers, we train | σ | models under different noise levels. We first compute the W h,l,σ under different noise levels σ , then sort each row in W h,l,σ , and add them together as W h,l avg = 1 | σ | ∑ σ ∈ σ W h,l,σ , resulting in the final weight for each head. An example can be found in Fig 1c. In Fig 4, we present more results for different h and l , and we also present the heat map for the decode-only transformers in Figure 9.

5

10

15

20

25

0

1

2

3

4

5

excess risk

n_layer = 3,

= 0

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 4,

= 0

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 5,

= 0

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 6,

= 0

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

excess risk

n_layer = 3,

= 0.1

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 4,

= 0.1

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 5,

= 0.1

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 6,

= 0.1

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

excess risk

n_layer = 3,

= 0.2

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 4,

= 0.2

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 5,

= 0.2

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 6,

= 0.2

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

excess risk

n_layer = 3,

= 0.4

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 4,

= 0.4

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 5,

= 0.4

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

1

2

3

4

5

n_layer = 6,

= 0.4

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

in-context examples

0

1

2

3

4

5

excess risk

n_layer = 3,

= 0.8

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

in-context examples

0

1

2

3

4

5

n_layer = 4,

= 0.8

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

in-context examples

0

1

2

3

4

5

n_layer = 5,

= 0.8

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

in-context examples

0

1

2

3

4

5

n_layer = 6,

= 0.8

ols

ridge

h = 1

h = 2

h = 4

h = 8

1

2

3

4

3

2

1

j-th layer

h_num = 4, n_layer 3 avg

1

2

3

4

4

3

2

1

h_num = 4, n_layer 4 avg

1

2

3

4

5

4

3

2

1

h_num = 4, n_layer 5 avg

1

2

3

4

6

5

4

3

2

1

h_num = 4, n_layer 6 avg

1

2

3

4

5

6

7

8

i-th head

3

2

1

j-th layer

h_num = 8, n_layer 3 avg

1

2

3

4

5

6

7

8

i-th head

4

3

2

1

h_num = 8, n_layer 4 avg

1

2

3

4

5

6

7

8

i-th head

5

4

3

2

1

h_num = 8, n_layer 5 avg

1

2

3

4

5

6

7

8

i-th head

6

5

4

3

2

1

h_num = 8, n_layer 6 avg

0.15

0.20

0.25

0.30

0.35

0.40

0.45

0.1

0.2

0.3

0.4

0.5

0.1

0.2

0.3

0.4

0.5

0.6

0.1

0.2

0.3

0.4

0.5

0.6

0.1

0.2

0.3

0.4

0.5

0.6

0.1

0.2

0.3

0.4

0.5

0.6

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.1

0.2

0.3

0.4

0.5

0.6

0.7

From Fig 4, we can find that in most settings, each head contributes almost equally, while in the subsequent layers, there always exists a head that has a much larger weight than the others. This

indicates that in trained transformers for in-context learning, in the first attention layer, all heads appear to be significant, while in the subsequent layers, only one head appears to be significant.

Pruning and Probing Here, we also set n = 10 and q = 1 , with an evaluation data size of 8192 . To further support our finding from the Head Assessment, we first prune the model based on our computed head weight W h,l avg , where we keep all heads in the first layer, whereas we only keep the head with the highest score weight and mask the others. We then train the pruned model with the same method as before for 60000 steps. In Fig 5, 6, 7, 8, we provide the Pruning and Probing results for different numbers of heads h ∈ { 4 , 8 } and noise levels σ ∈ { 0 , 0 . 1 , 0 . 2 , 0 . 4 , 0 . 8 } . It can be found that in almost all cases, the pruned model exhibits almost the same performance in each layer, while being largely different from the single-layer transformer. This further supports the results in the Heads Assessment and indicates that the working mechanisms of the multi-head transformer may be different for the first and subsequent layers.

input

1

2

3

0

1

2

3

4

5

excess risk

n_layer 3 avg

pruned h = 4

h = 4

h = 1

input

1

2

3

0

1

2

3

4

5

n_layer 3,

= 0

pruned h = 4

h = 4

h = 1

input

1

2

3

0

1

2

3

4

5

n_layer 3,

= 0.1

pruned h = 4

h = 4

h = 1

input

1

2

3

0

1

2

3

4

5

n_layer 3,

= 0.2

pruned h = 4

h = 4

h = 1

input

1

2

3

0

1

2

3

4

5

n_layer 3,

= 0.4

pruned h = 4

h = 4

h = 1

input

1

2

3

0

1

2

3

4

5

n_layer 3,

= 0.8

pruned h = 4

h = 4

h = 1

input

1

2

3

i-th layer

0

1

2

3

4

5

excess risk

n_layer 3 avg

pruned h = 8

h = 8

h = 1

input

1

2

3

i-th layer

0

1

2

3

4

5

n_layer 3,

= 0

pruned h = 8

h = 8

h = 1

input

1

2

3

i-th layer

0

1

2

3

4

5

n_layer 3,

= 0.1

pruned h = 8

h = 8

h = 1

input

1

2

3

i-th layer

0

1

2

3

4

5

n_layer 3,

= 0.2

pruned h = 8

h = 8

h = 1

input

1

2

3

i-th layer

0

1

2

3

4

5

n_layer 3,

= 0.4

pruned h = 8

h = 8

h = 1

input

1

2

3

i-th layer

0

1

2

3

4

5

n_layer 3,

= 0.8

pruned h = 8

h = 8

h = 1

input

1

2

3

4

0

1

2

3

4

5

excess risk

n_layer 4 avg

pruned h = 4

h = 4

h = 1

input

1

2

3

4

0

1

2

3

4

5

n_layer 4,

= 0

pruned h = 4

h = 4

h = 1

input

1

2

3

4

0

1

2

3

4

5

n_layer 4,

= 0.1

pruned h = 4

h = 4

h = 1

input

1

2

3

4

0

1

2

3

4

5

n_layer 4,

= 0.2

pruned h = 4

h = 4

h = 1

input

1

2

3

4

0

1

2

3

4

5

n_layer 4,

= 0.4

pruned h = 4

h = 4

h = 1

input

1

2

3

4

0

1

2

3

4

5

n_layer 4,

= 0.8

pruned h = 4

h = 4

h = 1

input

1

2

3

4

i-th layer

0

1

2

3

4

5

excess risk

n_layer 4 avg

pruned h = 8

h = 8

h = 1

input

1

2

3

4

i-th layer

0

1

2

3

4

5

n_layer 4,

= 0

pruned h = 8

h = 8

h = 1

input

1

2

3

4

i-th layer

0

1

2

3

4

5

n_layer 4,

= 0.1

pruned h = 8

h = 8

h = 1

input

1

2

3

4

i-th layer

0

1

2

3

4

5

n_layer 4,

= 0.2

pruned h = 8

h = 8

h = 1

input

1

2

3

4

i-th layer

0

1

2

3

4

5

n_layer 4,

= 0.4

pruned h = 8

h = 8

h = 1

input

1

2

3

4

i-th layer

0

1

2

3

4

5

n_layer 4,

= 0.8

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

0

1

2

3

4

5

excess risk

n_layer 5 avg

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

0

1

2

3

4

5

n_layer 5,

= 0

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

0

1

2

3

4

5

n_layer 5,

= 0.1

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

0

1

2

3

4

5

n_layer 5,

= 0.2

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

0

1

2

3

4

5

n_layer 5,

= 0.4

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

0

1

2

3

4

5

n_layer 5,

= 0.8

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

i-th layer

0

1

2

3

4

5

excess risk

n_layer 5 avg

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

i-th layer

0

1

2

3

4

5

n_layer 5,

= 0

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

i-th layer

0

1

2

3

4

5

n_layer 5,

= 0.1

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

i-th layer

0

1

2

3

4

5

n_layer 5,

= 0.2

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

i-th layer

0

1

2

3

4

5

n_layer 5,

= 0.4

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

i-th layer

0

1

2

3

4

5

n_layer 5,

= 0.8

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

6

0

1

2

3

4

5

excess risk

n_layer 6 avg

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

6

0

1

2

3

4

5

n_layer 6,

= 0

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

6

0

1

2

3

4

5

n_layer 6,

= 0.1

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

6

0

1

2

3

4

5

n_layer 6,

= 0.2

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

6

0

1

2

3

4

5

n_layer 6,

= 0.4

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

6

0

1

2

3

4

5

n_layer 6,

= 0.8

pruned h = 4

h = 4

h = 1

input

1

2

3

4

5

6

i-th layer

0

1

2

3

4

5

excess risk

n_layer 6 avg

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

6

i-th layer

0

1

2

3

4

5

n_layer 6,

= 0

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

6

i-th layer

0

1

2

3

4

5

n_layer 6,

= 0.1

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

6

i-th layer

0

1

2

3

4

5

n_layer 6,

= 0.2

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

6

i-th layer

0

1

2

3

4

5

n_layer 6,

= 0.4

pruned h = 8

h = 8

h = 1

input

1

2

3

4

5

6

i-th layer

0

1

2

3

4

5

n_layer 6,

= 0.8

pruned h = 8

h = 8

h = 1

P-probing Here, we also set n = 117 and q = 11 , with an evaluation data size of 1024 . We choose n ≫ q such that the model can handle more queries ( q = 11 ) than those in the training ( q = 4 ) process.

C Proof for Section 4

C.1 Proof for Proposition 4.1

Proposition C.1 (Restate of Proposition 4.1) . There exists a transformer with 1 layers, h = d heads, d hid = 3 d and the input projection W E ∈ R ( d +1) × d hid such that with the input sequence E set as Equation 2.3 the first attention layer can implement Algorithm 1 so that each of the enhanced data { ̂ r i x i,j } i ∈ [ d ] can be found in the output representation H (1) :

Proof. Here we first explain the key steps of our constructed transformer: the model first rearrange the input entries with a input projection to divide the input data into d subspace W E , each subspace includes an entry of x and the corresponding y (step C.2), then use h parameters { W V i , W K i , W Q i } h i =1 to calculate h queries, keys and values (step C.3), and compute the attention output for each head and concatenate them together (step C.4), finally use a projection matrix W 1

rearrange the result, resulting the target output (step C.5):

. The detailed parameters and calculation process for each step are as follows:

we set W E ∈ R ( d +1) × d hid to rearrange the entries:

where ✶ [ k ] is an 1 × d hid vector with 1 at i -th entry and 0 elsewhere, such that

· we set W V i , W K i , W Q i ∈ R 3 × d hid for values, keys and queries:

such that the i -th head extract i -th entry of x and corresponding y

Then concatenate the output of each head { V i MK ⊤ i Q i } h i =1 together with residue:

Finally, W 1 is applied to rearrange the entries:

✶ ✶ ✶ where the first · · · implies the omitted d -2 vectors { ✶ [3 i ] | i = 2 , 3 , . . . , ( d -1) } , the second · · · implies arbitrary values, then resulting the final output:

in this way we construct a transformer that can apply Alg. 1 so that each of the enhanced data { ̂ r i x i,j } i ∈ [ d ] can be found in the output representation H (1) .

C.2 Proof for Proposition 4.2

Proposition C.2 (Restate of Proposition 4.2) . There exists a transformer with k layers, 1 head, d hid = 3 d , let { ( ˜ x i , ̂ y ℓ ( i ) ) } n +1 i =1 be the ℓ -th layer input data entry, then it holds that ̂ y ℓ ( n +1) = ⟨ w ℓ gd , ˜ x n +1 ⟩ , where w gd is defined as w 0 gd = 0 and as follows for ℓ = 0 , ..., k -1 :

Proof. Here we directly provide the parameters W ℓ V , W ℓ K , W ℓ Q ∈ R d hid × d hid and W ℓ 1 ∈ R d hid × d hid for each layer TF ℓ ,

As we set W ℓ 1 as the identity matrix, we can ignore it and then apply Lemma 1 in [3]. By replacing ( W ℓ K ⊤ W ℓ Q ) as Q i and W ℓ V with P i , then it holds that ̂ y ℓ ( n +1) = ⟨ w ℓ gd , ˜ x n +1 ⟩ , where w gd is defined as w 0 gd = 0 and as follows for ℓ = 0 , ..., k -1 :

D Proof of Theorem 5.1

To simplify the notations, we use ̂ w t to denote ˜ w t gd . We first prove that with a high probability, there exists a R ∈ R d × d such that R ̂ R = ̂ RR = I s , where I s = diag { a 1 , . . . , a d } with a j = 1 { j ∈S} . Lemma D.1. Denote R = diag { r 1 , . . . , r d } , where r j = ∑ d i =1 w ⋆ i Σ ij . Suppose n ≥ O (log ( d/δ )) , then for any δ ∈ (0 , 1) with probability at least 1 -δ , we have

where K := C ( max i Σ ii + σ 2 ) , where C is an absolute constant. Lemma D.2. Define the event E R by E R = { | ̂ r | i ≥ 1 2 | r i | , ∀ i ∈ S } . Suppose that n ≳ s log ( d/δ ) /β 2 , then P ( E 1 ) ≥ 1 -δ .

We define R by R = diag { r 1 , . . . , r d } , where r j is given by

It is easy to see R ̂ R = ̂ RR = I s . On the event E 1 , we have that ∥ ∥ R ∥ ∥ ≲ 1 /β . Hereafter, we condition on E 1 .

D.1 Bias-variance Decomposition

Let ˜ X = X ̂ R with ˜ x i = ̂ Rx i . For ̂ w t , we have

Hence, we have

We can decompose the risk L ( ̂ w t ) by

Next, we present some lemmas.

Lemma D.3 (Theorem 9 in Bartlett et al. [9]) . There is an absolute constant c such that for any δ ∈ (0 , 1) with probability at least 1 -δ ,

where r ( Σ ) = Tr( Σ ) /λ 1 .

Lemma D.4. With probability at least 1 -δ , we have

As a result, when n ≳ st 2 ( r 2 / 3 ( Σ ) + r ( RΣR ) ) · poly(log ( d/δ )) , with probability at least 1 -δ , we have

We define the event E 2 as follows:

where

By Lemma D.4, P ( E 2 ) ≥ 1 -δ . Hereafter, we condition on E 1 ∩ E 2 .

D.2 Bounding the Bias

On E 1 ∩ E 2 , we have

Lemma D.5. On E 1 ∩ E 2 , we have and

hold with probability at least 1 -δ .

By Lemma D.5, we obtain that with probability at least 1 -δ ,

where the last inequality is by η ≲ 1 / ∥ Σ ∥ ≲ 1 .

D.3 Bounding the Variance

Lemma D.6. On E 1 ∩ E 2 , with probability at least 1 -δ , we have

By applying Lemma D.6 to Eq.(D.6), we obtain that

Lemma D.7. with probability at least 1 -δ , we have

By applying Lemma D.7 to Eq.(D.7), we obtain that

D.4 Final Bound

Combining Eq.(D.5) and Eq.(D.8), we obtain that

when ηt ≃ 1 β · ( σ 2 Tr( RΣR ) log ( d/δ ) n + σ 2 s Tr( Σ ) log 2 ( d/δ ) n 2 ) -1 / 2 .

D.5 Proof for Appendix D

Proof of Lemma D.1. Since y i = ∑ d j =1 w ⋆ j x ij + ϵ i , then we have

and

Since x ji ∼ N (0 , Σ ii ) for any i, j , by Lemma 2.7.7 in Vershynin [46], there exists an absolute constant C such that x jk x ji is a sub-exponential random variable with

where ∥ · ∥ Ψ 1 denotes the sub-exponential norm and the last inequality comes from the definition of K . By applying Bernstein's inequality [46, Theorem 2.8.1], we have

where the last equality due to n ≥ O (log ( d/δ )) . We also note that x ji ϵ j is a sub-exponential random variable with ∥ x ji ϵ j ∥ Ψ 1 ≤ K. Hence, we also have

Combining Eq.(D.9), Eq.(D.10) and Eq.(D.11), we have

By definition of ̂ R and R , we obtain

which completes the proof.

Proof of Lemma D.2. By Lemma D.1, for any j ∈ S , with probability at least 1 -δ , we have

where the last inequality is due to the definition of β .

Proof of Lemma D.4. We can decompose ∥ ̂ R ̂ Σ ̂ R -RΣR ∥ 2 as follows:

Next, we proof the bound for I , II and III separately.

For term I ,

where the last line is due to triangle inequality. By Lemma D.3, with probability at least 1 -δ/ 3 , we have

holds with probability at least 1 -δ/ 3 , where the last inequality is valid since n ≳ K 2 s ∥ R ∥ 2 2 log ( d/δ ) . Combing Eq.(D.14), Eq.(D.15) and Eq.(D.16), we have

For term II , we can decompose II as follows:

For term II . a , by using Lemma D.3, we have with probability at least 1 -δ/ 3 ,

Similar to the proof for bounding I , we can obtain that

For term III , we have

where the last inequality is by Eq.(D.16).

Combining Eq.(D.17), Eq.(D.18), Eq.(D.19) and Eq.(D.20) and taking the union bound, we obtain that with probability at least 1 -δ ,

where the second last inequality is by aa ′ + bb ′ + cc ′ ≤ ( a + b + c )( a ′ + b ′ + c ′ ) for a, a ′ , b, b ′ , c, c ′ ≥ 0 . Here ˜ C cov = K ∥ Σ ∥ 2 ( ∥ w ⋆ ∥ 1 +1) + ∥ RΣR ∥ 2 + ∥ R ∥ 2 ∥ Σ ∥ 2 K ( ∥ w ⋆ ∥ 1 +1) ≲ √ s .

Proof of Lemma D.5. By the triangle inequality, we have

Following the proof of Lemma D.4, we can prove that with probability at least 1 -δ ,

where the last inequality is by E 2 . By Eq.(D.21), we have

Hence, we obtain that

where the last line by R ⪯ 2 β · I . For the term II , we have

where the second last line is by the fact that x (1 -x ) k ≤ 1 / ( k + 1) for all x ∈ [0 , 1] and all k > 0 .

Proof of Lemma D.6. Similar to the proof of Lemma D.5, with probability at least 1 -δ , we have ̂ R ( Σ -̂ Σ ) ̂ R ⪯ 1 t · I . Then we have

where the second last line is by ∑ t i =1 ( I -η ̂ Σ ) i -1 ⪯ t · I . By the fact that x (1 -x ) k ≤ 1 / ( k +1) for all x ∈ [0 , 1] and all k > 0 , we have

where the last inequality is by the fact that ∑ t i =1 1 i ≲ log t

.

Proof of Lemma D.7. First, we can decompose ∥ ∥ ∥ 1 n · ̂ RX ⊤ ϵ ∥ ∥ ∥ 2 2 by

Let z i = Rx i , then z i ∼ N ( G ) , where G := RΣR . For any i, j , by Lemma 2.7.7 in Vershynin [46], there exists an absolute constant C such that ϵ j z ji is a sub-exponential random variable with

By applying Bernstein's inequality Vershynin [46, Theorem 2.8.1], for any 1 ≤ i ≤ d , we have that

hold with probability 1 -δ 3 d , where the last inequality is due to n ≥ O (log( d/δ )) . By taking the union bound, we obtain that

holds for any i , with probability 1 -δ 3 . Then we have

In the same way, we can prove that with probability at least 1 -δ/ 3 ,

By applying Lemma D.1, with probability at least 1 -δ/ 3 , we have

By Eq.(D.25) and Eq.(D.26), with probability 1 -2 δ/ 3 , we have

By taking the union bound, we derive the desired result.

E Proof for Theorem 5.2

To simplify the notations, we use w t to denote w t gd

.

Lemma E.1. with probability at least 1 -δ , we have

where α ( n, δ ) = √ Tr( Σ )+log (1 /δ ) n + Tr( Σ )+log (1 /δ ) n . As a result, when n ≳ t 2 (Tr( Σ ) + log (1 /δ )) , with probability at least 1 -δ ,

Proof of Lemma E.1. By Lemma D.3, we have

holds with probability at least 1 -δ , where the last line is by the inequality that max { a, b } ≤ a + b for all a, b ≥ 0 .

We define the event E as follows:

By Lemma E.1, P ( E ) ≥ 1 -δ . Hereafter, we condition on E .

Bias-variance Decomposition Similar to Eq.(D.1), we have

In the same way, we can decompose the risk E ( w t ) by

Bounding the Bias

Similar to the proof of Lemma D.5, we have the following lemma.

hold with probability at least 1 -δ .

As a result, the bound of the bias term is given by

Bounding the Variance By using the same way of the proof for bounding the variance term of Theorem 5.1, we have the following lemma.

Lemma E.3. On E , with probability at least 1 -δ , we have that

Combining Eq.(E.5) and Eq.(E.6), we obtain that

when ηt ≃ ( σ 2 Tr(Σ) log ( d/δ ) n ) -1 / 2

F Proof for Theorem 5.3

To simplify the notation, we use ̂ w t to denote ˜ w t gd and w t to denote w t gd .

F.1 Proof for the upper bound of the excess risk

When Σ = I , by Eq.(D.2), we have

Following the proof of Theorem 5.1, it holds that

with probability at least 1 -δ , when n ≳ t 2 sd 2 / 3

Similar to the proof of Lemma D.2, we can prove that

with probability at least 1 -δ .

When Σ = I , by Lemma D.4, we have that

holds with probability at least 1 -δ , when n ≳ t 2 ∥ w ⋆ ∥ 2 1 d 2 / 3 β 4 . As a result, RΣR -β 2 t · I ⪯ ̂ R ̂ Σ ̂ R . Hereafter, we condition on the above events. For the bias term, we have

where the last line is by the definition of β . When t ≳ log ( σ 2 ns ) / ( 2 log ( 1 -ηβ 2 / 2 )) , we have

When ηβ 2 / 2 ≤ 1 / 2 , there exist a c > 0 , such that

Hence, the variance term is bounded by

where the last line is by ∥ w ⋆ ∥ 1 ≤ s · ∥ w ⋆ ∥ 2 2 = s and ηt ≲ log ( ns/σ 2 ) β 2 . Combining Eq.(F.1) and Eq.(F.2), we have that

when n ≳ t 2 sd 2 / 3 β 4 ≥ t 2 ∥ w ⋆ ∥ 2 1 d 2 / 3 β 4 and t ≳ log ( ns ) ηβ 2 . When w ⋆ i ∈ U {-1 / √ s, 1 / √ s } , β = 1 / √ s . In this case, we have that

when n ≳ t 2 s 3 d 2 / 3 and t ≳ log ( ns ) ηs .

F.2 Lower bound for Ridge Regression

When n ≳ d +log(1 /δ ) , by Lemma D.3, we have that 1 2 · I ⪯ ̂ Σ ⪯ 2 · I For the ridge estimator ̂ w λ = 1 n · ( ̂ Σ + λ · I ) -1 X ⊤ y , we have

By Lemma D.3, when 1 2 · I ⪯ ̂ Σ ⪯ 2 · I , with probability at least 1 -δ , we have

where the last line is due to the fact that ̂ Σ + λ I ⪯ (2 + λ ) · I . Lemma F.1. Given X such that 1 2 I ⪯ ̂ Σ ⪯ 2 I , it holds that

with probability at least 1 -δ , when n ≥ O (log (1 /δ )) .

Proof of Lemma F.1. We consider the singular value decomposition of 1 √ n X ⊤ : 1 √ n X ⊤ = UΛV ⊤ , where U ∈ R d × d is an orthogonal matrix, Λ ∈ R d × n is a rectangular diagonal matrix with nonnegative real numbers on the diagonal, V ∈ R n × n is an orthogonal matrix. Let { σ 1 , . . . , σ d } be the singular values of 1 √ n X ⊤ . Then we have

where ˜ ϵ = V ⊤ ϵ ∼ N ( 0 , I ) . By ][Lemma 22], we have

where the last line is valid since { σ 2 1 , . . . , σ 2 d } is the eigenvalues of ̂ Σ = 1 n X ⊤ X and 1 2 I ⪯ ̂ Σ ⪯ 2 I . By Eq.(F.3), we obtain that

where the last line is due to d ≥ O (log (1 /δ )) .

Next, we define the event E as follows:

By Lemma F.1, we have P ( E ) ≥ 1 -δ when n ≥ O ( d ) ≥ O (log (1 /δ )) . On E ridge , we have

When d ≳ n + log (1 /δ ) , by Lemma D.3, with probability at least 1 -δ , we have that d 2 · I ⪯ XX ⊤ ⪯ 2 d · I . Hereafter, we condition on this event. By direct calculation, we can decompose the excess risk by

∥

∥

∥

∥

(

(

)

-

1

)

∥

∥

∥

∥

2

2

∥

∥

∥

∥

1

n

(

)

-

1

2

∥

∥

∥

2

∥

E

w

⋆

[

E

(

̂

w

λ

)] =

E

w

For the first term, we have

where the last line is due to ( I -X ⊤ ( XX ⊤ + nλ I ) -1 X ) is a d -n space.

where the first line is by ( X ⊤ X + nλ I ) -1 X ⊤ = X ⊤ ( XX ⊤ + nλ I ) -1 and the last line is by d 2( d + nλ ) 2 · I ⪯ XX ⊤ ( XX ⊤ + nλ I ) -2 . By Tsigler and Bartlett [44, Lemma 22], we obatain that

holds with probability at least 1 -δ . When n ≳ log (1 /δ ) , we have ∣ ∣ ∑ n i =1 ϵ 2 i -nσ 2 ∣ ∣ ≥ nσ 2 2 holds with probability at least 1 -δ . Taking the union bound, we obtain that

⋆

I

-

Σ

̂

+

λ

I

Σ

̂

w

⋆

+

·

Σ

̂

+

λ

·

I

X

⊤

ϵ

.

F.3 Lower Bound for Finite-Step GD

We first consider the case where n ≳ d + log (1 /δ ) . Define the event E GD by E GD = { 1 2 · I ⪯ ̂ Σ ⪯ 2 I } . By Lemma D.3, P ( E GD ) ≥ 1 -δ . By Eq.(E.4), we have

where the second last line is by ̂ Σ ( I -( I -η ̂ Σ ) t ) -1 ⪯ Σ + 2 tη · I and the last line is by Eq.(F.4).

We then consider the case where d ≳ n +log (1 /δ ) . Define the event E ′ GD = { d 2 · I ⪯ XX ⊤ ≺ 2 d I } . By Lemma D.3, P ( E ′ GD ) ≥ 1 -δ . Following the proof of Zou et al. [58, Theorem 4.3], we have

where we use the results from Appendix F.2.

F.4 Lower bound of OLS

Let w ols be the OLS estimator. It is easy to see w ols = w 0 . Hence, we have

holds with probability at least 1 -δ .

G Additional Experiments

Here, we provide additional experiments on the decoder-only architecture and train models with different settings.

Training Decoder-Only Transformer In this experiment, we adapt the same input setting and training objective as in [20]. During training, we set n = 24 and k = 8 in Eq.(G.1) (where in y i , we use zero padding to align with x i ), d hid = 256 . We choose h = 8 and l ∈ { 4 , 5 , 6 } . 3 We then conduct heads assessment experiments on the trained decoder-only transformers with 10 in-context

examples, as in the previous settings. The result is shown in Figure 9. We can observe that the decoder-only transformer exhibits the similar weight distribution for each layer as the encoder-based models, indicating that our algorithm may extend to decoder-only based models.

0

1

2

3

4

5

6

7

0

1

2

3

j-th layer

h_num = 8, n_layer 4 avg

0

1

2

3

4

5

6

7

0

1

2

3

4

h_num = 8, n_layer 5 avg

0

1

2

3

4

5

6

7

0

1

2

3

4

5

h_num = 8, n_layer 6 avg

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.1

0.2

0.3

0.4

0.5

0.6

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Training Models with s = d = 16 Here, we adapt the encoder-only transformer and the same settings as introduced in B, but set s = d = 16 . We observe that in these cases, there is no distinct performance difference between models with different numbers of heads. As shown in Figure 3, when we set s = 4 , d = 16 , transformers with more heads ( h = 4 , 8 ) always perform better than models with fewer heads ( h = 1 , 2 ). However, in Figure 10, such a difference is unclear, which aligns well with the theoretical analysis. When s is close to d , a clear better upper bound guarantee, as ensured in cases where s ≪ d may not hold.

Training Models with non-orthogonal design To further demonstrate the applicability of our experimental results to more general non-orthogonal settings, we conducted additional experiments by modifying the distribution of x to N ( 0 , Σ) , where Σ = I + ζ S , and S is a matrix of ones. We varied ζ across the values [0 , 0 . 1 , 0 . 2 , 0 . 4] to validate our findings. The results are presented in Figure 11, which reveals patterns similar to those observed in orthogonal design settings.

5

10

15

20

25

0

2

4

6

8

10

excess risk

n_layer = 3,

= 0

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 4,

= 0

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 5,

= 0

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 6,

= 0

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

excess risk

n_layer = 3,

= 0.1

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 4,

= 0.1

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 5,

= 0.1

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 6,

= 0.1

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

excess risk

n_layer = 3,

= 0.2

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 4,

= 0.2

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 5,

= 0.2

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 6,

= 0.2

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

excess risk

n_layer = 3,

= 0.4

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 4,

= 0.4

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 5,

= 0.4

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

0

2

4

6

8

10

n_layer = 6,

= 0.4

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

in-context examples

0

2

4

6

8

10

excess risk

n_layer = 3,

= 0.8

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

in-context examples

0

2

4

6

8

10

n_layer = 4,

= 0.8

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

in-context examples

0

2

4

6

8

10

n_layer = 5,

= 0.8

ols

ridge

h = 1

h = 2

h = 4

h = 8

5

10

15

20

25

in-context examples

0

2

4

6

8

10

n_layer = 6,

= 0.8

ols

ridge

h = 1

h = 2

h = 4

h = 8

1

2

3

4

3

2

1

j-th layer

=0.0, avg

1

2

3

4

3

2

1

=0.1, avg

1

2

3

4

3

2

1

=0.2, avg

1

2

3

4

3

2

1

=0.4, avg

1

2

3

4

5

6

7

8

3

2

1

j-th layer

=0.0, avg

1

2

3

4

5

6

7

8

3

2

1

=0.1, avg

1

2

3

4

5

6

7

8

3

2

1

=0.2, avg

1

2

3

4

5

6

7

8

3

2

1

=0.4, avg

1

2

3

4

4

3

2

1

j-th layer

=0.0, avg

1

2

3

4

4

3

2

1

=0.1, avg

1

2

3

4

4

3

2

1

=0.2, avg

1

2

3

4

4

3

2

1

=0.4, avg

1

2

3

4

5

6

7

8

i-th head

4

3

2

1

j-th layer

=0.0, avg

1

2

3

4

5

6

7

8

i-th head

4

3

2

1

=0.1, avg

1

2

3

4

5

6

7

8

i-th head

4

3

2

1

=0.2, avg

1

2

3

4

5

6

7

8

i-th head

4

3

2

1

=0.4, avg

0.1

0.2

0.3

0.4

0.5

0.15

0.20

0.25

0.30

0.35

0.40

0.45

0.1

0.2

0.3

0.4

0.5

0.1

0.2

0.3

0.4

0.5

0.1

0.2

0.3

0.4

0.1

0.2

0.3

0.4

0.5

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

0.45

0.1

0.2

0.3

0.4

0.5

0.1

0.2

0.3

0.4

0.5

0.6

0.1

0.2

0.3

0.4

0.5

0.6

0.1

0.2

0.3

0.4

0.5

0.6

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

0.1

0.2

0.3

0.4

0.5

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.1

0.2

0.3

0.4

0.5

0.6

0.7

NeurIPS Paper Checklist

1. Claims

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

Answer: [Yes]

Justification: Main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.

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: In Section 7, we discuss the limitation and possible future works for this paper.

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: We provide proofs for every theories and propositions in appendix.

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

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: We provide experimental settings and training details in Appendix B.

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: [No]

Justification: the code is being organized and we can provide it at an appropriate time.

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: We provide experimental settings and training details in Appendix B.

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: The reported mean behaviors in our experiments are sufficient to support our theoretical results.

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 provide experimental settings and training details 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 conducted in the paper conform, in every respect, with 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: [NA]

Justification: no societal impact of the work performed.

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: [NA]

Justification: the 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: [Yes]

Justification: 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

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: [NA]

Justification: the paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details 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: [NA]

Justification: the paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

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: [NA]

Justification: the paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

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

We recognize that the procedures for this may vary significantly between institutions and locations, and 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 empirically studies how the different heads works across different layers of the transformer: while the first layer uses all heads, the later layers mainly relies on a single head. In addition, the authors also propose a preprocess-then-optimize algorithm.\n\nSTRENGTHS:\nThis paper is easy to understand and provides interesting theoretical results.\n\nWEAKNESSES:\nThe main flow of Section 4 and 5 is that, it (1) constructs a first layer which can perform the exact preprocessing considered in the proposed algorith, (2) studies the generalization of the preprocessing, (3) constructs a the later layers which can perform gradient descent, and finally (4) studies the generalization of the optimization. The paper uses a linear attention-only transformer. \n\nGiven the above, there are several main weaknesses regarding the contributions in this paper:\n\n(1) Missing theory or theoretical intuition to explain why multi-head attention works in the way observed in this paper. Given the series of in-context learning research which focuses on aligning multi-layer transformers with gradient descent, the contribution is limited.\n\n(2) The linear attention-only transformer can be improved to better align with the practical LLMs. In the recent year, various studies consider single/multi-head attention with softmax attention [1,2,3]. The authors are encouraged to add softmax attention in the paper. In particular, compared to the linear attention result in [4], the softmax result in [1,2,3] shows that softmax attention works differently from linear attention. The authors are also encouraged to study how MLP works in the system. These elements in the transformer may result in different behavior compared to the linear attention-only transformer.\n\n(3) The paper constructs a transformer which aligns with the preprocess-then-optimize algorithm, rather than showing that a trained multi-layer transformer indeed works in the way designed in this paper. The authors need to provide more evidences for this, e.g., [5]. The authors may also want to highlight their different to [5]. A convergence analysis similar to [1] and [2] is even better.\n\n(4) The generalization properties are built upon a linear model, whose analysis is supposed to be routine. The authors may also highlight the challenges when performing the analysis.\n\nGiven the above limitation, the theoretical contributions are not sufficient enough. There is no critical difference between existing theories in literature and the study in this paper.\n\nReferences:\n\n[1] Huang, Yu, Yuan Cheng, and Yingbin Liang. \"In-context convergence of transformers.\" arXiv preprint arXiv:2310.05249 (2023).\n\n[2] Chen, Siyu, et al. \"Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality.\" arXiv preprint arXiv:2402.19442 (2024).\n\n[3] Cui, Yingqian, et al. \"Superiority of multi-head attention in in-context linear regression.\" arXiv preprint arXiv:2401.17426 (2024).\n\n[4] Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. \"Trained transformers learn linear models in-context.\" arXiv preprint arXiv:2306.09927 (2023).\n\n[5] Ahn, Kwangjun, et al. \"Transformers learn to implement preconditioned gradient descent for in-context learning.\" Advances in Neural Information Processing Systems 36 (2024).\n\nQUESTIONS:\nI would like to raise my score if the authors could provide an extra detailed theoretical result for any one of the above weaknesses. The more the better. Thanks.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work presents an analysis of how Transformers perform in-context learning by experimenting with a sparse linear regression problem setup. The author’s analysis combines both empirical and theoretical analysis. They first examine the properties of real models using pruning and probing methods. From this, they identify a key property of the Transformers, that multiple heads are needed in the first layer but in the remaining layers only a single attention head is necessary. From this, they present theoretical analysis and propose a preprocess-then-optimize learning approach to explain the observed behavior in transformers. The authors also present results showing the effectiveness of this learning approach.\n\nSTRENGTHS:\nWhile this work is not in my area of expertise, I believe the work is well organized and clearly presented. I appreciate the authors combination of empirical and theoretical analysis, to demonstrate that the proposed learning method may actually be employed by real transformers in real tasks and are not only achievable through hand-crafted weights. The authors results are mathematically rigorous as well in supporting the proposed preprocess-then-optimize method.\n\nWEAKNESSES:\nThe main weakness I see in the work pertains to the simplifications made for the analysis. The authors also acknowledge these limitations in section 7. In summary, the analysis focuses on a simplified problem (linear regression) and uses a simplified version of a transformer, which is also done by other similar works. As a result, it is hard to say if a similar learning approach is learned by transformers in more practical real-world problems. It is also impossible to say how the removed layers may alter the learning process in a full transformer. Overall, I believe this work warrants an accept rating, though I will defer to the other reviewers if they have significant issues with the work.\n\nQUESTIONS:\nLine 108: should this instead read \"to perform a tractable theoretical investigation\"?\nLine 209: typo \"proprocessing\" instead of \"preprocessing\"\n\nLIMITATIONS:\nThe authors have discussed the limitations I mentioned in the section above. I do not see any clear potential negative societal impact for this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the mechanism of transformers under the in-context sparse linear regression problem. The authors reveal that the transformer pre-trained for this task has the first layer preprocessing the data, and the remaining layers implement gradient descent. More intriguingly, only one head in the second to last layers is dominantly utilized.\n\nSTRENGTHS:\n1. The paper is a nice extension of the recent line of work on ICL formulated by regression problems. The study of multi-head also provides an important direction.\n2. The theory is solid and clear.\n3. The authors have found many interesting phenomena in their setting. I especially find their study on multi-head very interesting and beneficial for LLM interpretability.\n\nWEAKNESSES:\n1. The P-probing lacks a controlled experiment. Should you also try regressing on the hidden states before the first layer? I understand that $h=1$ might be a controlled experiment, but it may still have some preprocessing on the $x$.\n2. Is Theorems 5.1 and 5.2 a fair comparison? Should the $\\tilde{w}^t_{gd}$ in Theorem 5.1 also include the parameters in the first preprocessing layer? My understanding is that we want to compare the excess risk of \"preprocessing+gd\" v.s. \" gd\". Now we are comparing \"gd with preprocessed data\" and \"gd\".\n3. The orthogonal design is an over-simplified setting.\n\nQUESTIONS:\n1. I find the observation that only one head dominates in subsequent layers very interesting. Does this also occur in other settings of ICL tasks?\n\nLIMITATIONS:\nThe orthogonal design is an over-simplified setting, which the authors do not properly address.\n1. It is necessary to include ablation studies with non-orthogonal design. It's fine to get different results without the orthogonality, but it would be important to report the findings.\n2. If the MLP layer is added, an alternative lasso method should be available that uses the closed-form solution under orthogonal design. Can the authors present the experiment results?\nOverall, I think there should be more ablation studies in different settings. Although a simplified setting is fine for theory, I want to see experimental results in more realistic settings.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work seeks to provide a deeper exploration of the use of multi-heads, at different layers in a Transformer, to perform in-context\nlearning tasks. More specifically, the goal of the paper is to experimentally discover additional insights on the interactions of multi-headed attention across layers. Subsequently, the authors find that multiple heads are primarily used in the first layer, whilst the remaining layers of the transformer typically leverage a single head. Furthermore, the authors provide a hypothesis for their observation before building a preprocess-then-optimize algorithm.\n\nSTRENGTHS:\nStrengths:\n- This work provides a timely analysis of the underlying mechanism of multi-headed attention in Transformer models.\n- The authors provide strong empirical and theoretical evidence to justify the insights derived from on in-context learning.\n- Additionally, the authors provide a novel Preprocess-then-optimize Algorithm for training Transformers.\n\nWEAKNESSES:\nWeaknesses:\n- As discussed by the authors, the reviewer's primary concern with this work is the lack of account for the remainder of the Transformer architecture. In particular, the role of feed-forward layers is sidesteped when considering the effects of the Preprocess-then-optimize Algorithm.\n\nQUESTIONS:\nThe reviewer would like to better understand how the authors think about feed-forward layers with respect to the observations made by the authors regarding the necessity of multi-headed attention in the first few layers. In particular, would alterations to the structure of the first few feedforward layers provide a beneficial impact for transformers?\n\nLIMITATIONS:\nThe authors have clearly stated any limitations with their current work. Additionally, the reviewer does not foresee any potential negative societal impact from this paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}