Attention Sensitivity Is Not Enough: Dissociating Attention-Level and Behavioural In-Context Learning under Fine-Tuning

Jinyuan Zhang; Peng He; He Hu; Yin Yuan; ShengShuo Jiao

Source-linked AI summary

Attention Sensitivity Is Not Enough: Dissociating Attention-Level and Behavioural In-Context Learning under Fine-Tuning

Jinyuan Zhang, Peng He, He Hu, Yin Yuan, ShengShuo Jiao

arXiv:2609.00064v1Aug 30, 2026cs.LGcs.AIcs.CL

TL;DR

The paper asks whether attention changes remain a trustworthy proxy for behavioural in-context learning when that proxy is optimised. It formalises paired attention and behavioural diagnostics, then shows in a controlled Llama-2-7B ablation that ICS nearly reaches its ceiling while ICL-GAP stays near zero and MMLU declines. The results motivate behavioural validation and examine gating and pretrained anchoring as partial safeguards.

Problem
Attention-based diagnostics are attractive for monitoring or preserving ICL, but it is unclear whether they remain faithful when used as optimisation targets.
Method
The paper pairs geometrically bounded matched-versus-mismatched attention sensitivity, ICS, with behavioural accuracy gap, ICL-GAP, and tests ICS maximisation plus gated and pretrained-anchored variants.
Results
ICS reaches 1.413, within 0.5 percent of its geometric ceiling, while ICL-GAP remains near zero and MMLU falls from 0.371 to 0.279.
Takeaways & Limitations
Attention-level ICL proxies should be validated against behavioural gaps before being used as training objectives.
Takeaways & Limitations
The evidence is limited to one base model, one probe design, one demonstration count, and mostly single-seed baselines.

Abstract

