pith. sign in

arxiv: 2605.16591 · v2 · pith:CUTBSQNVnew · submitted 2026-05-15 · 💻 cs.LG · cs.AI

How Few-Shot Examples Add Up: A Causal Decomposition of Function Vectors in In-Context Learning

Entang Wang , Yiwei Wang , Aleksandra Bakalova , Michael Hahn This is my paper

Pith reviewed 2026-06-30 19:04 UTC · model grok-4.3

classification 💻 cs.LG cs.AI
keywords in-context learningfunction vectorsfew-shot promptingattention mechanismscausal interventionstransformer modelsprompt composition
0
0 comments X

The pith

An n-shot function vector is well-approximated by a linear combination of example-level sub-vectors, with attention reweighting that favors informative demonstrations.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper establishes that the direction in activation space responsible for in-context task performance can be expressed as the sum of separate contributions from each demonstration. It further shows that the model does not treat these contributions independently but adjusts their relative strength according to the full set of examples seen so far. The adjustment occurs mainly through improved alignment between queries and keys in the attention layers rather than through changes to the values being written. This account explains both the reliable additive structure and the context-sensitive selection that together determine how few-shot prompts produce task behavior.

Core claim

Across tasks and models, an n-shot FV is well-approximated by a linear combination of example-level sub-FVs, suggesting additive and composable contributions from individual demonstrations. Beyond additivity, models contextualize individual examples' representations based on prior examples to adaptively reweight which demonstrations dominate the FV: attention shifts toward examples that are more informative and less ambiguous under the context. A causal decomposition separates Query-Key routing from Value updates, finding that contextualization's most consistent contributions to FV quality arise from Query-Key alignment—particularly in ambiguous settings—while Value-mediated effects are more

What carries the argument

The function vector (FV), a causal activation direction that drives task behavior on the ICL query, together with its decomposition into per-example sub-FVs and the attention-based reweighting that modulates their sum.

If this is right

  • Individual demonstrations contribute additively to the overall task direction extracted by the model.
  • Attention reweighting allows the model to emphasize clearer or more diagnostic examples once the full prompt context is available.
  • Query-key alignment supplies the most reliable improvement to FV quality when examples are ambiguous.
  • Value updates produce more variable and task-dependent changes to the resulting function vector.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • Prompt construction strategies could prioritize example clarity and ordering to maximize the additive signal rather than relying on emergent synergies.
  • The same decomposition could be tested on non-classification tasks to check whether additivity persists when the output space is structured differently.
  • If the linear approximation holds, one could predict the quality of an n-shot prompt from measurements on its constituent single-example prompts alone.
  • The separation of query-key and value effects offers a concrete target for interventions that aim to improve few-shot performance without retraining.

Load-bearing premise

The function vector identified by causal interventions remains a stable task-driving direction whose linear decomposition and attention effects hold across tasks, models, and prompt formats without large higher-order interactions.

What would settle it

A systematic measurement across many tasks showing that the residual norm between the observed n-shot FV and the linear sum of the individual example sub-FVs stays large would falsify the additivity claim.

Figures

Figures reproduced from arXiv: 2605.16591 by Aleksandra Bakalova, Entang Wang, Michael Hahn, Yiwei Wang.

