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

Authors: xingwu  chen
lei  zhao
difan  zou
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 09:16:47

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.

Main paper: 
 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. 
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:L(w) := E (x,y)∼P (⟨x, w⟩ -y) 2 .Moreover, we are interested in the excess risk, i.e., the gap between the population risk achieved by w and the optimal one: E(w) := L(w) -min w L(w).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:Attn(H, θ 1 ) = H + Concat[V 1 sfmx(K ⊤ 1 Q 1 ), • • • , V h sfmx(K ⊤ h Q h )],whereV i = W Vi H, Q i = W Qi H, V i = V Vi H and θ 1 = W Vi , W Ki , W Qi ∈ R d hid /h×d hid h i=1are learnable parameters. The MLP then applies a nonlinear element-wise operation:EQUATIONwhere θ 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.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:EQUATIONwhere{W Vi , W Ki , W Qi ∈ R d hidh ×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, EQUATIONy = W O • TF L • • • • • TF 1 • W E (E),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:E ICL := E (x,y)∼P ( y n+1 -y n+1 ) 2 -σ 2 .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).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• 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. 
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:W i,j = ∆E ICL(i,j) h k=1 ∆E ICL(i,k) . (3.1)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. 
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: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.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 havey = ⟨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 2i ] 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. 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 Ki H) ⊤ W Qi 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 Ki and W Qi 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 Vi 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 = 3d 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 thatEQUATIONwhere • • • in third row implies arbitrary values.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 projectionW (i) 1to 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 (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 = 3d, 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:w ℓ+1 gd = w ℓ gd -η∇ L(w ℓ gd ), where L(w) = 1 2n n i=1 (y i -⟨w, x i ⟩) 2 . (4.3)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. 
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 ηt ≃ 1 β • σ 2 Tr(RΣR) log (d/δ) n + σ 2 sTr(Σ) log 2 (d/δ) n 2 -1/2 , it holds that E w t gd ≲ log t β σ 2 Tr(RΣR) log (d/δ) n + σ 2 sTr(Σ) log 2 (d/δ) n 2 ,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 thatE w t gd ≲ log t • σ 2 Tr(Σ) log (d/δ) n ,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, assumew ⋆ 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 wherex 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 thatE w t gd ≲ σ 2 log 2 ns/σ 2 log 2 (d/δ) • s n + ds 2 n 2 ,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 thatE w ⋆ [E(w)] ≳ σ 2 d n n ≳ d + log (1/δ) 1 -n d + σ 2 n d d ≳ n + log (1/δ),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. 
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. (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.𝑇𝐹1 𝑥1 𝑥2 … 𝑥𝑛 … Embedding Projection 𝑥𝑛+1 … 𝑥𝑛+𝑞 𝑇𝐹2 𝑇𝐹𝐿 𝐸 = 𝐻 1 = 𝑥 𝑛+1 … 𝑥 𝑛+𝑞-1 𝑥 𝑛+𝑞 𝑦𝑛+1 … 𝑦𝑛+𝑞-1 𝑦𝑛+𝑞 K-step GD 𝑤 gd 𝑘 ො 𝑦𝑛+𝑞 𝑘 ℰ 𝑘 = ( ො 𝑦𝑛+𝑞 𝑘 -𝑦𝑛+𝑞) 2 -𝜎 2To 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. 
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.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. 
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). 
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). 
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). 
 
 Appendix: 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.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.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 . 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. 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.1 Proof for Proposition 4.1Proposition C.1 (Restate of Proposition 4.1). There exists a transformer with 1 layers, h = d heads, d hid = 3d 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) :H (1) = TF 1 • W E (E) =          x 1 x 2 • • • x n x n+1 • • • x n+q y 1 y 2 • • • y n 0 • • • 0 . . . . . . . . . . . . . . . . . . . . .          .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 Vi , W Ki , W Qi } 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):E =    x 1 x 2 • • • x n x n+1 • • • x n+q y 1 y 2 • • • y n 0 • • • 0    (C.1) input projection -----------→ W E ∈R (d+1)×d hid H =              x 1,1 x 2,1 • • • x n,1 x (n+1),1 • • • x (n+q),1 y 1 y 2 • • • y n 0 • • • 0 0 0 • • • 0 0 • • • 0 . . . . . . . . . . . . . . . . . . . . .              (C.2) compute Qi,Ki,Vi ----------------→ W V i ,W K i ,W Q i ∈R 3×d hid K i = 1 n        0 • • • 0 0 • • • 0 • • • 0 0 • • • y 1 • • • y n 0 • • •        ; Q i , V i =        0 • • • 0 0 • • • 0 • • • 0 0 • • • x 1,i • • • x n,i x (n+1),i • • •        (C.3) Attn(W E (E)) ---------------→ H+Concat{ViMK ⊤ i Qi}              x 1,1 x 2,1 • • • x n,1 x (n+1),1 • • • x (n+q),1 y 1 y 2 • • • y n 0 0 0 x 1,1 x 2,1 • • • x n,1 x (n+1),1 • • • x (n+q),1 . . . . . . . . . . . . . . . . . . . . .              (C.4) H (1) =TF1•W E (E) ------------→ W1∈R d hid ×d hid          x 1 x 2 • • • x n x n+1 • • • x n+q y 1 y 2 • • • y n 0 • • • 0 . . . . . . . . . . . . . . . . . . . . .          (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:W E =   1[1] 1[d + 1] 0 1[2] 1[d + 1] 0 • • • 1[d] 1[d + 1] 0   ⊤ ,where 1[k] is an 1 × d hid vector with 1 at i-th entry and 0 elsewhere, such thatH = W E E =                 x 1,1 x 2,1 • • • x n,1 x (n+1),1 • • • x (n+q),1 y 1 y 2 • • • y n 0 • • • 0 0 0 • • • 0 0 • • • 0 x 1,2 x 2,2 • • • x n,2 x (n+1),2 • • • x (n+q),2 . . . . . . . . . . . . . . . . . . . . .                 .• we set W Vi , W Ki , W Qi ∈ R 3×d hid for values, keys and queries:W Ki = 1 n        0 0 1[3i -1]        ; W Vi , W Qi =        0 0 1[3i -2]        ,such that the i-th head extract i-th entry of x and corresponding yK i = 1 n        0 0 1[3i -1]                        x 1,1 x 2,1 • • • x n,1 x (n+1),1 • • • y 1 y 2 • • • y n 0 • • • 0 0 • • • 0 0 • • • x 1,2 x 2,2 • • • x n,2 x (n+1),2 • • • . . . . . . . . . . . . . . . . . .                 = 1 n        0 • • • 0 • • • 0 • • • 0 • • • y 1 • • • y n • • •        , Q i , V i =        0 0 1[3i -2]                        x 1,1 x 2,1 • • • x n,1 x (n+1),1 • • • y 1 y 2 • • • y n 0 • • • 0 0 • • • 0 0 • • • x 1,2 x 2,2 • • • x n,2 x (n+1),2 • • • . . . . . . . . . . . . . . . . . .                 =        0 • • • 0 0 • • • 0 • • • 0 0 • • • x 1,i • • • x n,i x (n+1),i • • •        , V i MK ⊤ i Q i =        0 • • • 0 0 • • • 0 0 • • • 0 0 • • • 0 x 1,i • • • x n,i x (n+1),i • • • x (n+q),i           I n 0 0 0   •        0 • • • 0 0 • • • 0 0 • • • 0 0 • • • 0 y 1 • • • y n 0 • • • 0        ⊤        0 • • • 0 0 • • • 0 0 • • • 0 0 • • • 0 x 1,i • • • x n,i x (n+1),i • • • x (n+q),i        =        0 • • • 0 0 • • • 0 0 • • • 0 0 • • • 0 x 1,i • • • x n,i x (n+1),i • • • x (n+q),i        .• Then concatenate the output of each head {V i MK ⊤ i Q i } h i=1 together with residue:H + Concat[{V i MK ⊤ i Q i } h i=1 ] =              x 1,1 x 2,1 • • • x n,1 x (n+1),1 • • • x (n+q),1 y 1 y 2 • • • y n 0 • • • 0 x 1,1 x 2,1 • • • x n,1 x (n+1),1 • • • x (n+q),1 . . . . . . . . . . . . . . . . . . . . .              . (C.6)• Finally, W 1 is applied to rearrange the entries:W 1 =   1[3] • • • 1[3d] 1[2] • • •   ⊤ ,where the first• • • implies the omitted d -2 vectors {1[3i]|i = 2, 3, . . . , (d -1)}, the second • • •implies arbitrary values, then resulting the final output:H (1) = W 1 H + Concat[{V i MK ⊤ i Q i } h i=1 ] =          x 1 x 2 • • • x n x n+1 • • • x n+q y 1 y 2 • • • y n 0 • • • 0 . . . . . . . . . . . . . . . . . . . . .         .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) .There exists a transformer with k layers, 1 head, d hid = 3d, 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:w ℓ+1 gd = w ℓ gd -η∇ L(w ℓ gd ), where L(w) = 1 2n n i=1 (y i -⟨w, x i ⟩) 2 .Proof. Here we directly provide the parametersW ℓ V , W ℓ K , W ℓ Q ∈ R d hid ×d hid and W ℓ 1 ∈ R d hid ×d hid for each layer TF ℓ , W ℓ V = - η n    0 0 0 1   ; W ℓ K , W ℓ Q =    I d×d 0 0 0   ; W ℓ 1 = I d hid ×d hid (C.7)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:w ℓ+1 gd = w ℓ gd -η∇ L(w ℓ gd ), where L(w) = 1 2n n i=1 (y i -⟨w, x i ⟩) 2 .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∥ R -R∥ 2 ≲ K • s log (d/δ) n ,whereK := 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 byr j = 0 j / ∈ S, 1/ r j j ∈ S.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 .Let X = X R with x i = Rx i . For w t , we havew t+1 -Rw ⋆ = w t -Rw ⋆ -η • 1 n n i=1 x i x ⊤ i w t -y i = w t -Rw ⋆ -η • 1 n n i=1 x i x ⊤ i w t -x ⊤ i Rw ⋆ + ϵ = I -η Σ w t -Rw ⋆ + η • 1 n X ⊤ ϵ.Hence, we havew t = I -I -η Σ t Rw ⋆ + 1 n t i=1 I -η Σ i-1 X ⊤ ϵ. (D.1)We can decompose the risk L( w t ) byE( w t ) = E (x,y)∼P ⟨ Rx, w t ⟩ -⟨ Rx, Rw ⋆ ⟩ -ϵ 2 -σ 2 (D.2) = E (x,y)∼P ⟨ Rx, w t ⟩ -⟨ Rx, Rw ⋆ ⟩ 2 = Σ 1/2 R w t -Rw ⋆ 2 2 = Σ 1/2 R -I -η Σ t Rw ⋆ + η • 1 n t i=1 I -η Σ i-1 X ⊤ ϵ 2 2 = Σ 1/2 R I -η Σ t Rw ⋆ 2 2 Bias + η 2 Σ 1/2 R 1 n t i=1 I -η Σ i-1 X ⊤ ϵ 2 2 Variance . (D.3)Next, we present some lemmas.). There is an absolute constant c such that for any δ ∈ (0, 1) with probability at least 1 -δ,∥ Σ -Σ∥ 2 ≤ c∥Σ∥ 2 • max r(Σ) n , r(Σ) n , log (1/δ) n , log (1/δ) n ,where r(Σ) = Tr(Σ)/λ 1 . Lemma D. 4 . With probability at least 1 -δ, we have∥ R Σ R -RΣR∥ 2 ≲ √ s • poly(log (d/δ)) • r(RΣR) n + r(Σ) + r(RΣR) n + r(Σ) n 3/2 .As a result, when n ≳ st 2 r 2/3 (Σ) + r(RΣR) • poly(log (d/δ)), with probability at least 1 -δ, we have∥ R Σ R -RΣR∥ 2 ≤ 1/t.We define the event E 2 as follows:E 2 := ∥RΣR∥ 2 ≲ α(n, δ) ≤ 1/t , where α(n, δ) = √ s • poly(log (d/δ)) • r(RΣR) n + r(Σ) + r(RΣR) n + r(Σ) n 3/2 .By Lemma D.4, P(E 2 ) ≥ 1 -δ. Hereafter, we condition on E 1 ∩ E 2 .On E 1 ∩ E 2 , we haveBias = Σ 1/2 R I -η Σ t Rw ⋆ 2 2 = w ⋆⊤ R I -η Σ t RΣ R I -η Σ t Rw ⋆ = w ⋆⊤ R I -η Σ t R Σ -Σ R I -η Σ t Rw ⋆ I + w ⋆⊤ R I -η Σ t R Σ R I -η Σ t Rw ⋆ II .. (D.4)Lemma D.5. On E 1 ∩ E 2 , we haveI ≲ 1 tβ 2 and II ≲ 1 ηtβ 2 •hold with probability at least 1 -δ. By Lemma D.5, we obtain that with probability at least 1 -δ,Bias ≲ I + II ≤ 1 tβ 2 + 1 ηtβ 2 ≲ 1 ηtβ 2 (D.5)where the last inequality is by η ≲ 1/∥Σ∥ ≲ 1.Variance = η 2 Σ 1/2 R 1 n t i=1 I -η Σ i-1 X ⊤ ϵ 2 2 = η 2 n 2 ϵ ⊤ X R t i=1 I -η Σ i-1 RΣ R t i=1 I -η Σ i-1 RX ⊤ ϵ = η 2 n 2 ϵ ⊤ X R t i=1 I -η Σ i-1 R Σ -Σ R t i=1 I -η Σ i-1 RX ⊤ ϵ I + η 2 n 2 ϵ ⊤ X R t i=1 I -η Σ i-1 R Σ R t i=1 I -η Σ i-1 RX ⊤ ϵ II . (D.6)Lemma D.6. On E 1 ∩ E 2 , with probability at least 1 -δ, we haveI ≲ η 2 t n 2 • RX ⊤ ϵ 2 2andII ≲ ηt log t n 2 • RX ⊤ ϵ 2 2.By applying Lemma D.6 to Eq.(D.6), we obtain thatVariance = I + II ≲ η 2 t n 2 • RX ⊤ ϵ 2 2 + ηt log t n 2 • RX ⊤ ϵ 2 2 ≲ ηt log t n 2 • RX ⊤ ϵ 2 2. (D.7)Lemma D.7. with probability at least 1 -δ, we have1 n • RX ⊤ ϵ 2 2 ≲ σ 2 Tr(RΣR) log (d/δ) n + σ 2 sTr(Σ) log 2 (d/δ) n 2By applying Lemma D.7 to Eq.(D.7), we obtain thatVariance ≲ ηt log t • σ 2 Tr(RΣR) log (d/δ) n + σ 2 sTr(Σ) log 2 (d/δ) n 2 . (D.8)Combining Eq.(D.5) and Eq.(D.8), we obtain thatE( w t ) ≤ 1 ηtβ 2 + ηt log t • σ 2 Tr(RΣR) log (d/δ) n + σ 2 sTr(Σ) log 2 (d/δ) n 2 ≲ log t β σ 2 Tr(RΣR) log (d/δ) n + σ 2 sTr(Σ) log 2 (d/δ) n 2 , when ηt ≃ 1 β • σ 2 Tr(RΣR) log (d/δ) n + σ 2 sTr(Σ) log 2 (d/δ) n 2 -1/2 .Proof of Lemma D.1. Since y i = d j=1 w ⋆ j x ij + ϵ i , then we haver i = 1 n n j=1 x ji y j = 1 n n j=1 x ji • d k=1 w ⋆ k x jk + ϵ j = d k=1 w ⋆ k n n j=1 x jk x ji + 1 n n j=1x ji ϵ j . (D.9) 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∥x jk x ji ∥ Ψ1 ≤ C Σ kk Σ ii ≤ K,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 have1 n n j=1 x jk x ji -E[x 1k x 1i ] = 1 n n j=1 x jk x ji -Σ ki ≤ K • max log (d/δ) n , log (d/δ) n = K • log (d/δ) n , (D.10)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 have1 n j=1 x ji ϵ j ≲ K • log (d/δ) n . (D.11)Combining Eq.(D.9), Eq.(D.10) and Eq.(D.11), we have| r i -r i | ≲ K • log (d/δ) n d k=1 |w ⋆ k | + K • log (d/δ) n = (∥w ⋆ ∥ 1 + 1)K • log (d/δ) n .By definition of R and R, we obtain∥ R -R∥ 2 = max i | r i -r i | ≤ K(∥w ⋆ ∥ 1 + 1) • log (d/δ) n ≤ K s∥w ⋆ ∥ 2 2 + 1 • log (d/δ) n ≲ K • s log (d/δ) n ,which completes the proof.Proof of Lemma D.2. By Lemma D.1, for any j ∈ S, with probability at least 1 -δ, we have|r i -r j | ≲ s log (d/δ) n ≲ β/2 ≤ |r j |/2, (D.12)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:∥ R Σ R -RΣR∥ 2 = ∥ R Σ R -R Σ R + R Σ R -RΣ R + RΣ R -RΣR∥ 2 ≤ ∥ R Σ R -R Σ R∥ 2 I + ∥R Σ R -RΣ R∥ 2 II + ∥RΣ R -RΣR∥ 2 III . (D.13)Next, we proof the bound for I, II and III separately.For term I,I = ∥ R Σ R -R Σ R∥ 2 = ∥ R -R Σ R∥ 2 ≤ ∥ R -R∥ 2 • ∥ Σ∥ 2 • ∥ R∥ 2 ≤ ∥ R -R∥ 2 • ∥Σ∥ 2 + ∥ Σ -Σ∥ 2 • ∥R∥ 2 + ∥R -R∥ 2 , (D.14)where the last line is due to triangle inequality. By Lemma D.3, with probability at least -δ/3, we have∥ Σ -Σ∥ 2 ≲ ∥Σ∥ 2 • max r(Σ) n , r(Σ) n , log (1/δ) n , log (1/δ) n ≲ ∥Σ∥ 2 • max r(Σ) + log (1/δ) n , r(Σ) + log (1/δ) n . (D.15) By Lemma D.1, we obtain that ∥ R -R∥ 2 ≤ K • s log (d/δ) n ≲ 1 (D.16)holds with probability at least 1-δ/3, where the last inequality is valid since n ≳ K 2 s∥R∥ log (d/δ). Combing Eq.(D.14), Eq.(D.15) and Eq.(D.16), we haveI ≲ K∥Σ∥ 2 s log (d/δ) n • 1 + max r(Σ) + log (1/δ) n , r(Σ) + log (1/δ) n ≤ K∥Σ∥ 2 s log (d/δ) n • 1 + r(Σ) + log (1/δ) n + r(Σ) + log (1/δ) n . (D.17)For term II, we can decompose II as follows:∥R Σ -Σ R∥ 2 ≤ ∥R Σ -Σ R∥ 2 II.a + ∥R Σ -Σ R -R ∥ 2 II.b.For term II.a, by using Lemma D.3, we have with probability at least 1 -δ/3,II.a ≲ ∥RΣR∥ 2 • max r(RΣR) n , r(RΣR) n , log (1/δ) n , log (1/δ) n ≲ ∥RΣR∥ 2 • max r(RΣR) + log (1/δ) n , r(RΣR) + log (1/δ) n ≤ ∥RΣR∥ 2 • r(RΣR) + log (1/δ) n + r(RΣR) + log (1/δ) n (D.18)Similar to the proof for bounding I, we can obtain thatII.b ≲ K∥Σ∥ 2 s log (d/δ) n • 1 + r(Σ) + log (1/δ) n + r(Σ) + log (1/δ) n . (D.19)For term III, we haveIII = ∥RΣ R -R ∥ 2 ≤ |R∥ 2 ∥Σ∥ 2 K(∥w ⋆ ∥ 1 + 1) • s log (d/δ) n , (D.20)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 -δ,∥ R Σ R -RΣR∥ 2 ≤ I + II + III ≲ K∥Σ∥ 2 (∥w ⋆ ∥ 1 + 1) log (d/δ) n • 1 + r(Σ) + log (1/δ) n + r(Σ) + log (1/δ) n + ∥RΣR∥ 2 • r(RΣR) + log (1/δ) n + r(RΣR) + log (1/δ) n + ∥R∥ 2 ∥Σ∥ 2 K(∥w ⋆ ∥ 1 + 1) • log (d/δ) n ≤ (K∥Σ∥ 2 (∥w ⋆ ∥ 1 + 1) + ∥RΣR∥ 2 + ∥R∥ 2 ∥Σ∥ 2 K(∥w ⋆ ∥ 1 + 1)) • log (d/δ) n • 2 + r(Σ) + log (1/δ) n + r(Σ) + log (1/δ) n + r(RΣR) + log (1/δ) n + r(RΣR) + log (1/δ) n ≲ C cov • r(RΣR) + log (1/δ) n + r(Σ) log (d/δ) + r(RΣR) + log(d/δ) n + r(Σ) log (d/δ) + log 3/2 (d/δ) n 3/2 ≲ C cov • poly(log (d/δ)) • r(RΣR) n + r(Σ) + r(RΣR) n + r(Σ) n 3/2 ,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 haveR Σ -Σ R 2 = R Σ -Σ R + R Σ -Σ R -R + R -R Σ -Σ R + R -R Σ -Σ R -R 2 ≤ R Σ -Σ R 2 + R Σ -Σ R -R 2 + R -R Σ -Σ R 2 + R -R Σ -Σ R -R 2 .Following the proof of Lemma D.4, we can prove that with probability at least1 -δ, R Σ -Σ R 2 ≲ α(n, δ) ≤ 1/t, (D.21)where the last inequality is by E 2 . By Eq.(D.21), we haveR Σ -Σ R ⪯ 1/t • I.Hence, we obtain thatI ≲ w ⋆⊤ R I -η Σ t • 1/t • I • I -η Σ t Rw ⋆ = 1 t w ⋆⊤ R I -η Σ 2t Rw ⋆ ≤ 1 t w ⋆⊤ RRw ⋆ (by I -η Σ 2t ⪯ I) ≤ 1 t ∥w ⋆ ∥ 2 2 , (D.22)where the last line by R ⪯ 2 β • I. For the term II, we haveII = w ⋆⊤ R I -η Σ t R Σ R I -η Σ t Rw ⋆ ≲ 1 ηt w ⋆⊤ RRw ⋆ 1 ηtβ 2 ∥w ⋆ ∥ 2 2 ≤ 1 ηtβ 2 , (D.23)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 haveI = η 2 n 2 ϵ ⊤ X R t i=1 I -η Σ i-1 R Σ -Σ R t i=1 I -η Σ i-1 RX ⊤ ϵ ≲ η 2 tn 2 ϵ ⊤ X R t i=1 I -η Σ i-1 t i=1 I -η Σ i-1 RX ⊤ ϵ ≤ η 2 t n 2 ϵ ⊤ X R RX ⊤ ϵ = η 2 t n 2 • RX ⊤ ϵ 2 2, where the second last line is byt 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 haveII = η 2 n 2 ϵ ⊤ X R t i=1 I -η Σ i-1 R Σ R t i=1 I -η Σ i-1 RX ⊤ ϵ = η n 2 ϵ ⊤ X R   t i,j=1 I -η Σ i+j-2 η R Σ   RX ⊤ ϵ ≤ η n 2 • ( t i,j=1 1 i + j -1 ) RX ⊤ ϵ 2 2 ≤ ηt n 2 • ( t i=1 1 i ) RX ⊤ ϵ 2 2 ≲ ηt log t n 2 • RX ⊤ ϵ 2 2 ,where the last inequality is by the fact that t i=11 i ≲ log t.Proof of Lemma D.7. First, we can decompose1 n • RX ⊤ ϵ 2 2 by 1 n • RX ⊤ ϵ 2 2 ≲ 1 n • RX ⊤ ϵ 2 2 + 1 n • R -R X ⊤ ϵ 2 2.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∥ϵ j z ji ∥ Ψ1 ≤ Cσ G ii .By applying Bernstein's inequality Vershynin [46, Theorem 2.8.1], for any 1 ≤ i ≤ d, we have that1 n n j=1 ϵ j z ji -E[ϵ 1 z 1i ] = 1 n n j=1 ϵ j z ji ≲ σ G ii • max log (d/δ) n , log (d/δ) n = σ G ii • log (d/δ) n (D.24)hold with probability 1 -δ 3d , where the last inequality is due to n ≥ O(log(d/δ)). By taking the union bound, we obtain that1 n n j=1 ϵ j z ji ≲ σ G ii • log (d/δ) nholds for any i, with probability 1 -δ 3 . Then we haveI = d i=1   1 n n j=1 ϵ j z ji   2 ≲ d i=1 σ 2 G ii • log(d/δ) n = σ 2 Tr(RΣR) log(d/δ) n .In the same way, we can prove that with probability at least 1 -δ/3,1 n X ⊤ ϵ 2 2 ≲ σ 2 Tr(Σ) log(d/δ) n . (D.25)By applying Lemma D.1, with probability at least 1 -δ/3, we haveEQUATIONBy Eq.(D.25) and Eq.(D.26), with probability 1 -2δ/3, we have1 n • R -R X ⊤ ϵ 2 2 ≤ R -R 2 2 1 n X ⊤ ϵ 2 2 ≲ σ 2 sTr(Σ) log 2 (d/δ) n 2 .By taking the union bound, we derive the desired result.To simplify the notations, we use w t to denote w t gd . Lemma E.1. with probability at least 1 -δ, we haveΣ -Σ ≲ α(n, δ), (E.1) 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 -δ, Σ -Σ ≲ 1/t. Proof of Lemma E.1. By Lemma D.3, we have ∥ Σ -Σ∥ 2 ≤ c∥Σ∥ 2 • max r(Σ) n , r(Σ) n , log (1/δ) n , log (1/δ) n ≲ max r(Σ) + log (1/δ) n , r(Σ) + log (1/δ) n ≤ r(Σ) + log (1/δ) n + r(Σ) + log (1/δ) n (E.2)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:E := RΣR∥ 2 ≲ α(n, δ) ≤ 1/t .By Lemma E.1, P(E) ≥ 1 -δ. Hereafter, we condition on E.Bias-variance Decomposition Similar to Eq.(D.1), we havew t = I -I -η Σ t w ⋆ + 1 n t i=1 I -η Σ i-1 X ⊤ ϵ. (E.3)In the same way, we can decompose the risk E(w t ) byE(w t ) = Σ 1/2 I -η Σ t w ⋆ 2 2 Bias + η 2 Σ 1/2 1 n t i=1 I -η Σ i-1 X ⊤ ϵ 2 2 Variance . (E.4)Bounding the BiasBias = w ⋆⊤ I -η Σ t Σ I -η Σ t w ⋆ = w ⋆⊤ I -η Σ t Σ -Σ I -η Σ t w ⋆ I + w ⋆⊤ I -η Σ t Σ I -η Σ t w ⋆ II .Similar to the proof of Lemma D.5, we have the following lemma. Lemma E.2. On E, we haveI ≲ 1 tandII ≲ 1 ηthold with probability at least 1 -δ.As a result, the bound of the bias term is given byBias ≤ 1 ηt + 1 t ≲ 1 ηt . (E.5)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 thatVariance ≲ ηt log t • 1 n • X ⊤ ϵ 2 2 ≲ ηt log t • σ 2 Tr(Σ) log (d/δ) n . (E.6)Combining Eq.(E.5) and Eq.(E.6), we obtain thatE(w t ) ≲ 1 ηt + ηt log t • σ 2 Tr(Σ) log (d/δ) n ≲ log t • σ 2 Tr(Σ) log (d/δ) n , when ηt ≃ σ 2 Tr(Σ) log (d/δ) n -1/2F Proof for Theorem 5.3To simplify the notation, we use w t to denote w t gd and w t to denote w t gd .When Σ = I, by Eq.(D.2), we haveE( w t ) = R I -η Σ t Rw ⋆ 2 2 Bias + η 2 R 1 n t i=1 I -η Σ i-1 X ⊤ ϵ 2 2.Following the proof of Theorem 5.1, it holds thatVariance ≲ ηt log t • σ 2 log (d/δ) n + σ 2 sd log 2 (d/δ) n 2with probability at least 1 -δ, when n ≳ t 2 sd 2/3Similar to the proof of Lemma D.2, we can prove thatr i ≥ r i 2 ∀i ∈ S, r i ≲ 1 ∀i,with probability at least 1 -δ.When Σ = I, by Lemma D.4, we have thatR Σ R -RΣR 2 ≲ β 2 tholds 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 haveR I -η Σ t Rw ⋆ 2 2 ≤ R 2 2 • I -η Σ t Rw ⋆ 2 2 ≤ w ⋆⊤ R I -η Σ 2t Rw ⋆ ≲ w ⋆⊤ R I -η RΣR - β 2 t • I 2t Rw ⋆ = i∈S (w ⋆ i / r i ) 2 • 1 -η (w ⋆ i ) 2 - β 2 t 2t ≤ s • 1 -ηβ 2 /2 2t ,where the last line is by the definition of β. When t ≳ log σ 2 ns / 2 log 1 -ηβ 2 /2 , we haveBias = R I -η Σ t Rw ⋆ 2 2 ≤ σ 2 n . (F.1)When ηβ 2 /2 ≤ 1/2, there exist a c > 0, such thatlog 1 -ηβ 2 /2 ≥ cηβ 2 /2.Hence, the variance term is bounded byVariance ≲ ηt log t • σ 2 log (d/δ) n + σ 2 ∥w ⋆ ∥ 2 1 d log 2 (d/δ) n 2 ≲ σ 2 log 2 ns/σ 2 log 2 (d/δ) β 2 • s n + ds n 2 , (F.2)where the last line is by ∥w ⋆ ∥ 1 ≤ s • ∥w ⋆ ∥ 2 2 =w λ = 1 n • Σ + λ • I -1X ⊤ y, we haveE w ⋆ [E( w λ )] = I -Σ + λI -1 Σ w ⋆ 2 2 + 1 n • Σ + λ • I -1 X ⊤ ϵ 2 2 ≥ 1 n • Σ + λ • I -1 X ⊤ ϵ 2 2.By Lemma D.3, when 1 2 • I ⪯ Σ ⪯ 2 • I, with probability at least 1 -δ, we haveE w ⋆ [E( w λ )] ≥ 1 n • Σ + λ • I -1 X ⊤ ϵ 2 2 = 1 n 2 • ϵ ⊤ X Σ + λI -2 X ⊤ ϵ ≥ 1 n 2 (2 + λ) 2 • ϵ ⊤ XX ⊤ ϵ,where the last line is due to the fact that Σ + λI ⪯ (2 + λ) • I. Lemma F.1. Given X such that 1 2 I ⪯ Σ ⪯ 2I, it holds that1 n X ⊤ ϵ 2 2 ≳ σ 2 d n ,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 have1 n X ⊤ ϵ 2 2 = 1 √ n UΛV ⊤ ϵ 2 2 = 1 √ n ΛV ⊤ ϵ 2 2 = 1 √ n Λ ϵ 2 2 = 1 n d i=1 σ 2 i ϵ 2 i ,where ϵ = V ⊤ ϵ ∼ N(0, I). By ][Lemma 22], we have1 n X ⊤ ϵ 2 2 -E 1 n X ⊤ ϵ 2 2 ≲ σ 2 max    d i=1 σ 4 i log (1/δ) n , max i σ 2 i log (1/δ) n    ≲ σ 2 max d log (1/δ) n , log (1/δ) n , (F.3)where the last line is valid since σ 2 1 , . . . , σ 2 d is the eigenvalues of Σ = 1 n X ⊤ X and 1 2 I ⪯ Σ ⪯ 2I. By Eq.(F.3), we obtain that1 n X ⊤ ϵ 2 2 ≥ E 1 n X ⊤ ϵ 2 2 -σ 2 max d log (1/δ) n , log (1/δ) n = σ 2 d i=1 σ 2 i -σ 2 max d log (1/δ) n , log (1/δ) n = σ 2 d n -σ 2 max d log (1/δ) n , log (1/δ) n (by 1 2 I ⪯ Σ ⪯ 2I) ≲ σ 2 d n ,where the last line is due to d ≥ O(log (1/δ)).Next, we define the event E as follows:E ridge := 1 2 I ⪯ Σ ⪯ 2I, 1 n X ⊤ ϵ 2 2 ≳ σ 2 d n .By Lemma F.1, we haveP(E) ≥ 1 -δ when n ≥ O(d) ≥ O(log (1/δ)). On E ridge , we haveE w ⋆ [E( w λ )] ≳ σ 2 d (1 + λ) 2 n . (F.4)When d ≳ n + log (1/δ), by Lemma D.3, with probability at least 1 -δ, we have that d 2 • I ⪯ XX ⊤ ⪯ 2d • I. Hereafter, we condition on this event. By direct calculation, we can decompose the excess risk byE w ⋆ [E( w λ )] = E w ⋆ I -Σ + λI -1 Σ w ⋆ 2 2 + 1 n • Σ + λ • I -1 X ⊤ ϵ 2 2.For the first term, we haveE w ⋆ I -Σ + λI -1 Σ w ⋆ 2 2 = E w ⋆ I -X ⊤ XX ⊤ + nλI -1 X w ⋆ 2 2 = (1 - n d )E w ⋆ ∥w ⋆ ∥ 2 2 , (F.5) = 1 - n d (F.6)where the last line is due to I -X ⊤ XX ⊤ + nλI -1 X is a d -n space. 1 n • Σ + λ • I -1 X ⊤ ϵ 2 2 = ϵ ⊤ XX ⊤ XX ⊤ + nλI -2 ϵ ≥ dn 2(2d + nλ) 2 • 1 n n i=1 ϵ 2 i , (F.7)where the first line is by X ⊤ X + nλI -1 X ⊤ = X ⊤ XX ⊤ + nλI -1 and the last line is byd 2(d+nλ) 2 • I ⪯ XX ⊤ XX ⊤ + nλI -2. By Tsigler and Bartlett [44, Lemma 22] , we obatain thatn i=1 ϵ 2 i -nσ 2 ≲ σ 2 n log (1/δ) + σ 2holds with probability at least 1 -δ. When n ≳ log (1/δ), we haven i=1 ϵ 2 i -nσ 2 ≥ nσ 22 holds with probability at least 1 -δ. Taking the union bound, we obtain thatE w ⋆ [E( w λ )] ≳ 1 - n d + σ 2 • dn 2(2d + nλ) 2 ≳ 1 - n d + σ 2 n (1 + λ) 2 d . (F.8)We first consider the case where n ≳ d + log (1/δ). Define the event E GD by E GD = 1 2 • I ⪯ Σ ⪯ 2I . By Lemma D.3, P(E GD ) ≥ 1 -δ. By Eq.(E.4), we haveE w ⋆ [E(w t )] = E w ⋆ I -η Σ t w ⋆ 2 2 + η 2 1 n t i=1 I -η Σ i-1 X ⊤ ϵ 2 2 ≥ η 1 n t i=1 I -η Σ i-1 X ⊤ ϵ 2 2 = η 2 n 2 • Σ I -I -η Σ t -1 -1 X ⊤ ϵ 2 2 ≳ η 2 n 2 • Σ + 1 ηt • I -1 X ⊤ ϵ 2 2 ≳ σ 2 η 2 d (1 + 1/(ηt)) 2 n ,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 ⊤ ≺ 2dI . By Lemma D.3, P(E ′ GD ) ≥ 1 -δ.E w ⋆ [E(w t )] ≥ E w ⋆ I -X ⊤ XX ⊤ + n ηt I -1 X 2 2 + 1 n X ⊤ XX ⊤ + n ηt I -1 ϵ 2 2 = 1 - n d + σ 2 n 1 + 1 ηt 2 d ,where we use the results from Appendix F.2.Let w ols be the OLS estimator. It is easy to see w ols = w 0 . Hence, we haveE w ⋆ [E(w ols )] ≳ σ 2 d n n ≳ d + log (1/δ) 1 -n d + σ 2 n d d ≳ n + log (1/δ),holds with probability at least 1 -δ.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. 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.EQUATIONTraining 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. 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.Here for the ease to presentation and comparison, we slightly abuse the notation of E(w) by extending it to w t gd , although E(w) is originally defined for the estimator for the raw feature vector x.We also tried other settings with fewer heads or layers, but even with delicate hyperparameter tuning, decoder-only transformers with fewer heads or layers consistently failed to learn how to solve our sparse linear regression problem. A possible reason is that decoder-only transformers first need to learn the causal structure[35] and then apply an optimization algorithm to the in-context entries, which is more challenging than our encoder-based settings. 


