Capturing In-Context Learning Dynamics with Task Operators

• Each ICL inference processes all demonstrations, which can be computationally expensive; fixed activation-vector methods can fail on complex multi-step reasoning over input content.Source: §1 · pp. 1–2
• TO represents demonstration effects as affine transformations of attention-head outputs and replays them through dynamic updates to the attention output projection.Source: §1 · p. 2
• Causal analysis identifies sparse, task-specific ICL circuits, while averaging operators from disjoint demonstration batches supports many-shot scaling.Source: §1 · p. 2
• Across four language models and eight tasks, TO consistently outperforms evaluated prior methods and substantially closes the gap between zero-shot and few-shot performance.Source: §1 · p. 2

• Theoretical accounts describe ICL as implicit Bayesian inference [36] or implicit fine-tuning through attention and meta-gradients, while mechanistic studies identify task-dependent heads and representations [24] [12] [7].Source: §2 · p. 2
• Reusable task representations include weight-space task vectors [16], activation-space task and function vectors [13, 33], and other representations derived from demonstration activations [21] [18] [19].Source: §2 · pp. 2–3
• Many prior compression methods replay knowledge through fixed additive interventions, whereas task operators combine multiplicative scaling and additive context contributions so replay depends on the test-time hidden state.Source: §2 · p. 3
• Richer intervention methods introduce learned or optimized components, or estimate transformations from activation statistics; TO instead derives input-dependent affine parameters directly from the ICL forward pass.Source: §2 · p. 3
• Causal circuit methods localize model computations, and this work uses causal analysis to identify sparse, task-specific sites storing ICL knowledge for zero-shot replay.Source: §2 · p. 3

3 Task Operators in In-Context Learning

• TO transfers ICL behavior to zero-shot inference through dynamic interventions at multiple model sites, comprising affine extraction, replay, and sparsification.Source: §3 · p. 3, Figure 1 · p. 4

3.1 Knowledge Extraction

• The ICL effect on each attention head has an exact affine identity, valid for arbitrary inputs, attention patterns, and context spans; its reparameterization is approximately equivalent to an output-projection update.Source: §3.1 · p. 4
• Operators are extracted across token positions and layers, with content-token and generation-position operators averaged; parameters are averaged across validation examples and can also be averaged across demonstration sets.Source: §3.1 · pp. 4–5

3.2 Knowledge Replay

• At zero-shot inference, extracted operators are injected at mapped query-template and content positions, then installed at generation positions during autoregressive decoding.Source: §3.2 · p. 5

3.3 Knowledge Sparsification

• Sparsification measures each site's causal effect by perturbing it to identity and evaluating the remaining fraction of operator NLL gain; site importance is the loss in recovery across validation samples.Source: §3.3 · p. 5

4.1 Experimental Setup

• Experiments cover four LLMs and eight lexical, algorithmic, and reasoning tasks, comparing TO with four extraction baselines, zero-shot learning, and 8-shot ICL.Lexical and algorithmic evaluations use 128 held-out samples; reasoning uses the standard test sets.Source: §4.1 · p. 6

4.2 Performance Comparison

• TO leads replay baselines on nearly every task-model pair, recovering much of the ICL gain on complex algorithmic and reasoning tasks.The exception is where ICL itself adds little or reduces performance relative to zero-shot learning.Source: Table 1 · p. 6, §4.2 · pp. 6–7

4.3 Sparse Circuit Analysis

• TO importance is often sparse but task-specific: Uppercase concentrates on template tokens, Reverse spreads across template and content tokens, and Deduplicate focuses on middle-layer outputs.The importance maps cover algorithmic tasks on Llama3.1-8B; the patterns help explain why template-focused baselines fare better on Uppercase than on Reverse and Deduplicate.Source: §4.3 · p. 7, Figure 2 · p. 7

4.4 Scaling Analysis

• With K=8 fixed, TO exceeds TV at every m on every task; performance is largely recovered by m=8, while larger m moves TO toward ICL even when accuracy falls.On Deduplicate, increasing m from 8 to 32 improves TO; on Reverse it slightly decreases TO toward the ICL result.Source: Table 2 · p. 8, §4.4 · p. 8
• As K grows, TO can decline on Reverse and Deduplicate, whereas batch-averaged TO is non-decreasing across tasks and can exceed ICL at K=64 on Deduplicate.TO (batch) averages operators from disjoint 8-shot batches; the figure fixes m=32 on Llama3.1-8B.Source: Figure 3 · p. 8, §4.4 · p. 8