Figure 1
Figure 1. Figure 1: Overview of FV formation in in-context learning. (a) Linear superposition: the n-shot function vector (FV) is approximated as a weighted sum of example-level sub-FVs. (b) Attention reweighting: contextualization modulates attention over demonstrations, changing the weights assigned to different sub-FVs. (c) Geometric refinement: contextualization further improves FV quality by altering the FV direction via… view at source ↗
Figure 2
Figure 2. Figure 2: Left: mean cosine similarity between the observed FV and the OLS reconstruction with/without contextualization (ctx/unc). Right: mean R 2 of the same fit. See Appendix F [PITH_FULL_IMAGE:figures/full_fig_p004_2.png] view at source ↗
Figure 4
Figure 4. Figure 4: A demonstration of an ambiguous prompt. The ambigu￾ous example is compatible with both character transformation and identity mapping. 4.1. Unambiguous examples are intrinsically more attractive Attention to examples is governed by Query–Key similarity: for a fixed last-token Query, higher attention to examples implies stronger QK alignment with their Keys. To isolate the role of per-example Keys, we first … view at source ↗
Figure 5
Figure 5. Figure 5: (a) Mean attention weight of FV heads on a normal task (COUNTRY–CAPITAL), attention is largely explained by recency bias. (b) Mean attention weight of FV heads on an ambiguous task (CHINESE AMBIGUOUS), FV heads consistently upweight the unambiguous demonstrations (here fixed at Ex2 and Ex4, marked by diagonal hatching) relative to ambiguous ones. Each experiment setting is averaged over 100 5-shot prompts.… view at source ↗
Figure 6
Figure 6. Figure 6: Mean attention weight of FV heads under symmetric Key patching on all attention heads on PRESENT–PAST AMBIGUOUS dataset. Unambiguous examples are fixed at Ex2 and Ex4, marked by diagonal hatching. Each experiment setting is averaged over 100 5-shot prompts. Error bars: 95% CIs from 1000 bootstrap resamples. See Appendix G [PITH_FULL_IMAGE:figures/full_fig_p006_6.png] view at source ↗
Figure 7
Figure 7. Figure 7: Mean attention weight of FV heads in datasets under uncontextualized and contextualized settings. (a) On a normal task (PARK–COUNTRY), contextualization primarily mitigates recency bias: attention becomes less back-heavy and the centroid shifts toward earlier demonstrations, without inducing strong example￾specific peaks.(b) On an ambiguous task (PRESENT–PAST AM￾BIGUOUS), contextualization instead sharpens… view at source ↗
Figure 8
Figure 8. Figure 8: Shapley value decomposition (ϕ) averaged over experiment configurations (n-shots, positional controls) on gemma-2 and Llama-3 model families. Contextualizing QK usually has a stronger positive effect on FV quality, compared to contextualizing V. See Appendix I [PITH_FULL_IMAGE:figures/full_fig_p008_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: Query composition and its effect on QK alignment for ambiguous tasks. Panels (a, c) report the FV injection accuracy, while (b, d) show the attention proportion on unambiguous vs. ambiguous examples. The value T above each bar indicates the total attention mass allocated to all example segments. Results are averaged over 100 trials using 10-shot prompts. See Appendix J Fig.52-57 for the complete results. s… view at source ↗
Figure 10
Figure 10. Figure 10: A diagram of uncontextualized ablation. This diagram illustrates the intervention used to isolate ICL component-level contributions by severing specific attention pathways. We keep all edges within each example and all edges from prompt components to the last token tn+1 while zeroing out the rest of the edges. F. Linear superposition To quantify the extent to which an n-shot FV can be explained as a linea… view at source ↗
Figure 11
Figure 11. Figure 11: Full Linear Superposition Results. For each model, we display: (Left) the mean cosine similarity between the observed Function Vector (FV) and the OLS reconstruction; (Right) the mean R 2 of the fit across all layers. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p019_11.png] view at source ↗
Figure 12
Figure 12. Figure 12: Diagram of the Key patching process in Section 4.1. (a) Patching ambiguous Key to unambiguous Key. (b) Patching unambiguous Key to ambiguous Key. 21 [PITH_FULL_IMAGE:figures/full_fig_p021_12.png] view at source ↗
Figure 13
Figure 13. Figure 13: Normal Tasks: gemma-2-2b mean attention weights of individual examples on FV heads across 3/5/10-shot settings. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p022_13.png] view at source ↗
Figure 14
Figure 14. Figure 14 [PITH_FULL_IMAGE:figures/full_fig_p022_14.png] view at source ↗
Figure 15
Figure 15. Figure 15: Normal Tasks: gemma-2-27b mean attention weights of individual examples on FV heads across 3/5/10-shot settings. 3-Shot 5-Shot 10-Shot COUNTRY–CAPITAL Ex1 Ex2 Ex3 Position of examples in prompt 0.0 0.1 0.2 0.3 0.4 0.5 Mean attention weight 0.24 0.28 0.37 Ex1 Ex2 Ex3 Ex4 Ex5 Position of examples in prompt 0.0 0.1 0.2 0.3 0.4 0.5 Mean attention weight 0.15 0.17 0.17 0.20 0.22 Ex1 Ex2 Ex3 Ex4 Ex5 Ex6 Ex7 Ex8… view at source ↗
Figure 16
Figure 16. Figure 16: Normal Tasks: Llama-3.2-1B mean attention weights of individual examples on FV heads across 3/5/10-shot settings. 23 [PITH_FULL_IMAGE:figures/full_fig_p023_16.png] view at source ↗
Figure 17
Figure 17. Figure 17: Normal Tasks: Llama-3.2-3B mean attention weights of individual examples on FV heads across 3/5/10-shot settings. 3-Shot 5-Shot 10-Shot COUNTRY–CAPITAL Ex1 Ex2 Ex3 Position of examples in prompt 0.0 0.1 0.2 0.3 0.4 0.5 Mean attention weight 0.11 0.14 0.20 Ex1 Ex2 Ex3 Ex4 Ex5 Position of examples in prompt 0.0 0.1 0.2 0.3 0.4 0.5 Mean attention weight 0.09 0.10 0.11 0.13 0.14 Ex1 Ex2 Ex3 Ex4 Ex5 Ex6 Ex7 Ex… view at source ↗
Figure 18
Figure 18. Figure 18: Normal Tasks: Llama-3.1-8B-Instruct mean attention weights of individual examples on FV heads across 3/5/10-shot settings. 24 [PITH_FULL_IMAGE:figures/full_fig_p024_18.png] view at source ↗
Figure 19
Figure 19. Figure 19: Ambiguous Tasks: gemma-2-2b mean attention weights of individual examples on FV heads across 3/5/10-shot settings. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p025_19.png] view at source ↗
Figure 20
Figure 20. Figure 20: Ambiguous Tasks: gemma-2-9b mean attention weights of individual examples on FV heads across 3/5/10-shot settings. 26 [PITH_FULL_IMAGE:figures/full_fig_p026_20.png] view at source ↗
Figure 21
Figure 21. Figure 21: Ambiguous Tasks: gemma-2-27b mean attention weights of individual examples on FV heads across 3/5/10-shot settings. 27 [PITH_FULL_IMAGE:figures/full_fig_p027_21.png] view at source ↗
Figure 22
Figure 22. Figure 22: Ambiguous Tasks: Llama-3.2-1B mean attention weights of individual examples on FV heads across 3/5/10-shot settings. 28 [PITH_FULL_IMAGE:figures/full_fig_p028_22.png] view at source ↗
Figure 23
Figure 23. Figure 23: Ambiguous Tasks: Llama-3.2-3B mean attention weights of individual examples on FV heads across 3/5/10-shot settings. 29 [PITH_FULL_IMAGE:figures/full_fig_p029_23.png] view at source ↗
Figure 24
Figure 24. Figure 24: Ambiguous Tasks: Llama-3.1-8B-Instruct mean attention weights of individual examples on FV heads across 3/5/10-shot settings. 30 [PITH_FULL_IMAGE:figures/full_fig_p030_24.png] view at source ↗
Figure 25
Figure 25. Figure 25: Key patching experiment results of all ambiguous tasks on gemma-2-2b. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p031_25.png] view at source ↗
Figure 26
Figure 26. Figure 26: Normal Tasks: gemma-2-2b Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p032_26.png] view at source ↗
Figure 27
Figure 27. Figure 27: Normal Tasks: gemma-2-9b Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p033_27.png] view at source ↗
Figure 28
Figure 28. Figure 28: Normal Tasks: gemma-2-27b Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p033_28.png] view at source ↗
Figure 29
Figure 29. Figure 29: Normal Tasks: Llama-3.2-1B Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p034_29.png] view at source ↗
Figure 30
Figure 30. Figure 30: Normal Tasks: Llama-3.2-3B Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p035_30.png] view at source ↗
Figure 31
Figure 31. Figure 31: Normal Tasks: Llama-3.1-8B-Instruct Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p035_31.png] view at source ↗
Figure 32
Figure 32. Figure 32: Ambiguous Tasks: gemma-2-2b Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p036_32.png] view at source ↗
Figure 33
Figure 33. Figure 33: Ambiguous Tasks: gemma-2-9b Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p037_33.png] view at source ↗
Figure 34
Figure 34. Figure 34: Ambiguous Tasks: gemma-2-27b Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p038_34.png] view at source ↗
Figure 35
Figure 35. Figure 35: Ambiguous Tasks: Llama-3.2-1B Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p039_35.png] view at source ↗
Figure 36
Figure 36. Figure 36: Ambiguous Tasks: Llama-3.2-3B Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p040_36.png] view at source ↗
Figure 37
Figure 37. Figure 37: Ambiguous Tasks: Llama-3.1-8B-Instruct Influence of contextualization on mean attention weights of individual examples on FV heads. Each row displays the mean attention for a normal task across 3-shot, 5-shot, and 10-shot settings before and after contextualization. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p041_37.png] view at source ↗
Figure 38
Figure 38. Figure 38: Shared PCA visualization of Query and Key vectors extracted from the Top-1 FV head on gemma-2-2b. (a) Normal tasks. (b) Ambiguous tasks. For each panel, each category of Q/K vectors is randomly sampled from 100 different prompts. The Q vectors correspond to the Query of the final (prediction) token in the prompt. The K vectors correspond to the Key of the single token within each example that receives the… view at source ↗
Figure 39
Figure 39. Figure 39: Diagram of the Value patching process in Section 5. (a) Patching V under uncontextualized ablation, replacing uncontex￾tualized V with contextualized V, keeping Q/K fixed. (b) Patching V under contextualized ablation, replacing contextualized V with uncontextualized V, keeping Q/K fixed. Both patching experiments are done on all attention heads. In the main text (Sec. 5), we report aggregated Shapley tren… view at source ↗
Figure 40
Figure 40. Figure 40: Causal Decomposition of Contextualization Gains on gemma-2-2b. (Normal Tasks) For each task and shot count, the left plot shows absolute FV injection accuracy F(QK, V) across four factorial intervention settings: (1) Uncontextualized F(0, 0); (2) QK-contextualized F(1, 0); (3) V-contextualized F(0, 1); (4) Full contextualized F(1, 1). We plot the marginal effects: QK@Vunc := F(1, 0)−F(0, 0), QK@Vctx := F(… view at source ↗
Figure 41
Figure 41. Figure 41: Causal Decomposition of Contextualization Gains on gemma-2-9b. (Normal Tasks) This is the detailed result of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p045_41.png] view at source ↗
Figure 42
Figure 42. Figure 42: Causal Decomposition of Contextualization Gains on gemma-2-27b. (Normal Tasks) This is the detailed result of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p045_42.png] view at source ↗
Figure 43
Figure 43. Figure 43: Causal Decomposition of Contextualization Gains on Llama-3.2-1B. (Normal Tasks) This is the detailed result of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p046_43.png] view at source ↗
Figure 44
Figure 44. Figure 44: Causal Decomposition of Contextualization Gains on Llama-3.2-3B. (Normal Tasks) This is the detailed result of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p046_44.png] view at source ↗
Figure 45
Figure 45. Figure 45: Causal Decomposition of Contextualization Gains on Llama-3.1-8B-Instruct. (Normal Tasks) This is the detailed result of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p047_45.png] view at source ↗
Figure 46
Figure 46. Figure 46: Causal Decomposition of Contextualization Gains on gemma-2-2b. (Ambiguous Tasks) For each task and shot count, the left plot shows absolute FV injection accuracy F(QK, V) across four factorial intervention settings: (1) Uncontextualized F(0, 0); (2) QK-contextualized F(1, 0); (3) V-contextualized F(0, 1); (4) Full contextualized F(1, 1). We plot the marginal effects: QK@Vunc := F(1, 0)−F(0, 0), QK@Vctx :=… view at source ↗
Figure 47
Figure 47. Figure 47: Causal Decomposition of Contextualization Gains on gemma-2-9b. (Ambiguous Tasks) This is the detailed result of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p049_47.png] view at source ↗
Figure 48
Figure 48. Figure 48: Causal Decomposition of Contextualization Gains on gemma-2-27b. (Ambiguous Tasks) This is the detailed result of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p050_48.png] view at source ↗
Figure 49
Figure 49. Figure 49: Causal Decomposition of Contextualization Gains on Llama-3.2-1B. (Ambiguous Tasks) This is the detailed result of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p051_49.png] view at source ↗
Figure 50
Figure 50. Figure 50: Causal Decomposition of Contextualization Gains on Llama-3.2-3B. (Ambiguous Tasks) This is the detailed result of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p052_50.png] view at source ↗
Figure 51
Figure 51. Figure 51: Causal Decomposition of Contextualization Gains on Llama-3.1-8B-Instruct. (Ambiguous Tasks) This is the detailed result of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p053_51.png] view at source ↗
Figure 52
Figure 52. Figure 52: Query-patching FV Injection Accuracy (AccQ) and Attention Proportion for Ambiguous Tasks on gemma-2-2b. Each shot compares the Query-patching FV injection accuracy (left) with the attention mass allocation (right). This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p054_52.png] view at source ↗
Figure 53
Figure 53. Figure 53: Query-patching FV Injection Accuracy (AccQ) and Attention Proportion for Ambiguous Tasks on gemma-2-9b. Each shot compares the Query-patching FV injection accuracy (left) with the attention mass allocation (right). This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p055_53.png] view at source ↗
Figure 54
Figure 54. Figure 54: Query-patching FV Injection Accuracy (AccQ) and Attention Proportion for Ambiguous Tasks on gemma-2-27b. Each shot compares the Query-patching FV injection accuracy (left) with the attention mass allocation (right). This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p056_54.png] view at source ↗
Figure 55
Figure 55. Figure 55: Query-patching FV Injection Accuracy (AccQ) and Attention Proportion for Ambiguous Tasks on Llama-3.2-1B. Each shot compares the Query-patching FV injection accuracy (left) with the attention mass allocation (right). This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p057_55.png] view at source ↗
Figure 56
Figure 56. Figure 56: Query-patching FV Injection Accuracy (AccQ) and Attention Proportion for Ambiguous Tasks on Llama-3.2-3B. Each shot compares the Query-patching FV injection accuracy (left) with the attention mass allocation (right). This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p058_56.png] view at source ↗
Figure 57
Figure 57. Figure 57: Query-patching FV Injection Accuracy (AccQ) and Attention Proportion for Ambiguous Tasks on Llama-3.1-8B-Instruct. Each shot compares the Query-patching FV injection accuracy (left) with the attention mass allo￾cation (right). This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p059_57.png] view at source ↗
Figure 58
Figure 58. Figure 58: Query-patching FV Injection Accuracy (AccQ) for Normal Tasks on gemma-2-2b across few-shot settings. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p060_58.png] view at source ↗
Figure 59
Figure 59. Figure 59: Query-patching FV Injection Accuracy (AccQ) for Normal Tasks on gemma-2-9b across few-shot settings. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p060_59.png] view at source ↗
Figure 60
Figure 60. Figure 60: Query-patching FV Injection Accuracy (AccQ) for Normal Tasks on gemma-2-27b across few-shot settings. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p061_60.png] view at source ↗
Figure 61
Figure 61. Figure 61: Query-patching FV Injection Accuracy (AccQ) for Normal Tasks on Llama-3.2-1B across few-shot settings. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p061_61.png] view at source ↗
Figure 62
Figure 62. Figure 62: Query-patching FV Injection Accuracy (AccQ) for Normal Tasks on Llama-3.2-3B across few-shot settings. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p062_62.png] view at source ↗
Figure 63
Figure 63. Figure 63: Query-patching FV Injection Accuracy (AccQ) for Normal Tasks on Llama-3.1-8B-Instruct across few-shot settings. This is an extension of Main Paper [PITH_FULL_IMAGE:figures/full_fig_p062_63.png] view at source ↗
Figure 64
Figure 64. Figure 64: Functional robustness and geometric stability of Query–Key alignment under semantic intervention. Rows correspond to the five ambiguous task families detailed in Appendix A. Columns 1–3 report downstream F V injection accuracy after Q patching across 3, 5, and 10-shot contexts. Column 4 presents the mean inner product between Q and K vectors on the Top-1 F V -head, averaged over 100 trials. The matrix vis… view at source ↗
Figure 65
Figure 65. Figure 65: Each row represents a model, with columns showing: (1) The cosine similarity between QK fixed uncontextualized and contextualized FVs (cos(F Vunc, F Vctx)), demonstrating that Value contextualization refines the task vector within its existing subspace; (2–4) present 2 × 2 cosine similarity matrices for ambiguous tasks (at 3, 5, and 10 shots), illustrating the cross-alignment between the {F Vunc, F Vctx} … view at source ↗
read the original abstract