Decision: This paper presents an interesting analysis of how transformers utilize multi-head attention for in-context learning, focusing on a sparse linear regression task. The key findings are that multiple attention heads are utilized in the first layer for data preprocessing, while subsequent layers primarily use a single head for optimization, supported by theoretical and emipirical explanations. All reviewers agree that the submission contributes significantly, and the comments have been addressed after the rebuttal. The AC concurs with the decision to accept the submission.

Reviews: review: paper_summary: paper_summary: This work seeks to provide a deeper exploration of the use of multi-heads, at different layers in a Transformer, to perform in-context
learning 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.
strength_weakness: strength_weakness: Weaknesses:
- 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.
questions: questions: The 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?
limitations: limitations: The authors have clearly stated any limitations with their current work. Additionally, the reviewer does not foresee any potential negative societal impact from this paper.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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$.
2. 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".
3. The orthogonal design is an over-simplified setting.
questions: questions: 1. I find the observation that only one head dominates in subsequent layers very interesting. Does this also occur in other settings of ICL tasks?
limitations: limitations: The orthogonal design is an over-simplified setting, which the authors do not properly address.
1. 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.
2. 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?
Overall, 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.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: The 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.
questions: questions: Line 108: should this instead read "to perform a tractable theoretical investigation"?
Line 209: typo "proprocessing" instead of "preprocessing"
limitations: limitations: The authors have discussed the limitations I mentioned in the section above. I do not see any clear potential negative societal impact for this work.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: The 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. 

Given the above, there are several main weaknesses regarding the contributions in this paper:

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

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

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

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

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

References:

[1] Huang, Yu, Yuan Cheng, and Yingbin Liang. "In-context convergence of transformers." arXiv preprint arXiv:2310.05249 (2023).

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

[3] Cui, Yingqian, et al. "Superiority of multi-head attention in in-context linear regression." arXiv preprint arXiv:2401.17426 (2024).

[4] Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. "Trained transformers learn linear models in-context." arXiv preprint arXiv:2306.09927 (2023).

[5] Ahn, Kwangjun, et al. "Transformers learn to implement preconditioned gradient descent for in-context learning." Advances in Neural Information Processing Systems 36 (2024).
questions: questions: I 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.
limitations: limitations: NA

Comments: comment: I thank the authors for their responses and additional experiments. I will increase my rating from 6 to 7 as I support accepting this work.
Thanks for the authors' responses. I would like to raise my score to 7.
I appreciate the authors in addressing my comments. I have raised my score from 4 to 5.