In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention
Authors: Jianliang He
Xintian Pan
Siyu Chen
Zhuoran Yang
Venue: ICML 2025 poster (2025)
Keywords: Multi-head Attention
Mechanism Intepretation
Training Dynamics
Expressive Power
Primary area: deep_learning->attention_mechanisms

Abstract: We study how multi-head softmax attention models are trained to perform in-context learning on linear data. 
Through extensive empirical experiments and rigorous theoretical analysis, we demystify the emergence of elegant attention patterns: a diagonal and homogeneous pattern in the key-query weights, and a last-entry-only and zero-sum pattern in the output-value weights. 
Remarkably, these patterns consistently appear from gradient-based training starting from random initialization. 
Our analysis reveals that such emergent structures enable multi-head attention to approximately implement a debiased gradient descent predictor --- one that outperforms single-head attention and nearly achieves Bayesian optimality up to proportional factor.  
We also extend our study to scenarios with anisotropic covariates and multi-task linear regression. 
Our results reveal that in-context learning ability emerges from the trained transformer as an aggregated effect of its architecture and the underlying data distribution, paving the way for deeper understanding and broader applications of in-context learning.

In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention
Jianliang He
Xintian Pan 2 Siyu Chen
Zhuoran Yang |
Abstract
with few examples instructions within the input con- text, without the need for explicit retraining or fine-tuning, Unlike traditional approaches that rely on extensive fine- tuning. ICL enables LLMs to infer patterns directly from input sequences and generalize to unseen examples a single forward pass (Huang et al., 2022). This emergent be- havior is largely attributed to the self-attention mechanisms in transformers, which allow the models to dynamically cap- ture intricate relationships within the data. Recently, there has been growing interest in exploring transformer based ICL in structured learning settings, such as linear regression, using synthetic data (Bai et al., 2024; Akyürek et al., 2022; Ahn al., 2023a; Fu al., 2023; Mahankali et al.. 2023). These studies not only shed light on the inner workings of transformers but also deepen our understanding of their statistical properties and optimization dynamics, offering a foundation for more interpretable and efficient Al systems.
We study how multi-head softmax attention mod- els are trained to perform in-context learning on linear data. Through extensive empirical exper- iments and rigorous theoretical analysis, we de- mystify the emergence of elegant attention pat- terns: a diagonal and homogeneous pattern in the key-query weights, and fast-entry-only and zero-sum pattern in the output-value weights. Re- markably, these patterns consistently appear from gradient-bart based training starting from random ini- tialization. Our analysis reveals that such emer- gent structures enable multi-head attention to ap- proximately implement a debiased gradient de- scent predictor one that outperforms single- head attention and nearly achieves Bayesian opti- mality up proportional factor. We also extend our study to scenarios with anisotropic covariates and multi-task linear regression. Our results re- veal that in-context learning ability emerges from the trained transformer as an aggregated effect of its architecture and the underlying data distribu- tion, paving the way for deeper understanding and broader applications of in-context learning. Our code is available at https: / gi thub. Xint iantianaan annoPan//ia Lt///ILL_Inng linear.
Prior work has extensively studied how transformers per- form in-context learning for linear regression, exploring various attention architectures and learning strategies. has been shown that single-head linear attention effectively implements preconditioned gradient descent linear data (Zhang et al., 2024; Von Oswald et al., 2023; Akyürek et al., 2022). This analysis has been extended to single-head soft- max transformers (Li et al., 2024a; Huang et al., 2023) as well as multi-head attention models (Chen et al., 2024a;c; Deora et al., 2023; Kim & Suzuki, 2024). Beyond linear regression, recent studies have investigated how transform- ers learn feature representations more complex settings, extending in-context learning to nonlinear data distributions (Huang et al., 2023; Yang et al., 2024; Kim & Suzuki, 2024; Chen et al., 2024b).
1. Introduction
Large language models (LLMs) built on transformer ar- chitectures (Vaswani, 2017) have revolutionized artificial intelligence research. A key capability of these models is their ability to perform in-context learning (ICL) (Dong et al., 2022; Brown et al., 2020), which refers to learning and adapting to new tasks or concepts simply by being provided
Despite significant progress in understanding transformer- based in-context learning (ICL), critical gaps persist in characterizing the training dynamics of multi-head softmax transformers for these tasks. While prior work has yielded insights under constrained theoretical regimes- the neural tangent kernel (NTK) framework (Deora et al., 2023), mean-field approximations (Kim & Suzuki, 2024), and specialized initialization protocols (Chen et al., 2024a)- fundamental questions about their broader behav- ior remain unanswered:
'Department of Statistics and Data Science. Yale University Kuang Yaming Honors School, Nanjing University. Correspon- dence to: Jianliang He Kjianliang.he he ww@yahe.eduuuuu edl,, Zhuoran Yang <zhuoran. edl...du>.
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

In- In-Context Linear Regression with Multi-Heed Softmax Attention
Dynamics
Head 1 gK
Head 1 DV
Deggrn Carment of J0
Head
Head
Positive Heads (2)
-201
Negative Heads
041
1ixe
125x
190c
Steps
OV Dynamics
!tead
Itead 2 DV
Head
Head
Headth
Sum of Heads
+1.0
-445
Negalive Heads (1)
Steps
(a) Leamndd Weights of 2-Head Attention.
(b) Training Dynamic of Attention Weights.
Figure 1. Visualization of main results. Panel (a) presents the heatmaps of 2-head attention matrices trained over 5, L 40. The trained model learns one positive and one negative head. KQ and OV share the same sign for each head, while the parameter scaling is homogeneous across heads. Panel (b) plots the dynamics of average diagonal values in KQ matrices and the last entry of OV vectors for a 2-head attention model along the training, zooming into the opposite pattern.
(i) How do the weights of randomly initialized multi-head softmax transformers evolve during training on linear
We identify two global patterns emerging: (a) The KQ weight matrices develop diagonal structure for the block associated with the covariates and the query. (b) (ii)
ICL tasks?
(ii) What is the final learned transformer model and what are its statistical properties under standard training?
The diagonal entries of the KQ weights and the effec- tive OV weights share the same sign. These patterns lead to the automatic grouping of positive and negative heads, allowing the model effectively capture both positive and negative components. See Figure 1.
(iii) Does multi-head attention confer measurable advan- tages over single-head architectures for linear tasks?
(iv) Does softmax attention outperform linear attention in linear ICI- PCI-particllarly given the latter's 's proven equiv- alence to gradient-basse based optimization (c.g., Von Os- wald et al., 2023)?
(iii) In the trained model, the diagonal KQ weights across all heads are nearly homogeneous in magnitude, and the average effective OV weight is approximately zero We provide theoretical insights into how these patterns emerge through training dynamics.
(v) How do the insights from linear ICL with isotropic covariates extend to scenarios involving anisotropic covariates or multiple linear regression tasks?
(iv) We discover the positive and negative heads jointly approximate a one-step debiased gradient descent- letting multi-head transformers outperform single-head models on linear and the learned model is prov- ably nearly Bayesian-optimal up a constant factor.
We answer these questions through both empirical and theo- retical analysis of training multi-head softmax transformer on linear regression task. We characterize the training dy- namics, demystify the emergence of distinct weight patterns in learned models, and elucidate the underlying mechanisms. We summarize our contributions as follows.
(v) Furthermore, we show that softmax attention outper- forms linear attention in linear ICL tasks, as it learns an algorithm that successfully generalizes to longer sequences at test time.
(i) We empirically analyze the training dynamics of multi-head softmax transformer on single linear re- gression task. We find that in each head, the key-query (KQ) and valuut-value (OV) matrices exhibit a univer- sal structure that implements a kernel regressor, where the kernel determined by the query and covariates.
(vi) We extend our analysis to anisotropic covariates and multi-task settings. In the anisotropic case, we find that the learned model the learned model implements a pre-conditioned version of debiased gradient descent. In the multi-tase setting, the model behavior depends on the task-to-head ratio. Interestingly, when the num-
2

---

In-Context Linear Regression with Multi-Heed Softmax Attention
ber of tasks exceeds the number of heads but remains below twice that number, we observe a superposition phenomenon, where individual heads encode multiple tasks simultaneously.
0 - (0(h, v(h,, kth,, JNEMH C IR(++1)x(++));
H
TFo (Zebd) = Zeod +
Σ 1=1
Zehd
smax o msk Leba K (1)) (1)) Zeba), (2.2)
Related Work. Several works have probed transformers on linear-regresson tasks to elucidate their ICL. capabilities. Empirically, Garg et al. (2022) demonstrate near Bayes- optimal performance, and Von Oswald al. (2023) show that a simplified linear transformer implements gradient- based inference algorithm. Theoretically, one-layer linear attention provably learns preconditioned gradient descent via training dynamics analysis in Zhang et al. (2024) and loss-landscape analysis in Ahn al. (2023a). Chen et al. (2024a) provide the first insight into the standard softmax attention, proving convergence to kernel regressor under specified initialization. In addition, Bai et al. (2024) char- acterizes the expressive power of transformers for various linear regression algorithms. Extensions to more complex settings include two-stage least squares for endogeneity (Liang et al., 2024), adaptive sparse regression (Chen et al., 2024c), EM-based mixture modeling (Jin et al., 2024), multi- step pudient-descenn (Gatmiry et al., 2024), and analyses of the nonlinear softmax component in regression tasks (Aksenov et al., 2024; Sun et al., 2025).
where smax(.) denotes the column-wiss softmax opera- tion and mak(-) dontes the element-wise causal mask, i.e. mak(-);s 1(i < j) 0o 1(i 2 j). Since the (K, Q) and (O, V) matrices always appear in tandem, we simplify the notation by grouping them into single matrices KQ and OV, respectively (Elhage et al., 2021). Following the embedding in (2.1), we extract the prediction By from the (d + 1, L + 1)-th entry. The goal is to estimate the con- ditional expectation E Eigg xy] 3 zq. Accordingly, we define the model output
iy := (Zgi ( De) (EEL)) = TFo TFo(Zebd)dtI.LH E E.
Compared to the standard decoder-only transformer model (e.g., Vaswani, 2017), we omit layer normalization and posi- tional embeddings. This choice arises because a single-layer architecture does not suffer from vanishing or exploding gra- dients. In addition, the associated linear regression problem is sermutation-invarite in {(1,, we))e(()) making posi- tional information unnecessary.
More related works and notation is provided in $A.
KQ and OV Circuits. Note the computation in (2.2) only depends on matrix products k Q(h) and y re- spectively. Following Elhage al. (2021), we refer to these matrices as the KQ and OV circuits and simplify the nota- tion as K KQ(h) ((h)" Q(h) and OV(M) (l)) In specific, KQ circuits characterize to what extent query token attends to key token, and OV circuits determine how token contributes to the output when attended to.
2. Preliminaries
We consider the setting of training transformers over the task of in-context linear regression, following the framework widely considered in literature (e.g., Garg et al., 2022).
Data Distribution and Embedding. Let Jy E Rdd denote the f-th covariate i.i.d drawn from distribution Pa and iid. let Tig ~ Px represent a new test input. The coefficient 3 E Rd is sampled from P... and the corresponding response is generated as y/ = 3 xig Eg, where the noise eg i.i.d sampled from (N o2). To perform ICL., we embed the dataset {(ze, ye) bejLe) along with the test input Fy ~ Pa into a sequence Zeod E (I+++)) formatted as follows:
Notice that the last entry of 2y is zero and we use the last row of ((h) to get Ty. Thus, in terms of computing dn we can write the KQ and OV circuits as
KQ =
KQN (M)
II
:
:]
-
-Ovar ov
*
*
KQ
(h)
(M)
21
Zodd =
Z
[q]
[1
Vi
22
U2
EL
UL
Lg
0
(2.1)
Here we use "x" to denote the entries that do not affect Dy. In the above two-by-two block matrix format, the top-left blocks are d-by-d matrices. See Figure 6 for complete depiction of the transformer architecture.
where we also denote by X = [1,, x2. I E Lxd
and y = [y,, ................................ 8/] T E RR In this paper, we focus on the isotropic case where P.. = NN Ta) and P3- = (0, La/d), and the non-isometric cases are in $B.3.
Training Setup. To investigate how transformers learn to solve the linear regression problem in context during pretraining, we examine both the training dynamics and loss landscape. Specifically, we train the model by minimizing the population mcan-squared error:
Multi-Head Softmax Attention. The attention model is a seneence-to-sequene model that takes Zebd as input and outputs a sequence of the same shape. We focus on the lae- softmax attention with H heads, parametrized by
L(0)= = ( = EE(u, TFo )a+1 )4+1,441) (2.3)
3

---