4.5 Ablation Studies

• The all-position operator performs best overall; simple tasks rely mostly on template positions, while Reverse, Deduplicate, and reasoning tasks need content and output positions.On the latter tasks, Reverse relies more on content positions, while removing output positions causes a larger drop on Deduplicate.Source: Table 3 · p. 9, §4.5 · p. 9
• At template positions, scaling is essential and bias alone is insufficient; on non-reasoning tasks, either component alone collapses to zero accuracy.The two components are individually insufficient but jointly effective, consistent with the observed affine structure.Source: §4.5 · p. 9, Table 3 · p. 9

• Task Operator captures ICL knowledge as affine transformations of attention head outputs and replays them during zero-shot inference through dynamic output-projection updates.Source: §5 · p. 9
• Across four language models and eight tasks, TO outperforms all evaluated prior methods and substantially narrows the zero-shot-to-few-shot performance gap.Source: §5 · p. 10
• The analysis finds sparse, task-specific critical sites and shows that averaging operators from disjoint demonstration batches enables many-shot learning without expanding the context window.Source: §5 · p. 10

A Limitations and Broader Impacts

• At very large K, TO can become less stable; averaging operators from independently sampled demonstration batches mitigates this, while evaluation remains limited to text-only decoder models and selected task categories.The authors suggest reduced attention mass on non-context positions may increase replay parameter variance; multimodal reasoning is a proposed future extension.Source: A Limitations and Broader Impacts · p. 16
• Reusing extracted operators may reduce context-window use and inference cost when the same task is queried repeatedly; the authors advise validating operators, especially for sensitive or biased demonstrations.The work focuses on mechanistic understanding of ICL and task-specific sparse circuits.Source: A Limitations and Broader Impacts · p. 16

B Dataset Details

• Non-reasoning demonstrations use a shared Input/Output template, and lexical and algorithmic evaluations each use 128 held-out test records.Reasoning prompts append “Let’s think step by step.”Source: B Dataset Details · p. 16

B.1 Lexical Tasks

• The lexical tasks are English-to-French single-word translation and English singular-to-plural inflection, with pools of 739 and 173 word pairs, respectively.The translation pool is split into 579 demonstrations, 32 validation, and 128 test records; the inflection pool into 13, 32, and 128.Source: §B.1 · p. 16

B.2 Algorithmic Tasks

• The three synthetic algorithmic tasks use an eight-letter alphabet and cover deduplication, uppercasing, and sequence reversal.Each split contains 1,024 demonstrations, 256 validation, and 128 test records.Source: §B.2 · p. 17

B.3 Reasoning Tasks

• All three reasoning tasks use chain-of-thought formatting, with held-out test sets from their respective benchmarks.The benchmark sizes are 1,319 for GSM8K, 500 for MATH-500, and 198 for GPQA-Diamond.Source: §B.3 · pp. 17–18
• GSM8K evaluates grade-school arithmetic reasoning; its 1,024 demonstrations and 256 validation records come from training, while its test set contains 1,319 records.Targets end with “The answer is N.”Source: §B.3 · p. 17
• MATH-500 evaluates competition-level mathematics across seven subject areas, using 1,024 demonstrations, 256 validation records, and 500 test problems.Demonstrations and validation records are drawn from the Hendrycks MATH training split.Source: §B.3 · p. 18
• GPQA-Diamond comprises 198 graduate-level physics, chemistry, and biology test questions, with demonstrations and validation records drawn from GPQA-Main after removing Diamond questions.Answer choices are shuffled per record with a fixed seed.Source: §B.3 · p. 18

C Implementation Details

• Experiments use Qwen3-4B, Qwen3-8B, Llama-3.2-3B-Instruct, and Llama-3.1-8B-Instruct.The models are instruction-tuned variants from the Qwen3 and Llama families.Source: C Implementation Details · p. 19
• TO extraction uses K=8 demonstrations and m=32 answer-excluded validation prompts, averaging per-head parameters across prompts and applying category-level operators to content and generated tokens.Template token positions receive independent per-head pairs; the all-positions, all-layers variant patches every decoder layer without site pruning.Source: C Implementation Details · p. 19
• The comparison includes Task Vectors, Function Vectors, In-Context Vectors, and Conceptors, alongside zero-shot and standard 8-shot ICL references.All baselines use the same K=8 demonstrations and prompt templates as TO.Source: C Implementation Details · p. 19
• Decoding is deterministic at temperature 0.0, with 64-token budgets for lexical and algorithmic tasks and 512-token budgets for reasoning tasks.Generation stops at the first newline for lexical and algorithmic tasks and at “\nInput:” for reasoning tasks.Source: C Implementation Details · p. 19

