Research map / Architectures and efficiency / Attention and transformers
Provable In-Context Vector Arithmetic via Retrieving Task Concepts
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Provides an optimization theory showing that a nonlinear residual transformer trained by gradient descent on QA data learns task-vector arithmetic for single-token factual-recall ICL, while ICL-style training fails to retrieve clean task vectors.
Problem
How a nonlinear residual transformer trained with gradient descent and cross-entropy on QA data can perform factual-recall in-context learning via task-vector arithmetic (akin to Word2Vec), and why word-label ICL training fails to do so.
Approach
Models hierarchical concepts (orthogonal high-level task vectors plus low-level label features) and a one-layer softmax attention + MLP with residual and layer normalization; proves that QA-style training makes the value/key-query projections retrieve the task vector, yielding vanishing 0-1 test loss, whereas word-label ICL training memorizes low-level features and keeps constant test error.
Key results
- QA-trained transformer's test loss converges to zero (below any epsilon), while word-label ICL-trained and QA-ICL-trained models keep a constant test error of about 0.2 in simulations.
- Theorem shows that with ICL-style training the hidden output is a hybrid of task vector and low-level features, with test loss Theta(1), under the stated conditions.
- OOD results: the QA-trained model generalizes to unseen dictionaries whose task vectors lie in the conic hull of training task vectors and to prompts with up to 3 co-task concepts.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train on QA sentences instead of Word-Label ICL prompts | Nonlinear residual transformer · Word-Label ICL test prompts | test error: 0.2 -> zero | Helped |
| Train on QA sentences instead of QA-ICL prompts | Nonlinear residual transformer · QA-ICL test prompts | test error: 0.2 -> zero | Helped |
| Train on Word-Label ICL prompts rather than QA sentences | Nonlinear residual transformer · Word-Label ICL test prompts | test error: zero -> 0.2 | Hurt |
| Train on QA-ICL prompts rather than QA sentences | Nonlinear residual transformer · QA-ICL test prompts | test error: zero -> 0.2 | Hurt |
Surprising findings
- Word-label ICL training fails to yield a clean task vector and instead memorizes low-level label features, so test error stays constant even with longer training.
- Harmful overfitting here stems from low-level feature memorization rather than noise memorization as in vision models.
Limitations
- Limited to single-token factual-recall settings; does not cover multi-token reasoning.
- Idealized hierarchical concept model; real task vectors are not exactly equal to concept vectors and dictionaries are polysemous and evolve during training.
- Experiments are small-scale synthetic simulations rather than real LLMs.
Experimental setup
Datasets: synthetic hierarchical concept word-label ICL data synthetic QA sentence data
Benchmarks: synthetic ICL test distributions synthetic QA-ICL test distributions
Metrics: 0-1 test loss cosine similarity of hidden vector with task vector training/test loss curves
Baselines: Word-label ICL-trained transformer Word2Vec
Builds on: Merullo et al. 2024 (vector arithmetic in LLMs) Hierarchical linear concept geometry (Park et al.) Trained transformer ICL theory (Li et al.)
Closest papers by meaning
- Emergence and Effectiveness of Task Vectors in In-Context Learning: An Encoder Decoder Perspective (ICML 2025 · similarity 0.89)
- Distributional Alignment as a Criterion for Designing Task Vectors in In-Context Learning (NeurIPS 2026 · similarity 0.89)
- Vector-ICL: In-context Learning with Continuous Vector Representations (ICLR 2025 · similarity 0.88)
- Provably Transformers Harness Multi-Concept Word Semantics for Efficient In-Context Learning (NeurIPS 2024 · similarity 0.88)
- In-Context Learning with Representations: Contextual Generalization of Trained Transformers (NeurIPS 2024 · similarity 0.88)
- Technical Debt in In-Context Learning: Diminishing Efficiency in Long Context (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.