Provable In-Context Vector Arithmetic via Retrieving Task Concepts

Thank you for acknowledging our theoretical contribution and recognizing our analysis non-trivial.

Q:Confusion over Theorem 1

• We'd make it in the theorem statement explicitly by referring Algorithm 1.
• 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$.

Q:Room for presentation

We highly appreciate the feedback:

• 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.
• MLM: MLM loss refers Masked Language Model loss.
• $w^i_k$: It should be replaced by $w_k$ with $i$ as a typo.
• 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].
• 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]$.
• Line 181: We’d replace $W_KT$ by $W_KT_l$.
• 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).

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

Q:arXiv2410.23501

Thanks for sharing the paper—aligns with [5-6], backing up our data modeling, and we’ll cite it accordingly.

Q:Merit of the problem & Real-world impact

We'd like to emphasize that our models, grounded in empirical observations, offer significant theoretical merit:

• 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.
• 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.
• 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].

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.

Reference (arXiv identifier)

[1]2305.16130 [2]2412.06538 [3]2309.14316 [4]2305.16380 [5]2406.01506 [6]2403.03867 [7]1601.03764 [8]2105.15134 [9]2405.01964 [10]2012.09816 [11]2202.05928 [12]2410.09605 [13]2402.15607 [14]2402.01258 [15]2310.01975

Thank you for acknowledging our theory as empirically-supported and sound. We appreciate your professional review and address your concerns below.

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.

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

Q:The study of geometric relationships of the trained models seems to disappear

A:We would like to remark:

• 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].
• 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.

Q:Broader impact may be limited, as it is unclear how well this transfers to the motivating settings of task vectors in pretrained LLMs

A:We would like to remark our impact on elucidating observed LLM phenomena:

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

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.

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.

Should 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!

Reference (arXiv identifier)

[1]2412.06538 [2]2309.14316 [3]1601.03764 [4]2105.15134 [5]2402.15607 [6]2402.01258 [7]2406.01506 [8]2403.03867 [9]2410.09605 [10]2305.16130 [11]2012.09816 [12]2202.05928 [13]2212.04089 [14]forum?id=vRvVVb0NAz

Thank you for your thoughtful review!

Q:Chicken-and-Egg Dilemma & real-world impact

We'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.

1. 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.
2. 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].
3. 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.

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

Q:Experimental details

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

Q:Room for presentation

We highly appreciate the feedback and would add:

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

Thanks again for your feedback! We welcome further discussion and appreciate your time!

Reference (arXiv identifier)

[1]2403.03867 [2]2305.16130 [3]2012.09816 [4]2202.05928 [5]iclr.cc/virtual/2023/poster/11533 [6]2405.01964 [7]2001.04413 [8]2310.01975 [9]2410.09605 [10]2412.06538 [11]2309.14316