Research map / Architectures and efficiency / Attention and transformers
Unifying Attention Heads and Task Vectors via Hidden State Geometry in In-Context Learning
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.
Proposes a geometric framework that measures separability and alignment of query hidden states across layers to explain ICL in classification, linking Previous Token Heads (separability) and Induction Heads (alignment) to task vectors.
Problem
Prior ICL mechanistic work studies attention heads or task vectors through end-task ablations, without a unified account of how they shape query hidden states layer by layer to produce the output.
Approach
Shows theoretically that classification accuracy is upper-bounded by maximum separability of query hidden states, achieved when hidden states align with label unembedding vectors (output alignment and directional alignment); then measures these layer-wise and attributes them to PTHs and IHs via ablation and task-vector injection.
Key results
- ICL improves classification accuracy over zero-shot mainly through alignment; separability gap between ICL and zero-shot is small despite ~80% accuracy difference (Llama2-70B).
- Ablating top 10% Previous Token Heads substantially reduces separability, while ablating Induction Heads severely disrupts alignment measures and output accuracy; random ablation has negligible effect.
- Injecting IH-derived task vectors into zero-shot hidden states restores the ICL-specific spike in alignment; rank-10 denoising of hidden states yields accuracy gains exceeding 10% in early post-transition layers.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ICL demonstrations instead of zero-shot prompting | Llama2-70B · Classification datasets | classification accuracy | Helped |
| Increase the number of demonstrations from 4 to 24 | ICL · Classification datasets | classification accuracy | Helped |
| Select demonstrations using kNN retrieval | ICL · Classification datasets | classification accuracy | Helped |
| Replace demonstration labels with uninformative symbols | ICL · Classification datasets | classification accuracy | Hurt |
| Apply rank-10 SVD denoising to hidden states | Llama2-70B · Classification datasets, intermediate post-transition layers | output accuracy: exceeding 10% | Helped |
| Train a logistic classifier on last-layer hidden states | Pretrained models · Classification datasets, zero-shot | accuracy: over 70% | Helped |
| Train a logistic classifier on ICL hidden states | Pretrained models · Classification datasets, 8-shot ICL | accuracy: ~10% | Helped |
| Apply rank-5 denoising before unembedding | Pretrained models · Classification datasets, final layer | classification accuracy | Mixed |
| Remove the top 10% of PTH outputs (ablation) | Llama2-70B · Classification datasets | separability score | Hurt |
| Remove the top 10% of PTH outputs (ablation) | Llama2-70B · Classification datasets | composite alignment | Hurt |
| Remove the top 10% of IH outputs (ablation) | Llama2-70B · Classification datasets | alignment measures | Hurt |
| Remove the top 10% of IH outputs (ablation) | Llama2-70B · Classification datasets | separability score | No clear effect |
Surprising findings
- Separability of hidden states far exceeds the accuracy of direct label decoding, so a logistic classifier on last-layer states improves zero-shot accuracy by over 70% across models.
- In final layers, label-related directions are filtered out and label-irrelevant semantics are reintroduced, slightly reducing output alignment.
Limitations
- Generation-task analysis is preliminary.
- Analysis is limited to inference-time behavior in pretrained models; training dynamics of the phase transition are not studied.
- Large models (>10B) quantized to 4-bit.
Experimental setup
Datasets: SST-2 SUBJ TREC SNLI RTE CB Famous People
Benchmarks: SST-2 SUBJ TREC SNLI RTE CB
Metrics: accuracy separability score output alignment (logit lens accuracy) singular alignment variance-based alignment mean-based alignment composite alignment effective dimension
Baselines: zero-shot prompting random head ablation
Builds on: induction heads task vectors logit lens
Closest papers by meaning
- Capturing In-Context Learning Dynamics with Task Operators (NeurIPS 2026 · similarity 0.90)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.90)
- Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability (ICLR 2024 · similarity 0.90)
- Which Attention Heads Matter for In-Context Learning? (ICML 2025 · similarity 0.89)
- Emergence and Effectiveness of Task Vectors in In-Context Learning: An Encoder Decoder Perspective (ICML 2025 · similarity 0.89)
- The mechanistic basis of data dependence and abrupt learning in an in-context classification task (ICLR 2024 · similarity 0.89)
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.