Capturing In-Context Learning Dynamics with Task Operators
In-context learning (ICL) lets LLMs adapt to new tasks from demonstrations, but reprocessing all examples for every query is costly and prior compression methods fail on complex, input-dependent tasks. This paper introduces Task Operator (TO), a training-free method that extracts per-head affine transformations — scaling weights and bias vectors — from the ICL forward pass and replays them during zero-shot inference by updating attention output projections. By treating the ICL effect as a dynamic, state-dependent modification rather than a fixed activation vector, TO substantially narrows the gap between zero-shot and few-shot performance while revealing sparse, task-specific circuits.
The paper makes a concrete advance by deriving an exact per-head affine decomposition of ICL's effect on attention outputs and converting it into a replayable operator. Empirically, TO is the strongest training-free ICL compression method tested, and the sparsification analysis offers genuine interpretability. Still, the claim that this establishes a "direct mechanistic link between in-context learning and parameter updates" overstates the finding: Equation (4) is an exact reparameterization of activations during ICL, not evidence that the model literally performs gradient descent on its weights during inference.
Two things hold up well. First, the affine decomposition in Section 3.1 is mathematically exact, and the cross-sample stability analysis in Appendix D strongly supports averaging extracted parameters across prompts. Second, the causal sparsification analysis in Section 4.3 clearly shows that ICL knowledge is not diffusely distributed but instead concentrates in sparse, task-specific circuits, which directly explains why prior fixed-location extraction methods fail on complex tasks.
The central conceptual leap — that the affine form means ICL is "approximately equivalent" to fine-tuning the output projection — is algebraically true for activations but does not imply the model is internally executing a parameter update. The paper also admits that TO becomes unstable when the number of demonstrations $K$ is large because the per-head scaling factor $w_t^{(\ell,h)}$ approaches zero, and the proposed remedy of batch averaging is somewhat post-hoc. Additionally, the evaluation scope is narrow: only four decoder-only models and eight tasks, all text-only, leaving open whether the extracted operators generalize to longer contexts, multimodal inputs, or more open-ended generation.
A second concern is that on several reasoning benchmarks the absolute gains of TO over strong baselines are modest, and in cases where ICL itself hurts (e.g., GPQA on Llama models), TO mirrors that failure. While the authors interpret this as faithfulness, it also means TO cannot recover useful task structure when the underlying ICL signal is weak or harmful, limiting its value as a standalone deployment tool.
The comparisons to prior methods are fair and well-controlled: all baselines use the same $K=8$ demonstrations, matched prompt templates, and answer-excluded validation samples for hyperparameter search. Table 1 demonstrates that prior fixed-vector methods collapse on algorithmic tasks like Reverse and Deduplicate, whereas TO recovers a substantial portion of the ICL gain, supporting the claim that dynamic, input-dependent interventions are necessary. That said, on reasoning tasks the baselines are often competitive with TO in absolute terms, and the paper's own analysis notes that on GPQA and MATH500 for both Llama models, where ICL gains are small or negative, TO shows correspondingly limited or no improvement. This weakens the universal claim that affine operators are always the right abstraction.
The authors release code and specify exact model checkpoints, generation parameters (temperature 0, stop sequences), and extraction hyperparameters ($K=8$, $m=32$). The extraction formulas, including the numerically stable logsumexp form for computing $w_t^{(\ell,h)}$, are given explicitly. Reproduction is nevertheless demanding: the full operator requires forward passes over all layers and positions, plus causal sparsification ablations that compute teacher-forced NLLs across many sites. Independent reproduction would be straightforward for a well-resourced lab but prohibitive without access to large-memory GPUs and the exact prompt templates, which are fortunately detailed in Appendix B.
In-context learning (ICL) enables language models to perform new tasks from demonstrations without weight updates. However, every ICL inference requires processing the full set of examples, resulting in inefficient deployments, and how ICL works mechanistically is not fully understood. Prior work compresses ICL into fixed activation vectors extracted from specific layers or positions, but these input-independent interventions fail on complex tasks where the output depends on fine-grained interactions with the input. By analyzing the ICL forward pass, we show that each attention head's output is an affine transformation of its context-masked counterpart, and that the parameters of this transformation are empirically stable across samples for a given task. Building on this, we introduce Task Operator (TO), which replays this transformation as an analytically derived update to the attention output projection. Across lexical, algorithmic, and reasoning tasks, TO achieves the best overall performance among prior methods and substantially narrows the gap between zero-shot inference and ICL. We further show that the extracted knowledge concentrates in a task-specific sparse circuit across layers and positions, and that averaging operators from disjoint demonstration batches enables effective many-shot scaling without expanding the context window. Our code is available at https://github.com/gzxiong/task_operator.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.