In-Context Linear Regression with Multi-Heed Softmax Attention
where the expectation is taken w.r.t 3 ~ P and (re. Me) Wi ~ Pz @ ()) for all & E [ U {q}.
Head 1 QK
Head OV
64
3. Empirical Insights
80
Although previous studies have explored how transformers perform ICL under the linear regression framework, much of this research has been limited to experimental analyses (Garg et al., 2022), linear transformers (Von Oswald et al., 2023), or single-head attention (Huang et al., 2023). This leaves a gap in understanding multi-hean softmax attention models, which are used in practice. To this end, we first conduct experiments to empirically investigate how multi-
2
-12
-15
-10
-14
(a) Learned Weights of Single-head Attention.
Dynamics
Dynamics
Average Diagnaal Elemants of
Head 1
Head 1
head softmax attention models learn to solve ICI with linear data. See $C.1 for experimental setup and additional results.
0:20
Head 2
Head 2
Head 3.
Head 3
0:2
Head 4
LI0-
Sum of Head 1&3
Observation 1. For any number of heads with H 2 1, in the trained one-layer multi-head attention model, the KQ and OV circuits take the following form:
2:90
-.30
(QO))
(M) 0
[ OJ (1)) (3.1)
-0.0
59c
3000
250x Juhc 300c
189c
200c X0.
4000
Steps
(b) 4-Head KQ Dynamics.
(c) 4-Head OV Dynamics.
Moreover, KQ and OV share the same signs within each head, ie., (igg(((()) (i) We refer to this prop- erty as sign-matching. In some cases, dummy heads may emerge, where (f) x 0 and ( x O. Notably, these patterns develop early in training and remain consistent throughout the optimization process.
Figure 2. Learned pattern of single-head attention and training dynamics of 4-head model. Panel (a) plots the heatmap of 1-head attention trained under the same setting as in Figure la. Panel (b), and (c) plot the dynamics of average diagonal values in KQ matrices and the last entry of OV vectors for 2-head attention model along the training. Head 2 the positive head, Head 1&4 are coupled to act as the negative head, and Head 3 is dummy.
We illustrate the KQ and OV circuit patterns for both one- head and two-head attention models in Figures 2a and la, respectively. For the OV circuits, we only plot the transpose of the last rows, i.e., (h) and (V)(A) These figures show that the trained multi-head softmax attention models with various numbers of heads exhibit consistent pattern:
(w ......... (. (T)).. Under such a structure, the transformer predictor takes the form:
H
(h)
(v. smax(w(M) xxa))
h=1
(a) For each KQ circuit, the top-left d--y-d submatrix is diagonal and proportional to an identity matrix, i.e.,
H
Σ Ve exp(w(()) xg t=1 E exp(w((t) xx zn))
Σ (())))
E R. (3.2)
kQ(i) (h) Ia for some (h) e IR Moreover, the first entries of the last row are all approximately zero, i.e., KQ = TT.
Thus, each head acts separate kernel regressor, and thus the attention model can be interpreted as the sum of kernel regressors (see $E.1). Here, 1/w((h) plays the role of bandwidth of the kernel. Each term in the sum in (3.2) is a weighted sum of responses {ye Jee(z) in the ICL. samples, and the weights are computed based on the similarity be- tween the test input xy and the sample covariate xy 's. The notion of similarity governed by the parameter (h))s.
(b) The last column each OV circuit, as vector in Rd+1 admits a last-entry-only pattern. That is, only the last entry is non-zero, which is given by (b( E RR
(c) We observe that (h) and (M) always have the same sign. Thus, each head can be categorized into either positive or negative head, depending on the sign of (h)) Positive and negative heads are defined H4+ = [h: (h) > 0} and H_ {h: (h)) 0}.
The single-head case has been studied in (Chen et al., 2024a), which shows that the single-head model functions as a kernel estimator to solve linear regression, where w()) = 1/Va and ( x /d when L is large. We replicate this finding in Figure 2a, where w() xs 0.52 zs 1/vd. Addi- tionally, we show that such a pattern is shared by multi-heat
The pattern of KQ and OV circuits shows that the weight matrices of the learned transformer essentially are gov- emed by 2H numbers I = ((("),...,,(((1))) (H))) and w =
4

---

In-Context Linear Regression with Multi-Heed Softmax Attention
attention models with H 2 2.
Head
... Muiti head TF
Head
43
Single-head
Moreover, when / > 2, as we show in Figure 13, this pat- tern persists. Additionally, we see dummy head (Head 4 in this case), i.c., a head whose KQ and OV matrices are close to zero. This means that this head does not contribute to the prediction. However, the appearance of dummy heads is not guaranteed, even with a large number of heads, as it results from the randomness of the optimization algorithm. As we will show later, dummy heads do not affect the statistical properties of the learned model and thus can be ignored in
GD Est
. Hoad
-9- Deblased GD Est
Ridge
9.
o o
Prompt
Prompt Length
(a) Loss of Attention Model with (b) Loss of Attention Models Different Head Number.
statistical analysis.
and Canonical Estimators.
Finally, in Figure 4 and 13, we plot the KQ and OV circuits of two-head and four-head attention models, respectively, along the training trajectory. We observe that the patterns in (3.1) emerge early during training and persist throughout.
Figure 3. Loss comparison of different models with different num- bers heads and the canonical estimators. The models are trained over L = 40 and evaluated over different input lengths.
When H 2 2, (w, pi) of the learned attention model satisfies that
Observation 2.
Comparing Figures la with 2a, we observe that in the two- head attention model, the magnitude of KQ circuits (~ 0.13) is smaller than that of the single-head attention (~ 0.52), and the magnitude of OV circuits (~ 3.50) is larger than that of the single-head attention model (~ 1.42). Hence, the two-head attention model learns different predictor than the single-head attention model.
(i) Homogeneous KQ scaling: The scaling of the top- left diagonal submatrix of each K KQ is nearly identical across all positive and negative heads, i.e., w Rs y for all E H+ U.-. Zero-sum OV: The sum of Jeemn is approxi- mately zero, i.e., (p,, Ix) ) x O.
More interestingly, revisiting the patterns learned in Fig- ures lb, 2b, and 2c we see that the two-head and four-head attention model learns the same predictor. To see this, note that Figures 1b and 2b have the same scaling in the KQ circuits (~ 0.13). Similarly, the magnitude of OV circuits aligns across positive and negative groups (~ 3.5). This implies that the four-head attention model ultimately con- verges to the same predictor as the two-head model. More broadly, this phenomenon holds across all multi-head atten- tion models with H 2 2 which consistently learn nearly identical predictors. Furthermore, study the statistical errors of these learned attention models, we compare their ICL loss in the test time, defined in (2.3). In particular, we consider length generalization where Zebd in (2.1) involves a different length that might be different from the training data with L = 40. We plot the ICL losses of various models in Figure 3a, which reveals the following findings:
By examining Figures la more closely, the two attention heads in the two-head attention model converge opposite heads, i.e., (1)) ~ (2) and y((1)) s (2) Moreover, as we show in Figures 2b and 2c, similar patterns emerge in the four-head attention model. In particular, Head 4 is a dummy head, i.e., (4) R (u(+) = O. Head and Head are negative heads and Head is a positive head. Thus, we conclude
(a) For non-dummy heads, the scaling of the non-zero entries of the KQ circuits is nearly identical across heads, i.e.,
| se y for all h H4 H.H.
(b) For the OV circuits. it holds that (u, 1H) ze 0 such that Enen, (1())) R Σnen_ (10))
In single-head attention, the attention head is either positive or negative. The multi-head attention models have both positive and negative heads, with special patterns in (w,
(a) Multi-head attention outperforms single-head attention while maintaining nearly identical performance across different the number of heads used.
Observation 3. In terms of the ICL prediction error, the two-head attention model outperforms single-head model. Moreover, multi-head models with H > 2 exhibit similar performance, closely approximating that of the vanilla gra- dient descent (GD) predictor T75ed = --1 Ete zi - Me- Furthermore, all softmax attention models can generalize in length during the test time. In addition, the statistical error of the multi-head attention model is comparable to the optimal Bayes estimator up to a proportionality factor.
(b) All softmax attention models generalize in length and the ICL loss decreases L increases.
Moreover, we plot the ICL losses of the trained attention models and the classical statistical estimators in Figure 3a. This figure demonstrates that multi- head attention models closely track the performance of both the vanilla GD pre- dictor and its debiased variant (see $4.3). This suggests
5

---

In- Cn-CConextt Linear Regression with Multi-Heed Softmax Attention
consider the parameterization in (3.1) to interpret the core aspects of model training. With this reparameterization, the evolution of the model model can be analyzed by tracking (w, p), and the transformer predictor follows (3.2). Moti- vated by the analysis in (Chen et al., 2024a), we present a refined argument for approximating the population loss in (2.3) under the reparameterization in 13.1). As we will see later, gradient flow based on this approximate loss reveals the empirieal observations from theoretical perspective.
Proposition 4.1 (Informal), Consider an H-head attention model parameterized by (3.1) with dimension d E Z+ and sample size L E zz, Suppose that > log L and param- eters (w. (1) C R satisfies that ||wl|s0 2 vlog L/d d and || E I. 1,1/5, then it holds that
Figure 4. Heatmaps of the KQ matrices and OV vectors along training epochs with H = 2. d = 5, and L = 40, trained over 10t steps with random initialization. KQ matrices (top) and OV vectors (bottom) form the diagonal and the last-entry-only patterns respectively during the early stages of training and then optimize within the simplified regime described in (3.1).
[(w, 14) = = - 2p' w + www' (ww + (1 (1+))) --1 -exp(dww> 1) + O(dH2.L- 1/5).
The formal statement and proof of Proposition 4.1 are de- ferred to $E.2, where the result extends beyond the case d >> log L and allows a more flexible trade-off between the scaling and the resulting approximation error. Proposi- tion 4.1 establishes that, under certain scaling assumptions, the true loss can be well approximated by simplified for- mulation in terms of (wz, μ1) with the approximation error vanishing as the ICL sample size L increases. In addition, we provide experimental validations for Proposition 4.1. As shown in Figure 5, the approximation can effectively cap- ture the true loss landscape (see Figure 5b.) and performs particularly well when w or μ is small (see Figure Sa). See $C.4 for more detailed explanations and additional experi- mental validations of Proposition 4.1. Hence, to analyze the parameter evolution of the transformer, it suffices to con- sider the gradient dynamics of the approximate loss, which is given by
that multi-head attention model approximately implements a version of GD algorithm.
The training of KQ and OV circuits across different heads follow similar trajectories under random initializations. In particular, the patterns found by Observations and 2 appear early in the optimiza- tion trajectory, and are preserved during training. More- over, the magnitude of (h) in positive negative heads first increases and then decreases to converge, while the Enen, (((h) monotonically increase until convergence.
Observation 4.
Despite random initializations, the parameter evolution fol- lows a highly consistent trajectory throughout training. As shown in Figure 4, the attention model quickly develops the pattern in (3.1), identified in Observation 1, during the early stages of training. The attention model continues to opti- mize the loss function with this pattern preserved throughout the training process. Moreover, by looking at the dynamics of (w, p),, we observe that the two properties identified by Observation 2 are also preserved during the training. Fur- thermore. the training dynamics of the KQ circuits are not monotonic. The magnitudes of (hh)'s first increase and then decrease to stabilize. Whereas the magnitude of OV circuits, i..., ΣnEn, (,()), increases steadily.
Ō(w, p)) 02 x²+(1= 4 Tw)²
+ = +1++02)) L- exp(dww T)p.
(4.1)
4.2. Training Dynamics, Emerged Patterns and Solution Manifold
In this section, provide detailed analysis of the evolu- tion of attention weights during the training process, with a focus on multi-head attention models. Starting from ran- dom initializations, we show that the parameters evolve to form a pattern consistently similar. To interpret such a phe- nomenon, we provide a theoretical explanation and follow the training dynamics based on the approximated loss in (4.1). Our analysis focuses on the training process using gra- dient descent (GD). For a learning rate r) > O, the parameter update at step t E N is given by
4. Mechanistic Interpretation
4.1. Reparametrization of Attention Model
As shown in Observation 4 in $3, the attention model devel- ops diagonal-only pattern in KQ circuits and a last-entry only pattern in OV circuits during training. This structure emerges early in the training and then continues optimizing within this regime (see Figure 4). Thus, is sufficient to
Pet1 <- P4- V D(H,, W441 +- wi wi- V Vo (.[( W.).
6

---

In- Cn-CConexxt Linear Regression with Multi-Heed Softmax Attention
and the updating scheme of w shares similar form as
47
Wi+1 t + + (1 - (1 + (1 + ++)+== (pe, wi)) pe
Sign-Matching Term
2n (1 0²) 00
d
(k 1)! (pe, w Me 2 ws uoth1
(a) Loss Difference C E. (b) Landscape of True Loss.
High-Order Terms
Figure 5. Experimental validation for Proposition 4.1 with 2-head attention model under (w, -w) and (μ, -H) with L - 40, d = and o² 0.1. Panel (a) plots the heatmap of difference between the actual and approximated loss [(u,w) w) - E(H, Panel (b) illustrate the actual and approximated loss landscape, respectively. The gray arrow highlights the "river valley" loss landscape where the loss decreases gradually that direction towards minimum.
In the early stage of training, the higher-order terms are negligible due to small initializations, and the sign-matching terms and zero-sum OV terms dominate. For the zero-sum OV term, it is straightforward to see that it encourages (Pe, 1 1H) = O. For sign matching terms, (Wi: 14) remains small in the initial stage such that (1 + (1+0) o?) dd-'). (pr, < 1. Thus, 14, and wi are updated in each other's direction, gradually aligning to share the same sign. Also, the sign-matching terms contribute to the gradient update until the following optimality equality is established:
Following this, the training dynamic undergoes two stages:
In Stage I, two patterns emerged-sigte matching and zero-sum OV, which optimizes the low-order terms mainly due to the small initialization and are maintained along training (see (4.2)). In addition, optimality equality in (4.3) holds approximately, which facilitates subsequent evolution.
(1 + (1+a²) dd-1)- (pe, wi) = 11 (4.3)
In this stage, we ignore the contribution of the high-order terms the gradient both μr, and we have small magni- tude. After (4.2) and (4.3) are establised, the training enters the second stage where the high-order terms dominate.
In Stage II, the (M)* s converge to homogeneous KQ scal- ing in (4.4) to optimize the remaining high-orde order terms. Based on the signs, these heads group into positive, nega- tive and dummy heads. Finally, μ (1h)'s fall into the solution manifold in (4.5) defined by scale of (h).
Stage II: Emergence of Homogeneous KQ Scaling. As shown in Figure Ib, after some training steps, these heads are grouped into positive, negative, and dummy heads. For dummy heads, all parameters degenerate to zero, indicating that these heads do not contribute to the final predictor and output O. For positive and negative heads, their KQ parameters share the same scaling with different signs:
Stage I: Establish the Sign-matching and Zero-sum OV Pattern. Consider starting from a small initialization and assume that || S 1/Vd and || I evd hold for all N with constant e e (0, 1). As shown in Figure lb, at the very beginning of the training process, the model quickly develops the sign matching and zero-sum OV pat- tern, summarized below: for all h E [H], it holds that
(N) sign(w(h) with 7 > O, vh E [H]. (4.4)
To understand how the homogeneous scaling develops, which significantly simplifies the predictor given by the multi-head attention (see $4.3), revisit the population loss and focus on the regime where (4.2) and (4.3) are well- established after Stage I. Note that
(n, Ir) = 0, and (ign(w(")) = sign(a (h). (4.2)
To understand the mechanisms, we examine the gradient update. By applying the Taylor expansion over "xp(dwow) at w = 0H,, the gradient update of μ at step followz
E(wr, me) = 0² +1++02) (pr. 1z)
Σ d
+ (1-(me + 02) L-1
M441 + 1e + 2m 1-++11++0 aal=-))((aa (H, we
where the first two terms are constants given the zero-sum and transitional optimality. Hence, the remaining high-order terms dominates the dynamic, which satisfies
Sign-Matching Term
21y(1 02) L
(ur, 11) +
(per ai yuu M
k!
dl (Mes lwoe WK)2 + Σ
dl (lpell I weyl 1ok2
Σ
Zero-sum OV Term
Kt
High-Order Terms
KE )
ke
7

---

In- Cn-Context Linear Regression with Multi-Heed Softmax Attention
de
4.3. Expression, Approximation and Optimality
(leel 2
M
(12g-11E2+))
In this section, we examine the statistical properties of the learned attention models. Our experiments show that, re- gardless of the number of heads, attention models consis- tently learn the diagonal-only and last-entry-only patterns in (3.1) while using different predictors. For multi-head attention with H 2 2. heads are grouped into positive, neg- ative and dummy heads as discussed in $4.2. Positive and negative heads are coupled to approximate gradient descent (GD), consistent with the findings from linear transformers (Mahankali et al., 2023; Ahn et al., 2023a; Zhang et al., 2024), but with an additional debiasing term. As we will show below, this is a result of the fact that the limiting val- ues of attention parameters fall in the solution manifold P with a small 7. Moreover, when L is much larger than d. debiased GD coincides with vanilla GD, which means that multi-head attention learns the same algorithm the linear attention in this regime.
where the first equality follows sign(w>())- = sign(!!)) and the last one uses non-negativity of even--power terms and Hölder's 's inequality for odd-power terms. Given zero-mean Mr, the lower bound is attained when (i) the ween-power terms reach the O-minima by taking | = y 1 since Kye, |um] ( = = (Me, 1 0 for all & E {2g: q E Z+} }. (ii) the odd-power terms attain the lower bound if and only if [pal @ we | x l for all k E {2q - 1: q E -+} such that [wee] x 1g, leading to the homogeneous pattern.
Solution Manifold. Suppose patterns in (4.2) and (4.4) fully emerge. At this convergence stage, for arbitrary scale y, with a bit of abuse of notation, the approximated loss can be rewritten as:
[(y,)) =2 0² + (1 (1- 11
+ +1++02) L- sinh(d??)) ||a||
Multi-head Attention Predictor. Consider an arbitrary scale y > 0, then for any parameters on the solution mani- fold (ws, (u) C Sq. the attention predictor takes the form:
where we use the Taylor series of hyperbolic sine functions. Following this, the loss is quadratic with respect to || II thus for any fixed scale 7 > 0, the minimizer satisfies that
L
V exp(7 27 mg)
F4 =
Σ
||a|| (72 + o2) --1 sinh(d??)) -..7---. - == 2
Σ exp(7 . rr. zzg)
hEW
Here we define He, as one half of the optimal value of || Hence, we know that the limiting w and μu are characterized by a solution manifold, given by g = {-x7,>0, where
ye expl (xxp(-7- 22 zzg)
+ Σ
Σ exp(-x x xzg)
bEA
Then, for any H 2 2, the multi-head attention model im- plements an equivalent predictor as 2-head model with Welf = -1,,-) and Hetf = (EnEn, 11(t) En w()) _n))))) with
Sn = (((w))) S CR.. (())) = 7. siguu(u
Σ m0a) NEH
E w(l) mn) (4.5) NEW
En(h)
Σnen, 1(()) =- -- En_ (1()) pyy For small scale y > O, by linearizing exp(.) around we have
We remark that the leamed scale y depends on the learning rate, training steps, and batch size (if SGD or Adam is applied) due to the "river valley" loss landscape near the global minimum (see Figure 5b). This phenomenon can be attributed to edge of stability (Cohen et al., 2021), behavior commonly observed in neural network training. In this regime, gradient descent progresses non-monotonically, oscillating between the "valley walls" of the loss surface and failing to fully converge to the minimum.
2 Dy ~
Σ ur 3 Fgg + th (n) = T.
y F2 rqg
T.
where we denote 3, = x2 - + Ze and the approxi- mation also uses the first-order approximation of reciprocal (see Remark F.1). Thus, multi-head attention emulates the one-step pre-conditioned gradient descent' optimize the empirical loss 1 Ei= (BT It- ye)² with initialization So O, learning rate 2 and the debiased covariates Zy 's. In experiments, the learned 7 is very small (~~ 0.1), which ensures the validity of the approximations. This observa- tion suggests that, despite the nonlinear nature of softmax, the model can still capture the inner linearity. To achieve
Although the learned 7 can vary slightly, j always lies within its corresponding solution manifold .y The optimal value of y and the resulting bransformer-based ICL predic- tor are derived later in $4.3. Furthermore, to explain the evolution of the parameters, we provide a more detailed analysis of the gradient flow of training a two-head attention based on the approximate loss in $D under a well-specified initialization, where the full gradient flow dynamics can be characterized explicitly.
Under the isotropic setup, the pre-conditioned GD coincides with the standard GD due to the identity covariance matrix. Here, we refer to it as pre-conditioned GD to align with experimental ob- servations in the anisotropic case (see (B.3), where clear evidence indicates that it utilizes a pre-conditioner.
8

---

In-Context Linear Regression with Multi-Heed Softmax Attention
this, the softmax attentions work in the linear regime by taking small-scale 7 and canceling out the constants via coupled positive-negattve heads. As shown in Figure 3b, the learned predictor performs nearly identically to the de- biased GID predictor, which also matches the standard GD predictor when L is large since I. Ei=1 Fe x B[zz] = Op and 1 Et= Ue E De + ce) 0.
a result, the optimal attention model that minimizes C(w. p) coincides with the debiased GD estimator
09 (n"),, where we define nf" =
-z)
(ii) (Bayes Risk) Consider the high-dimension
regime
where + oc and d/L. + &, where & E (0, x)) is a constant. Suppose the noise level 0² > 0 and 5 is
sufficiently small such that o² + > 1. Then
Different from multi- head models, single-head attention has less expressivity, and the learned predictor takes the form of
Single-Head Attention Predictor.
El (n))) BayesRisk 5,02
I $1+0 $1+0 ((+EO))
L Ve explx x/ xg) In 1.
where BayesRiskg, denotes the limiting Bayes risk.
(=1 E explw Ty xn))
The proof of Theorem 4.2 is deferred to SF. In part (i), we show that softmax transformer can represent the linear pre-conditioned GD predictor with small scale y Part (ii) analyzes the approximated global minimizer corresponds to the debiased GD predictor. Finally, in part (ii), we compare the learned predictor with the Bayes risk and show that it is nearly optimal in the proportional regime, particularly when observations are noisy or the dimensionality of the covariates is high.
By minimizing the approximate loss, it is shown in (Chen et al., 2024a) that the minimizer is given by w' = 1/va and n vd. (1 e (1 o²) dd.- We remark that the single-head attention precisely implements Nadaraya- Watson estimator with Gaussian RBF kernel, when the co- variates are sampled from sphere, i.e., 22) = =d-11 Moreover, the optimal parameter w" corresponds to a band- width with value d2/4, Hence, we see why having one ad- ditional attention head significantly improves the statistical rate in in-context linear regression as shown in Figure 3 one-head attention corresponds to nonparametric predictor while a multi-head attention yields parametric predictor.
Discussions on Linear Transformer. Building on our analysis, we can show that by working in the linear regime, linear and softmax attention can achieve comparable expres- sive power. However, linear transformers rely on fixed nor- malization factor, limiting their ability to generalize across different sequence lengths, whereas softmax transformers dynamically adjust for flexible length generalization (see B for detailed discussions).
In the following, we provide formal statement of the above argument and establish the optimality of the learned predic- tor in a high-dimensional asymptotic regime.
Theorem 4.2. Consider a larger parameter space f 2 y defined
Extentions. We extend our analysis to anisotropic and multi-task regression settings. In the anisotropic case, trained multi-head attention continues to learn a precon- ditioned gradient descent predictor. For multi-task learning, the number of attention heads determines distinct learning strategies, ranging from weighted kernel regression to inde- pendent preconditioned gradient descent predictors for each task, revealing an intriguing superposition phenomenon (see $B.3, SC.2 and C.3 for details).
5 =[(w, M)): min min(m14+111 17--11 1, By > O, w (h) T sign (wr (1) th E [H] (4.6)
and denote ΣnEn, (()) ut and ΣnER 100h p_..
(i) (Approximation) Let 17 > O be a constant learning rate and let 8 E (0, be a given failure probability. For any scaling 7 > O, we define jı n/ (27). Con- sider a multi-head atlention with no dummy head and O E S Moreover, we set p+ =-M- B. Then, when L Z 1og(1/8) and 7 S. (Lg(/)) with probability at least 1 - 8, we have
5. Conclusion
Our study provides a comprehensive analysis of multi-head softmax transformers for in-context learning in linear re- gression, uncovering key structural patterns that emerge during training. We characterize the model's development and demonstrate its implementation of preconditioned gra- dient descent, shedding light on its underlying optimization dynamics. These findings enhance our understanding of transformer- based ICL, highlighting its advantages over single-head and linear attention models and offering valu- able insights into its broader applicability.
(0( - 1
(n)) (m)) s 0(
where Ö(:) omits logarithmic factors. In particular, suppose we drive to zero while keeping já = n/ (27),
the resulting (0) coincides with (n))
(ii) (Optimality) Consider minimizing Ō(w, n)) over J.
The minimum is attained in Jr, imizer can be obtained by taking 7 - 0+ in (4.5). As
. In particular, the min-

---

In-Context Linear Regression with Multi-Heea Softmax Attention
Impact Statement
Cohen, J. M., Kaur, S., Li, Y., Kolter, J. Z., and Talwalkar,
A. Gradient descent neural networks typically occurs at the edge of stability, arXiv preprint arXiv:2103.00065, 2021.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here
Cui, Y., Ren, J., He, P., Tang. J., and Xing, Y. Superiority of multi-head attention in in-context linear regression. arXiv preprint arXiv:2401. 17426,,... 2024.
References
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn to implement preconditioned gradient de- scent for in-context learning. Advances in Neural Infor- mation Processing Systems, 36:45614-45650, 2023a.
Deora, P., Ghaderi, R., Taheri, H., and Thrampoulidis, C. On the optimization and generalization of multi-head attention. arXiv preprint urXiv:2310.12680 2023.
Dobriban, E. and Wager, S. High-dimensional asymptoties of prediction: Ridge regression and classification. The Annals of Statistics, 46( 46(1):247-279 2018.
Ahn, K., Cheng, X., Song, M., Yun, C., Jadbabaie, A., and Sra, S. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv preprint 100000002, 2023b.
Dong, Q. Li, L., Dai, D., Zheng, C., Ma, J., Li, R., Xia, H., Xu, J., Wu, Z., Liu, T., et al. A survey on in-context learning. arXiv preprint arXiv:23010002
2022.
Aksenov, Y., Balagansky, N., Vaina, M. L. C., Shaposh- nikov, B., Gorbatovski, A., and Gavrilov, D. Linear transformers with learnable kernel functions are better in- context models. arXiv preprint arXiv:2402. 10644, 2024.
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., et al. mathematical framework for transformer circuits. Transformer Circuits Thread, 1(1):12, 2021.
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. What learning algorithm is in-context learn- ing? investigations with linear models. arXiv preprint arXiv:2211.15661 2022.
Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield- Hatfseld-Dodds, Z., Lasenby, R., Drain, D., Chen, C., al. Toy models of superposition. arXiv
preprint arXiv:2209.1055 2022.
Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. Trans- formers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural infor- mation processing systems, 36, 2024.
Fikioris, G. Spectral properties of kac-murdock-szege
matrices with a complex parameter. Linear Algebra and its Applications, 553:182-210 2018.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,
Fu, D., Chen, T.-Q., Jia, R., and Sharan, V. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310. 17066, 2023.
Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Fu, H., Guo, T., Bai, Y., and Mei. S. What can a single attention layer learn? study through the random features lens. Advances in Neural Information Processing Systems, 36, 2024.
Chen, S., Sheen, H., Wang, T., and Yang. Z. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality, arXiv preprint arXiv:2402. 19442, 2024a.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. What can transformers learn in-context? case study of sim- ple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
Chen, S., Sheen, H., Wang, T., and Yang, Z. Unveiling induc- tion heads: Provable training dynamics and feature learn- ing in transformers. arXiv preprint arXiv:2409. 10559, 2024b.
Gatmiry, K., Saunshi, N., Reddi, S. J., Jegelka, S., and Kumar, S. Can looped transformers learn to implement multi-step gradient descent for in-context learning? arXiv preprint arXiv:2410.08229,, 2024.
Chen, X., Zhao, L., and Zou, D. How transformers utilize multi-head attention in in-context learning? a case study on sparse linear regression. arXiv preprint arXiv:2408. 04532, 2024c.
Giannou, A., Yang, L., Wang, T., Papailiopoulos, D., and Lee, J. D. How well can transformers emulate in-context newton's 's method? arXiv preprint Sit::2400303888 2024.
Cheng, X., Chen, Y., and Sra, S. Transformers implement functional gradient descent to learn non- non-linear functions in context. arXiv preprint arXiv:23/2.06528, 2023.
10

---

In-Context Linear Regression with Multi-Heea Softmax Attention
He, J., Chen, S., Zhang, E,, and Yang. Z. From words to actions: Unveiling the theoretical underpinnings of Ilm-driven autonomous systems. arXiv preprint arXiv:2405.19883, 2024.
Lin, Z. and Lee, K. Dual operating modes of in-context learning. arXiv preprint arXiv:2402./8819 2024.
Liu, J., Shen, D., Zhang, Y., Dolan, B., Carin, L.., and Chen, W. What makes good in-context examples for gpt-3? arXiv preprint arXiv: 222.0000..0,,,. 2021.
Hu, X., Zhang, F., Chen, S., and Yang, Z. Unveiling the sta- tistical foundations of chain-of-thought prompting meth- ods. arXiv preprint arXiv:2408. 14511, 2024.
Mahankali, A., Hashimoto, T. B., and Ma, T. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint 23x7::2307.03576 2023.
Huang, Y., Chen, Y., Yu. Z. and McKeown, K. In- context learning distillation: Transferring few-shot learn- ing ability pre-trained language models. arXiv preprint arXiv:2212.1067 2022.
Marchenko, V. A. and Pastur, I. Distribution eigenval- ues for some sets of random matrices. Matematicheskii Sbornik, 114(4):507-536, 1967.
Huang, Y., Cheng, Y., and Liang. Y. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
Min, S., Lewis, M., Zettlemoyer, L., and Hajishirzi, H. Metaicl: Learning to learn context. arXiv preprint vrXiv:21001594// 2021.
Jeon, H. J., Lee, J. D., Lei, Q. and Van Roy, B. An informationntheoretii analysis of in-context learning. arXiv preprint arXiv:2401.15330 2024.
Nichami, E., Damian, A., and Lee, D. How transform- ers learn causal structure with gradient descent. arXiv 2024.
Jin, Y., Balasubramanian, K., and Lai, L. Provable in- context learning for mixture of linear regressions using transformers. arXiv preprint arXiv:2410. 14183, 2024.
preprint arXiv:2402.1477
Nie, F., Chen, M., Zhang, Z., and Cheng, X. Improving few-shot performance of language models via nearest neighhor calibration. arXiv preprint arXiv:2212002216 2022.
Kim, J. and Suzuki, T. Transformers learn nonlinear fea- tures in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv:2402.01255, 2024.
Sheen, H., Chen, S., Wang, T., and Zhou, H. H. Implicit regularization gradient flow on one-layer softmax at- tention. arXiv preprint arXiv:2403. 08699, 2024.
Kim, J., Kwon, S., Choi, Y., Park, J., Cho, J., Lee, D., and Ryu, E. K. Task diversity shortens the icl plateau. arXiv preprint arXiv:2410.05448, 2024.
Sinii, V., Nikulin, A., Kurenkov, V., Zisman, I., and Kolesnikov, S. In-context reinforcement learning for variable action spaces. arXiv preprint arXiv:23/2.13322, 2023.
Laurent, B. and Massart, P. Adaptive estimation of quadratic functional by model selection. Annals of statis- fics, pp. 1302-1338, 2000.
Ledoit, O. and Péché, S. Eigenvectors of some large sample covariance matrix ensembles. Probability Theory and Related Fields, 151(1):233-25 264,244,32222. 2011.
Sun, H., Jadbabaie, A., and Azizan, N. In-context learning polynomial kernel regression in transformers with glu layers. arXiv preprint arXiv:2501.1813 2025.
Li, H., Wang, M., Lu, S., Cui, X., and Chen, P.-Y Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint ar 2tV::2402.16607.. 3, 2024a.
Tripuraneni, N., Jin, C., and Jordan, M. Provable meta- learning of linear representations. In International Con- ference on Machine Learning, pp. 10434-10443. PMLR, 2021.
Li, Z., Cao, Y., Gao, C., He, Y., Liu, H., Klusowski, J. M., Fan, J., and Wang, M. One-layer transformer provably learns one-nearest neighbor in context. arXiv preprint arXiv:2411.10830 2024b.
Vaswani, A. Attention all you need. Advances in Neural
Information Processing Systems, 2017.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pp. PMLR, 2023.
Liang, H., Balasubramanian, K., and Lai, L. Transformers handle endogeneity in in-context linear regression. arXiv preprint arXiv:2410.01265. 2024.
Lin, L., Bai, Y., and Mei, S. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. arXiv preprint arXiv:2310.08560 2023.
Wang. L., Zhang, Z., and Dobriban. Inference in ran- domized least squares and pea via normality quadratic forms. arXiv preprint arXiv:2404 0.0000002..... 2024.

---

In-Context Linear Regression with Multi-Heed Softmax Attention
Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02066 2021.
Yang, T., Huang. Y. Liang, and Chi, Y. In-context learning with representations: Contextual generalization of trained transformers. arXiv preprint arXiv:2408. 10147, 2024.
Zhang, R., Frei, S., and Bartlett, L.. Trained transform- learn linear models in-context. Journal of Machine Learning Research, 25(49): 2024.
Zhang, Y., Zhang, F., Yang. Z., and Wang, Z. What and how does in-context learning learn? bayesian model averag- ing, parameterization, and generalization. arXiv preprint arXiv:2305.19422 2023.
Zhang, Y., Singh, A. K., Latham, P. E., and Saxe, A. Train- ing dynamics of in-context learning linear attention. arXiv preprint arXiv:2501. 2025.
12

---

In-Context Linear Regression with Multi-Head Softmax Attention
A. Additional Background and Related Works
Notation. For some n E N++ let [n]] = {i E 1 s For vector e E Rd, denote by O the element- wise Hadamard product and uck (uf .uf...... vg)) Let -1p denote the fn-norm and softmax operator is defined as Let sign(-) denote the sign function that returns 1, -1 or 0 based on the sign of input R. Let 1a and O. denote the all-one and all-zero vector of size d and denote Ia as the d-by-d identity matrix. For vector E Rd and indices set I C [d], we define VI. = (us)ET E IR 21, For two functions
smax(e) = (smax(v)))) where smax(u)- = exp():)/> Σ, exp(vy).
(z) 2 0 and g(z) 2 0 defined on æ E Rt, we write f(x) S g(x) or f(x) as ((g(z)) if there exists two constants c >> 0
such that f() S c. g(r), we write fx) x g(z) or f(x) = (z)()) if f(x S g(z) and g(x) S f(x).
A.1. Related Works
In-Context Learning (ICL). LLMs exhibit strong reasoning abilities, with their ICL capability playing a crucial role in their performance. Unlike fine-tuned models customized for specific tasks, LLMs demonstrate comparable capabilities by learning from informative prompts (Liu et al., 2021; Min et al., 2021; Nie et al., 2022). The theoretical understanding of ICL remains an active area of research. One line of research interprets ICL, as a form of Bayesian inference embedded within transformer architectures (Xie et al., 2021; Zhang et al., 2023; Jeon al., 2024; Hu et al., 2024). Another line of work focuses understanding how transformers internally emulate specific algorithms address ICL tasks (Garg et al., 2022; Nichani et al., 2024; Chen et al., 2024a; Fu et al., 2024; Sheen et al., 2024; Li et al., 2024b). Among these works, some focus on ICL for classification problems or learning with a finite dictionary (Sheen et al., 2024; Huang et al., 2023; Yang et al., 2024; Nichani et al., 2024), while another line of work examines how the attention model performs in-context linear regression (Von Oswald et al., 2023; Zhang et al., 2024; Chen et al., 2024a;c; Zhang et al., 2025). Meanwhile, multiple works demonstrate the remarkable ability of the model in feature learning (Kim & Suzuki, 2024; Yang et al., 2024; Huang et al., 2023) and decision-making (Sinii et al., 2023; Lin et al., 2023; He et al., 2024). Beyond these works, many researchers seek to uncover the internal ICL procedure of transformers, typically in more complex algorithms (Fu et al., 2023; Giannou et al., 2024; Cheng et al., 2023; Lin & Lee, 2024). Other works investigate the training process with multiple tasks (Kim et al., 2024; Tripuraneni et al., 2021), which are broadly related.
Comparison with Related Work. We provide a detailed discussion of the differences between our work and that of Chen al. (2024a) and Cui et al. (2024), which are among the most closely related studies.
Chen et al. (2024a) considers multi-head attention in the context of multi-task linear regression, where the number of heads matches the number of tasks. Under specialized initialization, they show that each head independently learns to solve a distinct in our framework. In contrast, our setup allows multiple heads to be flexibly allocated to a single task, enabling a more expressive and complex model architecture. As a result, while Chen al. (2024a) shows that the single-head model learns a nonparametric, kernel-type predictor with scaling of KQ parameters as 1/Vd. we demonstrate that the multi-head model instead learns parametric gradient descent predictor with KQ converging to 0+. This not only recovers the known results for linear attention (Zhang et al., 2024) but also reveals that multi-head softmax attention can outperform the single-head one by effectively encoding the linear architecture through an explicit approximation.
effectively reducing to a single-head per task setup, which corresponds to the single-head case
Cui et al. (2024) also identifies the diagonal KQ patterns with potentially positive and negative values in two-head softmax attention, and observe identical performance when the number of heads exceeds two. We go further by quantitatively characterizing the learned model. Specifically. we reveal detailed sign-matching. homogeneous KQ magnitudes, and zero-sum OV patterns for head counts beyond // 2. and show that multi-head softmax attention learns to implement a gradient descent predictor. From a theoretical perspective, Cui et al. (2024) adopts full-model parameterization and conducts a loss landscape analysis. In contrast, we begin by establishing an approximate loss and then develop a comprehensive explanation based on training dynamics, function approximation, and optimality analysis. Our results go beyond the core argument in Cui et al. (2024) regarding the superiority of multi-head over single-head attention: we not only compare the testing loss, but also explicitly demonstrate that single-head attention learns a nonparametric kernel regressor, while multi-head attention learns more powerful parametric gradient descent predictor.
A.2. Interpretation of Transformer Architecture
Interpretation of Attention Model. Note that (2.2) is a standard multi-head attention layer with a residual link. This function can be regarded as a mapping from a sequence of I + 1 vectors in RR++1 i.e., (ze]ee(c) U (zz), to new sequence
13

---

In In-Context Linear Regression with Multi-Head Softmax Attention
@
Znn
Zeda
Zees
smax + mak
Zeba
(a) token embeddling
Z
smax mak
z
K2* *
circuit consolidatio
A
(C) causal mask
Z
smax
z
(di (f. + 1, eew-entrreadddom
Figure Illustration the derivation from full multi-head attention architecture in (2.2) the simplified model. graph, we show the specified token embedding and read-out function with consolidated KQ and OV circuits. We use to denote the ineffective parameters due to read-out function or the 0 embedding for the query token.
of L + 1 vectors in Rd+1, In particular, each token position j. the query, key, and value in the lithh head are given by Q(4)zy, K(h)22, 2j. and V(h)2,, respectively. To get the output at position j, we compute the similarity between the query Q()22, and all the previous keys to get {(K(()) 2,, Q(*) 2g) hiss- These similarity scores are passed through the softmax function smax(-) to form probability distribution over all previous token positions [/ - I]. This distribution weights values (V(M) bij to compute the attention output. The outputs of all H heads are aggregated by the output matrices {Oth) heem Combining with the input Zebd from the residual link, we get the output in (2.2).
Note that when computing the output at each token position j. we only use the key Q(h),, 2, to look at the queries and values before position j. Thus, the softmax output dimension varies with j. forming probability distribution over [j for each j E [L + 1]. In the model (2.2), this enforced by the causal mask mak(-) before the softmax function. With the causal mask. only the first j - entries in the output vector of the softmax function nonzero. Compared to the standard decoder-only transformer model (e.g., Vaswani, 2017), we omit layer normalization and positional embeddings. Layer normalization is unnecessary because a single-layer architecture does not suffer from the issue of vanishing exploding gradients. In addition, the regression problem is permutation invariant {(T,, yr JEELL)- making positional information unnecessary.
B. Discussions and Extensions
The previous sections demonstrated that when training a multi-head attention model linear ICL data, the learned weight matrices exhibit significant patterns. These structured weights enable the model to implement debiased GD estimator. As we theoretically established, these patterns emerge from the interaction between the transformer architecture and the underlying data distribution. this section, we further investigate how different components contribute to this phenomenon. Specifically, we examine the role of the softmax function the attention model and the impact of isotropic covariate distribution. Additionally, we extend the multi-task in-context regression setting and study how the interplay between the number of tasks and the number of heads affects model behavior in our empirical study.
B.1. Comparison: Linear vs. Softmax Transformer
In this section, we study the connection between linear and softmax transformers. Proposed by Von Oswald al. (2023), linear transformer simplifies the architecture by removing the nonlinear activation and the causal mask, given by
1 Σ (hh Zera Zoba (h)) Q(h)) Zead.
LinTFo(Zebdl Zood +
(B.1)
k=1
Recent research has focused on understanding the inner mechanisms of transformers through the simplified linear trans-
14

---

In-Context Linear Regression with Multi-Heea Softmax Attention
formen(e.g.. Von Oswald et al., 2023). For the linear ICL task, Zhang et al. (2024) show that linear attention can be trained implement the one-step GID estimator with just one head. This raises the following question:
In linear ICL tasks, does softmax attention offer advantages over linear attention?
We argue that the softmax attention models are more advantageous than linear attention due to their enhanced expressive power. In fact, any H-head linear attention can be approximately implemented by a multi-head softmax attention with 2H heads, when the token embeddings are centralized. Specifically, consider 2Hah softmax attention model parameterized as
(-1 27
TFo ZZebd
Σ
O(h) Zebd smax amax((-1)7 Z K(h) Q(h) Zeebd);
h= (0,11)
where we the same parameters as in LinTFa(-) and set 7 a small rescaling constant. We define Zehd= L-- Zoba T. which averages the token embedding across the L token positions. Using similar approximation when 7 is close to zero, we have
TFo (Zeba) LinTF 8 (Zeba)
Σ (1) Zebd Zeba Q()) Zebd-
(B.2)
From (B.2), the output of softmax attention closely matches its linear counterpart when the second term is small, e.g., token embeddings are centralized such that Zebd R Og, which is the case in in-context linear regression.
Length Generalization. While softmax attentions use more heads, they benefit from dynamic normalization. Notice that the linear atten- tion in (B.1) requires a fixed normalization factor 1/L which depends on the sequence length L. In contrast, the parameters of softmax attention do not involve explicitly, but it effectively produces a I/L normalization factor when 7 sufficiently small, as shown in (B.2). The property of dynamic normalization enables softmax attention trained on linear ICL data to naturally generalize to different sequence lengths. During testing, a trained softmax model can process inputs with many more demonstration examples than it saw during training, and the ICL prediction error decreases as more examples are provided. In contrast, a trained linear attention model cannot generalize to longer sequences easily because its weights explicitly contain the normaliza- tion factor T/T. When the sequence length changes, linear attention would require different model parameters.
TF
Multi Head Softmax TF
Single Head Softmax
1.4
Loss
0.6
0.4
60
80
We conduct an experiment to test the length generalization ability of both model types. We train these models with T, = 40 demonstration examples and test them with L' examples, where L' ranges from 10 to 100. The ICL prediction errors in Figure 7 show that both single- and meat softmax attention models generalize effectively longer inputs, while the linear attention fails to do In particular, the prediction errors further decrease when the number of demonstration examples L' increases beyond L = 40. In contrast, the prediction error of linear attention starts to increase after T/' 40.
Prompt Length L
Figure 7. ICL prediction errors of linear attention, single-hcad and multi-hcad softmax attention trained on 40 but evaluated over different lengths. Note that softmax attention can generalize in length because the prediction error decays with more examples. In con- trast, linear attention cannot generalize in length and the prediction error starts to increase when the number of examples exceeds L.
B.2. A Alation Study: Alternative Activations Beyond Softmax
Building on our comparison between linear and softmax transformers, we now further examine the role played by the softmax activation function itself. Consider normalized activation _2d + d defined by letting (v)) f(4)/ Σ f(v)) for all i E [d]. where v e ddd is the input vector, 1/2 is the i-th entry of v/ and a(v); is the i-th entry of the output (v)). The function f RR. + R. can be any suitable univariate function, with softmax being the special case where f(-) = exp(.).
15

---

In-Context Linear Regression with Multi-Heea Softmax Attention
In $4.3, identified specific patterns in the KQ and OV circuits of trained softmax attention models. A key property promotinh these patterns is the exponential function underlying softmax. These patterns allow the learned transformer implement sum of kernel regressors where the exponential function serves as the kernll. Using the first-order approximation exp(z) s 1 + z, we proved that the model approximately implements the debiased GD. This raises questions:
(a) Suppose we use a different nonlinear activation in the attention, do we expect similar patterns the attention weights? (b) Does this transformer also implement debiased GD?
To answer these questions, we conduct additional experiments on the two-head attention model-using the same setup as $3-by replacing softmax with other normalized activations f that satisfy the first order approximation f(x) 1 Cy x for some constant C₁ > O. We test fi (z) = 1 Cx, fz (z) = (1 + Ca) and fz( + tanh(z) with C f0.5, 0.8, 17. Their first-order coefficients Cy are given by Cn C, Cf2 2C and Cf. respectively.
Head 1KKK
Hoad 1.0V
Read Lg
Fledd DVV
Heedd 1 OK
Head 2 CVV
221
220
Head OK
Head 2 DV
Head 2.05
Hoad
OK
Head CV
113
-15
(a) Activation f(x 1 +
(b) Activation f(x) (1 + 0.5z)².
(c) Activation f(x) + lanh(x).
Figure 8. Heatmap of KQ matrices and OV vectors of trained two-hcad attention using alternative activation functions. The trained models
consistently exhibit the same patterns shared by softmax attention, though the parameters of different models converge to different values.
Head
Postive Heads (2)
Head
Pasitive Heads
Positive Heads 2
Negative
Negative Heads
NegativeHesde Heads (1)
Saace
Shece
Steps
340
Zv Dynamia
Dyyammmss
OV Dyramis
com. Head
Head 1
Head
227
Head 2
Postive
Head 2
Sum Heads
Positive Heads (2)
Head
Heads
Positive Heads (2)
600
Negitive Heads
Negative Hes de (1)
335
Negative
(a) Activation f(z) 1
(b) Activation f(x) = (1 0.5zr)
(c) Activation f(z) + tanh(x).
Figure 9. Comparison of the training dynamics of KQ and OV parameters for two-head attention model with various activation functions. In particular, Figure (c) illustrates the dynamics of the case f(-) 1 tanh(-), which are closer those of the vanilla softmax attention model.
16

---

In-Context Linear Regression with Multi-Heed Softmax Attention
Activation function
1 Ca
(1 Cz))
exp(x)
+ tanh(x)
C = 0.5
0.8
C = 0.5
0.8
0.1267
0.1504
0.3677
0.2411
0.1979
0.2882
0.1810
0.1452
1
3.5006
3.3362
2.3963
2.2819
2.2221
1.7561
1.7487
1.7448
Y)eff 0.8871
1.0035
0.8810
0.8802
0.8794
1.0122
1.0128
1.0132
Table Comparison learned parameters with different activationss for two-head attention model. The table reports the learned (c)) 1, |. and the effective learning rate 7/of 2C3- | |/(*),, where Cý denotes the coefficient in the first-order Taylor expansion of
each activation function f. The exponential activation corresponds to the standard softmax attention. Despite different scaling patterns the individual parameters, the effective learning rates remain remarkably consistent across all activation functions, demonstrating that trained models implement similar debiased GD predictors regardless of activation choice.
- [1+ x)²
. 1 + tanh(x)
0.6
. 0.8x
0.7
(1 0.Bx)*
.a. SoftMax
10 0.5x
0.5x)²
. GD Est
0.5
+
GD Est
+ Deblased GD Est
...
Deblased GD Est
Debiased GD
0.5-
LOss
0.4
0.4
0.4
0.3
03
0.2
0.2
20
100
60
80
60 Prompt Length
Prompt Length
Prompt Length
Figure 10. Comparison of two-head attention models with different activation functions and the corresponding canonical predictors vanilla GD and dehiased GD predictor-with effective learning rate Vesr 2C,- - (10 1- |/ As shown in the figure, the ICL. prediction error attention models with different activations closely matches of the vanilla GD and debiased GD predictors. While the normalized quadratic activation exhibits slight deviations from the GD predictors for small L, the loss curves ultimately converge L increases.
In Figure 8, we plot the KQ and OV parameters for all these attention models. We observe the consistent positive-negative pattern in the attention weights. That is, the first two observations of $3 are still true. This implies that all of these learned attention models implement sum of kernel regressors. Moreover, by examining the magnitude of the KQ and OV parameters, we observe that the magnitude of w = (w))), (2) is still relatively small while the magnitude of μ = (u(²), 2²) is large. This is also consistent with the softmax attention. Mathematically, this means that the transformer predictors can be written as
0g 1101 (E) Ve f(ww(1)) (1 x xz)
Σ f--f((
f(-(ww(()) xt. (eTxx.) |w\ II xzg
Teff Σ yc o Zzg.
Ele f((((1)) Io. zzg)
Ele f
where Teff = 20,, | 11 and By is the debiased covariate and the approximation can be similarly derived as in Remark F.1.
In Table 1, we report the limiting values of (1)))), ( and 7jeff of these models. We observe that while the values of [ and [3) differ across different activations, the values of Netr are all close to one. This suggests that all these models effectively learn the same debiased GD predictor 07 (n"). Here 1) (1 (1 + o²) d/L)-1 x 1 when d/L 0.
Moreover, in Figure 10 we report the ICI, prediction errors of these learned multi-head attention models, together with the error of the debiased GD with corresponding effective learning rate. We observe that these eгror curves are highly consistent with debiased GD. Notice that all these models can readily generalize in length, which seems a benefit of the normalized activation attention. Thus, we expect that Theorem 4.2 can be generalized to attention models based on normalized activations in general, as long as f(x) + Cy around x = 0. Finally, we remark the choice of the bias term 1 is just
17

---

In-Context Linear Regression with Multi-Heea Softmax Attention
for simplicity, and the argument can be readily generalized other positive real numbers thanks to the normalization effect in softmax attention.
Finally, we plot the training dynamics of the KQ and OV parameters in Figure 9. Interestingly, the training dynamics across different activations exhibit slightly different behaviors. The behavior of + tanh(x) ) is similar to the softmax, where ( first increases and then decreases. The behavior of the other two cases in Figure 9 seems more complicated. We defer their study to future work.
In summary, show that for the class of multi-head attention models normalized activations, as long as the activation satisfies the first-order condition f() = 1 C5 I, the first three observations in $3 remain valid. That is, the attention weights share the patterns and the limiting attention models learn the same debiased GD predictor. However, the training dynamics are more subtle, which seem to rely on the choice of activation function.
B.3. Extension to Anisotropic Covariates
In the following, examine how the data distribution contributes to the observed patterns in the learned attention weights, To this end, we focus on the anisotropic case where the covariates are sampled from centered Gaussian distribution with a general covariance structure. In particular, we investigate the following questions:
(a) Can multi-head attention solve in-context linear regression with anisotropic covariates? (b) Do observe the same patterns in the learned transformer? Does the learned transformer implement GD algorithm approximately?
Head 1 OK
Head 1 OV
Inv. Covariance Matrix
611
155
1.0
035
055
1.0
-15
C115
Head QK
Head 2 OV
Optimal KQ (LinTF)
0.15
1.0
0.00
0.5
605
-10
-0.5
-1.0
-0.15
Figure Heatmap of KQ matrices and OV vectors trained on anisotropic covariates. The left two Figures illustrate patterns of trained two-head attention model, and the right Figures plot the inverse of covariance matrix E-1 and the conjectured optimal KQ pattern Σ-1 in (B.6). The learned KQ matrices exhibit a dominant tridiagonal structure with small values the (v, + 2) and (i + 2, positions all i E [d]. suggesting the model deviates from strict implementation of pre-conditioned GD using Σ-1 since the inverse covariance matrix should he tridiagonal. However, this pattern aligns with the conjectured structure --1 indicating the model implements a pre-conditioned predictor with structural adjustment.
In our experiment, we let the covariate distribution be P., = (0, E), where a C'ac-Murdcck-Szeee -Zegö matrix (Fikioris, 2018) with parameter = 0.5. That is, the -th entry of Σ is E,g = أ We train a two-head attention model with and 40.
18

---

In-Contest Linear Regression with Multi-Heea Softmax Attention
Transformer Learns Sum of Kermel Regressors. We observe that the diagonal pattern of the KQ circuits disappears, but the pattern of the OV circuits in Observation 1 in $3 persists. Moreover, as in Observation 2 in $3, the weight matrices of the two heads sum to zero, which means they also split into positive and a negative head. In particular, the learned weight matrices can be written as
(-1)
K Q(A)
O(h = 0 (-1) ml)
(B.3)
Here > O is small scaling parameter, Hoy is defined in (??), and f E R is positive definite matrix. That is, the first head is positive and the second is negative. The properties of sign-matching and zero-sum OV still hold. Moreover, the magnitude of the OV circuits is roughly the same as in the isotropic case. While KQ matrices still have a small magnitude, they are no longer proportional to an identity matrix. See the first column of Figure 11 for details.
As result, the learned transformer still implements a sum of two kernel regressors:
L (
y expl7 27 (hrigg
Ve . expe 22. (trigg
)
Dg = Myy
Σ
(B.4)
EL -xple 22 (.x,))
Ei=1 exp( -? x (
t=1
where the kernel function is induced by the bivariate function F(x, , (t) 1/7 exp(-7xpppp x T f2x'). Since y is small, we can similarly perform a first-order Taylor expansion to exp(-) in (B.4), which implies that By in (B.4) is close to pre-conditioned version of the debiased GD:
L
27 Mi
Fy
Σ
ye
S.gg
(B.5)
L
t=1
Thus, the effective learning rate, 17 = 27 . pyy is the same as in the isotropic case.
Fre-Conditioning Matrix. It remains determine the pre-conditioning matrix Q, Recall that in the isotropic case, we prove that multi-head softmax attention recovers the estimator found by the trained linear attention. For the anisotropic case, it is proved that linear attention finds pre-conditioned GD predictor (Ahn et al., 2023b), which is given by
yygd 9q =
Σ t=1
yl Fe
Fgg with D
(1 + T.
1
tr(2) + do²
la-
(B.6)
L
Besides, as we will show in $G.4, E corresponds to the optimal pre-conditioning matrix for the pre-conditioned GD, which minimizes the ICI. risk. We conjecture that Ω is close to Σー1, since (i) two-head softmax attention model and the linear attention when 7 small, and (ii) Σ enjoys optimality over all pre-conditioning matrices. In the right column of Figure 11, we plot Σ and Σ-!. A closer examination of Figure 11 shows that Ω is closer to É-1 than Σ-!
Note when T. is sufficiently large, Σ is close Σ. Thus, while we cannot prove Ω = Σ-1, we know that when L. is sufficiently large, the transformer predictor is equal to L- >/. ye xx I-1xg. Also see Figure 14 in SC.2, which shows that 9 Σ, and E are close.
B.4. Extension In-Context Multi-Task Regression
In the following, we consider the multi-task in-context regression, where the response variable a vector. In this setting, as we will show below, depending on the number of heads and the number of tasks, the learned transformer may exhibit different patterns.
Task Formulation. We first introduce the data generation process of multi-task linear regression as follows, where each task has its own set of features.
Definition B.1 (Multi-task Linear Model). Given d E Zt, we assume the covariate x E Rd is independently sampled from P.. and let E 2 be a fixed signal parameter. Let N E Z+ denote the number of tasks. For each task n E [N], let S., S [d] denote nonempty set of indices for task E [N]. Let 3s., and xs., denote the subvectors of 3 and indexed by Sn. We define the response vector y (y1.-..., UN]] NR- by letting yo. = 13, S.S. + en for all [N], where
6n rid. NN(,
19

---

In-Context Linear Regression with Multi-Heed Softmax Attention
In other words, each task n only uses features in Sn and the response is generated by a linear model with the corresponding signal parameter (ss, and covariate xs,,- Now we introduce the ICL. problem for multi-task regression, which is almost the same as in the definition in $2.
Definition B.2. Suppose that the signal set [SN] is fixed but unknown. To perform ICL, we first generate 3 P and then generate L demonstration examples (xr bet) where xt lid. P and Me is generated by the linear model in Definition B. with parameter B3. Moreover, we generate another covariate Ii ~ Px and the goal is to predict the response y4 RN for the query zq.
Note that for each sequence Zeba E (((+1)x(L+N) defined as in (2, the signal set {S,, [NE[N] is shared while the signal parameter 3 is randomly generated. Thus, we expect that the trained transformer model should be able to bake the shared structure } [N into the model. Chen et al. (2024a) study a special case of multi-task scenario with H = .N. They assume that, up to a rotation in y, each S, does not overlap with each other, and Une(N) S, = [d]. In this special case, each head learns to solve a separate linear task. That is, for each task there exists head h(n) such that the h(n)-tt )-th head solves the task n using the single-head attention estimator given in (??). Moreover, {f(n) [((n))n([]] ne[N] is permutation of [M]. Thus, in this special setup, the learned transformer model is a sum of N independent kernel regressors, where the n-th regressor is based on the S, -ubbeecoo of xy and the n-th entry of ye-
overlap with each other, which exhibits more complicated behaviors in the learned transformer model. In our experiment, we study simple two-task scenario with d = 6, L 40, a² = 0.1, and N 2. We set S {1, 2, 3, 4] and S2 {3, 4, 5, 6), which have two overlapping entries. For brevity, we defer the details of the experiment results and their interpretations to SC.3. We present the main findings as follows.
In contrast, we allow the signal sets (S,)neia
Global Pattern. In the multi-task setip. the KQ and OV matrices are (N + d) x (N + d) matrices. Similar to $3, we focus on the effective parameters that do not meet the zero vector in Zebd. i.e., the top-left (d+ + N) x d submatrix of KQ matrix and the last N rows of OV matrix:
KQ (k)
KQUN- =
KQ (h)
*
E (((++))((++.))
ov = Lov
ov (h)
E (
(B.7)
Here, (A E Bdd and Ov E RNN, Regardless of the number of heads, we observe the following consistent global pattern: There exists (h) E Rd and (b) e R such that
K KQ (h) = (h)
ovlh) diag(a(h)
(h) = OV K KQ
(h) 21 = ON
Vh E [+]
(B.8)
Therefore, the KQ circuit exhibits diagonal-only pattern in the top d submatrix with other entries equal to zero, and the OV circuit has non-zero values exclusively in the last N x N matrix. Following this, the trained transformer essentially behaves like a single-task model, where each head uses softmax function to compute similarity scores (h)
p(() smax(X- (1)) @ xxg) E R',, then produces a weighted response (x( Et Pe ye, where Pe the ê-th entry of (h), In other words, the learned transformer model is a sum of H kernel regressors, but the weights computed by the kernel are used to aggregated the vector-valued responses.
Local Patterns. However, different from the single task case, here each (h) does not corresponds to an all-one vector. Rather, w(th) can have different magnitudes in different entries, which reflects the tension of feature selections across different tasks. Intuitively, focusing on each attention head h, the value of (h) determines the similarity scores p(A), If we only use it to solve the first task, we expect that (h) only has zon-zero entries in Sy. Similarly, if we only use it to solve the second task, we expect that (h) only has non-zero entries in Sz. But when it contributes to solving both tasks, the problem is tricky. How much effort does this head put into solving task one and task two? This will be reflected in the magnitude of (h) in the overlapping entries.
We run experiment with multi-head attention models with Il ranging from 1 to 4. Depending on the number of heads, the trained transformer exhibit four different local patterns:
= 1: Single-head attention learns one weighted kernel regressor, assigning higher weights to the informative entries, i.e., the entries that are shared by more S,, 's. On our simple case, these entries are S, nS₁ = (3, 4).
20

---

In-Context Linear Regression with Multi-Heed Softmax Attention
(ii) H = 2: The two-head attention learns act as one weighted pre-conditioned GD predictor with also higher weight assigned to the informative entries.
(iii H 2 2N: The attention heads are clustered into N groups corresponding to N tasks, where each group contains at least 2 heads. Within each group, the attention heads learn to solve single task, say, task n, via the same mechanism as in the single-task setting. But the main difference is that these heads learns to use the true support S, to solve the task. Thus, this group of heads learns to implement debiased GD predictor based on data ( Vt. [ee where Tt,,,, is the subvector of zy indexed by Sn and Ve,n is the n-th entry of ye. In particular, the knowledge of { is encoded in the KQ circuit, which is reflected in the support of parameter (*) Overall, the transformer learns to implement N independent debiased GD predictors at the same time.
(iv) 2 < H < 2N: In this case, the learned pattern varies in repeated runs. In our experiment with H = 3 and N = 2, we observe that an interesting head that simultaneously serves as a positive head for task one and a negative head for task two. Interestingly, in this case, the model effectively implements N independent debiased GD predictors as in the previous case. But uses complicated superposition mechanism (Elhage et al., 2022) to achieve this.
C. Auxiliary Experiment Results
In this section, we provide additional experiment results support the discussions in $3 and $B.
C.1. Auxiliary Experiment Results for Section 3
Experiment Setups in $3. For all experiments the main paper, we set I. 40, 5, 0² 0.1 for data generation and employ the Adam optimizer with learning rate 7) 10-3, a batch size Tbbtch 256, and update for T = 5 x 10s training steps. The weight matrices are initialized using the default random initialization in PyTorch. Each experiment takes about 15 minutes on single NVIDIA A100 GPU. Furthermore, we found that the results are not sensitive to the choice of optimization algorithm, hyperparameters such as batch size and learning rate, or the configuration of the ICL data-generating process.
To back up the argument that multi-head attentions implement the same predictor, we consider the same setting as in $3 with various number of heads and visualize the training dynamics of the KQ and OV parameters. In particular, we consider H (2,3, 4), and set d 5, L 40, and 6² 0.1.
In Figure 12, we plot the training dynamics of the KQ and OV parameters for different numbers of heads. The first column shows ( [N and the second column shows (1l")) for different multi-head attention models. Our findings are consistent with theoretical results in I, 1 and can be summarized as follows:
(i) The patterns of the KQ and OV matrices are consistent across different numbers of heads, and are the same as in those described in Observation I and Observation 2 in $3. In particular, the property of homogeneous KQ scaling and zero-sum OV hold for all models.
(ii) More importantly, across different models, the limiting value of (1)'s and μ4 = =EhEN- (1())) are the same. Besides, - = -ENEH- ju(b) -14+. particular, for all non-dummy heads, the magnitude of ( ~ 0.13 for each model, and μ4 and | are both ~ 3.5.
(iii) When H > 2, we observe dummy heads, which corresponds to a head h E [H] such that ( ~ (a(l)) z 0.
By the observation (ii) above, we know that when H ranges from 2 to 4, the trained multi-head attention models all implements the same predictor as given in (??), which is a sum of two kernel regressors. Moreover, since the magnitude of the KQ parameters are small, as we show in (??), such a predictor is approximately equivalent to a debiased gradient descent predictor. Therefore, the experimental results support our theoretical findings in $B.1 that multi-head attentions learn to implement the same predictor.
Training Dynamic of Dummy Head. As a supplementary to Figure 4, which only plots the two-head cases, we provide the evolution of KQ and OV circuits for the four-head cases. Different from the two-head scenarios where the dummy heads will not occur due to the constraint of solution manifold S i.c., there should be at least one positive head and one negative head, there may occur dummy heads when H 2 3. In our experiment, we observe dummy head in the four-head model.
21

---

In- In-Context Linear Regression with Multi-Heea Softmax Attention
QK Dynamics
Dynamics
6:10
110
0:20
Head Head
Value the Last Elemert of
1.0
1.00
Head 1
Head 2
Positive Heads
Average of Diagnal Elemantss
0.23
6.00
Positive Heads (2)
2.00
0:3
-1.00
6.6-
Negative Heads (1)
Negative Heads (1)
6.30-
6.40
200x Stops
250k
300c
400k
-.00
506
100k
150c
2008 Stops
20
328c
2016
QK Dynamics
OV Dynamics
0.50
4.0
Aurorago of Diagnnal KCQ
Head
1.50
Head 1
Head
Head 2
Head
the Last Element
2.00
Head 3
Positive Heads (1&2)
0.20
Positive Heads (162)
1.00
Sum Head 1&2
6:0
2.00
Negative Heads (3)
0.20
Negative Heads (3)
-150
6.30
2008 Steps
Mikk
195a
-100
106k
199c
2004 Steps
OK Dynamics
OV Dynamics
Average of Diagiaal Boments KQ
Head
Head 1
0.20
Positive Heads (2)
Head
Volue of the Last Element of
Head 2
Positive Heads (2)
Head
2.00
Head
Dummy Heads (4)
Head
1.00
Head
Head 163
Dummy Heads (4)
3.00
6:0-
Negative Heads Headss (163)
Negative Heads (163)
1..0
1150
0.30
Steps
Figure 12. Comparison of training dynamics under the same setup with the number of heads ranging from 2 to The magnitude of the KQ parameters the same for all non-dummy heads in all models. The values of p+ and μ. also consistent across different models. This shows that all these multi-head attention models learn to implement the same predictor. Furthermore, in this experiment, we observe dummy head when H Dummy heads will appear when H 2.
We note that whether the dummy head appears and the number dummy heads are random and seem to depend on the initialization and the optimization algorithm. We also observe dummy heads for H 3 with different random seeds.
Furthermore, as shown in Figure 13, the KQ matrix of Head 4 demonstrates a chaotic pattern at the early stage and converges to zero-ne mato-x. In correspondence, the OV vectors of the dummy heads converge to zero vectors, ensuring the dummy heads will not contribute to the output. For positive and negative heads, the behavior is similar to two-head cases. These observations can also be seen from the last row of Figure 12.
C.2. Additional Results for Anisotropic Covariates in Section B.3
In the following, we provide additional details of the experiments on anisotropic covariates, which is introduced in $B.3.
A key property of the Kac-Murdock-Szegó matrix is that its inverse matrix a tridiagonal matrix. Let A = p²). --1, then the entries of A are given by
Kac-Murdock-Szegj matrix
A = (1 + p²) 1(i * and d) + li 1 or d), A A = =Pp. for aolal i [d].
22

---

In-Context Linear Regression with Multi-Head Softmax Attention
Step
Step
100K
500k
Head KQ
Step
500k
Step
1k
Step
100k 500k
Hese OV
Head
Step
Step 1k
Step
Step 10k
Step 100k Step 500k
Step
Step 1k
Step 2k
Step 10k
Step 100k Step 500k
Head
40200
4000
4300
Figure 13. Heatmaps of the KQ matrices and OV vectors along training epochs with H = 4, d= =, and = 40, trained over 7 = 5 10 steps with random initialization. The behaviors of the first three heads similar the two-head matrix of Head 4 (dummy head) demonstrates chaotic pattern the early stage and converges to zero caro-matrix. The OV vectors of the dummy heads converge zero vectors.
shown Figure 4. The KQ
The rest of the entries are zero. See, e.g. Theorem 2.1 in Fikioris (2018) for more details. We use this property examine the value of the KQ matrices learned by the attention model.
KQ Implements Pre-Conditioned GD with Ω E-1 We show (B.5) that multi-head attention learns implement pre-conditioned version of the debaised GD, where Q is the pre-conditioning matrix. Here 7 Ω appears in the KQ matrix of the positive head. In Figure 14, we plot the inverse of the top-left d x submatrices of the KQ matrices, and compare them with Σ and Σ. This figures shows that these matrices are all very close, which is the case when T, is large.
In addition, since the E-1 a tridiagonal matrix, its (1, i 2)- and (i (i+3) 2, 2,2,)-entrii i)-entiees are zero. As shown Figure 11, not only have significant values in the tridiagonal entries but also exhibit very small values the (i, i + and (i+ 2, e)-entries. This indicates that the trained model does not strictly implement pre-conditioned GD using E-1, In comparison, the adjusted
23

---

In-Contes Linear Regression with Multi-Heea Softmax Attention
Head 1 Inv. QK
Head 2 Inv. QK
10
10
2-
-0
3-
-5
3-
Covariance Matrix
Inv. Optimal (LinTF)
1.00
0.75
0.5
2-
0.00
0.0
-0.25
-0.00
-1.00
Figure 14. Comparison of inverse KQ matrices trained anisotropic covariates with the ground-truth covariance matrix V, and the inverse matrix of conjectured optimal KQ pattern Σ in (B.6). The learned KQ matrices exhibit an positive-negative pattern, similar to the isotropic case $3. Additionally, inverse KQ matrix nearly proportional ground-truth covariance matrix and Σ, suggesting that the model implements pre-conditioned GD predictor.
E-1 (see bottom-right of Figure 11), also exhibits small values beyond the tridiagonal entries due to the perturbation of tr(E)/L- Id. By comparing KQ matrices and Σ-!, we can show that the KQ matrices of two heads take the form Σ and -7 Σ with small scaling y, aligning with the conjecture in $B.1. Meanwhile, note that the difference between Σ- and Đ-1 is roughly O), which is negligible when L is large. Thus, it suffices to identify the estimator in (B.6) with simple pre-conditioned GD predictor L E We x -EE-aa xg in the low-dimensional regime where d/ L
C.3. Additional Results for Multi-Task Linear Regression in Section $B.4
In the following, we introduce the details of the experiments on multi-task linear regression, which discussed in $B.4.
Experimental Setups. Recall that the data generation process is defined in Definitions B.1 and (B.2). We let {( vr) denote the ICL examples, where Vy Idd and ye E N. We consider the isotropic case with P,, = =N Ia) and P = (0, I/dd. Moreover, in our experiment, we focus on the two-task case with N 2, d 6, L 40, o² 0.1. Moreover, we set S (1, 2,3, 4) and S2 {3, 5,5, 6), i.c., the features of the tasks have overlap {3, 4]. We train multi-head attention models with H E 21, 4).
The rest of the experimental setup is the same that in $3.
To simplify the notation, we let S* S n S2 and S S" Besides, for each
[N]. we let ye, denote the n-th entry of the response vector For any subset S C [d]. we let Fe.s denote the subvector of xe with indices in S. As discussed in $B.3, the trained transformer exhibits global patterns in the KQ and OV matrices. specific, by writing the KQ and OV matrices as block matrices as (B.7), we have (B.8) for all h E [H] In other words, each head captured by parameters (h) E d and (M) RR, As we will show follows, each four attention models illustrate one of the four possible modes of behavior respectively.
24

---

In-Context Linear Regression with Multi-Heea Softmax Attention
Head OK
Head CV
Mode (H 1): Single Weighted Kernel Regressor. In Figure 15, visualize the learned pattern of single-head attention. Upon closer examination, the third and fourth entries, corresponding to Si S2,, have a larger value (r~ 0.5) than other diagonal entries (~~ 0.42). That (1) (1) is, w (1) ~ 0.42 for j $ {3, 4] and wis w 0.5. In the general case, have
-12
(1) ws-- W. 1|
Wge of 118 w jí 1N.
(1)
Figure 15. Hcatmap KQ and OV matrices trained on multiple tasks with H = 1. The trained attention model learns implement a single weighted kernel regressor.
where W*, & and ji are three scaling factors.
This indicates that, the trained single-attention model learns a weighted kernel regressor. Specifically, for each task n E [N], the predictor is
L exp(#* (or,s (g.St) + 8 (FI.se.Ig.SS)). ye,
Ei= expl@* " + w
which applies uses a same set of weights to aggregate responses for all tasks.
Mode (H 2): Single Weighted Pre-Conditioned GD Predictor. As shown in Figure 16, the two-head attention learns opposing positive-negative pattern which is nearly identical to the single-task to Mode I, the attention model assigns more weight on the informative entries S (~ 0.13 vs. 0.12). In addition, the learned scale is much smaller than that in the previous single-head case, which is also observed in the single-task setting, The pattern learned by the attention model can be summarized as below:
observed in $3. However, similar
(1) wggh = - 25' 1..11. s.l. wge-
- w 1|s0// 1 = -2 = 1.N.
where 8","" 8 and ji are three scaling factors. Therefore, when = 2, for each task n E [N], the predictor approximately takes the following form:
2j0i0*
L
2j
L (Bsrs ye,
Fg.n
Σ
(5
ven
Here use the fact that W is small. Thus, the trained attention model solves the multi-task linear regression using weighted debiased GD predictor, which assigns slightly larger weight to S
Head OK
Head GV
Head OK
Head QV
412
435
-033
-015
Figure 16. Heatmap of KQ and OV matrices trained on multi-task problem with H the trained model, the KQ matrices exhibits diagonal-only pattern and the diagonal values have two magnitudes. The magnitude of {w/))) (3,43,4) is slightly larger than the rests of the diagonal entries. Besides, (?)) - and they proportional all-one vector.
25

---

In-Context Linear Regression with Multi-Heea Softmax Attention
Mode III (H 2N): N Independent Debiased GD Predictors. Recall for single task, at least two heads are required implement debiased GD predictor. Here, we explore regime where H 2 2N. Thus, the transformer has the expressive power use two attention heads to solve each task independently, via implementing debiased GD predictors
As shown in Figure 17, the experiment result four-head model supports this hypothesis. In the trained model, Heads 1 and are coupled to solve Task 1, focusing exclusively on S in KQ while contributing only to the prediction of Ug. This
be seen by noting that w| - (2) + s. (1) -w 0 (s
Pi - -m1 já, and n12 μ12
(2) 0.
(h) denotes the 2-th entry of ( Similarly, Heads and are coupled to solve Task 2, following the same structure. In the general case, consider H 2N with (2n 1)-th and 2n-th
Here St and j are two positive scaling parameters and fA
heads solving the n-th task together. For all [A]. the learned pattern satisfies
(2m) S,
3 1s
(2n-1) w Ss
(2m) w Sc
Op5gl
n (2n-1) = = - pp(2n) = it en
(2n)
where we use en E R to denote a canonical basis whose n-th entry is one and the rest of the entries are all zero. With this pattern, using the fact that à is small, for each task [M], the predictor approximately takes the following form:
1.
Dg.n
2pa L
Σ
St. Iq.S..) De,r
(C.1)
where 2jã = 1. This implements an independent GD predictor for each task n, using , Me In other words, the model bakes the true supports (Sn)nE(H) into the learned model weights.
Note that Mode and Mode II, the same weights are applied across different tasks, relying on all covariate entries. In contrast, when the model has sufficient expressive power, it learns allocate dedicated heads for each task and efficiently solve each task by using debiased gradient descent predictor that is restricted the true support S,, of the task.
OK
Head
OK
Head
Head OK
Head OV
Head OK
Head OV
0.22
-015
0.25
0.5
3
0.59
0.59
-032
Figure 17. Heatmap of KQ and OV matrices trained on the multi-task problem with H = The trained model learns to implement N independent pre-conditioned GD predictors by assigning each task to unique pair of heads, ensuring that each head is exclusively used for single task. Task-rrlevant entry information is encoded in the KQ circuit and OV controls the output task.
Mode IV (2 < H < 2N): Pre-Conditioned GD Predictors with Superposition. Unlike Mode and Mode II, where attention employs shared weights across all tasks, or Mode III, where each task assigned distinct, non-overlapping heads,
26

---

In-Context Linear Regression with Multi-Heea Softmax Attention
when 2 H 2N, observe an intriguing superposition phenomenon. Intuitively, superposition requires requires some heads solve move than task simultaneously, which means that may observe w()) that have both positive and negative values. Superposition enables the transformer model to approximate the performance of the N debiased GD predictors despite the limited head capacity.
Head QK
Head OV
HEAD
HEAD 2
HEAD 3
1.75
(h)
-0.0316
0.0779
-0,0837
1.50
1055
-0.0689
0.0000
0.0927
5.000
w
-0.0893
0,0870
0,0152
-3.9911
6.9204
-2.8739
P22
-7.0742
2.3371
4.7399
Head OK
Head 3 OK
Head 3 OV
0.015
0.0
0.533
0.00
00008
-0.
-0.0
-0.159
-0.02
-0999
-0.08
Figure 18. (Top left) Learned parameters of three-head attention model trained on multi-task problem. By direct calculation, we can show that learned parameter follow the patterns described (C.5) (C.6), indicating that the model effectively emulates the pre-conditioned GD predictor despite its limited expressive power via superposition. (Remaining) Heatmap of KQ and OV matrices trained on multi-task problem with H for Task 2 and negative head for Task 1. The trained model learns implement the N independent debiased GD predictors with superposition.
this case, each head contributes to tasks, and in particular, Head positive head
Figure 18 shows the learned pattern of three-head model, trained on the two-task in-context linear regression problem. We have the following observations:
(i) All the three heads contribute to solving both regression tasks. This can be seen from the fact that (h) does have entries for all h.
(ii) Head is negative head and Head 2 is positive head, (2) (1) are negative, while (0(2), (2) are positive.
(iii) Head has both positive and negative components, illustrating the superposition phenomenon. Head 3 is a negative head for Task and a positive head for Task 2. The first two entries of (3) are negative, the last two entries are positive, and the middle two entries are close to
(iv) Head put more effort in solving Task 2, as #2 (1) has a larger magnitude than pi Similarly, Head 2 focus more on Task 1. Head is rather more balanced.
These observations lead mathematical depiction of (a) ( he(H) as follows. Let denote the set of all possible grouped divisions {SS* Si..: } with S S, S* for n E (1. Under our setup, we have Sf [1,2], S (5. 6). and S* {3, 4). The parameters of the three-head model be summarized follows:
H (h) = ON,
w
(h) 11s
v(h,S) E [H] 2.
h=
27

---

In-Context Linear Regression with Multi-Heed Softmax Attention
In other words, the supports of the tasks split [d] into non-overlapping groups, and the value of the KQ parameters
are the same within each group. Here, wh) denotes the subvector of w() with entries in S and ( is a scaling factor. Moreover, the OV parameters sum to a zero vector, generalizes single-task setting. Let ji(h)) denote the learned value of p(h), Then Head h contributes to the prediction of Task n via the parameter pnn), and we have E ((h)) 0 for all n.
By plugging the values of (w(t), ,e [he(m] into the transformer model, for each task n E [M], its output can be written as
exp( whth) O xxg)) Ve,
Dgn
(C.2)
t=1 E exp((x w(h) o xz)))
where (1)) 3 3.g denotes the Hadamard product. By the construction of 9, we have (0,, (hh) o xzg)
Σ
(h) (FL,S, Tg,))
(C.3)
SE
When L is sufficiently large, we combine (C.3) and apply Taylor expansion (C.2) to obtain that
H
L
In.n
23
(h)
Me.
1+
Σ
20 (Te, Fnss))
h=1
SE9
(1) (1e
Σ
St (n ye,ns
(C.4)
I.
1=18E9
the last equality follows from the fact that Σ Mn = 0. That is, Jyn can be viewed as a sum of [ independent debiased
GD predictors, each of which restricted to the support S.
Furthermore, we carefully examine the learned values of (w( y (h) ( JeELHn which are listed in Table 18. We observe that the following conditions hold:
"
ST.
Σ
(1) (h)
S5.
=0,
(C.5)
h=1
H
H
#
S
(N)
#1
Σ
St
(h)
O.
(C.6)
h=1
h=1
h=1
As result, Jg. in (C.4) can be simplified as
ye,rs
Vn E [N].
This estimator is identical to that in (C.1). This means that, when the same estimator sum of two independent debiased GD estimator as the case when Hf 2.N. But here the transformer model achieves by exploiting the superposition. That is, each head is used to solve both tasks, and the aggregated effect is that the model can approximate the performance of the N independent debiased GD predictors. Here we only study a toy case with H- = 3 and N = 2. We believe that the superposition phenomenon is general phenomenon that occurs when 2 < H < 2N. But more a larger N, the superposition phenomenon becomes more complex and entangled. For the general multi-task setting, characterizing the training dynamics and solution manifold of the multi-head attention with superposition an important direction for future work.
3 and N the trained transformer model learns
C.4. Experimental Results for Loss Approximation
In the following, we examine the difference of the true population loss of ICL and the approximate loss in (4.1) in terms of the training dynamics.
28

---

In-Context Linear Regression with Multi-Heed Softmax Attention
Intuitions of Approximation and Scaling in Proposition E.I. In the proof of loss approximation, the main rationale is that the high-order terms are negligible when the prompt length I. is sufficiently large under certain conditions. For fixed d, the intuition is that p = smax(w Xxxg) should be distributed relatively "uniformly" across the entries, resulting in quite small norm lpll2 for k 2 2. As shown in Figure 19a, with fixed d, the norm concentrates around 0 as the length L increases. However, this property may not hold in the high-dimensional proportional regime, i.e., d/ [. - & (see Figure 19b). To address this issue, we impose scaling conditions ||all & log L/d to offset the effects of high dimensionality and ensure that the norm remains small (see Figure 19c).
16-
L-40, d=5
4- d-5
L= d=5
L-100, d-5
L- 200, d-25
L-200, d-25
L-500, d=5
800, d-100
12
L- 8-00, d-100
6-
10.
Frequency
Frequency
Frequency
8-
1-
0.6
0.6
0.2
5.6
0.0
Norm
Norm
(a) Low--dmenssion, 0.5.
(b) High-dimension, w = 0.5.
(c) High-dimension, w = 0.35.
Figure 19. Distribution of norm ||pll. Figure (a) depicts the low-dimensional regime with a fixed d = 5. As the sequence length increases, the distribution of lIpll converges to singleton at zero. Figure (b) and (c) illustrate the proportional regime d/L 1/8. As increases and the scaling factor w decreases, the norm |pl| becomes increasingly concentrated around zero.
Training Dynamics of True Loss and Approximate Loss. To further evaluate the performance of approximate loss in (4.1), we compare the training dynamics given by these two different losses. Consider the four-head attention model with 5. L 40, o² 0.1. By comparing Figures 20a and 20b, observe that the approximate loss effectively captures the evolution of parameters seen in the full model. Additionally, we note that the actual training dynamics appear less "regular" due to the finite batch size (as opposed to the population risk used in the approximate loss) and internal perturbations arising from parameters at non-target positions in the full model.
D. Analysis on Training Dynamics of Two-head Attention
In this section, we offer a more comprehensive analysis of the gradient flow on the approximate loss defined in (4.1) with a two-head attention model. This analysis provides justification for the evolution of the parameters (w)()) ((()))())))))) ((()) in Figure 1b. For simplicity, we consider the gradient flow of C/2 with a two-head attention, where E is the approximate loss function. We let @ denote the vector of all parameters, i.e., 0 (n, 12,, where μ (nL²), (2) E IR² and w (2()),((1))) IR... Consider the gradient flow algorithm. We let 6e denote the parameter at time t. The gradient flow dynamics is characterized by an ordinary differential equation (ODE):
0,0: = (eC(()))
where ((0) = 2(0)/2.
Here we slightly abuse the notation by writing E/2 as L, which is the case only in this section. For simplification, consider the following initialization.
Definition D.1 (Initialization). Consider two-head attention model with symmetric initialization such that μ (a, -a) where a > 0 is sufficiently small constant defining the initial scale.
Under initialization in Definition D.I, gradient flow starts from a point satisfying sign-matching and homogencous scaling (see $4.2). In the following, we show that these patterns, once emerged, are preserved along the gradient flow trajectory. This aligns with the experimental results shown in Figure 1b, where we train the full parameters of the attention model from random initialization.
29

---

In-Context Linear Regression with Multi-Heea Softmax Attention
QK Dynamics
Asproximated KQ Dynamics
Head
Head
Head
Head
Average Daggnal Elements
Head
0.23
0.38
Head
0.02
0.2
0.80
-0.2
0:33
0.20
28
138
37C
52t
Steps
15%
93c
1128
1336
150k
2000
4300
Steps
1000
803:
13000
OV Dynamics
Approximated OV Dynamics
3.52
Head
Head
Head
Headd
Head
2,00
Head
Head
Heads
Sum Grouped Heads
0.02
the
3.00
3.00
3350
26
14K
30x
10x
63c
3128
173K
1508
2003
4020
6200
3000
13000
Steps
Steps
(a) Training Dynamics of Attention Weights.
Approximate Training Dynamics.
Figure 20, Comparison of training dynamics between the true loss and the approximate loss for four-head attention model. The approximate training dynamics closely capture behavior of the true dynamics, providing effective characterization of the whole training process.
Preservation of Symmetry and Homogeneous Scaling. Based on gradient calculations and the definition of gradient flow, the time derivative parameters satisfies that
0efl: = we M4 We we A exp(dwiw,) ) pe,
Oewr pe M4 We fe dd (exp(dwrw)) o (Hip) ))
where define (1 d²)/L for notational simplicity. Let exp(d (w()))))) exp(d (w( and u (We, 1), (We, 11, 82) E RR3 By direct calculation and careful decomposition, we have
Oeuur 1) = (1 - u1 wi) (wr, 1) - \ exp(d wi (1) (/(2)) (uz, 1)
expp( (w())))) (Me, 1) + X Mi St = (gr. u.),
(D.1)
(w 1) (1 u wt) (pt, 1) dd m4 1 pi (2) exp(d w( (?) (wi, 1) dd wi (1) (14) 1)
dd )) exp(d (u(2))2) (wo (2) exp(d (2)))) (14 pt (1) (Pr: 1) (g2,2, Me).
dA w(
(D.2)
Here observe that the dime-derivatives of both (14, 1) and (wr, 1) are linear functions of u4, and let I1.t and 92. denote the linear coefficients, respectively. Furthermore, it is easy to see that
0,s: = 2d exp(d (d (2) w arw 2d exp(d- (w(2))2) 2 Opu) (2)
2d (1))) 8,m! st + 2d exp(d (w) (u())))) d.us (wi, 2d exp(d (w) (2) ( ae(ur, 1) (ga,2 we).
(D.3)
30

---

In-Context Linear Regression with Multi-Heed Softmax Attention
In other words, the time-derivative of 8 is also a linear function of uy, with coefficient given by gs.t- Here the last equality uses the fact that can write 0r 1) = (9 Ue) as shown in (D.2). We note that 91,t, 92.t, and 9a, are functions of up-
To show that ( = - and w()) = (2) are preserved along the gradient flow, note that based on (D.1), (D.2) and (D.3), we can characterize the dynamics of uy with matrix ODE d.ue- = Gr where G, = [g 92. 93.1 E R By solving the ODE, we have u, = w exp (Jo G,ds), where w is a constant vector that reflects the initial condition, i.e. 0. Therefore, u, 0 at any time t. In (2) wi + we = 0. See Figure 21a for an illustration, which plot the full evolution particular, this implies that u + 14 of μl and w.
Mo w. Note that the initialization in Definition D.1 indicates that ug
(2)
Hence, track the dynamics of we) by only focusing on pe (1) and w!") For notational clarity, we let y and o denote Me and ( respectively. Then the evolution of (ftt (t) is given by the following ODE system:
0,g. 0t - - Xxpr (exp(dei) - exp(-dei))
(D.4)
0.e. = 4p= - - l: - d'ppler (exp(de²) +
Such an ODE system does not admit a closed-fore form solution. but can be solved numerically. To gain some insight, in Figures 21a and 21b, we plot the full dynamics of ju and w with L = 40, d = 5 and a = 0.1, as well as the dynamics in the early stage of gradient flow. In addition, in Figure 21c we plot the evolution of the loss (0,) as a function of t, as well as the ratio spr/ (:. In Figure 21d, we plot the dynamics of sp:, together with yp: @: and y0 (et). Here (e:) denotes the optimal value of Pt that minimizes the approximate loss function, when de fixed. See (D.17) for its definition. In addition, we additionally plot the dynamics of μ and w in the case with d 10 in Figure 21. We observe almost identical behaviors as in the case with d = 5.
By examining these figures, we observe that the evolution of the ODE system in (D.4) exhibits the following three phases:
Phase I (Exponential and Synchronous Growth). From small injtialization, both for and e grow exponentially fast in t at nearly the same rate. This leads to rapid reduction in the loss E. Moreover, during this stage, yo stays roughly equal to @t, as shown in Figure 21b. In particular, their ratio /x/ Q: stays close to one, which is shown in Figure 21c.
Phase II (Slowed Growth and Peak Formation). As s and or increase, their growth rate in time decreases, i.e., the exponential growth in both Ω and @ stops. Moreover, d eventually reaches its maximum value (d2-1/2). Although the parameters no longer grow exponentially, the loss decreases sharply until d: begins to decline. Moreover, the ratio sp// /: does not grow much from one. This phase ends when @: attains its maximum value, which corresponds a critical of the ODE in (D.4), under the condition that pp: Pt- As shown in Figure 21b. by the end of this phase, (t and o: part ways, with Pt keeping increasing while @t begins to decrease. Throughout this phase, the ratio s:/ Qt does not increase much from one, shown in Figure 21c. In addition to the sharp drop of the loss, the product yo: 2t increases rapidly, as shown in Figure 21d.
Phase III (Convergence). After & reaches its peak, it decreases zero while Sr continues increasing. Their ratio (Pt/ et thus keeps increasing. In addition, the loss converges to the minimum value and the value of the product yr: 2 increases and converges. Moreover, the last phase exhibits an interesting phenomenon: the difference between fPr and (ex) gradually becomes negligible. This motivates us to use (et) as a surrogate of ye for the analysis of the limiting behavior of the ODE system.
In the following, we analyze the ODE system defined in (D.4) under the high- high-dimensional regime where d/L + & and d L + oo. Recall that we define x = (1 +
(1+02)/L /²) T., and thus dd + (1 + o²) E. Here we additionally assume that & is a small constant such that the limit of d' is less than one.
Phase I: Exponential and Synchronous Growth. At the beginning of gradient flow, both Φe and Qt quickly escape from the starting point, maintaining an almost identical growth rate while keeping their magnitude small. During the first stage, &P e: holds due to the same small initialization (see Figure 21b for an illustration). In this case, we have exp(de?) - expp-ddř? = O(e?) and exp(dee) + exp(- Rs 2. Ignoring the high-order terms of st and et, we can approximate (D.4) by a much simpler system of differential equations:
O:pt Rd Ct,
0.g: R Cc:
(D.5)
31

---

In- In-Context Linear Regression with Multi-Heea Softmax Attention
(4)
0.
11)
12
(2
(1)
0.4
(1)
Value
((2
Value
0.-
-044
08
8000
10030
Time
Time
(a) Full dynamics.
(h) Early-stage dynamics.
16-
2.00
(1)) /100)
120
15
((0)
10
3..
0.2
4p
(g)
3.25
6.2
232
6.75
6.-
6.35
0.2
Time
15
Time
20
Loss C(O,) and ratio sp/ et-
(d) Product fee le and yp (ec).
Figure 21. Gradient flow with respect the approximated loss Ž/2 with L 40, d = = and 0² 0.1 under initialization in Definition D.I with initialization a 0.001. Figure (a) shows the full evolution of the four-dimensional dynamic system terms of and along the gradient flow. Figure (b) provides closer view of the dynamics during the early stage. with colors highlighting the three phases. Specifically, g: first reaches its maximum (~~ 0.57) which on the order of O( /lolg d/d) (~ 0.56) and you keeps increasing. Figure (c) and (d) present the evolution of loss (())), spx/ gt, Prer, and 1p" (@:) to track the relative behavior of parameters. Here yp (e:) is defined in (D.17).
shaded in different
whose solution given by
(t: R a J eads xo a exp(().
@t R eep(((),
(D.6)
Thus, both for and d: grows exponentially in I.
Moreover, despite the exponential growth, the ratio between Pp and remains at the initial value, i.c., po/e0 for pretty long time, resulting in a synchronous growth at the initial stage. To see this, note that we have
d log(pp//:)) = Ôe log st Ô log e: 1/4: 0,e: - 1/er- d.et
Lt Pt - 2( vt) - A (exp(de?) exp(-def))
+ dXip? (exp(def) + exp(-dgi) ~ 2d) 5. > 0.
(D.7)
Here the last approximation step is derived from the facts that = 0. and that exp(de?) cxp(- do3) = 2. As a result, the ratio p:/e increases in t. Moreover, recall that we define (1 + o2)/L, which is small when large. As result, when is small, the ratio sp/ de increases in at rather slow rate. More specifically, by solving the ODE above, yt/ lt can
32

---

In- In-Context Linear Regression with Multi-Heea Softmax Attention
(4)
14 (1)
12
(2)
0.50
(1)
(4)
0.2
Value
((2
Value
(2)
0.25
-0:50
8000
10030
Time
Time
(a) Full dynamics.
(h) Early-stage dynamics.
2.08
11 /w001
22
0.0
375
Z(0)
0.5
2.-
(0)
1.55
0.20
300
6.75
0.30
0.0
6.5
6.05
(c) Loss and ratio.
(d) Product and ppT (2c).
Figure 22. Gradient flow with respect to the approximated loss Z/2 with = 40, 10, H = 2 0² 0.1 under initialization Definition D.I with initialization and w along the gradient flow. Figure (b) provides a closer view the dynamics during the early stage, with areas shaded different colors highlighting the three phases. Specifically, g first reaches its maximum (~~ 0.45) which the order of O( O(ylog d/d) (~ 0.47) and yo keeps increasing. Figure (c) and (d) present the evolution of loss, si/ Dt. Pelt, and sot to track the relative behavior of parameters.
0.001. Figure (a) shows the full evolution of the four-dimensional dynamic system terms of
be approximately written as a function of Pt, i.e.,
Pt ~= exp
2dl / o.ds) xe exp 2dda2, [' 'eexp(2sdss) exp (dxa² exp(2t)) = + dda? exp(2t) = 1 + dd o.
Qt
(xp(2))
(D.8)
Here, the second and the last approximation step uses (D.6), and the third approximation step uses the fact that t is small. As shown in Figures 21b and 21c, in the first phase, while both &rr and increase in I, their ratio remains substantially close to one when 0:.
In addition, Figure 21c also shows that in the first phase, the loss started decrease alongside the rapid growth parameters. To show this rigorously, we compute the dimevderivative of the loss function. Let L, denote L(0.), which is given by
C. [())) (1 + 03) 2pree + 22p2et + Xxp (exp(de?) exp(-de?))
Then we have
O,0, = -2(1 - 24pr0t) (pp: 0:2: + o: 0(pp)) + 2Xip: Dept (exp(de))eep((ddlle - exx((deo))
+ 4Xdipe e: Oee: (exp(de?) + exp(-de;)
(D.9)
~ 40² (2:tet- 1) + SAd o et -4ei 4a² exp(2t).
33

---

In-Context Linear Regression with Multi-Heed Softmax Attention
Here, the first approximation is based on (D.5) and the facts that Pt @t and that exp(def) exp(-def) 2. In the second approximation, we only keep the dominating term, -4gh and remove the high-order terms that are negligible. In the last equality, we plug in the closed-form of Pr in (D.6). Hence, solving this differential equation, we have
L R Lo - 4 1." a² exp(22)) Lo 2
where Co [())) is the initial value of the loss. That is, the loss decreases exponentially in t. However, since the magnitude of 2: remains small despite its exponential growth, the loss function does not decrease significantly during this phase.
In summary, during Phase I, sp and d grow exponentially while keeping their ratio close to one. In addition, the change in both the loss L, and the ratio spr/ @: is proportional to e². indicating exponential changes in both quantities. The end of Phase is defined by
a. ma Rt min 0x) s /d-1/2/22
(D.10)
when 9o and de is sufficiently large such that approximations for exp(dgī). exp(- ) and the high-order terms are no longer valid. Hence, the loss has decreased at least by (2d)-1 and &t lt is approximately bounded by + A/4 1+ + 1/L during the first phase.
Phase II: Slowed Growth and Peak Formation. In the second phase, we focus on the evolution after time 2 T1, which defined in (D.10). In this phase, we cannot ignore the high-order terms of you and @t. In particular, we use the first-order approximation (xp(d exp(-dg; x 2do7. Moreover, We still adopt the approximation that yo 22 Qc- To see this, note that (D.7) now becomes
a log(ver/:)) R 2dd et + dhip (2+ (2+ = dd :s.
(D.11)
Here we plug in @t fp and the second-order approximation that exp(de)) + exp( (dp(- 2 de et. Similar to (D.8), we can solve (D.11) and conclude that the ratio sPt/ De remains close to one, as long as Stt and et remains small.
Now we focus on the dynamics of fp and or. By (D.4), the dynamics of (pt can be simplified to
O.pe R a: - 2 - Xup: . 2dge R 4p: 2(d2 + 1) Pi:
(D.12)
where the approximations follow from exp(dg?) exp(-de? 2dg and po: et- Solving this differential equation, for all for all time t during this phase (t 2 TT), we have
St R {2(d) + 1) + <(TI) exp(e +())} 1/2 1//2 with ((t) a exp( 2(ds + 1).
(D.13)
Note that here we directly solve the differential equation 0:p:: = S 2(dA+ 1) . with initialization fr, = a: . (xp((1) in closed-form. Based on (D.13), we see that the growth rate of yor decelerated compared with the exponential growth in Phase I. But the increase in the magnitude of yo is much faster thanks to the much larger initial value a exp(()) d 1/2 //22//2/ > a.
Note that d shares similar behavior since their ratio remains around one given its small growing rate. In particular, the dynamics of et can be approximately characterized by
0.00 == Pp - 2. d'sfer (2+ + et) R 0: - 2(d 1) o² d= o7.
(D.14)
where we use the second-order approximation exp(deř) exp(-dgř + d2 and sp Qr. Furthermore, to better understand the relative behavior between (et and et, we can caleulate the time derivative of their difference p 2c. In particular, combining 12) and (D.14), we have
dep: 0.e: RJ d
of > 0,
(D.15)
where we use the fact that (pe R Qt. Hence, the gap between e and @t increase in t and thus approximation for R Qe would be violated by the end of this phase. Note that with the increasing of @t. depi and dee are quickly decreasing due to the exponential growth of exp(dee) in (D.4). Based on (D.15), we can see that d.ee < S.sp, such that & will first reach the critical point, i.c., Oegt= = 0. which marks the end of Phase II. Define
T2 = max E RR+ Sel:: (dA.exp(det + dd . exp(p + 22 S 11).
34

---

In-Contest Linear Regression with Multi-Heea Softmax Attention
To characterize the scale of et by the end of the second phase, given Dt R Ppt, we have
2 e (dA exp(de + dd exp(=dee + 2) 2 e exp(de
(D.16)
since we assume dA (1 d2) { is of constant order. Moreover, note that na characterizes the time when e hit its peak value since et increases monotonically during the first stages. And as will show later, @: continues to decrease
thereafter. Thus, conclude that
max Q: - lrz -r2= -0:=0 (Vollod log loglllg//// d)/d) Ö(1/Va), KER!
by solving the inequality in (D.16) given sufficiently large d. More specifically, we solve the equation dA of -ex((dd²,) = 1 and use the expansion of Lambert-W W function. Here Ö(-) omits logarithmic factors. Therefore, we prove that the value of O: forms a peak, with magnitude of O(1/Vd). This aligns with the observation in Figure 21b.
Phase III Convergence Next, we provide heuristic analysis of the last phase showing that decays to zero asymp- totically. In particular, we employ a two-timescale analysis by first pushing fr: to its limit, enabling us to eliminate you by treating it as a function of o;. Recall that
f = (1+ 2pter + 24700 xup (cxp(de)) -
Note that the loss is quadratic with respect to pr. By finding the critical point of La with respect to yr, we know that the limiting value for for a fixed o: is given by
pf = (0 = 2eî + > (exp(dee) (xp()
(D.17)
Moreover, shown in Figure 21c, in this phase, the evolution of Pp and yi are almost identical. Motivated by this observation, to characterize the limiting behavior of lt, we replace yot by oot in the ODE in (D.4). This leads to an ODE only involving d::
de = 1p (1- (2 dd (exp((de)) +
p.e.).
(D.18)
Using first-order approximations
exp(det) - exp(-det) x 2dg2, and
exp(det) + exp(- 22 - det
in (D.17) and (D.18), we have p0; ~ (2(d) 1) e:)-1 and
{
+ dd (2 + et)
d²
dee: ~ 2(d) 1) 2:
o..
2(dA 1)
4(dA + 1)2
Solving this ODE in closed form, we can characterize the evolution of et and ipt (et) by
RJ 1 & (1 (1-12) m2)
-1/2
(1 1 d (1. T2)
1/2
~ 2(d
Vt E (02,
112
- 2(d) 1)
As mentioned above, when t is sufficiently large, we regard that Poī is almost equal pt. Furthermore, we know that of z log d/d.. d. Under the high-dimensional regime where d/L. - &, we know that dA o2) ) < 1, which can be regarded as constant. Thus, @ decreases and converges to 0+ at rate of (((dvi)-1), i.e., loo 0+. While şr increases at rate of (dv7) when / increases, and thus grows to infinity. Moreover, while & increases, the product 24:::: remains almost constant and converges to 1/(1 +11+ (1+o²)- E) during this phase.
In summary, through a heuristic two-timescale analysis, we prove that @: decreases to zero after reaching its peak. And yo: increases to infinity while maintaining Seett at a constant level.
35

---

In-Context Linear Regression with Multi-Heed Softmax Attention
E. Implications of Learned Pattern of QK and OV Circuits
E.1. Transformer Model as a Sum of Kernel Regressors
Recall that we use smax(-) denote the softmax function and use Zebd E (R((+I)((L+)) denote [ zy]. where Z contains the first L demonstration samples and 2 = (x] O)T contains the test input. Here the L demonstration samples are generated from noisy linear model with parameter which itself is drawn from distribution Pg. The transformer output TF (Zeba) is specified in (2.2), which is (d + 1) (I. + 1--dimensional matrix. Finally, the prediction head outputs the (d + 1, L + 1) entry of TFa (Zaad), denoted as Ng- which is used to predict the desired response 3 zg. Notice that the (d 1, L + 1)-th entry of Zeba is the (d 1)-th entry of 2g. which is equal to zero. Due to causal masking, 2 does not attend to itself and only attend to the previous tokens, namely Z. Thus, we can simplify the (T. + 1)-th column of TF (Zedd) to
H oo smax (z kth)" Q(h) E R((+1))
(E.1)
Compared with (2.2), we replace Zebd with Z in the value part of attention, because causal masking ensures that the output of the softmax function is a probability distribution over [L]. Then, Dy is the last entry of the vector in (E.1). Under the simplified parametrization given in (3.1). we have
z K (h))" Q(h1 = (M) Xxg E R
Oll) Z - n 1h) y ERL
e [H].
Thus, we can simplify Bg to the following form:
1
Σ (h))) (V., smxx (w(h). Xxrg ) = Σ (h) p()))
0g
(E.2)
h=1
where we define ych) smax (w)t). Xx.) for all h [H). Here, ((h) and (h) y are the attention score and value of the h-th head at position (T. + 1), respectively.
Interpretation of (E.2) as a Sum of Kernel Regressors. Consider the h-th head, where (d)) is probability distribution over [L., and the probability assigned to each is proportional to exp(w(h) xx ag). Then, the output of the 1-th head is equal to
explw xz Øg)
Σ ye (1
where 8(zg,
(E.3)
Σt= explw xx] sig)
Here R(Zy, w) is kernel function that captures the similarity of n4 and and the parameter 1/w plays the role of bandwidth of the kernel. Notice that the kernel A in (E.3) is slightly different from the standard Gaussian radial basis function (RBF) kernel, which is defined as exp(-w- ||=r, - zell3), where w is a parameter. These two kernels coincide when when Fg and zt are on the unit sphere. However, in our case, Pa is supported on IRd and thus these two kernels are different. Nevertheless, each term in (E.3) is still Vadaraya-Watson kernel regression predictor. Moreover, the intuition of such an estimator is clear: the output in (E.3) is weighted sum of the responses [eele in the demonstration data, and the weights are determined by the similarity between the test input 2, and the demonstration input zg. In summary, when the parameters of the transformer are given by (3.1), the prediction Bg is a sum of H kemne regressors.
Interpretation of Heali--HaadAll Attention Predictor with (w, C Ly, As shown in the empirical observations in $3 and the solution manifold in $4.2, parameters (ww(()) pu (ee))) found by the transformer model are in the solution manifold Sy defined in (4.5) for some y > 0. Thus, Dy in (E.2) can be simplified as sum of two kernel regressors:
0g Σn My :i (K(mg, - A(zr 22; 7)))
(E.4)
The intuition behind such estimator is as follows. The kernel R(z 7) assigns larger weight to z, when a. xy is large, i.e., when I. and I, are similar. The kernel A(Zy We; -)) is large when x. xr is small, i.e., when z. and xy are dissimilar. Thus, the predictor in (E.4) is a sum of two kernel regression predictors based on datasets {ze, Be heeee and {-xz, - Ve hee[r) respectively. It is also reasonable why both these two kernel regressors appear. Specifically, (Iz, yr) and (-ze, -yt) have the same distribution when It ~ (O, Ia) and ye is generated from a linear model. Thus, we can use both (re, Ve hete and (-xe, - ye ]ee(z) to learn the parameter (3.
36

---

In-Context Linear Regression with Multi-Heea Softmax Attention
Difference between Single-Head and Multi-Head Softmax Attention. As shown in Chen et al. (2024a), when H 1, the learned transformer implements a Nadaraya- Watson kernel regressor in (E.3), with [w(1) 1/vd. Moreover, the KQ and OV matrices of the trained single-head model exhibit positive negative only pattern. In contrast, when H 2 2, the multi-head attention mechanism leams to behave as a sum of kernel regressors, achieving better performance due to its greater expressive power. Furthermore,the trained multi-head model displays a coupled positive-negative pattern and works a regime with a significantly smaller magnitudes of KQ parameters. Therefore, despite the form of the sum of kernel regressors, the multi-head attention indeed learns to approximate a variant of the debiased GD predictor, which better captures the linear structure of the problem compared the single-head case (see $4.3 for details).
E,2. Proof Proposition 4.1: Loss Approximation
For better understanding the approximation and scaling conditions in Proposition E.I, we provide a more detailed discussion alongside the simulation results in $C.4.
Proposition E.1 (Formal Statement of Proposition 4.1). Under the parametrization in (3.1), consider the H-head attention
model with dimension d E z and sample size L E Z Suppose that the parameters (w, p)) S I2" satisfy one of the
following conditions:
|| S 10.1.llgg L/ max (d, log L}.
max{Ilall 11111/0) S L --//2+3/10
where À > O is an absolute constant defining the error level. Then, it holds that L(w, p)) + o² 2μ'
T w + n (ww"" + 02) L- exx((dww)) +) L->)
Here (-) omits some universal constant.
The informal statement in Proposition 4.1 follows by assuming d > log L and setting À 1/5. A more refined trade- off between the scales of w and j can be achieved by carefully selecting k, the scaling factor for w, i.e. ||w||oo S kn log <y/log L/ max dd log LJ, where k is now set to a constant 0.1.
Proof of Proposition E.I. We prove this proposition in three steps:
In Part 1, we first expand the population loss L(w, pi) into multiple terms, where each term involves expectations over the randomness of both coefficient 3 ~ P/3 and ICL samples Zebd-
In Part 2, we simplify these terms by conditioning on Ty and taking an expectation over the randomness of the rest 3 and { ye) be(r)) within the context sequence, which gives us reformulation of the loss in terms of || and moments of the attention probabilities.
In Part 3, we isolate the high-order moments from the low-order moments in the reformulated loss. We show that low- order terms can be well approximated by a polynomial of OV weights and exponential QK weights up to O(dH L-~>) error, and high-order terms can also be uniformly bounded by OddH2 L->) under appropriate scaling conditions.
STEP 1. EXPAND POPULATION Loss INTO MULTIPLE TERMS.
Recall that under the simplified parametrization given in (3.1), the prediction of the transformer model, Dy. is given by (E.2), where we let p^A) smax xx,)) be the attention score of the h-th head. Then, the population loss defined in (2.3)
can be expanded as follows:
H BT (TT2 Σ
27 Tp(t))
[(w, = (0)
(b)
(XB+c)
+ +02
4=1
E (B xz )
22 3 Tg
Σ
((0) (X3 + c) p(()
(X3
h=1
(ii
E
(1 (1
(XB c) pb) (XB + c) T2(0))
+o2.
(E.5)
M=1
M=1
(iii
37

---

In-Context Linear Regression with Multi-Heed Softmax Attention
We let e = (f1....,
Lid. ........... fL) E RL with & ~ (00 o²) be the noise terms in (yehee(L)- Notice that (d)) E IR_ and X E Rexd Recall that |3 ~ (0, Ia/d) and Fg ~ (0, Ia), then we have
i E [88 Zq )
E b. 88 x0] =
(E.6)
For the remaining two terms, we can simplify using a similar argument to marginalize 3 ~ NO,, Ia/d). Specifically, for term (ii), we have
H
(ii)
B Eg
3
T
x
p(()
hml
(h)
(1) E
BB x pmh)
BB
Σ
EE
I
T
x
p21)))
(E.7)
For term (iii), using the independence between and (X; Vig. B), we have
(iii)
(h
(
E pous (X33 x + + DON
H
H
H
(b) (0h')
.E OO' X33TX x p
(h) Holl' E spors EE pp
h=
H
H
(h) (ul1))
H
E pon XXT p())
+0²
1e (1) (h') (a
E 100000 you]
(E.8)
d
h=1 M==
h=
M=1
where the last equality holds because EEle"" 1 1 and [[ =aa/d /d.
In the following, to further simplify the expectation terms in (E.7) and (E.8), we decouple the randomness of X E R.xd and F. E TRd, where (are keeL, and Ty are i.i.d. random vectors sampled from (0, Ta). To this end, we take conditional expectations with respect to X given Zgg.
STEP 2. TAKE CONDITIONAL EXPECTATIONS GIVEN Zig.
Recall that p(A) = (1) Xx.)). the following, for any l E [L.) we let Pé and (smax(w(()) Xx,))e denote the l-th entry of the attention probability (b)) Also note that X == [x= .I.] with at iid. NO, Ta). For any random vector
(h)
x ~ (O, La) and any differentiable functiin g. Stein's 's Lemma states that [[xg(()] = E [V)]( By this lemma, for all
[H). holds that
E[XT (1)
| szg
E B[xxe (s Xa.g)le [xg]] =
Σ 1=1
E Uze ze (smax(w Xxrg)le [
(h) Fgg B[p (h) (p! (+) [g]] = (1)) Fgg w(N) 2gg E (1 xgg]
(E.9)
Here, the second equality follows from Stein's 's Lemma, the third equality follows from the derivative of the softmax function,
and the last equality follows from the fact that DL= Pt (h) 1. Combining (E.9) and (E.7), conditioning the value of zg.-
we can write the second term in (E.5) as
H
Σ
/d. ||0,122 (
Σ
( 11,11//d (///dEEllp())) "1Il"") xg).
(E.
where we use (ii) x to denote the counterpart of (ii) with vg fixed. It remains to simplify the third term in (E.5). In particular, we need to handle terms of the form Elp(') xX T (h) myy. To this end, we introduce the following lemma.
38

---

In-Context Linear Regression with Multi-Head Softmax Attention
Lemma E.2. Consider random variable X = [x1,..., [E] E RLdd with Dg ~
N(0, Id) for all & E [T]. Let w, e R 1t
and U. E R be fixed parameters. We define smax(w Xv) and p- = smax( Xu). Then, it holds
EF Xx = d...EEp
Elp' 7] + w - ت5 |101| IB[(1 - (2l22 (1 10:1
+ 2w 110112 E1-p"p + 50 p: ||e|5] + -55 1| Bi-p po2 + p + p. 1P
+w. 36 ||0122 Fly p - pT pTp02= p" oe + ( m.
Here p0 = p p denotes the element-wise square of E R
Proof of Lemma E.2. See $G.1 for detailed proof.
To compute Elp(>'' EXTp(M) N.l. we apply Lemma E.2 with w(()), i3 = =00,, and e Ig. Then, for any
h, h' E [H], we obtain that
Elpla> XX p(h)] xg4
= d. Elp(h) 1 p (14)) [ (ww))) 110o1:: TB((1 - (1 (1) - -10(0011)) xzg
+ 2() |, E[- El p(a)02 + p(N)T pia). (1 11 xzg
+ 2(w / 11 El El-p(t)) (h/y002 + ph(h pih'. 1p""">11 xzgl
+ (() wh')) || (DD p(h') (h" p000 -p00" (()) (D) p(00) zgl
= = d Blp(") phh)) zzg + w (w) || 12 + | (Fgg)
(E.11)
where we use (x.g) to denote the high-order terms in (E.11). Here the high-order terms contain the product of | and expectations of polynomials of ((h) and p(k') with degrees at least two.
Combine (E.5), (E.6), (E.7) and (E.8) and the form of conditional expectations in (E.9), (E.10), and (E.11), write the conditional expectation of the loss function given Tig into polynomial of (1r*) )eea)) Specifically, we have
H
2
L/x, (w, 10) := :=E
Σ
(hi y p(())
Fgg
h=1
2 12
H
=1 1 o²
(1) (1)
h=1
|
H
# (1) 2e (1) (1+ 0²)
r (1)
Elp(N) p0l') 1 xz
==
1t=1
2 1r 1
|
H
H
(
Elllp(A) x.y]
( (h)))
(n.))
(E.12)
Notice that Lis, (w, (1) above is function of Eq.. In the following, we analyze the expected value of each term in (E.12).
STEP 3. MARGINALIZE Iy AND ELIMINATIONS OF HIGH-ORDER TERMS.
Let = (1))) direct calculation, we have
(T)) E 2 and w = (1) (H)) E HR Since I.g ~ (O, Ta), we have ETI = d. By
H (()
H
H
E
Σ
= 1 w.
E
(h) (u (w)
(1 (W)
(E.13)
h=1k=1
In the following, we show that ((p p p(h0) with large can be approximated by exp(d (())) (Step 3.1), and the remaining terms are mostly high-order and thus can ignored (Step 3.2).
39

---

In-Context Linear Regression with Multi-Heea Softmax Attention
Step 3.1. Approximate Elp(h) p(O] under Large L. under large I. This enables control the term
Next, we approximate Elp p(h?1 by exp(d w())(())
H (N)
(1+
ΣE
1h0 ( h=1W==1
( p(a)
of the loss [(w,/ pi), given in (E.12). Our derivation is based the following insight. When w(l) is small, for p(), approximate the denominator in the softmax function by L.. Viewing p(h) function of wh)), the first two terms in the Taylor expansion is:
(1) Pt = T. . w()) xe xg)) +1+
E [T..
Here, we approximate the denominator Ee= exp(w(() xz xg) in the softmax function by We can similarly write down the Taylor expansion for p(h) in terms of (h)) Thus, we have
Blp p(u)) [
1
B + (h)) xx xg) ) (1 + (h) xx xzg) [gg
L
= = (1 +1++0 (1) (w)) 11=g112)
This implies the following approximation scheme:
Blp(1)) (h') [x) ~ --1 (1 + w (w') x xTTz))
-- exp(w xx xg).
(E.14)
Similarly, by taking expectation with respect to By in (E.14), we have
Elp(h) D( ~ L- (1 (1)) B[r, (lll -- (1 (1 + (h) x L- exp(d w(l)))
In the following, we rigorously prove the above approximation. We prove by leveraging a lemma from Chen et al. (2024a) quantify the accuracy of the approximation in (F.14) and the large deviation of Tg- We first present the lemma below.
Lemma E.3. Let Tq E Rd be any fixed vector that satisfies max x l/ | ag 13 S t log L, where > 0 constant and W E RRddd We E IRdxd are two parameter matrices. Let = smax XW W.,)) and p = smax( (XWxx) where
X == [2,...., T E Lxd with xe itd (O, Ia), holds that
(i). Elp P ng1 - L- exp x W WR vx,)
((1-2(1-))):
(i)). max [Epp T 202 xgl. Ellp|| 12 p p xgl. B((p" 2) xz11] O(L-2(1-1)))
where e vT/2+ 3/ 1//T) When Ig ~ NN(00 and | S log L/ max (d, log L}, we further have
(ii). E (EpT [zg] --1 exp (r, W WTw,,,)) [x (3.--(3)))
Note all these three arguments are at population level and thus hold with probability one.
Proof of Lemma E.3. See Lemma B.2 and Lemma B.3 in Chen al. (2024a) for detailed proofs. The first items in Lemma B.2, and the third item is in Lemma B.3.
The first item in Lemma E.3 quantifies the error of the approximation in (E.14). The second item shows that the high-order terms of p and p are small. Additionally, the third item quantifies the squared error of the approximation given x, ~ (0, I2). Furthermore, we note that the approximation errors in Lemma E.3 depends on e, which depends on the parameter 5. To make the approximation errors small, we need to choose small 7 such that the resulting is small.
40

---

In-Context Linear Regression with Multi-Heea Softmax Attention
To apply this lemma pl) and pd²), we set W w La and W w(h) Ia. To satisfies the requirement that zy bounded in the lemma, control the tail behavior of || using the concentration results for x By setting == log L in Lemma G.I, we have
P (1 > 5 log L}) < P || > d 2 dlog L + 2 2log L) SL'
(E.15)
In the following, we assume that E H satisfies the following condition:
(C1) ||olo S kV <v/lo L/ max/d, log L},
where x denotes parameter to be determined later. Based on any w satisfying (C1), define the following good event where ( are all bounded:
Egood
{vh E [M), ( 2g12 s 5x2 log I.}.
(E.16)
Then, when ws satisfies (C1), by (E.15) we have
P(Egood) P (Ilwll | > 5n log L) S (1xgg12 5 maxfd, log L}) s L-
Besides, when Egood holds, the premise of Lemma E.3 holds with 5 5k². Now we are ready to work on Elp(h) p(O)
We decompose the ([p(()) plk0) as follows:
(l)(h) pO =E(ph) pON' 1 (Egood)] E [p
(M) p(h') (Egoodd)
S = p(N). 1 (Egood) L-1
(E.17)
where the last inequality follows from the fact that p(a)" and ((h') are probability distributions over [Z], and thus p(h)" p(h) < 1. Furthermore, we decompose the first term on the RHS (E.17) as
E (p p(h') (Egood) E (DONT p(h') --1 (w |,, 1 (Egood)
(iv)
+ - (W) |0 1 (Egood) exp(d w(l)
(v)
+ L- exp(d (1)))
(E.18)
We control term (iv) by applying Lemma E.3-(i) with 5x², which implies that there exists and absolute constant C such that
(iv)
<
E E (p(h)
phh?= - --1 exp(w 10 (10,)111) 1 (Egood)
Zg
]]
C1 L 201-0)
(E.19)
Moreover, here is function of specified constant x that controls the scaling of ||w/|oo. defined as
e = 5/2 3/(1 + V 1 (5m2)-1).
(E.20)
To handle term (v), we write
(v) L- -e.L-llexppp exp(d (()) (())))) (E| ( / d)) (Egood) - -1)
(110,113
Let denote the random variable (xp(w(() (h) (h')
|2 11 - d)) 1 (Egood), which depends on Ty only. Motivated by (E.15),
we consider the cases where I ([]lr,, 113 2 log + 2 log T.) true or not separately. Note that when this event true, we have
T < exp(kw(()) (2 dog L + 2og L)).
(E.21)

---

In-Context Linear Regression with Multi-Heed Softmax Attention
When this event does not hold, we have
exp(d w w(()) T= exp(w(l) ||0)1))
1 (Egood) S exp(5n log L):
(E.22)
thanks to the event Egood (w) defined in (E.16). Therefore, we conclude that
(v) = -- 1 exp(d (d w (h)) E -1]
-- exp(d w (1) (h)) EJr 1 (112011 d 2/d log L)
+ 1. exp(d (h)) B|T. .1(1]=l - d> 2 dlo I. + 2 log L.]
<L exp (
(w (2vdlog L 2og L)) - 1)
+ L. + -
exp (5k² log L.)
(E.23)
Here, the last inequality follows from (E.21) and (E.22), Note that when (C1) is true, we have
(())) (2v d + 2 L) S x²² log L/ max log L} (2vdlog + 2 log L) S 4x log L, exp(d (ild-w(h)) (h)) v exp( (K² log L/ (d, log L} d) - L
Thus, in this case, have By substituting the arguments above back into (E.23), obtain that
(v) S L 1+xx
+ L -+45k2
< 21 2L-1++
(E.24)
Combining (E.,7), (8. (E.19) and (F.24), have
El (0)) L exp(d (1)))))))) + C L L-2(1-e) + 21 2222-1465 145-1 + ---
= L exp(d (()) + C1 L -21-rr) + 3L 1456²
where we define a = 2//5/2+3 3/(1 + v + 1/ which is a function of the specified constant x, and C1 an absolute constant coming from (E.19). Thus, we conclude that
H
n (h) (K))
((())
El D(NO)
A= h'=l
=1
H
(1) (h') exp(d ((())) 111130 H² (G L ---((1)) 3T. 11--2))
(E.25)
h=1M=1
This enables us to bound the corresponding term in the loss function L(ws, 4), as given (E.
Step 3.2. Bound E [|
| r.y]] under Large I.. Next, we bound the term
1 2
Hf Σ
E
d
in the decomposition of the loss C(w,//) 14). Recall that || ~ X2. Thus, using the second moment of x2: we have 1 [ d (d + 2). Now, applying the Cauchy-Schwand Schwartz inequality, we have
h
H p (w))))
E ||اا ll Bllld =01]
v
[[ E (B]Ip( x²1)
d
w
(Vie) [B[ (B[ll
[.y] L-1 exp exp(du(()))))) -1 (k)
42

---

In-Context Linear Regression with Multi-Heea Softmax Attention
Here, in the last inequality, we use Recall that & b to denote that C. b for some absolute constant C 0. By applying Lemma E.3-(iii) with 5x², we obtain the following inequality:
E V a) la|, where X is any random variable and is constant.
(Ellp((
1| x11] - L-1 ())()) S C2 --3(1-))
where C2 is an absolute constant and is defined in (F.20). In addition, under the scaling condition (C1), for all [H] we have exp exp(dw (dw) 2 L= Therefore, we conclude that
Σ a (
wh) (h). E |[ Elle (0 12
xmll] z |alloo ||oloo dH (L (L-13-c)/2 + L --1t))
S |Allo dH (L- + L --1+m²)
(E.26)
where the last inequality follows from the fact that x S 1.
Step 3.3. Bound High-order Terms 11 Ja, ( (22g) as in (E.11), which comes from the decomposition of Epp("' XX T p(()) [g). As shown in (E.11), Ta, (xg) is composed of terms such as
under large L. Recall that we define high-order terms Th, (ziy)
Elp 2o2 xgl.
Elppl p P xgl. and
E[ 2) xgle=
where we denote {p. p) = (p(h) (0') Also notice that, by (E.11), each term in Th.n (xxg) involves has multiplicative factors (2, or (2) or (h)) Following a similar argument as in Step 3.2, we apply Cauchy-Schie Schwartz incquality to bound the high-order terms as
H ΣΣ
H En (
(b)
(N))
E1121111 Tnw S d. (N)
( (xng) )) xzy) xxy))²]
(E.27)
A=1
h=1
Then we apply Lemma E.3-(i)) with 5k2 the terms in (E.27), which implies that
[TTR, (Ir,) x I C3 |wll L- 2(1-)) L -
where is defined in (E.20). By (E.27), we have
"
ΣΣ
(1)
E11 (x)) S 11 dH2 L 2(1--c)
(E.28)
h=lk==
This inequality establishes an upper bound on the last term in (E.
STEP COMBINE EVERYTHING AND CONCLUDE THE PROOF.
Now we take the expectation with respect to Tyg (E. and get
H
2E xz
(N) (()) L
((w. =1 0²
M==
(E.13)
(E.
2 or
+ +(++0)
(h) (h'). (lp p(ay)++
(k) (M)
E
EE
(
1
d
(F.25)
(F.26)
H
(1) (h)
Jn (xa))
(E.28)
43

---

In-Context Linear Regression with Multi-Heed Softmax Attention
Here list the number of inequalities that bound each term. Combining (B.13), (B.25), (E.26) and (E.28), we have
H
C(w, p)) =1+ o² 22 w (H w +
Σ
# (h) u (h)) exp(d w()) + Err,
=1+ 0² 24 w + 1 Ww + (1+ -- exp(duiw 1) Err..
where we error term Err, is bounded via
Err, & 11m1180 H². (I,-2(1-e) + ---1+5 + |1al| dH (I.-3-e))/2 + I, I--
+ 1allo dH² L - -(1(-))
<
11 dH² max L -- 2(1-e) L (3(3e))22 L-
(E.29)
Here é is defined in (F.20) and x is a small constant. By taking K. = 0.1 such that e ~ 0.69 - 0.7, the approximation error
can be further simplified as Err, & max(]1al|o,,, [elll//ll)) dlf². L 1-2(1-6) I max || 11 dHf². L L-3/5, Hence,
for given error level) > 0, if we assume y E R satisfies
(C2)
max (llallss [I $ I, 2/243/33/10
then we can control the error by Err, OddH2 ......--)) I->>). This completes the proof.
F. Proof of Theorem 4.2
In this section, we prove Theorem 4.2, the main result of this paper. We present separate proofs for each part of the theorem in the following subsections. Recall that we use A = (w. p) E R to parameterize the multi-head attention model with H heads, where the weight matrices satisfy (3.1). Furthermore, recall that we define the parameter spaces e and 9 (4.5) and (4.6), respectively, with 5 C L.. We define ΣICR, (1()) = 14 and ENER 4 = μ to simplify the notation. Also recall that we define the debiased GD predictor with learning 77 as
7 T
(1)
ye Ie Fgg
with
Ee = 8t
T.
Wt.
(F.I)
Note that here we use Ee denote the centered covariate.
In the following, present proof of Theorem 4.2. In $F.1, show that the multi-head attention model approximate the debiased GD predictor with a small error. In $F.2, we show that minimizing the approximate loss in (4.1) over the parameter space s leads to a multi-head attention model that converges to the debiased GD predictor. In $F.3, we show that the debiased GD predictor is asymptotically Bayes optimal up to a proportionality factor.
F.1. Proof of Theorem 4.2-(i): Approximation
Proof of Theorem 4.2-(i). In this proof, we regard the learning rate as constant. Let >> 0 be scaling parameter and let ja be defined such that 2/17- = 1/. Since we assumed that μ44 _ j with já > 0 and (w, E L, the multi-head attention estimator takes the form
expl1 2x7 xg)) Ele exp(1 - 27 xgg)
Ye exp( -p 2 xzg) Ei=1 exp( - x7 xzg)
(0) = -.
(F.2)
and the debiased GD predictor is defined in (F.1) with 2yū. Following this, we the decompose difference between transformer (0) and debiased GD @ed (1) as
# 1
(8)
T (n))
1
Σ
(exply) 2o xzg) -expl-y- 27 xzg)= 27 27 xz)) 4
L
Em1
44

---

In-Context Linear Regression with Multi-Heed Softmax Attention
L El exp(y Te rg)
I.
L. Ei=l exp(y 2/ [9))
exp(r - 2z (g) y1 +
L
Σ
Iq
Ve
E=1
t=1
f== (xpppxy-111 x4) -
L
exp(-7- kn)) ye
+
Iq
Me
L Ei= exp( (- x. Øg)
(i + (ii)
(ii..
(F.3)
Here we define the three terms as (i), (ii) and (iii) to simplify the notation. Before delving into details, we first define three auxiliary functions over (7, x) E I22 as below:
$ Φ) exp()>) exp( 72) 2yyz, 2b (7, = exp(1x) 1 -yE, (1,) exp(1x)
By Taylor expansion, we can obtain that d1 (7, and d2(7, x) are both (??2x²) when y is small, and vb(7,z)) O(w)). We use Φ1 to bound term (i) in (F.3) as follows:
Ilulle di -11(1 max [o] zz)) VL
Ilvlla vL
Σ
(1 x7 xg)= ve S max |0 (7, 27 2g)1- EE[L]
(F.4)
EE[L)
Here the inequality results from the Seuchh-Sehwarrd inequality and the last equality follows from the fact that (1 -) is an even function and monotonically increasing for y > 0. Next we consider the term (ii). To simplify the notation, we define Sexp= = Ele exp(7 2. xgg) and then we have
E-22,,0T mg)
Σ
1 1 x x/ Fgg exply 1.p(7.xfI.1)) x.u)) I. +
Σ
T D/ Fgg 2=1
Sexp-
(F.5)
E=1
Based on the definition term (ii) in (F.3), we have
(ii) =
I SSxp
exp(r - 2 xx)== y8
+
7
4
x
Fq Vt
L Sexp
L
1=1
Σ 22 2T xzg) 7
Σ
Σ explp 1t xZxz)) xq) y1
Σ
zz Zz
z
Fz y
L Sexp
L
8=1
t=1
22 (7, x 200)
Σ
exp(y x1 xzg) ye
L Sexp
t=1
E xt zgg
exp(7 2h xq) ye
+ T
x xq
ye-
(F.6)
L
8exp
t=1
Here, the first equality follows from (F.5). To simplify the notation, we detine
(ii.a) ( 1=1
( exp(7 [g)g)= we). (ii.b) =
1 1
Σ
2g) ye-
L
t=1
Then, the first term in (F.6) corresponds to L/ eepp (i.a). Furthermore, to handle the last two terms in (F.6), we denote
I = L/L. Ete Ie and ō = 1/L. Et le. Then, we have
7 L Seap
L Σ all 1=1
mg)
Σ exp(( x xgzz)) Vt
7
Σ
L
x
T
xg
L
t=1
(L/8er (ii.b) + (1 - L/80x)) 1)
= =7.8 F In
(F.7)
Combining (F.6 and (F.7), we can upper bound term (ii) as follows:
(GD) < T./8exp ((ii.a)) + - 1x T snll (iibb)
+ + 10 rgl (8exp/ / 1). 107))
45

---

In-Context Linear Regression with Multi-Heea Softmax Attention
Furthermore, can easily that exp(7-xx xrg) 11 (x xig 0) for all C, TTTTr xgg s maxxe]4] |al vigl and ltl 1lvlle/vI. Thus, can upper bound (ii) as
(ii) [L 22,, (((ia)) + max faa/ nzl EE|L)
-1
{(aD)) + ((i.e)
[|vlla
}).
(F.8)
vL
1=1
We first consider (ii..). Note the Φ satisfies (22(1, 20] (2 |z.)), and for any fixed y O, (2(77,)) .) is monotonically increasing in on [0, oo). Following similar argument as in (F.4), we have
llellz vL
max ( zal) exp (7 max 2x) x0)) [E
(ii.a) - Φ2 77s
(F.9)
SEIL]
Similarly, for term (ii.b) and Sxxp/ L - 1, using the facts that (16(7,)) S (7, 1-11-11 and viTr, )) > 0 monotone increasing on [0, 0o) for y > O, we have
[iibb)] S < (7. max | mnl) rE(r.
L 1/L. SSepp L s u17
aar |xTxa/] mol) (F.10)
max RE(T)
VL"
Therefore, combining (F.8), (F.9) and (F.10), obtain
[ii s
1 (x By >
> 01) 01)
(dz
(7
max [x) z0l)) exp (r rE(r.
max (E(T]
| mal))
|1
I
VL
||
+ +
7.
max |of x4g) LE(L]
max
| zgl RE[L]
VL. S
(F.11)
For term (iii) in (F.3), by symmetry, can derive similar bound as in (I) by changing the parameter 7 to -y, Note that ( I (2 and (eb(-%,) S vhy,, 1. given > O. In addition, exp(-7 26 xg)) 2 1 (xx] z4g < 0) for
all é E [L]. Thus, similar to (F.11), we have
(iii) -
(
Σ
1a7 Dig S =0))²
{
(1
22 (7,
max |a, mgl) exp (1 EE[L]
max 1 nal) rE[L] vT.
llv1l2
L
1=1
llB/l2
27 1b
7, max ( xal) [E[T.]
max LE(L)
(F.12)
VT.
Combining (F.3), (F.4,, (F.11) and (F.12), establish an upper bound on (0) (g) using j x To simplify the expression, we apply the concentration concentration results, stated in Lemma G.3. With probability at least 1 8, these three inequalities hold simultaneously:
max [x xrl S v12d. log(go (8g(8//)), FE[T]
1
dh < x1+ poly(d 1-122 - (1, Σ t=1
I
(F.13)
L
Σ
(o)
0)
< min { (log(8/0)/L,1///)
E=1
When L > 16 log(8/8), we have 1/L Σle 1x 24 0) [0.25, 0.75], with high probability. When this is the case,
1/L. El= (a] Fg > 0) also lies [0.25, 0.75). Following (F. 13) and the upper bound on (i) (F.4), we have
r |00| I vI+02 . x- $1 (7:
max ( xzl] EE[L]
v max | .r,l = O O0VI0011//. ..d d. log loggo (8L/5)), kE
(F.14)
where O-) hides constant terms that are independent d, L, and 8. Now, we consider the case where y is sufficiently small constant such that y Ci (Vd. log(L/o)) and thus exp (> maxxely] saa] xi < 1. Similarly, using the upper bound on (ii) given in (F.11), we have
(i)) $ v1+ 0² ... o2 (5((7
max 1xT nal) eex -exp(7 (7. max 1 xzgl EEIL)
?
EE[L]
46

---

In-Context Linear Regression with Multi-Heea Softmax Attention
+ V1+ + 0² 2
max [E[L]
[ T/l)
max ( Zyl RE(L)
/ 1 0².
max EEIL]
2 24 12 10V0+(1++)) . (8L(()))
(F.15)
We can also get the same upper bound for (iii), using (F.12). Recall that we let = (w. nu) be the transformer parameter with
w = -7) and = (iz, -M), where 7, A > 0 with r = 2 j., which is a constant. Combining (F33), (F.14) and (F.15),
we conclude that
(0) 7ed (1)) 2 (10) (i)) + (ii))
(F.16)
where Ö-) omits logarithmic terms. Therefore, we complete the proof.
Remark F.1. Here, we also provide a heuristic but simpler derivation of the debiased GD predictor using first-order Taylor expansion expansion. Note that
Σ Yr exply x7 xa) {=1 Et exp(y 27 xg))
L
Σ ye -1 L+ Y- 2.zg Iq
1. Σ (1+7 t=1
yt
L -.
?
Jq
tig Dig
t=1
where the second approximation results from the first-order approximation of reciprocal and ignoring the higher-order O()2) terms. Similarly, we can show that
L yE exp(-y-2)
L
(1-
Σ
> 8=1
ye
Dg + 7 > x2
L
EL expl 7 x xg))
By subtracting the two terms above, we can get the desired form of approximation, i.e., (0) 2o (T), (0) is given in (F.2).
(1)) where 2 and
F.2. Proof of Theorem 4.2-(ii): Optimality
Proof of Theorem 4.2-(ii). Recall that for fixed > O, we consider the parameter space defined as
f =[(w;)): Vy > 0, (M) y (l))(w((t) for all h E [H], min 1}.
For notational simplicity, we denote μ4+ = Enen- (()), μ = =ke# th and μa = Σnew VHLUH n respectively as the sum of OV parameters for the positive, negative and dummy heads. With this parameterization, the transformer estimator takes the form
We exply 5. xzg)
y1 . exp(ppp) ...00---..0p0 zz)) Ei exp(- xxp(-7xx1] x.g)
(0)
+ p
+ Pd T,
8=1 Ei exp(7 27 zz))
t=1
and thus it suffices to consider three-head attention with w = (,,-7,0) and (J4+, p!_ M- Ma). Recall that we introduce the approximate loss C(w,, (1) in (4.1), which takes the following form:
Ō(w, u) = - 2/ wo + 1 (ww" + +1++02) --1 exp(dww ()
= -11 n Tw)² + (1+ x2) L-1 11 exp(dww
Consider the problem of minimizing this loss function minw,r Blw, nn). For notational simplicity, we let wa = (7, -2) and Na = ()4, denote the parameters of active heads and let pa denote the OV parameters of dummy heads with w/ 0. Then, we can write the loss as
Ilway pas pa) = (1 wa) + (1+02) --1 (Nr exp(dwaw M + 12 2pz 14 1),
(F.17)
where the last term 2pa p._ is due to the interaction between the OV of dummy heads and the OV of active heads. In the following, we find the minimizer of E(waa Pay Ma). We first optimize pa E R with wa and pa fixed. Note that E is quadratic
47

---

In-Context Linear Regression with Multi-Heed Softmax Attention
with respect to pdd By direct caleulation, obtain that the optimal value of pa is pd -wo 1. Plugging pa in to the loss C in (F.17), we have
E(was Ho: (2)) = (1 Mad wa)2 (1 + +2)) L- (ur exp(dwaws
)ua )un - (4 1)))
(F.18)
With slight abuse of notation, we write L(wa: Mas Ma)) as L(way Ma) in the sequel. By rearranging the terms, rewrite the above loss as function of Wa and Mai
Elwar Mo) = 1- 2wd Me + Ma (wawd + (1 + 0²) - (cxp(dwaw.) 11 M)
(F.19)
We note that the two-diment -vo-dimensional vector is orthogonal to 1 = (1,1). Moreover, wa and 1 form an unnormalized orthogonal basis of IR2 To simplify the notation, we define A(waa) = WaWp + (1 + 0²) I.-1, (exp(dwwow) 117). Let us decompose exp(dwaus) (dw into
exp(dwow)) = exp(d)²) + exp(-da2)
exp(d>2) exp(-d722 WaWl 272
= cosh(d??)) 11t T + sinh(dy?) Wowa
7/2
Thus, A(wa) allows the following rank-one decomposition:
Wowl 11
A(wa) Wawl +(1+ + 02) --1 (cosh(d 11T + sinh(dy7).
72
= (1+ L-1 sinh(dd)) }) wown + +1+ L- (cosh(d)7)) - 1) 11
In other words, Wa and 1 are just the (unnormalized) eigenvectors of A(wa). In particular, since sinh(x) 2 0 and
cosh(2) - 2 0 for all 2 0, the matrix A(w,) is positive semi-definite. Thus, minimizing the objective in (F.19) is just convex optimization problem. Optimizing this quadratic function over Na with wa fixed, the optimal value of Mas denoted by Mo (wa), is given by
(w(ww) = ((Wa))--aa
(F.20)
Plugging (wa) in (F220) into E, we obtain the closed-form expression of the optimal loss with respect wa:
E(wa) = I(was s Na (Wa); Mī) = 1 - w A(wa) Wa.
(F:21)
Here we abuse the notation slightly by writing the loss function as a function of wi only.
From the discussion above, see that the inverse of A(wa) is given by
A(wa)- (1+ L- sinhhdd)ddd)) //
Wawy Iwa 1
+ C. 11
(F.22)
where C. > 0 is a number depending y. Noting that wa is orthogonal to the all-one vector 1, combining (F.22), we can simplify (F.21) as
[(wa)= =1 (1 (1+ (1+ (1 I-1 .in.h) (p) 2))²
Consequently, minimizing the loss E(wa) is equivalent to finding the minimal value of sinh(d??) (d>2)-1, By simple calculus, we can show that sinh(x)/x monotonically increasing on R+ and thus the minimum value of sinh(d??))(dd>>) achieved at 7 = 0+. Then we observe that the loss [(wa) is a monotonically increasing function of y, and the optimal value of Wa is " (-"),""). Therefore, the optimal value of the loss is given by
1
(1+02)
(F.23)
(++1) d/L (1+0)).d+
48

---

In-Context Linear Regression with Multi-Heed Softmax Attention
In addition, the optimal choice of Pa and pa under the limit 7 - 0+ satisfy
lim (1t wi) = lim [A(W)) t 7-0
y-50+
1 -1 ((wa) TWT)
wF)2
lim (1 (1 (1+ (R) - sinh(d7?))
7-20
73
| l1
(1+ (1++ - L .d)
(F.24)
(w 1
lim (ā(2)) a +
= lim (Ma 1) = + (1 + ) - d) 7-00
= 0,
where uaa = A(wa) w is parallel to (-1, 1). Recall that we let 0g denote any element in Tx That is, for any h E H4UH.
[w)kka| = 7/ p4 = - = Mey where Moy is defined in (??). Moreover, for any dummy head with h & H4 H__HHhh 0. Moreover, by the definition of My, we have
22 2y My 72 (1+ [2) L-1 cosh(dd7) 1+ (1+ L-1.d
1
as 7 - 0t. Therefore, by the construction of pa and (F.24) have shown that the global minimizer of L over f is attained in .y,, with y - 0t. In other words, the attention model with parameter (wa paa MG) is exactly the model with parameter 0,-.
Furthermore, recall that we have established an upper bound on the difference between transformer predictor and debiased GD predictor in the proof of Theorem 4.2-(i). Let 1j" denote (1 + (1 o²) I.-1 d) " and let 17, 27 My- Given fixed input { (ire,ye)eeep reara) U {. as shown in (F.16), for any sufficiently small y > 0, we have
Gy (6,) De (n-)) & 0(V1+ d.d.
Taking the limit y + 7 0+, prove that the predictor that minimizes the approximate loss E over Fr, (0,.).. coincides with the debiased GD predictor T2d (nr)) This completes the proof.
E3. Proof of Theorem 4.2-(iii): Bayes Risk
In Theorem 4.2-(1), we have shown that the difference between the transformer predictor and debiased GD predictor close when 0 = (w, (u) S and 7 is small. In the following, we show that the risk of the debiased GD predictor 775 (17*) is comparable to the Bayes optimal predictor, where 17" denotes the optimal learning rate. This implies that the transformer model approximately learns the Bayes optimal predictor.
Proof of Theorem 4.2- To prove this theorem, we use the vanilla gradient descent predictor as the bridge. In the first step, we characterize the risk of the vanilla GD predictor. Then, in the second step, we prove that the risk debiased GD predictor is close to that of the vanilla GD predictor. Finally, in the last step, we establish the risk of the Bayes optimal predictor.
Step 1: Risk of Vanilla Gradient Descent Predictor. follows:
We define the predictor given by the vanilla gradient descent as
(n)
" "7 I. y Xx.
Σ ye I Dgg
L.
where 77 > 0 is the learning rate. We characterize E( (n)), the fiirisk of vanilla GD predictor as follows. By direct calculation, we have
E B((Trad (n) vz) 31 = E (" y Xx. -3 -3Tag) zzg)
2
=113 +
+ 0² E
x
(X3
(x3+c)
E || ạ x x 18) 8
+ 17
112
E
X 2
+o2
(F:25)
i
(ii
49

---

In-Contes Linear Regression with Multi-Heed Softmax Attention
where in the second equality, take expectation over 2ig
(0, Ia). Furthermore, we have
1
2n
|||
3
|
+
n
E
x
x
L
2
211
Σ
L
B[(xT
11²
Σ
L
1 (1) (x]
=
T,
T.2
t=1
1=
21/²
Σ
[ Kj. 3 xa 13]
L2
where we use 3 ~ NO, Ia/d). By moment calculations, we have [x 3) = |||||||1 = and
E ||=l|| (x] d E || tr(
1
E3
|
= d + 2,
d
E my B x5 $ - E [ tr(zya)))
B [x( =1
d
for all 1 $i Combining the results above, we have
n L(d+2)+ 2np L(L - 2m + L (d + L 1).
1- 2
(F.26)
2
Moreover, based on the i.idd assumption of the demonstration examples, we have
1
(ii) Σ E=1
[llee
zell] L
B = dL 0².
(F:27)
Combine (F.25), (F226 and (F.27,, the f2-risk of the vanilla GD predictor is given by
(T
vu) = 1 o² 2n 112 + L + 1).
This is quadratic function of 17, and the optimal learning rate that minimizes this risk is given by
1) vgd.- = o/(d((+0))+ +
(F.28)
The optimal f2-rs is given by
L d(1 + 02) +
E (1)d)))))))))))-)))))))))))))))) =1
(F:29)
Step 2: Risk of Debiased Gradient Descent Predictor. direct calculation, we have
Now we compare gad (1) and (y) in terms of the fz-risk. By
E [
(n) - 34) 21 - E[ (T (n) - 3g
= E (1 ugd (1) - y1 1, 1 Xx / L2 vo)
(1g (n) - yo) ?
T E (x) 1r 1[ 1[xx,) xr,) )²1
E [2, y Xx - -BT xr)
T,,
1r 1[xx,)
1,2
E
||x 12 1. 1
2n 13
E
1. 1. xxx T ')] + L
2
E 1, 1[x
La
(iii)
(iv)
In the sequel, we define Tn = X T 1,. Then, Fg ~ NN L- Ia). By direct calculation, we have
(iii = F. x 1, I[ + |11] = = x 1 1
+ + [x
1, 17c11]]
= = TT [ + E [ 1 12) E 1
50

---

In-Context Linear Regression with Multi-Heea Softmax Attention
1 d E || dl = ( +))
(F.30)
Moreover, for term (iv), we have
(iv) E [X + 6) 1,
(X3+
E B X 1 X X3]
+ c _ 11 x[xxTe]
4 Σ
d E 1. [ + E x 1 d
1
E + + B [1
E=1
where the last equality results from XX X. Furthermore, by direct calculation, have
Bi(r] x2) = Bl|aa +
[[{TT1)) +
E1 x[ = d (d + L 1).
ISAEELL
ISKIJEL
Based on the arguments above, we obtain closed-form expression of (iv):
(iv) (d + 1 d
(F.31)
Following a similar argument, for term (v), have
(v) = E + e) 1s 1[ = E x x111 = d B| L.
(F:32)
Combining (F.30), (F.31) and (F.32), we have
E [ (n) vn) [ (n) - yn) y,)² d(1 )²) + 2 2L
2n I
(d(1 +0f) T. I++1).
I
When is bounded by a constant, given L - oc and d/L - E, we know that
E [ (n) ya) (1) (Tyg (n) y,)
0
L with d/L E.
Hence, we have ( (n)) + E ( (17) uniformly for all bounded n,. Thus, under the proportional regime, the optimal learning rate of gRd (n) coincides with that of gord (7), which is given by (F.28). Furthermore, the asymptotically optimal learning rate and the limiting & risk are given by
L
17
lim d/L d(1 + 02) + L+1
E (1+03)+ +1'
1000
(1+102
lim ((T o²
(F:33)
d/ 4/L
&. (1+ +1'
Lox
Here we take the limit with do and d/ L + & in (F:28) and (F.29) respectively.
Step 3: Risk of Bayes Optimal Predictor. As comparison, we also consider the Bayes risk under this setup. Note that under the Gaussian prior over 3, it is classical result that the Bayesian optimal estimator of yg given {(I2, ye) re))e((4] U Ux is given by ridge regression, which takes the form of Tridge = Fridge Here, we let Gridge, denote the ridge estimator of B. and let BayesRisk,,or o² denote the Bayes risk of 3 gridge, where subscripts and o² are used to indicate the dependence on in the limiting regime and the noise level σ². The ridge estimator and the Bayes risk are given by
Fridg= (XT X Ala) y.
BayesRisk
[[ By 1/
The connection between the ridge estimator and the Bayes optimal estimator as follows. With the Gaussian prior (O 72 Ia) for some parameter 7, using the Gaussian likelihood of the linear model, calculate the posterior distribution of The Bayes optimal estimator of / is given by the posterior mean, which takes the form of gridge with
51

---

In-Context Linear Regression with Multi-Heed Softmax Attention
0²/72. Thus, in our case where B ~ NN(0, Ia/d), the Bayes optimal estimator coincides with the ridge estimator with A d
In the following, we derive the Bayes risk which largely follows the main proof in Dobriban & Wager (2018). We set do² hereafter. Conditioning on X. the risk of Vy is given by
E [ - yny² X = EI(Bride,, 2q B x²)²)²
X] + 0² = =I -112 X + o²,
where the expectation is with respect to the randomness of zgg and B. By direct calculation, we simplify B[ - y/g X] as follows:
B Grigee 0112 X]
do +0
- -11 x] =E[](((XX + 1o) 1x) x 12)0812 x
+EIII(X X do² I1] - c11 X
=do ir ((XTX X + do? do)=)))) ?)
+ 0² tr ((XTx X + do² 1a) x(x do2 dodd21a) 1a)
0². ((X X + do² Ia) 1)
where the third equality follows from the fact that (X T X + do [8) -1 X X Ia = do 122..(XT) + (do and the last
equality is obtained by plugging in X X = (X X + do Ia) - do² I By denoting Ed,t. = L and Σ = XTX/I. we can rewrite the equation above as
312 x| = 1+ Ed ((S o² E 1a) 1a)
E gride
(F.34)
The asymptotic behavior of the trace above can be tracked using random matrix theory (e.g., Ledoit & Péché, 2011; Wang et al., 2024). Let M denote the Stieltjes transform of F (Marchenko & Pastur, 1967). where F is the limiting spectral distribution F(()= 1 Ej=1 (I) t). Then,
Ed.2
(E + Ed,L 1a)- AM, E. m=((-22) with Ea. -.
(F:35)
Under the isotropic case where ze iid NO, Ia). mF has an explicit expression
m=(t) = (1 E - -1- (1- - 42t)/2Ct
for all & E CVRt. Based on (F.34) and (F.35), we have
( 02 11 + 402 (02 + 5-1-111 1)²)
BayesRisk
(F:36)
Combining (F.33) and (F.36), the excess risk of D7 (17* compared to the Bayes risk satisfies that
E ( gd (n))) BayesRisk
0² + {(1 <(1+02)+1 +1
0² -0-
1 2 40 + + -02
2
= 2/1 (02 +1+
-E(1+ 1)
25 1. (0² + 5 + 1+ /
40 (0² 1 (5(1 + 02)+
SE {11
+ { (1 +
where the inequality results from 040 + (0² + 1)2 2 o 22 since we assume 0² + 2 1. Furthermore,
since BayesRisk 2 2 x² and we assume 0² 0, then it holds that
E $1+20 (0) -1+
BayesRisk
which gives that 7od (n²* is near optimal when & or o² is large and then we complete the proof.
52

---

In-Context Linear Regression with Multi-Heed Softmax Attention
G. Proof of Technical Lemmas
G.I. Proof of Lemma E.2
Proof of Lemma E.2. We prove this lemma by applying the Stein's Lemma, which states that B|z g(x) = E [[)) holds for x (O, Ia) and any differentiable g-
STEP 1: DECOMPOSITION USING STEIN'S S LEMMA.
To apply Stein's 's Lemma, rewrite the Epp' XX p] into the summation of sequence of functions of Gaussian random variables:
EF xx
pl = E B pj X
p
E
Pe
pi It
Bil x 112 pepel Σ
+
Σ ISEAKEL
E (x) -pepkl.
(G.1)
Here we let and Pr denote the ε-th entry of softmax vectors p and p, respectively. In the sequel, for any E [d] and any i E [L], we let Ft,j denote the j-th entry of xę. For notational simplicity, we let dt,5 denote as the partial derivative of function with respect to Tz.j- By applying the Stein's Lemma twice, we can show that
d
B 3: pcppl
=
Elpepe] +
Blar Pe dr.Pel +
Σ
Blue pe -c,pel
j=1
==1
d
d Elpe Sb + Σ Elpe oZ pe.
= d E pePel + Blor,ppe ds, +
Σ
(G.2)
1=1
j=1
Here the first equality is obtained by applying Stein's 's lemma with g( = Ft,j applying Stein's Lemma twice, we have
Pe. Similarly, for any k + l, by
d
E| (xe, σx) pePpl
E E[aks Pe 0e,sKk] +
Σ
E[zn. Ps despel
j=1
Elpe Oc.PPll + Σ
Elpx OR. Se,jpl
jml
BiJ de,pPll +
Σ E BlO de,pel.
(G.3)
STEP 2: DERIVATIVES CALCULATIONS AND SIMPLIFICATIONS.
By (G.2) and (G.3), write (G.1) of expectations involving p. p, and first- and second-order their derivatives. By direct calculations, the first-order derivatives of p given by
te,Pe = w Uy (pe - př),
Ox - Uy PEPk-
(G.4)
The second-order derivatives of are given by
2? (1 (1-2pt) (pt - -p))- On Oe, 0t,, Or. U 2pt) ŌR.iPe = _w² 11. (1 2pe) PePa-
8i = w uy (1 - 2pr) O2,5Pt Pe = w2.
(G.5)
(G.6)
We can similarly derive the first- and second-order derivatives of p.
In the following, we conclude the proof of this lemma by substituting the calculations above back into (G.2) and (G.3), and simplifing the expression. First, note that, we view Epp" X y) as bivariate function (w, W) is quadratic
53

---

In-Context Linear Regression with Multi-Heed Softmax Attention
function. The constant term is given by
d E Elpepel d 1.- P].
(G.7)
Thus, we can write Elp" XX p] as
p] = Ao w² + A w w. + A2 _02 + d Blp"7 y]
(G.8)
for some constants Ao; A1, A2 to be determined. By examining the derivatives of p and notice that these coefficients are given respectively by
Ao
Blpe Pe, +
Σ
Σ
Elpr Sz.
(G.9)
f=1 ==1
1sttksLj=1
A1 BlOE, De.pell +
Σ
Σ
B Pe, + BlOk, de,pel].
(G.10)
A-1
L
A2 =
Bip 0 j0 +
Σ
Blpe dc.pRRl
(G.11)
|=1j=
For Ag, combining (G.5), (G.6), and (1), we have
A2 = 2²2 ||0 12 EElp" T pl - 2 Elp" T 1oo2)
~. 11012
E [peee Pel +
E Pepee Pk
+ 2052 1|011/
Elpepe 17 +
Σ 1SERKEL
Elpepe pkl
By direct calculation, we have
L Σ
ElpePe Pe] +
Σ
Eppepe. .
Σ
prPt
Pk
=Ely BiyTT) p),
IStENEL
k=1
Elpape pi +
Σ
Epp.e
P
E
pcpt
Ely P- TT 112))
ISEAKEL
where in the first equality, we use the fact that p is prabability distribution. Combining the above three equalities, we can simplify A2 as
A2 = =2 |10112 (Elp 5j - 1pp p 102
25. || Elp T.. TT 1r] + 2w 1111 Elp TT 5 1ل
2 ||01| El-pT po2 + +p [p1133
(G.12)
Following the same argument, we can show that Au defined in (G.9) can be simplified as
Ag = 2 | E pJ 02 TT p. (1r18)
(G.13)
Then, it remains consider A1. Note that by the first equality in (G.4), we have
ElOr,PP 0z,5el ww.
0z Elpe Pe
p Pe Pe
P2
p F2]
j=1
54

---

In- In-Context Linear Regression with Multi-Heea Softmax Attention
wi. 1|e:12 Elp p p ppp p'po2 + wie 111111
>
E [pip
ē=1
By the second equality in (G.4), for any k + E, we have
E j de, wū, 2z E EleeP= pePk]
E , de, ww. oj Elpe Pe
P4 Pt - PE -.pppp- p pil.
Thus, combining (G.10) with the above three equalities, can simplify the first part of A
E B Se.SPe)
Σ Σ
E
Bsokpee dr.jāk]
wis |1012 Elp p p T po2
p-p p 02
1|013
E0 (pepe) +
Σ
E Elpept paPkl
ISEEKE
wi | Elp p p² po2
p0 p0
(p 1)
(G.14)
Similarly, the second part of A1 be written
B B Despel +
Σ
Ei Sr.jPi]
15t4k$L
ww
1|
Blye Dr pē -
Pt. P
2 P2 P
wil 1101112
Σ Elpe: Ph D2 Pk
22 Pe Pk
p P
1St4K5L
wi:
||0:12
Elp" T 1, p - 11p11 . 12 - 1712 p 1. + 11-12-171131
w ||0:12
[[ (Iel2) (1-11P1))
(G.15)
Thus, combining (G.14) and (G.15), we can write A₁ as
A w - ||01| Bpp T p pp Finally, by combining (G.7), (G.8), (2. (G.13), and (6.16), we complete the proof.
D" 12
(y p) + (1 ll( (1 11]
(G.16)
G.2. Proof of Concentration Arguments
Lemma G.1 (x2 (2-connentrationn) Consider random vector (#1...., Id) E Rd with each entry i...d. sampled from
(O 1). For any positive 1 > 0, the following inequality holds:
F[| 2 d 2 + 2t) S exp(-t),
P(I < d 2vdt) exp((tt).
Also, holds that P|[ 2 d (1 + t)²) < exp(-dt? /2).
Proof of Lemma G.I. See Lemma in Laurent & Massart (2000) and Lemma A.1 in Bai (2024) for detailed proofs.
Lemma G.2 (Binomial concentration). Consider Binomial random variable Binom(n, with n E N and p E (0, 1), then for any E (0, 1), holds that
P(X 2 (1+1) mp) <xpp(-nnp/3,, P(X 2 (1 t) np) S (nnp2))
Proof of Lemma G.2. The result follows a direct implementation of Chernoff's bound.
55

---

In-Context Linear Regression with Multi-Heea Softmax Attention
Lemma G.3. Consider xg ut No, Ta), 3 4 N(0, da/d) and ye = 3 Ig + eb where of utd. any 8 E (0, 1), with probability at least 6, the following events hold simultaneously:
N/O,, with L E N. For
(i). max |z]2,g nal I 122 log(81/));
(ii) 1. El v $ 1 + ²² poly(d 1/2 log(1/))):
(iii). 1. (z 2 > 0) 2 - min min
Proof of Lemma G.3, We first consider the following good events hold:
& = -11llll s /6d || < 1 + 2d- 1og(4/8)
and we can upper bound P(E1)= < 8/4 using Lemma G.I. Note that for fixed Tg. we have x Iq Lt NoO, 212g by applying a Gaussian tail bound, we have
and hence
P (33 [L] 12 Zg
2 t xgg) S
P (1x) Egl 22122 s 2L exp (-11/11x,,11))
Hence, max 1o zgl S v12d (8g(81//) with probability greater than 5/4 under good event & Similarly, we have = B It +e
sid. (O || + )²)). Based on Lemma G.I, it holds that
PP (1yl3 L (1I 2 + 0²) (1+ t)² (3) S exp(-It²
Following this, under good event E with probability greater than 1 5/4, we have
Σ v S (1 + 5d 5/2 1/2 log(4/0)) -(1 5T. 1/2 log(4/5))
t=1
=1 =1 0² poly (d 1/2 I 1/2 log(1/6))
Furthermore, since 27 2g idd A(0, (=))(2) and P(z) xq 0 xg)) 1/2 for all Zg. then using the binomial tail bound in Lemma G.2, we can show that
1 ] t/2 xg) < 22epp(-AL/A P 1 0)
and hence 1/2 E/= 1(x7 x > 0) < min L 1og(8/8), with probability at least - 45. Combining all the
arguments above, we can complete the proof.
G.3. Proof of Anisotropic Pre-Conditioned GD Estimator
For simplicity, we consider a symmetric pre-conditioning matrix which proved to be the global optimal parametrization for pre-conditioned GD, which can be extended to the general case with more careful arguments (e.g.. Ahn et al., 2023a).
Lemma G.4. Suppose 3 (O, Ia/d), X [II,.... FT.TT E R with Dz sid N(0, Σ) and ye BT x, + c, with eg idd
AOO,, o²). For any I' Rdddd, define T := To (T) L-1 (X3 + E) XT-1xp Consider the ICL prediction
risk E[(Tyed - ?] as a function T. Then, the minimizer of [ | given by
r" (I= 1) Σ (tr(E) + do // Ta-
Proof of Lemma G.4. For any matrix I' e dd and the risk of corresponding estimator follows
2 E ( + XT In 33 -Txx.)) xzq
E E
t
56

---

In- Ca-Context Linear Regression with Multi-Head Softmax Attention
E tr (XTXXT
La)
1a)2)
o² B(tr(T-IX
XT-XX--1 Σ) to?.
(G.17)
L
iii
Following the decomposition above, by simple calculations, the second term follows
(ii) = E - XT-
L ur ((E))1)²)
(G.18)
Furthermore, note that the first term takes the form
2 ii) tr(E)) - Eftr((X T XI [T +
1
(x
(XT (X-))x--1
2
1
(2-))) + T (
(G.19)
where the second equation follows XTX Et= faaa and we can show that
[[(x T x) Σ B( 19 +
Biraz] 2,2, = L B[( ) L(L - 1). E2.
1stRiSL
Specifically, for all (t,5) E [d] [d]. the (i, entry of the first matrix satisfies that
d
( 12gl Σ 2 E +
Σ Ela Fe Bl = Eig tr(E) + E. 1=1
where the first inequality uses the Isserlis' theorem and thus we have B1 tr(E)=+ 222 For notational simplicity, denote I' = I Combining (G.17), (G.18) and (G.19), it holds that
E (TNyed - = 02
tr Σ
2
tr
+ (r(E) + do²
tr (f-2)
L+1
tr (I dL
d
d
dL
tr(2)
2
1 dt. ir ((((((E) + do²)
=0² +
Ir (
Ia (L + 1) E} 2-
d
Since + + Ia (L+11 + 1). E is symmetric and invertible, the minimizer given by
-t I. (tr(EE) + do²) Ia + (L 1)- E, s.t. r r== (1+ Σ + tr(E) + do²
Ia
due to the quadratic form and then complete the proof.
57


Reviews:
[{"review_id": "tM88pvlcy3", "reviewer": "Reviewer_M7UK", "summary": "This paper investigates the training dynamics of multi-head softmax attention in in-context learning. Through experimental analysis, the authors discover two key patterns: (1) QK weight matrices develop a diagonal structure, with diagonal elements being nearly uniform across all heads, and (2) QK weights and effective OV weights share the same sign, with their average approaching zero. Based on these findings, they propose a simplified training model with 2H parameters and describe the training process in two distinct stages. The authors also analyze the expressive power of their simplified multi-head attention model, comparing it with single-head attention, and validate their findings through extensive experimentation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the positive feedback and appreciation of our work. Here's our response to the questions.\n\n**On Heuristic Derivation and Two-Head Gradient Flow.**\nWe acknowledge that in Section 4.2, we employ a heuristic derivation to explain how the observed patterns emerge during gradient optimization, using an expansion-based argument. To establish a more rigorous foundation, we analyze the gradient flow in Appendix C to explain why these patterns persist once they emerge and to characterize the trajectory of the parameter evolution. \n- (`Multi-head Attention Effectively Acts as Two-head Model`) Furthermore, our experiments (see Figure 3) demonstrate that a multi-head model essentially implements the same estimator as a two-head model. As detailed in Equation (4.5), on the solution manifold the multi-head model ultimately separates into two distinct head types, determined by the signs of $\\mu^{(h)}$ and $\\omega^{(h)}$. Within each group, the heads behave equivalently to a single head with $\\omega = \\omega^{(h)}$ and a $\\mu$ equal to the sum of the $\\mu^{(h)}$ values within that group. Thus, analyzing the dynamics of a two-head model provides all the essential insights.  \n- (`Dynamics are Similar Under Symmetry`) Besides, note that if we initialize a multi-head model with identical absolute values for $\\mu$ and $\\omega$ (see Definition C.1) and randomly assign each head a positive or negative sign, then as the number of heads increases, each type group is highly likely to contain nearly the same number of heads. By symmetry, tracking the effective KQ and OV parameters within each group is equivalent to analyzing a two-head model.\n\n\n**On Discrepancy between Theoretical Predictions and Experimental Results.** \n- (`Experimental Justification: Small KQ already Behaves as the Limiting Model`)  As shown by the theoretical results in Section 4.3, when KQ approaches 0, the model exactly implements the GD predictor. Therefore, we compare the performance—measured in terms of mean squared error (MSE)—of the learned model with small KQ against that of the GD predictor. Figure 3(b) demonstrates that the loss curves for the learned transformer model and the GD predictor coincide, indicating **identical performance** and confirming that the approximation (see Section 4.3) holds well even for a small constant level KQ. Moreover, because the model is trained via gradient optimization, the **negligible loss gap** between the small KQ multi-head attention and its limiting case suggests that further optimization from a small constant KQ to the limit of 0 would require an exceptionally long time —likely exceeding the number of training steps used in our experiments.\n- (`Potential Meachanism: Edge of Stability`) In practice, the particular magnitudes of KQ parameters learned by the optimization algorithm depends on particular choices of learning rate, training steps, and batch size (if SGD or Adam is applied). This is partly due to the *\"river valley\"* loss landscape near the global minimum (see Figure 5(b)). From the theoretical perspective, this phenomenon can be attributed to the **edge of stability** (e.g., [1]), which is commonly observed in neural network training. In this regime, gradient descent progresses non-monotonically, oscillating between the “valley walls” of the loss surface and failing to fully converge to the minimum. Experimental results indicate that choosing a *smaller learning rate*, a *larger batch size*, and *sufficiently many training steps* can drive the KQ parameters to eventually approach zero.\n\nWe would thank the reviewer for the nice questions and we will incorporate these clarifications in the revised version.\n\n[1] Cohen, J. M., Kaur, S., Li, Y., Kolter, J. Z. and Talwalkar, A. (2021). Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065.", "claims_and_evidence": "The claims are well-supported by both theoretical analysis and empirical evidence.", "methods_and_evaluation_criteria": "The methodology and evaluation criteria are sound and appropriate.", "theoretical_claims": "While I haven't examined all proofs in detail, the theoretical framework appears well-documented and logically sound.", "experimental_designs_or_analyses": "The experimental designs are robust and well-constructed, providing sufficient evidence to support the paper's conclusions.", "relation_to_broader_scientific_literature": "This research contributes to our understanding of large language models by examining how multi-head softmax attention networks learn to perform in-context linear regression, offering valuable insights into LLM training processes and mechanisms.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "**Strength**:\n\n1. The paper effectively uses experimental analysis to reveal parameter patterns in trained softmax attention models, providing insights into both training dynamics and operational mechanisms.\n2. The research offers a thorough comparative analysis of softmax and multi-head attention architectures.\n\n** Weaknesses**:\n\n1. The analysis of multi-head attention training dynamics relies heavily on intuitive reasoning, with Appendix C limiting its formal analysis to two-head attention only.", "other_comments_or_suggestions": "If there are any misunderstandings on my part, please point them out, and I will reconsider my evaluation of this work.", "questions_for_authors": "Regarding the apparent discrepancy between theoretical predictions and experimental results: While the analysis in Appendix C suggests w should approach 0, Figures 2(b) and 1(b) show convergence to a non-zero constant. Could you explain this apparent contradiction?", "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": "zhI8hnMQ0c", "reviewer": "Reviewer_Y3fQ", "summary": "The paper provides an comprehensive understanding of multi-head softmax attention in conducting ICL for linear regression tasks. Through empirical investigation, it observes the specific patterns for the optimal parameters within in the multi-head transformer structures and the superiority of multi-head over single-head. Based on the observations of the specific patterns, it reparameterizes the attention models to conduct theoretical understanding of the training dynamics and expressive power of multi-head softmax attention. It shows that multi-head softmax attention emulates preconditioned gradient descent and achieves near-Bayesian optimality in in-context learning. Additionally, the paper extends its findings to non-isotropic covariates and multi-task regression and also provide comparison with linear attention.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and assessments. Here's our response to the comments.\n\n**On Comparision with Existing Work.**\n- **Comparison with [1].**\n    - (`[1] corresponds to single-head case in our paper`) [1] considers the multi-head attention for **multi-task** linear regression, where the number of heads equals the number of tasks. Under a specialized initialization, [1] proves that **each head learns to handle one task**,  corresponding to the **single-head case** in our paper. In contrast, we consider a more flexible setup in which multiple heads can be freely allocated to a single task, allowing the model architecture to be more expressive and complex. As a result, while [1] shows that the single-head model learns a **nonparametric** predictor with KQ scaling as $1/\\sqrt{d}$, we demonstrate that the multi-head model learns a **parametric** GD predictor with KQ converging to $0^+$. Hence, the analysis are **fundamentally different**.\n    - (`[1] also adopts reparametrization`) Note that [1] assumes decomposable weights (Definition 3.1 in [1]), which ensures that the eigenvector space remains fixed. Thus, the analysis in [1] also employs a **reparameterization**, showing that diagonalizability is preserved along the gradient flow and focusing on the eigenvalues, i.e., the diagonal entries in the isotropic case.\n- **Comparison with [2].** \n    - (`Empirical Insights`) [2] only identifies the diagonal KQ patterns with potentially positive and negative values for **two-head** attention, as well as identical performance when the number of heads exceeds two. We further **quantitively** figure out the detailed **sign-matching**, **homogeneous KQ magnitude**, and **zero-sum OV** patterns beyond $H=2$, and reveals that the multi-head softmax attention learns to implement **GD predictor**, which are not covered in [2]. Thus, we establish a more complete understanding.\n    - (`Theoretical Insights`) We provide a comprehensive explanation by analyzing *how patterns emerge* through **training dynamics**, as well as *how the trained model operates* via a **function approximation and optimality analysis**. While [2] considers full-model parameterization, it focuses solely on a loss approximation similar to our Proposition 4.1. However, it does not capture the training dynamics and function approximation----that we develop based on the loss approximation results. Moreover, the main result in [2]—the superiority of multi-head over single-head attention—is a **natural corollary** of our findings, since we  show that single-head attention learns a nonparametric kernel regressor, whereas multi-head attention learns a parametric GD. By applying the approximate loss from Proposition 4.1 with $H = 1$ and $H = 2$, we can reproduce the result. Finally, [2] only analyze KQ parameters with OV parameters fixed to the corresponding optimal solution, while we analyze with both KQ and OV parameters.\n- **Comparison with [3].** [3] focus on the binary classification tasks while we work on the regression task, marking a clear distinction. \n\n**On Real-Word Implications.** Understanding what transformers learn for specific task serves as a starting point for investigating general ICL meachanism. The theoretical and emppirical insights lay the foundation for understanding how deep models learn language.\n\n**On Transformer Architecture.** In this paper, we aim to understand the mechanism of the **standard** multi-head attention architecture and the weighted sum of head is beyond the scope. Also, $\\mu^{(h)}$'s naturally act as \"weights\" for each head, and are expected to learn the \"optimal\" weight during the training. Thus, we do not anticipate any improvement.\n\n**On Reparametrized Model.** As shown in Observation 4 of §3, the attention model develops a diagonal-only pattern in the KQ circuit and a last-entry only pattern in the OV circuit during the **early stages of training** and then continues optimizing within this regime. This indicates that a diagonal parameterization is sufficient to capture the core behavior of model training. Importantly, we **do not impose any diagonal constraints** on the KQ parameters—the diagonal pattern **emerges naturally** during **full-model** training and persists throughout (see Figure 4). Moreover, if the KQ parameters are initialized as a diagonal matrix, this structure is preserved over the course of training.\n\nWe thank the reviewer for pointing out the related work, and we will incorporate these comparisons in the revised version.\n\n[1] Chen, S., Sheen, H., Wang, T., & Yang, Z.. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. COLT 2024.\n\n[2] Cui, Y., Ren, J., He, P., Tang, J., & Xing, Y. Superiority of multi-head attention in in-context linear regression. 2024.\n\n[3] Li, H., Wang, M., Lu, S., Cui, X., & Chen, P. Y. How do nonlinear transformers learn and generalize in in-context learning? ICML 2024", "claims_and_evidence": "In general, the claims made in the submissions are clear and well-written. The paper provides empirical evidence for the observed specific patterns for the optimal configuration of multi-head softmax attention and the superiority of multi-head attention. For theoretical results, it also provides some empirical validations and proofs of the statements.", "methods_and_evaluation_criteria": "The paper did not propose methods / evaluation criteria.", "theoretical_claims": "I have checked the sketch of the proofs for the main theorems which are shown in Appendix. I did not observe mistakes in the proofs.", "experimental_designs_or_analyses": "I have checked the details of the experiments setting, results and analysis in Section 3 (and some further explanations and discussion in Appendix). In general, the empirical settings are reasonable and align well with the design of theoretical settings, and I did not observe particular issues.", "relation_to_broader_scientific_literature": "The key contribution of the paper, claimed by the authors, includes identifying the optimal configuration of multi-head softmax attentions for in-context linear regression, analyzing the training dynamics and expressive power, and demonstrating its advantages over single-head models. The analysis builds on prior studies of ICL for linear regression using transformers. \n\n**However, a major concern regarding this paper is its novelty, as many of its contributions appear to have been covered by previous works.** For instance, [1] has already analyzed the training dynamics of multi-head softmax attention in ICL, discussing optimality through upper/lower bounds on the ICL loss and providing results based on Bayesian risk analysis. Additionally, [2] has presented similar findings regarding the optimal configuration, such as the diagonal structure of the QK matrix. [2] has also compared multi-head and single-head attention and extended the discussion to external scenarios such as non-Isotropic covariates (correlated features).\n\n[1] Chen, S., Sheen, H., Wang, T., & Yang, Z. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. COLT 2024.\n\n[2] Cui, Y., Ren, J., He, P., Tang, J., & Xing, Y. Superiority of multi-head attention in in-context linear regression. arXiv preprint arXiv:2401.17426. 2024.", "essential_references_not_discussed": "This paper focuses on investigating the training dynamics and expressive power of multi-head softmax attention. [2] investigating similar topics but the paper did not cite and compare with [2]. In addition, given that [1] also investigates the training dynamic and optimality, the paper did not provide a very clear explanation about the difference between this work and [1]. [3] also have a close relation to the content of this paper. These works are highly related and important, and it is necessary to discuss them for a careful literature review.\n\n\n\n[1] Chen, S., Sheen, H., Wang, T., & Yang, Z.. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. COLT 2024.\n\n[2] Cui, Y., Ren, J., He, P., Tang, J., & Xing, Y. Superiority of multi-head attention in in-context linear regression. 2024.\n\n[3] Li, H., Wang, M., Lu, S., Cui, X., & Chen, P. Y. How do nonlinear transformers learn and generalize in in-context learning? ICML 2024", "other_strengths_and_weaknesses": "Strength:\n1. In general, the paper is well-written and easy to follow.\n\n2. The analysis is comprehensive, as it involves both empirical and theoretical results of multi-head softmax attention, and cover different perspectives, including optimal model configuration, training dynamics, expressive power, comparison between multi-head / single-head attentions, comparison with linear attention, and extensions to non-isotropic covariates and multi-task regression.\n\nWeakness: \n1. As indicates in the previous review section, the novelty of the paper is limited, as the main contribution has been covered by previous papers [1][2].\n\n2. The paper observers the optimal configuration of softmax multi-head attentions by only empirical observations. In contrast, [2] provides theory showing the diagonal property of the optimal QK weight and other configurations. Additionally, compared to [1] and [2], the theoretical analysis in this paper reparameterizes the attention parameters into just two vectors, significantly simplifying the complexity of the theoretical derivation.\n\n[1]  Chen, S., Sheen, H., Wang, T., & Yang, Z.. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. COLT 2024.\n\n[2] Cui, Y., Ren, J., He, P., Tang, J., & Xing, Y. Superiority of multi-head attention in in-context linear regression. 2024.", "other_comments_or_suggestions": "1. It is suggested that the authors can provide some practical implications of the findings in this paper for in-context learning on real-world text data. Understanding this is important, as the theoretical analyses focus on simplified settings, and their relevance to real NLP tasks remains uncertain.\n\n2. The paper considers that for multi-head attentions, the outputs of all heads are summed. Would incorporating a learned weighted sum for each head further improve the prediction? Intuitively, it can provide more adaptability by allowing the model to assign different importance to different heads. Additionally, the paper states that increasing the number of heads beyond H = 2 provides no additional benefit. If a weighted sum is used, could further increasing the number of heads lead to performance improvements?\n\n3. The training dynamics analysis is based on a reparameterized model. If I understand correctly, it assumes that throughout the entire training process, the QK matrix is constrained to always have a diagonal structure, with only its diagonal values being trained. Although empirical results suggest that the optimal configuration after training is indeed diagonal, this simplification deviates from the actual training process, potentially oversimplifying the learning dynamics and failing to capture the complex interactions that may emerge in real transformer training.", "questions_for_authors": "N/A", "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": "AgI1cgxmhV", "reviewer": "Reviewer_bTEk", "summary": "The authors find that a multi-head softmax attention effectively becomes a two-head softmax attention network that approximates linear attention better than a single-head softmax one. The advantage of softmax based attention over linear attention is that one does not have to train separately for different context lengths. The experiments bears out these observations. There is also some analysis to training dynamics.\n\n## Update after rebuttal\n\nI maintain my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their positive feedback and appreciation of our work. Here's our clarifications on theoretical analysis.\n\n**On Informal Theoretical Claims.** We would like to clarify that most of the theoretical arguments presented in this paper are rigorous, with the exception of the analysis for Stage I in Section 4.2, i.e., theoretical derivation of feature emergence via gradient optimization. While both Stage I and Stage II admit a Taylor expansion-based argument, the analyses are derived from different perspectives. We will provide further clarification as follows.\n\n- (`Analysis for Stage I`) In Stage I, we perform a gradient-based analysis to identify the driving components in the gradient that lead to the emergence of patterns. We acknowledge that in this stage, we employ an informal argument by focusing on the first-order approximation of $\\exp(\\cdot)$, leveraging the small-scale initialization that renders higher-order terms negligible. To support this argument, we compare the informal theoretical conjectures with experimental results, which show perfect alignment.\n- (`Analysis for Stage II`) In contrast, during Stage II, we analyze the loss landscape to understand the convergence manifold. Assuming that the patterns identified in the first stage have already emerged, we examine the higher-order terms in the Taylor expansion of the loss function to demonstrate that homogeneous KQ scaling is a necessary condition for loss minimization. We emphasize that the expansion-based argument in this stage is rigorous for the following reasons: (a) the parameter scale remains small during the final stage, allowing the series summation to be analytically tractable, and (b) we derive the optimality condition for each term of the form $\\langle\\mu_t,\\omega_t^{\\odot k}\\rangle$ for $k\\in\\mathbb{Z}$, without truncation, ensuring that the result is both rigorous and concise.", "claims_and_evidence": "The claims are justified.", "methods_and_evaluation_criteria": "It is mostly a theoretical paper, needing simple experimental evidence, which was provided.", "theoretical_claims": "The theoretical claims involve informal but reasonable Taylor expansion based arguments.", "experimental_designs_or_analyses": "Figure 7 provides empirical evidence of the arguments being right.", "relation_to_broader_scientific_literature": "In recent past, there has been many papers analyzing in-context learning in toy problems. Many of these involve single-head attention. There are not many papers among them tackling multi-head attention. Although some of the arguments are informal, there is value in this work.", "essential_references_not_discussed": "I did not notice obvious omissions.", "other_strengths_and_weaknesses": "I think the analysis is useful for many but the key results were relatively obvious to me.", "other_comments_or_suggestions": "This seems to be well-written paper.", "questions_for_authors": "I do not have any.", "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:

Meta-review: This paper investigates the training dynamics of multi-head softmax attention in in-context learning. The reviews for the paper are mostly positive. While reviewers expressed concerns regarding novelty and relationship with earlier works, they were mostly resolved during the rebuttal process. I recommend acceptance and encourage the authors to address all the reviewer's concerns in the final version.
Decision: Accept (poster)