The Initialization Determines Whether In-Context Learning Is Gradient Descent
Authors: Shifeng Xie
Rui Yuan
Simone Rossi
Thomas Hannagan
Venue: Accepted by TMLR ()

Abstract: In-context learning (ICL) in large language models (LLMs) is a striking phenomenon, yet its underlying mechanisms remain only partially understood. Previous work connects linear self-attention (LSA) to gradient descent (GD), this connection has primarily been established under simplified conditions with zero-mean Gaussian priors and zero initialization for GD. However, subsequent studies have challenged this simplified view by highlighting its overly restrictive assumptions, demonstrating instead that under conditions such as multi-layer or nonlinear attention, self-attention performs optimization-like inference, akin to but distinct from GD. We investigate how multi-head LSA approximates GD under more realistic conditions—specifically when incorporating non-zero Gaussian prior means in linear regression formulations of ICL. We first extend multi-head LSA embedding matrix by introducing an initial estimation of the query, referred to as the initial guess. We prove an upper bound on the number of heads needed for ICL linear regression setup. Our experiments confirm this result and further observe that a performance gap between one-step GD and multi-head LSA persists. To address this gap, we introduce $y_q$-LSA, a simple generalization of single-head LSA with a trainable initial guess $y_q$. We theoretically establish the capabilities of $y_q$-LSA and provide experimental validation on linear regression tasks, thereby extending the theory that bridges ICL and GD. Finally, inspired by our findings in the case of linear regression, we consider widespread LLMs augmented with initial guess capabilities, and show that their performance is improved on a semantic similarity task.

Published in Transactions on Machine Learning Research (12/2025)
The Initialization Determines Whether In-Context Learning Is Gradient Descent
Shifeng Xie
shifeng. Stteoeooo-eeeeecm-ppo
Telecom Paris
Institut Polytechniqqe de Paris France
Rui Yuan
yy42606r Ggmaail com
Lensi Labs, Poris
France
Simone Rossi
simonc.rossiBeurccom.ff
EURECOM
France
Thomas Hannagan
thomas.hannagen watellontt..ttteeeeeoooooooooooe
Stellantis
France
Reviewed OpenReview: https: openreview. net/ forum? id= fugSKLDt.s
Abstract
In-context. learning (ICL) in large language models (LLMs) striking phenomenon, yet its underlying mechanisms remain only partially understood. Previous work connects linear self-attention (LSA) to gradient descent (GD), this connection has primarily been established under simplified conditions with zero-mean Gaussian priors and zero initialization for GD. However, subsequent studies have challenged this simplified view by highlighting its overly restrictive assumptions, demonstrating instead that under conditions such as multi-layer or nonlinear attention, relf-attention performs optimization-like inference, akin to but distinct from GD. We investigate how multi-head T.SA approximates GD under more realistic conditions specifically when incorporating non-zero Gaussian prior means in linear regression formulations of ICL. We first exlend multi-head LSA embedding matrix by introducing an initial estimation of the query, referred to as the initial guess. We prove an upper bound on the number of hcads noeded for ICL lincar regression sctup. Our experiments confirm this result and further observe that a performance gap between oue-step GD and multi-head LSA persists. To address this gap, we introduce L.S,, simple generalization of single-head LSA with trainable initial guess Dq- We theoretically establish the capabilities of LLA and provide experimental validation linear regression tasks, thereby extending the theory that bridges ICL. and GD. Finally, inspired by our findings in the case of linear regression, we consider widespread LLMs augmented with initial guess capabilities, and show that their performance is improved on a semantic similarity task.
1 Introduction
Large language models (LLMs) exhibit the interesting phenomenon of in-context learning (ICL), whereby models adapt to new tasks from a few input-label pairs presented in the context, without parameter updates (Brown al., 2020; Dong al., 2024). This capability has motivated extensive efforts clarify the underlying mechanisms. A prominent line of work interprets ICL. in simplified linear regression settings implicitly
1

---

Published in Transactions on Machine Learning Research (12/2025)
performing gradient descent (GD) within a forward pass of linear self-attention (LSA) (Garg et al., 2022; Von Oswald et al., 2023).
However, this equivalence has mostly been established under restrictive assumptions, notably zero-mean Gaussian priors for regression weights and zero initialization for GD. Recent work indicates that these conditions are fragile: Zhang et al. (2024b) showed that introducing non-zero mean prior produces persistent gap between LSA and GD, undermining previous guarantees (see Fig. 1). From modeling perspective, the assumption of zero-mean prior is quite restrictive: it corresponds to a learner that believes all task-specific regression weights are centered around the origin in parameter space. In realistic pre-training regimes, transformers are exposed to broad data distributions and acquire a shared bias across tasks, which naturally captured by a non-zero prior mcan. These findings raise fundamental question: "under what conditions can LSA failkfully recover GD, and when does it fundomentally fail?"
Loss Leta 6D
20-
10-
1000
2000
3000
4000
Training Steps
In this paper, we revisit the ICL-GD connection under more realistic assumptions, explicitly incorporating non-zero prior means and system- atically analyzing the role of attention heads and initialization. Our study reveals that the decisive factor is the initialization of the query's prediction, which we term the initial guess by. Misalignment between By and the prior induces persistent gap that cannot be resolved by simply increasing the number of heads. Motivated by this observation, we propose 1,-LAA an architectural extension that incorporates trainable initialization mechanism, thereby restoring equivalence with GD even in the non-zero mean setting.
Figure 1: Training and evalua- tion loss curves of LSA with a non-zero prior mean. The dashed red line denotes the baseline loss achieved by one-step GD.
Our analysis is closely related to the concurrent work of Zhang et al. (2024b), who study a linear transformer block (LTB) obtained by composing an LSA layer with linear multi-layer perceptron. In the non-zero mean setting, they show that an optimally trained LTB implements a preconditioned, near-Newton update with a learnable initialization. In contrast, we ask to what extent the ICL GD correspondence can be recovered with LSA architectures. From this viewpoint, B--LSA can be seen as minimal extension that acts only on the inpul-side init ialization, isolating the role of the query's 'ss initial guess without adding extra layers or a large number of additional parameters.
Contributions. This work makes the following contributions:
We prove that when regression weights have a non-zero mean, multi-head LSA cannot in general replicate one-step GD, even with arbitrarily many heads, establishing a fundamental limitation of the ICL-GD correspondence.
We show that the query initialization yg is the decisive factor: misalignment induces a persistent gap, while correcting y, suffices to recover GD even with a single head.
3. We propose LSA, an extension of LSA with a trainable initialization vector, and demonstrate both theoretically and empirically that it restores equivalence with GD in the non-zero mean setting.
4. We provide proof-of-concept experiments showing that introducing explicit initial guesses improves ICL performance in LLMs, thereby linking our theoretical results with practical prompting strategies.
Scope. Our analysis focuses on linear regression with linear self-attention, a simplified but analytically tractable setting. Within this framework, we identify precise conditions under which LSA diverges from gradient descent and propose LSA as principled correction. These results provide a foundation for extending analysis to richer transformer architectures, including softmax attention and multi-layer models.
2

---

Published in Transactions on Machine Learning Research (12/2025)
1.1 Related Work
Theoretical studies on ICL have analyzed its mechanisms to understand how LLMs effectively learn from contextual examples (Brown et al., 2020). ICL can be framed as an implicit Bayesian process where the model performs posterior inference over a latent task structure based on contextual examples, performing a form of posterior updating (Xie et al., 2022; Falck et al., 2024; Panwar et al., 2024; Ye et al., 2024). Alternatively, a more recent perspective suggests that ICT, in transformers is akin to gradient-based optimization occurring within their forward pass. Von Oswald et al. (2023) demonstrate that self-attention layers can approximate gradient descent by constructing task-specific updates to token representations. They provide mechanistic explanation by showing how optimized transformers can implement gradient descent dynamics with a given learning rate (Rossi et al., 2024; Zhang et al., 2025). While this work provides a new perspective on ICL, it limits the analysis to simple regression tasks and it simplifies the transformer architecture by considering a single-heat head self-attention layer without applying the sfmx(-) function on the attention weights (also known as linear attention). Ahn al. (2023) extend the work of Von Oswald al. (2023) by showing how the in-context dynamics can learn implement preconditioned gradient descent, where the preconditioner is implicitly optimized during pretraining. More recently, Mahankali et al. (2024) prove that a single self-atlention layer converges to the global minimum of the squared error loss. Zhang et al. (2024b); Wang et al. (2025) also analyze a more complex transformer architecture with a (lincar) multi-layer perceptron (MLP) or softmax after the linear self-allention layer, showing the importance of such block when pretraining for more complex tasks. In a related direction, Cheng et al. (2024) show that transformers can implement functional gradient descent to learn non-linear functions in context, further strengthening the view of ICL as gradient-based optimization.
Recent works have also raised important critiques of the ICL to GD hypothesis, questioning both its theoretical assumptions and empirical applicability. For example, Shen et al. (2023; 2024) point out that many theoretical results-sus as those in Von Oswald et al. (2023)-reth on overly simplified settings, including linearized attention mechanisms, handcrafted weights, or order-invariant assumptions not satisfied in real models. Giannou et al. (2024); Fu et al. (2024) demonstrated that in a multi-layer self-attention sctting, the intornal iterations of the Transformer conform more closely lo the second-order convergence speed of Newton's 's Method. Therefore, the interpretation of ICL needs to be examined under more realistic assumptions.
In this work, extend the above lines of rescarch by emphasizing more realistic priors, specifically, non-zero prior means. While Zhang et al. (2024a); Mahdavi et. al. (2024) explore broader prior distributions by analyzing covariate structures or modify the distribution of input feature, our focus instead lies on the interplay between a non-zero prior mean and the capacity of LSA to emulate GD. We note that while Ahn et al. (2023); Mahankali et al. (2024); Zhang et al. (2024b) provide compelling theoretical analyses, their work does not include experimental validations. In doing so, our study builds upon and generalizes the prior-zero analyses found in Von Oswald et al. (2023); Ahn et al. (2023), illuminating new challenges and insights that arise when priors deviate from zero, both theoretically and empirically.
2 Preliminaries
We use E Rd and y E R to denote a feature vector and its label, respectively. We consider a fixed number of context examples, denoted by C' > 0. We denote the context examples as (X,y) e RCxd xxd x RC, where each row represents a context example, denoted by (x,, v1), i E [C]. That is,
y/
X def
E R Cxd and y def
E R.
(1)
[
UC_
To formalize an in-context learning (ICL) problem, the input of model is an embedding matriz given by
def
X
T
Xq] E ((d++))(C++)
E
yq
where X, E Rdd is a new query input and By is an initial guess of the prediction for the query xg. The model's 's output corresponds to a prediction of y E R. Notice that the embedding matrix in equation 2 is
3

