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-05-20 19:15 UTC · model grok-4.3

classification 💻 cs.LG cs.AI
keywords in-context learningfunction vectorsfew-shot promptingattention mechanismscausal decompositionlanguage modelsquery-key alignment
0
0 comments X

The pith

An n-shot function vector approximates the sum of one-shot sub-vectors from each example, with added context-driven reweighting of their contributions.

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 internal direction a model uses to perform a new task after seeing n examples can be closely reconstructed by simply adding the directions obtained from each example in isolation. This additivity holds across multiple tasks and model sizes and implies that individual demonstrations contribute in a composable way rather than interacting in complex nonlinear ways. The analysis further shows that the model does not treat every example equally; it shifts attention toward examples that become more informative once earlier ones are seen, especially when some demonstrations are ambiguous. A finer causal breakdown isolates that these reweighting effects operate mainly through improved alignment between queries and keys in the attention layers, while value updates show more variable influence. The resulting picture explains how a short prompt assembles a usable task representation from separate pieces.

Core claim

Across tasks and models, an n-shot function vector is well-approximated by a linear combination of example-level sub-function vectors, indicating additive and composable contributions from individual demonstrations. Models further contextualize each example's representation according to the examples that precede it, adaptively reweighting which demonstrations most strongly shape the overall function vector. The dominant causal contribution to this reweighting comes from query-key alignment in attention, particularly when examples are ambiguous, whereas value-mediated updates produce more heterogeneous effects.

What carries the argument

The function vector, defined as the causal activation direction that steers the model toward the demonstrated task on a query, together with its decomposition into per-example sub-vectors and the attention-based reweighting that adjusts their relative influence.

If this is right

  • Multi-shot performance could be predicted from single-example measurements if the linear approximation continues to hold.
  • The order in which examples appear should affect which ones dominate the final function vector because of the contextual reweighting.
  • Targeted improvements to query-key alignment during training or inference could strengthen in-context learning on ambiguous inputs.
  • Demonstrations that remain ambiguous even after earlier context is provided contribute less to the overall task direction.

Where Pith is reading between the lines

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

  • Similar additive decompositions might appear in chain-of-thought or other structured prompting formats if the same internal mechanism is at work.
  • Training objectives that reward better contextual reweighting of examples could improve few-shot robustness without increasing example count.
  • The pattern may extend to non-language settings such as vision or multimodal models if function-vector-like directions exist there.
  • Redundant or mutually inconsistent examples should produce smaller net function vectors under this additive-plus-reweighting view.

Load-bearing premise

That the measured shifts in attention and query-key interactions are direct causes of better function-vector performance rather than incidental byproducts of the intervention technique.

What would settle it

If an intervention that prevents the observed attention shifts toward less ambiguous examples leaves the function vector's task performance unchanged on ambiguous queries, the causal account of reweighting would be falsified.

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

2 major / 3 minor

Summary. The manuscript claims that across tasks and models, an n-shot function vector (FV) in in-context learning is well-approximated by a linear combination of example-level sub-FVs, indicating additive and composable contributions from individual demonstrations. It further reports that models contextualize individual examples via attention shifts to adaptively reweight demonstrations (favoring informative, less ambiguous ones), and a causal decomposition separates Query-Key routing from Value updates, finding that Query-Key alignment provides the most consistent contributions to FV quality, especially in ambiguous contexts, while Value-mediated effects are more heterogeneous.

Significance. If the results hold, the work offers a mechanistic, testable account unifying additive superposition with context-dependent attention reweighting in ICL. Strengths include consistent empirical patterns across tasks/models and the use of causal interventions to probe the decomposition. This could inform prompt design and LLM interpretability. The additivity approximation is less vulnerable to intervention concerns, but the contextualization and QK-dominance claims depend on clean separation of effects.

