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 2024 poster (2024)
Keywords: in-context learning
transformers
deep learning theory
learning theory
Primary area: 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
Xingwu Chen
Lei Zhao" University of Pennsylvania
The University of Hong Kong xingvu@connect. hewianneecchuu.haa... hku.hk hk
leizhao7@tharton upenn edu
Difan Zou
The University of Hong Kong dzou@cs .huu.
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 anal- yses accordingly. However, these works mostly focus the expressive power of transformers by designing specific parameter constructions, lacking comprehen- sive understanding of their inherent working mechanisms post-training. this study, consider sparse linear regression problem and investigate how trained multi-head transformer performs in-context learning. We experimentally discover that the utilization of multi-heard 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 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, 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 understanding the more intricate mechanisms hidden within trained transformers.
prove that such a preprocess sesprocess-then-opphee then-optimize
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 33
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 tackle challenging tasks related to memorization 33
reasoning 24
Equual contribution.
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

28 14].
function approximation 22 [3), causal relationship [35,, and simulating complex circuits 21 B1. These endeavors typically aim to enhance our understanding of the capabilities and limitations of transformers when configured with varying numbers of heads [6] 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 11 47). These investigations have primarily focused on linear regression tasks with a Gaussian prior, demonstrating that transformer with L layers can mimic L steps of gradient descent on the loss defined by contextual examples both theoret- ically and empirically 3 [3]. These observations have immediately triggered 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 [TS] Newton methods [2 19 and ridge regression g However, these theoretical works 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 noisy linear model with sparse ground truth w E Rd with ||w*|lo llo S 8 < d, and train a transformer model with multiple layers and heads. While a line of works also investigates this problem (20 g 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 the experimental findings, we conjecture that muti-layer transformer may exhibit preprocess-then-opiiiiz algorithm on the context examples. Specifically, transformers utilize all heads in the initial layer for data preprocessing and subsequently employ a single head in sub- sequent 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 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 prepro- cessing 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 implemented using multiple heads in the first layers and single head in subsequent layers.
then-optimize algorithm can be

---

We further validated our proposed mechanism by comparing the performance of the preprocess- then-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-optinee 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 terms of excess risk.
To further validate our theoretical framework, we conducted additional experiments. Specifically, we performed probing on the output 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 I. Linear Regression. We consider sparse linear models where (x, ~ P Plin, is sampled as x ~ N(O, M), = (w",x) + (0, where the Σ is a diagonal matrix and ground truth w Rd satisfies [|w S Then, we define the population risk of a parameter as follows:
((w) == (x.))>> [(xx, w) - =
Moreover, we are interested the excess risk, i.e. the gap between the population risk achieved by and the optimal one:
El ((w) == L(w) min 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,, h hm] E Rm It computes different queries, keys, and values, and then concatenates the output of each head:
Attn (H,,HHH 01) H + Voncat(V,sfmxx
Q1). V (K,,, Q:)),
where V, = Wv; H, Q: = Wo, H, V, = Vv, H and 0, = { (W,; K, We, E RRda//hedd learnable parameters. The MLP then applies nonlinear element-wist operation:
i=1
MLP(H, 82) = W, WiReLU(W,Attn(E 01))
(2.1)
where 02 = (W1, W2 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 tractable theoretical investigation 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]5] [3 32 3. In particular, the i-th layer TF, performs the following update on the input sequence (or hidden state) ( as follows:
H (H W, (H(i-)) Concat((V,MK, Q.Vii)) M ==
In
0
R pmxm
0
(2.2)
where (W,, WK.. WQ. E R Anid xdd and W1 E Rdus are learnable parameters, note that ignore the ReLU activation in 2.I. so merge the parameter W, and W₂ into matrix W, Besides, the mask matrix M is included in the attention to constrain the model focus the first n To adapt the transformer for
in-context examples rather than the subsequent m--n queries [3] 32 (2 544
introduce additional linear layers We R (d-11) 1)x x dha and Wo E 4aax for input embedding and output projection, respectively. Mathematically, let E denotes the input sequences with n in-context example followed by q queries,
solving sparse linear regression problems,
X1
X2
Xn
Xx++1
Xn+ty
E
(2.3)
y1
32
0
3

---

Input
Output Projection
Multi-Head Attention
subsequent layers l
||w/lo S X - 2(OE)) o(
Add Norm
FFN
Y1- (xL, w. ).
X
optimizing
Add & Norm
In Context Sporse Lincor regression
Multi-Head Attention
First loyer [1=
X1 ... In 1 ... xntq
y1 Yn 0
Embedding Projection
Input
preprocessing
(a) Overview of the experiments, including task, data, transformer architecture, and our insights.
Jayer
prandd
00000
42
Wntortest cxamptes
1ith head
14th layer
(b) ICL with Varying Heads
(c) Heads Assessment
(d) Pruning and Probing
Figure I: Experimental Insights into Multi-head Attention for In-context Learning
Then model processes the input sequence E, resulting in the output ý E R x(a+a));
= Woo o TFL .... TF, o We
here, L is the layer number of the transformer, and Vitn is the prediction value for the query X+tn- During training, we set q for efficiency, and for inference and theoretical analysis, we set q = and define the in-context learning excess risk EICL as:
&CCL= (x ( 1n+1) 0²
Notations. For two functions f(x) 2 0 and g(x) 2 0 defined on the positive real numbers (r > 0),
we write f(x) S g(x) there exists two constants c, xu > such that Va 2 Ty. f(x) S c g());
write (x) 2 g(x) if g(x) & [(x);
write f(x) g(x) if ((x) S g(x) and g(x) S f(x). If
f()) S g(x). we can write f(x) as (g{()). We can also write write f(x) as Ö(g(z)) there exists
a constant k > 0 such that f(x) S g(x) lo)"(z). (z).
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 81. the hidden mechanism behind the trained transformer for solving this problem remains unclear. To this end, we design series of experiments, utilizing techniques like probing [5] and pruning (7 to help us gain initial insights into how the trained transformer utilizes multi-head attention for this problem. For all experiments in Sections and 6 we choose an encoder- based architecture as the backbone (see Figure Ia), set the hidden dimension dnid to 256, and use the input sequence format shown Eq.2..3) where d 16, s
N(O,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), we know that the j-th head at the i-th layer corresponds
to the subspace of the intermediate output from (j - 1) dnia/h / to j dnia/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 AECC(())) Then we normalize the risk changes in the same layer evaluate their relative importance:
AEIC((E.) Ek A&
Was
(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 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 ( to evaluate the prediction performance for different layers. An example can be found in Figure Id we can find that the "pruned transformer" and the original model exhibit almost the same performance for each layer. Additionally, compared 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 preproces-thhen-ptiiia 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 the preprocessed in-context examples using the subsequent layers.
We note that (24] adapts similar two-phase idea to explain how transformer learning specific functions 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 In-context Examples
First, the multihead attention is designed to facilitate 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, possible data preprocessing method is 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 x on the complement of the nonzero support, as a data preprocessing procedure, and perform projected gradient descent obtain the optimal solution.
5

---

In general, the nonzero support of w is intractable to the learner, so that one cannot perform idealized masking- 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= wf + Çı, implying that ri == E[r,y] BI>i- wwa;; - xid] + B[[sx = [x[[ if considering independent data features. Then it is clear that such correlation between the feature and label will be nonzero only when |m=]| * 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 provided with n in-context examples {(x,, Di) such correlations can be estimated accordingly: Pj- = + Di=1 FigTii, which will be further used to perform the data preprocessing on the in-context examples. We summarize this procedure in Alg.
Algorithm Data preprocessing for in-context examples
1: Input: : Sequence with {(x; v:)) (xx, 2 as in-context examples/queries. 2: for k I,..., do
3: Compute X5 by Xx = Rxx, where R = diag Fi, Fa,....ffal Fa), where fj is given by
(4.1)
4: end
5: Output: Sequence with the preprocessed in-context examples/queries {(x,, v.) ((R,, 0)
The preprocessing procedure aligns well with the structure of multi- 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 (W k, H) "W),Wa, Wq. H. 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,. Here, WK, and Wq. 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 Wv; 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. [ We formally prove this in the following theorem. Proposition 4.1 (Single-layer multi-head transformer implements Alg. D. There exists a single-layer transformer function TF with heads and daia 3d hidden dimension, together with an inpul embedding layer with weight Wg E Rd xd that can implement Alg. [ Let E be the input sequence defined in Eq. 2.3) and x; Rx be the preprocessed features defined in Alg. it holds that
81
X2
X,
X Xn+1
Xn+t
H = TF, W
=
y1
y2
yn
(4.2)
where... ... in third row implies arbitrary values.
4.2 Optimizing Over Preprocessed In-Context Examples
Based 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] [5] 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 linear projection w to rearrange the dimensions of the sequence processed by the multi-head attention, resulting in the hidden state (+) of each layer. We refer 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 fx, Vi H= Then, the following proposition shows that the subsequent layers of transformer can implement multi-step gradient descent on the preprocessed in-context examples { {(4, Ne.......
Proposition 4.2 (Subsequent single-head transformer implements multi-step GD). There exists transformer with k layers, head, daia = 3d, let Vn+1 be the prediction representation of the L-th layer, then it holds that Km+11) (w ga data feature, w is defined as wad = 0 and as follows for È = 0, O....,k k -1:
( xn+1), where Xn+1 Rx,+1 denotes the preprocessed
s+1 gd
gd
where L(w)
Mi (w,x;))
(4.3)
The proof of Propositions 4.1 and 4.2 can be found in Appendix [ 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 pegrroces-ttennopttiii then-opptimization 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 Preproces-hhen-optiniz Algorithm
In this section, will develop the theory to demonstrate the improved performance of the preprocess- then-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 and respectively.
We first denote w gd the estimator obtained by (-step GD {(x.. v=) which can be viewed as the solution generated by the & 1-layer transformer based on our discussion in Section 4 and w as the estimator obtained by 1-step GD on first need to redefine the excess risk of GD on {(x,, 3k 1= Note that in our algorithm, the learned predictor takes the form x + [Rx, w W Consequently, the population risk of parameter w is naturally defined as ( gd) = 1/2 E (x.x))>P ( [([Rx, W'a) - -] and the excess risk is then defined as E ((w)= ) == I(w)= min L( (w) Next, we provide the upper bound of the excess risk for E {(w(a) and E E(w')) (d respectively.
((x; 3/) )== Before presenting our main theorem, we
Theorem 5.1. Denote S = {i w * 0} and R = diag(ri,....,?) where T3 = Σi=
w Eig. Suppose that there exist B > such that mintes |r,l 2 B, |R| 21 E|ln ||w ll2 ~ O(1) and
n 2 1/32 ..s (Tr 2/3
E + (r(RE))) poly(log (d/5)). Then set 17 S 1/ RER and
(TER log (d/8) o'sTr(E) log" (d/8)) -1/2
nt
3
n
n²
it holds that
logt
a 2Tr TTr (RER) log (d/8)
22sTr( (E)((E)))))))) (E) log² (d/8)
( gd
19
n
n²
with probability at least
Theorem 5.1 provides an upper bound on the excess risk achieved by the
algorithm, where tuned learning rate 17 to balance the bias and variance error. Then, it can be seen that the risk bound is valid (T.RR))) / /n - O and Tr(E)s/n² + 0 when n - co. This can be readily satisfied if we have ||w'|l and Tr(E) be bounded by some reasonable quantities that are independent of the sample size n, which 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 implics better performance. This implics the ability of the proposed preprocess-then tippooesstheen-opiiiize for discovering and leveraging the nice sparse structure of the ground truth.
'Here for the ease to presentation and comparison, we slightly abuse the notation of E(w) by extending it to We although E(w) is originally defined for the estimator for the raw feature vector
7

---

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 |ا|||
[[ [12 = O(1) and 2 ((Tr(Z) + log (1/6)). When S
1/ Eاك and ntt (="1r12) log (d/6)
1/2
it holds that
o
Tr (T log (d/6)
E(w'd) S logt
n
with probability at least 1- 5.
We are now able to make rough comparison between the excess risk bounds in Theorems 5. and 5.2 Then, it is clear that (( ga) - E(w requires Tr(RER) z Tr(E) and s/ /nnn/ 3² S nn. Specifically, we can consider the case that E to be diagonal matrix, assume W ~ Uf-1/ 1/v
has a restricted uniform prior for E S and minges En 2 1/k for some constant k > 1, we can get B 2 ()/(kk²), thus Tr(RER)/ S x² Liwy #o En and s/(n2 B²) S 2² //2.. Note that | = d, then if the covariance matrix Σ has flat eigenspectrum such that Eies Σn Lield Σn Tr(E), we have TrRER)/33 < Tr(E) and s/(n S //22 if 8 of min VTt)). This suggests that the preprocess-theme the-tooiimmzaaiionn algorithm can outperform the standard gradient descent for solving a sparse linear regression problem with of min vn)).
To make more rigorous comparison, next consider the example where I; Lid. N(O, 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| probability from the set ........ for E S, ||w0ll2 = ((1) and n 2 3
selected such that each element is chosen with equal ... ... and wt ~ U{-1/v8,, 1 /&) has restricted uniform prior
Then there exists a choice of and such that
(
ds²
El ((wa)) S log (ns/o²) log² (d/0)
n
with probability least 6. Besides, let W be the ridge regression estimator with regularized parameter X, and wals be the OLS estimator, it holds that
odd
n 2 d log (1/8) d d2n n log (1/6),
Ew- Bw-[E[(w] 2
2
d2
with probability at least 1 8, where E (WA, Wols: w }.
It can be seen that for wide range of under- mander-arramtterizeer and over- peee-parameterieed cases, w has smaller excess risk than ridge regression, standard gradient descent, and OLS. In particular, consider the setting o² 1, in the over-parameterizer setting that dl 2 n, the excess risk bound of eeeprocess-thennoptiiiie is (ds/n²²) which also outperforms the ?²(1) bound achieved by OLS, ridge regression, and standard gradient descent the sparsity satisfies condition can be certainly removed as E(w! ga also has a naive upper bound Õ(1).. In the under- parameterized case that S n, it can be readily verified that the data preprocessing can lead to (s/nn) excess risk, which strictly better than the (dd/n) risk achieved by OLS, ridge regression, and standard gradient descent. Moreover, it is well known that Lasso can achieve (s/n)) excess risk bound the setting of Theorem 5.3 Then, by comparing with our results, we can also conclude that the
O /d) (in fact, this
algorithm can be comparable Lasso up to logarithmic factors when d S n, while becomes worse when d Z
Experiments
In Section 3
we conduct several experiments, and based on the observations, we propose that a trained transforner can apply a preprocess-then-ond then-optimize algorithm: (1) In the first layer, the transformer can apply preprocessing algorithm (Alg. on the in-context examples utilizing multi-heal attention.

---

Trained Transformer
Embodding Profeciio
extract
entries
(b) P-Probing result for Transformers trained varying heads and data settings
Preprocessing Probing
0d pmhgg
9d
pd
powgd
14
5e x5 Tane
stips
50.155
Stops
(a) P-probing
(c) Directly apply Alg. on input data for gradient descent
Figure Supporting experiments for our preproces--then-opiimie algorithm and theoretical analysis
(2) In the subsequent layers, the transformer applies gradient descent algorithm on the preprocessed data utilizing single-head attention. While the second part is supported by extensive theoretical analysis and experimental evidence [57] 53 B) here we develop a technique called preprocessing probing (P (P-probing) on the trained transformer to support the first part of our algorithm. We also directly apply Alg. Ton the in-context examples and then check the excess risk for multiple-step gradient descent verify the effectiveness of our algorithm and theoretical analysis.
P-probing: To verify the existence of preprocessing procedure in the trained transformer, we develop a "preprocessing probing" (P-probing) technique on the trained transformers, as illustrated in (Fiiueeee 2a] For trained transformer, we first set the input sequence as Eq. (2.3). where the first examples {x) have the corresponding labels {y) and the following q query entries only have { in the sequence. Then, we extract the last q vectors in the output hidden state H 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 { from H besides, the excess risk computed by the preprocessed data has better upper bound guarantee compared 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 HD 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. [ on the input data (x,, y: / and then implement gradient descent on the example entries (x, 91 1 and compute the excess risk with the last query (Rn+1, Bn+1). we refer this procedure 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 D gd = and tune the learning rate v) for each model by choosing from [0,10-1, 10000--- 10-3,, 10-4 10-5 10-6] with the lowest average excess risk.
Based on Figure 2b can observe that compared 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 lower risk. This aligns well with our experiment result Figure 2c] where compared to gd, the data preprocessed by Alg. 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 provide 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 propose preproces--then-ppmine
attention for in-context learning. Based on our empirical investigations, we 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 explored. First, our work focuses on the case of sparse linear regression, and it may be beneficial implement experiment for more challenging 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.
Acknowledgements
We would like to thank the anonymous reviewers and area chairs for their helpful comments. This work supported by NSFC 62306252, Guangdong NSF 2024A151501 Hong Kong ECS awards 27309624, and the central fund from HKU IDS.
References
[1] Jacob Abemnethy, Alekh Agarwal, Teodor Vanislavov Marinov, and Manfred K. Warmuth. 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 httppss//pprcceeednggg mlr prwpppressss........ v237 abernethy24a html
[2] 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.
[3] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers lear to imple- ment preconditioned gradient descent for in-context learning. Advances Neural Information Processing Systems, 36, 2024.
[4] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2022.
[5] Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes. arXiv preprint 1611v: 2016.
[6] Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes, 2017. URL httptt::/oopereeiiiiiieeee net//orrmmiddry/777/a
[7] Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.2, knowledge manipula- tion. arXiv preprint arXiv:2309. /4402, 2023.
[8] 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.
[9] 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(45):30063-3007 2020.
10

---

[10] 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. nethneuron- explainer/paperr/ndex (Date accessed: 14.05. 2023), 2023.
[11] Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, al. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, page 2, 2023.
[12] 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.
[13] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi- head softmax attention for in-context learning: Emergence, convergence, and optimality, arXiv preprint arXiv:2402. 19442, 2024.
[14] Xingwu Chen and Difan Zou, What can transformer learn with varying depth? case studies on sequence learning tasks, arXiv preprint arXiv:2404.0160) 2024.
[15] Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:23/2.00558,, 2023.
[16] Yinggian Cui, Jie Ren, Pengfei He, Jiliang Tang, and Yue Xing. Superiority of multi-head attention in in-context linear regression. arXiv preprini arXiv:2401.17422 2024.
[17] Nan Ding, Tomer Levinboim, Jialin Wu, Sebastian Goodman, and Radu Soricut. CausalLM is not optimal for in-context learning, September 2023.
[18] Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs, Advances in Neural Information Processing Systems, 36, 2024.
[19] Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers Learn Higher-Ore Order Optimization Methods for In-Context Learning A Study with Linear Models, October 2023.
[20] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What Can Transformers Learn In-Context? A Case Smdy of Simple Function Classes, August 2023.
[21] Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped Transformers as Programmable Computers, January 2023.
[22] Angeliki Giannou, Liu Yang. Tianhao Wang, Dimitris Papailiopoulos, and Jason D. Lee. How Well Can Transformers Emulate In-context Newton's 's Method?, March 2024.
[23] Andrey Gromov, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel A Roberts. The unreasonable ineffectiveness of the deeper layers. arXiv preprint arXiv:2403. / 77877, 2024.
[24] 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 Learning with Representations, October 2023.
[25] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
[26] Tokio Kajitsuka and Issei Sato. Are Transformers with One Layer Self-Attention Using Low- Rank Weight Matrices Universal Approximators?, July 2023.
[27] 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 httpp:://peerreiies.t http://wwwwweeeeeeeeeeooorrtrrrr

---

[28] Yingcong Li. Muhammed Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in-context learning. In International Conference on Machine Learning, pages 19565-19594. PMLR, 2023.
[29] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In International Conference on Machine Learning, pages 19689-19729. PMLR, 2023.
[30] David Lindner, János Kramár, Sebastian Farquhar, Matthew Rahtz, Thomas McGrath, and Vladimir Mikulik. Tracr: Compiled Transformers as Laboratory for Interpretability, November 2023.
[31] Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Trans- formers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations, 2023. URL httpttt//ppptreeeeeeeeee
[32] Arvind Mahankali, Tatsunori B. Hashimoto, and Tengyu Ma. One Step of Gradient Descent is Provably the Optimal In In-Context Learner with One Layer of Linear Self-Attention, July 2023.
[33] Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization Capacity of Multi- Head Attention in Transformers, October 2023.
[34] 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.03885 2024.
[35] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprini arXiv:2402.147333 14735, 2024.
[36] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkín, 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.
[37] Onkar Pandit and Yufang Hou. Probing for bridging inference in transformer language models. arXiv preprint arXiv:2104.09400, 2021.
[38] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 18):9, 2019.
[39] Shokichi Takakura and Taiji Suzuki. Approximation and Estimation Ability of Transformers for Sequeece-to-Sequeee Functions with Infinite Dimensional Input, May 2023.
[40] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transform- ers as Support Vector Machines, September 2023.
[41] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer, July 2023.
[42] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint (rXiv:2302.13977 2023.
[43] 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.
[44] Alexander Tsigler and Peter L Bartlett. Benign overfitting in ridge regression. Journal of Machine Learning Research, 24(123):1-76, 2023.
[45] 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,
12

---

[46] Roman Vershynin. High- dighhsiimeniinal probability. University of California, Irvine, 10:11, 2020.
[47] Johannes Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mord- vintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, May 2023.
[48] Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking Like Transformers, July 2021.
[49] 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 's transform- ers: State-of-the-art natural language processing. arXiv preprint arXiv: 1910.03771, 2019.
[50] 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 Learning Representations, 2023.
[51] 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.
[52] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference. In International Conference on Learning Representations, 2021.
[53] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306. 09927, 2023.
[54] Ruiqi Zhang, Jingfeng Wu, and Peter L Bartlett. In-context learning of linear transformer block: Benefits of the mlp component and one-step gd initialization. arXiv preprint arXiv:2402.14951, 2024.
[55] 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 Learning Representations, 2024. /tRp:ttts://opeerreee wteeriiwwweeeeeeeeeeeeeeeeeeeeee net/ fttt//ooomiidddeo Ass IuHnmHX
[56] Zeyuan Allen Zhu and Yuanzhi Li. Physics of language models: Part 3.1, knowledge storage and extraction. arXiv preprint arXiv:2309. 14316, 2023.
[57] 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-54 -4656, 2021.
[58) 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 Neural Information Processing Systems, 35: 2909-122990 2022.
13

---

A Additional Related Work
In addition 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 488 (55, probing the behavior of individual layers [37 [11 [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.
parallel, a body of work has also investigated how trans- formers learn these algorithms, i.c., 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 al. 3 analyze the training dynamics of 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. 1 extended these findings from linear attention to softmax settings, with [3] revealing that trained transformers with multi-head attention tend to utilize different heads for distinct tasks in various subspaces. Additionally, Tian et al. AI Li et al. 229] study the convergence of transformers on sequences of discrete tokens. Gromov et al. P3] Men et al. 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.
Training dynamics of transformers
B Additional Details for Sections and 6
Architecture and Optimization We conduct extensive experiments on encoder-only transformers with dhid 256, varying the number of heads h E {1, 8,8,, layers / {3, 5, 5,6,, and noise levels E {0, 0.1, 1, 0.2, 0.4, 0.8). For the input sequence, we sample ~ N(O, I). For w, we first sample w ~ N(O, I) E R21 and randomly choose & = 4 entries, setting the other elements to zero. Note that We don't apply positional encodings in our setting, as no positional information needed in our input setting. To further support our
algorithm, we also try decoder-only architecture(Figure 9. train models on other settings like stardand linear regression task = 16 (Figure 10p and non-orthogonal data distributions (Figure IT) as comparisons in Appendix G During training, we
12 and g 4, with batch size of 64. We utilize the Adam optimizer with a learning rate y = 10 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 trained and evaluated with the same training and evaluation dataset. We 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 X, E {1,10-1 10 10 2,103, 10 the lowest risk, as in [20).
ICL with Varying Heads
1000 *) respectively for
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.
From Figure 3
Heads Assessment Here, we set 10 and q 1, with evaluation data size of 8192. For model with h heads and layers, we train |a] models under different noise levels. We first compute the Woile under different noise levels σ, then sort each row in Wh.le, and add them together as Who Σ Eoso Wh.l.a resulting in the final weight for each head. An example can be found in Fig In F we present more results for different h and and we also present the heat map for the decode-only transformers in Figure 9
14

---

Jayer 3. 1.om
Layyer 4. o O
Hayer 5, comm 0
n. n_layer a
CXCOSS
J.byer
0_Jayer
lass
Iasea rdga
ndge
nage
EXCOSS
1.0
layer =. 0.2
layer 0.2
Bayyer 0.2
Dalayer 0.2
ENCOS5
byyer
Jayer
Jayyyr
Jayee
11000
8-0
byyer
3.0 0.8
Jayer 4 0.8
Jayer 3, 0.0
0.8
layer
0.8
1-2
20
examples
In-context
to 22
context examples
context examples
Figure 3: ICL with varying heads, layers and noise levels
num layar 5 avg
layer
1642
.umm 4, n.yyer
3.5
65
00
1th byerr
0.5
-1.
-220
h_num njlayer
-35
layer
62
5,
mum 8, layer avg
-64
-24
04
14th head
-01
1:th 10.00
Figure 4: Head Assessment with varying heads, layers
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 much larger weight than the others. This
15

---

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 Why 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 7 8 we provide the Pruning and Probing results for different numbers of heads 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.
{4.8] and noise levels
{0, 0.1, 0.2, 0.4, 0.8). can be found
K. 3.2
2.
NJwe-33--0 2-
and
Rad hel
Ford
premee 1
pranet
pranec
protech
pranet =
prancc
Itt
1th
Ith
Figure 5: Pruning and Probing, layers
pracche
pratch
pranch
praeei
Laye Eart
y
pracched
layer
Ihh kyyer
15h kyyer
1th Saper
Byer
14h layer
Figure 6: Pruning and Probing, 4 layers
16

---

22
FAn 1-44
Frash
pratch --44
Clee Fards
1 1a
1-th leyer
Hiht kayer
1ch layer
I-h Sayer
Hift layer
14h layer
Figure 7: Pruning and Probing, 5 layers
2.2-32
Nove
is
Franch
presech
5th
Figure 8: Pruning and Probing, layers
P-probing Here, we also set 117 and 11, with evaluation data size of 1024. We choose n > @ such that the model can handle more queries (q = 11) than those in the training (g 4) process.
C Proof for Section
C.1 Proof for Proposition 4.1
Proposition C.1 (Restate of Proposition 4.1 There exists transformer with 1 layers. h = d heads, dnid 3d and the input projection W E IR (d+ 1)x such that with the input sequence E set as Egquation 2.3the first attention layer can implement Algorithm 7 so that each of the enhanced data {Fixe,) hee(d) can be found in the output representation H
x1
X,
X,+1
X 0
O() TF, WE(E
1/1
y/2
0
Proof. Here first explain the key steps of our constructed transformer: the model first arrange the input entries with input projection to divide the input data into d subspace We. each subspace includes an entry x and the corresponding y (step C.2 then use h parameters (Wv, W W.. Wa. calculate queries, keys and values (step C.3 and compute the attention output for each head and concatenate them together (step C.4), finally projection matrix Wi
17

---

rearrange the result, resulting the target output (step C.5
X1
X2
Xn
Xx+1
Xn++
(C.1)
y1
1/2
V/
0
1,1
22
n,
(n+11),1
(n++),11
input projection
31
H
0
(C.2)
Wi
0
0
compune V.
1
K;
0
0
0
Qi, V
0
0
WV. W.. Wi; ERax
n
V/
0
X1,i
Xn.i
((n+1),5
(C.3)
X1,1
X2.1
Xn.1
X(n+1),1
X(n+g).11
Attn( W (E
y1
1/2
yn
0
X1,1
X2,1
Xn,1
X(n+1),1
((+tg)11
(C.4)
H+Concat/ MM,, Q.
X1
X2
X,
X +1
Xnto
H =T)
Vn
0
0
(C.5)
Wi this
The detailed parameters and calculation process for each step are as follows:
We E E (dd+1)xdha to rearrange the entries:
[[d] [[d+1] 0) T
WE = [[1]] [[d+1] 1] [2]
1|d
1]
where [k]] is daid vector with at i-th entry and 0 elsewhere, such that
X1,1
X2,1
Xn,1
X(n+1).1
((++)).1
1/1
1/2
Yn
0
0
0
0
0
= W
X1,2
X2,2
Xn.2
X(n+1),2
(n++g).22
we set Wv, WK., K. Wo. E 3x for values, keys and queries:
W K.
0
Wv. Wo.
0
1
1]
1 2]
such that the i-th head extract i-th entry and corresponding y
X1,1
X2,1 ...
Xn,1
X (n+1)).1
0
y1
1/2
On
0
K,
0
0
0
0
0
7
1
1]
X1,2
X2,2
Nn,2
X ((n+ 1).2
91
3n
0
X X1,1
X2,1
Xn,1
((n-1).1
5/1
1/2
Un
0
0
0
Q V
0
0
0
0
| 2]
X1,2
X2.2
X,,2
((n 1).2
X1,4
X n.
X(n+1),
:
18

---

0
V,MK Q:
0
(
0
0
X1,
Xp.á
X(n+1),5
X ((++)),
0
0
0
0
0
0
0
0
1/n
0
0
X1,
Xn,i
X(n+1),,
(()+)),
0
X1,1
Xn,
X(x+1),
((++)))
Then concatenate the output of each head {V,MK) Q Q together with residue:
X1,1
X2,
Xn.
x ((n+1),1
((x+g),1
y1
y/2
yn
H Concat((V +
X1,1
22,1
X
X ((+1))1
x(++)),1
(C.6)
Finally, W, is applied to rearrange the entries:
W1 = 133
[3d]
1[2]
... )..)
T
where the first ... implies the omitted 2 vectors {1[3i]]i 2,3..... (d 1)}. the second implies arbitrary values, then resulting the final output:
X1
X₂
X, Xn+1
Xx+y
H W, H (V, Q.)
Vn 0
0
in this way we construct a transformer that can apply Alg. so that each of the enhanced data (fxx, can be found in the output representation (1).
C.2 Proof for Proposition 4.2
Proposition C.2 (Restate Proposition 4.2p. There exists a transformer with layers, head, dnid 3d, let {(R,, ( be the l-th layer input data entry, then it holds that D(n+1) (w! ge xx+t). where Wed is defined as w = 0 and as follows for = -0...,, 1:
W f+1
e gd
vV(W (d), where ī(w)=
(ys x:))2 2n
1
1=1
Proof. Here we directly provide the parameters W We K for each layer TF..
W E RRddax dhis and W E Rhi this
W 7 O 1) W W n
[axd 0
0 0):
W Idwxxdaa Xdis
(C.7)
As we set W as the identity matrix, we can ignore and then apply Lemma in 3 By replacing (we' w:) as Q and we with P,, then holds that V(t41) (w Xn+1). where Wed is defined as w.O gd 0 and as follows for = O,.... I 1:
1 n
f+1 w ~vVi where ī(w)=
Σ (v. (w. xx:)²²
2n
i=1
19

---

D Proof of Theorem 5.1
To simplify the notations, we use W; to denote W We first prove that with a high probability, there exists a R E idddd such that RR = RR 1,, where 1, diag((1. (at..,.... with aj [(jEs).
Lemma D.I. Denote R (ri. , ra), where rj Ei= w Eig. Suppose n 2 Oloo (d/6)),
then for any (0, 1) with probability at least 1 6, we have
log (d/0)
|R R K.
where K Ca-CCmmmm (max;; Ei + )²),, where C is an absolute constant.
Lemma D.2. Define the event En by En {1A, 2 Hrd. vi E s}. Suppose that n 2 s log (d/5)/B², then P(E 2 5.
We define R by R diag(Tin (T,.... Tal, where 7y is given by
0 1 S,
1/5 jES.
is easy to see RŘ = RR I,. On the event &,, we have that condition on E.
R - 1/3. Hereafter, we
D.1 Bias-variance Decomposition
Let x Xй with x; Rx,, For Wt, we have
1
W++1 Rw We -Rw" 17.
[ y1)
i=1
1
w. - -Rw" -17
X1
X
W. - x, Rw +c)
i=1
(I ()- ( (w: - Rw*) + + XTc.
n2
Hence, we have
115) (1)))-11-))) Rw* +
W = (1-
Σ
Σ
t
II -
i-1 TTc. 6
(D.I)
11
We decompose the risk I(W:) by
E(W:) = ~x (Rx, w.) - (Rx, Rw") -e) )
- -02
(D.2)
=x ( (Rxx, w.) - (Rx, Rw 27
)
/ (w (iww-=-w") 1/2
2
E2/2R (1(--1z) Rw +
1
Σ
I
X c
i=1
2
E - n
Rw*
+n²
Σ R
-
X
(D.3)
.
2
Bias
Variance
Next, we present some lemmas.
20

---

Lemma D.3 (Theorem 9 in Bartlett et al. 19). E (0, 1) with probability at least 6.
There is an absolute constant such that for any
r(2)
r(z)
log (1/6)
log (1/6)
Σ - E
S
E max
7
71
72
where r(Z)
TT(E)/A1-
Lemma D.4. With probability at least 6, we have
((RER) Vr(E) + (RERR)
r(E)
RER RER E Vs. poly (d/5))
n2/22
+ As a result, when 2 st ( (RER)) poly(log (d/6)), with probability at least 1 - 6,
we have
RER RER|2 S 1/t.
We define the event &g follows:
82 82== RER Sa(,, (0 111).
where
ō(n,5) = V poly(log (d/5)) (V
V 12
Vr(E) + +RERE
r(E)
31222
n
By Lemma D.4 P ((&)) 1- 5. Hereafter, condition ε n 2
D.2 Bounding the Bias
On & n &g, have
2
Bias
1/2 n1 Rw*
22
ww' R n RER I --1 Rw'
= = n R E - Σ) R ([I I --)) Rw + w TH(RII v RSR I - n/Ei "Rw"
(D.4)
Lemma D.5. On & n E2 we have
1 1,32
1
and
II < 1 nt
hold with probability least 5.
By Lemma D.5 we obtain that with probability at least 1 - 8,
Bias 11+11 11
+
S
(D.5)
1.2
TT
n1.322
where the last inequality is by 1/ L
21

---

D.3 Bounding the Variance
8
i-1 XTe
2
Variance- = o/?
EV/2
Σ
[ n
71
XR
[-
RER
>
(--
RX
n²
i=1
1=1
112
XR
[I 1
R
Σ - E R
(1-
RX
n
i=1
==1
+
112
XR
Σ
(I- n1
1-1
RER
I 1ー
i-1
RX c.
(D.6)
n
Lemma D.6. On & O E2, with probability at least
8, we have
net 1
RX
n²
and
n log n²
?
II
Rx T
2
By applying Lemma D.6 to Eq. (.6).
we obtain that
Variance = I <
2 n²
RX
+
xyl yllog! L
R
1/2
S
nlllgg n²
Rx
9/2 (D.7)
2
n²
Lemma D.7. with probability at least 1 - 8, we have
Tr (EER) log (d/o)
o'sTIr E) log² (d/o)
RX
n
n12
By applying Lemma
D.7 to Eq. D.7..
obtain that
TTTT(RER) log (d/8)
o3sTr(E) log" (d/o)
Variance S 1yt logt-
(D.8)
n
2
D.4 Final Bound
Combining D.5 and Eq D.8
we obtain that
1
[ (TTRR(ER) log (d/8)
o'sTr(E) (eTT(())ogg log (d/i)
E(W.)
nttt
nt52
n
n²
log t
(RERR log (d/(0)) (2sTr(E) log² (d/5)
n
n²
when 1ft (RER) log
of Tr)) (d/66) - 1/2
n
n2
D.5 Proof for Appendix D
Proof of Lemma D.I
Since yi = Ej= wHas + css then we have
n
Ft Σ wtask + t.
d
w
2
+
1 n
Σ
Jztl- (D.9)
12
n 1=1
n
n j=1
j=1
22

---

Since 2js N(0, Ea) for any i.j, by Lemma 2.7.7 in Vershynin [46], there exists an absolute constant C' such that 2y4 is a sub-exponential random variable with
| CV S K,
where || llw, denotes the sub sub-exponential norm and the last inequality comes from the definition of K. By applying Bernstein's 's inequality [46, Theorem 2.8. I1, we have
1 71
Job
10 MILLTI
j=1
log (d/o) log (d/6) n
$ K max
n
K log (d/6)
(D.10)
where the last equality due to 2 Ollog (d/5)). We also note that Tyity is sub-exponential random variable with | I K. Hence, we also have
log (d/ö)
IjcCj
SK
(D.11)
7
==1
Combining Eq.DD. Eq. (D100)
and Eq. Sq.ID.II. we have
log (d/5) lotil + K.
log (d/5)
log (d/8) n
ral S K
(1lw*l1 111 1)K
n
n
By definition of R and R, we obtain
log (d/6)
R Rl12 = max ir - rrsl K(Il 111 + 1).
n
(Vsssl1++)++) log (d/6)
s log (d/5)
S K
which completes the proof.
Proof of Lemma D.2 By Lemma D.1 for any E S, with probability at least 1 - ō, we have
log (d/8)
Ir - Fyl S
$ $1/2 < 1ry1/2,
(D.
n
where the last inequality is due to the definition of B.
Proof of Lemma D.4
We can decompose
|RER
RER
|l2 as follows:
- = R.R
RER
+ R ER. RER RER RER 1
RER
RER
|ER
RER
12 |EE
RER (D.13)
Next, we proof the bound for T, II and IIT separately. For term 1,
-= RŘR RER||1 11 ( R E
< Î R|12 |ا1ز - ||R|||2
< R R|22 (1Σ) + -11 - >112) (
(IR|/2 + |R (R(l)
(D.14)
23

---

where the last line is due to triangle inequality. By Lemma D.3 with probability at least 6/3, we have
r(E)
log (1/6)
log (1/6) n
Σ 2
max
n
r(E)
log (1/8) r(E) log (1/5)
max
(D.15)
n
71
By Lemma
D.I
we obtain that
log (d/5)
| - R112 S K
21
(D.16)
holds with probability at least 1 1-6 6//3 where the last inequality is valid since n 2 K2s log (d/6).
Combing (D.44. Eq. (9.1 and Eq. (D.16), we have
log (d/8)
r(E) + log (1/8)
r(E) + log (1/6)
I S K| Eا1
+ max
n
n
log (d/6)
r(E) log (1/8)
r(E) log (1/8)
$ K >
1
(D.17)
n
n
For term 11, can decompose 11 as follows:
R
R
l2
R
E
E
R|12+
+
R
$ E
( R R))12.
TT.a
11.b
For term II.,, by using Lemma D.3
we have with probability at least 1 - 8/3,
(RER)
r(RER)
log (1/6) log (1/5) n
IL. 2 |RER||2 max
11
11
Y1
(RER) log (1/8) r(RER) log (1/8)
< |ERR| max
n
(RER) + log (1/0)
(RER) + log (1/8)
< <|ER||2
(D.18)
Similar to the proof for bounding I, we can obtain that
(1 1+
log (d/o).
( + log (1/8)
r(E) + log (1/8)
IL.b S K|EE|2\
(D.19)
n
n
For term ITT, we have
log (d/0)
II RE R R 1l2 S
[|/||| [112k([[ 1))
(D.20)
n
where the last inequality is by D.16.
24

---

Combining Eq. (.17). Eq. (.18).
and Eq. D.20 and taking the union bound, we obtain
that with probability at least 6,
N RERR| 2 I++ II + III
< K| (2(1 + DD log (d/ō)
r(z) + log (1/6) 1+
r(E) + log (1/6)
n
( + log (1/6)
RERR + log (1/6)
+ |RER|2
log (d/5)
+ |R|22 2 w 11a + 1).
- (1/12
11 + 1) + |RER| 12 + ||| 2 (1) w/ 11 + 1))
log (d/o)
rr{E) + log (1/6)
r(E) + log (1/6)
2
(RERR + log (1/8)
(RER) + log (1/5)
r (((ER + log (1/6)
(2) log (d/0) + (REER) + log(d/8)
$ Coor
n
+ r( (Ev log (d/5) 3/2 (d/5)
n2/2
((RER)
Vr(E) + +R(ER
r(E)
$ &Cov plog (d/8))
n
n
n2/22
where the second last inequality is by an' +bb +oo S (a + b+ c) + for b.b.
0. Here Ceov = K| [1]]([[
|| + 1) + |RER|| + |R /|
V8.
0
Proof of Lemma
D.5
By the triangle inequality, we have
Ñ Σ 2) [R
R Σ Σ R R
E E) iR (Ř-R) R) (f. R) (I- 15--pe
R -E)
î R
Σ - 8) (n ((n-R) R) |
< R {
+
|
5) 5) R (RR-RR], | - ( 2
| ( R) (- {)
+
Following the proof of Lemma D.4 we can prove that with probability least 1- 8.
R [
Σ - 5) R
a( < 1/t,
(D.21)
where the last inequality is by &g. By Eq. (D21.
we have
R -
E R = 1/t I.
Hence, we obtain that
IS w T R I (2II-15)) -1EE
1/1. I. (1- 'R''
Rw-
w R I- (1)EE) 21 Rw*
1 w" 'RRw"
(by (I (1- ID
||
(D.22)
25

---

where the last line by R
R 3 I. For the term II, we have
II W R(I T n RER (I-nE) Rw
S 1 w' 'RRw"
11l
71,32 l1
1
112 -
1
(D.23)
nt,332
where the second last line is by the fact that x01 - x))* S 1/(k + 1) for all x E [o, 1] and all
k > 0.
Proof of Lemma D.6
Similar to the proof of Lemma D.5 with probability at least 8, we have
R R(I - 1 5) R
3
I. Then we have
XR
I
1123
1-1
R
Σ
R
I - n))))
i-1
Rx
n
i=1
112
XR
I 1
i-1 Σ
I
(1-n
1-1
RX E
In2
1=1
n²t
XRRX
n²
E
n2t
2
n²
|
T
e
'2'
where the second last line is by Ei=1 (I-15)) 1-1
31.L I. By the fact that x01 - x))* S 11/(k + 1)
for all x E [0, 1] and all k O, we have
2
XR
I (1-)E)
1-1
R.R.
I (--)E)
i-1 RXT,
n2
1=1
n
E XR
Σ
1 11
itj-2 nR
Rx
n1
h-=1
S 1
1
)
|X
6
112
+ j
2
1j=1
n
2
Rx
n2
2
1=1
ntt log L
R
T
n
where the last inequality is by the fact that Ei=
1 S logt t.
2
Proof of Lemma D.7
First, we can decompose
- Rx T
c/2 by
1 n
R
RX
S
RX
+
(R-R)
R x e
Let z, Rx, then z; N(G). where G == RER. For any i, by Lemma 2.7.7 in Vershynin 46). there exists an absolute constant C such that Ezz is sub-exponential random variable with
ر, S Cov/Ga
26

---

By applying Bernstein's 's inequality Vershynin 46
Theorem 2.8.1]. for any 1 S i S d, we have that
J E1c1
S V Gi max V log (d/6) log (d/o)
log (d/6)
-OVGH
(D.24)
hold with probability 1 union bound, we obtain that
where the last inequality is due to n 2 (llo(d/))), By taking the
32
1
log (d/6)
S
SovGu.
n
holds for any i, with probability
8 3
Then we have
log(d/6)
(TT((R log(d/8)
1= Σ
EjZ
G
t=1
1=1
In the same way, we can prove that with probability at least 8/3.
TTr (d/))
(D.25)
By applying Lemma D.1 with probability least 1 8/3, we have [ R ||= 2
8 log (d/6)
(D.26)
n
By (.25) and Eq. D.26). with probability 1 - 2. we have
1 (R-R)
s R - x
o2sTr(E)) log" (d/o)
n2
By taking the union bound, we derive the desired result.
E
Proof for Theorem 5.2
To simplify the notations, we w, denote w w'dd Lemma E.1. with probability at least 1 - 6, we have
Σ - Σ S 8),
(E.1)
where an,, 8) -
1r(E)-log(1 (1/8)/8)
Tr. () +og (1/6), As resuls, when (Tr(E) + log (1/6)).
with probability at least
S Σ
1/t.
Proof of Lemma
E. By Lemma D.3
we have
r r(E)
r T(E)
log (1/8) log (1/8)
Σ Σ
E| max
r(E) + log (1/8)
r r(E) + log (1/8)
E max
r r(E) + log (1/6)
Σ + log (1/5)
(E.2)
holds with probability least 1 5, where the last line is by the inequality that max {a, b) < a + for all a, 6 > 0.
27

---

We define the event & as follows:
= RERII2 2 a(n, R 1/1
By Lemma E.1 P(E) 2 1 -. 8. Hereafter, we condition on &
Bias-variance Decomposition Similar to Eq D.I. we have
W: I [I n )
w'
I (-
X
(E.3)
In the same way, we can decompose the risk E(w,) by
E(w,) = 1// (I (I- و w
11
1/2
i-1
T
(E.4)
i=1
Variance
Bounding the Bias
Bias = w+ (1 (I- Σ II II(-E wv
=w [I- -w E Σ) (I- w +w t
TI-
m
D
n
n
w'
Similar to the proof of Lemma Lemma E.2. On E, we huve
D.5
we have the following lemma.
1 1 1 1
and
S Ttt
hold with probability at least 6.
As result, the bound of the bias term is given by
1 1 1
Bias <
(E.5)
1tt
t
T
Bounding the Variance By using the same way of the proof for bounding the variance term of Theorem 5. we have the following lemma.
Lemma E.3. On E, with probability at least 5, we have that
T2Tr(E) log (d/8)
Variance S ntt log
1xxT
XT c
S nyt logt.
(E.6)
n
Combining (..) and Eq. E..), obtain that
myt logt TTr log (d/o)
a TTr(( log (d/8)
E(w,) S nt
< logt
n
when 1p6 ~ (o"Tv(E) log (d/o)
-1/2
Proof for Theorem 5.3
To simplify the notation, W, to denote w and to denote w gd.
28

---

F1 Proof for the upper bound of the excess risk
When Σ = I, by Eq. D.2). we have
E(w,) = R
12
Rw
+
t1²
R
I
m
i-1
X
1=1
Bias
Following the proof of Theorem 5. it holds that
Variance I ny logt a² Plg (d/5) +
o²
dd (d/5)
n
n²
with probability at least
6, when n 2 s2
Similar to the proof of Lemma D.2 we can prove that
P
Vi S,
Pi P Vi,
with probability at least - 8.
When E I, by Lemma D.4 have that
8a S t
RER RER
holds with probability at least 1 - 6, when n 2 17/2/19 As result, RER Hereafter, we condition on the above events. For the bias term, we have
- 7 - RER.
R
| (I--1 1 Rw"
(1- n Rw
T
21 Rw*
w T
R
I 17
RER
82
t {.1)) Rw
5²2 t ))
Σ (w) 1P:)
E5
(1 11,32 1 /2) 20
where the last line is by the definition of B. When t 2 log a² ne
/ (2 log (1 - (1-182/2)) n/3312)) we have
o2
Bias R
AR(I-n13 1 n)
Rw
(F.1)
n
When n/32/2 1/2, there exist a c 0, such that
log (1 - (1-182/1) 2 c7132/22 /2.
Hence, the variance term is bounded by
log (d/o) o²
o²||w* 1dd (d/8) (
Variance S Tyl log
n
n²
σ Looog 2 (ns/o2)
log? (d/o)
ds
(F2)
82
n2
where the last line is by || || 5 s. /| = and nt < log (ns/d²)
Combining Eq. F.I and
82
Eq. F22, have that
E(w,) S o² +
log? (ns/o²) log2 (d/o)
ds
oog?? (ns/o²) log? (d/6)
(1
ds
+
n
82
n
n²
82
n
n²
29

---

when n 2 Lade/s
B 2
fllw"
ga and
log (n=) na
When m E U(-11//6,1/33) 8 = 1/v5. In
this case, we have that
(() S lolog" (ns/o²) log" (d/o). (s + ds²
n?
when n 2 d2/3 and
log (ns)
F.2 Lower bound for Ridge Regression
When n 2 d + log (1/6), by Lemma D.3 we have that 1/2 I Σ 2- For the ridge estimator
WA = n (S+A.I) X y, we have
Ew [E(w)]]
I- S 1 MI) 1
w
1 1. (É+
($+A-1) x
2
2
2
( ( + x
n
By Lemma
D.3 when 1 1 E 2 I, with probability least 1 8, we have
- (S (S+A-1) -1 x
E (g(wx)) 2
n² e x [ +1) -2
X
1
c TXX",
n2 + x)2
where the last line is due to the fact that Σ + AI 3 (2+ A) I.
Lemma F.I. Given X such that i
Σ
< 21, il holds that
odd
x
with probability at least 1 when n 2 Olog (1/6)).
Proof of Lemma F.I We consider the singular value decomposition of J. L 1 UAVT. vn
where U Raxd is orthogonal matrix, A Rdxa is rectangular diagonal matrix with non- negative real numbers on the diagonal, V e R=xn is an orthogonal matrix. Let {oI...., Ta] be the singular values of X Then we have
x
UAV
AV
n
Vn
2
72
d
Aa
0
1=1
where ē V N(O, 1). By ][Lemma 22], we have
Ei of log (1/0)
max; of log (1/0)
X
E
X
0² max
dlog (1/8) log (1/0)
0² max
(F3)
30

---

where the last line is valid since (oi...,, og} is the eigenvalues of Σ 1 By F.3... we obtain that
X and I = 21.
Vdlog (1/0) log (1/6)
X
E
X
0² max
71
2
"
d
v dog (1/0) log (1/8)
0² max
0
71
1=1
v dlog (1/6)
log (1/6)
(by 1 3 21)
max
71
72
o n'
where the last line is due d (log (1/8)).
Next, we define the event & as follows:
1 I 21,
1
2
o²d n
Erigge
-x
n
By Lemma F.1 we have PPE) 2 1 8 when 2 O(d) 2 Oloo (1/6)). On Eridge, we have
o²d
E [E(w>)] 2 (1+ 1)
(F.4)
When d Z n + log (1/8), by Lemma D.3 with probability at least 1 - 8, we have that I R 2d I. Hereafter, we condition on this event. By direct calculation, we can decompose the excess risk by
3 XXT
E [E(w)) = Bw- || (I- (1-($+))
S AI)
Σ
w
+ n
1
( (S+A-1)
"x
x E 2
2
For the first term, we have
Ew ( I - (E+ + 21)
-1
E w'
2
Ew
I --x
x
(XX nnI) x) x) w w/11
(1
Eww ||w* 12
(F.5)
d
1 d
(F.6)
where the last line is due to X- (XX
+ +ndI) -xx
a n space.
| (Σ+A-1)
2
=e =xx (XXT + +nAl) -2
x
dn
2
(F.)
2(2d + nA)²
where the first line is by (xTx X + ndI) X
x (XX + nAI) and the last line is by
2(d+nA)² I XX (XXT + nAI)
-2 By Tsigler and Bartlett
44
Lemma 22], we obatain that
Σ & no?
S 0² v (1/6) + o²
holds with probability at least 1 8. When n 2 log (1/8), we have |E 6
- 2
ng² 2
holds
with probability at least 8. Taking the union bound, we obtain that
Ew- [[())]] 2 d + dn
2(2d + nA)a 2 1
n
d + +
(F.8
21+
31

---

F33 Lower Bound for Finite-Step GD
We first consider the case where n 2 d + log (1/6). Define the event EgD by ECD =
{1
{ 1 E 21). By Lemma D.3 P(EGD) 2 1 11.. 6. By Sq.EE.4, we have
2
E [E(w,)] = Ew
(T-n)_)
w
+12
I 1) X c
2
1=1
1
i-1
2
I 1/
XTc
i=1
2
-1
172
3 -- (I- (T- /155) (1)))
X
n
7/2
1
-1 x
1)
T
n2
Tyl.
c
2
r7ad
(1 1/ 1n(
where the second last line is by E 8(1 811- (1- n (1-(1_))))
3
2 by
and the last line is by Eq. F.4..
We then consider the case where dl 2 n log (1/8). Define the event EoD 1
I 2dI
By Lemma
D.3 P(E/D))
2 1 - 6. Following the proof of Zou et al. [58] Theorem 4.3). we have
Ew- (E(w)] 2 Ew
I I-X
XX
71 7tt
I
X
+
X
XXT
+
n
I
C
n
nyt
2
02n
n
= 1
+
d
(1 +
1
2
d
where we use the results from Appendix F2
F,4 Lower bound of OLS
Let Wols be the OLS estimator. It is easy to see Wols
Wo. Hence, we have
o'd
n z d - log (1/8) d 2 n - log (1/5),
E [E( [[Wo)))]] 2
dnn
holds with probability at least 1- 6.
G Additional Experiments
Here, we provide additional experiments on the decoder-only architecture and train models with different settings.
Training Oncoder-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. (.II (where in ys+ we use zero padding to align with x,), dnis = 256. We choose = and E {4. 5, 6} We then conduct heads assessment experiments on the trained decoder-only transformers with 10 in-context
'e 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. possible reason is that decoder-only transformers first need 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.
32

---

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.
E
X1
/1
X2
1/2
Xn
yn
L =
y1
ys)
(G.1)
k=
num B. Jayer avy
num 5. n. Jayer avg
n. byyr
Figure 9: Heads Assessment for decoder-only transformers
Training Models with d 16 Here, we adapt the encoder-only transformer and the same settings introduced in B but set = 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 4, d 16, transformers with more heads (h 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 is close to d, clear better upper bound guarantee, as ensured in cases where 8 & may not hold.
Training Models with non-orthogonal design To further demonstrate the applicability of experimental results to more general non-orthogonal settings, we conducted additional experiments by modifying the distribution x to N(O, Σ), where Σ I (S, and a matrix of ones. We varied across the values [0, 0.1, 0.2, 0.4] validate our findings. The results are presented in Figure 11] which reveals patterns similar to those observed in orthogonal design settings.
33

---

layer
layg
Tayer 5.0
layer
layer
Jayer 0.1
Jayer
layer
61
MAX
context examples
contoct exampies
In n.contaxt atamples
In co.conext examples
Figure 10: Train Models with d 16
-1
2-0.2
Ce0.0. me
(+t.A me
Figure 11: Train Models with x ~ N(O, Σ), where Σ = I+ (
34

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect 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 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 include aspirational goals as motivation long 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 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 encouraged create 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 few datasets with 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, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or 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, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
35

---

judgment and recognize that individual actions in favor of transparency play an impor- tant 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 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 cross- referenced.
All assumptions should be clearly stated 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 short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
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, 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 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 a specific model and empirical evaluation, may be necessary either make it possible for others 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 replicate the results, access a hosted model (e.g.. in the case of 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
36

---

(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is new model (e-g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d)
We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, it may be that access to the model limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions 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 NeurlPS code and data submission guidelines (https://nips.co cc/ public/guides /wtoe//uuiissssioooooooooooo onloiicy))o for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
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. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Detalls
Question:
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, 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
Guidelines:
37

---

The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary 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
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: The reported mean behaviors in our experiments are sufficient 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, confi- dence intervals, statistical significance tests, 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 should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars reported in tables 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
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide experimental settings and training details in Appendix B
Guidelines:
The answer NA means that the paper does 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 well estimate the total compute.
38

---

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http://pws..aoooo...../oo//////a
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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal 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 no societal impact of the work performed.
authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., 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 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 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 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 (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
39

---

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 (c.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 such risks.
Released models that have 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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators 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?
Question:
Answer: [Yes]
Justification: the creators 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 used and, possible, include 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 (c.g.. website). the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pppppsaiithhoeee thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
40

---

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 the dataset/code/mmodel as part 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 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 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 research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According 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 the subjects, and whether Institutional Review Board (IRB) approvals (or 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 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 NeurlPS 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.
41


Reviews:
[{"review_id": "KTQPBvn9KT", "reviewer": "Reviewer_fd5H", "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\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.", "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": "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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strengths:\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.", "weaknesses": "Weaknesses:\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the time and effort you spent on thoughtful reviews. We address your comments below:\n\n> **Q1**: *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?*\n\nThank you for your insightful questions about the feed-forward layers. In transformers, we can observe that (ignoring the layer norm) the attention mechanism is the only module capable of linking tokens across different positions, while the FFN layer performs a token-wise (nonlinear) mapping without considering contextual information. For learning from context, especially in our linear sparse regression settings where contextual information is crucial for transformers to identify appropriate subspaces for sparse linear regression problems, additional (nonlinear) mapping for each token is less important. In our theoretical analysis, we primarily focus on the attention-only transformer without the FFN layer, focusing on the role of multi-head attention, such simplification is widely adopted in theoretical analyses, as seen in [1], [2], and [3]. To demonstrate the limited role of the FFN layer in our specific setting, here we provide comparative results of the linear attention model with and without MLP layers in the attached PDF.\n\nWhile the FFN layer's role is limited in our setting, we acknowledge its potential beneficial impact when combined with the attention layer for more complex nonlinear or real-world tasks. Transformers may utilize similar preprocessing-then-optimize algorithms across different layers. Although this is somewhat beyond the scope of our paper, insights from other works provide valuable perspectives. For instance, [4] demonstrates that transformers tend to first learn appropriate representations using MLP layers before learning from context through the attention layer. Additionally, [5] and [6] use experiments to show that parameters in deeper layers are less critical compared to those in shallower layers. We believe our findings are both interesting and warrant further investigation in future works.\n\n\\[1]Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., & Vladymyrov, M. (2023, July). Transformers learn in-context by gradient descent. ICML2023\n\n\\[2]Ahn, K., Cheng, X., Song, M., Yun, C., Jadbabaie, A., & Sra, S. (2023). Linear attention is (maybe) all you need (to understand transformer optimization). ICLR2024\n\n\\[3]Ahn, K., Cheng, X., Daneshmand, H., & Sra, S. (2024). Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023 \n\n\\[4]Guo, T., Hu, W., Mei, S., Wang, H., Xiong, C., Savarese, S., & Bai, Y. (2023). How do transformers learn in-context beyond simple functions? a case study on learning with representations. ICLR2024\n\n\\[5]Gromov, A., Tirumala, K., Shapourian, H., Glorioso, P., & Roberts, D. A. (2024). The unreasonable ineffectiveness of the deeper layers. arXiv preprint arXiv:2403.17887.\n\n\\[6]Men, X., Xu, M., Zhang, Q., Wang, B., Lin, H., Lu, Y., ... & Chen, W. (2024). Shortgpt: Layers in large language models are more redundant than you expect. arXiv preprint arXiv:2403.03853.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VW4eCD7Az4", "reviewer": "Reviewer_uBeS", "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.", "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": "The 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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. 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.", "weaknesses": "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$.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the time and effort you spent on thoughtful reviews and comments. We address your comments below:\n\n> **Q1**: *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$.*\n\nThank you for your insightful suggestions. In our preprocessing algorithm, the transformer employs multiple attention heads ($h>1$ ) to process each subspace differently, thereby enhancing the optimization steps for subsequent layers in our sparse linear regression task. However, a single-head transformer ($h=1$ ) can only process the entire space uniformly, lacking the ability to differentiate between subspaces. Consequently, single-head transformers have very limited \"preprocessing ability,\" so we believe using $h=1$ serves as a reasonable \"no preprocessing\" comparison.\n\nWe appreciate your suggestion to use hidden states without attention module for preprocessing as another point of comparison. We have provided additional results in the attached PDF.\n\n> **Q2**:*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\".*\n\nSorry for causing the misunderstanding.  We would like to clarify that the comparison between Theorem 5.1 and 5.2 are made in terms of two “mechanisms” (corresponding to multi-head transformer and single-head transformer), which we believe is fair and the model parameters do not need to be included. In particular, based on Propositions 4.1 and 4.2, we conjecture that the working mechanism of multi-head transformer is “preprocess-then-optimize”, i.e., performing GD on the preprocessed data, while the working mechanism of the single-head transformer is conjectured to perform the “purely-optimize mechanism” on the original data [1]. When applying our results to $L$ layers transformers, $t$-step gd with preprocessed data ($L = t+1$) truly requires just one more layer than  $t$-step gd ($L = t$), However, factor $L$ only appears in the log term of our bound (Theorem 5.1&5.2), hence it doesn’t affect our results (which are interpreted in orders). \n\nTo this end, it is fair to compare GD with preprocessed data and GD to demonstrate the superiority of our preprocess-then-optimize mechanism. We will make this clear in the revised version.\n\n[1]Ahn, K., Cheng, X., Daneshmand, H., & Sra, S. (2024).Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2024\n\n> **Q3**: *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\nThank you for your interest. We believe a similar phenomenon occurs in other settings of ICL tasks as well, and transformers may utilize similar preprocessing-then-optimize algorithms across different layers. While this is somewhat beyond the scope of our paper, we can gain insights from experimental results of other works. For instance, recent studies by [2] and [3] demonstrate through experiments that parameters in deeper layers are less critical compared to those in shallower layers. Additionally, [4] shows that multi-head attention sublayers exhibit low-rank structure in large language models. We believe our findings are not limited to simple sparse linear regression settings but can be extended to more complex ICL and other real-world tasks, and such phenomenon is worth further investigation in future works.\n\n[2]Gromov, A., Tirumala, K., Shapourian, P., & Roberts, D.A.The unreasonable ineffectiveness of the deeper layers.arXiv preprint arXiv:2403.17887.\n\n[3]Men, X., Xu, M., Zhang, Q., ... & Chen, W.Shortgpt: Layers in large language models are more redundant than you expect.arXiv preprint arXiv:2403.03853.\n\n[4] Li, G., Tang, Y., & Zhang, W. LoRAP:Transformer Sub-Layers Deserve Differentiated Structured Compression for Large Language Models.ICML2024\n\n> **Q4** *The orthogonal design is an over-simplified setting, which the authors do not properly address.*\n> > **Q4.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.*\n> > **Q4.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?*\n\n[**A4.1**]: Thank you for your suggestions. Our experimental and theoretical results can indeed be extended to non-orthogonal designs. To validate this, we conducted additional experiments by modifying the distribution of x to x∼N(0,Σ), where Σ=I+ζS, and S is a matrix filled with 1. We varied ζ across [0,0.1,0.2,0.4] to further verify our findings. The results, which are consistent with those presented in Sections 3 and 6, can be found in the attached PDF. We will incorporate a more detailed discussion of these non-orthogonal design experiments in our revised manuscript.\n\n[**A4.2**]: While it's true that there exist specially designed transformers capable of solving the lasso problem using a closed-form solution, we find that the trained transformers are more likely to use other algorithms. Here we compared the performance of linear attention transformers with and without MLP layers in the attached PDF, which shows that the inclusion of MLP layers does not significantly impact the results for this particular problem. Although different initializations may affect the results, we believe these results are sufficient to demonstrate that our choice to use a simplified attention-only transformer for theoretical analysis is reasonable and can provide valuable insights for real-world models.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "THolkkuKDD", "reviewer": "Reviewer_Swt3", "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.", "questions": "Line 108: should this instead read \"to perform a tractable theoretical investigation\"?\nLine 209: typo \"proprocessing\" instead of \"preprocessing\"", "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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "While 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.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments! We sincerely appreciate the time and effort you've dedicated to providing thoughtful reviews. We've addressed your concerns as follows:\n\n> 1. *The theoretical analysis is based on a simplified transformer maybe hard to apply the results to more practical real-world problems*.\n\nWhile our analysis uses a simplified transformer model, we believe this widely adopted simplification captures the essence of the model and provides valuable insights applicable to real-world transformers. In the attached pdf, we use additional experiments to demonstrate that MLP layers are less critical in our sparse linear regression setting, and our theoretical and experimental results can be extended to various data distributions. Moreover, based on the recent studies about the parameter distribution of real word transformers, our results are promising to be extended to more practical settings, for example [1] and [2] demonstrate through experiments that parameters in deeper layers are less critical compared to those in shallower layers, additionally, [3] shows that multi-head attention sublayers exhibit low-rank structure in large language models. We believe our findings are not limited to simple sparse linear regression settings but can be extended to more complex ICL and other real-world tasks, and this phenomenon is worth further investigation in future work.\n\n\\[1]Gromov, A., Tirumala, K., Shapourian, H., Glorioso, P., & Roberts, D. A. (2024). The unreasonable ineffectiveness of the deeper layers. arXiv preprint arXiv:2403.17887.\n\n\\[2]Men, X., Xu, M., Zhang, Q., Wang, B., Lin, H., Lu, Y., ... & Chen, W. (2024). Shortgpt: Layers in large language models are more redundant than you expect. arXiv preprint arXiv:2403.03853.\n\n\\[3] Li, G., Tang, Y., & Zhang, W. (2024). LoRAP: Transformer Sub-Layers Deserve Differentiated Structured Compression for Large Language Models. ICML2024\n\n> 2. *Line 108: should this instead read \"to perform a tractable theoretical investigation\"? Line 209: typo \"proprocessing\" instead of \"preprocessing\".*\n\nThank you for pointing out our typos and sorry for any confusion they may have caused, the \"intractable\" in Line 108 should be \"tractable\" and \"proprocessing\" should be \"preprocessing\", we will correct these typos in our revised version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "f8v1HODt24", "reviewer": "Reviewer_Bazq", "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.", "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": "NA", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper is easy to understand and provides interesting theoretical results.", "weaknesses": "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. \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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the time and effort you've invested in reviewing our work. We've addressed your questions and concerns as follows:\n\n> 1. *The difference between our work and other works towards aligning multi-layer transformers with gradient descent, like [3,4,5].*\n\nThank you for your insightful reference paper, we will add more discussion about them in our revised version. We want to highlight that our work is not just focused on the expressive power of transformers (such as their ability to implement novel optimization algorithms). Instead, we aim to answer the question: \"How do trained transformers tend to use multi-head attention for in-context learning?\" We first gain insights from trained transformers (in Section 3) and then propose an algorithm to explain our observations (Sections 4 and 5). While [4,5] focus on single-head attention and [3] explores the expressive power of single-layer transformers, our work seeks to understand the hidden mechanisms of trained multi-layer transformers through a novel perspective. We believe our findings are both interesting and extensible to more real-world tasks and complex models, warranting further investigation in future work.\n\n> 2. *Theoretical intuition for why multi-head attention works in the way observed in this paper, and the linking between softmax attention and linear attention*.\n\nAn empirical intuition for the benefit of multi-head attention in sparse linear regression is that it enables the transformer to process different subspaces of the input data using separate attention heads, as discussed in Section 4.1. Here we can provide a more theoretical explanation for the working mechanism of multi-head attention in sparse linear regression, building on [2]. Consider I = d/k, representing I different tasks, each focusing on a sparse linear regression problem with k assigned dimensions, and $g_i$ is the same for different tasks i (i.e. the task can be seen as a combination of I sparse linear regression problem). Based on the results for ms-attn in [2] (eq 4.3 & theorem 4.2), the optimal solution involves each attention head processing a specific subspace while zeroing out other dimensions, aligning with our Proposition 4.1.\n\nAssuming sufficiently large $d_i, d$ and $L$ (assumptions 2.1 and 2.2 in [2]), the parameter for $w_{qk}^h$ is scaled by $1/\\sqrt{d_i}$, allowing softmax attention to converge to linear attention:$x̂_{ij} \\approx\\frac{\\sqrt{d_i}}{N} \\sum_{k = 1}^{n} (\\frac{\\langle x_{kj} ,x_{ij}\\rangle}{\\sqrt{d_i}} + 1) y_{k}=\\langle\\frac1N \\sum{x_{kj} y_k}, x_{ij}\\rangle + \\frac{\\sqrt{d_i}}N \\sum{y_k} $\n\nHere, we omit the term $\\frac1{1 + e d_i \\phi_i L^{-1} }$  for $u_i$ (eq4.3 in [2]), as it's constant across i in our setting. The left term $\\frac1N \\sum{x_{kj} y_k}$ can be interpreted as our preprocess coefficient $\\hat{r}_i$ (defined in eq4.1 in our paper), while $\\frac{\\sqrt{d_i}}N \\sum{y_k}$ remains constant across different tasks i.\n\nWhile our results for the first layer's multi-head attention resemble those in [2], we emphasize two key differences: 1) we elucidate the role of multi-head attention in multi-layer transformers through a new preprocessing perspective, and 2) we provide both theoretical and empirical explanations for the advantages of multi-head attention, demonstrating a potential integration mechanism between the first layer and subsequent layers. Although the convergence behavior of deeper layers remains an open question, our theoretical intuition above suggests a potential approach: first assuming that subsequent layers optimize data through gradient descent, then prove that the first layer converges to utilizing multi-head attention for data preprocessing, we believe this warrants further investigation in future research.\n\n\n> 3. The role of MLP in our setting.\n\nIn our linear setting, learning from context with linear attention is widely adopted in theoretical analyses. While linear attention incorporating MLP layers can implement complex algorithms, it's challenging for a trained transformer to adopt the exact solution we've constructed. In the attached PDF, we compare the performance of linear attention-only transformers with and without MLP layers, demonstrating that the difference between these models is not significant, thus justifying our simplification.\n\n> 4. Challenges in our analysis on the linear model\n\nAlthough we are analyzing a linear model, when we analyze the convergence of Preprocessing+GD, we encounter some non-trivial challenges. We point out two main challenges encountered in analyzing this problem as follows: 1) In general linear models, data are assumed to be independent, but in our setting, the data preprocessing involves the utilization of all training data points, making the preprocessed data become non-independent. This raises a great challenge in proving the concentration-related results. 2) In the analysis for the preprocessed data, the preprocessing matrix is not commutable with either the empirical data covariance matrix or the population data covariance matrix, thus many prior techniques in the literature cannot be applied and we need to develop new techniques accordingly. This also raises many challenges in obtaining reasonable risk bounds in our analysis.\n\n--- \n\nFinally, we want to emphasize that large language models are intricate systems, and the theoretical understanding of transformers remains in its early stages. Similar to early physics investigations, our approach of gaining insights from carefully designed experiments and then explaining our observations is a reasonable path toward this goal. While we acknowledge numerous areas for future research, we believe our observations and results are both interesting and meaningful, contributing to our understanding of the mechanisms behind large language models, both empirically and theoretically.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely appreciate the thoughtful reviews and comments provided by all reviewers. Below, we address the main points raised, details can be found in corresponding blocks for each reviewer:

- Reviewer fd5H focused on the role of MLP layers in our setting and their potential benefits for other tasks. In our  sparse linear regression setting, nonlinear tokenwise operations are less critical. While recent works suggest MLP layers may benefit other nonlinear tasks and real-world tasks; here we provide additional experiments comparing linear transformers with and without MLP layers to demonstrate that our theoretical analysis without MLP layers is reasonable.

- Reviewer uBeS primarily questioned the experimental details and results when extending our setting to different distributions. We clarify that our experimental and theoretical results can be extended to various data distributions and provide additional experiments to support this claim.

- Reviewer Swt3 inquired about the extensibility of our experimental findings and theoretical results to other settings. We reference recent experimental results on large language models to demonstrate the potential for our findings to be applied to more realistic scenarios.

- Reviewer Bazq acknowledged our experimental findings but expressed interest in 
extending our theoretical analysis to other aspects, such as training dynamics and more complex problem settings like softmax attention. We highlight that our focus is on how trained transformers utilize multi-head attention for in-context learning, and these aspects are a bit beyond the scope of this paper. To address the reviewer's concern, we provide a theoretical explanation showing that the working mechanism we proposed for the multi-head first layer is similar to the theoretical analysis for softmax multi-head attentions. We also highlight the differences between our paper and other works. We believe our findings are both interesting and extensible to more real-world tasks and complex models, warranting further investigation in future work.

Additional experiments in the attached pdf:

Fig 1 compares the performance between linear attention transformers with and without MLP layers, where MLP is initialized with Xavier initialization. We chose the performance with 10 context examples. Although there exist some delicately designed transformers that can conduct more complicated algorithms for lasso, the trained transformer tends to use other solutions. The difference between models with and without MLP layers is minimal in our linear setting, justifying our simplification for theoretical analysis for sparse linear regression problem.

Fig 2 addresses Reviewer uBeS's interest in experimental results with different data distributions. We conducted additional experiments by modifying the distribution of x to x ∼ N(0,Σ), where Σ = I + ζS, and S is a matrix filled with 1. The results show that under different data distributions, our experimental findings remain consistent with orthogonal settings (Fig. 2a). Our pre-gd algorithm still outperforms gd, and our theoretical analysis can be extended to different data distribution cases (Fig 2b, 2c).

Fig 3 responds to Reviewer uBeS's interest in p-probing results using H₀. While we believe our experimental choice of H₁ for different heads is reasonable, using H₀ is could also be another option. Note that under the transformation with the first layer, the data distribution is shifted. Consequently, it is expected that the result using H₀ differs from that using H₁, h = 1. Nevertheless, the results still demonstrate that multi-head attention can achieve lower risk, justifying our preprocessing algorithm for multi-head attention.

Meta-review: 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.
Decision: Accept (poster)