In-context learning (ICL) lets large language models adapt to new tasks from demonstrations, and fine-tuning can erode this behaviour. Many preservation diagnostics inspect attention: if attention changes when demonstrations change, the model is treated as context-sensitive. This paper asks how far that proxy can be trusted once it is optimised. We formalise In-Context Sensitivity (ICS), the average row distance between last-token attention on matched and mismatched demonstration prefixes, and pair it with ICL-GAP, the behavioural accuracy gap between the same prefixes. In a controlled four-arm ablation on Llama-2-7B, an ICS-maximising regulariser (\armKL) drives ICS to 1.413, within 0.5% of its geometric ceiling. The behavioural readout tells a different story: ICL-GAP stays near zero and MMLU accuracy moves from 0.371 to 0.279, a Goodhart dissociation of the bounded attention proxy. Endpoint statistics locate the mechanism: attention grows sharp and near-disjoint across prefixes yet routes to formatting and demonstration-body tokens rather than labels. A random-label protocol confirms that the behavioural probe family retains dynamic range at the same checkpoints. In a constructive sweep, behaviour gating partially mitigates the effect, while objectives anchored to pretrained computation hold the high-MMLU, moderate-ICS region that divergence maximisers leave. The main lesson is diagnostic: attention-level ICL proxies earn their place as training targets only after validation against behavioural gaps.
Figures & tables
Figure 1: Under the ICS-maximising stress test, the attention proxy approaches its ceiling while ICL-GAP stays near zero and MMLU declines. This divergence motivates testing the proxy against behavioural ICL under optimisation.
Figure 2: Diagnostic framework. Matched and mismatched prompts feed the same fine-tuning pipeline, while the attention-level ICS and behavioural ICL-GAP probes are computed on a held-out pair set disjoint from any regulariser pairs. The architecture also shows the two constructive directions studied later: behaviour-gated diagnostics and anchoring to the pretrained computation.
Arm
Match
Perm.
None
Gap [95% CI]
F_A
0.663
0.360
0.492
0.302 [0.235, 0.371]
F_M
0.656
0.381
0.492
0.275 [0.206, 0.346]
F_None
0.623
0.435
0.492
0.188 [0.115, 0.260]
F_ICS-Max
0.692
0.340
0.492
0.352 [0.292, 0.413]
Table 1: RQ0 random-label A/B protocol at step 5,000 . All arms retain a clear matched-vs.-permuted behavioural gap on this protocol, so the near-zero controlled QA ICL-GAP cannot be attributed to a universally insensitive behavioural probe family.
Arm
ICS ( ↑ )
ICL-GAP ( ↑ )
MMLU acc. ( ↑ )
ECE ( ↓ )
Pretrained Llama-2-7B
0.516
—
—
—
F_A
0.492
—
0.338
0.434
F_M
0.482
—
0.342
0.406
F_None
0.500
—
0.375
0.407
F_ICS-Max
1.413
−0.010
0.279
0.231
Table 2: Final controlled-ablation results at step 5,000 . ICS is the attention-level diagnostic from Equation 3 , bounded above by 2≈1.414 . ICL-GAP is the behavioural quantity from Equation 7 . MMLU accuracy is on a 240 -item probe (chance 0.25 ). The pretrained baseline ICS is 0.516 .
Figure 3: Step- 5,000 outcome in the (ICS, MMLU) plane. F_ICS-Max saturates the attention proxy near the 2 ceiling but moves into the low-MMLU region; anchored variants remain near the pretrained ICS while preserving MMLU.
Figure 4: Per-step trajectories of ICS, ICL-GAP, MMLU accuracy, and MMLU ECE across 5,000 fine-tuning steps. The trajectory view makes the dissociation temporal: F_ICS-Max drives ICS upward while ICL-GAP remains near zero and MMLU falls.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
F_A
F_M
F_None
F_ICS-Max
Step
ICS
MMLU
ECE
ICS
MMLU
ECE
ICS
MMLU
ECE
ICS
GAP
MMLU
ECE
100
0.474
0.375
0.158
0.470
0.375
0.160
0.474
0.354
0.147
0.513
−0.020
0.371
0.136
500
0.485
0.367
0.220
0.480
0.354
0.230
0.485
0.371
0.215
1.225
+0.005
0.367
0.274
1000
0.487
0.346
0.310
0.481
0.342
0.305
0.490
0.371
0.305
1.292
+0.025
0.304
0.473
1500
0.488
0.342
0.360
0.481
0.342
0.350
0.493
0.371
0.350
1.309
+0.050
0.333
0.429
2000
0.490
0.338
0.395
0.482
0.342
0.380
0.495
0.371
0.380
1.319
+0.040
0.317
0.452
Appendix
Table 3: Full trajectories for the controlled proxy-stress experiment, sampled every 500 steps.
Labels
N
Match
Perm
None
Gap [95% CI]
A/B
480
0.669
0.362
0.492
0.306 [0.231, 0.379]
X/Y
240
0.762
0.225
0.504
0.537 [0.450, 0.621]
foo/bar
240
0.642
0.338
0.496
0.304 [0.200, 0.400]
Appendix
Table 4: Pretrained behavioural ICL screening on random-label binary episodes. CIs are paired bootstrap intervals over episodes.
Arm
ICS (→2)
ICL-GAP
MMLU
ECE
Pretrained
0.516
—
0.371
—
F_ICS-Max (stress test)
1.413
−0.010
0.279
0.231
F_BICS (this work)
1.397
+0.020
0.283
0.388
Appendix
Table 5: F_BICS at β=10,δ=0.05 attenuates the ICS-vs.-behaviour dissociation but does not eliminate it. The arm still reaches ICS=1.397 ( 98.8% of the 2 ceiling), but achieves a +0.030 ICL-GAP improvement over F_ICS-Max . We attribute the residual saturation to the softness of the gate; a hard 1{⋅} gate with a straight-through estimator is a candidate for further evaluation.
Configuration
ICS (→0.516)
ICL-GAP
MMLU (→0.371)
ECE
Pretrained
0.516
—
0.371
—
F_None (controlled)
0.500
≈0
0.375
0.407
F_ICS-Max (stress test)
1.413
−0.010
0.279
0.231
F_anchor , λ=0.01
0.501
−0.040
0.338
0.472
F_anchor , λ=0.05
0.500
−0.010
0.367
0.375
F_anchor , λ=0.20
0.500
−0.020
0.350
0.374
Appendix
Table 6: AnchorTune λ -sweep at seed 42 and three-seed final at λ⋆=0.05 . Pretrained ICS is 0.516 and pretrained MMLU is 0.371. At λ⋆ , the three-seed AnchorTune result is MMLU 0.370±0.058 .
Figure 5: ICS and MMLU at step 5,000 on Llama-2-7B. The dashed vertical line marks the 2 ceiling. F_ICS-Max saturates ICS while MMLU approaches chance, and F_BICS partially follows the same pattern under the smooth β=10 gate. F_anchor at λ⋆=0.05 and the logit-anchor baseline LogitAnchor keep ICS near the pretrained value while maintaining MMLU within the observed seed variation.
Figure 6: Per-step trajectories of ICS, ICL-GAP, MMLU accuracy, and MMLU ECE across 5,000 training steps for five representative arms. F_ICS-Max and F_BICS rapidly drive ICS toward the 2 ceiling. F_anchor , LogitAnchor , and WeightAnchor keep all four quantities near the pretrained baseline. ICL-GAP fluctuates around zero throughout; the single-checkpoint sampling scale for the 200 logged paired examples is about 0.046 .
Arm
ICS
ICL-GAP
MMLU
ECE
Pretrained
0.516
—
0.371
—
F_None
0.500
≈0
0.375
0.407
F_ICS-Max
1.413
−0.010
0.279
0.231
LogitAnchor (KL-to-logits)
0.494
−0.020
0.383
0.373
WeightAnchor (weight ℓ2 , λ=10−7 )
0.500
−0.010
0.363
0.419
F_anchor (ours, λ⋆=0.05 )
0.500
−0.010
0.367
0.375
Appendix
Table 7: AnchorTune compared with two capability-preservation baselines and the F_ICS-Max regulariser. AnchorTune values are the seed-42 result for direct comparison with the single-seed baselines; Table 6 reports its three-seed mean.
Figure 7: (MMLU, ICL-GAP) Pareto scatter at step 5,000 . Divergence regularisers (red X for F_ICS-Max , orange for F_BICS ) sit in the low-MMLU region; anchor regularisers (blue/green for F_anchor , LogitAnchor , WeightAnchor ) cluster around the pretrained reference point (black star). The two regulariser families form distinct clusters on this plot.
Anchor target
ICS
ICL-GAP
MMLU
ECE
Pretrained
0.516
—
0.371
—
None (= F_None , controlled reference)
0.500
≈0
0.375
0.407
Mid-band attn rows ( F_anchor )
0.500
−0.010
0.367
0.375
Last 4 layers attn rows ( F_anchor_last4 )
0.510
−0.015
0.358
0.377
LM head ℓ2 only ( F_anchor_lmhead )
0.494
−0.020
0.375
0.378
Appendix
Table 8: Anchor-target ablation. All three variants keep ICS within ±0.020 of pretrained while maintaining MMLU; the lm_head-only ℓ2 variant without an attention anchor matches pretrained MMLU, weakening the claim that mid-band attention rows are uniquely necessary for MMLU recovery.
In-context learning (ICL) allows language models to adapt their behavior from instructions and demonstrations in a prompt, but accuracy alone does not reveal how contextual information controls their predictions. We study the acquisition and trainability of ICL with a controlled diagnostic that separates copying a nearby demonstration from integrating information across the prompt. Across 17 models, about a third prefer the copied answer, a few are undecided, and the rest prefer the rule-consistent answer, which goes with higher accuracy and with control by the requested value rather than the closest conflicting example. Checkpoint analyses show that models begin prioritizing task-relevant prompt tokens over conflicting examples long before their accuracy improves beyond chance level. In matched training experiments, synthetic documents that define endings and then query them shift models toward compositional responses on the controlled diagnostic, whereas fine-tuning on definition exercises built from rewritten real documents improves accuracy when prompts reassign label meanings. The results suggest that ICL is not acquired as a single, uniform capability. Instead, prompt control, answer preference, construction without candidates, and transfer can develop along distinct trajectories.
In-Context Learning (ICL) enables pretrained LLMs to adapt to downstream tasks by conditioning on a small set of input-output demonstrations, without any parameter updates. Although there have been many theoretical efforts to explain how ICL works, most either rely on strong architectural or data assumptions, or fail to capture the impact of key practical factors such as demonstration selection, Chain-of-Thought (CoT) prompting, the number of demonstrations, and prompt templates. We address this gap by establishing a theoretical analysis of ICL under mild assumptions that links these design choices to generalization behavior. We derive an upper bound on the ICL test loss, showing that performance is governed by (i) the quality of selected demonstrations, quantified by Lipschitz constants of the ICL loss along paths connecting test prompts to pretraining samples, (ii) an intrinsic ICL capability of the pretrained model, and (iii) the degree of distribution shift. Within the same framework, we analyze CoT prompting as inducing a task decomposition and show that it is beneficial when demonstrations are well chosen at each substep and the resulting subtasks are easier to learn. Finally, we characterize how ICL performance sensitivity to prompt templates varies with the number of demonstrations. Together, our study shows that pretraining equips the model with the ability to generalize beyond observed tasks, while CoT enables the model to compose simpler subtasks into more complex ones, and demonstrations and instructions enable it to retrieve similar or complex tasks, including those that can be composed into more complex ones, jointly supporting generalization to unseen tasks. All theoretical insights are corroborated by experiments.
Xuhan Tong, Yuchen Zeng, Jiawei Zhang
Department of Computer Sciences, University of Wisconsin-Madison · Microsoft Research
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.