In-context learning (ICL) excels at new tasks from minimal examples, yet we still lack a mechanistic explanation of how few-shot prompts shape a model's function vector (FV)--a causal activation direction that drives task behavior on the ICL query. Across tasks and models, an $n$-shot FV is well-approximated by a linear combination of example-level sub-FVs, suggesting additive and composable contributions from individual demonstrations. Beyond additivity, we show that models contextualize individual examples' representations based on prior examples to adaptively reweight which demonstrations dominate the FV: attention shifts toward examples that are more informative and less ambiguous under the context. Finally, a causal decomposition separates Query-Key routing from Value updates, finding that contextualization's most consistent contributions to FV quality arise from Query-Key alignment--particularly in ambiguous settings--while Value-mediated effects are more heterogeneous. Together, these results unify additive superposition with context-dependent attention reweighting into a mechanistic, testable account of how few-shot prompts implement tasks.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

0 major / 2 minor

Summary. The paper claims that n-shot function vectors (FVs) in in-context learning are well-approximated by linear combinations of example-level sub-FVs, that models contextually reweight individual examples via attention based on prior context, and that a causal decomposition isolating Query-Key routing from Value updates shows Query-Key alignment (especially in ambiguous cases) as the dominant contributor to contextualization effects on FV quality.

Significance. If the linear decomposition, reweighting effects, and causal separation hold with the reported consistency across tasks and models, the results would unify additive superposition accounts of ICL with context-dependent attention mechanisms into a single mechanistic framework, offering testable predictions about how demonstrations contribute to task behavior.