D Stability of Extracted Parameters

• The stability analysis relies on averaging extractions across prompts at fixed token positions and across token offsets within content or output categories.It separately checks cross-demonstration-set stability.Source: D Stability of Extracted Parameters · p. 19

D.1 Cross-sample stability

• 80.7% of sites have CV(w)<0.1 and 92.1% have cos(b)>0.9 across twelve model-task cells, supporting cross-sample averaging.Llama models have higher fractions at CV(w)<0.1 than Qwen models (85–92% versus 70–77%).Source: Table 4 · p. 20, §D.1 · p. 20

D.2 Cross-token stability within a category

• Broadcast matches or exceeds per-rank across all settings, including Translation and Reverse with complete offset coverage.Translation matches exactly (∆=0); Reverse favors broadcast by 2–10 pp. On GSM8K broadcast matches or outperforms in three of four models (gaps of −5 to −28 pp), with Llama3.1-8B the exception (+2.3 pp).Source: §D.2 · pp. 20–21, Table 5 · p. 21

D.3 Cross-demonstration-set stability

• Scaling-factor correlations exceed 0.96 in every model-task cell, averaging 0.984; bias-vector cosine similarity ranges from 0.81 to 0.95 across cells.The operators are extracted from non-overlapping demonstration pools from the same task.Source: §D.3 · p. 21

E Task Operator Performance with Sparsified Circuits

• For Qwen3-8B on Deduplicate, accuracy rises from 65.6% with the dense operator to 82.8% at p = 20%.The passage suggests low-importance sites can introduce interference rather than useful task information.Source: E Task Operator Performance with… · p. 22
• At p = 1%, performance generally collapses, while retaining roughly 20% of sites recovers nearly full dense-operator accuracy on Uppercase and Reverse.This indicates that replay needs a coordinated set of interventions, rather than a single dominant site.Source: Figure 4 · p. 22, E Task Operator Performance with… · p. 22
• Validation-NLL selection sweeps p across thresholds, but adds tuning and may overfit validation noise; the main experiments therefore use the full operator at p = 100%.The full operator remains competitive across tasks and models.Source: E Task Operator Performance with… · pp. 22–23
• 12 model–task cells: validation-selected sparse operators preserve or improve test accuracy while retaining 67% of sites on average.The largest gains are on Reverse and Deduplicate, where pruning improves accuracy by up to 9.4 percentage points.Source: E Task Operator Performance with… · p. 22, Table 7 · p. 23

F Robustness to Prompt Template Choice

• The evaluation uses three matched formats: canonical Input/Output, Q/A, and compact arrow prompts, with TO and TV extracted and evaluated under each same template.The experiments cover one representative task per category: Translation, Deduplicate, and GPQA.Source: F Robustness to Prompt Template Choice · p. 23
• TO applies affine operators across layers, attention heads, and position types, making replay less dependent on any single prompt token or extraction site.The reported robustness applies when extraction and evaluation use the same template.Source: F Robustness to Prompt Template Choice · p. 23
• Across matched templates, TO stays close to ICL on Translation, remains well above TV on Deduplicate, and is broadly comparable across GPQA templates.The table reports test accuracy for the canonical Input/Output, Q/A, and compact arrow formats.Source: F Robustness to Prompt Template Choice · p. 23, Table 8 · p. 24

G Cross-Task Transfer of Extracted Knowledge

• The study extracts an operator from GSM8K using K = 8 demonstrations and replays it on MATH-500 and GPQA-Diamond while keeping each target prompt target-specific.The comparison includes zero-shot inference, standard in-domain operators, and full ICL on the target task.Source: G Cross-Task Transfer of Extracted… · p. 24
• On GPQA-Diamond, transferred operators remain competitive with zero-shot inference for most model–task combinations, but transfer is more sensitive to target format.The source and target differ in domain and answer style; one Llama model improves over its in-domain operator.Source: G Cross-Task Transfer of Extracted… · p. 24
• GSM8K operators transfer strongly to MATH-500, staying close to in-domain operators and sometimes competitive with or better than full ICL.The table compares transferred operators with zero-shot inference, in-domain operators, and standard ICL.Source: G Cross-Task Transfer of Extracted… · p. 24, Table 9 · p. 25