Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Authors: Dake Bu
Wei Huang
Andi Han
Atsushi Nitanda
Qingfu Zhang
Hau-San Wong
Taiji Suzuki
Venue: ICML 2025 poster (2025)
Keywords: Task Vector; Residual Transformer; Factual-recall In-Context Learning;
Primary area: theory->deep_learning

Abstract: In-context learning (ICL) has garnered significant attention for its ability to grasp functions/tasks from demonstrations. Recent studies suggest the presence of a latent **task/function vector** in LLMs during ICL. Merullo et al. (2024) showed that LLMs leverage this vector alongside the residual stream for Word2Vec-like vector arithmetic, solving factual-recall ICL tasks. Additionally, recent work empirically highlighted the key role of Question-Answer data in enhancing factual-recall capabilities. Despite these insights, a theoretical explanation remains elusive. To move one step forward, we propose a theoretical framework building on empirically grounded *hierarchical* concept modeling. We develop an optimization theory, showing how nonlinear residual transformers trained via gradient descent on cross-entropy loss perform factual-recall ICL tasks via vector arithmetic. We prove 0-1 loss convergence and show the strong generalization, including robustness to concept recombination and distribution shifts. These results elucidate the advantages of transformers over static embedding predecessors. Empirical simulations corroborate our theoretical insights.

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Dake Bu"
Wei Huang? Andi Han 23
Atsushi Nitanda 456 Qingfu Zhang'
Hau-San Wong' Taiji Suzuki
Abstract
dation models. A growing academic perspective high- lights that the core strength of LLMs lies in their markable In-Context Learning (ICL) capability (Lu et al., 2024), enabling them to infer underlying tasks or func- tions from input demonstration pairs (Minegishi et al., 2025b)-akin to the concept of function references in pro- gramming (Bernays, 1936). For example, given the prompt Japan Tokyo France Paris China, the expected output "Beijing" reflects the underlying function "Coun- try's ' Capital" applied to the final query.
In-context learning (ICL) has garnered signif- icant attention for its ability to grasp func- tions/tasks from demonstrations. Recent studies suggest the presence of a latent task/function vector in LLMs during ICL. Merullo et al. (2024) showed that LLMs leverage this vec- tor alongside the residual stream for Verd2 like vector arithmetic, solving factual-recall ICL tasks. Additionally, recent work empirically highlighted the key role of Question-Answer data in enhancing factual-recall capabilities. Despite these insights, a theoretical explanation remains elusive. To move one step forward, we propose theoretical framework building on empirically grounded hierarchical concept modeling. We de- velop an optimization theory, showing how non- linear residual transformers trained via gradient descent on cross-entropy loss perform factual- recall ICL tasks via vector arithmetic. We prove 0-1 loss convergence and show the strong gen- eralization, including robustness to concept re- combination and distribution shifts. These re- sults elucidate the advantages of transformers over static embedding predecessors. Empirical simulations corroborate our theoretical insights.
Role of Task Vector. A growing body of research seeks to uncover the internal mechanisms underpinning ICL, shedding light on how LLMs retrieve function/task from demonstration. Recent studies employing advanced prob- ing techniques have identified the emergence of task vector (Hendel et al., 2023; Merullo et al., 2024; Yang et al., 2025) (or function vector (Todd et al., 2024; Ka- hardipraja et al., 2025)) within the latent representations of LLMs during ICI. deduction. This vector appears around the 15-19th layers in models like GPT-J-6B (Wang & Komatsuzaki, 2021). Formally, Hendel et al. (2023); Merullo et al. (2024) propose that, given the prompt T [x,, (x1). X2, fxx2). ... Xyueryl, the LLM O constructs the task vector in the earlier layers, denoted as (T), based on its understanding of the task/function message in T. Furthermore, Merullo et al. (2024) revealed that, certain factual-recall ICL tasks indeed correspond some latent representation (xx where bgueryy (Xquery) is the query-encoded residual stream in deeper layers. That is, the model O can execute the ICL task through simple vector addition:
1. Introduction
(T) querry (Xquery):
Transformer- baased Large Language Models (LLMs) (Vaswani et al., 2017) have ushered in new era of foun-
"Work completed during internship at RIKEN. 'Department of Computer Science, City University of Hong Kong. Hong Kong SAR *Center for Advanced Intelligence Project, RIKEN, Japan SSchool of Mathematics and Statistics, The University of Sydney, Australia "Institute of High Performance Comput- ing (IHPC), Agency for Science, Technology and Research (A*STAR), Singapore 'Centre for Frontier AI Research (CFAR), Agency for Science, Technology and Research (A*STAR), Sin- gapore 'College of Computing and Data Science, Nanyang Technological University, Singapore 'Department of Mathemat- ical Informatics. The University of Tokyo, Japan. Correspon- dence to: Wei Huang <wei. v vr@rikenjp>>, Hau-San Wong <eshswong @ittt..du.hk...
pe (f (XQqury) -. I
pe pe(ag
b query ag. T)
(1)
pa T)d(ao),
where po (ag T) denotes the model's confidence in rec- ognizing ag as the task vector a/ (T). The formulation in Eq. (1) highlights the pivotal role of residual streams. How- ever, existing theories on ICL either overlook the residual term entirely (Tian et al., 2023; Kim & Suzuki, 2024; Bu et al., 2024a) or handle it in an unnatural manner (Nichani et al., 2024). Additionally, while some studies suggest that Question-Anan Answer (QA) training data are crucial for en- abling L.MSS to retrieve factual knowledge (Allen-Zhu &
Proceedings of the 1nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Provable In- Mn-Context Vector Arithmetic via Retrieving Task Concepts
Prpectional Coultry Cappltsane
Sopeeteee Syene
Projectin Hedden States for Chine Butirr
Progection of sidden States for takuen and Jouns
2012 to son
0,003
city dircttinn [e-aeNs)
criggn direction teaenk
(a) Country-Capital (2D).
(b) Country-Capital (3D).
(c) Symbol-Origin (2D).
(d) Symbol-Origin (3D).
Figure I. Visualization of task vector and word embedding. The 22nd layer's vector embeddings of GPT2-medium the 2-D and 3-D projection spaces are shown, where the task vector ('o (_task") is extracted from some internal layer, following Merullo et al. (2024). For the task of retrieving a country's 's capital (a-b) and the task of retrieving a symbolic creature's origin (c-d), the latent embeddings of the LLM demonstrate the approximate relationship "x + o_task = y", with 'x' containing littlc components aligned with "o_task".
Li, 2024; Zhang et al., 2025a), there currently no theo- retical framework to substantiate this claim, nor expla- nation of how it facilitates factual-recall ICL, These gaps naturally raise the following research question.
erties, as illustrated in Figure 1. In factual-recall tasks, answer/label vectors y tend to align more strongly with the task vector ag than query word vectors æ, satisfying x + ag Ri y. Here, ag can be interpreted as the represen- tation of high-level task concepts (aligned with the x-axis in Figure 1), while the components of æ orthogonal to a.- represent low-level concepts (aligned with the y/z-axes in Figure 1). This interpretation is consistent with Park et al. (2025), which demonstrates that LLMs exhibit hierarchical linear concept geometry after generative pretraining, with word representations residing within polytopes. These in- sights inspire the modeling approaches in Section 2.1.
Research Questions 1
How does a non-linear residual transformer, trained via gradient methods with realistic cross-entropy loss on QA data, naturally perform factual-recall ICL in the vector arithmetic style described by Eq. (1)?
Analogy to Word2Vec. As noted by (Merullo et al., 2024), the vector arithmetic in Eq. (1) mirrors that of Word2Vec, a static word embedding model (Mikolov et al., 2013). For instance, the representation arith- metic "France Paris Poland Warsaw" can be interpreted "France Paris" represent- ing the task vector ag. which performs the function "get capital (+)," while representations of countries like "Poland" or "Japan" act as bguer.
Our contributions are summarized as follows:
1. We develop an optimization theory demonstrating that transformers with nonlinear softmax attention, MLP, layer normalization, and residual connec- troined via Gradient Descent (GD) with cross- entropy loss-can effectively perform factual-recall ICL in a vector arithmetic manner, grounded in empiri- cally motivated data modeling. Our analysis shows that the transformer retrieves the high-level task/funccion concept through attention-MLP, which, when combined with any embedded query vector within the same high- level task concept, yields the correct corresponding swer vector. Despite the inherent non-convexity of the learning problem, we establish the asymptotic behaviors of the optimization dynamics and prove the convergence of ICL test losses.
Recent work by Wibisono & Wang (2023) examined the relationship between one-layer bidirectional attention opti- mized via Masked Language Model loss through reparam- eterization and the Word2Vec method. However, they show that such BERT-like models function as an inexact approx- imation of the Word2 Vec method, exhibiting non-trivial er- rors. Their analysis failed to delve into the the optimization process, recognize the task vector mechanisms, or clarify the comparative advantages of transformer over Word2Vec. This raises the following question for further investigation.
2. We demonstrate a clear learning scenario separation be- tween training QA data and on Word-Label demo- pair ICL data, both framed within hierarchical con- cept modeling. While prior theoretical works assume both training and testing on Word-Label pair ICL data (Zhang et al., 2024; Kim & Suzuki, 2024; Chen et al., 2024a; Bu et al., 2024a), we show that this approach
Research Question 2
In the context of (1), what are the primary strengths of transformers over their Word2Vec predecessors?
To address these questions theoretically, we first observe that the latent geometry of L.MS exhibits intriguing prop-
2

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
fails retrieve the high-level task vector and instead leads to harmful memorization of low-level features. In contrast, we prove that training on QA data enables the model to effectively learn the task vector, achieving ar- bitrarily small error when tested on either Word-Label pair ICL or QA-ICL distributions.
et al. (2024b); Marks & Tegmark (2024) provided both the- oretical and empirical evidence that L.MSs tend to represent independent concepts and facts in a linear manner.
In- Mn-Context Learning Theory. Recent work has studied how transformers perform in in-context learning across vari- ous settings (Zhang et al., 2024; Chen et al., 2024b; Kim & Suzuki, 2024; Nichani et al., 2024; Chen et al., 2024a). However, these analyses often overlook the crucial role of the residual stream and the phenomenon of in-context vec- tor arithmetic. Moreover, they typically rely unrealis- tic assumptions, such as linearized or combined QK atten- tion (Kim & Suzuki, 2024; Nichani et al., 2024), or on im- practical loss functions like squared or hinge loss (Chen et al., 2024a).
3. We provide theoretical guarantees for compositional generalization in in-context learning. focusing on task vector emergence and OOD robustness. Specifically: (i) we show that transformers can regress task vec- tors directly from demonstration pairs, without requir- ing explicit query, and can apply these vectors compo- sitionally at test time via arithmetic manipulation; (ii) we prove that models trained on structured QA data can generalize to unseen task prompts or dictionaries by leveraging conic combinations of learned high-level task vectors and novel orthogonal low-level and task- irreevant features; and (iii) we establish the model's adaptability to distribution shifts in prompt content and length. These bonuses shows transformer's superiority over traditional methods.
Additional Related Work could be found in Appendix A.
2. Problem Design and Intuition
Notation. We denote the Bernoulli distribution with pa- rameter p. which represents a discrete distribution over {o, 1), as Ber(p). We use R2, Rddd to denote vector and matrix space, and utilize R to denote identity ma- trix. For two sequences an and bns On O(b,)) indicates that there exist constants C' > 0 and N > 0 such that |a,l S Clb,al for all n 2 N. Similarly, an =l,,)) means b, = (an). and an = =(b,,) signifies both On = (b₂)) and an = 12(b,).. We define span(v1, .2...., vK) as the linear space spanned by vectors U1 02 ................................ Uk. and conic(v), 12,,...,
Remark 1.1. Humble Remark. We humbly acknowledge that the scope of this study is confined to illustrating the merits of task vector arithmetic mechanisms observed in some single-token factual-recall ICI. tasks, modeled based on prior empirical and theoretical observations. We do not claim this as a universal explanation for how LLMs han- dle complex multi-token or factual tasks. Nevertheless, to the best of our knowledge, this is the first theoretical work to explicitly consider the pivotal role of the resid- ual stream and layer-wise normalization in ICL, contribut- ing to the growing body of research on transformers and ICL. These nonlinearities, alongside softmax attention and cross-entropy loss, introduce significant challenges in ana- lyzing optimization dynamics.
vx) the conic hull, which includes all non-negative linear combinations of these vectors. We use for the l2 norm and || || for the Frobenius norm.
2.1. Hierarchical Data Modeling
In this section, we present our data modeling based on the observations of task vector arithmetic in rrcall ICL illustrated in Figure 1. We found the near-orthogonal prop- erties in Figure 1 coincide with Park et al. (2025), which suggests that LLMs encode high- and low-level concepts in an approximately orthogonal manner. Specifically, we treat the task vector as a high-level concept representation, while orthogonal components represent task-specific low- level concepts. Details are delayed to Appendix C.
1.1. Related Work
Task Vector Mechanism. This line of research empiri- cally validated the emergence of task vectors (Hendel et al., 2023: Todd et al., 2024; Liu et al., 2024; Yang et al., 2025). Especially, Merullo et al. (2024) revealed that large language models (LLMs) implement certain single-token factual-recall tasks through vector arithmetic. However, no existing study explains why and how gradient-update trans- former models naturally implement this mechanism.
Hive-Level Task Concept Vector. There are K high-level binary concepts, each denoted by 2 E {0, 1}. We define the steering vectors a Rd represent 2 = 1 over 2 = 0, where all K vectors are mutually orthogonal. This orthogonality ensures that for any distinct kn h2 ka, ka E [K]. ax, - ai, 1 au. - -kk, signifying the independence between concepts, consistent with findings in recent work (Park et al., 2024; 2025; Jiang et al., 2024b; Marconato et al., 2025; Liu et al., 2025). For example, the concept of "capital"" can be considered independent of "gender""
Storage and Retrieval of Factual Knowledge. Recent re- search has explored how LLMs perform factual recall or associative memory tasks (Nichani et al., 2025; Cabannes et al., 2024; Allen-Zhu & Li, 2025). Allen-Zhu & Li (2024); Zhang et al. (2025a) empirically show the impor- tance of QA data to improve LLM's capability to retrieve factual knowledge. Additionally, Park al. (2024); Jiang
Low-Level Task-Specific Concept Vector, For each high-
3

---