---

Published in Transactions on Machine Learning Research (12/2025)
slight extension to the commonly used embedding matrix, e.g. presented in Von Oswald et al. (2023), where Dg is set to be zero by default. Its interpretation will be clearer in the next two sections.
Linear regression tasks. We formalize the linear regression tasks as follows. Assume that (X,y, are generated by:
First, task parameter is independently generated by w ~ N (w., Ia), where N (w., Ia) is the prior, and w, is called the prior meon.
The feature vectors independently generated by Xq. X1,... xe 1.d. N( Ia).
Then, the labels are generated by y = (w, xg); and y: = x,), E [C], with no noise.
Here, w, Rd is fixed but unknown and governs the data distribution.
A linear self-attention. We consider a linear self- attention (LSA) defined as
fLSA (((+))x(C++) +R,
E E+ WW (E" (WK) (WK)TWEE) -1,-1'
(3)
where WK Wº, WP, W R(d-1)x(d+1) are trainable parameters, [ J-1,-1 refers to the bottom right. entry of a matrix, and WM def Ic 0f mask matrix. Our lincarized self-attention removes softmax, LayerNorm, and nonlinear activations. Consequently, the update is affine function of low-order context aggregates (e.g.. XTX, X v), which enables closed-form analysis of initialization effects while preserving the in-context lcarning sctup.
0
ICL risk. We mcasure the ICL risk model f by the mean squared error,
R(f) def (((((E) - 1)
(4)
where the input E is defined in equation 2 and the expectation is over (equivalent to over X, y, and xg) and y. The performance of different models are characterized by the ICL risk.
3 Muldi-Head Linear Self-Attention
In order to improve the performance of linear self-attention (LSA), we consider the multi-hoad extension. Let e N be the number of heads. Similar to equation 3, we define the output of each transformer head as
head, E) def W.WW
WWW EW (w ) C
W) h
E (H],
(5)
where WK O W and wy are trainable parameters specific to the h-th head. The multi- LSA & function defined as
K
fH-LSA (E)
def
E Σ
n=l head,, (E)
(6)
-1,-1
Standard multi-head attention concatenates head outputs and applies a linear projection Wo Algebraically, Concat(headı,, head W equals Σ W. head, after absorbing Wo into per- head projections (W,"). We therefore use a sum without loss of generality, keep the model dimension (d+1), and retain per-head contribution after reparameterization.
We emphasize that both the single-head LSA fLSA and the multi-head LSA fH-LSA share a common structural property: the bottom-right entry of the output matrix corresponds the prediction for the query point Ey which can be interpreted an initial guess Vy refined by an attention-basse update. In the special case of linear regression with zero prior mean, i.e. W. 0, the choice yy 0 introduces non-trivial prior for the
4

---

Published in Transactions on Machine Learning Research (12/2025)
initial guess, already observed by Von Oswald et al. (2023). The empirical role of this initial guess in the multi-head setting will be further analyzed in Section 5.1.3.
We denote by
FHSL def
{Wk W W w " h=
the hypothesis class associated with multi-head LSA models with H heads. Our first theoretical result establishes an invariance of the optimal in-context learning risk with respect to the number heads once it exceeds the Teature dimension.
Theorem 1. Let d E N, and consider the hypothesis classes F(a+1) LSA and F(a+2) LSA corresponding to maltti-head LSA models with H d and H + attention heads, respectiveky. Then
inf R(f) inf
R(f).
where R() is the ICL risk defoned in Eq. (4).
While the full proof of Theorem provided in Appendix A.1, we outline the key intuition here. Each attention head contributes a rank-one update to set of (d 1) matrices that fully describe the model. Collectively, these matrices live in space of dimension (d 1)3. A single head provides (d 1) degrees of freedom, so once the number of heads reaches d d+I. 1, the parameter space already has enough capacity to span the entire target space. In fact, with d+1 1 heads one can explicitly construct any target configuration, which means the model is already maximally expressive. Since adding further heads simply amounts to appending zero-contributing heads, the hypothesis class does not grow beyond d+1 1 heads, and the achievable risk remains unchanged. In Section 5, we provide empirical evidence supporting this theoretical result across a varicty of model configurations.
Relation to concurrent work. Theorem 1 is capacity statement for linear self-attention (TSA)): once the number of heads reaches H d+1, the hypothesis class and the attainable ICL risk no longer improve by adding heads. This contrasts with results for softmax attention, where (Cui et al., 2024) give exact risk formulas for single/mun magge//ault-head ICL, and show that as the number of in-conlext examples C' grows, both risks scale as O11/C) but multi-head achieves a smaller multiplicative constant when the embedding dimension is large an improvement in performance constants rather than capacity. Complementarily, (Chen et al., 2024) study trained multi-layer transformers and find that multiple heads matter primarily in the first layer, proposing a pepproces-then-opttin softmax and multi-layer architectures), whereas Theorem isolates an expressivity saturation specific to single-layer LSA.
mechanism; their conclusions concern learned utilization patterns (with
Next, we explore the convergence of multi-head LSA. Inspired by the analysis of Ahn et al. (2023), we analyze the stationary point the ICL risk for multi-head LSA functions.
Theorem 2. Let H E N and consider the hypothesis class FH-LSA of meulii-head LSA models with context size C 0o. Then the in-contert learning risk R(f) admits no non-trivial stationary point in parameter space. More precisely,
VR()) + for all f E FH-ISA
for every choice of parameters {WK, WR. WV. W except in the case where the prior mean vector vanishes, w, = 0.
Theorem states that when the context size C -+ oo, the gradient of the multi-head LSA's ICL risk R(JH-LSA remains non-zero for the entire parameters space long w. # 0. This result highlights fundamental limitation of multi-head LSA under non-zero priors: no choice of weights WK w
h W w. and W with h e [H] can minimize the ICL risk in the infinite-context limit. See Appendix A.3 for a detailed discussion of how the results change when the context size C is finite.
Relation concurrent work. Although previous works such as Ahn et al. (2023) and Mahankali et al. (2024) provide analytical solutions corresponding to stationary points of the ICL risk, these results are derived under the assumption that the prior mean w. 0. In this special case, the gradient of the ICL. risk can
5

---

Published in Transactions on Machine Learning Research (12/2025)
vanish, allowing the existence of stationary point. Our analysis generalizes this observation: we prove that when w, * 0, the gradient of the ICL, risk remains strictly non-zero for all weights as context size C + oo, thus precluding the existence of stationary points. We adopt C' + oo as an asymptotic approach, as done by Zhang et al. (2024a); Huang & Ge (2024). Our analysis targets the asymptotic regime C + oo, where finite-sample correlation lerms vanish and the gradient remains strictly non-zero for w, * 0, hence no non-trivial stationary points exist. For fixed, finite C, an additional finito-sample correction decaiiinn inversely with C' can partially cancel the leading gradient, producing apparent stationary points or plateaus in practice. As C' grows, these effects fade and the behavior converges to the asymptotic prediction, matching our experiments.
Finally, even though such a stationary point exists with finite context size, we still cannot imply that the stationary point is the global optimum, as the ICL risk multi-head LSA (((H-LSA) is not convex, presented in the following lemma.
Lemma 1. For ony H N, the in-contert learning risk
R(),
FH-LAA
is not conver in the parameters (WK, W W. Th=1'
Because R(SH-LSA) is non-convex, any stationary point that arises, even at finite context. sizes, does not. guarantee global optimum. In other words, one may encounter local minima or saddle points that satisfy the stationary condition without minimizing the overall ICL risk.
4 34- -Linear Self-Attention
To address the performance gap between one-step GD and multi-head LSA, we introduce y,SLSA a general- ization of single-head LSA.
4.1 Formulation of LSA
Our approach builds upon the GDD-transformer developed by Von Oswald et al. (2023); Rossi et al. (2024), which implements one-step GD in a linear regression setup when the prior mean w. is zero. The original formulation is defined by the weight matrices
0 w
9. w==we = Ia W = -
w =
where represents the GD step size. From the standard LSA formulation equation 3 with the given embedding equation 2, we derive
(E) = V4 - (w. -y )Xx:
(8)
where the initial guess y4 = 0 = w Xg is fixed for any query Xg. and the prior mean w. is zero. See the derivation of equation Appendix for the completeness. Notably, we retain the terms for yg and w. facilitate future extension non-zero scenarios. Rewriting the equation equation 8 with Be w Xg yields
fisA (E) (w, C x (XW, - -y) Tx,-
The red term represents the gradient of the least-squares loss in linear regression. Consequently, fLSA (E) becomes equivalent to linear function f(x,) w Xxy where w is the one-step GD update initialized at the prior mean W..
For the more general case with a non-zero prior mean w., we relax the condition on the initial guess yg- By allowing Ba to be a linear function of te, specifically w Xny we obtain the prediction of the linear regression task with a given query X,
W. ---2xx (Xw. Xy.
(10)
which still implements the one-step GD update. Given this, we now define LSSA

---

Published in Transactions on Machine Learning Research (12/2025)
Definition 3 (yg-LSA). We define Vy
LSA with a flezible initial guess embedding matrix
E
X
Xy
[(++))x(++)) with B4 w X,g
(11)
where w E IRd is trainable parameter and Yu is the initial guess. The y4 y--LSA function is defined as
fu. (x,, def fisA(Ew).
(12)
The Bg-LSA extends the standard LSA by introducing an additional parameter in the embodding, enabling better aligmment with the query's 's initial guess. The trainable parameters of By -LSA now include W W W w and w, with imputs X, y and Xg.
4.2 Analysis of U,-L
Similar the analysis of multi-head LSA, we first examine the stationary point of yg-LSA
Theorem For a Mo- LSA function in equation 12 with a non-zero prior mean w. and contert size C' oc, the weights (W W WP, WV, w.) in equation 7 with W. constitute stationary point of [R
Theorem asympt.ticc in the context length C: when C + xo, the gradient vanishes at the weights Eq. (7) with w w., For finite C, each gradient component differs from its infinite-C value by correction of order 1/C. Thus w, behaves as an approximate stationary point whose residual gradient (and the resulting bias) decays as C grows, explaining the small plateaus occasionally observed at finite C. Similar to multi-head LSA, we cannot conclusively determine that this stationary point represents the global optimum. This uncertainty comes from the non-convex nature of the L,A ICL risk, as established in the following lemma.
Relation to concurrent work. Unlike Ahn al. (2023) who show that single-layer LSA attains one-step preconditioned GD under zero-mean preor-mh 4 establishes that with non-zero prior mean, one-step GD is still recovered without MLP by introducing trainable query initialization D/ w Xy- In contrast. to Zhang et al. (2024b), where LTB (LSA+MLP) realizes identifies input-side initialization as the minimal mechanism that closes the ICL-GD gap within LSA.
via the MLP, our result
Lemma 2. The ICL risk of Ug LSA R(fu,-LSA) is not conver.
While the non-convexity prevents definitive proof of global optimality, our empirical investigations in Section 5.2 suggest an intriguing hypothesis. Notably, conjecture that the stationary point identified in Thcorem 4 may indeed be the global optimum. Empirical evidence indicates that the performance of one-step gradient descent serves as lower bound for -,SA
An additional noteworthy observation LSA's relationship to the linear transformer block introduced by Zhang et al. (2024b). Unlike LSA, LTB combines LSA with a lincar multilayer perceptron (MLP) component. Critically, the global optimum of LTB implements Newton step rather than one-step gradient descent. This approach fails to bridge the performance gap between onc-step GD and single-head LSS and requires significantly more parameters through the additional MLP, in contrast to -ySLSAss A's more parsimonious approach of introducing a single vector parameter w. See Lemma 3 in Appendix B for more details.
5
Experiments
For experiments in Sections 5.1 and 5.2, focus on simplified setting where the LSA consists of single lincar self-attention layer without LayerNorm or softmax. We generate lincar functions in a 10-dimensional input space (d = 10) and provide C = 10 conlext examples per task. We endow the LSA parameters with ICL capability by minimizing the expected ICL risk [[(fo(E) - y) over random tasks. Each training step is an Adam update of (WQ, K wV WP) (and w for Yo -LSA) using freshly sampled (X,y,,, y); at test time, no parameter updates are performed. We train for 5000 gradient steps. Further implementation details are provided in Appendix C.I.

---

Published in Transactions on Machine Learning Research (12/2025)
5.1 Multi-head LSA
5.1.1 Multi-head LSA with Varying Numbers of Heads
We investigate the ICL risk (evaluation loss) of the multi-hean LSA under dif- ferent numbers of attention heads in the setting of non-zero prior mean and By is fixed at zero (details in Table 1). Fig. 2a illustrates the loss curves over the course of training for several head configurations, while Fig. 2b summarizes the final evalu- ation losses
Namber leace
Loaa Mamber dieedded
Las moadke 1
Marbeer
+D Lrdo
function of the number observe of heads. From these results, that increasing the number of heads up to d+1 1 (here d = 10, see Fig. 2b) sub- stantially enhances the in-context learning capability of multi-head LSA, as reflected by pronounced reduction in the final evaluation loss.
1000
2000
3000
4000
8
10
12
14
Training
Number Heads
(a)
(b)
Figure 2: Training loss of multi-head LSA with different numbers of attention heads, In (a), we visualize the training loss curves for models with different head configurations, each curve shows the expected ICL. risk during parameter training (Adam updates of (WQ, W WV W); no updates at test time). While (b) shows the final trained loss function of the number of heads.
However, adding more than d + heads yields negligible further improvement, in- dicating a saturation effect beyond this threshold. This confirms our results in Theorem Notably, even at d + heads, the multi-head LSA model does not converge to the one-step GD baseline loss, suggesting that while additional heads can capture richer in-context information( nnformation(Crosbie Crosbie & Shutova, 2024), they alone are insufficient for achieving full parity with the one-step GD performance in non-zero prior means setting. In other words, one-step GD loss serves as strict lower bound of the ICL risk for multi-head LSA empirically.
5.1.2 Effect of Prior Mean w. in Multi-Head LSA.
We investigate how the prior mean w.: which represents the mean weight of the generated linear function, affects the per- formance of multi-head LSA when the number of heads is fixed above + and yy is fixed zero. Fig. 3a shows the loss curves for different values of ||w.]l: while Fig. 3b presents the final trained loss as a function of ||?.
Loss
82.8
4-
Our results demonstrate that even when the number of heads sufficiently large (i.e., 2 d+1, 1 reaching the optimal multi- head LSA configuration), multi-head LSA only matches the loss of onc-step GD when the prior mean w. is zero. For non-zero prior means, a systematic gap re- mains between Multi-Head LSA and one- step GD. Furthermore, this gap increases linearly with the squared lo norm of the
2.5
0
3000
2000
3000
4000
10
Training Steps
[w.
(a)
(b)
Figure 3: Training loss of multi-head LSA under different
prior means w., (a) Training loss curves for different values of
[[w]]] I]. (b) Final trained loss as a function of [|w,||? Multi-head T.SA matches the one-step GD loss only when w, = 0; for w, A
the gap grows approxinately linearly with ||, l|.
8

---

Published in Transactions on Machine Learning Research (12/2025)
prior mean, |w, II? indicating that the prior mean significantly impacts the optimal loss and that larger deviations from zero result in a larger discrepancy from the GD baseline.
8
Bacal 00.0
Acsal F F-.0
Dazml 5.00
-100
Dazmml -253
7-
Fazam 3.00
14
Fream -..0
2.7
Dream 4.47
Fa 4.80
12-
6-
Fr 5.30
Fuzmm 5.00
Dusse 6075
Frazml 5.000
15
20
Lo55
5-
Loss
Loss
CD Loss
12
4-
10
3-
4-
2
1000
2000
3000
4000
1000
2000
3000
4000
Training Steps
Training Steps
1 Irn
Figure 4: Training and final loss of multi-head LSA under different initial guess configurations. Left Training loss curves for various 111/ bi Middle Final trained loss as function of llve bi b Right Upper Training loss curves for various || and Right Lower Final trained loss as function of "uessll Multi-head LSA reaches the GD loss only when both the linear guess component and the bias vanish (Vg w Xy and no offset).
5.1.3 Effect of yg in LSA
To investigate the effect of the initial guess Vy, contained in the embedding matrix equation 2 on the in-context. learning ability of multi-head LSA, we decompose My into two components:
y4 Xg guess biass
We set the prior mean w. to zero and number of head is d + 1, then conduct two separate experiments: (1) varying Yg_gees while fixing blas 0, and (2) varying blas while fixing Yg_guess O. This al- lows to isolate the contribution of each component, to the model's behavior.
Lone
3.0
Gradent Norm Sim larty
Eiftrends
2.5
blas
Nor Differesse
Simil anty
Loss
2.5-
-0
2.0
As shown in Fig. 4, multi-head LSA only converges the same loss as one-step GD when Ya_guess (i.e., equal to the prior mean) and blas 0. In all other cases, systematic gap remains between the loss of multi-head LSA and one- step GD. Moreover, this gap is di- rectly proportional
0.5
0
1000
2000
3000
4000
0.0
Training Steps
Training Steps
(a)
(b)
Figure 5: Training loss and sensitivity analysis of LSA. (a)
Training loss curves of LLS and one-step GD. (b) Model behavior metrics including prediction norm difference, gradient norm difference, and cosine similarity.
g 2| (the squared fo-morm of the guessed component) and [g bi (the squared bias term). These findings suggest that deviations in Bg from the

---

Published in Transactions on Machine Learning Research (12/2025)
optimal initialization introduce persistent discrepancy in multi-head LSA's performance relative to one-step GD, regardless of the training of multi-head LSA.
5.2 U4-LSA
In this section, we aim to empirically validate whether Vy-LSA, introduced in Section 4, aligns with onc-stop GD across different prior settings. Fig. presents the training loss of yy-LSA. Throughout Fig. 5a the dashed "GD Loss" curve the in-context risk of the predictor obtained by one GD step initialized the prior mean Wo w.: W1 wo T - y), (GD(X,) x, Wi, and the plotted baseline is RGD-Istep E ((GD(xx)) y)
In Fig. 5a, we compare the convergence of Ly one-step GD, demonstrating that regardless of the prior configuration, 3,-LSA effectively matches the GD solution. Fig. 5b provides detailed evaluation of prediction norm differences, gradient norm differences (defined in Appendix C.2), and cosine similarity between the models. The results confirm that yy-LSA exhibits strong alignment with one-step GD in both loss convergence and gradient analysis.
5.3 LLM experiments
Through theoretical and experimental analysis, we hypothesize that providing an initial guess for the target output during the ICL significantly improves the model's ability to refine prodictions. Specifically, we posit that initial guesses as prior for optimization, guiding the model to more accurately. To validate this hypothesis, we conduct experiments leveraging widespread LLMs, demonstrating the efficacy of initial guesses in improving prediction accuracy,
Our experiments utilize Meta-LLaMA--31138 Instrucct (Grattafiori et al., 2024), Qwen/Qwen2, S0dven2.5-7B-nnsiiee (Yang et al., 2024; Team, 2024) and the STS- Benchmach dataset (English subset) (May, 2021). Each prompt is presented in conjunction with a context comprising 10 labelled examples, where each example included a pair of sentences and its correct similarity score. A lightweight guess model is used to generale initial guesses for both the query and context examples. These guesses are included in the prompts provided to the LLM model, framed as prior guess. The model's task is to predict a similarity score for the query pair, explicitly improving upon the initial guess. For evaluation, we calculate the mean squared error (MSE) between the predicted and true similarity scores, comparing the models with and without initial guesses. More details are in Appendix C.3.
Ink al Guess Enor
Lama Guess Errrr
Mean Squared Error
CWen Baseine Emor
ne Error
6 Guess Model Training Epoch
10
Figure 6: Error Comparison Two pre-trained models show consistently improved ICL performance on a sen- tence similarity task when prompted with non-trivial initial guess.
The results demonstrate that the inclusion of initial guesses signif- icantly enhances the performance of LLMs in ICL tasks. As shown in Fig. 6, incorporating initial guesses into the context reduce MSE under all experimental conditions. Comparative analysis of the LLaMA and QWen models further underscores the generality of this approach, as both models consistently benefit from the inclusion of initial guesses. These findings follow our hypothesis that initial guesses enhance ICL by providing an initial guess for refinement.
6 Conclusion
In this work, we have theoretically and empirically studied the extent to which multi-head LSA approximates GD in ICL, under more realistic assumptions of non-zero prior mcans. Our analysis establishes that while increasing the number of attention heads to + suffices to reach the minimal ICL risk the linear setting, the model fundamentally fails to reach stationary point when the prior mean is non-zero and context. size
10

---

Published in Transactions on Machine Learning Research (12/2025)
grows. This limitation is further connected with the initial guess y4, whose misalignment with the prior induces a persistent optimality gap, even when the number of heads is sufficient. To solve this, we introduce V,LLA, an LSA variant with a trainable initial guess, and show both theoretically and empirically that it bridges the gap between LSA and one-step GD in linear regression. Finally, we illustrate that incorporating an initial guess also benefits ICL large language models, showing how this approach can be also used in more common settings.
Limitations. While our analysis is limited to linear regression tasks and simplified architectures without nonlinearities, normalization, or softmax, these assumptions are standard across much of the theoretical literature on in-context learning and mechanistic interpretation of transformers. The theoretical results rely the infinite-context limit, which, although analytically tractable, diverges from practical settings where context size is finitc. Additionally, while yy-LSA closes the gap with onc-step GD in controlled experiments, its applicability to complex real-world tasks remains contingent effective mechanisms for estimating learning initial guesses. The LLM experiments suggest empirical benefits, but further exploration is required to assess gencralizability across diverse tasks, model familics, and training regimes.
References
Kwangjun Alm, Xiang Cheng, Hadi Daneshmand, and Suvrit. Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Therty-seventh Conference on Neural Information Processing Systems, 2023. URL 7, and 16.)
(Cited on pages 5,
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Nyya Sutskever, and Dario Amodei. Language models are few-shot learners. H. Larochelle, M. Ranzato, R. Hadselll M.F. Balcan, and H. Lin (eds.), Advances in Newral Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL https: https: p/ttope///proceedings cwtppper.ccppp.cccc http:///e..aaaaooo//o//o///////0
pdf. pdf. (Cited
on pages 1 and 3.)
Xingwu Chen, Lei Zhao, and Difan Zou, How transformers utilize multi-head attention in in-context learning? case study on sparse linear regression, 2024. URL. https://arxiv. orw.ooggg/aa.......o..oo... 04532. (Cited on page 5.)
Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.). Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 8002 8037. PMLR, 21-27 Jul 2024. (Cited on page 3.)
Joy Crosbie and Ekaterina Shutova. Induction heads as an essential mechanism for pattern matching in in-context lcarning, 2024. URL attps://arxiin wwtrriivv..or.eeeeo org/abs/////eeeoooorogggooeeeeer /207 07011. (Cited on page 8.)
Yingqian Cui, Jie Ren, Pengfei He, Jiliang Tang, and Yue Xing, Superiority of multi-heal attention in (Cited on page 5.)
in-context linear regression, 2024. URL https: /ttpss////xxvoo// abs wwwpaabss.....//
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li, Heming Xia, Jingjing Xu, Zhiyong Wu, Baobao Chang, Xu Sun, Lei Li, and Zhifang Sui. A survey on in-context learning. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference on Empirical Methods Natural Language Processing, pp. 1107-1128, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10.18653/ 20w22...mllmmmmmee.mm 2024. amnlp-main 64/. (Cited page 1.)
URL https: h/tp:///aaaaaaogg////// org/
11

---

Published in Transactions on Machine Learning Research (12/2025)
Fabian Falck, Ziyu Wang, and Christopher C. Holmes. Is In-Context Learning in Large Language Models Bayesian? A Martingale Perspective. In Forby-first International Conference on Machine Learning, 2024.
URL https://openreviie http:///pppaaeeeeeeeee///tottaap wett//oorrrrmttttt/00000 (Cited on page 3.)
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn to achieve socond-order convergence rates for in-context linear regression, 2024. URL htaps://arxiv.or/aabb/////2 17086. (Cited on page 3.)
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? A case study of simple function classes. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Newral Information Processing Systems, 2022. URL https: /tttps///pennrvieev.o net/ 2rrumidd=1w2222eet.. (Cited on page 2.)
Augeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos and Jason D. Lee. How well can 03183.
transformers emulate in-context newton's 's method?, 2024. URL (Cited on page 3.)
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Authony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, and al. The llama 3 herd of models, 2024. URL nttps://arxis org/abs/2407 21783. (Cited pages 10 and 29.)
Ruomin Huang and Rong Ge. Task descriptors help transformers learn linear models in-context. In ICML 2024 Workshop on In- In-Contert Learning, 2024. URL https: http:p:://peerevveeeeeeet/// http://ww..etoooouuoadddooSSoDD
(Cited on page 6.)
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint urXiv:1412.69900 2014. (Cited on page 30.)
Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One Step Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention. In The Twelfth International Conference on Learning Represenlations, 2024. URL http:::/opperrevis..eeee///// http://wwweettormmmmmdaaaaaaoaoa (Cited on pages 3 and 5.)
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Revisiting the equivalence of in-context learning and gradient descent: The impact of data distribution. In ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 7410 7414. IEEE, 2024. (Cited on page 3.)
Philip May. Machine translated multilingual sts benchmark dataset., 2021. URL hhilipMay/sssb-multte multi-mt. (Cited on pages 10 and 29.)
Madhur Panwar, Kabir Ahuja, and Navin Goyal. In-Context Learning through the Bayesian Prism. In The Twelfth International Conference on Learning Representations, 2024. URL https: htttss://openrviee...r forum?:d-HXSujdasSonn (Cited on page 3.)
Nils Reimers and Iryna Gurevych. Making monolingual sentence embeddings multilingual using knowledge distillation. Tn Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 11 2020. URL on page 30.)
2004. 09813. (Cited
Simone Rossi, Rui Yuan, and Thomas Hannagan. Understanding in-context learning in transformers. In
The Third Blogpost Track at ICLR 2024, 2024. URL (Cited on pages 3 and 6.)
Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Position: Do pretrained transformers learn in- context by gradient descent? In Forty-first International Conference on Machine Learning, 2024. URL https://openrevier:
http:://ppaareeeeeeeeee//rrrrraa net.foouuridddaawwccaaaaa (Cited on page 3.)
12

---

Published in Transactions on Machine Learning Research (12/2025)
Shiyu Shen et al. Towards understanding in-context learning via training dynamics. arXiv preprint
arXiv:2305. 14743, 2023. (Cited on page 3.)
Qwen Team. Qwen2.5: A party of foundation models, September 2024. URL https://qvenln https////qeennnngitubu..... blog/quen2.5/ (Cited on pages 10 and 30.)
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 35151-351 35174. PMLR, 23 29 Jul 2023. URL press/v202/von- www.ooooooo....aaarraeoea html.. (Cited on pages 2, 3, 4, 5, 6, and 22.)
mml.
Aaron T Wang, William Convertino, Xiang Cheng, Ricardo Henao, and Lawrence Carin. On understanding attention-based in-context learning for categorical data. In Forty-sccond International Conference Machine Learning, 2025. URL https://openreviee http://ww.eeetoooom/dd////a//aaa
(Cited page 3.)
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An Explanation of In-context Learning Implicit Bayesian Inference. In International Conference on Learning Representations, 2022. URL (Cited page 3.)
http:////ppeeeeeeeeeeeeeerrrrrra https://openreview nettfforumiiddRa/FFOO
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng Hc, Junyang Lin, Kai Dang, Keriing Lu, Keqin Chen, Kexin Yang, Mei Li, Mingleng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yugiong Liu, Zeyu Cui, Zhenru Zhang, and Zhihao Fan. Qwen2 technical report. arXiv preprinl or 24:: 10611, 2024. (Cited on pages 10 and 29.)
Naimeng Ye, Hanming Yang, Andrew Siah, and Hongseok Namkoong. Pre-training and In-context Learning IS Bayesian Inference la De Finetti. In JCLR 2024 Workshop on Matheratioal and Empirical Understanding of Foundation Models, 2024. URL attps://openrevieet http:////eeeeeeeeeeeeeeee///aaaa wwt..e//ooouuuniidttuutuaooo (Cited on page 3.)
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-555 155, 2024a. URL http://jmlr wwwlll.ogg//pppee org/pappers/255222-- 1042. html.
(Cited on pages and 6.)
Ruiqi Zhang, Jingfeng Wu, and Peter Bartlett. In-context learning of lincar transformer block: Benefits of the MLP component and one-step GD initialization. In The Thirly-eighth Annual Conference on Neural Information Processing Systems, 2024b. URL on pages 3, 7, and 22.)
(Cited
Yedi Zhang, Aaditya K Singh, Peter E. Latham, and Andrew M Saxe. Training dynamics of in-context: learning in lincar attention. In Forty-second International Conference on Machine Learning, 2025. URL /ttpss//operreviev..
httppp//pppppareee...eee///eeett http:/ooouumttddaaa//l11..... (Cited on page 3.)
A Proofs of Section 3
For the sake of completeness and seff-containment, we restate the theorems and lemmas shown in Section 3 and provide their full proof in this section.
13

---

Published in Transactions on Machine Learning Research (12/2025)
A.1 Proof of Theorem 1
First, let's redefine the notations used in Theorem and restale the theorem. We write the input of model as an embedding matric given by
x T
E
def
X yq
E
[g(++)x((++),
(13)
where X, X,y, Xyx My are defined in Section 2. The multi-head linear-self attention (LSA) function is defined as
def
H head,, (E)
fH-LSA(E)
E
(14)
-1,-1
where the output each transformer head is defined as
head, (E) def C WW
EW" (ETN"(EETTW
h E [M]
(15)
The trainable parameters WK h W W and W are specific to the h-th head, and WM def mask matrix, to ignore the query token when computing the attention scores. Let's define by
[e
0
0 is
FH-LSA def LAAL
{ [ W W. W " A-1
the hypothesis class associated with multi-head LSA models with H heads. Finally, we measure the TCL. risk of model f by the mcan squared error,
R(f) E((((E) - 3)
(16)
where the expectation is taken over the data distribution (and effectively over the embedding matrix E defined in equation 13).
Now we are ready to restate and prove Theorem 1.
Theorem 1. Let d E N, and consider the hypothesis classes F(d+1)-isA and F(d+2) LSA corresponding to multi-head LSA models with H = d 1 and H = d 2 attention heads, respectively. Then
-SA
inf
R(f)
inf
R(f),
FEF(A=2)-1SA
where R(f) is the ICL risk defined in Eq. (4).
Proof. To simplify the notation, lot's introduce a couple of additional definitions. For cach head h E [H]. the product of the output projection W and the value projection W can be written without loss of generality
as
w.w. def b
R(d+1)x((+1),
where by IR++1 is the last row of the matrix, and the block denotes entries that have no influence on the ICL risk. Then, let's rewrite the product of the key and query matrices as
(W W def A ((d+1)x(++1)
and denote its column decomposition by
A = |a]
a
where at IRd+1 for each E (d 1).
14

---

Published in Transactions on Machine Learning Research (12/2025)
With this notation, the contribution of all heads to the attention mechanism can be expressed in terms of the matrices
M
(a
R(++))((++)
i [d 1].
Each M, a (d+1) x (d (d+1) 1) real matrix. The space of such matrices, R(++1)x(d=1) has dimension (d+1)² The collection
(M1, M2,..., Mc+1)
is thus an element of the Cartesian product
(R((+I)x(a+1))d11
with dimension dim (R(++))((+))) = (d + 1)². Hence, the set of all possible tuples (Mi...., Ma+1)
can be identified with a vector space of dimonsion (d +++1 1)
We now compute the number of parameters available per head. For a fixed head h, the parameters that influence the construction of M, are (1) the vector by which contributes (d+1) free parameters, (2) the family of vectors a ....., 12411 which contributes (d + 1)(d + 1) free parameters. Therefore, in total one head contributes (d + 1) + (d + () (d + 1)(d+ 2) degrees of freedom. With H heads in lotal, the dimension of the parameter space 22 is dim(S1)) = H (d + ()
Suppose H 2 d Then
(d 1)(d+) 2 (d+1)(d +)(++))
Since (d+ 2 (d+ +1+1), 1), we obtain
H (d () 2 (d (d+1) 1)
This inequality shows that, when H 2 d+1, the parameter space has dimension at least as large as the target space. In particular, there is no dimensional obstruction to surjectivity the mapping from parameters (b,, a to matricos
To demonstrate that the mapping is indeed surjective once H 2 d+ 1, we now construct explicitly any desired collection of matrices (4....,,a-i))
Fix i E [d + 1]. Let e e Bat1 denote the standard basis vectors of Rd+1, For each h E [d + 1], set
b eh. a = [M],.
where [.]h] denotes the h-th row of the matrix M. For h > d + 1, we may set by = and at 0, so that those heads contribute nothing. With this choice of parameters,
d+1
Σ (a > eb (M.) M..
Thus, every M, is exactly reproduced, and therefore every tuple (MI,..., M)+1) is realizable when H 2 d+ 1.
We have shown that with H = d + 1 heads, the model can realize any element of the target space, and therefore the hypothesis class is saturated. Adding additional heads H d cannot enlarge the class of realizable functions. For this reason, for any H 2 + 1, we have
inf JEF(112) LSA
R(f) < inf
R(f).
JEF(i1) LSA
Finally, observe that F(a+1)-LSA S F(42))-SAA, since (d + 1)-head model can be viewed as (d + 2h-h model with the additional head parameters set to zero. Consequently, it follows that the only possibility is that
inf FEF(a-1)-1sM
R()) =
inf
R()),
which concludes the proof.
15

---

Published in Transactions on Machine Learning Research (12/2025)
A.2 Proof of Theorem 2
Theorem 2. Let H N and consider the hepothesis class FH-LSA of mulli-head LSA models with context size - oo. Then the in-contert learning risk R(f) admits no non- trivial stationary point in parameter space. More precisely,
VR()) # for all J E FF
for every choice of perameters (WA, W W. W except in the case where the prior meun vector vanishes, w. 0.
The proof of Theorem 2 is based on the analysis of Ahn et al. (2023).
Proof. Step 1: Simplify the risk function and compute its gradient
We first derive explicitly the expression of multi-head LSA's ICL risk and simplify it. The key idea is decompose the ICL risk into components. That is,
R(fH-LSA)
equation
((HH-LSA(E)) -)
with y w Xg and NNww., Ia),
2
equation
E
E +
head (E)
W Xg
-1,-1
27
H Σ
equation
E
E +
WWWW
EWN
E (W))
W E
w X2
Since the prediction of fH-LSA is the bottom right entry of the output matrix, only the last row of the product W.. wy contributes to the prediction. Therefore, we write
W.WK def
E (R(++))x(++))
byj
where by Rd+1 for all h E [H], and * denotes entries that do not allect the ICL risk.
To simplify the computation, we also rewrite the product ( Tw and the embedding matrix as
(Wk)
del A E ((d+))x(d+1)
del
zo zct] E R(d+1)x(C+1)
E
Z
Z2
where
A
def
a
a
a
with at. a Rd+1
2;
def
Xi
E
Rd+1
for all i
E [C], and 20+1
def Xx yg
E Rd+11
yi
We define
G def 1
Σ z,z
1
EWEME E ((d+1)((4+1)
and def = +c,
C
where Rdd (O, Ia) is the noise.
16

---

Published in Transactions on Machine Learning Research (12/2025)
Then the ICL risk can be written
R(JH- LSA)
f1
+
by GA - WT Xx1
h=1
14 +
Σ
bj G a
ag
a/t1
w Xa
E
Σ bj
Ga xp
+ bb Ga
-wil -wixx Xg
where [ the i-th coordinate of the vector Xy-
Furthermore, know that, for all h [H] and E [d 1],
bj Ga
Tr Ga Tr (G (G,
where v) Tr (UVT) is the Frobenius inner product for any squared matrices U and V. Hence, by using the linearity of the Frobenius inner product, rewrite the ICL risk as
R(HH-LSA)
U4 + Σ
(a
Uq +
Σ
ΣΣO
b,
w(ll) xp(1]
h=1
E
G, Σ ba
baletee)
11
Σ
G. Σ (a())))))
m0(o]
xg(i]
f=1
where w[[] is the i-th coordinate of the vector W.
By reparametrizing the ICL risk, using composite function, we have
(G+ R(SH-15A) E G. x. ( (a+ (G, Ma+1) Ua + (G, w(() ( x,,1)
)
27
(17)
where
M; def
by (a.)) T E (((+1)x((++))
for all i [d 1].
h=1
Recall x, (O, Ia). Thus, both G and w are independent to x1 for all i [d], and x,li] NO 1) are i.i.d.
Expanding equation 17 yields
R(JH-ISA) Eo [1+ (G, Mat1)) vi]
Σ
((G, M.) w(()²] Ex, [x][]]]]
G.w
= -Ba( (G, Ma vi]
Σ
((G, ((G,M) M.) - w(b1)]]
+
G
1=1
(,(M)),
(18)
17

---

Published in Transactions on Machine Learning Research (12/2025)
where
((M)) def E G (i) -[1)]]
for all i E [d]
def Eg (1 + (G. Ma+))) vu]] (a+i(Ma+1)
Thus, the ICL risk equation 18 is decomposed into (d+ 1) separated components L2 with i E [d + 1]. Each component function of M, To compute the gradient of ((H-LSA)), we can first compute the gradient of each component. with respect to M, for i [d]. That is,
(M,, = 2B [[G, M,) G| - a, [[]]] V Mati (Ma-1) = 22 [[1 + (G, Ma+1))
for i [d]
(19)
(20)
Step 2: Compute B[ M, G], E [ in equation 19
Recall that W ~ ~N(W,,, Ia) and xj iid (0, Ia) are independent for all j E [C], Dj w x,, TT,, and w =, +
with é ~~ (0, Ia).
For G, [W]]]] in equation 19 with i E [d], we have
E G W 1
Σ
E
[[] x,x,
B [W]]] yjXy]
E-
[w[i] [i] vÿ
w
w.x,
In particular, for each block of the above matrix, we have
E w.x,
[ Xyx,1 = Ex [w]]]] [x,x]] w. Ad
E x,
[ yクXد] = B..x, [W]i] [W]]][Txxxx)) X,X}]
= Ee,x, [(w.[i] + clil) + c) x,x,] =w,[w. + +
E- W.X,
W [1] 191
-E w.x, [W]]] x,x} w] E [W]]] w
= E. [( [iw + E [i)) w, + é) (w. + <)] = [.]][[/w.. + d 2),
where e; denotes the standard basis vector with zeros in all coordinates except the i-th position, where the value is 1.
Combining the above three components, we have
w [a E G W[G [w. + C.)
w + e; [[ 1| + d
(21)
Now we compute E [[ M.) G| for i E [d].
We start by calculating the expected value of the product of elements in matrix G. That is, for all m, n,P, q E (d 11,
C
E Gm Gpyl
2z [m]z, [n]zapppza [g]2
C
where Gmn is the value of matrix m-th row and n-th column position for all m,n E [d + 1]. By expanding the summation, we have
E |mm Gpyl
Blgysmllyynnaaaaaaalllll C E (m Is jtk
Σ
E
C(C -1C-1)
B [zze/p/z2ll + E iz [m]zz] [m] C m
2 E [z) [m/2 [n]] E (22
when C - oo.
To compute E. [z (m)zi [n]]]
18

---

Published in Transactions on Machine Learning Research (12/2025)
1. For m, E [d], have 21(m] x1 [x(n and z [1[n] x1[n]. Thus, E [z [m//z//] [n]] 6mns where ó is the Kronecker delta.
2. For E [d] and d we have z1 [x]] y1- Thus, E [z] [m] [ra]] = EE [xx [m|x]] w] [m[]].
3. For d + 1, we have E 22 (m) [v]] = E x w = B
||w.|1 + d.
We denote
M
La w
||w-|l (ld++1)x(++))
(22)
By using equation 22, when C² - 0o, we have
E [G mn Gpg] Mmn Mpn
(23)
By linearity of the Frobenius inner product, we have
B || GJ = MM) M. Combining the above equation with equation 21, equation 19 becomes
(24)
W.
w. t
(MM 2 (M, M,) M - [(w, [i]. e,)
w. + d 2)
=2(M, M,) M 211 M 2N (2 (M,, M,) 2w. ( 2N,
(25)
where
def 0 e, 2w
[il
Notice that M full rank and the rank of N smaller or equal to 2. Thus, for any M; E (R(++))(d++1) we have
M (M,)) #0.
A.3 Finite-context corrections and dependence on the context size C In the proof of Theorem 2 we work in the infinite-contxtt limit C' - oo and use
E E Gpq] Mmn
(26)
where G E (dd+))x(d++) defined in equations (22) (24). For completeness, we now make explicit how equation 26 is obtained from the finite-C expression and how the resulting gradients are modified when is finite.
Recall that
G
def 1
2iz
Zj
def
Xj ff
E
Rd+1
C J=1
For any indices m, m.n.p.q [d+ 1] we have
1 C
Gmm
Σ
zj
Gpq
Σ
z (pl //////
C
j=1
k=1
and therefore
1
Gmn Gpq
2j [m |zj
nn l [p) (q).
C
J=1 k=1
19

---

Published in Transactions on Machine Learning Research (12/2025)
Taking expectation and separating the cases j + k and j k yields
E
mm ppq
[3 m 2j (n/ps///
C
C
j-k=1 jfk
C(C
1)
Elz1
[m 21 [n]
Eizz|p/zzll
Elz [ n/z1
[p|z1 [q]]
C
-1-C)
1
Mmn M +
M
C Tmmpq:
1
(27)
where we introduced the fourth-order tensor
def E [m]zm///zzz [p]z]
T mmppq
Equivalently,
E Gmm Gpq] = Mmn Mp +
C mmpq:
mmpq
del
T mmpp
Mm M pq:
(28)
that in tensor form
E|G & G|
M M
A.
C
Using the linearity of the Frobenius inner product, the quantity that appears in equation (19) can be written a8
E GG = (M, U₁.
for some matrix U, p(d+1)x(d+1) whose entries linear combinations of the tensor coefficients Amapy- Substituting this expression into equation 28 and combining with equation (21) gives the finite-C gradient
2
(M,,) (2(M, 2ww - 2N
C U,
(29)
Equation (25) corresponds exactly to the leading term in equation 29 when C' oc, since in that regime U is bounded while the U; correction vanishes. For any fixed finile C, however, the additional term & U provides an O(1CC) perturbation of the gradient. This perturbation can partially cancel the leading term (2M,M;) 2w" [i)M 2N and may create (non-global) stationary points in parameter space, which is consistent with the non-convexity discussion.
Implications for Theorem 4. The same finite-C decomposition is also relevant for the proof of Theorem 4 in Appendix B.3. There, the gradients with respect to the parameters b, aj, a4+1 and v'j] are expressed in terms - of expectations of products involving G (see the expressions preceding the verification that b 1
aj
ej a4+1 = 0, = w form stationary point in the C limit). Using equation 27 equation 28, each
such expectation can be written as its infinite-context value plus an O(1/C) correction. Consequently, for finite C every first-order derivative at w wn takes the form
DR(Fu.-LSA 00
dR (fU) + 0 (
0 E {b, aj. a4+1, vi]},
00
where Roo denotes the risk in the limit C + co. Thus, w w an approximate stationary point whose residual gradient decays at rate (1/CC) as the context size grows. This clarifics how the exact correspondence with one-step gradient descent established in Theorem 4 approached increases, and how small but non-zero biases may appear in practice when C is finite.
20

---

Published in Transactions on Machine Learning Research (12/2025)
A.4 Proof of Lemma 1
Proof. From equation 25, can compute the Hessian of the function (((M), that is,
VM, (M,, 2M.
We verify that M is posilive semi-defnite, Indeed, let Rd and E R. We have
n
u
M
u
equation
u
u
w
La
l|w, 11
d
u
u
u + uw.
iu'
u
[ T 11 + u(I) w+ll + d)
||]| 2uw, u + u² (1[w, + d) ||u + ww.ll 1w +
du² 2 0.
Since M is positive semi-definite, we have the function C; is convex with respect to M,.
From equation 18, we know that
d+1
R(fH-LSA) Σ (.M))
i=1
Each function L, is function of M,. We denote
R(JH-LSA) = f(M .. .Ma+1).
Then the Hessian of the function f with respect to variables MI.... Ma is a block diagonal matrix, each block on the diagonal is VM, (M,. 2 0. Therefore, the function is convex with respect to Mi,..., MM++1
Lastly, M; Σ ( for [d + 1]. To simplify it, we can consider only head. That is, M. = (a bilinear function, which is known to be not convex with respect to bj and ay.
To conclude, the ICL risk R(fH-LSA) is a composite function with convex function and non convex functions, which implies that R(JH-LSA) is not convex.
B Proofs of Section 4
B.1 Derivation of equation
Here provide the derivation of equation 8. Recall
WV
0
w
W
[
WP
- I
-1
0
C
So we get the LSA simplified as From the standard LSA formulation equation 3 with the given embedding in equation 2, have
def
def
X
xg
K
0
def
0
W
x
w Xg
4q
fLSA(E) =
E W+W'
-1,-1
21

---

Published in Transactions on Machine Learning Research (12/2025)
In this case, have
0 (ww x -y (w. x -y Xx
0
VW" M K Q)
and LSA recovers the result in Von Oswald et al. (2023), which performs one-step GD on the update of the linear regression parameter initialized at w. = 0 with Bg 0 w Xyi
1 [w: x -y )xx,
fisa y/ C
<(w, c (Xw. Now.-y) -x)) Xy.
that yields equation 8.
B.2 LSA is a Special Case Linear Transformer Block
In this section, we show that y--LSA defined equation 12 is a special case of linear transformer block (LTB) presented Zhang et al. (2024b), which is mentioned in Section 4.
LTB combines LSA with linear multilayer perceptron (MLP) component. That is,
SiTB ((++))((4)) R
(30)
E
Wº
Wi
E WWWW (WW)) WyE 7
11-1
where W1, W2, W WV W and Wº are trainable parameters for fLTB, and
X y
Xg
E ((++))((++1)
E
for X Rxxd y E RC and Xg Rd. Notice that there is initial guess yg involved in this embedding matrix E.
We denote the hypothesis class formed by LTB models as
FLTB def (LTB:
W WK W WY WP, W.. W.
where fLTB is defined in equation 30. Then we have the following lemma. Lemma 3. Consider fua -SA defined in equation 12. We have
Jy,- E FLTB-
Proof. Let w E RRd. For all X E Rxd y E RC and Xg E Rdd, we have
yxx,ss = fisA(EW) Ew + WWW E W
(WeK)TTORp (ET (W Ew) -1-
with
T X, (R(4+1)((++)) Ew
w Xg
We aim to find (W"), (WP) y,, W1, We for fLTB such that for (X., = ffTe(E) with
X
E (((+1)((++1)
E
y
0
22

---

Published in Transactions on Machine Learning Research (12/2025)
Let choose W2 la+1 and
Ia
w
W1 =
(31)
w T
with + /|wl?, then 22 W W, and W E ((d+1)x(d+)) is invertible.
u
Indeed, let u E I&d and u e I such that W
0. So we have
uw 0, w u +
From + uw =0, 0, we have u uww Plugging it into w TTu cu =0 0, we obtain
(c (o-1/w/>)ee )u=
Since c + /// we obtain a = 0. Thus, u -uW = 0. This implies that W is invertible.
Next, we consider the following matrix
Ia
0 (d++1)
0
W3
w 0
Let
(WPY- = W.'WWP
(W = W W3.
(W9) = W&W, W3.
(WV)= W W3
We show that fv,-LsA(X,y,,) x,y) (LT Indced, by using Xw y, we have
W&E
w
X
XT
W X2
X,
=
-Ew-
So
(LT((E) wi (i((B: E 10+72 Wi wf WoW""" T w9
1
[W1E WW TW En
-w Xx + ((d W WW WM TW En
Thus, we conclude fr, LSA E FLTB.
C
B.3 Proofs of Theorem 4
The risk (loss) function with learnable vector is given by:
R(f,, LSA) (S)) = E ( Au(E)) C+1.H1
AHH(D))HH.HHHt +vTx, wTx,)) x²)
23

---

Published in Transactions on Machine Learning Research (12/2025)
Similar Appendix A, we rewrite the risk:
R(f -LSA) =E ((1 [11+ + b² Ga + + GA. - )))xx)) 2]
= =E [(1 + b" b" Ga + Gb"GA.d - -w xy]
GA.
(G, + (G, ba T |vbj + vis) - -mb)) 2
E
We define, for each j:
t (G, bay) + (G, ba balt)) víil + vijl w[]]
Then
LLSA LSA =
Σ ==1
E B(g)].
Step 1: Gradient for parameters
We list the first-order partial derivatives with respect b, ay, au41, and v(i]- j is from to d
Gradient w.r.t.
0t, 0b
Ga + vijl Ga
0i
(t))
21
2t, (Gay v[j G Ga#t1).
bb
0s
d
Σ
E 21 (Gaj + v(s] Gatti)]
Gradient w.r.t. aj
0tj = G b. da,
0 (1)) 21, (G
da,
Only the j-th term depends on aj, so
0f LSA Ōaj
LSA
2E(1, (GTb)].
Gradient w.r.t. ag=1
Øt, v(5 (G..).
0as+1
(tg) 2t, (vij) G b).
Da.+1
24

---

Published in Transactions on Machine Learning Research (12/2025)
0fy,-LSA 0an+1
>
E L vii (G'b)).
j=1
Gradient w.r.t. v[j]
We have
i, bb + ofjl (b Gast1 +1) - -w[i] + w.bll).
Øt, Ou[j]
(b G Gaz-1+1).
0f ,, LSA Ou[j]
2 (b (bTGa aast1 + +1]
Step 2: Plug in One Step GD
we verify when b W.
ej
a, =
Ad+1 = 0, v= w., the gradients equal to zero
we define w w W., We have the following intermediate formula:
b" Ga = w 1]
x T yx
x y
[e
Ei= [w" xxxx""" w x
[[]]
E
w
x
0
C
obb](bb Ga441) =0
CC t,
L1=1 C
w X,X, [5] w(pl
1
x [j] y/xx []]
Gaj C
Gradient w.r.t. b
0S, LSA
2
>
E
Ei=l
w T
x,,+00] ww/l)
1
x [j] yix; [j]
bb
C
==1
Calculate each part:
wiil 0
Σ
x li] O,
w[D]
1
Σ
C
-
1
Σ
( + w))xx -wbl Σ w xxxxx//]==> = 11
Σ
w T xxxxxxxxxx = O.
25

---

Published in Transactions on Machine Learning Research (12/2025)
E C
w" xx,,111 = E
E
w"
x
Σ
(w:
w"
x;xx
[5]
wxxx|1]
C
C
1 E
(Σ
xt(s xxx
c²
k=1
compute E x1اي]
(Een xx|i| xxx }
when i + k,
[|xx =1 1.
when i= k
xxbxxx,Di(x>>x)> x.) x Σ
= d+2
Because E[x[ji((x*x] x [[x[) Empty [mx])>x(m]] F[x||1 3.
Σ x1xx) (Σ xx|i| xsx = C(C-1)+
if we have very large C, we have:
x 11l _ C
ri
wx
T Σ yxxlii1 1.
C
So that
-1n 0 ab
Gradient w.r.t. aj
-LAA = 2t(t, (G (TT b)] = =B[ E da,
-
Σ C
XxxX w)
x
-w(l)
c
|xi w
C
X w (w" >xxx w/ )= B[-
w//]]]]
- [
E[ [[[]
Li=1
]
C
+
compute En w x [5]
x x
w We aim to compute the expectation:
C
C
C
E
w
Σ
Σ
xxx
w
i=1
k=
26

---

Published in Transactions on Machine Learning Research (12/2025)
First, expand the product inside the expectation:
)
C
x
w T
Σ xxX|م]|
w Σ
Σ
x, x Xt.
w
Taking expectation:
c
C c
w
x x
>Σ
Σ
E
k=1
Case 1: i k
Since X; and xk are independent:
E w x
x
x w xs| = E [wTxx x, x, E Xxf w xb3
Given X; (a)
E x
E [ w xx] W.
Thus, for + l:
E w x, x w- xx| = ww/]..
There are C(C - 1) such terms, contributing:
C(C - 1) C(C 1)e3-
Case 2: i ==/h k
For i i= k:
E [ x, x, x w xx] E [(w"'x xxx [)1x]]]]]]
Using properties of Gaussian vectors:
E xx) = 2w, w + [|
where ey is the j-th standard basis vector. There are C such terms, contributing:
C(2w,w + |[w])2,))
Adding contributions from both cases:
(ixxxx)) w = 1w w(2w, + ///l),,))
w Σ
xxxx
1=1
Simplifying:
= CCC w[
Thus, the expectation is:
C
E T
Exxx(i
Σ
xxxx
km
when C is large E
E En w
Txx,
-ي x w
=
27

---

Published in Transactions on Machine Learning Research (12/2025)
De compute Eg w x;x; [5] 2 mif w
Σ Li=1 w x;
2k=1 k= (w
w xxxXX w
C'
C
From our previous experience, we only need calculate case when k +
Σ w
Ek=1 (w? + +w -w" xxxxx w = [w[j]ww] + +w")ed =w. [7]
E
T x,,
C
- 0
So that we have
daj
Gradient w.r.t. ax-1
pf.0-1Sh LSA
2 2 k [ = 2 > B w.bi (GTb)].
Dad+1
5m1
we already have So that we have
- da,
2E(t, ty (G
85 LSA
0
daal1
Gradient w.r.t. vljl
0f
21 (b G Ga +1]
2E
-i=
c
w
x
w(il)
( wx x, +1]
Oulj]
C
2
i=
w
x
w(il) =0
E nh x x [il
2
اi=l T x|, w(l))
we still only consider the case x k
E
2E
==
T
[
w(1))
w x 1i| = -2E[wb]
w(D) w"]
C
W.
[e3]
verify that b
a
ad+1 0 W., is stationary point for loss f
1
0
B.4 Proof of Lemma 2
Proof. Based on the proof of Lemma 3, consider the following matrix
Ia
E R ((d+1)((++))
w
Now for any fy,-LSA'S inputs (X,y. xx), x.y,))), by using X w y, have
X
Xg
WE
E
X
Xg_
28

---

Published in Transactions on Machine Learning Research (12/2025)
Thus,
(1,- ((sA(Ew)> (
By using Lemma with one-single head, we know that R(f(sA) non-convex. Thus, we conclude that -1S)) is non-couvex, as is a composite function with a non-convex function R(Jisa) and a linear function.
C Details of Experiment
C.1 Implementation Settings.
The experiments use JAX to implement and train the LSA models. We set the learning rate to b 5 x 10 10-4 and batch size of 2,048. single linear attent layer used, without any LayerNorm or softmax operations. We will release our code repository upon publication to facilitate reproducibility.
Table I:
Overview of the experimental setups. Each experiment modifies one factor (number of attention heads, prior mean, or y4) while holding the others fixed.
Experiment
Number of Heads
Prior Mean
3/
Head Section 5.1.1
Varies
[2, 2, .......)
Prior Mean Section 5.1.2
11
Varies
1/ Section 5.1.3
11
[0,0,
Varies
C.2 Detailed Metric Definitions
Prediction Norm Difference The prediction norm difference measures the discrepancy between the outputs of LLSA and one-step GD (faD). Given a test input Xg, we define the difference as:
(lv,, faD((x)))
This metric quantifies how closely -SA approximates the predictions of the explicit one-step GD solution.
Gradient Norm Difference The gradient norm difference assesses the deviation between the sensitivity of the model predictions to the input. Given the gradient of the output with respect to the input Xg, we compute:
[ofo((x) 05y. DSu--LSA(Xg)) Ox,
dx,,
This metric evaluates whether LLA captures the same local sensitivity as onc-step GD.
Cosine Similarity The cosine similarity measures the angular aligument between the gradients of the two models. It defined as:
8top(xy) 8x4
0fug (iag-1sa(x.))) 9x,
8fan(xa) Ox,
05 (x,) Ox,
A cosine similarity of 1 indicates perfect alignment between the two models, while lower values suggest deviations in the learned representations.
C.3 LLM Experimental Settings
We conducted our experiments using the STS-Benchmart dataset (English subset) 2021), which consists of sentence pairs labelled with semantic similarity scores ranging from to 5. The LLM used in our study Meta-I.LaMA-3.II 1--.33.1------sssrnrrrraaataa et al., 2024) and QQvan/QQeen255-7-IInsstte (Yang al., 2024;
29

---

Published in Transactions on Machine Learning Research (12/2025)
Team, 2024). The model's generation parameters included maximum of 150 new tokens and deterministic decoding.
The guess model was trained to generate imitial similarity score guesses. It consisted of a two-layer feedforward architecture, taking as input the concatenated embeddings of two sentences computed by the Sentence- Transformer model all-MiniAM-T.&-22(Reieeee embeddings 16-dimensional space with ReLU activation, followed by second layer that outputs single scalar value the predicted similarity score. The model was trained using Adam pptimizer(Kingnan 2014) with learning rate of 1e-3 and mean squared error loss function. Training performed over 10 epochs, with batch size of 8. Sentence embeddings were dynamically computed during training. The loss for training the guess model was computed as the MSE between the predicted and ground truth scores.
& Gurevych, 2020). The first layer mapped the concatenated
For each prompt, a context was constructed by randomly sampling 10 labelled examples from the dataset. Each labelled example included two sentences, ground truth similarity score, and an initial guess for the similarity score generated by lightweight guess model. The query example included two sentences and its guessed similarity score and an explicit instruction for the LLM to refine the guess and provide a similarity score between 0 and 5.
To evaluate the effectiveness of the initial guess, we calculated the MSE between the LLM's predicted similarity scores and the ground truth scores across 100 experimental runs. The baseline performance, derived from the initial guesses provided was compared to the refined predictions generated by the LLM.
30


Reviews:
[{"review_id": "Voxje81afi", "reviewer": "Reviewer_3NnV", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper builds on the results of previous work that examined In-Context Learning (ICL) in simplified linear regression as a gradient descent step within the forward pass of a linear self-attention (LSA) mechanism. Earlier studies considered linear regression with a zero-mean Gaussian prior and constructed LSA weights such that the equivalence between ICL and gradient descent on linear regression would hold. Later, in Theorem 4, they showed that LSA behaves like gradient descent only if the query’s initial guess is a learnable parameter that matches the data’s prior mean.\nEmpirically, they first confirmed the effect of different initial guesses of the model in a toy example. In more realistic large language model (LLM) experiments, they used another model to provide a good initial guess in the initial prompt. Their results showed that a better initial guess leads to better in-context learning performance overall.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Yes. The claims are supported by clear theoretical arguments and empirical validation. The paper’s core proof demonstrates that even infinitely expressive LSA cannot emulate one-step GD under non-zero prior means, revealing a concrete limitation in prior equivalence claims. The accompanying experiments—both toy linear regression and LLM prompting—test and confirm this dependence on initialization, making the evidence both accurate and convincing.\nMy only concern is with the novelty of the theoretical results presented in this paper. I have asked for more clarification later.", "audience_yn": "Yes", "audience_explanation": "The results convincingly demonstrate that the deviation between ICL and GD stems not from limited expressivity but from a structural dependence on initialization. Theorem 1 shows that once the representational capacity of linear self-attention saturates at\nH=d+1, adding more attention heads no longer brings it closer to gradient descent. This result matters because it establishes that the ICL–GD gap in this simplified setting is not an expressivity issue, but an inherent limitation of the mechanism itself.\n\nAlso they show that a zero-mean linear regression cannot be equivalent to the existing ICL of linear self-attention. However, The “no stationary point” impossibility of theorem 2 in the zero mean case is only for when the context goes to infinity C→∞. \n Then they show that an adaptive initialization can theoretically help with this mismatch. However, there are practical implications to a good initial guess:\n\n- It implies the model can learn how to start thinking about each query, a kind of meta-initialization.\n\n- It also changes the training problem: the model now must optimize not just transformations but also the inputs to its own inference.\n\nThis limits the applicability of the proposed framework empirically. However, the theoretical insights backed by valid experiments, gives a better understanding of ICL and is useful for the community.", "requested_changes": "My main concern is the novelty of the work given the current literature. I have take a look at: https://arxiv.org/pdf/2402.14951 after writing this review. The theoretical results sounds to be very similar. Can you state exactly what is the distinct theoretical contribution of yours? Also if a similar theorems exist in the literature, you must elaborate the exact distinction of your theorem with them in the main text. I don’t see this in your paper.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "75cyJQurgB", "reviewer": "Reviewer_tukB", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper investigates the connection between in-context learning (ICL) and gradient descent (GD) in a more realistic setting than prior work, specifically by relaxing the zero-initialization assumption commonly made in previous studies. The work primarily extends von Oswald et al.'s framework to non-zero initialization scenarios using linearized self-attention (LSA).\n# Contributions:\n\n- A theoretical proof demonstrating that when regression weights have a non-zero mean, multi-head LSA cannot generally replicate one-step GD, even with many heads. This is a significant theoretical result for understanding the fundamental limitations of LSA-based ICL.\n- Theoretical upper bounds on the performance of LSA-based models based on number of heads.\n- Empirical validation of the theoretical findings, including validation of prior mean w* impact.\n\n# Strengths:\n\n- Important theoretical contribution that advances understanding of when ICL can and cannot faithfully replicate gradient descent\n- Generalizes existing work to more realistic settings (non-zero initialization)\n- Empirical validation supports theoretical claims\n- The mathematical proofs appear to be sound\n\n# Weaknesses:\n\n- Significant clarity and presentation issues that make the paper difficult to follow, especially for readers new to ICL research\n- Critical missing information: dimensions of weight matrices (Wq, Wk, Wv, Wp) are not specified, making mathematical verification difficult.\n- Insufficient explanation of linearized self-attention and its relationship to standard self-attention. \n- Experimental setup is poorly explained (e.g., confusion about \"training steps\" in Figures 2 and 3 when ICL typically implies no training)\n- Lack of intuition for key design choices (e.g., why multi-head outputs are summed rather than concatenated)\n- Heavy reliance on external references to understand core concepts.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The mathematical proofs appear to be sound, and the empirical evidence supports the theoretical claims, including validation of the prior mean w* impact. The core theoretical contribution regarding the impossibility of multi-head LSA replicating one-step GD with non-zero mean regression weights is convincing. However, the evidence is not presented clearly. The lack of dimensional specifications for key matrices makes independent verification challenging. The experimental setup lacks clarity, particularly regarding what “training” means in the context of ICL (Figures 2 and 3). While the claims appear to be supported, the presentation issues significantly hamper the ability to fully assess the evidence.", "audience_yn": "Yes", "audience_explanation": "The paper addresses a fundamental question in the machine learning community: understanding when and how in-context learning relates to gradient descent. This is highly relevant given the widespread use of transformer-based models and growing interest in understanding their mechanisms.", "requested_changes": "1.  Specify matrix dimensions explicitly: Provide clear dimensional specifications for all weight matrices (Wq, Wk, Wv, Wp) in the linear self-attention formulation. This is essential for readers to understand and verify the mathematical claims, particularly why updates are rank-one.\n2. Clarify experimental setup: \nProvide a clear, detailed explanation of the experimental methodology, particularly:\n- What \"training steps\" mean in Figures 2 and 3 in the context of ICL (which typically involves no training)\n- What exactly is being trained and how it relates to the ICL framework\n- Make the experimental protocol self-contained and understandable for readers not deeply familiar with this specific line of work\n3. Explain linearized self-attention thoroughly: Add a comprehensive explanation of linearized self-attention in the introduction or early sections:\nIts relationship to standard self-attention ,Why and how it differs from traditional attention mechanisms ,Intuition for why it's an appropriate model for studying ICL-GD connections\n4. Clarify multi-head aggregation design choice: Explain why multi-head attention outputs are summed (with each head having its own projection matrix) rather than using the standard concatenation approach in transformers. Provide motivation for this architectural choice.\n5. Improve accessibility for ICL newcomers: The paper currently requires extensive familiarity with prior work. Add more intuitive explanations and self-contained descriptions of key concepts to make the paper accessible to a broader audience\n6. Add intuition throughout: The paper would benefit from more intuitive explanations accompanying the mathematical formalism. For example, provide intuition for why the rank-one update structure emerges and why non-convexity of R(fH−LSA) is expected. \n7. Improve figure captions: Make captions more informative and self-explanatory. For instance, \"loss curves over the course of training\" should specify what type of training and what is being optimized.", "broader_impact_concerns": "This is primarily a theoretical paper investigating the mathematical relationship between in-context learning and gradient descent. The work does not raise obvious ethical concerns or require a Broader Impact Statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "nyi4wmmVXj", "reviewer": "Reviewer_2Dxp", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the setting of in-context  linear regression with linear self-attention (LSA) where the underlying parameter $\\hat{w} \\sim N(w^* , I_d)$ with $w^* \\neq 0$. Previous work [1] has shown that when $w^* = 0$ the global minimizer of one layer LSA is one step of pre-conditioned gradient descent. In this paper, the authors focus on the case of non-zero $w^* $ and show how in this case the loss is non-convex and thus does not have a global minimizer. They then introduce a learnable parameter in the input $w$, and when $w = w^*$ the LSA is in fact one step of gradient descent for specific choice. They then perform experiments to show the gap of one step of gradient descent and the trained model, while also exploring this as a function of the number of heads. They also perform some experiments in language models where they add some extra initial guesses in the prompt (these values are attained from another smaller model).\n\n[1]: Ahn, Kwangjun, et al. \"Transformers learn to implement preconditioned gradient descent for in-context learning.\" Advances in Neural Information Processing Systems 36 (2023): 45614-45650.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Most of the claims are clearly stated.", "audience_yn": "Yes", "audience_explanation": "I think this work provides some understanding on how even changing the initialization affects what these models are learning.", "requested_changes": "Regarding the experimental section in figure 5 what is the GD loss correspond to? Since we have different initializations for w, the loss of one GD step should be different. I would like the authors to clarify what exactly this GD loss is. I would also like a clarification paragraph about the differences between $C \\to \\infty$ and $C$ being a constant. How the results would be affected?", "broader_impact_concerns": "No ethical concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)