minor comments (2)
  1. [Abstract] Abstract: the phrase 'across tasks and models' is used without enumerating the specific tasks, models, or prompt formats tested; adding this list would improve clarity and allow readers to assess the scope immediately.
  2. The term 'FV quality' is used in the final sentence without an explicit operational definition or metric; a brief parenthetical or reference to the relevant evaluation section would resolve ambiguity.

Simulated Author's Rebuttal

0 responses · 0 unresolved

We thank the referee for their positive assessment of our work and the recommendation of minor revision. The referee's summary accurately reflects the paper's core claims on the linear decomposition of n-shot function vectors, attention-based reweighting of examples, and the causal isolation of Query-Key versus Value contributions.

Circularity Check

0 steps flagged

No significant circularity in derivation chain

full rationale

The provided abstract and reader summary describe empirical approximations (linear combination of sub-FVs) and causal interventions separating Query-Key from Value effects. No equations, fitted parameters, or self-citations are quoted that reduce any claimed prediction or decomposition to a definition or input by construction. The central claims rest on observable attention shifts and intervention results rather than self-referential fitting or imported uniqueness theorems. This is the expected non-finding for a paper whose internal logic does not collapse to its own inputs.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Only the abstract is available; no explicit free parameters, axioms, or invented entities can be extracted. The central claims rest on the prior existence of function vectors as causal directions, which is treated as background.

