Provable In-Context Vector Arithmetic via Retrieving Task Concepts

• URL: http://arxiv.org/abs/2508.09820v1
• Date: Wed, 13 Aug 2025 13:54:44 GMT
• Title: 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,
• Abstract summary: We show how nonlinear residual transformers trained via gradient descent on cross-entropy loss perform factual-recall ICL tasks via vector arithmetic.<n>These results elucidate the advantages of transformers over static embedding predecessors.
• Score: 53.685764040547625
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• On Multi-Step Theorem Prediction via Non-Parametric Structural Priors [50.16583672681106]
  In this work, we explore training-free theorem prediction through the lens of in-context learning (ICL)<n>We propose Theorem Precedence Graphs, which encode temporal dependencies from historical solution traces as directed graphs, and impose explicit topological constraints that effectively prune the search space during inference.<n>Experiments on the FormalGeo7k benchmark show that our method achieves 89.29% accuracy, substantially outperforming ICL baselines and matching state-of-the-art supervised models.
  arXiv  Detail & Related papers  (2026-03-05T06:08:50Z)
• Outcome-Based RL Provably Leads Transformers to Reason, but Only With the Right Data [4.344634631420729]
  We analyze the policy gradient dynamics of single-layer Transformers trained via Reinforcement Learning.<n>We prove that despite training solely on final-answer correctness, policy gradient drives the Transformer to converge to a structured, interpretable algorithm.
  arXiv  Detail & Related papers  (2026-01-21T16:36:19Z)
• Provable Benefit of Curriculum in Transformer Tree-Reasoning Post-Training [76.12556589212666]
  We show that curriculum post-training avoids the exponential complexity bottleneck.<n>Under outcome-only reward signals, reinforcement learning finetuning achieves high accuracy with sample complexity.<n>We establish guarantees for test-time scaling, where curriculum-aware querying reduces both reward oracle calls and sampling cost from exponential to order.
  arXiv  Detail & Related papers  (2025-11-10T18:29:54Z)
• Understanding Task Vectors in In-Context Learning: Emergence, Functionality, and Limitations [19.539276425108987]
  This work proposes the Linear Combination Conjecture, positing that task vectors act as single in-context demonstrations formed through linear combinations of the original ones.<n>We show that task vectors naturally emerge in linear transformers trained on triplet-formatted prompts through loss landscape analysis.<n>We predict the failure of task vectors on representing high-rank mappings and confirm this on practical LLMs.
  arXiv  Detail & Related papers  (2025-06-10T17:59:31Z)
• Illusion or Algorithm? Investigating Memorization, Emergence, and Symbolic Processing in In-Context Learning [50.53703102032562]
  Large-scale Transformer language models (LMs) trained solely on next-token prediction with web-scale data can solve a wide range of tasks.<n>The mechanism behind this capability, known as in-context learning (ICL), remains both controversial and poorly understood.
  arXiv  Detail & Related papers  (2025-05-16T08:50:42Z)
• Argumentation Computation with Large Language Models : A Benchmark Study [6.0682923348298194]
  Large language models (LLMs) have made significant advancements in neuro-symbolic computing.<n>We aim to investigate the capability of LLMs in determining the extensions of various abstract argumentation semantics.
  arXiv  Detail & Related papers  (2024-12-21T18:23:06Z)
• Rethinking Associative Memory Mechanism in Induction Head [37.93644115914534]
  This paper investigates how a two-layer transformer thoroughly captures in-context information and balances it with pretrained bigram knowledge in next token prediction.<n>We theoretically analyze the representation of weight matrices in attention layers and the resulting logits when a transformer is given prompts generated by a bigram model.
  arXiv  Detail & Related papers  (2024-12-16T05:33:05Z)
• Provably Transformers Harness Multi-Concept Word Semantics for Efficient In-Context Learning [53.685764040547625]
  Transformer-based large language models (LLMs) have displayed remarkable creative prowess and emergence capabilities. This work provides a fine mathematical analysis to show how transformers leverage the multi-concept semantics of words to enable powerful ICL and excellent out-of-distribution ICL abilities.
  arXiv  Detail & Related papers  (2024-11-04T15:54:32Z)
• Interpreting Affine Recurrence Learning in GPT-style Transformers [54.01174470722201]
  In-context learning allows GPT-style transformers to generalize during inference without modifying their weights. This paper focuses specifically on their ability to learn and predict affine recurrences as an ICL task. We analyze the model's internal operations using both empirical and theoretical approaches.
  arXiv  Detail & Related papers  (2024-10-22T21:30:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.