from arXiv · show

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 \emph{In-Context Sensitivity} (ICS), the average row distance between last-token attention on matched and mismatched demonstration prefixes, and pair it with \emph{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.

1 Introduction

The paper tests whether attention changes remain a trustworthy proxy for behavioural ICL when used as an optimisation target. In a controlled Llama-2-7B ablation, ICS nearly reaches its ceiling while ICL-GAP remains near zero and MMLU declines.

Motivation: ICL adapts models to new tasks from demonstrations, but fine-tuning can weaken this behaviour.Attention divergences are inexpensive and differentiable, yet they do not by themselves establish improved task behaviour.
Diagnostic framework: ICS measures matched-versus-mismatched last-token attention distance, while ICL-GAP measures the corresponding held-out accuracy difference.ICS is geometrically bounded; ICL-GAP is the behavioural quantity the preservation method aims to keep positive.
Main result: 1.413 ICS, within 0.5 percent of the ceiling, coexists with near-zero ICL-GAP and MMLU falling from 0.371 to 0.279 under ICS maximisation.The controlled four-arm ablation uses Llama-2-7B and 5,000 training steps from a shared checkpoint.
Main result: The 2.74× increase in attention divergence does not translate into stronger behavioural context use.This dissociation motivates validating attention-level diagnostics against behavioural gaps before optimisation.
Contributions: The paper combines bounded ICS with behavioural validation and tests structural, gated, and pretrained-anchored responses to proxy failure.Endpoint statistics and a structural construction examine how sharp, disjoint routing can maximise ICS without selecting answer-relevant tokens.

2 Related Work

The paper situates its diagnostic at the intersection of mechanistic accounts of ICL, behavioural studies of fine-tuning drift, and concerns about attention as evidence. It frames ICS maximisation as an attention-level instance of surrogate over-optimisation.

ICL mechanisms: Mechanistic ICL research links adaptation to induction heads and attention-based updates, while other work finds apparent ICL can reflect memorisation or label priors.These perspectives motivate distinguishing genuine task adaptation from superficial context sensitivity.
ICL under fine-tuning: Fine-tuning studies primarily track ICL through held-out accuracy or matched-versus-mismatched behavioural gaps.This paper adds an attention-level diagnostic alongside that behavioural family.
Attention as evidence: Attention maps are cheap and differentiable evidence, but attention weights need not track the computation behind a prediction.The paper tests this concern under optimisation rather than only interpreting a fixed model.
Component-level attribution: The controlled component arms repurpose attention-only and MLP-only tuning to ask which subnetwork carries ICL drift; neither alone changes ICS substantially in this setup.The comparison complements mechanistic circuit localisation and parameter-efficient tuning.
Auxiliary objectives: ICS reaches within 0.5 percent of its geometric ceiling while behavioural ICL remains near zero, giving Goodhart-style surrogate failure an attention-level instance.The comparison extends familiar concerns about optimising measurable proxies beyond reward-model settings.
Preservation methods: The paper covers freezing and anchoring families of ICL preservation, while replay is discussed but not completed in the main comparison.AnchorTune is described as anchoring the trainee to its own pretrained checkpoint.

3 Methodology

The methodology contrasts matched and mismatched demonstration prefixes using held-out attention and behavioural probes, then deliberately maximises ICS to test whether proxy gains preserve ICL. It also evaluates behavioural gating and pretrained attention anchoring.

Framework: The diagnostic framework separates attention-level context sensitivity from behavioural ICL effects using matched and mismatched prompts.The probes are computed on held-out pairs disjoint from regulariser pairs.
Problem setting: A matched prefix preserves input-label correspondence, whereas a mismatched prefix preserves the token template while breaking that correspondence through label permutation.The same contrast underlies both attention and behavioural measurements.
Diagnostic metrics: ICS averages distances between last-token attention rows across matched and mismatched prefixes over a probe distribution and layer band.For Llama-2-7B, the reported band is layers 8–16.
Diagnostic metrics: The attention-distance geometry reaches equality only when matched and mismatched rows are one-hot on disjoint coordinates.This establishes the proxy’s fixed geometric ceiling and its disjoint-support extremum.
Behavioural validation: ICL-GAP is the matched-prefix accuracy minus mismatched-prefix accuracy, with MMLU accuracy and ECE reported as additional outcomes.The central test is whether high ICS implies positive ICL-GAP under optimisation pressure.
Proxy-maximisation stress test: The stress test adds an ICS-maximising auxiliary objective to cross-entropy on held-out regulariser pairs.The design treats proxy improvement without behavioural improvement as direct evidence of Goodhart exposure.
Constructive variants: B-ICS gates the attention proxy with a differentiable behavioural sentinel, while AnchorTune anchors matched-prefix attention rows to the pretrained model.The soft behavioural gap is only a differentiable surrogate and is not claimed to agree pointwise with accuracy gaps.
Constructive variants: Unlike divergence maximisation, AnchorTune is minimised at the pretrained attention pattern rather than along the disjoint-support ridge.This encodes preservation of pretrained computation instead of maximisation of attention divergence.

4 Experiments and Analysis

The experiments show that ICS can be driven near its geometric ceiling without producing behavioural ICL, while attention routes increasingly toward non-label tokens. Behaviour-gated objectives only partly mitigate this dissociation, whereas anchoring to pretrained computation preserves stronger general ability.

RQ0: All arms retain a matched-versus-mismatched behavioural effect on the random-label protocol, ruling out universal probe insensitivity.The authors keep this validation separate from the harder controlled QA probe, whose ICL-GAP remains near zero.
RQ2: The 2.74× increase in attention divergence does not yield behavioural preservation: ICL-GAP ends at −0.010 after fluctuating around zero during training.Its logged trajectory has mean 0.005 and standard deviation 0.020 across 50 checkpoints.
RQ1: ICS reaches 1.413, within 0.5% of its geometric ceiling, while controlled QA ICL-GAP remains near zero and MMLU degrades.The controlled ablation starts from the same checkpoint and trains for 5,000 steps; unregularised arms remain near the pretrained ICS.
RQ2: Near-ceiling ICS produces sharp, nearly disjoint attention, but top-1 routing shifts from labels toward formatting and demonstration-body tokens.Top-1 concentration rises from 0.18 to 0.94, support overlap falls from 61% to 4%, and label-token routing falls from 58% to 26%.
RQ2: The divergence objective is indifferent to semantic value, allowing attention-polarised solutions whose selected tokens carry no answer information and whose behavioural ICL can remain absent.The constructive argument gives useful and behaviourally vacuous disjoint-support solutions the same proxy score.
RQ3: Mild behavioural gating only partially mitigates the effect, while objectives anchored to pretrained computation remain in the high-MMLU, moderate-ICS region.The comparison supports anchoring as the safer family in this setup, without establishing AnchorTune as uniquely superior.

5 Discussion and Limitations

Attention divergences remain useful diagnostics when read from an unoptimized model, but optimizing them requires behavioural validation. The evidence is limited in model, probe, demonstration, and seed coverage.

Attention divergences should not be discarded, but using them as objectives requires a behavioural guard such as ICL-GAP.
The evidence covers one base model, one probe design, one demonstration count, and mostly single-seed baselines.

6 Conclusion

The paper shows that an attention diagnostic can improve under fine-tuning without preserving the behaviour it is intended to measure. Behavioural gating partially offsets saturation, while anchored objectives remain closer to pretrained computation in this setup.

ICS reaches within 0.5% of its geometric ceiling while controlled QA ICL-GAP remains near zero and MMLU falls from 0.371 to 0.279.
Endpoint analysis links the dissociation to sharp, disjoint routing toward semantically weak tokens.
Mild behavioural gating partially offsets saturation, whereas anchored objectives stay closer to the pretrained operating region.
Training trajectories report ICS, MMLU accuracy, and ECE every 500 steps, with ICL-GAP during training only for F_ICS-MAX.

B Contamination Postmortem (F_ICS-MAXv1)

The contamination postmortem identified shared training and evaluation data in the first F_ICS-MAX run. The corrected rerun separated loaders and used genuinely held-out evaluation data.

F_ICS-MAXv1 used one ICLEvalDataLoader for both regularization pairs and ICS evaluation pairs.
The fix split the regularizer and evaluation loaders, added answer_idx, and wired evaluate_icl_gap into the trainer loop.
Table 3 reports full controlled proxy-stress trajectories sampled every 500 steps.

C Probe-Layer Ablation

Probe-layer checks show that ICS saturation is not confined to the original mid-network band. The unregularized checkpoints retain near-zero behavioural gaps, while the pilot is not quantitatively comparable to the controlled experiment.

Alternative early, mid-late, and late probe bands produce ICS values of 1.401, 1.418, and 1.396.
All three bands lie within 0.5% of the original 1.413 and within 1.5% of the geometric ceiling.
F_A, F_M, and F_NONE each finish within ±0.04 of zero on ICL-GAP.
The pilot used different data and hyperparameters, so its endpoint is excluded from quantitative comparisons with the controlled experiment.

F Controlled Proxy-Stress Training Details

The controlled stress test specifies separate probe, training, and evaluation sets, then examines whether smooth behavioural gating suppresses proxy optimisation failures.

Evaluation design: 500 paired examples form the ICS/ICL-GAP evaluation set, while a separate 500-pair set is used by the ICS-maximising regulariser.The evaluation pairs use 4-shot matched or mismatched prefixes; MMLU and GSM8K are separate capability probes.
Evaluation design: MMLU uses 240 items across 57 subjects, GSM8K uses 50 items with 8-shot chain-of-thought prompts, and ECE uses 15 equal-mass MMLU bins.
Behavioural validation: Matched demonstrations outperform swapped demonstrations with paired bootstrap confidence intervals excluding zero across all three random-label screening runs.The screening uses additional single-token label pairs to rule out a single label-pair artefact.
Behavioural gating: The smooth B-ICS gate substitutes a differentiable soft gap for the non-differentiable discrete ICL-GAP, but this substitution is explicitly heuristic rather than a theorem.
Behavioural gating: For β = 10 and δ = 0.05, the smooth gate bounds the ICS-only gradient component by σ(−0.5) ≈ 0.378 on the bad ridge, while the gap-only component can select directions away from it.The bound is described as loose because both terms are bounded by their worst-case product.

K AnchorTune: Memory and Compute

AnchorTune uses a frozen pretrained reference and cached attention targets, creating substantial memory overhead but only a modest per-step compute increase.

Memory: The frozen Llama-2-7B reference occupies approximately 13GB, while trainee weights, gradients, and AdamW state bring the pre-activation total to approximately 65GB.
Memory: Cached anchor attentions add approximately 370MB for 500 pairs, 9 layers, 32 heads, and 640-token sequences.
Compute: AnchorTune adds one trainee forward pass per optimiser step and runs at approximately 5 seconds per step versus approximately 4 seconds for unregularised fine-tuning.The reported overhead is 25% relative to unregularised fine-tuning and parity with F_ICS-MAX.
Practical alternatives: Quantising the reference or storing it on CPU with a disk cache are proposed alternatives for reducing in-memory reference-model overhead.The CPU-resident option raises startup time from 3 minutes to approximately 15 minutes on the reported hardware.
Experimental scope: The constructive section evaluates behaviour gating, AnchorTune’s λ-sweep, Pareto comparisons, and anchor-target ablations as responses to the observed dissociation.

L.1 Behaviour-Gated Sweep

Smooth behaviour gating only partially mitigates divergence-driven proxy saturation, whereas anchored objectives preserve capability performance and remain near the pretrained diagnostic regime.

Behaviour-gated sweep: F_BICS reaches ICS = 1.397, MMLU = 0.283, and ICL-GAP = +0.020, compared with F_ICS-MAX’s ICS = 1.413, MMLU = 0.279, and ICL-GAP = −0.010.The behavioural advantage is +0.030, but MMLU degradation remains essentially unchanged.
Behaviour-gated sweep: The smooth β = 10, δ = 0.05 gate remains too weak as a training target, motivating stronger gating or a hard indicator with a straight-through estimator.
Behaviour-gated sweep: The diagnostic B-ICS guard returns near-zero for F_ICS-MAX because its ICL-GAP is below δ, while remaining near-pretrained for F_ANCHOR.
AnchorTune sweep: At λ⋆ = 0.05, three-seed AnchorTune gives MMLU = 0.370 ± 0.058, ICS = 0.500 ± 0.001, and ICL-GAP = −0.012 ± 0.017.These values are reported as within seed variation of the pretrained model.
Anchored comparisons: All three anchored arms recover MMLU within seed noise and keep ICS near the pretrained range, while divergence regularisers occupy the low-MMLU region of the Pareto plot.LOGITANCHOR anchors predictions, WEIGHTANCHOR anchors weights, and F_ANCHOR anchors mid-band attention rows.

L.4 Anchor-Target Ablation

The anchor-target ablation weakens the claim that mid-band attention rows are uniquely responsible for MMLU preservation, while preserving the distinction between divergence and anchored objectives.

Layer ablation: Anchoring only the last four layers gives MMLU = 0.358 versus 0.367 for full F_ANCHOR, a 0.009-point reduction within seed noise.
Component ablation: The LM-head-only ℓ2 variant reaches MMLU = 0.375, matching the pretrained reference of 0.371 within seed noise without anchoring attention rows.
Interpretation: The results indicate that proximity to the pretrained checkpoint in weights, the LM head, or attention rows can preserve MMLU under this fine-tuning configuration.
Interpretation: AnchorTune’s residual distinction is that its anchor target coincides with the diagnostic ICL signal, rather than producing the largest absolute MMLU value.
Constructive summary: Across the constructive experiments, anchored regularisers maintain MMLU while divergence targets move toward high ICS and low MMLU; AnchorTune recovers 99% of the F_ICS-MAX-to-F_NONE MMLU gap.