major comments (2)
  1. [Causal Decomposition] Causal Decomposition section: the claim that Query-Key alignment effects are the primary causal driver of FV quality improvements (particularly in ambiguous settings) requires stronger evidence that the intervention isolates QK routing without confounding from entangled attention components or indirect reweighting; if QK and Value pathways share downstream effects, the reported dominance may reflect correlated side effects rather than direct causation.
  2. [Contextualization and Reweighting] Contextualization results: while attention shifts toward more informative examples are reported, the analysis would benefit from explicit controls confirming these are not artifacts of the specific intervention method or model routing, as this is load-bearing for the claim that models adaptively reweight demonstrations beyond simple additivity.
minor comments (3)
  1. Clarify the precise extraction and normalization procedure for sub-FVs and how linear coefficients are fitted, including any regularization or constraints applied.
  2. Add statistical details such as confidence intervals or p-values for the approximation quality of the linear combination across the reported tasks and models.
  3. Ensure consistent notation for FV, sub-FV, and attention components throughout; define all acronyms on first use in the main text.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for their constructive comments on our work. We provide point-by-point responses to the major comments below and will make revisions to address the concerns about evidence strength in our causal decomposition and controls for contextualization results.

read point-by-point responses

  1. Referee: [Causal Decomposition] Causal Decomposition section: the claim that Query-Key alignment effects are the primary causal driver of FV quality improvements (particularly in ambiguous settings) requires stronger evidence that the intervention isolates QK routing without confounding from entangled attention components or indirect reweighting; if QK and Value pathways share downstream effects, the reported dominance may reflect correlated side effects rather than direct causation.

    Authors: We appreciate the referee raising this important point about the isolation of effects in our causal interventions. In the manuscript, we separate Query-Key routing by intervening on the pre-softmax attention scores derived from QK alignments while preserving the Value vectors, following established methods in attention mechanism analysis. Our results indicate that these QK interventions yield more reliable improvements to FV quality, especially under ambiguity, compared to Value interventions which are more variable. To mitigate concerns about confounding or shared downstream effects, we will expand the revision with additional experiments that quantify the independence of these pathways, including measuring residual effects after QK intervention and vice versa. This will clarify that the dominance of QK is not an artifact of correlation. revision: partial

  2. Referee: [Contextualization and Reweighting] Contextualization results: while attention shifts toward more informative examples are reported, the analysis would benefit from explicit controls confirming these are not artifacts of the specific intervention method or model routing, as this is load-bearing for the claim that models adaptively reweight demonstrations beyond simple additivity.

    Authors: Thank you for this feedback. We note that the attention reweighting analysis is based on direct observation of attention patterns in the standard forward computation, without any interventions applied. To address potential artifacts from model routing or specific setups, we will add explicit control experiments in the revised manuscript. These include randomizing example orders, using shuffled contexts, and comparing to baseline models with fixed attention. The results of these controls will be presented to demonstrate that the shifts toward informative and less ambiguous examples are indeed adaptive and not spurious, thereby supporting our claims about contextualization beyond additivity. revision: yes

Circularity Check

0 steps flagged

No significant circularity; empirical measurements and interventions are self-contained

full rationale

The paper's central claims rest on empirical approximations of n-shot function vectors as linear combinations of sub-FVs, observed attention shifts, and causal interventions separating Query-Key from Value effects. These are obtained via direct measurements and interventions on model activations rather than any derivation that reduces to fitted parameters or self-defined quantities by construction. No equations or results are presented as predictions that are statistically forced by prior fits within the paper. Self-citations, if present, are not load-bearing for the core additivity or decomposition findings, which are tested across tasks and models with external falsifiability through interventions.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

Based solely on the abstract, the central claims rest on standard mechanistic-interpretability assumptions about function vectors as causal directions and the validity of activation interventions; no explicit free parameters or new entities are described.

axioms (1)
  • domain assumption Function vectors are causal activation directions that drive task behavior on the ICL query
    Invoked in the definition and causal analysis of FVs throughout the abstract.

pith-pipeline@v0.9.0 · 5717 in / 1185 out tokens · 50750 ms · 2026-05-20T19:15:02.385992+00:00 · methodology

discussion (0)

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