pith-pipeline@v0.9.1-grok · 5717 in / 1139 out tokens · 27996 ms · 2026-06-30T19:04:06.683603+00:00 · methodology

Review history (2 revisions) →

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

12 extracted references · 4 canonical work pages · 1 internal anchor

  1. [1]

    Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y ., Chen, A., Conerly, T., et al

    URL https://arxiv.org/abs/2510.1 0425. Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y ., Chen, A., Conerly, T., et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. URL https://tr ansformer-circuits.pub/2021/framework /index.html. Gu, Z., Zou, H. P., Liu, A., Chen, Y ., Zha...

    work page doi:10.18653/v1/2023.findings-emnlp.624 2021

  2. [2]

    Lu and Mary Letey and Jacob A

    URL https://openreview.net/forum ?id=w7LU2s14kE. Jiang, G., Jiang, C., Li, Z., Xue, S., Zhou, J., Song, L., Lian, D., and Wei, Y . Unlocking the power of function vectors for characterizing and mitigating catastrophic forgetting in continual instruction tuning. InThe Thirteenth Inter- national Conference on Learning Representations, ICLR 2025, Singapore, ...

    work page doi:10.1073/pnas.2502599122 2025

  3. [3]

    Min, S., Lyu, X., Holtzman, A., Artetxe, M., Lewis, M., Hajishirzi, H., and Zettlemoyer, L

    URL https://openreview.net/pdf?i d=8p3fu56lKc. Min, S., Lyu, X., Holtzman, A., Artetxe, M., Lewis, M., Hajishirzi, H., and Zettlemoyer, L. Rethinking the role of demonstrations: What makes in-context learning work? InProceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arab Emirates, Decembe...

    2022

  4. [4]

    In-context Learning and Induction Heads

    URL https://aclanthology.org/2022. emnlp-main.759/. Nguyen, K. A., im Walde, S. S., and Vu, N. T. Distin- guishing antonyms and synonyms in a pattern-based neural network. InEACL (1), pp. 76–85. Association for Computational Linguistics, 2017. URL https: //aclanthology.org/E17-1008/. Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T...

    work page internal anchor Pith review Pith/arXiv arXiv doi:10.18653/v1/2024.emnlp-main.192 2022

  5. [5]

    Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J

    URL https://openreview.net/forum ?id=WZH7099tgfM. Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J. M., Bengio, S., and Nakkiran, P. What algorithms can transformers learn? a study in length generalization. InThe Twelfth International Conference on Learning Representations, 2024. URL https:// proceedings.iclr.cc/paper_files/pape r/20...

    2024

  6. [6]

    Why are the findings about linear superposition in Section 3 then nontrivial? Conditioned on fixed attention weights, each head output is linear in V

    The function vector is the output of attention, which by definition is a linear combination. Why are the findings about linear superposition in Section 3 then nontrivial? Conditioned on fixed attention weights, each head output is linear in V . But the weights themselves areprompt- dependentnonlinear functions via softmax(QK ⊤), and Q is produced by multi...

  7. [7]

    on normal datasets, we do not observe a single universal alignment driver

    The task suite is restricted to classification-style mappings, failing to represent complex behaviors like creative generation or complex algorithmic tasks. We follow relevant prior work to focus on short-horizon, discrete mappings because their computational cost is desirable, and ambiguity can be introduced in a controlled way. We agree that open-ended ...

    work page arXiv 2022

  8. [8]

    The attention outputPn i=1 ai ·ψ(x i, yi) can be viewed as a FV , in thatϕ(xn+1,·)evaluates the task for other queriesx n+1 as well

    For unambiguous examples, ψ(x, y) only depends on the task. The attention outputPn i=1 ai ·ψ(x i, yi) can be viewed as a FV , in thatϕ(xn+1,·)evaluates the task for other queriesx n+1 as well

  9. [9]

    better than chance performance

    Whenever a prompt has both ambiguous and unambiguous examples, attention is lower on the ambiguous examples than the unambiguous examples. Here, by “better than chance performance” we rule out solutions where regularization is so strong that the optimal solution provides a constant answerϕ(·). 2Allowing infinite attention logits is an idealizing assumptio...

    2024

  10. [10]

    Write u:=ϕ A(xn+1), v:=ϕ B(xn+1), and, without loss of generality, assumeF A(xn+1) = 1andF B(xn+1) =−1

    For any unambiguous ICL query xn+1, we have {FA(xn+1), FB(xn+1)}={−1,1} . Write u:=ϕ A(xn+1), v:=ϕ B(xn+1), and, without loss of generality, assumeF A(xn+1) = 1andF B(xn+1) =−1. Define m:= u+v 2 , s:= u−v 2 . 68 How Few-Shot Examples Add Up Then (u−1) 2 + (v+ 1) 2 2 =m 2 + (s−1) 2. By (8), we have|s| ≤L √ S. Therefore, (u−1) 2 + (v+ 1) 2 2 ≥inf |s|≤L √ S ...

  11. [11]

    for unambiguous examples,ψ(x i, yi)depends only on the task 2.ψ pA =−ψ pB 3.ϕis linear on the line betweenψ pA andψ pB

  12. [12]

    First, for any choice of S:= max x,y ∥ψ(x, y)∥2 2 and L, there is a model satisfying 1–4 that makes (6), (8), and (9) tight

    in prompts containing both ambiguous and unambiguous examples, attention is zero on the ambiguous examples We want to show that any global optimum that has above-chance performance must satisfy these features. First, for any choice of S:= max x,y ∥ψ(x, y)∥2 2 and L, there is a model satisfying 1–4 that makes (6), (8), and (9) tight. Take ψA to be any vect...