Provable In-Contes, Vector Arithmetic via Retrieving Task Concepts
level concept žky k E [K] there is associated low-level binary concept We E {0, 1}. represented by two semanti- cally opposite vectors bb and b& E Rd, where We = 0 corresponds to -b and wa = 1 to bg. As noted in (Park et al., 2025) in the latent representation of I.LM, the high- level task concept vectors are orthogonal to the low-level ones, i.e. aA₁ I DE_ for all ki, l22 E [K], and by are mutual-orthogonal..
ICL vs. QA Training. Indeed, training and testing solely the ICL data is unrealistie, despite its popularity in the- oretical studies (Zhang et al., 2024; Kim & Suzuki, 2024; Chen et al., 2024a; Bu et al., 2024a). A more practi- cal approach involves generative pretraining or Question- Answer (QA) pretraining or fine-tuning (Allen-Zhu & Li, 2024; Zhang et al., 2025a). Notably, our concept data mod- eling naturally arises from generative pretraining (Park et al., 2024; Jiang et al., 2024b; Park et al., 2025). Fur- thermore, as highlighted by Allen-Zhu & Li (2024); Zhang et al. (2025a), QA data plays a crucial role in enhancing a transformer's 's ability to retrieve relevant factual knowl- edge. Additionally, Merullo et al. (2024) leverage QA data to extract task vectors. Motivated by these insights, we in- corporate QA data into our training framework, which we formalize as follows.
Word-Label Pair ICL Prompt Distribution (Pr). The prompt T (x,,,,,"'. XJ,YJ, X E Rdx(22+1) and label y++ are generated as follows. Each prompt's co- task concept, which is shared within the demo-pairs within the prompt, is sampled as kT ~ Unif[K], with label in- dicator ykr.d ~ } and noise terms Et.x: fiy ~ N(0, I). Word-label pairs in demonstrations and queries follow:
Σ (Ia - ak + yku - bb) + E1,:
QA Sentence Distribution' (PoA). We model our QA- type sentence S := [x yy as follows. Each sentence is associated with task concept lis ~ Un] The word x and label vector y are constructed similarly to the word- label pair distribution Pr. The QA prefix XOA is created by combining common tokens Vn,1:M (M task-irrelevant common tokens) and the task vector akg which can appear anywhere before x as the last column of x To capture semantic variability,, noise vectors fI:M: &x NN( of.1) are added to the common tokens. Formal details are pro- vided in Appendix C. An illustrative example with M = 5 to show our modeling intuition on concept k is:
X: :=
RENT.
Y1
Σ
(ax + yk.s bu) + Eny.
REYT.
for 1 E [J. To model word polysemy, we define Dr. C [K] as the sets of latent concepts associated with Xi and yı, respectively, each set includes the shared co-task concept kr along with possibly task-specific but contextually irrelevant concepts. The noise term & cap- tures semantic variation, and the task anchor T, ag facil- itates concept retrieval, as illustrated in Figure 1. with va set to 0.1 (see Appendix C). The query word is defined X ExENT [1 (Ia- ag + 34, bx) + E.+1,x. fol- lowing Eq. (2). The expected label, by contrast, is given by yJ+1= = akr + U+r.J+1 bkr, since only semantics tied to the shared co-task concept kr contribute to prediction under task KT. An illustrative example with = to show our modeling intuition is
What Pa.1 +Ei
is the
capital Big
y ze tebig ang tobes
Pa.2 +t2
+ha
+Ei
+
Here, we assume that a specific position ms 4 E [M] encodes task vector akg encoding the high-level task mes- sage. This formulism is inspired by the empirical findings in Allen-Zhe & Li (2024) and serves role similar to the relation token described in Nichani et al. (2025), but ours differs in its ampirically-support. concept modeling.
Japan Sakura y1
France
Rooster Y2
China X3
X1
X2
QA-ICL Prompt Distribution. (PT)) A natural ex- tension of the above two distributions is that we can re- place the word-based demonstration by QA-based demon- stration in the task specific ICL prompt, namely Toi- = .A [x] .y.... ,x y,,, J+1 E (Xx((++)(M++)) ~ PTo..
where the shared co-task concept ky in this context is "National Symbol", and the expected label vector is y3 = P Panda-the symbol associated with China under the task kr. While the words "China" and "Panda" may carry semantics relevant to other tasks (e.g., country capitals or animal categories), only the "National Symbol' meaning is expected to contribute to ys, illustrating how the model se- lectively attends to tesk-relevant semantics in the presence of polysemy. If the prompt is modified to
2.2. Residual- Transformer Model
The structured "up-word-down-llbbe embedding E(T)
I1, FF Equery proposed in prior theoretical y!. 0
y1.
work (Bai et al., 2023; Zhang et al., 2024: Huang et al.,
Japan Sakura France y1
Iris Y2
China Xa
X2
2024a; Kim & Suzuki, 2024; Chen et al., 2024a; Bu et al.,
2024a) is designed for simplified, structured, and residual- free transformers, where words are only available at-
then the expected label vector becomes y = Peony, as the shared co-task concept Hy now corresponds to "Na- tional Flower". This illustrates how the shared co-concept steers correctness during inference.
Indeed, the formula the sentence can be factual statement (c.g. S = Ox y] = [The au, of, x, y)) other than a QA. in 24
4

---

Provable In-Contes, Vector Arithmetic via Retrieving Task Concepts
tention matrices and labels are restricted to value/combine matrices. This residual-free setup deviates from real-world scenarios. In contrast, we consider the case where the [T1.-. Tz] E
we assume ||ag| |a|| ||bاl| ||b|| || ll 1,
common setup for theoretical studies (Tian et al., 2023; 2024). Formal details of U are provided in Appendix C.
prompt T = [21: V1. ,.., Equery RdxL (I. = 2J + 1) is directly processed by a non-linear transformer with a residual stream as follows.
The training data p is sampled from training distribution P, with a sample size of N. To address the scale differ- ence between the gradients of attention and MI.,, we adopt a smaller learning step for MLP characterized by a scale factor ęv- The initial weight matrices w and K are sampled independently from Gaussian distribution (O,, oj II), a more realistic choice than scaled identity or overly constrained initializations used in prior work (Li et al., 2023b; Bu et al., 2024a; Chen et al., 2024a). Sim- ilarly, ( is initialized as NN of I), consistent with standard practice in recent theoretical analyses of Trans- formers (Tian et al., 2023; Jiang et al., 2024a; Li et al., 2025a; Yang et al., 2024b).
L-1
(T) = Σ Wr (WwTT)) (Wo)T)) E Rd.
1=1
he- = WolN((a,,o(T) + TT,
where (s(+)) denotes the column-wise softmax operation, Wo == Laxa. LN(z) z/|zz|22 is the l2 layer-wise nor- malization, and he is the vector output of transformer.
Connection to Word2Vec Arithmetic. When || = ||bæ||| for all k,k' E [K], and both the cardinalities of &r.1, and the noise magnitude are sufficiently bounded, we have an approximate identity:
Test Setup. During testing, we examine the probability that the label vector is the most-likely token to be selected from the disrupted token dictionary tion P where the noises are sampled from N/ o²² II) (either word-based prompt or QA sentence- based prompt)
Ys+1 R OKT XJ+1;
(3)
the test prompt distribu-
where components unrelated the current co-concept kt are negligible compared to the dominant term 1.la,, = byy, which governs the logits when the model attends to the vocabulary dictionary. Therefore, if the transformer 0 can extract the high-level task vector ho (T), then adding it to any word vector within the same task concept should yield the task-specific label vector in the context of argmax sampling. This aligns with empirical findings by Merullo et al. (2024), which show that adding various embedded query words, such as "Poland" or "China", the vector Oeity which captures the function get.capital (+) in the latent space- ppaducproduces the correct capital city.
exp (ue hha
Lp- = -Bp. (ky F argmax(
1)]
ExcTr ky exp (u he)
-Ep. [1k of argmax u ho)].
where by is the index of y in the total dictionary U. The whole procedure is in Algorithm I.
Algorithm 1 Training algorithm
Input: Training distribution P. Test distribution Pt, Training size N, step size nyv, scaled parameter qv, stopping criterion and total epochs T.
Training Setups. Define 0- == [WK, Wo, Wv). We minimize the Legreuulari cross-entropy loss by gradient descent (GD) at each time step
Initialize the model ((0) = (w
w(0) K w(
Sample training data pir ~ P. for = 0. 1, T do If Lg. (0(t)) S E stop else continue. Update model parameters:
exp (u he)
Lpx (0) = Epps log( ELELTK+ exp (u ho)
(u.ho))) }..
(4)
where K denotes the number of tasks, K' the number of task-irrelevant tokens, and us, refers to the vector corre- sponding to y in the token embedding matrix U, which contains K both high- and bi-label low-level concepts, noise-free and single-task-speeific word and label tokens, as well as K' irrelevant tokens-the number "7K + K"" is thus by K sets of bb, 0. =b,, } as well as K' irrelevant tokens. Notably, our analysis focuses exclu- sively on the semanties associated with the prompt's shared co-task encoded in given word or label token.² 2 Normal- ization ensures that each token u; has the same length, en- abling fair comparisons during sampling. For simplicity,
Wy Wy nuv V with LB, (0(1),
WR (++1) Wx (1) V W L1, (0(1))
We (++1)= = Wo (1) 1/ with LB, (0())
end for
3. Theoretical Results
In this section, present our main theoretical results. Condition 3.1. Suppose the following holds for some suf- ficiently large constant C >> 0:
1. Dimension d satisfies d 2 CM2 log( K" M²
While it is tractable to assume that all dictionary entries fol- low the structure in Eq. (2), doing would require additional as- sumptions and care regarding how fixed polysemous words relate to their corresponding labels.
Training sample size N sufficiently large N 2
C max(KK (K log(t), M
KK' log(H)
5

---

Provable In- In-Context Vector Arithmetic via Retrieving Task Concepts
Tain Loss
Test Losses
(W [W_KS.0
W.TWWV
3.25
Ran Losa
1.3
Not Loss
100
3.00
3.
MUJURUUA
[ 2.75
8
BUJUKO_UU
0.4
8
2.25
02
2.00
1000
2000
3000
4000
5000
2003
2000
2000
4000
5003
1000
2000
3000
4000
5000
1002
2000
2000
4000
5000
Epoch
(a) Training loss
(h) Test loss
(c) u W WoU
(d) Z W.
Figure 2. Training dynamics over the ICL prompt training distribution. (a-b) Training and test losses; (c-d) projection values of key-
query and value matrices. In (d), Wv overfits to low-level features b2, resulting in persistently constant (0.2) test error shown in (b).
3. Dictionary size K, K" satisfies K 2 C'log( }), K' > C' max K}.
Training on Pes: When testing on a sample with task concept k* E [K], the model approximately retrieves the appropriate task vector before adding the residual stream:
4. The standard deviations of Gaussian initializations satisfies 01 S min + C, (lo 102
}. and do $
cos o axx) = (1), cos ,o u) = d(1), (7)
min (log( (K 1)= (log( c.
for all u E {(a)aafk-(K3) U (b. U [ux WE(KK))
d=
5. The noise level op satisfies o, I
c
The first result shows that under our setting, when train- ing and testing on ICL-type data, the transformer fails to reliably retrieve high-level co-task information from ICL or QA-ICL prompts. Consequently, the resulting hoa.) be- comes hybrid vector, contaminated by low-level features that disrupt task vector addition with residuals, validated in Figure 2(d). Unlike harmful overfitting in vision models (Frei et al., 2022; Cao et al., 2022), which stems from noise memorization, the harmful overfitting here arises from low- level feature memorization, due to their unintended co- occurrence in ICL-style data. In contrast, the second re- sult shows that training on QA data enables more accurate retrieval of co-task vectors, which is inherently due to the absence of b in XQA. as illustrated in Figure 3.
6. Scaled parameter qv and learning rate 1) satisfy Cafd
o7d oi 2ida C(M-1)=
Clog(d 2 Voll K3 S qu 1
Log(o
'd log(
n S min{ of Ky log
The conditions on d, N, K ensure that certain concentra- tion inequalities hold and that the learning problem is ad- equately overparameterized (Chatterji & Long, 2021; Frei et al., 2022; Cao et al., 2022; Kou et al., 2023). The con- dition on K' controls the impact of contributions from answer-irrelevan dictionary tokens, though it can be re- laxed at the cost of more intricate analysis. The con- ditions on dg and 01 regulate the model's 's initial bias and ensure that gradient descent updates the model effectively. The condition on Op guarantees that the gradient flow is only mildly affected by noise, which is reasonable given the typically high ägnal-to-noise ratio in language data. Fi- nally, the conditions on 1) and qv are technical assumptions necessary for the optimization analysis.
The following proposition further validates the conse- quences of harmful vs. benign task vector retrieval. Proposition 3.3 (Test Losses Disparity). For Ve 0,
under Condition 3.1, for any test distribution p* E (Pi, PTA") ", with the prompt length? satisfying J
(eg0//e))) D (lg((K) and the noise level o = O((K log(())
)
The followings present our primary results regarding the retrieval of task vectors and the convergence of test loss.
with probability at least 1 - 5, there exists t
O "v Kf log( it holds that
Theorem 3.2 (Task Vector Retrieval). Under Condition 3.1, let test ICL prompt distributions include Word-Label ICL Prompt P and QA-ICL Prompt PT, * Then, for the gradient descent iterates in Algorithm I, with probability at least 1 8, there exists t = oo such that:
Training on PT or P The test loss satisfies (p(00))) =
Training on Pex: The test loss satisfies Lp- (0(t)) S Furthermore, we have two bonuses
Training on PT or Ptr When testing on a sample with task concept K** E [K), before adding the residual stream, the model generates hybrid vecior with both high- and low-level components:
-Task Vector Regression from Demonstrations. The transformer could approximately retrieve the task
"This condition ensures that other task concepts within the prompt do not interfere with recognizing the prompt's co-task. It can he relaxed by assuming lower probability of words being as- sociated with task concepts outside the prompt's co-task concept or by limiting the number of task concepts per word.
cos cosshgir..o or )x O11), cos(hgu,, bu.) O(1) (6)
6

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
vector (i.e. satisfies Eq.(7)) solely from demo-pairs, without requiring a query word.
model's fact retrieval capability generalize to rare en- tries. While their settings emphasize fact retrieval, our framework centers on the retrieval of task vectors, a sin- gle token subcase of theirs. Indeed, despite constrained on factual-recall task, the results offer a foundational perspec- tive on why simple task vector arithmetic holds potential for a range of downstream applications, including concept erasure (Ilharco et al., 2023). mitigation forgetting (Jiang et al., 2025), model editing (Li et al., 2025b), and model merging (Wortsman et al., 2022; Matena & Raffel, 2022; Yang et al., 2024a; Lee et al., 2025; Yoshida et al., 2025; Cao et al., 2025; He et al., 2025).
Task Vector Arithmetic in Factual Recall. Given a test sample with co-concept k², the intermediate output highh. can be added to any query with co- concept ke, which also leads to L-p- (0()) Se. c.
The first bonus parallels findings in Hendel et al. (2023); Yang et al. (2025), demonstrating that transformers can re- gressively infer task vectors from demo-pairs. The second aligns with empirical studies Hendel et al. (2023); Merullo et al. (2024), where Merullo et al. (2024) add extracted task vectors from current prompt to queries outside the prompt and still obtain accurate predictions.
For Distribution Shift, when multiple co-task concepts ap- pear in prompt, the transformer would likely form hy- brid task vector:
The following proposition examines the model's capability to address OOD unseen ICL, tasks.
hocero R
Σ wo.k KEK
wa.k > 0, Σ wook = 1. (8) KEA
Proposition 3.4 (OOD Generalization). Under the condi- tions of Theorem 3.2, suppose the transformer is trained on Pos. At test time, for any distribution P E <(Pi.Po '1. the model satisfies:
This suggests soft, weighted integration of concepts, where w denotes the model's inferred likelihood of task k by softmax operation. Such behavior closely resembles Bayesian Model Averaging (BMA) in topic models-an in- terpretation that recent work has proposed as an underlying mechanism of ICL in LLMs (Blei et al., 2001; Xie et al., 2022; Wang et al., 2023):
1. Dictionary Shift Adaptability. The model gener- alizes to unseen dictionaries Unew with high-level (Ka S K), low-level (Kī S (d-Kat and irrele- vant (K; S <d-K' (aiic(faaheeet )EE(K)) and ||a |1 -11b,, || -1,,,, || == 1 for all hi E K2 his E Kg. and lo E K2 with {ak }. (b ). and } mulually orthogonal.
Kg) concepts, where at, E
conic(fas
pa (Xqurr) = T) = -spPo z, T) po T) dz. (9)
Here, z represents the latent concept variables inferred by the model, linking Eq. (9) to Eq. (1) in the factual- recall ICL setting. The integral over corresponds a conic combination of task vectors {ax)kek as described in Eq. (8), computed by the transformer. Moreover, fo (X guerr)) aligns with hore,o ,o + Xquery in our framework.
2. Distribution Shift Adaptability.
The model accommo- dates test prompts containing multiple co-task con- cepts K... C. [K] with [C] < 4. The test prompt length J and the distributions over XT. and YT, for all 1 E [J + 1] are arbitrary, provided that Eg 1T,,,,, E3 Drr,1- = /10).
4. Proof Idea
This proposition contributes to the ongoing discussion of LLMs' linear latent geometry and OOD extrapolation, par- ticularly in relation to Question 1.4 of Reizinger et al.
The high-level intuition of proof is as follows:
For any E >> O, to establish the convergence of the test loss defined by argmax sampling in Eq. (5), the pri- mary task is prove Eq. (6) and Eq. (7) in terms of constant disparity. Additionally, the required success- ful cosine similarity (i.e., how well hoco oo approxi- mates the task vector) must scale with op / /log(e-1) mitigate noise disruption during testing.
(2024), which ask whether the linear latent geometry aids the OOD extrapolation, superior to static embedding meth- ods. Prior works (Li et al., 2023b; Bu et al., 2024a) explore OOD benefits in multi-task classification but assume new label features share consistent signal (e.g., =1). In con- trast, our framework allows diverse lowelevel features (c.g.. bi = =, - beg). highlighting the role of task vector arith- metic in retrieving task message rather than memorizing label-level features, consistent with Todd et al. (2024) on compositional task vectors. Under a Dictionary Shift sce- nario where b² 1 U, word-label pairs involving unseen low-level task concepts (i.e., x = a. + bb + &x and y= == a4 b + Eyy benefit from the model's high-level task retrieval capability. This phenomenon resonates with the "celebrity helps minority" effect in Allen-Zhn & Li (2024), where the presence of tigh-freeuency entities enables the
The cosine similarity in Eq. (7) is closely tied to the learning dynamics of matrix projections along the di- rections of task vector ag for k: [K]. For QA training, the remaining task is to characterize the learning dynamics along these directions (specifically, aj w as and (W(2) an) (w ax)) and the negli- gible updates of other projections.
For ICL-type training. we show that there is only one key difference: the unexpected co-occurrence- induced neg-nggl growth of some |b] (T(T)) be).
7

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
This drives MLP falsely memorize low-level features, which results in imperfect task vector retrieval and is responsible for the constant-level test error.
from either Pes. Pr or P Then for Vk E [K] and
No. there exists some C1-2. 22 C1-2 > 0 such that
ag W a& 2 a2d 12 (12K KR
Cingvt
Main Challenges. We begin by observing that the gradi- ents of WK. Wq. and Wy can be computed as follows:
(11)
ag W ak S oyd1/2K + <2llg
Congvt
512K+K72 001
Epg, [0w [s[(0]= Eng
M m (WoSen). x))
Furthermore, when pir ~ Pp or p ~ 0A for Vy E [=1}, there exists ky E [K] and some constanl 1: Cs such that
0
Epg, (@w, La(0)) - -Brg
Σ of Wx.See - π) TSe.))
12 Goord M
12 w
812 V2log
II Moy Eearx-rr Me (St)
5
Erg, (0w, -Bra
Sx)
Coord 2log(
(12)
ST
8 Jo
M
(10)
where -I = v - UT |e|| | 'j a coefficient defined in Eq. (30), represents the concatenated softmax weights for positions I,....L L 1 in the sentence, and Wi denotes the likelihood of the model's output matching the k-th dictionary token uą (See Eq. (30)). The primary chal- lenge in our analysis stems from the fact that many com- ponents in the gradient computation, such as Luxy of and [[ySST]] II. exhibit intricate monotonicity properties and varying rates of evolution across different phases. This complexity arises inherently from the non-linear interac- tions between the softmax operation, layer normalization, residual connections, and the cross-entropy loss, making it challenging to precisely characterize the dynamics. To address this, we derive separate upper and lower bounds for the discrete gradients of both the attention and MLP components in each phase. Specifically, we introduce a set of six differential equations (see Appendix D.I) as contin- uous surrogates, which effectively capture the upper and lower bounds on the evolution rates of the attention and MLP components across different phases.
Additionally, Lemma F.1 ensures that other types of pro- jections remain constrained near their initialization scale. Also, the update of attention projections along key direc- tions satisfies the recurrence order fz+1 = fa + a( - fta where a is a constant. The accumulated outcome of such updates can be upper bounded by an exponential function and lower bounded by quadratic function, as demon- strated via integration of its continuous-time analogue in Lemma D.5. The precise statement is provided below.
Lemma 4.2. Under Condition 3.1, suppose Eq. (49) holds at iteration t S Ti Consider Pr sampled from either P Por- Pr or PT Then for Vk E [K], there exist some posi- live constants Ca, C4, during [l4, Ti}. where l4 is defined in Lemma F.7, it holds that
Can aolltollttttoo - 14- -1]
(w Q ax) ) ((w@( as)
ofM2K2
44 -4ooog116(2K + k')3/")oadd''
(w aw)) ( ((w as) - Cu((4Vlog(16(2)) + K''² //)
30
30/d C2Cer
2
+
2
exp(
ofdK²
4.1. Phase 1: Linear Growth of MLP and Accelerating Growth of Attention
In this phase, we analyze the regime where a w V a2 = ([[Wy)ST and || [T]] = d(o)d). Given the known order of the denominator in the update rule for Wv V (Eq. 10), the update of a w( V ag can be linearly bounded by update a++1 = a₂ + b, where b is constant. The ac- cumulated update is further upper bounded by the inte- gral of its continuous-time counterpart, as characterized in Lemma D.I. When ptr ~ PQs. the terms bb (TT)) by remain negligible, since QA contains no by- In con- trast, under per ~ Pr or PTA co-occurrence asymme- tries between demo-pairs and query inputs with respect to some =bb induce nontrivial growth of some bb wT(T)) These behaviors are formally summarized below.
4.2. Phase 2: Decelerating Growth of MLP and Attention
In this phase, ag ( a determines the order of /WSTT/, which, as the denominator the gradient, leads to deceleration in the updates of both a w V as and the attention projections. As in the previous phase, there is nontrivial growth in some |b ((T)) byl when trained ICL-type data (i.e., pir Pr or P(A) due to occurrence asymmetries. The MLP dynamics slow to a sublinear rate, with their analogue captured in Lemma D.3. In contrast, the evolution of the attention projections is bounded between e 1/t and e' as in Lemma D.2 and Lemma D.6. The formal statement as below. Lemma 4.3, Under Condition 3.1, during t1 STST== 2(n of K M log(?)) where i S Tr. Consider pr
v
V byl.
Lemma 4.1. Under Condition 3.1. during t S Ti < tr. where lo is defined in Lemma F5. Consider P sampled
8

---

Provable In- In-Context Vector Arithmetic via Retrieving Task Concepts
Train Loss
Test Losses
[W] (. (.)
UCT V.
5.2
3.4
Bain Looa
1.00
Losa CL
25
0.20
0.75
MUJURUUUA
8.:15
0.30
20
BUJUKO_UU)
0.10
2.0
0.25
E005
24
350
3000
1500
2000
3.00
500
1030
1500
2000
500
3500
2000
0.00
500
3000
1520
2020
Spech
(a) Training loss
(h) Test losses
(c) u W K
W Q
(d) a Wy
Figure 3. Training dynamics over the QA sentence training distribution. Legends are consistent with those in Figure 2. Compared to the ICI--trained model in Figure 2, Wy trained on QA data does not overfit to low-level features by (sce (d)), leading to improved generalization and reduced test errors in (b).
sampled from either Pos. Pr or PTT Then for Vk [K], there exists some C7 such that
6. Conclusion, Limitations, and Future Work
This work presents an optimization theory for non-linear residual transformer's task vector arithmetic mechanisms under empirically-
a w/
/ (tv - fl)
+ (a with P.)?
ax
MK
hierarchical modelings. Our analysis framework provides notable merits for such mech- anisms, including flexible autoof-distribution generaliza- tion and the ability handle multi-concept words, supe- rior to traditional static Word2 predecessors lacking the flexible generalization potential.
a WWWSa ax S
wt 10---1 [1] + (a) wif a.)"
CanAv
K
Kía W Wit ag)
aw) (W2(a)) s Coagdeniooy ngut
ofdKMs 207 2
(w)))g) (w/a:) (2)) 2 Cia abofde MPqy"
Dopax
However, our framework is currently limited to single- token settings, whereas real-world factual retrieval often requires multi-token reasoning, demanding more sophisti- cated probing techniques to uncover internal mechanisms beyond simple geometric relationships (Allen-Zhu & Li, 2024; Zhang et al., 2025a). Morcover, while our analy- sis is grounded in empirical modeling, it remains idealized and does not investigate how task vectors naturally emerge in the deeper layers of transformer models (Knittel et al., 2024; Yang et al., 2025). In practice, the representation of high-level concepts is not strictly equivalent to the cor- responding task vectors-e.g. the Qeity in Merullo et al. (2024) does not exactly match the latent representation of "Capital", Additionally, real-world vocabulary dictio- naries consist of output embeddings of natural language to- kens (c.g., words or labels) that are inherently polysemous and evolve over the course of training, rather than remain- ing static. This dynamic evolution of semantic representa- tions likely co-occurs with the emergence of retrieval capa- bilities, potentially in self-reinoorcing manner.
Con 0RMPK² 4Css -4Csylogg16 1og(16 + /
where A is defined in Lemma F.9. Furthermore, when p" ~ PT or p" ~ -oA for Vy E {+1], there exists ky E [K], Wk E [K] we have
16, W! bel, 167. b= O(a W( ax).
(13)
After a sufficient number of epochs beyond T*, for mod- els trained PoA. the task vector ag- becomes the dom- inant direction in hæro o = ( V (1) during test- ing when encountering ICL prompts with A* as the co-task concept. In contrast, for models trained on pir Pr and p"r PH the components bbb,, remain negligible, which supports Theorem 3.2. Building on these learned projec- tions, the proofs for the propositions follow.
5. Experiments
In this section, we present simulations of Algorithm 1. For comparison, we use the same parameters for both ICL- trained and QA-trained models and plot the I-sigma er- ror dynamics, as illustrated in Figures 2 and 3: = 2.
Besides, LLMs often rely on more sophisticated or unex- plainable operations task vectors for complex tasks, be- yond simple vector arithmetic (Hendel et al., 2023; Todd et al., 2024; Zhang et al., 2025a). While the ground-truth leveragggg function f(r, ax-) (Hendel et al., 2023) may lack a closed-form, the residual architecture in transformers may implicitly ap- proximate such a function. Future work could explore how these components enable LLMs to interact with task vec- tors across broader tasks.
K' 100, d 3000, n 200, M 30, L L 30. " 5, vv = 10-5 do 10-3, 01 5 x 10-3 Op og 10-2. The QA-trained model is trained for T = 2000 epochs, while the ICL-trained model under- goes a longer training process with T' = 5000 epochs. De- spite the extended number of iterations, the ICL-trained and QA-ICL-trained model (Figure 2 and Figure in Appendix B) maintain a constant test error, whereas the QA-trained model (Figure 3) converges to zero rapidly. These findings validate main theorems.
9

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Impact Statement
Bietti, A., Cabannes, V., Bouchacourt, D., Jegou, H., and Bottou, L.. Birth of transformer: A memory viewpoint. In Advances in Neural Information Processing Systems, volume 36, pp. 1560-1588, 2023.
This paper presents work whose goal 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.
Blei. D., Ng, A., and Jordan, M. Latent dirichlet alloca- tion. In Dietterich, T., Becker, S., and Ghahramani, Z. (eds.), Advances in Neural Information Processing Sys- tems, 2001.
Acknowledgements
We thank the anonymous reviewers for their instrumen- tal comments. DB and HW are supported in part by the Research Grants Council of the Hong Kong Special Administrative Region (Project No. CityU 11206622). WH was supported by JSPS KAKENHI Grant Number 24K20848. TS was partially supported by JSPS KAK- ENHI (24K02905) and JST CREST (JPMJCR2115), This research is supported by the National Research Founda- tion, Singapore, Infocomm Media Development Author- ity under its Trust Tech Funding Initiative, and the Min- istry of Digital Development and Information under the Al Visiting Professorship Programme (award number AIVP- 2024-004). Any opinions, findings and conclusions recommendations expressed in this material are those of the author(s) and do not reflect the views of National Re- search Foundation, Singapore, Infocomm Media Develop- ment Authority, and the Ministry Digital Development and Information.
Bu, D., Huang, W., Han, A., Nitanda, A., Suzuki, T., Zhang, Q. and Wong, H.-S. Provably transformers harness multi-concept word semantics for efficient in- context learning. In The Thirty- Anmual Confer- ence on Neural Information Processing Systems, 2024a.
Bu, D., Huang, W., Suzuki, T., Cheng, J., Zhang, Q. Xu, Z., and Wong, H.-S. Provably neural active learning suc- ceeds via prioritizing perplexing samples. In Proceed- ings of the 41st International Conference on Machine Learning, volume 235, pp. 4642-4695, 2024b.
Cabannes, V., Simsek, B., and Bietti, A. Learning as- sociative memories with gradient descent. In Interna- tional Conference on Machine Learning, pp. 5114-114111 PMLR, 2024.
Cao, S., Wu, M., Prasad, K., Tian, Y., and Liu, Z. ParamS' deltaS for direct mixing: Post-train large lan- guage model at zero cost. In The Thirteenth In- ternational Conference on Learning Representations, 2025. URL httph:tt//////tt http////ppe/eeeeeee/eeeeeeeeeeet n forum? i id=vgbd200nGp.
References
Ahuja, K. and Lopez-Paz, D. A closer look at in- context learning under distribution shifts. arXiv preprint arX1v:2305.16774, 2023.
Cao, Y., Chen, Z., Belkin, M., and Gu, Q. Benign overfit- ting in two-layer convolutional neural networks. In Ad- vances in Neural Information Processing Systems, vol- ume 35, pp. 25237-25250, 2022.
Allen-Zhu, Z. and Li, Y. Physics of language models: Part 3.1, knowledge storage and extraction. In Proceedings of the 41st International Conference on Machine Learning, volume 235, pp. 1067-1077, 2024.
Chang, X., Li, Y., Kara, M., Oymak, S., and Roy- Chowdhury, A. Provable benefits of task-specific prompts for in-context learning. In The 28th Interna- tional Conference on Artificial Intelligence and Statis- tics, 2025. URL https::/t//.cm httpp::///peerevviee net./ Berumiid=INPP15HaT77
Allen-Zhu Z. and Li, Physies of language models: Part 3.3, knowledge capacity scaling laws. In The Thirteenth International Conference on Learning Representations, 2025. URL httpt:://ppnreeeeeeeeeee http:///epeeeeeeeee//eesee net formmm?rm id=FxNNiUgtfa.
Bai, Y. Chen, F,, Wang, H., Xiong. C., and Mei, S. Transformers as statisticians: Provable in-context learn- ing with in-context algorithm selection. In Advances in Neural Information Processing Systems, volume 36, pp. 57125-57211, 2023.
Chatterji. N. and Long, M. Finite-sample analy- sis of interpolating linear classifiers in the overparam- eterized regime. Journal of Machine Learning Re- search, 22(129):1-30, 2021. URL /ttp: org/paperz httpppppppp///////aaaaa 222/2222222222 111.9774........
Bernays, P. Alonzo church. an unsolvable problem of ele- mentary number theory. american journal of mathemat- ics, vol. 58 (1936), pp. 345-363. The Journal of Sym- bolic Logic, 1(2):33-74, 1936.
Chen, S., Sheen, H., Wang, T., and Yang, Z. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. In COL.T, 2024a.
10

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Chen, S., Sheen, H., Wang, T., and Yang, Z. Unveiling induction heads: Provable training dynamics and fea- ture learning in transformers. In Advances in Neural Information Processing Systems, volume 37, pp. 66479- 66567. Curran Associates, Inc., 2024h.
on Learning Representations, 2023. URL https: // openreview net
Jiang, G., JIANG, C., Li, Z., Xue, S., ZHOU, J., Song. L.,,
Lian, D., and Wei, Y. Unlocking the power function vectors for characterizing and mitigating catastrophic forgetting in continual instruction tuning. In The Thir- teenth International Conference Learning Represen- tations, 2025. URL https: http:://ppppeeeeeeeeeee net./ forum?id==c8QQQEXX
Fan, Y., Yadlowsky, S., Papailiopoulos. D., and Lee, K. Transformers can learn meta-skills for task generaliza- tion in in-context learning, In NeurlPS 2024 Workshop on Compositional Learning: Perspectives, Methods, and Paths Forward, 2024.
Jiang, J., Huang, W., Zhang, M., Suzuki, T., and Nie, L. Unveil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization. In The Thirty-eighth Annual Conference on Neural Infor- mation Processing Systems, volume 37, pp. 135464- 135625, 2024a.
Frei, S., Chatterji, N. S., and Bartlett, P. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. Proceedings of Thirty Fifth Conference on Learning Theory, volume 178. pp. 2668-2703, 2022.
Garg, S., Tsipras, D.. Liang. P. and Valiant, G. What can transformers learn in-context? A case study of simple function classes. In Advances in Neural In- formation Processing Systems, 2022. URL https: / //openzeview wttprrrrrree.eeefffooooaaaaa--22
Jiang, Y., Rajendran, G., Ravikumar, P. K., Aragam, B., and Veitch, V. On the origins of linear representations in large language models. In Proceedings of the s1 International Conference on Machine Learning, volume 235, pp. 21879-21911, 2024b.
Ghosal, R., Hashimoto, T., and Raghunathan, Under- standing finetuning for factual knowledge extraction. In Proceedings of the 41 International Conference Ma- chine Learning, volume 235, pp. 15540-15588, 2024.
Jin, X., Ren, X., Preotiuc-Piette Pietro, D., and Cheng, P. Data- less knowledge fusion by merging weights of language models. In The Eleventh International Conference on Learning Representations, 2023. URL https:// openreview net /
He, Y., Hu, Y., Lin, Y., Zhang, T., and Zhao, H. Localize- and-stitch: Efficient model merging via sparse task arithmetic. Transactions
Machine Learning Re- search, 2025. ISSN 2835-8856. URL https:// www/////oouneee wwt.ooormmmmimueeeedeee dwwwd0--cw----oo 86d.
Kahardipraja, P., Achtibat, R., Wiegand, T., Samek, W., and Lapuschkin, S. The atlas of in-context learning: How attention heads shape in-context retrieval augmen- tation. arXiv preprint arXiv: 2505. /5807, 2025.
Hendel, R., Geva, M., and Globerson, A. In-context learn- ing creates task vectors. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 93 188- 9333, 2023.
Kim, J. and Suzuki, T. Transformers learn nonlinear fea- tures in context: Nonconvex mean-field dynamics on the attention landscape. In International Conference on Ma- chine Learning, pp. 24527-24444 24561. PMLR, 2024.
Huang, R. and Ge, R. Task descriptors help transformers learn linear models in-context. In ICML 2024 Work- shop on In-Context Learning, 2024. URL https: / /popnreeveew.... net wtte/foouummeee wtt////rrmm/amm//mmm//// iineiid--4see 4SfCIIDD.
Knittel, J., Gangavarapu, T., Strobelt, H., and Pfister, H. Gpt-2 through the lens of vector symbolic architectures. arXiv preprint arXiv: 2412.07947, 2024.
Huang, Y., Cheng, Y., and Liang, Y. In-context conver- gence of transformers. In Proceedings of the 41st Inter- national Conference on Machine Learning, pp. 19660- 19722, 2024a.
Kossen, J.. Gal, Y., and Rainforth. T. In-context learn- ing learns label relationships but is not conventional learning. In The Twelfth International Conference on Learning Representations, 2024. URL https:// openreview net forum? wwwoumiiid-ee -iee:id=PPPP Midd=PIIA7pgd55
Huang, Y., Wen, Z., Chi, Y., and Liang. Y. Transformers provably learn feature-position correlations in masked image modeling. arXiv preprint arXiv: 2403.02233, 2024b.
Kou, Y., Chen, Z., Chen, Y., and Gu, Q. Benign overfit- ting in two-layer reLU convolutional neural networks. Proceedings of the 40th International Conference on Machine Learning, volume 202, pp. 17615-17659, 2023.
Ilharco, G., Ribeiro, M. T., Wortsman, M., Schmidt, L., Hajishirzi, H., and Farhadi, A. Editing models with task arithmetic. The Eleventh International Conference

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Kwon, S. M., Xu, A. Yaras, C., Balzano, L.., and Qu, Q. Out-of-distribution generalization of in-context learning: A low-dimensional subspace perspective. arXiv preprint WrXi:22505.44808, 2025.
Liu, Y., Gong, D., Cai, Y., Gao, E., Zhang, Z., Huang, B., Gong, M., van den Hengel, A., and Shi, J. Q. predict therefore i am: Is next token prediction enough to learn human-interpretule concepts from data? arXiv preprint arXiv:2503. 009880 2025.
Lee. S., Liu, J., Wang, Q., Wang, J., Cai, X., and Wu, Y. Dynamic fisher-weighted model merging via Bayesian optimization. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Associ- ation for Computational Linguisties: Human Language Technologies (Volume I: Long Papers), pp. 4923 4935, 2025.
Lu, S., Bigoulaeva. Sachdeva, R., Madabushi, H. T., and Gurevych, I. Are emergent abilities in large language models just in-context learning? In Proceedings of the 62nd Annual Meeting of the Association for Compula- tional Linguistics (Volume 1: Long Papers), pp. 5098- 5139, 2024.
Li, B., Huang, W., Han, A., Zhou, 7., Suzuki, T., Zhu, J., and Chen, J. On the optimization and generaliza- tion of two-layer transformers with sign gradient de- scent. In The Thirteenth International Conference on Learning Representations. 2025a. URL https:// openreview net htrr//orrmmm/m Seneuum?id--pe 97 roggDPmk2.
Marconato, E., Lachapelle, S., Weichwald, S., and Gre- sele, L. All none: Identifiable linear properties of next-token predictors in language modeling- In The 28th International Conference on Artificial Intelligence and Statistics, 2025. URL https://///// /openreviee.... net/ net/forumg Rid- XCMIIImmImmee...
Li, H., Wang, M., Liu, S., and Chen, P.-Y. A theoreti- cal understanding of shallow vision transformers: Learn- ing, generalization, and sample complexity. In The Eleventh International Conference on Learming Repre- sentations, 2023a. URL https:/ openreview. net/forum? deriddddjcciaannnee
Marks, S. and Tegmark, M. The geometry of truth: Emergent linear structure in large language model rep- resentations of true/false datasets. In First Confer- ence on Language Modeling, 2024. URL https:// openreview net forumm poorummiddaaaiii aajjyHjjjj..
Matena, M. S. and Raffel, C. A. Merging models with fisher-weighted averaging. In Advances in Neural In- formation Processing Systems, volume 35, pp. 17703- 17716, 2022.
Li, H., Zhang, Y., Zhang, S., Chen, P.-Y,, Liu, S., and Wang, M. When task vector provably ef- generalization anal-
fective for model editing? ysis of nonlinear transformers. In The Thirteenth International Conference on Learning Representa- tions, 2025b. URL /ttl:ttts//:/peeeeiieee /wwp//ppweeee...eeo///.. net// forum? forum?id-vRrviie VRVWVbbNN....
Meng. K., Bau. D., Andonian, A., and Belinkov, Y. Locat- ing and editing factual associations in gpt. In Advances in neural information processing systems, volume 35, pp. 17359-17372, 2022.
Li, Y. Li, Y., and Risteski, A. How do transformers learn topic structure: Towards mechanistic understanding. In Proceedings of the 40th International Conference on Machine Learning, pp. 19689-19729, 2023b.
Meng, X., Zou, D., and Cao, Y. Benign overfitting in two-layer ReLU convolutional neural networks for XOR data. In Proceedings of the 41st International Confer- ence on Machine Learning, volume 235, pp. 35404- 35469, 2024.
Li, Y., Rawat, A. S., and Oymak, S. Fine-grained analysis of in-context linear estimation: Data, architecture, and beyond. In Advances in Neural Information Processing Systems, volume 37, pp. 188324-133366 2024.
Mcrullo. J., Eickhoff, C., and Pavlick, E. Language mod- els implement simple word2vec-style vector arithmetic. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 5030-5047, 2024.
Liang H., Balasubramanian, K., and Lai, L. Trans- formers handle endogeneity in-context linear regres- sion. In The Thirteenth International Conference on Learning Representations, 2025. URL https:// openreview. openreview.nett/orrmm www.ooorrummiiidooooeeeeeee dd-ofh3322221..
Mikolov, T., Yih, W.-L, and Zweig, G. Linguistic regular- ities in continuous space word representations. In Pro- ceedings of the 2013 conference of the north american chapter of the association for computational linguisties: Human language technologies, pp. 746-751, 2013.
Liu, S., Ye, H., Xing, L.,, and Zou, J. Y. In-context vectors: Making in context learning more effective and control- lable through latent space steering. In Proceedings of the 41st International Conference on Machine Learning, volume 235. pp. 32287-32307, 2024.
Minegishi, G., Furuta, H., Iwasawa, Y., and Matsuo, Y. Rethinking evaluation of sparse autoencoders through
12

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
the representation of polysemous words. In The Thir- teenth International Conference on Learning Represen- tations, 2025a. URL https:/ http:///ppnneeeeeeeeeeeeeeeeeeee forum? HHUHSS
Shen, W. Zhou, R., Yang. J., and Shen, C. On the train- ing convergence of transformers for in-context classifi- cation. arXiv preprint arXiv:2410.11778, 2024.
Tian, Y., Wang. Y., Chen, B., and Du, S. S. Scan and snap: Understanding training dynamics and token composition in I-layer transformer. In Advances in Neural Informa- tion Processing Systems, volume 36, pp. 71911 19911-7994, 2023.
Minegishi, G., Furuta, H., Taniguchi, S., Iwasawa, Y., and Matsuo, Y. Beyond induction heads: In-context meta learning induces multi-phase circuit emergence. arXiv preprint arXiv: 2505.16694 2025b.
Nichani, E., Damian, A., and Lee, J. D. How transformers learn causal structure with gradient descent. In Proceed- ings of the s International Conference on Machine Learning, pp. 38018-38070, 2024.
Tian, Y., Wang, Y., Zhang, Z., Chen, B., and Du, S. S. JoMA: Demystifying multilayer transformers via joint dynamics of MLP and attention. In The Twelfth In- ternational Conference on Learning Representations, 2024. URL https: //tp/////e///eeeee////me net/ forum?
Nichani, E., Lee, J. D., and Bietti, A. Understanding factual recall in transformers via associative memories. In The Thirteenth International Conference on Learning Rep- resentations, 2025. URL https:// https://openrevvew openreview. net/forume wwwp:iiwwwwwwwwww........
id=1bJGRGNYCE.
Todd, E., Li, M., Sharma, A. S., Mueller, A., Wallace, B. C., and Bau, D. Function vectors in large lan- guage models. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: / //operrview. net meet/foruuee htt///orummm///eee ddd-AwixtyMaaese
Oko, K. Song, Y., Suzuki, T., and Wu, D. Pre- trained transformer efficiently learns low-dimensional target functions in-context. In The eiiht Annual Conference on Neural Information Processing Systems, 2024. URL hwtps::ttss/cc.... hatttss//openzevieew net/forum? 1d=uHcG5Y6fdB.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, u., and Polosukhin, Atten- tion is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017.
Wainwright, M. J. High-dimensional statistics: A non- asymptotic Viewpoint, volume 48. Cambridge university press, 2019.
Pan, J., Gao, T., Chen, H., and Chen, D. What in-context learning "learns" in-context: Disentangling task recog- nition and task learning. In The 61st Annual Meeting Of The Association For Computational Linguistics, 2023.
Wang. and Komatsuzaki, A. Gpt-j 6b: A bil- lion parameter autoregressive language model. https://github..cm mesh-transformert jjan 2021.
Park, K., Choe, J., and Veitch, V. The linear representa- tion hypothesis and the geometry of large language mod- els. In Proceedings of the 41 International Conference on Machine Learning, volume 235, pp. 39643-39666, 2024.
Wang, Q. Wang, Y., Wang, Y., and Ying, X. Can in- context learning really generalize out-of-distribution tasks? In The Thirteenth International Conference on Learning Representations, 2025. URL https:// openreview net
Park, K., Choe, J., Jiang Y., and Veitch, V. The geome- try of categorical and hierarchical concepts in large lan- guage models. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL ht ats::: //openzeviews /aat/ffooumtt httpr//orrumm2..ss//2222
Wang, X., Zhu, W., Saxon, M., Steyvers, M., and Wang. W. Y. Large language models are implicitly topic mod- els: Explaining and finding good demonstrations for in- context learning. In Workshop on Efficient Systems for Foundation Models @ ICML2023, 2023.
Raventós, A., Paul, M., Chen, F., and Ganguli, S. Pretrain- ing task diversity and the emergence of non-bayesian in- context learning for regression. In Advances in neural information processing systems, volume 36, pp. 14228- 14246, 2023.
Wibisono, K. C. and Wang, Y. Bidirectional attention as mixture of continuous word experts. In Uncertainty in Artificial Intelligence, pp. 2271-2281. PMLR, 2023.
Reizinger, P., Ujváry, S., Mészáros, A., Kerekes, A., Bren- del, W., and Huszár, F. Position: Understanding LLMs requires more than statistical generalization. In Forty- first International Conference on Machine Learning, 2024. URL https:/ http///pppeeeeeeee//eee// net forum? dd=pVyOchMUBa.
Wortsman, M., Ilharco, G., Gadre, S. Y., Roelofs, R., Gontijo-Lopes, R., Morcos, A. S., Namkoong, H., Farhadi, A., Carmon, Y., Kornblith, S., and Schmidt, L. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference
13

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
time. In Proceedings of the 39th International Con- ference Machine Learning, volume 162, pp. 23965- 23998, PMLR, 2022.
Zhang, Y., Zhang, F., Yang, Z., and Wang. Z. What and how does in-context learning leam? bayesian model eraging, parameterization, and generalization. In The 28th International Conference on Artificial Intelligence and Statistics, 2025b. URL https: / openreview. net/ forum -ide-iddBB00 B500F 0Fc60.
Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Rep- resentations, 2022. URL https:/ openreview.
Yadlowsky, S., Doshi, L., and Tripuraneni, N. Pre- training data mixtures enable narrow model selection capabilities in transformer models. arXiv preprint arXiv:2311.00871, 2023.
Yamagiwa, H., Oyama, M., and Shimodaira, H. Discover- ing universal geometry in embeddings with ica. In Pro- ceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 4647-4673 2023.
Yang, E., Wang, Z., Shen, L., Liu, S., Guo, G., Wang, X., and Tao, D. Adamerging: Adaptive model merging for multi-task learning. In The Twelfth International Conference Learning Represen- tations, 2024a. URL https://openreview n forum? did- 6NgD30Y.
Yang, H., Kailkhura, B., Wang, Z., and Liang, Y. Train- ing dynamics of transformers to recognize word co- occurrence via gradient flow analysis. Advances in Neural Information Processing Systems, volume 37, pp. 46047-46117, 2024b.
Yang, L., Lin, Z., Lee, K., Papailiopoulos, D., and Nowak, R. Task vectors in in-context learning: Emergence, for- mation, and benefit. arXiv preprint 0092422501.09240, 2025.
Ye, N., Yang, H., Siah, A., and Namkoong, H. Pre-training and in-context learning bayesian inference a de finetti. arXiv preprint arXiv: 2408.03307 2024.
Yoshida, K., Naraki, Y., Horie, T., Yamaki, R., Shimizu, R., Saito, Y., McAuley, J., and Naganuma, H. Mas- tering task arithmetic: Stttuujp
key indicator for weight disentanglement. In The Thirteenth In- ternational Conference on Learning Representations, forum?
2025. URL
id=1VwWi6zbxs.
Zhang, R., Frei, S., and Bartlett, P. L. Trained transform- learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
Zhang, Y. Heinzerling, B., Li, D., Ishigaki, R., Hitomi, Y., and Inui, K. Understanding fact recall in language mod- els: Why two-stage training encourages memorization but mixed training teaches knowledge. arXiv preprint arXiv: 16178. 2025a.
14

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
A. Additional Related Work
Task Vector Mechanism. This line of research investigates the emergence of task vectors through controlled experiments. Hendel et al. (2023) and Todd et al. (2024) independently observed the emergence of task or function vectors in tasks such as translation and factual recall. Knittel et al. (2024) demonstrated that GPT-2 effectively learns mechanisms analogous to vector symbolic architectures, where different blocks communicate by writing and reading nearly orthogonal vectors from the residual stream. Liu et al. (2024) proposed framework that recasts in-context learning as in-context vector manipulation. Yang ct al. (2025) explored the emergence and benefits of task vectors in in-context learning (ICL). Finally, Merullo et al. (2024) revealed that large language models (LLMs) implement certain single-token factual-recall tasks through vector arithmetic. However, no existing study explains why and how paddient-uddate transformer models naturally implement this mechanism.
Applications of Task Vectors. Owing to the inherent sparsity in LLM representations, task vector arithmetic has emerged as efficient and interpretable method for model merging (Lee et al., 2025; Yoshida et al., 2025; Cao et al., 2025), ex- plicitly via arithmetic mean (Wortsman et al., 2022; Ilharco et al., 2023), Fisher information weighting (Matena & Raffel, 2022), regression-basee mean (Jin et al., 2023), or learned merging weights (Yang et al., 2024a). Recently, He et al. (2025) introduced a localized merging approach to better exploit task arithmetic. Beyond merging, Li et al. (2025b) demon- strate task vectors' effectiveness for model editing, and Jiang et al. (2025) reveal their promise in mitigating catastrophic forgetting.
Storage and Retrieval of Factual Knowledge. Recent research has explored how large language models (LLMs) per- form factual recall or associative memory tasks (Meng al., 2022; Nichani al., 2025; Cabannes et al., 2024; Ghosal et al., 2024; Allen-Zhu & Li, 2024; Marks & Tegmark, 2024). These studies span a wide range of topics, including the expressiveness of LLMs (Bietti et al., 2023), the optimality of memory storage and optimization dynamics (Cabannes et al., 2024; Nichani et al., 2025), and the empirical mechanisms by which LLMs encode factual knowledge and their inherent limitations. Allen-Zhu& & Li (2025) empirically demonstrated the memorization capacity of transformer language models of varying sizes trained on synthetic factual recall tasks, observing near-linear scaling with the number of parameters. (Allen-Zhu & Li, 2024) found that QA data is essential for learning to retrieve factual knowledge in biography data. Sub- sequently, (Zhang et al., 2025a) further showed that training on QA data, when mixed with Bio data, would increase both quantity and importance of shared parameters compared to training upon biography and QA data in sequential manner, and (Kahardipraja et al., 2025) revealed that different attention heads play distinct roles-in context heads comprehend instructions and retrieve relevant contextual information, while parametric heads store relational knowledge. Additionally, Park et al. (2024; 2025); Jiang et al. (2024b); Marks & Tegmark (2024); Marconato et al. (2025); Liu et al. (2025) provided both theoretical and empirical evidence that L.LMs tend to represent independent concepts, words, and factual statements in a linear manner. Minegishi et al. (2025a) showed the Sparse autoencoders (SAEs) have their limits for representing polysemous contexts, while deeper layers of LLM contributes to distinguishing the polysemy. These observations inspire further investigation into the underlying mechanisms of factual knowledge retrieval.
Empirical Investigations of ICL Mechanisms. Garg et al. (2022) showed that Transformer-batt in-context learning (ICL) is robust to distribution shifts, motivating a series of studies on cut-of-distributionn (OOD) generalization (Yadlowsky et al., 2023; Raventós et al., 2023; Ahuja & Lopez-Paz, 2023; Kossen et al., 2024; Pan et al., 2023; Fan et al., 2024; Wang et al., 2025). Yadlowsky et al. (2023) and Wang et al. (2025) analyze mixtures of function classes and show that Transform- ers struggle to generalize unseen ones; Kwon al. (2025) further extend this to mixtures over multiple distributions and provide optimal parameterization for regression. However, these works largely overlook ICL's core capability: inferring the underlying function/task behind demonstrations function reference). This aligns with the perspective that identifying high-level latent concepts corresponds to task inference, a view supported by empirical findings (Wang et al., 2023). While BMA-based models (Zhang et al., 2025b; Ye et al., 2024) offer a theoretical lens, they rely on unrealistic assumptions that distort Transformer architectures for kernel regression. Empirical work has instead uncovered that LLMs encode task identity via function or task vectors during ICL (Hendel et al., 2023; Todd et al., 2024; Yang et al., 2025; Merullo al., 2024), and that multi-head attention-only models exhibit multi-phases in forming task circuits (Minegishi al., 2025b). Further, Kahardipraja et al. (2025) identify distinct roles of in-context and parametric heads in function comprehension and knowledge storage. Yet, these works do not systematically connect task vector mechanisms with BMA-style inference or with OOD generalization in ICL, leaving a gap that we aim to fill.
Theoretical Investigations of ICL Mechanisms, A recent line of theoretical research investigates how transformers learn in-context in various scenarios (Zhang et al., 2024; Tian et al., 2023; Nichani et al., 2024; Chen et al., 2024b; Liang et al.,
15

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
2025; Shen al., 2024; Huang & Ge, 2024; Chen et al., 2024a). However, these studies often rely on idealized assumptions, such linearized query-key (QK) combinedd attention mechanisms (Zhang al., 2024; Tian et al., 2023; Nichani al., 2024; Chen et al., 2024b; Oko et al., 2024; Zhang al., 2025b), employ simplified loss functions like squared or hinge loss (Chen et al., 2024a; Huang et al., 2024b; Li et al., 2023a; 2024; Chang et al., 2025). Furthermore, rather than investigating how transformers inherently process task demonstrations T, these studies adopt specialized embedding schemes, such as the "up-word-down-label" approach:
E(T) =
01: JJ+ y1,
Equery
yJ:
where words and labels are processed separately by attention and MLP layers (Bai et al., 2023; Li et al., 2023a; 2024; Bu al., 2024a; Chang et al., 2025). In contrast, we consider the case where the prompt T ( y1.-. yJ, [Qqery [T,. .......7. Tr] d (L = 2J + +) is directly processed by a non-linear transformer with a residual stream. Besides, in these work's s construction, there must be similar label-related (low-level) patterns in the demonstration pairs, and essentially not showed the function reference of transformer models during ICL. For example, given the prompt Japan Tokyo Francee Paris s China, the expected output "Beij jing" reflects the underlying function "Country' s Capital" applicd to the final query, not dependent on prior low-level label semantics in "Japan" or "France". In addition, none of these theoretical works consider the weal-world latent geometric relationship between words and labels depicted in Figure 1, nor do they explain the observed in-context vector arithmetic.
B. Additional Experiments
Tain Loss
Test Losses
Ku)"T (W Qvv)
W'
3.25
3.00
Lass
0.4
U.K.V.U
OUKVOUU
3.00
3:75
K.K.V.KKK
8 5.755
3:50
OUK.UUU
2.50
0:35
0.0
2.25
2.00
1000
2000
3000
4000
5000
3.00
1000
2003
3050
4000
5000
-50
1000
2000
2000
4000
5000
-02
1020
2000 1000
5050
Basch
Baach
(a) Training loss
(b) Test losses
(c) WWWWo
(d) u Wrv
Figure 4. Training Dynamics over the QA-ICL Training Distribution. The definitions of legends follows Figure 2. In (d), Wy overfits to low-level features bg. resulting in persistently constant (0.2) test error shown (b).
We also conducted experiments on the scenario when the training is conducted on QA-ICL data. For comparison, we adopt
the similar parameter configuration :K = 2, K' 100, d = 3000, n = 200, M = 30, L = L 30, = 5, qv = 1e -5 Oo- = le-3 01 = 5 le op- = = = 12- and = 5000.
C. Model Details
In this section, we present the missing details of our learning problem.
C.1. Training Algorithm
Our training algorithms for transformer is presented in Algorithm 1.
C.2. Hierarchical Data Modeling
To model the empirically observed task vector arithmetic mechanism in factual-recall ICL as illustrated in Figure 1, where the answer/label vectors y tend align more strongly with the task vector ag than with query word vectors σ, draw inspiration from recent work Park et al. (2025) on hierarchical concept geometry. This work observes that LLMs tend to encode high- and low-level vectors in an approximately orthogonal manner within their latent space, and that concept vectors at the same level, which are semantically independent, also exhibit mutual orthogonality observation that is further supported by both theoretical and empirical studies Jiang et al. (2024b); Yamagiwa et al. (2023); Park et al. (2024). Furthermore, Figure Park et al. (2025) suggests that word representations reside within polytopes, aligning with the
16

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
empirical illustration in Figure 1. Motivated by these findings, natural theoretical modeling approach is to treat the task vector high-level concept vector, while the components orthogonal to the task vector correspond task-specific low-level concept vectors, which we define as follows.
Definition C.I. Word-Label Pair ICL Prompt Distribution (PT). The prompt T (I, J,,YJ, XJ+1 i| (xx(2/+1) and y+1 R is generated as follows. Each word-label pair in T shares high-level task concept with uniform probability, and each pair has an equal chance of possessing positive or negative low-level label feature. Specifically, for all I E [J + 1], the word and label vectors are defined as:
kT ~ Uni Mord ~ Uni[[+1], E ~ (N(0.
TT. {MT) Uui [K] {KT}: B, ~ Ber ())((K)), Bi = 1}.
Tr = UiEE [K] {KT} B, ~ Ber(K -K) ",, B3 = 1}.
yka Unif[=1], wka E Nr, \ {KT).
X; : Σ
Fa au + you! bk:) E1,x
(14)
kE Yiy = Σ
akky + yeugtu bREN) ) + Ely Vly < J +1
ki, EXIriyy
YJ+1 := Okr + U+r bkr
where the noise vector & accounts for the inherent semantic inexactness of latent representation for language. The expected label vector of the prompt, yJ-1, is constrained by the co-task concept by of the prompt and remains noiseless to ensure an accurate semantic representation of the concept-specific answer. The term Ta ag acts as an anchor, guiding the transformer to retrieve high-level task concepts, with x, |as / l| S o(1). For simplicity, we set Fa :- 0.1.
Remark C.2. To model polysemy, each word vector X12 (l2 E [J + 1)) and label vector yı, (ly E [J] is associated with multiple task concepts. However, all tokens a prompt share the co-task concept kr. For example, "Apple" may relate to both "Color" and "Species" task concepts, with labels "Red" and "Fruit", respectively. Likewise, "Fruit" may indicate "Outcome" in contexts unrelated to "Color". To reduce the freedom of our learning problem, assume the ground truth label follows single-concept noiseless form: YJ+1 Okr + 1++,1++1 bKT-
This is reasonable, as kT naturally disambiguates label meanings. For instance, while "Fruit" could mean "Outcome", under the "Species" task, its meaning is unambiguous.
Failure of ICL-trained Transformer. Under the above definitions, we conduct experiments across various parameter settings in the procedure of Algorithm 1. However, as illustrated in Figure 2, the test loss always fails to converge, remaining at a constant level. Moreover, the projection dynamics indicate that the value matrix Wy learns low-level concepts in a disorganized manner, suggesting that hT, becomes hybrid vector containing both agr and b Consequently, this compromises its role approximating task vector: the presence of low-level concept components ebst, [+1}, within (T)) disrupts the vector arithmetic in Eq. (3), biasing it toward specific ebay. This leads to degraded predictions, particularly when X, contains -ebar.
Reason of Failure: ICL Data vs. QA Data. At high level, the failure arises from unintended co-occurrences between low-level features in the demonstration pairs and those in the query words. In contrast, QA data contains only the task vector and irrelevant tokens in XQA, avoiding the possibility of such co-occurrences. According to Yang et al. (2024b), transformers tend to encode fixed co-occurrence patterns via gradient descent. Moreover, training and testing solely on the ICL data distribution- thoughh common in theoretical studies (Zhang et al., 2024; Kim & Suzuki, 2024; Chen et al., 2024a; Bu et al., 2024a) is unrealistic in real-world settings. A more practical training paradigm involves generative pretraining or QA-style pretraining/fine-tuinng (Allen-Zhu & Li, 2024; Zhang et al., 2025a). Notably, our concept data modeling naturally arises from generative pretraining (Park et al., 2024: Jiang et al., 2024b; Park et al., 2025). Furthermore. QA data has been shown to enhance transformers' factual retrieval capabilities (Allen-Zhu & Li, 2024; Zhang et al., 2025a). and has been effectively leveraged to extract task vectors (Merullo et al., 2024). Motivated by these observations, we incorporate QA data into our training framework, which we formalize as follows.
17

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Definition C.3. QA Sentence Distribution' (Pox). Building on the findings of Allen-Zhu & Li (2024); Nichani et al. (2025), which imply that incorporating high-level task messages or relation tokens within questions significantly enhances transformer's s ability to retrieve relevant factual knowledge, we model our QA-type sentence S,, == [x yn] E R2x(M+2) from distribution PoA as follows.
ks ~ Unif[K], yks.. yo ~ Unif[+1], Sn,1:M, Sx, ~ o(
Xs. (kg) Ukk E [K] (ks.) Bx B,r(K-)), =
yk, 1ni vh E Xs. \ (
Xn = E
Ta an,, + yeu bar)) + En,x
Kn
(15)
Yn = akg + Vn bks
ms.
Uni[[M], i1:M Unif[K'], Vn,I:M == Viaao
x = -M + En,1, ... Vn,ms, -1 + En,m,, -1, absn + En,msn'
Vn,ms, +1 + En,ms,+1 +1'" Mn,M + En,M, x] E RR
where K" denotes the number of task-irrelevant tokens, Vn,1:M denotes common tokens (e.g., task-irrelevant tokens) that appear uniformly across different tasks k E [K], and the position of ag can be arbitrary before x.
An illustrative example with M on concept i is:
What is the ag of to ta tn MA zn ay +cb, MM toby
y
The index (n. m). where n E [N] and m E [M], denotes the sample index (n-th) and the position index (m) for the noise vector En.m. Complexity of QA and ICL Tasks. In our setting involving retrieval over hierarchical concept graph, one natural way to measure task complexity is by the model's difficulty in confidently producing its argmax prediction. A sim- ple metric is CTT) = 1/ max,, Pe (y/ (T, where higher C(T) indicates greater complexity. QA tasks tend to be simpler: a cue word (e-g.. "capital" in "What is the capital of Japan?") guides 0 to retrieve the relevant task in col- laboration with the query word, often resulting in lower C(T). In contrast, Word-Labell tasks lack such explicit cues, requir- ing 0 to infer the task underlying the word-laber pair. For instance, given the prompt T = [Japan, Sakura, China], the model may assign non coiviaivial confidence to both "Panda" and "Peony" (e.g., struggling in choosing "National Symbol" or "National Flower" as the tasks). This reflects the polysemy-induced challenge posed by the hierarchical conceptual knowledge embedded in natural tokens.
QA-ICL Prompt Distribution. (A A natural extension of the above two distributions is that we can re- place the word-based demonstration by QA-based demonstration in the task-specific prompt, namely TQa := [x]*, OA . ,x YJ, x E (gdx(++1)(M+2) PT..
Success of QA-Trained Transformers. Unlike ICL-trained scenarios, we observe the success of QA-trained dynamics, as shown in Figure 3, where the column and row projections of the value matrix Wy primarily align with the task vector directions. As result, he.o closely approximates the task concept vector, leading an arbitrarily small test loss for both ICI. and QA-ICI. data after sufficient training epochs. Interestingly, depicted in Figure 4, training a transformer with QA-ICL data induces harmful hybrid task vector generation, yet its test error remains at similar level to that observed when training on ICL data. These findings motivate a rigorous theoretical analysis of the underlying optimization dynamics and the test convergence guarantees, which will be explored in detail in the following section.
ICL and QA Tasks. Given prompt T or an QA sentence S, where we denote the last token the residual query vector: by querry with some y [=1]. k [K]. We expect the transformer model can return the original low-level semantic label vector y by by argmax sampling from the undisturbed normalized token dictionary
01 + b
a1 b
Fa a1 bi Fa a1 b a1
b
b
U
a: + by
||a
by Il'
| a1 + bull'
|| a1 bill |a
|b 1l |
(16)
bx
bx
VK'
[u []] (d*(7*+
|xxl
[xx] |l
P1ll
|
'Indeed, the formula of the sentence is unnocessarily a QA, but can be a factual statement (c.g. = [x0A, y) = The, ans of, x, y)).
Me
18

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
The normalization of the noiseless token in the dictionary is ensure that every token ug, vk E [K] enjoys the same length fairly. Similar empirically driven dictionary approaches have also been adopted in recent theoretical work (Tian et al., 2023; 2024). For simplicity here we let ||ax]] ||a]| = |bull ||bl |m = 1. Given label vector y aa + y bb: + &. its
corresponding dictionary vector is defined as:
(2x
ky 1) -11) - 6 )1)))
Bky LN [
+ 2
7
where the dictionary index is determined by:
ky = 7k y/ + 1
2
Connection to Word-2-Vector Arithmetic. In the context above, when AT [1r,cl and the noise length are limited, approximately have:
YJ-1 akr + XJ+1
(17)
Therefore, if the transformer can extract the high-level task vector h (T), then adding it to any word vector within the same task concept should yield the task-specific label vector the context of argmax sampling. This aligns with empirical findings by Merullo al. (2024), which show that adding various embedded query words, such as "Poland" or "China," to the vector Jeny which captures the function get_capita (+) in the latent space- procuep the correct capital city the answer.
D. Preliminary Lemmas
D.1. Continuous Flows
In this section, we present several lemmas whose continuous flows serve as surrogate processes to bound the discrete training dynamics. While similar but simpler approaches have been explored in Tian et al. (2023); Meng al. (2024); Bu al. (2024a), our analysis considers significantly larger number of more complex flows. This extension is necessitated by the inherent complexity of our model, which incorporates softmax attention, layer-wise normalization, and residual stream structures-featuree that make our setting more realistic than prior work.
The key idea of using surrogate continuous flows is to leverage the monotonicity of differential equation derivatives. By comparing earlier derivatives with their integrals over a time interval, we establish upper and lower bounds for the discrete process via its continuous counterpart.
Lemma D.1. Suppose that a sequence a, to S t t1 follows the iterative formula
ar41 az
for some constant 6 > 0. Then holds that
= bt + ao-
for all S15 h-
Proof. The proof is direct by the monotonicity of the linear function and Comparison Theorem.
Lemma D.2. Suppose that a positive sequence by > 0,t, St t2 follows the iterative formula
ab ct + d'
b b
for some constant a > c > 0, d > 0. Then it holds that
b
+ d g
b 5 (
[t
(h
+
d)
c
for all 1 S1S 12
19

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Proof. Consider continuous-time sequence 3i 1 defined by the integral equation
12 Iv dr, St, = bt, 2t = Fe, + a {
(18)
cT + d
Here, is an increasing function over t, such that
dra
O.E
Ft,
bir-
(19)
dt
CT + d'
To solve this ODE system, by separation of variables first have
log((s)) = d log((t + c
which then leads to
log(a:) log((t d 2)²) constant.
To finalize the constant, extend 21, bi, into the formula, and have
bb (41 + 2)
(L + d c )
Tt
which is unique by the monotonicity of log function and the positiveness of a, c, d. Note that >> define f(t) (++g)
ctddd
To determine whether ft) is increasing, we compute its derivative (t) and analyze its sign. We use the quotient rule:
('(1) g'(t)(ct + d) -((()
(ot d)2
where g()) (t. + 2) a/c Thus,
g((1) -
+ d ī-1
Substituting g() and (t),, we obtain:
1 (t
d a (ct d) (ct 6 d)²
(t + 2) a/c
f'())
We can simplify the numerator:
f((1) (t+ 2) - --1
at + d) c (1+22)
(ct + d)2
Focus on the bracketed term:
(el + d) - c 2
=
Expanding A:
A -cl -d ct cd
Simplify:
A = (a - c)t + ad
d.
Therefore, given 1. (a e)t > 0, since
> 2. ad
0. as a > c ensures
@ >
20

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
[+1 azt dr, Thus, A 0, implying f'(1) 0, and S(t) is strictly increasing.
[+1
0. CT
21+1 24 +
dr
2 24
d
ct +
[++ (+ ct+d
where the inequality by examine the monotonicity of f(t) b1
By Comparison Theorem of two positive series,
obtain 6, Tt (1)+t)) [[ [(+ -
Lemma D.3. Suppose that a sequence e, 2 0,4 S $ follows the iterative formula
c++1 +
for some constant d > 0. Then it holds that
xt S 4 S
+ St.
Ct1
for all t1 <t 12. Here, 21 = ddlt 11) + ci,,
Proof. The proof strategy follows Lemma C. in Meng et al. (2024) despite the consideration of different ODE processes. The key the decaying and positive nature of g(x) 1/2 as well as the examination of integration. Consider a continuous- time sequence Ht, > defined by the integral equation
F2 dl 22 dr Tt, Ct2 / d1 IT
(20)
Note that x, is an increasing function of t, satisfying
dr; d It, Cta
(21)
dt 21
By solving this ODE system, we have
d(t t1) + c,,
Ft
which is unique. We first show the lower bound of ct by (20) we have
++1 dr
J++1 vz + d
dr
d
r. d
It
Ft
where the inequality is by the increasing nature of and the decreasing nature of g(z) = 1/x. Then by Comparison Theorem we have Te I az. On the other side, we have
a = ct + Σ d
1=11 c
1
Sc + Σ
Ir
1=t1
+ Σ
d
ct + C1
Fr
-=
dr
S ct +
+
On
b1 Fr
d dx, c2
d
ct
+
+
21.
Cn
Ct.
d
+ Zt.
Ct
21

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Here, the first inequality is by o Iy; the second inequality is by the definition of integration well the decreasing nature of g(x) = 1/2, the third equality is by Eq.(21). The proof is completed.
Lemma D.4. Suppose that a sequence Et, eo 0,0 s La follows the iterative formula
61+1 2 et bec.
for some constant b O. Then holds that for all 0 S S t3.
6x 2 -beos + en
Proof. It is obvious that e, is decreasing, and since co O, the decreasing speed during t S t3 would not exceed -beo- By the monotonicity of the linear function, the result follows directly from the Comparison Theorem.
Lemma D.5. Suppose that sequence fir fts 0, < S t, follows the iterative formula
se+1 alt- tz)
for some constant a >> 0. Then it holds that
(1ft-1-fot 1z)² 1] Stg + afts [tt - (ta
< S Stg exp (al2 /2)
for all S t4-
Proof. Obviously ft is increasing sequence. First D.3, by the definition of integral, the monotonicity of g(x) exp( x)) and Comparison Theorem, the continuous flow I. satisfying
+ 13 by Substitution Method. Similar Lemma
dr. a. To fta S 14 - 13
ds
would be the upper bound fa. By solving the ODE and consider the monotonicity of g(x) exp(x), the unique upper bound result follows.
For the lower bound, observing that
-1 +
Σ a(T - 1a)
r=ts
2 +
Σ
a(T - 1z)S
(22)
rmts
d-1
251h +
alT - -tz)sesdr
1/11
=fi + afta [tt - 2 1]
Here, the first inequality is by the increasing nature of ft; the second inequality is by the definition of integral and the increasing nature of g(x) = az; the last equality by the integration of linear function. The proof is complete.
Lemma D.6. Suppose that a positive sequence > 0,4 < L S 12 follows the iterative formula
agt
g++1 = gt + (bl + cc
for some constant a, b, c and satisfies 2c. Then it holds that
9t 2 It. e Vostor Rogio
for all 1 SLSt. 12
22

---

dr, Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Proof. Consider continuous-time sequence 3i > 4 defined by the integral equation
ax, (br + c)²
2t 24,
dr, It, = 941
(23)
Here, is increasing function over t, such that
drt dt
ast (bl + c)2²
2t, = gt.-
(24)
Then, by separating the variables have
dlog((,,) = -0
d
bt +c
Therefore, it holds that
log(z1) = b(bl + + constant.
a
To settle the unique constant, we extend Tt, = St, into the formula, and have
9t, b(8t 13)
which is unique by the monotonicity of exponential function and the positiveness of a, c, d. Note that a < 2c, define (t) e ककोी, the quotient rule and have:
(bs+c)² To determine whether f(t) is increasing, we compute its derivative f'(t) and analyze its sign. We use
(bt + c) (bt + c)2 < 0.
As such, it holds that
1+1
ptl ax
ax,
Tt+1 It
+
dr S x4
(br + c)²
(bt + c)2
महोro (66+e]a
where the inequality is by examine the declining monotonicity of f(t) positive series, St 2 It, baftee +o0 immediately holds.
By Comparison Theorem of two
D.2. Concentration Inequaities
Lemma D.7. Suppose that 8 and d (2log((2K + K))N2M²//)). Then with probability least 1 8. for VE, E/ E ( (aam))neNN),an (e) E [M]. E (ax).e,k] (bs).E U [UN]NERR
ofd 2
< |E2l < ofd
2
7(NM)2
(E:, Ex>I s 20, V dlog
3(2K + K'NNM )
(E: u)سl < ||u|120p V 2log(
5
Proof. See Lemma B.2 in Cao et al. (2022) for proof.
Lemma D.8. Suppose that 8 and d 2(log (2K + K2)NN2M2/8)
Then with probability least 8, for
23

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
vu, U. (as) U bbs) U Uew Jeekeey E E (In.m)nE[N,,mE[MM
| 10 11 S od
8(2K + K2
w V u/ 1 V 2lo
)o | |0/1
5
0.990 21|u|| 3 - u/1² < 1.01o 2
8(2K + K?}
u| < <4\
2 log(
6
W (1) u S 44 K |
I0 W
V & S 2/22lg(16(2K 2
15/ W &/1 & (/2 + KM
01 2 Op dd I (1) 4/12 & 00
W (0) E:1l - (/)
[6/ W (0)
0 E:ll 2l.2log((662KK 1
K')///00
I&
W W/ v E/ - 16 Tog/ (16(K + k')3/8)07002002
Proof. The proof strategies follow Lemma B.3 in Cao et al. (2022).
We analyze the Frobenius norm squared of ( where w (0) w is d x d matrix with entries independently sampled from of). The Frobenius norm squared is:
W 00 V
12
Σ W
(0) Each W2,, is an independent chi-squared random variable with mean of and variance 20f. The sum W [ is thus of d² independent chi-squared variables. By Bernstein-type bounds for sub-exponential random variables, when d = Ω (Vlog(1/0)) we have:
( S 2d
with probability at least 6/8. For Yu, E {a.h.ejk] U (babee U (pw Jeepy it holds that T T o1 | || Then by Gaussian tail bound and union bound, we have
u
a T (
u - 2 log
8(2K + K'
)a 1/11/1111011 < v 2 log
8(2K + K2)²
)o ||||||||||||
8
8
with probability at least 1 - -5/8. Besides, for Vu E {anteejk U Ubb.)EEkK, we notice that w (0) (0) T w wo!))
u ~ N of /
Then it holds that u Bernstein tail bounds (for appropriately large and union bound, we have
σ1 2| 12 where x2 is a chi-square distribution with d degrees of freedom. By
(0) σ1 | < 0(2o 2 /dlog(16(2K + K')/8)
with probability 1 1-88 8/8. By appropriately configuring d = 2[logg(2 K /6),, we have that with probability at least 5/8 we have
0.99I1 2 |||||| <
W
(0) v
"/|" 2 < 1.01o || u|/ 2d.
Combining the obtained results we obtain the following bound
8(2K + K')²
1 w w|ll < 2 2
log(
||0|//Va /Vd
8
24

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
On the other hand, notice that for Vu # {as)ae(k) U (bs)se(n) mean O. Given that
W (0)
W (0)
is sub-exponential variable with
Cor(W V v, w u) E
] O.
thus
Var(u W WV v) = B
b) E ( w) ( u) | ( v. v); ( / m)
Σ
Tr ww = of | || 1|
Therefore, for Vu + w E {as, bokke by Bernstein inequality and union bound, with an appropriately large we have
e T W W u 0| s 012/ + K2)2K + K" - 1)/8)of |||[b]))
I 4Vdlog(16(2K + K1)3//off ||/l|//
with probability at least - 8/8. Finally, given that o/ S |1i12 - 30/d/2 in Lemma D.7, the rest can be proved using the same strategy.
By the union bound the proof is completed.
Lemma D.9. Suppose that > and = Q(log(2K + M²))N2 Then with probability at least - 8, for
VX E (Q, K), u, u E (as) {bx).c(K] U (WR)
1 x 11 & 2df
8(2K + K')
10 WTW (0) 2 S / 2 log(
)oo [o|||||||u|l
uull 2 d/2 - W (0)
x u u1 S 300 2 || d//2,
I0 (0)
w -x u < 1olg( K3)2/500
(0)
||]
l W
- 4/ - 4 - K')3/8003 1/2 /22//)/) |||||||||||
I0 (( (0)
&i 2 2log K
+
σo 00'op Op d < [W]
(0) x & - 300 20 dd
I T O [ El - 16 log 1og( + K1)3/8)07
Proof. The proof strategies follow Lemma D.8.
Lemma D.10. Suppose that 8 > O, and N S2 K (1))) K (2log(1/8)) Let N denote the index set of sampled QA data points P (or ICpe data points P where the high-level task concept is lig k E [K] and the task specsiiiiiii low-level semantic real-valued label is y)g = = [+1]. Then with probability least 1 - 5, we have
10 2)N
11100-2)0
<
W <
(25)
2K
2K
for Ak E [K], E [+ This would further suggest
(1 - 2. 10 3)
SINIS (1 + 2- 2-2210 N)N
(26)
K
K
where [V6 W n Widl represents the index set of sampled QA data Pi with their high-level task concepts as kg = k E [K]. Moreover, for Vy {11), there exists ks E [K], such that
IN N/(22) 5- 10 1N(K))
(27)
25

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Proof. The proof strategy follows Lemma G.3 in Bu et al. (2024b). See | binomial random variable with proba- bility (2K)-1 and number of experiments N:
W
(in (2K)-1).
The expectation and standard deviation are given by:
N
1 2K
2K
ok =
2K
By Exercise 2.9.(a) and (b) in Wainwright (2019) and (2K)-1 1/2, with probability at least 8, we can directly have W (2(2)) 0.001 by Hoeffding Inequality well appropriately large N.
On the other hand, define the event:
A 11N N/(2K) < 10 (N(KK)
Observing that the probability of As, & E [K] occurring is lower bounded by absolute constant denoted as Therefore, considering K trials and with the condition K = (2(lo((1/)) by the linearity (independence of different tasks), we have:
P
= =1 -
n
Ak
k=1
k=1
1 K 2
Thus, with probability at least there exists some E [K] such that:
W - N/(22) 10 10
Lemma D.11. Suppose that 8 > O, and N = S(KK" log(1/6)/M) For vk E [K], E [K"), denote VN as the number of common token Vy appearing in sample set N& C PiA (or PT)). Then with probability at least 1 8, we have
0.99N M
1.01NM
| S
(28)
KK
KK'
Proof. It is worth noting that [ can be viewed as binomial random variable with probability (K')-1 and the number of trials [ Using the bounds 0.998N [V < 1.002N from Lemma D.10 and applying the same proof strategy, the result follows.
K
E. Detailed Gradients
For clarity, in this section we omit the notation of sample index and iteration index (t).
Lemma E.1. The gradients of W K. WQ: Wy can be computed as follow.
M
Epp, 1ww Ls(( -Ep
Lukky
(WeSen) (S(e) - -T) ]
==1
Epg. (.WLL((0) -Ep
buy WR.SSE, T) (Ser.)
(29)
1=1
uby - Tx Wile (St)
II
Sm)
E Lc(0)) -Epg
Sy
26

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
where ,,, T RR, (wi,... , T Additionally, for forvj [L- 1], k E [7K + K") as well any vector y, u, IR,, the coefficients Tt wk; buyy . and operator IT u defined
exp(e W Wo
Ty == Noltmax(WWSSeen (Waaee)""WWWee (WQSeL)) M exp(u (x) + in LN j=1 My W Se,))
Ej exp(e, W WxSe;))
Wk
EKElTK+RY exp(u o (x) + i. LVQ j=1 Ty WvSe;))
(30)
{luky
EKE[TK wa::) ) (W (WVST,ee)) (WySS)((WyS(T (WyS(T-TT) -Te,)))
Luty. J
WyST
u
17 x
| 2
Proof. For W E (Wx, We. Wr), we have
Ls(())) - -.wEpy [ur ho log( Σ
exp(u ho))
kE(TKKKK
=-wwppy [ur, ((x) + LN
Tt Wy
loog(
exp(ux ((x) + LN
Ty Wy Se, ))))
(31)
kE[TK+K'
=-Epg, (Ow u, _, LN
Ty Wy,,,)
M exp(u ((x) + LN
- dw log(
J, W Wv
kε[TK*K]
where
exp(e] S W W Se,),
T = =We (WoSer)) = Eie(L-I) exple, S W W Se.)
(32)
We compute uw LN T. Sy,)e) as follows. By the fact that the trace a scalar is the scalar itself, we directly have
M du LN
M W,Se)) = tr(d(u, LN WvSe,)))
Σ o,
1=
j=1
27

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Then we
M d(up LN(
M
Jj Wv = tr(Owx (upy LN( Ty W. Wesee))dWa
f=1
==1
M
+ ir(Owa T (4y LN(
Tj W.WySo We)))WWW
j=1
M
+1r((Oww T (4 LN
T (WSeeee)))WW)
j=1
(33)
M
= =ir((Ow N (u LN
Tj WvSe,))) dWx) (dWk))
M
+ +Ir(Owa (uky LN(
T Wr Se
==1
M
+ +1r((Ow, (uly LN
F Wv-ee))) (dWr) (dWr))
As such, we can compute
M tr(d(u), LN
M
W WvSe))) tr d(LN( Σ
TÚ Sv;)))
1=1
1=1
(34)
Σ 1= Tt W Se,
M
F.. W We
=tr(ul, d(
1)
Σ
tr(u,
M d(
|> T. W. We,e,
Σ n. W WySe,
1=1
28

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Then we have
M
3 Wy y Sej
3 W v Se,
tr(ukky d(
tr U dl
L-1 1i=1
W Sel
my Wr Dies T. Wv
j=1
M
d(TT) Wv)Se, Σ
T. WrSeill
tr(uky
|
T Se;
==1
M
WySe,d(1 Σ T. WvSe;
Σ
tr(up,
1E T. Sv.e
/=1
(u W WrSe, + T. dl (Wrr)e,
Σ
ir(
|E T, WvSell
M
TTy WySej Eieeie-
ST dd के Wr)Sei)
Σ 1=1
tr(
(35)
2/1 2/ i=1 T. Se, 3
1--1
M
( T (e) + TJU d(Wr See)
Σ
ir(
1E T, We.lll
1=1
M
2J
Ty By Wvse, Te W W Se:)
d(Ti))
1Σ We
J Exte(l-1 S,,u, W WySe,mmle ST (d(W)))v W]d(Wv)))
tr(
2
-ー= T Wy
M
M > 1=1
(tr(Sine, (dx,)) + -ttrrr dWv))
=
where
E EEELL- Teee WiWW)S
<dn,
= (ub, Wp WySee)(
DEL- Ti Wv Seill EitgA (W (W,ST,,, Wy
1E n. Wy
= (u WvSe;)
|E T Se; Σt,e((211 Seiu by We S WT
n. Se
2 -i n. W., 1 n. We.eell
ri
Bky
T Exte(h--1)
WySeeu Wy TTTT
(36)
2 1=1 T
W. 1
ΣΣ 1=1
n W
Se,ll
TT. E1,E((L-1]
न (W,Se,e) S WTye TT
Ty(ye]ST)
2 LL-1 n. We
1 म, W Se,lll
Jj ELtELL- K J W WeSeee.TT) TT
by
2 x- n: W Se: 1
Σ D m W. Se
29

---

M Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Then for Vj e [L 1]. we investigate d(Ty):
d(TT)) d exp(e W
Eiel exp We Se,
T d( W WKe
Ty EieLL-1((exp(et W W (e JSe
(37)
W We,))
(d - Σ (olWWWR)SSe
EE|--1
By E3.(3 (35),337) we have
LN( d
Σ
Ty
Wv,)))
(r Σ
Suby
Tzer ST
Σ (
1=1
j=1
1E[L--1
+
r Wky
dWv)
(38)
M
tr
Suyy S)S(e)
Σ Tejee ( (d(W3Wk))
M
+ tr
dWv)
uky
/=1
As that
d(w wx) = d(w Wx + W)d(Wk)
Therefore, by the Cyclic Invariance Property of operator trace and (31), we have
M
h-1
(Owx (up LN Ty Wy WvSe;))) = Σ
Suy S,
Σ
mejees W
1=1
==1
That is,
M
M
4-1 Suy W WoSer(ez
awx (u LN
W WvSe)) = Σ
Σ
me; TST
(39)
j=1
Similarly, we have
M Ow (ul, LN
1-11 M T. WvSe)) = Σ S x, WxSSee - Σ me ee
M
L
(40)
j=1
By (33, 35), we can also obtain
M Ty WvSe)) Σ
Owr (u LN
( ruyy )T
(41)
Similarly, to compure dw log( Eke exp(u ((x) +N M
Tú W Se,ee ))), investigate
logg( Σ
exp(ui (x) + of LN
ny W, We
kE(TK+K
30

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
which can be calculated
exp(up (x) / LN(E TTy (dd, T, Se,))
EME[TK+ exp(u (x) - u LN(_,=1 "M WvSe,))
(42)
M
Σ KE
wid(u (I(JULLNV LN(
Σ ==1
Ty TyWySe;).
where
(xp(u]((x + LN
M Ty Wv WwJ)))
Wk =
Σ
Wk = 1.
(43)
Σ E (x) + un LN( 3, WvSe))) KE[TK+K]
On the other hand, by the derivative results we got from d(u,, |N(E M my Wy;)))), directly have similar outcome for Uk, & E [7K + K:
M
M
1-1
Owx (ug LNO
Σ
Ty WySe;)) =
Σ
Suz W WoSen(es -
Σ
mTee)
T
J=1
Ø=1
1=1
M
M
L-1
aw (1 LN(()
T W, SyJ)))))=
Σ
Cú Tj W WKS((e)
Σ
me
/=1
j=1
i=1
Ow (1 LN()
Ty WT WW,Se))
Σ
(m) ( TT
1=1
j=1
(44)
Eile(L-) 1] T. Wv S W.ueeTTT
M Tyu.CSST.
Σ n. Wy
n, W Seil
1=1
Jj ExelL 11 7. W S. W ST
2|| EE
m. W.SSe 2
1E Se,
As such, by (q.(31, 43) we have
M Epg 10w Ls(0)] = -EEp [(1 - Why
L-1
Suzy n, W WoSer
Σ
T TS
f=1
i=1
Σ
cd. n W WoSei(ej Σ
me; T
Wk
k
j=1
Rpy, 1ow [s((]] -Rpp [(1 wky)
Suky n. WKS(es
(45)
f=1
Σ T
cl n. W Sk
Wk
J=1
=
- Shu W
Σ
GA
31

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Define [1. " LL-1, 0J RL and [- wakk
T [0sk] By the definitions in Eq.(36), have
Epg, (dw Lx()] = Epy, (WQSeL)) ae
M (((1 wky Jky ENjRy wx (W (WrSzzej))
(S( e,))
(WyST) (W (WTS(TT S(T TTE,))
WyST
M
Epg
E
(WoSer.) (S(+= e;))
Qu
j=l
M
(46)
Epy, 1 Lo() Epg
WxS(T= - -e))
j=1
why Juy
EKSRy wRUN) (WVST,ee)) (Ser. 1
(WVST) (WyS WVST - eq -ml)e
WW 1
M
E
Juyy WRS -e)) (Se
where the busy .f is defined as
{((1 - Wky JRR
Exyin ( ((WV:)(W))))) -z,e))) -wes(
(47)
Ww.ST| 1
Similarly,
M
M Wk k
Epy, 1ow, Ls(0)
-Epy [(1 why
Σ
E(rl)) )
j=1
[f -u:ee,s
WE Σ Er.
T. Wy We,eW "s De ST
ST
-BBp,,
kikk
Wr
wATy Emieie- m. W, Sc, (y un) W. " Wy
2 W STll
WyST
(48)
(((1 (((1-way Jucy - Exdky w (W.
((1 -Wy Juk Excpy (ST)
[
Epg.
(1(1 ([(1-wby)) Juny Exnr, w (WrS+))(W,Sn))
((1 Wky Juky - Exatkt w
ST
(ST)
WW
(1-wky)usg Juy Lipkyy (ST)
SW)
Egg. g
[W
where the projection operator II is defined as
=
v - v
TT(
|u]||
u
|u]]
u
32

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
F. Training Dynamics: QA data
In the following sections, we consider training is on QA data, and assume results in Appendix D.2 all hold with high probability. Following Definition C.3, we denote:
Sn Pi the n-th QA sentence sample in P QA
Yn the label vector (last column) of Sn.
ks E [K] as the co-task concept of S,,,
ykk,, Yn the task-specific low-level semantic real value label,
lyy = 2ks + #1+ as the dictionary index of yn.
2
ms, as the position index of aksn in sentence Sa.
mn,m as the attention score over the m-th position in S,, at the iteration 4
(t) w n as the weight of dictionary vector ug in S,, defined in Eq. (30) at iteration 4
buy mn as the coefficients defined in Eq. (30) at iteration t,
No. PA the index set of QA samples with their high-level task concept ky A: E [K] and task-specific low-level semantic real value label Mks y E [=1].
VV the number of common token My appearing in sample set No c PQA
F.. Phase 1: Linear Growth of MLP and Accelerating Growth of Attention
In this phase, the MLP is upper and lower bounded by linear continuous counterparts. Meanwhile, the evolution of attention upper and lower bounded by exponential and quadratic counterparts, respectively.
Lemma F.l. Under Condition 3.1, during S Ti ((ngv) (ofdk)) we have
w v au = ola, 122),
ag
8(2K + KT)
|a w u/, | TW
1 ( 0'1 = O(V 2 2log(
)o1),
5 = =Olofdene ofd),
(49)
1w aks.
K (kks))
16(2K + K')2
(w((xx)) x w!! (U, mss En.-ms,)) 1 O(V log(
)ozd
for Vk [K]. E [N], E [7],-m) ms # ms. E [M]. u E (asl U
(bb,) (bb)) U Uww hrejk U UnnnmbnnnNN
and u # o' E (b.).e[K] U (Em.m)ne(N At the end of this stage, we have
a ( ax 2 d/
8(2K
K')2
+
V 2lo
8 )01
M
(50)
(W (T) a 202 of
(W! az) ag)
Cs(
M2qv
M2
for some positive constants C1-3-
We first examine the initialization. At = 0, the results in (q.) be directly derived based on the orders of the initialized products and the norms presented in Lemma D.8, Lemma D.9, well as the small initialization do = 0dd-1/2)
low noise op = O[[]u]|d-11/) and overparameterization condition (2M² log( K²2 N² )) in Condition 3.1:
(W
(0) Q xn.)
( (aks. + En.me, }1, I(W Q x..)) TW (Un--mg, ms, + En,-ma. ) S 12 /2 + K')
8(2K + K')
8(2K + K'2
ag (0) ax S 2log(
)o1 < o(o)d
Iu w u/. la w - 2og
00
33

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
for vk [K] n [M E [7], -s, + ms.. E [M] (K))) WEEKK U (w..))/ INENNNN By direct calculations, it holds that
(0)
(
# n, 1
M
M
0.980 1/2 - I
Sn T R
(0)
S 1.020rd'/2
W (0) Sm
(0)
i m (0)
(0) (0)
Sn V S,
8(2K + K)
kyn
S 8V 8 l
1/vd=o(0.01)
(0)
I S n
(0)
5
(kky..
EKc[TK+ W u
S, n
161 2log( 8(2K K2)²
)/ /vd==
W T (0)
( (0) V S. TT
(0) (0)
V S Tn. mm
(0)
1.05
(M - 1
()(44/lgg log [6((2 + /2)//0/d-1/2
1.11
I V S Ta 1
I S. (0)
0.95M
0.95.
M
Since these scales are crucial to our analyses of gradients, we first introduce the following results based on Eq. (49) and (50) at iteration / T.. Lemma F2. Under Condition 3.1, suppose Eq. (49) and (50) hold at any iteration t S Tí. then
(Tt) 3 An,
S 0
+ (M- - 1)l-))egeee o ) < 0.95, an, ms, 2 0
(TT)
1+ (M- 1)e=odderr0 o7a
| S |1 00 13)
w! S (4)
1
0
1 (M I)e ogganv ofd
N (1)
v ST (t)
8(2K + Ki)2
(51)
Sn,m
1S 15((8/ 2
)d-1/2)
W S 1'11
IW (W (())
(W (6)
(Wlyn
ERETK+RY Un.
00
|W V S T.
1+ (M- 1)e a3do"d ofu
(W ST m
( W Tn Cm o(1) < 0.993 < 1,
W S (0)
WW V S
(9)
for vk E [K], *k + k E [K],
N C [N], E (7), ms, # ms, E [M]
(Tr) Proof. Suggest (49) and (50) holds at f, then Tn,ms, S O 11e oz ofa); O - 3a. ofa
1+(M-1)e o
Tn, by qv Q log(a, -d- logt N
Furthermore, as we see that Iu w Wiul ul, la ((!)) u is feeble compared the growths of ag ( ag suggested in E4.(49) and (5,),,), thus the ag ( Sn mn,,s, ems, would be primarily responsible to the changes of [ S, ("(1) I. which is the change of the term
(Ti)
add
(e)
ja (ak cos (w()) V ag ax) ( II) T (w ax) (m (a) W V ax)
n,mm
where cos (W()) (1) v ak ax) a w ag Note that if a. w v
(1)) 0 at t = O. a v a will increase and thus
IW azll
as
(a w-()) )2 would decrease, which might lead the decrease of | remain at d(a during (Tm. )²_ (a w a2)² < O(lw V S 111) + (Th))s.)) (a w V as) - O(a,d1/)). This apparently holds during t < Ti by (49). Therefore, we can safely conclude that 0.95g, 1/2 | S. !0) 1/2.
(1) ST (2)) The scale of | S, ("12 would
34

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Then, based the scales in Eq.49 well Condition 3.1 3.imiiryyyyy can directly compute the last three inequality by the definition of un (uyya Lakss
yks, bks.. )):
7k+s W V
S (t)
o(oid
S O
S n 1
1+ (M 1)e oode 1)e-odde (01d1/2)
0
1+ (M -)e-oodee ofd
U7-k W
V S [W (c) V S (t)
5m.m w STTT (t)
8(2K + +2)
< 888 2 2
)d )dd11/
V Sn TT
5
[Uyy
EkElTK - K Ukk)
v S ())
S @ o(oid!/2)
1
IW S S (1)
(2)1d)12)
1+ (M - 1)e- ofdeni o7
=0(
1+ 1+ (M 1)e ogde o'o
(W (W S.m ( w Sn Tn. Cm
(
5 0
w 1 S T.
(+)
ST (9)
(1 + (M- - 1)e agdelv of
(M - 1)(44 /1g((6(22 + k/33/8/d-112 1/2)
+ 0
0.95M (1 + (M 1)e 1)e-odde= oP2)
s 0
(1 + (M - -e-oadee ofdg3 ) 1.
The remaining is prove (49) during & TT. To facilitate the proofs, we introduce the following auxiliary lemmas.
Lemma F3. Under Condition 3.1, for the whole iteration I S T = 2i1 " of'o,KMe log( we have
wi.k - (-)) o(1),
(52)
for Vn E [N], E [K].
Proof. By the definition of wakk see that for the whole iteration / T there exists some constants Cf C22 0,
(t)
exp(u (xn) + ug LN(E
m= Tn.m W 4
Sem)) ))
(0)
V
w
Σ Eke[TK+K]] exp(up
(xm) + u, LN(E j=1 nn W (0) Sem))
M
< 0 (7K + K') exp (-1) ) 1 0
exp (1.5)
7K + K" = 0
13
C4
= o(1),
(53)
K"
10 wnn 2 (7K + K') e1.5
7K + Ki) = 0
0.08
C²
K" = o(1)
Here, the inequality is by the definitions of un and LN well as the large number of common token condition K" = ((K)
in Condition 1.1. This would suggest that (t) = O + always hold during the whole iteration S T
Based on this results, now we examine the evolving formulas of u T ( Vu, v E (astve(k) U (balve(k) (Ne]eejk]) &,. E/ E (MM in the period < Ti based on Eq. (49).
35

---

Provable In-Contes Vector Arithmetic via Retrieving Task Concepts
Lemma F.4. Under Condition 3.1, suppose Eq. (49) holds at iteration S Tú, then og W as increasing such that
79v (M (t)
w(++1) 0 V ax a ( og - 0( Σ
2 S 19)
BEN;
(54)
a w u a w u/. 100 (((+1) 2 v T 1%.
a (w(++++)) e La Tw 0'1 (.0 w(
as a w az).
for vk E [K], E [K].uu E {a.)s U (bslec(K] U (5.m InMM and v * o' (bstae(k) U (SMM)M)
Proof. To examine the last inequality, we first examine the update of W V. By Lemma E. denote.Ng C Pox as the index set of QA samples with their high-level task concept ks = k E [K] and task-specific label real value yks = E [=]] for
vi E [K] W + ke n E NR, we have the following nethoort relationship
a (Emmmm, ((An (En', + vn) + En' ms, + txx) (T) ag < <0(-
3(2K + KK'NM,
- O(op V 2
M
5
= =
(Uky.'
ΣLeTTH +e Wen (wlk)))) (W
S T (t)
w S.. n (1)
w V SR ,
(0)
aa (URya
5 10
| V Sw (1 (t)
||
|| S.. |
(
| v S. /|
= o((00001
(55)
Here,
the first
result
is
by
ag
(Fqq
e (
(1) Tn.ms,
1+(M- 146 - - ofa O oa ov
O
((+)
well as the low noise condition
14 -fa fa
mg.
op < d dd1122/0 /C in Condition 3.1; the second result is due to (51). The first (0.001) can actually be smaller if require larger C' in Condition 3.1, but for the simplicity of presentation, here choose feasible one. We latter would show that (Tt))) // )
(6) T ag would maintain the scale during the whole iteration.
36

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Then we have
Miys ENg +4 W (ST (0)) Og. WWy 2 )
W k V
ag= - q
W az = 199 Bpg
IW S nn
ENeTTA Hel) Un (STn (1)) T
ndv E
Σ
Σ
(WvS., m )))) )
N
I S (
RE[K] vE|=
Miya -UERT KI wakk us
Tag. (S (1i) ag
= wV
Σ
Σ
e([
(WySa 12
N
W 11
vE AEN
(UkYn
SAppR4KyWA.eu (W SAm
w S (1) WW
(1kya
79V
N
SATT (1) [I((Thmmm a + Em Em (m..m))))) Tag)
v REV?
(1/2)(M R.me, n.m, + toMM (MT (t)
a ma )Vllog(3(2K + (
7IV N
M|[W S (2)
NEN
7av FM ms, =0.001)
= (
2N A
BEN,
(56)
Here, the first and second equalities are by definition; the third equality is by (5.(; the forth equality is by definition of IT in (30); the fifth equality is by the definition of uy (uky, = LN(aks., + yks,, bks, }), Lemma D.8 as well as
Eq1( the last equality is by the low noise condition in Condition 3.1. Note that r7 S olgEtoRdt (2 K /log( ))
Similarly, we have
1 II Why ENCTTA ky) AR (Wips ( 1
) aj
bb W (++1) V ag b W/ agl
nav
Epr
(
10
May, Exe(r & (ST ) ag. II
MV
Σ
Σ
(WVS.
N
S
kE(K DE[+
b Waya ExETK Kn) Wa (ST ) ag
D0V
Σ
Σ
O(l
Sa
N
SA 1
[=[=IREEN?
(ukya -E us) (W Sam
W
b/6 (Usy
ngv
Σ
e(Il
| (0
| s.m
N
[[ SA (8) (0)1((=(()) (0)
A,m,, ag + Σm (T En.m)) azz))
vE[=IJ RENT
ngv (0.001 + 0.01))/2 - (1 (0.001 + 0.01)) (t)
n.ms, 0.001)/2
< O Σ
| S 12'
neN;
1av (0.01) (0) T R.MS
e([ Σ
ag w((+1) ag a w ag.
2N
SATE
REN
(57)
Here, the deduction of the first- to-t-to-forthh equality follows (56( the fifth inequality is by Lemma D.10, the definition of Us (Myya LN(aks, + yks. bkg, )), (q1( as well as overparameterization condition 2(M log(K" N2 (PPM228)
37

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
The value 0.001 can be further reduced by increasing C in Condition 3.1. However, for simplicity, we opt for illustrative choice here Similarly, for vk K E [K") E {α*)s U (bs)sE[K] U (Ea, ME(() and + o' E {bbS)OEKK] mee we would have the following
a (((-1)
0 W u S 0
nyv nyv (0.001)
M W SATT
BEN
nyv wwv(0..))
u T W V
a 2 T w a
m. 2N W ( S 1
nEN,
(58)
ngv (0.00))))))00)
w w 1 W
1
2N
T (
(0000000
AEN
n.ms, 0j + Em (T(") En,m)) (aam))()a TT)) that would appear if is on the right Here, (0.001) is by the term (S u (
side, Lemma D.8 and the low noise condition op = Odd-1/2) in Condition 3.1; (0.01) is by the third and forth inner products in our (51) as well as overparameterization condition d = R(M2 log log(K'2 @<2N*MM* N²M² Lemma D.8 and the low noise condition op = O(d 1/2) in Condition 3.1, as well as the balanced property of the QA data deduced Lemma D.10.
Besides, it holds that for vk [K] kf E [K'], we have
(ukya (ukya
Exetr+ an)) (W (W("', s,,
w
| W/TW(C-1 (1- ax 1/ W axl
74V
Σ
E(1
W (
N
/ Sam (2) [[(TT)
ag Em (1 (1)
m Ea,m)) aa))-
1E1+1] BEMY
A/lL (0.01)
0.001)
2NM
SA
REN;
I w((+1) by W. bul < 0
Σ
Aav 7 (0.01)
2N (W SaT
REN;
(59)
where the left (0.01) by (Eq.(51)) as well as the overparameterization condition 22M log( 22 N2 in Condition 3.1. The right term, (M (M + 0.001) or (0.001), follows from the calculations of (ST (1)) T (1)
((T A.ms ag + Em ( TT u), Tu), u E fax, by bk]ee( based on e e S Tn.ma. S E
(t)
1+ 1+ ofa ).
O e 2 1 1nf
Th, mng. < 0 +++) in (4.),,) as well as the low noise condition o, = 0(d-1/2) in Con-
1+M-e -f
dition 3.1. addition, denotes ag Vn.m: by Lemma D.11 and K" S2(M) in Condition 3.1, as well as the strategy
38

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
above, the following holds
(uby -E EkeLTK wa.)n))
w S 10)
ell 0 (ukya
lag w UE a& W
nI N
I| $A
I 2
I| V
SR
(5) (1(((2
( (T Mnm m) TTp))
(1)
DE[ RENT
nAV (M T A.maary
+ +00001)
1av M (1)
(0.01)
S Σ
1) s 0
2NM 7 10 1).
2NM 7 1
REVerivig
BEN;
7Av (0.01) (0)
+ 0.001)
16% ((+1)) VK b% W uwl S
2 S 1
RENErINU
nav MT (0)
Bammar (0.01)²
1),
W $ 1
AEN
(60)
where New C PüA denote the index set of QA samples with Vs in the sentence, and my denotes the position index of Vs when E Nage
1+(M-1)e ofaa* o07 ) O -6 of
By (56),,) (57), (58), (59), (60), ) S nn.ms.
(e)
S O
1+(M-1)e aaacn of
S
(2) Tm. ( xy that for W u, u. u {ak) U [bs)) U (ereee U (Enm (omm)ne|NN Jne N [
only when = az, vk E [K] the products would grow in a non-negligible manner. Therefore, by definition of ms, and low
noise condition on = Od2-)/2), the primary contributor to the change of (W("s, Smm (2) is the increase in Tn,ms, and ag w ag
(t)
The following lemma examines two scenarios for the upper and lower bounds of the growth evolution of a w
V az Lemma F.5. Under Condition 3.1, suppose Eg. (49) holds at iteration I Ti I 1/ where 12
(nav)- (g (dd//2 V log( 8(2K+KPP) Then for wk
[K] and
No. there exists some
C1-2, --22-1--22 such that
Cinget
V 2 8(2K KT)2
C2nqyt
8(2K + K12
Jon S ag wil) as oTd/22K + + loog(
)o1,
OK
5
8 O2y log(
(61)
GGnggt
CIV2
log(
ag ((()) S nt Tn, Eman
(t)
C2Cenqut ofdK
8(2K +K)2
oddMRK
Mdl/2
W (1) S ग. (
0,d1/2
Proof. By Lemma D.8 and Condition 3.1. At t. t= 0, by Lemma 1.10 and (51) we have
ag a a w(( as 2 2
(62)
2MK(2,
Therefore, by Lemma D.1 and Lemma D.8, for some C₁ we have
a W(O (t) as 2 C
a/dl/2MK -1 2
nyyt
8(2K + K")²
)01.
8
for 0 t < Ti. Define the positive estimate of the time point when the lower bound evolution of a w( ag has
potential to reach (22 as t == (πqv) " / 22log( ((2) Note that t < t thus we
39

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
have Ti S li.
Collaborating with the non-decreasing nature of ag [ ag, the non-increasing nature of g(x) = 1/2,
14 -aa =fa ) O e e-08de"v" offa
( S an,ma, I (
1+ - tofa S -n;-ma msn I E (M)) for some C > 0, we
(1)
directly have
0k W S. ma, Cman.
(1)
(1)
2 C1 n,m., (C1
1qpt oRdMK
V 2log(
8(2K + K')²
//1.02d//1)
W V T
8 8(2K + K!)
navt C1
22
)/(d
ojdMaK
In terms of the upper bound, by similar approaches, for some C2, C2 C2 > 0 we have
C2nqvt 0id-/2K
8(2K + K')
0g W as
22log(
)o
ag T W (t)
C2Canqvt ojdK
C2 2 log( 8(2K K')?
V ins Eman
W nm
0/21/2
for 0 t t 8 == ( o7d /2 2 log( 8(2K+K'2 )). Note that as d (2MM log( K'2N= 2s2 5
by
Condition 3.1, we have to >> 0. We can thus appropriately choose Ti s tz.
We also have
CT G0
8(2K + K')2
ag ( as
2 2loo((
)01
M
C2
8(2K K')²
- 22log(
00
M
Lemma F.6. Under Condition 3.1, suppose Eq. (49) holds iteration t S T1, then for Wk: [K] and Ng. it holds that that
( ( V
S.
Smmn (1)
W
S ms, enmea < ( )²)
p M
<
(63)
for 0 Stt Tr.
Proof. In terms of the lower bounds, in the scenario (1) and (2), for the numerator it holds that
T
w W (1) nim,, en,m, 2 @((W() Snat m.ms, en,ms,)) )2)
2 0 22
(W!)) (1) S.en,my )²)= = O M
1 (W(") (ax + En.m..))
where the inequalities are by w S
= Eme(ar) ( Sn a...m.e... Eq.(51) and O
M
1+M "fcc ozu O o. 070
O
1- 1aMM-11j--fac-0 =fac" "o
< ms
( Therefore, the results are valid by consid-
ering the lower bounds of w E < 30, pap 2d, and upper bound
W S, (o) = O(a,d1/2) by (5.
40

---

Provable In-Contest Vector Arithmetic via Retrieving Task Concepts
Now consider the upper bound. Observe that
(W (+) V S n
W S Tn,ma, Cn.ms..
Tn.mm, Cn.mg,, )
S, T (0)
I S 12
Σ m (W Sn (t)
T( Sn Tn Cn,m)
n.ms, en,ms,
IW Sn T
1012
mn,
(t)
)² ms,
(0)
M
M
)8/
log
16(2K + R
oo/d1/2
= = )²).
The first inequality is by the positive estimate of the products (W Sn ms, Cn.ms. ) (W and upper bounds of |
(1)
mn the lower
(2)
(t) S, Th" well as / a& the last inequality is by (q.(5
Lemma E7. Under Condition 3.1, suppose Eg. (49) holds at iteration S T1, then for vk E [K], u E {as) U {br) setg m (MM and v o o' E [bb]se[[]] (SnmmIne MM]M) there exist some positive constants Cz, Ca, Ca, Co, such that
ag S =0at==0 O, after at most 13 (nav)- of 2 8(2K 8(2 j) iterations, a TW! dj would get positive. During this period, have
T
(( (1) Q ag) (w ag)l S 4v (og(16 + k/
(64)
During 0 I ag w v a S ((ladd within [t,, Til. where t4 satisfying S t4 S t3. holds that
a: ( ax) Car [1 14 - 1)2
1ogg(12 + k')
0RM²K²
aw ag) ( T ax) CCc((A/oog11())) k')9/8)o2 30rd )
3o3d
C202r
(65)
exp
oidk2
(W& a)) (W k u), (w) TW]())) (w (wi 0') = 0( (oog + K1)2//002d77
Furthermore, it holds that
@ (W (T)) ( (W ag) 20 of
Can of
40's 40344 (og((662K 22)
Mºqv
(66)
(W ag) T ((
(T)) a;) de lojd
Proof. For vy E [K], observing that
w(++1)
Q a= w Q a = n/Er
Σ mE[M]
w W!! (m - T (S, agl
Fuyy m
(ukya EKETKKKK Wn uk) S
N
W S (t)
KE(K ME[M] [ nENV
(W S (1) W ( S an
w S, (em (S,)(S Tagl.
| V S ( (1
W ( ST (1)
(67)

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
and also it holds that
(ukyn ExE(TK+ Wn,k4)) (WV Tn.m
W
11)
W
Σ
N
S, (1)
kE[K mE[M]
ve[i:] NENE
(1 (W 1 V Sn (0)
// S (e)
Tn en w w SeetiiSS - TT agl.
| V S n (1)
(1) T 11
(68)
Therefore, we have
(w(+"aa)) (W - -W"l) T(w aw) =(W(+t)
a =
w q 9 + q ax) (W (
a
w w k (Ka) ( (w 4a)
-W a -Wl ) TTa)))(()) "( aj w ag)
(69)
+ +(W ax)) T (ww (1+1 a -Wo (a))
+ +W a T (W 11+1)
a w ag))
Here,
(uky. Exe K un)) (W Sp mm
(1- (W S T
(W(+t)) a; TT a a=)==
Σ
N
[W S.Th
|W S. (t)
KCIK
WELLL
REN?
ME[M]
W V S. Tn. Cn.m
(wRas) S - (S. azl
| S. T
ExepKIKK wn))us) ((W() S mm em)
(W""S, (
!
Σ
O
N
IW S
|| S.
VE[+1]
RENT
McMM
W S. Tm Gn.m
(W ag) Twi T. (S, aill
( (2) S. 19
(Uly. -KEEPK+ un:) (WW( Sa S eman
(W S. m (0)
"
Σ 10
m
|/ Ww(Os, S (0
#EL+1 NEXY
W!
(1)
ma. Cn.ms. (W ag)
W (e - - (S.e(()e)))))
S.
K
(70)
Here, the second equality is by the near-orthogonal relationship
3(2K + K'NNM,
(Sw)e a a (ax + (n, S O(o,, V 2
}) = o(..00)),
for vx E [K], k E Ne: the third equality follows from o 222og
3(2K+KK)NM
< ((0000)), well
as the non-increasing nature of af (W//s Snorth (t) mmsg Coma.) ) ,n E NY, which is due the significant increasing behavior of a. (w(")) ag) as well as the scale differences of gradients characterized in (56) (57), (58), (59), (60).
IW ( 11
42

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Similarly,
(ukyn
EKE[TK+K un T T(W V Tn Cmso )
(wit+))
K ag -W ax)
((Wa) az)
7
Σ @(
(1
N
S nr 4
DEILII BEAT
(71)
(W S. n )
W
Sn ms, en.men ). (W Q az) W Sn
(0)
I S TT (+)
Spp (Cmsn
(1)) ()))Ta))))
(ukyn EKEITK+RI (W. S., (1))
By definition of u small scale of
Roingn eman
( 1 S
IW s..ll
aa npa -1) .. O
-M De ". of I mm,
S O (
Lemma D.11 and K" = (M)), through
M
similar derivation techniques in (57), we have
(witt)) Q a²) K (a))=a)))-) az) - -W" (x) (w a =e(( ax) (w ax - W az))
(ww(tt)
(W (wak) az)
(w!+")) a w az))
(uiyn -ERETT Wn, u (W( S, Anmms emaa
(t)
7
Σ
@e[
N
| S 21
WE[+1] NEN?
(W S. T. (9)
W (0)
Sp Tnimen Cn.mBn
(1
V
|W S. Th
((W R ag)
Twill "illss
T (Sner+1)) aj
K
+ (wl))
(t)
Q
ag)
w S (S. - -x!!)) Ta:]))
7 E ENEN, a wi =0(
(w('s, (+))
V Thimen Cman
I 1 (1)
IW S, (11
W Snin.ss enmma ) -1 - n (1)
n.mss )a ((W()) Tw!!)
Iw S
+ (w Jag)
(72) Hence,
Therefore, we found of (w!)) T w (we)) T (1) K ag) is non-decreasing after a T w
ag of (wl)
K ag terms in the gradients, which is nonnegative. v ag got strictly positive. Similarly, we have
w(t
K
(1) " EnEN, a W S ms, emsn (1 (W S ( ax) (w (( az)
(w(t))) az) (witt))
-0(
N
|w 1
1
| WW T("1
w V
(t)
S, Tn.ma, Cn.man (1
((() nnms. )a (2(W T ( )ag)
I S, R.
" Σ DiNN a (++1) (x)) - T( a.) - (W(+')) T (W
mnmm Cmen
(W "W S, 100 Ww!"S,TT!
N
|w S 11
w/ S, ग. n,ms, Cn.ms,
(1 -1-(())) (()
n.m.. )a (2(W!)) T w (a:))
W V
S m
(73)
43

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Then we can control (W Q a&) T (W
k ag) by
TW (1) k ai) (w To T(w k ag) - (wi+ a;)) T (w(tt1) @ a2) - (W
1 ai) (W ai))
(W(+))
2
(74)
(W K ag) T
(Wk (w(++)) (1+1)
ag)) 10-(WW) R a;) T( ag)
K
2
By similar approaches in (57), (58), (59), (60), it holds that for Vk E [K], E [K'] u [fK))) U (b.JoE]k] U (5.m MeM and v * v' E [b.)sejk] U (eL,) the growing of (wo) ( (w"?)) u). (W(")) ( ( v), (w T) (wlke e') is relatively negligible.
Now we check the lower bound of (W ag) T T( K ag), in the worse case before starts to increase. By Lemma F.5, the a V ag would got positive after at most la (nyv) (TT
2 log ((244K)))) Q ag) T( k ag) will in- K ag) would be its by Lemma D.9); in situation (ii,, when the K ai) decrease to be lower than 0, the situation would be back to (i). Therefore, the lower bound (ta)
iterations. During the period 0 S I t3. if (i), (W ag)a) (w ag) < O. (w
crease, otherwise (ii) it will decrease. In situation (i), the lower bound of (W Q ag) T(wkk lower bound of initialization (i.e. -4V1og(16(2K + K'Y3/8oogd"
(We)) a (w
(t)
of (w(()) ag) T (W ag) would be -4V110((16(2K + /1)2/)) Worth noting that the (W Q ag) (W ag) and (W (1)
(O)T(WW (w
(t) ag)) also decrease in the same period. By Lemma D.4, the lower bound would be (1 blz)eo.
ag)
where = 8/2 1 Vlog( $(2K))))))) )), eo = minx=qx(( (1) az) (W ag)) c.d/2. By σ1 S Og (log(((2( ((2**)'²² ) in Condition 3.1, we would have (wo) Q TT) ((W ag), (W K ag) T T( a;) 2 oz.//4
8
Next, we examine the increasing evolution of (wo)
Q ag) (W ag). During the period a T a S 0.010,d af-
ter a ( ag becomes strictly positive at S 14 S 13 = ) ofd//2M(( - 2 log(
$22+ iterations
(i.e., during la < < Ti), by Eq. (61), (74), Lemma D.5 as well as considering the positive estimate of the initial
(w(le) an)" T (wikk) ag), (w(h) ag) ((w()a) ag). (w a;) T ((w(ka) ag) at t4, we have
(t)) Q
(w (w k ag) s Ca((4vlog(16(22) + K'2//)0// + 3a
30 C2Conpavt )),
2 + 2 exp(
ojdk²
for some C4 > 0. Now we extend Ti -((m)))~oofd to this upper bound evolution, which equivalent to solving
the following system:
a + be
a 11((4 k2)3/8)ogd/12 + 3o0d )).
2
= C4 300d
b
2
C222C212
c oidk2
t
= =2,
We have
(w Q
ax)
T ((W(())))))) (1)
ax) S O(ffd(2 eat - of S Cso ofd
K
aRd
oid
for some C3 > O. That is, by gv 0
and = U2 log(o, zd- logt b.oo
it holds that
log(d 1 1og (
0.1ooofde of 49-4/ byg 1'1 )a/a1/2
(Ti) Tnmm.. < 0
Csofdee' - 18(2 Jazd1/2 + (M - 1)c -0e-ll--//inn 10-21/1og 10/244'2 jogd1/2
< 0
1+ (M- -1)e 1)e-ooaaev o75
o73
<
0.95
44

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
1- (M-1)e "ge of )1 (74), and Lemma D.5, it holds that
Define A (1
(0.097))(0.05) 0.0048, by Eq. (61), (63),
1+(M-1)e "gac" ofa
o -t-4 -1
(W ag) (W a o
0RM2K2
CanPavoō 14 4--122 -11
11
1 *K1)//800d21
0RM²K²
for some C3 0.
Similarly, we drag the Ti to the lower bound via the following system:
a + bb2;
Can 4034 (16662 -K))7800/81122
oPM22
Qva
oRM²K2
1 Ti -1
holds that
Curiavo oRM2R2
4C33 10g5/ + K3)3/80oddd2 + +C' 0.o3iR
(W (W ((W ((x))
M2qy
oo Car
²C Maqv
4C3 4C3y TCg/6og(16(2K + k')2//0003
MRM2
for some C. 0.
Proof. Proof of Lemma F.I. From Lemma F.5 and the condition T1 S Lg we derive the first result in Eq. (49) as well as the first lower bound in Eq. (50).
The second result follows from I Lemma F.4, where the updates of |ax: (-( ul,, to (/ v/ and & w constrained by the low-noise condition, the symmetry of the low-level real-value label distribution, and the large dictionary assumption detailed in Condition 3.1.
v o' are effectively
Finally, the third and fourth results in Eq. (49), as well as the second lower bound in Eq. (50) are established using Lemma F.7, together with the choice of Ti.
F22 Phase 2: Decelerating Growth of MLP and Attention
In this section, we would witness the decelerating growth of MLP and Attention.
Recall that the main gains in the first section is the growth separations denoted in Eq.(49) and (50) in Lemma F.1:
C/
8(2K + K')²
ag W (Ti)
as
log
)a O(M 1(M-1),
M
8(2K + K')
a w ( u//.
(T()) (1)
w/. le w (T)
0' - o(V==(0 2
8 )01),
for Vk E [K].u u E {as)a<ke(K) U (b.).e[k] U (Deleett U {En.m Ine(e) and w # o' E (b.).ejk U (Ti)
(En.mIn By = Q2M log( K" (²² )) in Condition 3.1, the scale of at (:(T))
V ag greatly surpass others. (t) mainly come from
However, as we see Lemma F.2 that the role of a w as in the growth of w/( S
(t) (1 ja (ag cos (W()) ag ax) Il) (w ax) (m(() ja {a w ax)
45

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Since og ( ag on the initial scale of [|W()s S, in (51) holds:
((a) during S Ti, it follows from the above formula that 0g W ag has only a mild influence (!) ((01d-1/2) as shown in Lemma F.2. Consequently, the small cosine similarity
ag WT/ Sn n (t)
S 0
W S, n (+)
1 + (M 1)e ofde ofa
for vk E role of w ag in the volume ||W V axll. To achieve this, based on the scale separations of gradients shown in Lemma F.4, we define
[K],n E M₂ C [M] 8 E [7], ms, + ms. E [M]. In this section we would witness the increasing
W a2
( w. a - (a W a - a / a a
V. ==
where (0) V.
denotes the residual components of w as w au orthogonal to au- That is,
w a2 W a + (a w (t) V a ag w
(0) ax VV
(0)
(75)
Similar to Lemma F.1, for all k E [K], the following lemma reveals that during the second phase (T1 S 1 S T*), the term ag ( V ag exhibits substantially faster growth compared to other cross-terms of the form T w( v, where {u, * {a,, a.) )e(K] This growth pattern establishes ag w ag the dominant component in determining the order of [ v ax:l- Furthermore, "v.a! an = 0, and the contributions from lI |l and (() asymptotically negligible, being bounded by the order of ( axll-
(1)
( V ag become
VVa
Lemma F8. Under Condition 3.1, during t Sts T = 22 '9 of where t S Ti. we have
8(2K K2) )o1),
la WWuul, l0 T b/, v0 - WW/''1= O(V 2log(
5
(100) 1l, (00 T "w( as W axll),
Via
(76)
af w as 2 v Ct
ngv -1-t t1)
+ (a (( ax)²,
MK
W v ax S Cs ng (t)
+ (aj ((( ax 2
Cangy
a
K
K(a
w(( as)
for Vk [K] [N], E [7], m,s + ms, E [M], u and u # s' E {bx]se[K) U [Ex.mlne(N)mmm()
U UMW
Furthermore, aftert t 2 T 211 qv KMd (og())) it holds that
4 0g w v a = 223(1 + 150; logl - (1 5
(77)
| | (oV K w ax)
for Ak E [K].
Proof. The first two results can be derived through analogous reasoning as in Phase 1. Specifically, the updates of a; Tw Tw( u. o w( and v. w(( e become progressively weaker. This weakening is attributed to the diminishing mag- nitude of u T w (t) v S, +!" / T (!) || for n E No. which results from the increasing dominance of as T (W("S, (t) Additionally, the already weak updates of a (w u, o ( V v, and v T w!!") vn in Phase further contribute to this effect. Consequently, the component norms and products V. |l. (()) (0) ization.
Vv. V ag are comparatively feeble with initial-
46

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
(0) After Ti, the terms in (5,( maintain negligible due to the increasing of "n" m ( [[]]W S n.
a W V dg.
I S. 2
( (l) as well as ag ( ag = a
W aw).
Besides, W u, Yu E {as) U [E[] U (mw hEERRy U (fn. (SN.mIn[MMMmM(MM will always preserve some volume over some directions after initialization, it holds that cos(ag, w V u) I constant < 1 during the entire iterations 7*. Therefore, there exists some 0, such that
us.)uu) - (URYA EneTTK Kg uh, (W SATT (t)
(4) W
S TT (4) 2 12c>0 0,
of ((ukya
S (9)
S.
KE[7K+K]
for all the iterations / ST. Similar to Eq.(56), we then have
T ELETK+RY wa.u (ST ( ak (1))
a (w(+1))
11+ ag 0g W.W
az
Σ
ee
(WyS
N
I| S (1)
vE[=1]
AENT
(IT n.ma aj + +m (T (am)) a
naV
Σ
@[
N
|| SA T
=E[=1| REA
(78)
Mm pn
+ +pM (M (1) -m, (32 + K
=
N
W S
(0) 1gv (M RR m.
1 D).
2N
REN
In this phase t 2 Ti. the ag W y ag is comparable to the potential maximum norm of the initialized its order is larger than those ST. with ak, u ( ag. Therefore.
(0) |
V S T (0)
I|
and
| ST (T i
(aa w az)² + 2( (t)
n.ms, [ a w ()) ; " @([ (
(79)
= e((a w )²)
(1) Then by Tn.ms..
0 [ it holds that
ag W (((
a a W(O au 2 A
1bqV
2M w as)
for vk E [K]. By the positiveness and monotonicity of ag (!()) ag as well as Lemma D.3, for Wk E [K] and some C7 > 0,
we have
ag W a 2 1 (t
(a) W(( w()))
MK
during 4 T t ST. Define T3
Md + 15a, Md log( {) + M V 2 log( 812K+K00
2n1 1 qv KM log( {)), we can choose > Tz. Subsequently, extending t Ti and aj w.(T)) a. d,d) (dd)sT TT
//2 2 log( x or in this formula, such that for some t 2 T* we have M
ag (1 A a 2 2(3(1+ + 150 ()g( W (0) V axr Il) = 2(3(1 + 150 (o
47

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
for vk [K]. Similarly, define T3-
M / 82 + oid), and choose
T 2 Ta $ T = 22(n qv of KM log( By || / l3 I 2d in Lemma D.8, for VV 2 Tz. holds that
T W a% 2 l aid Σ (a W ax) = 2 ( ( 11:).
0x
KE[K]
As we see that u / v v only grow if u = v vk E [K], we have /W//r= Ø VaK w as), e [K]. Similarly, by Lemma D.3, for some Cs > 0 we also have
ag w V as S C3 ngv (t t1)
+ (aa WiT ax)²+ + Kla, ((
Cangv
K
fort L ST SIST.
Lemma E9. Under Condition 3.1, during SIST ay Kk oKK log( where S Tú, we denote
(8) A min, Tn,ms,
=min M
}).
Then it holds that
1)e 1)e--11--880 -8 -F off (dloe(?)) Toqy
az) (W K u) (W v) (W 1 v), (W (W W = Ol - k183/8)0gd
(W az) T (w az) Scoonde av ofd
M i0ey
ofdk M
aw) T( a;) 20 ofopd
20
K (
noyt +1)
(80)
Mºqv
Can K2
4C3 /log(16(K + k2)2/8)08d'12,
for vk [K], E [N], [7], -ns ms. E [M] u (K)) U (ew IEELKY In
and v # v'é E (baket[K] U (N
Furthermore, after t 2 T = = 1 -1 of K d ()) it holds that
(W az) (w ax)) s O of of
1)
Tuo )
((T)) <0((
1
)
(81)
1 (M 1)e ofdear' (dd((()))))) 1047
for vk [K].
Proof. First, observing that
(1 (wl's, (0)
W( Sn mn, en.ma..
(e)
Σ Empmes
(W y Tn T T(W S,, mn. en,ms.))
| (
| V x
w S, n (1)
Empms, (W
(0) Spam mn.m)
T(W!!) (1) S, Tnmms, Cn.ms,)
(1)
O(la W ax)²)
O(2 m
(1) Tm ( (ofd))
@((1 Thims, MThhms. (afd))
O(aa (() ))
@((a (( a2)
(82)
for Vn E Nr. Here, the first equality is by definition; the second equality is by Eq.(79); the third equality is by the feeble growths of fask-irreevann components denoted in Lemma F.8 and [ a @(fdd). Therefore, similar to Eq.
48

---

Provable In-Contes, Vector Arithmetic via Retrieving Task Concepts
for vk [K] we have
" "EEN, a W/ (1)
y Sn ms,ms, Cmsn
(W WWVS mT )
ag) (wlt)ag) (Wkk az) - - (WWWaa) -0(
I S S 11
| /WWS S 1
W SW enmms (1 (t)
)a ((W())
(0) ww Tw0
[W SUSS,,, S
ma,
+ +W W K )az)
K
=0( (1 - ms..a. ) (Tn.man )
)a ((W()) /w/++
K(a W oz)²
(W K) Ti )aa)
(w
ax) (W ai) - (W( ag)) T(Wel az)
(w(++1)
=0(
2
(w + K az) T ((W(tt)) (1+1)
ax) (0 og) (wl)) k qg)
2
(83)
where the first equality is by Eq.(72); the second inequality is by (82) and (82); the last equality is by Eq.() Note
that 9(x) = x²2 (1 2 x E (0, 1) has pole
1/2, y 1/16, and is symmetry around the axis x 1/2. Therefore, k ag) in Lemma F.7, the upper bound of the
by Lemma D.2, Lemma F.8 as well as the upper bound of (w Q az) T (w())
evolution of (wl) Q a²)
T TW R ag), vk E [K] is the following system:
(w Q ax) TW ax) S 0((
)( + d )e (41 + 2
by
e
=
ToK )
11
b
CsoRdenr' old
nay C7 MK'
d
C₂ 20 ((rv) (lonv)-ooouu ofak)
+ ofd
4 = (1av))1of(d)
Here, obviously < by M 2 2 and qv 0/ < by Condition 3.1. we take bt, as the upper bound of (wo)) ag) Tww K az), Vh E [K] during the first phase; d is defined by the lower bound of a w-( & ag in the second phase by Lemma F88, with k w( ag == d/2 be the upper bound. Therefore, we have
(W Q aR)
ax) T (w a s de 07d
+
M + 1
T02V
oidK (M +1) M TopV
M +
dgo off
1ggtt
oidKM
for k [K]. Here Cg > is some constant. Extend 7 ==== qv of M log(H)) in the formula, we obtain
(Wo) az) (W k ax) < de of 0' (dd(d -
1849
As such, it holds that
(T*) I <0(( n.m.
1
1+ (M 1)e 13e-o azdecv " 2o94 (d ( Tepn (t)
To examine the lower bound, we first denote A = min, An.ms, = min 17
1+ (+ 110--10--3ac of (a (3) r007 H). Sec-
nav (1-1)) + (a. W (1) ax)² + ondly, to alleviate the complexity, we choose linear upper bound of a w as S v Ca
K
49

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Cangy 2 (aa) By the concavity of g(x) var +b cx 2 31, the linear upper bound after 21 is g(x)
(ax1 b)-1/2(xx 21) + g(x1). That is
Cangvo
dd11/2(t -a
t1)
0g W ak S
+ tod
2K
(84)
Canqvoi'd 1/2t
+o
2K
Therefore, by Lemma D.6, (65), (73), (74), (83), the lower bound of the evolution of (W T( ag), vZ E [K] is the following system:
C3n q 403 4C33/log(66222) + k3)2/8)o/d'2
(wo a T (W K a²) 2 Ogg, nuisor xwi)tt)
M2
(1 - A)
Canava 12
2K
apd?
ofoidh
C6
Magy
f = 0.
Here, a 2c by A2 A)² <
(M-1)²
and the minor gradient step 7 = of g o Ar M4
by Condition 3.1. Therefore,
(M (M=1)²)
holds that
(1)3
(W (w ag) ((W ax) (0a))2(0 200
ogofde Maqy
novo
La 1/2 neya
La 1/2
nev 2
Le 1/2
e
C3n q
4C3 / 1lg(166(K + K'Y//8003d'2
oRM2K?
Byy 2uy ko2 2)2 +2q5 +24v" 1220'KA2(--AA) ngvi 20fax
201 oz
M2qy
e
Can
403 (og(oog16(2K k')2/8)03d122
M2
for some positive constant C10-
G. Training Dynamics: ICL-type Data
This section also assumes the results in Appendix D.2 all hold with high probability. This section examines scenarios involving training on ICL data, using notations for T, consistent with Section F. The proof strategies for QA-ICL data are identical to those in this section. Therefore, we focus solely on the training dynamics for T., n E [N].
The projection evolution dynamics remain largely the same as described in Section F, with the key distinction being the evolution of b. w b.. This difference arises inherently from the gradient forms, where the unavoidable imbalance between positive and negative data for some concept E [K] can drive the growth of either bJ w((t) by or -b w(() bg. as suggested by Lemma D.10.
Lemma G.1. Under Condition 3. during & < -(((ngv) -oofk)) for Vy (+1), there exists ky E [K]. holds
50

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
that
1b b b 16 W. bu_ I. of W/P ax o(o)d!/2),
8(2K +K'
bb T ". bk 18x+ Tw axb |bx_ Tw axl O1V 2 log(
)01)
8
8(2K |K')a
bx+ TW/ul, 1bx_ TWWel, ul, as Tw ul. l T w u). l0 w 0' = O(V 2 2log
)01)
(85)
8
ast. (abT, )1 = O oid).
Iw Q xn.) T (Vn.mer + En.-mr.)1 =n 1 log( 16(2K + K)²
)o
8
for vk E [K].n E [ E [7]. T. # mT. E [M] u E {0s)s U [bb]sTkLEIK] U {uw bee U (MM and + o' E (b,1 U (SS.mIne(M)mmMMM At the end of this stage, we have
0x 2 Giord!/
8(2K + +22
0% W
2 2oog(
)01,
M
5
oBojdL
(W az)
(Wk*) (11)
01) 2c Maqu
C3(
+4 444log(16(22) (6g(16(2K + K1P//50o/d13))
of M2K2
(86)
Ciord!/2
8(2K + KK)2
1 w
bk 2
2 loog
)01
M
8
C5ord//2
8(2K + K')2
1b W.
bk_
2 log(
8 )0,
M
for some positive constants C1-5
Proof. The proof strategies follows Lemma F.1, despite differences in the growing of [b2, w (25)) bul and 16 w (Ti)) bs as follows.
By Lemma D.10, we see that for Vy E {+1) there exists ky E [K], such that WN N/(2K)) 2 5- 10-10- 10-11013/1100 (2K) Therefore, similar to (57), we have
E ERCPKIKT en (S (0) Tb
by (W S. )
[by, Wil Dk, bb, w bkyl = 1////15 EEpg
S. n
b Miya Eve KI) AA (ST i
(1)) ) bky
"a
Σ
A[l
(Wv S. m (2)
(87)
N
S o 10
1|[=1] REAT
(8)
ndV A.ms
s 10e
Σ
Oaa w ax. ak w (t) V (xx).
RENY WAT 2N SAn
(4)
The inequality is inherently due to the existence of bR, in Ta, while QA data did not possess by, in its sentence. Since we see that No is guaranteed to not balanced with probability 1 5, the growing of [b%, w (t+1) bb, bT, w y inevitable. The remaining proofs follow Lemma F.2.
5,
by. is
V
u {ast Lemma G.2. Under Condition 3.1, suppose Eq. (49) holds at iteration t S Tr, then for wk E [K],
(bs)oejK) U (5mmm (NE()) and v * v' E [(b))e[K] U (5 IN( there exist some positive constants Cz, C., Cs. Co. such that
'If TrTTTi W ag S Oatl=0 - O, after at most L3 = (T
2 log( M22+K'PB
iterations, af WiO V dg
would get positive. During this period, we have
(w() (1) Q ax)
(W k az) S 4/log(16(2K KK)
(88)
51

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
During 0 S ag W as s O(o)dt12) within [14, Til, where La satisfying S la S 13 holds that
CoreavoBll - 14 -14-)
(W (W ax) 2
(16 k3)2/8)o8a"12
0RM²K2
(W) (W ax) (.((( + k179/5)oft112+ 3ord
3ad
CoCon
(89)
+ 2
expl
))
oddk2
(W ( (Wk u), (W (W ( v')
0 -0 log(16(2K k1)
Furthermore, it holds that
oooid²
Can M/
(w T (21) K ag) 20B M2qq
- 40' 40gg/1g(16) K2)9/0)02d12 "
(90)
(w az) T (W K ag) $ Cso 'ofa
Proof. The proof strategies follows Lemma F.7.
Lemma G.3. Under Condition 3.1, during t S15 -00 'gr Kd (g where t Tú, we have
8(2K + Kn2
Ib__ T Ib,, Tw asl. 1bx_ Tw a&l = OOy 2 loo
)o1),
8
lbx
Tw(P uul. |bu Tw(uu), lax TW/uu u/, l0 T o 10 "w -O(
2 2log(
8(2K + K')²
)ai),
5
(1ttt)
Vaa
11,
(1) Vvatt
w/ ax = o[l axll).
(91)
ag w a# 2 c,, ngv ny t1)
+ (a. wif ax)?
MK
a w ax S VCs
nav (t -t))
+ (a wif)) ax)² +
Canqy
K
K(a (((h)) V ax)
for vk E [K],n E [M]. s E [7], mmr * mr. E [M], E {as)szeE[K) (ba) heeeey (GMM)) and v # v' E [b.),e[k] U (Gmmm)eetNMmmM
Furthermore, after & 2 T === === 1 .y of KM log(t)), it holds that
a (1) a# 22(311 + 150, y log( E d)
W 1 || O(VKa ( (()aa)) ax),
(92)
bj, w bul. ToT w b = O w/ as).
for vk E [K].
Proof. Based on Lemma G.I, the proof strategy follows Lemma G.3.
Lemma G.4. Under Condition 3.1, during 4 S S T* 22 of log(? where 4 S Tí, we denote
52

---

Provable In-Contest Vector Arithmetic via Retrieving Task Concepts
(8) A min, Tn,me,
min(
}). Then it holds that
1+(M-1)e - of (dloae )) T0gV
0') (W (W K u) (W ( v), (w ( 0') = 01 000oog111 K318/8)08d13) (W a (w ai) SC SCookde ofd
2
194t
M
j7dk M
aw) (w aj) 10 10 M2qv
ofoide
24v
K 4 2 (1-4)
may +1 20fuK"
+1)
(93)
oiM²K² 4Cs /l3g/1g + K
for vk E [K], E [M] [7),-m,, # mr, [M]] u E (0 ((w (S Ine and e +v E {ba] U (5n.m
qy Furthermore, after t 2 T = -( - KM log(t)), it holds that
(W TW ag) OOogde '07 (dlog) )) (T') M,
1oey) of
(94)
s 5
M 1 (M - 1ee---e---- aadeei (f 1867
for vk E [K].
Proof. The proof strategies follow Lemma F.9.
H. Test Loss Convergence
This section also assumes the results in Appendix D.2 all hold with high probability.
Proof. Proof of Theorem 3.2 and Theorem 3.3. For the ICL task, given the prompt T == yxF.FT. .xJ,J Ffxx+1 (23+()) we denote kyT_, = 7kir + yr. as the dictionary index for y f+1 kr as the task concept index, yr E [J + as the label indicator of y. Also we define
exp(e]+1 T (W(l)) w Te,)
(1) J.T
Vj E [2J].
Σ EEE(2 exp(e],, TT (W!!) TT)))
If p = Pr, then
exp (ux T (ho) CklTKKK1 exp (n, he)
Lp. =
* argmax
))]
-Epy(1(kyT + argmax u ho)l.
k
Ejetn W Tx (
are
-Ep-(1(kyT)) argmaxx of (xf+1
ESeen w WT
ee
EjEN w( x
n v
(1)
(95)
-Ep-(1kkg) * argmax
of (xF
EEEE w ( + x ))]
=Epi (1
Vr,J+1 +1 # argmax(ux)) T((
Σ
Fo akzz yksts bR EI41T
2
Σye( w V - + T3 T)
+
x + 2 y T) })]
|| ||EjEt)
(T 2j
x
(t)
53

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
where the equalities are by definition. By the isotropic probability property in (1.( and the definition of U in (1,6) we have
(RT + +r,++IbbT) ( (E.41.T) + (IrasT + +1,++1bbr)
Σjen W/ (c)
(T LTTx. + (6)
T2,
Lp- = [1p|1(11
V1.01
VI.01
11 ||jEM w (m 2/-11117 xxx + (6)
2 T)
1.1
(akr + (r,++1brr) ( (E++LT) (akr + (V.J++bbr)
ΣSEN w (T2y x ! (1)
V2
V2
V/2
EEEM W (T2 1T
T + 22 13,
(0)
Ep- [1.442 + (J+1,T(-0055 + 0.115gT,
Σjen W (T2 + T2)
1.41( abr + +vr
EEjEM w
( T22 TV
(2)
ΣjEr W (t)
(t)
> + (ast + VI.J+Ibbr) )
|EjE,t W T.
2 T.
Ep. [10.42 + L (541,T(-0.58560 0..885xxr + 104158T,++1bbt
Ejen w/ (m2)
1 T
+ 22. T)
(t)
[0.86arr 0.4be)
T ){
|EjE(1 w (g
1--11T x
(t)
(t)
S Ep. (1.5 EJ+rT( 0.7 (.5y
TT ΣjEry w( 12
(6)
F) ry,
(t)
> jak 0.5yT,
)}
|[EjEt W (0)
T
(96)
where the first equality is by the definition of dictionary U as well as the mutual-orthooonaaliyy of {a=)etkE[[K] {bs) U (uw heepeen which implies that the main competitor of the label-related dictionary token, namely
Gir +0r.++10e
tor. is the word-related dictionary token, namely 2. B.asy t tor.a-bbrr Tr tUr
Jerr -UT
V2
an tt the second inequality is by V2 < 1.415, 2((.01) S 1.43, 2/(1.11) S 1.41, >
V2
1, 1,1111,01, V1.01 < 1; the last inequality and equality are by direct calculations.
(0.1m (1) + ( as the accumulated softmax weights assigned aur in T, XT Uرeeu jr Tai, as the accumulated softmax weights assigned on those current prompts' co-task-irrelevant task vectors as in T, and ir = Ee Exr, Uk," by,) + yer b we have
Then denotes TakT
(t)
(t)
Σ (T (T + (10) ugGTT??) = - akr +
Σ ((
T Tas ax.) + iT + ET.
(97)
]E|j]
(1) Here, E9 = (nep) Eir)
(0)
(=)))x))) I). Let of
o EjEro (1T6) (2) It is straightforward to observe that the upper and lower bounds of of are 1 and 27. respectively, corresponding to the cases of extreme and uniform weights.
2
54

---

Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Consider the training distribution is POA. then have
Lp - Ep. (1.(///1,,,(0-77000 + --.5yT.++bbb
WV (Taiy aer + (Eneex)((k)
Tas, ak) + ir + Err
> a 045 T
Tai, ax,) + ir +xrill ) 0.5],
(1) [W (Tar aur (EneEdt
< Ep, (--00aa 0.5yr,J+1bby
(1) Taip
W akr
a
w (0)
ETl
0.5]
(
akrl
Tak
(t)
ak II) + (WWerrl Er
Ep. (1_-))1111111(..760000 0.5yT.J+1bb,
(0)
akt w/ aut 10 W Erl
Tarr
) - 0.5]
(98)
TTaIT + EK (NT) Maz, ) (/ V our II) + | Erll
S Epp, (1E),,1E)1, (-0.7ab + ++1r,,++1bk
(1) Tapr a W akr | w Erl
0.5]
(TarT + EEIEXT((TT! (1) < Ep.
[ (0)
acr (l) + aje w axT) + [/ Erll
Tai,
+ 0,5gr,++1bbr)
o w Erl
(0) RagT
a w.y auT
0.5]
Σ 2 Rak
( axp
(1 +
(0)
+ 1)
+
OuT
W
aur
(6) MapT
a W
aur
Here, the first inequality is by (96) and (97); the second inequality is by triangle inequality of norm as well as Lemma F.8 denoting the dominant growing of ap w air compared other projections; the third inequality is by the balanced growing of all ([w([) V | denoted in Lemma F88; the forth inequality is also by traingle inequality well Lemma F88 showing the main contributors of norm. Subsequently, by Gaussian tail bounds well as o O(K log( (2++1) -²) we that
4 2 1.50 log( () s $/4.
Pp- (4]+1,,(-007aa + 0.
(99)
Additionally, recall that vair Ejeen scenario, where the number of task concepts outside the prompt's co-task concept may exceed 1.1J (the number of a&r in the prompt). Given J = f2 (log (2) (22loggg (K)), the probability of this event is bounded by e/8. By Lemma F.7 and Lemma F.9, which characterize the progressive learning of task vectors and the limited learning of non-task vectors, we have of = O (KK log 12=+1))=
(2)
(0.1 Eje
(t)
T) and B, 11r( We first analyze the worst-case
This implies that the probability of noise vectors exhibiting significant components along the axes of any ak₁ or akr is constrained to less than e/8. follows that:
(0)
Tas
Pp.
(001)) 2 1/444
(100)
(t)
Take
Besides, by Eq, (99), a sufficient condition for the event in the last inequality of Eq. (98) is that the right-hand side
Ja W( Exl
Takp ley
(1)
will ase
log (
< 0.5 + 1.50 p1
(101)
|| azr
w Er
1.11
+ (8)
W abr
wake
WiO axT
55

---

Provable In-Contest Vector Arithmetic via Retrieving Task Concepts
h W Erl
(0.5 1,50 Soogyloog log (2)) By Lemma F.8, as well
This requires the denominator to be no less than
(1) Tait
ony WV acr WWill WV
la W Grl
IW
as of S op we have
NO ' " Klaxa)). Again, by Gaussian tail
Kapp "ky W auT
anp WV
"kr
bounds as well as of- =((( log( 2x+1)) *)
we have
Pp. ( WWWErl Er [
20 K Klog(
4
1 0.9 7,30
p
K Klog(
S e/4.
(102)
kT W air
2
1 30 /
K Klog(?)
Note that 1-20% / loR(
0.9-3.30, p /ogl :) 4 40 V loR(
0.9-7.30 v (o
1.1
Therefore, by Eq.(1
0.5+1.50; v (o)
1430 v (=
1-30; v log(
lapp W Exl
lI Srll der W
we bound the
in the numerator to be less than 20 K log( ()) and the
in the domina-
(1) FakT aN
wV air
tor be less than {
1 1-20 (lo(()))
1.1). Therefore, our last job could be showing 1.10
li akp |I
0.5+1.50 (lo(())
a W aiT
1 1 1-20 K (og()))
1.1), which can then ensure Eq.(1 holds with probability less than e.
0.5+1.50; y (+)
By Lemma F.8, of = O(K log( (2++)) - #)
and the Taylor expansion of g(x) = 0.9-7.3
1+
10
9
1000
81 J
((z) at zero,
we conclude that
agp w ONT
(23(1 150
V log((= )| W (0)
10
1000
2.2(
+
(1.50 log(
akrl
9
(103)
W akt
1
1- 20
K klog(
(1.1)
S/2 2
1.1).
W
(0. + 1.50 / log (1)
akr
kop
By (9.), (99), (100), (102), (103), by union bound we have
Lp. S..
Similarly, for QA Sentence distribution S ~ Por as well as QA Sentence based Prompt Dsitribution Tox ~ PT, the population loss convergence could be shown with the same strategy
Lew LpT S S.
Separately, when the training distribution on Pr or pa for prompts whose task concepts are li. with probability 1/2 >> 8, T has more words with bE than b b_k, This would lead b,, w(()) iT ignorable by Lemma G.3. In addition,
| || 's contribution to Eğe(J (!()) (m x
(0) T + (1)
T I| is also ignorable by Lemma G.3. Therefore, it's safe
draw the conclusion that with probability least 1 8,
(0(0) 0(1.
Through similar approaches we have Lpg (0(1)) = (0(1)= 0(1.
Proof. Proof of Proposition 3.4. The proof follows proof of Theorem 3.2 and Theorem 3.3, grounded on the learned knowledge on task vector ag, Ak [K]. We here omit the proofs for brevity.
56


Reviews:
[{"review_id": "VVJQOG73tw", "reviewer": "Reviewer_R7sM", "summary": "To study retrieval of task vectors in ICL, the authors perform a careful gradient descent analysis on residual self-attention modules (with nonlinearities and normalization) under a synthetic (but empirically-motivated) data distribution. They find that when pre-training on QA distribution (and testing on word-pair ICL distribution) the model achieves near zero test error w.h.p.; however, the model fails to accurately retrieve the task vector when pre-trained on the pair ICL distribution or a hybrid distribution. The transformer learns to extract the task vector from the context, then adds it to the final token embedding at the residual step, thus completing the ICL task in word2vec fashion. Accurate task vector retrieval explains the model's ability to perform ICL on low-level concepts unseen during training. The results are strongly dependent on the distribution of the prompt embeddings, which is assumed to have some particular structure (motivated by empirical observations) but the theory does not explain the process by which the prompt embeddings obtain this structure.", "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": "Thank you for your thoughtful review!\n\n**Q**:Chicken-and-Egg Dilemma & real-world impact\n\nWe'd like to emphasize that theories often relies on abstract, empirical-motivated models to enable tractable analysis and explore a model’s potential—an approach we adopt. However, we appreciate the opportunity to discuss how our theory connects to real-world multi-layer dynamics.  \n1. **Concept Geometry & Vector Retrieval**: Next-token prediction (NTP) shapes concept geometry in prerequisite layers[1]. [2] show that task vectors emerge in earlier layers during ICL, while arithmetic retrieval occurs in later layers.  \n2. **Theoretical Simplification & Justification**: The analysis of multi-layer dynamics is typically intractable. Akin to [3-4], which model simple feature-noise vision data justified by *ResNet’s deep latent space* properties[3], we simplify by directly modeling the resulting concept geometry from prerequisite layers. This is further supported by *layer convergence bias*—shallower layers typically converge faster[5]. \n3. **Alignment to Real World**: We offer the first optimization theory for a realistic transformer on QA data for vector-retrieval ICL and showing its *OOD edge* over word2vec, potentially addressing Question 5.1.4 in [6]. Yet, like [3-4], our outcome does not fully match real-world multi-layer dynamics. Empirical evidence suggests that shallow and deep layers most-likely *co-evolve*: even in a toy setting for approximating polynomial functions, [7] shows that mechanisms across layers must evolve simultaneously to reinforce each other—layer-by-layer training induces unwanted errors.\n\nBy studying this non-trivial learning problem in a **comparatively realistic** setting, we take an important step forward. A key future direction is to extend beyond [7] by analyzing a 2-3 block transformer on random spherical features, integrating NTP and QA training to model the *self-reinforcing* co-evolution of concept geometry and arithmetic retrieval based on empirical observations.\n\n**Q**:Experimental details\n\nOur experiments are conducted on synthetic data defined in Definition C.1-3;our model is defined in Section 2.2(no layers frozen);our hyperparameters are provided in Section 5&B;our Algorithm procedure is in Section C.1. \n\n**Q**:Room for presentation\n\nWe highly appreciate the feedback and would add:\n- **Ineq.(9)**: Simplified flows bounding projection updates are listed in Appendix D.1(Lemmas D.1-D.6), abstracting complex updates into sequences with simple constants (e.g., $a,b,c,d$). The idea is that since analyzing original complex gradient update in eq.(27) is too hard due to the nonlinearities of our problem, we instead split training into phases identifying key components that dictate update growth rates, avoiding tackling original dynamics directly. For example, in the first phase, $||W_VS_n\\pi||=\\Theta(\\sigma_1d)$ dominates, while $a_kW_Va_k$ starts small and the grows’ rate is then $a_{t+1}=a_t+b$ (eq.(60)), bounded by Lemma D.1’s continuous flow,yielding ineq.(9). In the next phase, once $a_kW_Va_k$ controls $||W_VS_n\\pi||$’s order, the latter, as a gradient denominator, slows $a_kW_Va_k$’s update (Lemma 4.3, simplified in Lemma D.3). Unlike [8], our problem’s pronounced nonlinearities demand more complex treatment. Per your advice, we’d add intuitive links between simplified flows and actual updates in both main text and appendices.\n- **Eq.(11)**: To explain our *harmful overfitting* in eq.(11), we contrast it with [6], where vision-inspired data(features+Gaussian noise) leads to noise memorization. There, high noise-to-feature ratios make the inner product of weights with noise significant in the gradient, causing harmful overfitting[6]. Our case involves memorizing low-level features: in ICL prompts, some demo pair's low-level features w.h.p. co-occur with query word's ones due to imbalanced frequencies in **finite** training sets, which produces small but non-negligible products in projection's updates, driving harmful memorization. Akin to [9]'s result—feature co-occurrence biases gradients toward memorization—in our case ICL pairs introduce unexpected co-occurrence. QA data, lacking low-level features before query words, would not memorize this co-occurrence and focus on task vector. [10]’s \"relation\" token mirrors task vector, though their artificial model (their eq.(13)) and data lack realism. Empirical backing includes [11], showing QA aids (multi-token) recall.\n- **Plots & K’**: We’ll update y-axes to \"projection\", use LaTeX formatting, and enhance captions. We’d fix eq.(14)’s ellipsis and explain $K'$ (number of task-irrelevant tokens) around eq.(13) with intuition in the main text.\n\nThanks again for your feedback! We welcome further discussion and appreciate your time!\n\n*Reference (arXiv identifier)*\n\n[1]2403.03867\n[2]2305.16130\n[3]2012.09816\n[4]2202.05928\n[5]iclr.cc/virtual/2023/poster/11533\n[6]2405.01964\n[7]2001.04413\n[8]2310.01975\n[9]2410.09605\n[10]2412.06538\n[11]2309.14316", "claims_and_evidence": "The claims are clear and the theoretical evidence appears convincing (I did not carefully check the proofs, though). The experiments could definitely be more convincing (see Experimental Designs section). \n\nI have some concerns about chicken-and-egg reasoning. In particular, the theoretical setup here assumes that the latent features that are inputs to the transformer layer *already have* the prerequisite structure which makes word2vec-style ICL possible. This structure then enables the model to solve ICL by simply extracting the task vector. However, it seems plausible to me that, instead, the model first learns orthogonal task and concept directions in the weight matrices, potentially by some other mechanism; these task and concept singular vectors might then be carried to previous layer weights via backpropagation, which pushes the latent representations towards having the observed structure. In other words, it's possible that W_Q, W_K, and W_V obtain the structure depicted in Fig 2 through a *different* mechanism, which *later* causes the embeddings to have the hierarchical structure observed by Park et al. If this were the case, it seems that it would invalidate the gradient descent analysis performed in this work. This possibility is not currently ruled out by the proposed theory. (Most likely, what is actually happening is something in between -- the embeddings and the attention weights gain structure in tandem, each reinforcing the other, throughout training. But I understand that such an analysis would be very difficult to do.)", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "The proofs are rather cumbersome and outside my wheelhouse, so I did not check them carefully. Sorry!", "experimental_designs_or_analyses": "Are the prompt embeddings synthetic, or are they taken from the learned embeddings from open-source models? If it is the former, then it may be useful to perform an experiment with real data, to show more convincingly that the prompt embedding assumptions are satisfied in practice.\n\nAn even more convincing experiment would be to pretrain a small language model from scratch, on a small dataset, and showing that the input embeddings converge to the desired structure quickly, and that the attention weights then learn the structure indicated by the theorems. Such an empirical result would also address the chicken-and-egg issue raised earlier.", "relation_to_broader_scientific_literature": "I don't know this area well enough to comment on this.", "essential_references_not_discussed": "I don't know this area well enough to comment on this.", "other_strengths_and_weaknesses": "This paper makes strong assumptions on the prompt embedding distribution; in return, the analysis is able to handle the full architectural complexities of transformers, including realistic prompt structure, softmax nonlinearities, layernorm, and cross-entropy loss. This seems to me to be a major strength of the paper (that the authors obtain analytical results in this complicated and highly nonlinear regime).\n\nProbably the greatest weakness is that the proofs are very difficult to follow. I understand that it may not be possible to simplify them. However, if that's the case, I think it will be very beneficial to provide a less rigorous derivation of the main results (for example, eq. 11, ineq. 9) in a simplified setting. This will likely help readers build intuition and make the key components of the result more transparent. I think the intuition behind the proofs should be provided in the main text as well, especially explaining eq. 11 in more detail. Why do the weight matrices learn singular vectors aligned with b in ICL and QA-ICL pretraining, but not in QA pretraining?", "other_comments_or_suggestions": "In Matplotlib, it's straightforward to use LaTeX formatting in the titles/axes/labels. It would greatly enhance the readability of the plots. Minor point, but some of the y-axes are labelled \"projection length\" but lengths must be non-negative. Maybe simply \"projection\" is appropriate?\n\nIn eq. 14, I believe there should be an ellipsis between \\nu_1 and \\nu_{K'}. Also, the meaning of K' is inferrible from eq. 14, but I can't find it stated in the main text.", "questions_for_authors": "I'd rate the current submission as a very weak reject -- the proof technique appears impressive, especially in its ability to handle realistic transformer architectures, but there are a few key weaknesses in the approach, as far as I understand. I would be happy to increase my score, pending some clarification or addendums.\n\nMy biggest concern is the chicken-and-egg dilemma raised in the Claims And Evidence section. A compelling explanation that rules out the alternative hypothesis would convince me to increase my score. Ideally, I'd love to see an experiment that addresses this, but I understand that it's probably a lot of extra work.\n\nI would also like to better understand the main ideas behind the major results. If it's possible to add an appendix where versions of ineq. 9 and eq. 11 are derived in a simplified setting (a nonrigorous derivation is fine, in my opinion), I would increase my score. I would also appreciate if the authors provide further explanation/intuition of this part of the result in the main text.\n\nThe last minor point would be to improve the readability of the plots and provide more description of the plots in the figure captions.", "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": "LDGgDdjQJ1", "reviewer": "Reviewer_mnid", "summary": "The authors study task vectors in the context of single-token factual-recall ICL tasks. In this context, they show that training on QA data enables learning a task vector which can effectively solve ICL problems for Word-Label and QA tasks. They additionally analyze this phenomenon theoretically.", "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": "Thank you for acknowledging our theory as empirically-supported and sound. We appreciate your professional review and address your concerns below.\n\n**Q**:Are QA and Word-Label tasks natural language or simplifications? The example in QA Sentence Distribution is confusing. The paper would benefit from showing some examples of the exact task sequences in each case.\n\n**A**:In our theoretical context, the QA and Word-Label ICL tasks are empirically-motivated abstracted simplifications. The example in lines 201-202 backs our modeling intuition—QA comprises irrelevant tokens, a task vector, a word, and a label—akin to [1]’s model (in its Figure 2) and inspired by [2]. Per your advice, we’d add illustrative examples to further clarify our ICL data's intuition: e.g., $T_1$=[Japan,Sakura,France,Rooster,China] with co-task “National Symbol” yields **y_3**=Panda, while $T_2$=[Japan,Sakura,France,Iris,China] with co-task “National Flower” yields **y_3**=Peony. These highlight the multi-task nature of our multi-concept modeling (e.g., “**x,y**=Japan,Sakura” fits ≥2 tasks), reflecting this sparse coding-type model suited for capturing *language polysemy*[3-4], more realistic than prior ICL theories over unrealistic data [5-6].\n\n**Q**:The study of geometric relationships of the trained models seems to disappear\n\n**A**:We would like to remark:\n- **LLM Geometry (Figure 1)**: Task vectors align differently with words and labels in trained LLM prerequisite layers, forming our data modeling for optimization theory, grounded in LLM concept geometry [7-8].\n- **Trained Model Geometry (Theorem 3.2 and the discussions below, Figures 2-4)**: the first layer’s output $h_{\\theta,0}$, formed by $W_V,W_Q,W_K$ before adding residual vector, varies by training data. QA training aligns $h_{\\theta,0}$ with the true task vector $a_{k^{\\star}}$, ensuring correct prediction when added to any task-specific **x** via residuals. However, when trained via ICL-type data, the $W_V,W_Q,W_K$ would *non-negligibly* memorize some low-level features—akin to [9]’s results. That is, $h_{\\theta,0}$ aligns with both $a_{k^{\\star}}$ and low-level $±b_{k^{\\star}}$ to some non-negligible extent, leading to constant test error w.h.p.. Detailed descriptions would be added to the figures' captions.\n\n**Q**:Broader impact may be limited, as it is unclear how well this transfers to the motivating settings of task vectors in pretrained LLMs\n\n**A**:We would like to remark our impact on elucidating observed LLM phenomena:\n- We elucidate why gradient methods yield task vector arithmetic in LLM[10], why QA boosts retrieval[2], and why transformers outshine word2vec—grounded in an empirically-motivated theoretical modeling akin to [11-12]’s analyses on feature-noise vision data (justified by ResNet’s latent space [12]), common due to the intractability of analyzing multi-layer dynamics.\n- Beyond this, our theory supports recent LLM task vector-arithmetic work (e.g., application of task vector in editing, unlearning, merging [13-14]), which assume vector arithmetic between pretrained and modified models *without explaining its origins in language model*. Though focused on single-token recall, our optimization theory, grounded in concept geometry, provides a foundational step. Section 6 outlines future work on complex mechanisms to further enhance this domain’s theoretical merit.\n\n**Q**:Is there a way to vary and study the complexity of QA and Word-Label tasks? In the natural language settings you draw from for motivation, we have a clearer sense of complexity for different QA pairs and, similarly, a clearer sense of concept hierarchy.\n\n**A**:In our context regarding retrieval over hierarchical concept graph, one way to measure a task’s complexity is the difficulty a model faces in achieving high confidence for its argmax-sampled prediction. A simple metric could be $C(T)=1/\\max_{y}p_θ(y|T)$, where $\\max_{y}p_θ(y|T)$ is model θ’s top answer confidence, and a higher $C(T)$ denotes greater complexity. QA tasks terms to be simpler: a keyword (e.g., “capital” in “What is the capital of Japan?”) guides θ to the task collaborating with query word, more-likely keeping $C(T)$ low. Word-Label task, in contrast, lacking this cue, requires the θ to infer the task behind the pair—given prompt $T_3$=[Japan,Sakura,China], the model might have non-trivial confidence over both Panda and Peony (e.g. the answers of $T_1$ and $T_2$ defined before). This stems from the polysemy-induced challenge due to the hierarchical concept knowledge encoded in the tokens.\n\nShould the reviewer wish to discuss any point further, we would be more than delighted to continue our productive exchange! Once again, we deeply appreciate the reviewer’s time and valuable comments!\n\n*Reference (arXiv identifier)*\n\n[1]2412.06538\n[2]2309.14316\n[3]1601.03764\n[4]2105.15134\n[5]2402.15607\n[6]2402.01258\n[7]2406.01506\n[8]2403.03867\n[9]2410.09605\n[10]2305.16130\n[11]2012.09816\n[12]2202.05928\n[13]2212.04089\n[14]forum?id=vRvVVb0NAz", "claims_and_evidence": "The theoretical claims are well-supported in the paper, albeit in a limited context and not for real-world LLMs, and the empirical claims have some support.", "methods_and_evaluation_criteria": "The methods, e.g. evaluating the test loss under different training settings, make sense but are difficult to understand concretely. It is unclear if the QA and Word-Label tasks in this context are natural language, or a simplification. It appears to be the latter, but this makes the example provided in the QA Sentence Distribution paragraph a bit confusing. The paper would benefit from showing some examples of the exact task sequences in each case.", "theoretical_claims": "I checked the proofs of Linear Growth and Decelerating Growth and did not notice any glaring errors.", "experimental_designs_or_analyses": "The experiments seem sound, but limited in scope.", "relation_to_broader_scientific_literature": "The paper relates to prior work identifying task vectors in the context of ICL, showing that demonstrations allow LLMs to internally construct a vector that expresses the task concept, as well as theoretical study of ICL.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\n\n* The authors provide a controlled and precise setting to study a model of task vectors in ICL\n* The theoretical conclusions are well-supported and there is some additional empirical evidence provided\n\nWeaknesses:\n\n* The study of geometric relationships seems to disappear when discussing the empirical results with the trained models\n* Broader impact may be limited, as it is unclear how well this transfers to the motivating settings of task vectors in pretrained LLMs", "other_comments_or_suggestions": "The authors discuss the geometric relationship between words and labels, and similar discuss the relationship between task vectors and task-related words in the context of pretrained LLMs. I would like to see some analysis of whether the same relationships exist in their simplified settings. When training on the QA-ICL dataset, do you notice similar relationships emerge?", "questions_for_authors": "Is there a way to vary and study the complexity of QA and Word-Label tasks? In the natural language settings you draw from for motivation, we have a clearer sense of complexity for different QA pairs and, similarly, a clearer sense of concept hierarchy.", "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": "tLUPTwGQ5g", "reviewer": "Reviewer_UbXD", "summary": "This work analyzes the optimization dynamics of transformer networks trained on in-context learning tasks via Gradient-descent on an L2-regularized cross-entropy loss. The analysis relies on a simple specific data-generating process which provides a way to formalize the notion of task concept vector. The results shows that the internal representation of a Transformer can recover the the task concept vector. Experiments validating the theory are performed.", "questions": "", "limitations": "", "rating": 2, "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 acknowledging our theoretical contribution and recognizing our analysis non-trivial. \n\n**Q**:Confusion over Theorem 1\n\n- We'd make it in the theorem statement explicitly by referring Algorithm 1.\n- As we claim for any ε>0, the varying item here is ε, with other parameters satisfying Condition 3.1 and the conditions in our statement. The existence of $t=O(g(ε,...))$ means for every (varying) ε, there exist a corresponding t upper bounded by $C\\log(1/ε)·(\\eta^{-1}q_{V}^{-1}σ_{1}^{2}d^{2}KM)$ for some constant $C$.\n\n**Q**:Room for presentation\n\nWe highly appreciate the feedback:\n\n- Eq.(1): We’d note $f = a^f + b$ as the label vector representation. Per [1], one task $f$ corresponds to one task vector $a^f$ in latent space, yielding $f(x_{\\text{query}})$ when added to $b(x_{\\text{query}})$, an encoded residual stream independent of demo-pairs in prompt $T$. Model $\\theta$ infers $a_θ^f$ from $T$, with $p_θ(a|T)$ reflecting model's confidence in recognizing $a$ as the task vector.\n- MLM: MLM loss refers Masked Language Model loss.\n- $w^i_k$: It should be replaced by $w_k$ with $i$ as a typo.\n- Line 189-201: M represent question length excluding query **x**. 201 explains our modeling intuition—QA includes irrelevant tokens, task vector, word, label—similar to [2] and modeled after [3].\n- Line 156-161 & Eq.(2): Every **x_l** - **y_l** pair in the prompt $T$ share a co-task $k_{T}\\in[K]$ but each pair can have its own low-level real label $y_{k_T,l}$ regarding $k_T$. The label vector of the prompt *only* depends on $k_T$ and the $y_{k_T,J+1}$ in query **x_{J+1}**. An illustrative example showing this intuition is $T$=[Japan,Tokyo,China,Beijing,France], **y_3**=Paris. The expected **y_3** only depends on the co-task “capital” and the semantic of France (not Japan or China). Under our modeling, **x_{J+1}** ≈$0.1a_{k_{T}}+y_{k_T,l}b_{k_{T}}$, and **y_{J+1}** ≈$a_{k_{T}}+y_{k_T, l}b_{k_{T}}$, modeled after Figure 1[1]. We'd avoid nested indices by let $l_x=l_y=l,l\\in[J]$.\n- Line 181: We’d replace $W_KT$ by $W_KT_l$.\n- Eq.(3): K’ is the number of irrelevant token in U, U is the dictionary matrix for cross-entropy training akin to [4] described in our 'Training Setups' in Section 2.2, 7 denotes K sets of {a±b, 0.1a±b, ±b, a} in eq.(14). \n\nWhile key details, such as Algorithm 1's procedure in Section 2.2, are included in the main text, we understand the importance of making our contributions more accessible—we'll add a notation table in a new appendix to summarize key definitions and intuitions for easier reference.\n\n**Q**:arXiv2410.23501\n\nThanks for sharing the paper—aligns with [5-6], backing up our data modeling, and we’ll cite it accordingly.\n\n**Q**:Merit of the problem & Real-world impact\n\nWe'd like to emphasize that our models, grounded in empirical observations, offer significant theoretical merit:\n\n- Modeled after [1] per Figure 1 as well as the concept latent geometry [5-6], our models are indeed sparse coding approaches suitable for capturing language polysemy [7-8] (see our 1st and 3rd responses to Reviewer mnid)—we successfully show the **OOD edge** of transformer over word2vec—addresses *Question 5.1.4* in [9] in theory.\n- Indeed, our empirically-grounded theoretical approach is common in theory (see our 1st response to Reviewer R7sM)—akin to [10-11] analyzing on feature-noise vision data, which is justified by the properties of ResNet’s latent space [10]. While [11] states the harmful overfitting over *vision* data is due to noise’s memorization, our harmful overfitting over ICL data is by falsely memorize co-occurrence of low-level features—akin to [12]'s result—feature co-occurrence biases gradients toward memorization.\n- We offer the **first** optimization theory for a *softmax-layernorm-residual* realistic transformer on empirical-motivated QA data trained by cross-entropy loss for factual recall ICL per [1], going beyond prior theory with idealized assumption like residual-free [13], QA-combined attention [14] with unrealistic loss functions (square or hinge loss) and oversimplified data. High nonlinearities of our problem induces complex gradient, for which we introduce *six continuous flows* in Appendix D.1 to capture its property in different phases, contributing to theory community as a method for treating complex optimization dynamics beyond [15].\n\n**Summary**: We thank the reviewer for acknowledging the theoretical questions tackled here interesting and our analysis non-trivial. We highly value the feedback, and should the reviewer have any further advice or wish to discuss any point further, we would be more than delighted to continue our productive exchange. Once again, we deeply appreciate the reviewer’s valuable time and comments.\n\n*Reference (arXiv identifier)*\n\n[1]2305.16130\n[2]2412.06538\n[3]2309.14316\n[4]2305.16380\n[5]2406.01506\n[6]2403.03867\n[7]1601.03764\n[8]2105.15134\n[9]2405.01964\n[10]2012.09816\n[11]2202.05928\n[12]2410.09605\n[13]2402.15607\n[14]2402.01258\n[15]2310.01975", "claims_and_evidence": "From what I understood of the paper, I believe the claims are supported.", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "I read the theoretical statements contained in the main paper. I list some of my confusions which may be due to the fact that I am not very familiar with learning theory and bounds holding with high probability and so maybe this is commonly used jargon.\n\nIn Theorem 3.2, at first I was confused since the statement does not explicitly mention which optimizer is used on which loss, but this is mentioned earlier in the text (not in an assumption environment nor anything like that). I think it should be more explicit in the theorem statement. \n\nI was  very confused by the use of the $O$ and $\\Omega$ notation in Theorem 3.2 and Prop 3.3. It's unclear to me what is varying/increasing here. Is it 1/epsilon? K? t? The representation dimensionality? This should be more explicit.\n\nAlso in Theorem 3.2, I was confused by a statement of the form \"there exists t = O(g(epsilon, L, M, ...))\". Are we saying that there exists a sequence t's that belongs to O(...) ? And again, what's varying?\n\nThese points are examples of the general lack of clarity in this work.", "experimental_designs_or_analyses": "No.", "relation_to_broader_scientific_literature": "I am not well versed in this literature.", "essential_references_not_discussed": "The authors might be interested in https://arxiv.org/abs/2410.23501.", "other_strengths_and_weaknesses": "Strengths:\n\n- I think the theoretical questions tackled here are interesting.\n- The analysis seems non-trivial. That being said, I do not think I am qualified to judge its novelty and relevance.\n\nWeakness:\n\n- Clarity is a big issue (see comments about theory). I provide some examples below:\n    - Eq (1) is weird… $a^f_\\theta$ is a deterministic function of T? Why is there a distribution p(a | T)? Same question for b… Also, are we saying a + b = f here? If so it should be written explicitly.\n    - Line 74 (left): What does the acronym MLM stand for?\n    - w^i_k is introduced without explaining what i is. Is there multiple low-level binary concepts for a given high-level concept? Is this what i is supposed to mean?\n    - Line 161 (right): Notation is unnecessarily heavy, see for example $a_{k_{l_y}}$ which has too many nested indices. This makes everything difficult to parse…\n    - Line 156 (right): I don’t understand, an expression is given for the expected label of the query but that expression contains a random element $y_{k_T,J+1}$. How is that possible? Also, I computed the expected value of y_j with j<=J and I don’t get the same expectation as for j = J+1. That’s a bit weird no? I was expecting the distribution of y given x to be the same across all examples j in the task T. Looking at (12) in the appendix didn’t help me understand.\n    - Line 189: The number $M$ appears for the first time, without explanations. \n    - Line 201: This illustrative example is unclear. \n    - Line 181 (right): Something’s off with that equation. You are multiplying a vector with a vector, and it yields a vector. I think the softmax should be indexed by $\\ell$.\n    - I don’t understand how you obtain Eq. (2)\n    - Eq. 3, what’s K’? Where does this \"7’’ come from? What is the matrix U? It seems many important details are left in the appendix. The contribution should be understandable from the main text. Another example is the algorithm shown only in Appendix C.\n- As mentioned above, many important details necessary to understand the contribution seem to be hidden in the Appendix.\n- I not sure the data-generating process is realistic or interesting. I understand that this kind of analysis is challenging without simplifying assumptions, but overall I believe the data-generating process could be motivated further. That being said, it is possible I did not understand it properly due to clarity issues.", "other_comments_or_suggestions": "-", "questions_for_authors": "See clarity issues above.", "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 presents a theoretical framework explaining how large language models perform in-context learning (ICL) for factual recall through task vector arithmetic, grounded in hierarchical concept modeling and optimization theory. The authors show that nonlinear residual transformers trained via gradient descent can generalize well under data shifts and concept recombinations, offering both provable guarantees and empirical validation. 

The motivation of the paper is interesting and the analysis is good with solid theory support. The reviewers have a few concerns about the clarity and accessibility of the theoretical proofs, particularly equations (9) and (11), suggesting that simplified derivations and stronger intuitive explanations in the main text would help, and they also raise issues regarding the empirical validation of theoretical claims, the limited discussion of geometric relationships in trained models, and the potential narrow scope of broader impact. During the rebuttal, the authors made an effort to address these concerns, and the reviewers were generally satisfied with their responses.
Decision: Accept (poster)