Technisch Samenvatting: Logit-Contribution Scoring Identificeert Non-Literaal Retrieval-Heads

1. Probleemstelling

Large Language Models (LLMs) synthetiseren antwoorden vaak op basis van de semantische betekenis van een relevante contextspan, in plaats van door letterlijk tokens te kopiëren. Hoewel bestaand werk op het gebied van mechanistische interpreteerbaarheid "retrieval heads" heeft geïdentificeerd die verantwoordelijk zijn voor contexttoegang, zijn huidige detectiemethoden fundamenteel beperkt tot letterlijke retrieval.

Bestaande detectoren (bijv. Wu et al., 2025; Fu et al., 2025) vertrouwen op attention-weight observables. Zij belonen heads waar de geattendde bron-token overeenkomt met de gegenereerde antwoord-token. Dit "token-matching" criterium identificeert succesvol heads die informatie kopiëren en plakken (literaal retrieval), maar faalt in het detecteren van heads die non-literaal retrieval uitvoeren. In non-literaal scenario's kan een model naar een bronfrase attenderen (bijv. "Eiffel Tower"), maar een ander antwoord-token synthetiseren (bijv. "Yuki") via zijn output-value (OV) circuit. Attention-gebaseerde scoring legt vast waar een head leest, maar negeert wat het schrijft, waardoor het cruciale mechanisme waardoor non-literaal synthese plaatsvindt, wordt gemist.

2. Methodologie: Logit-Contribution Scoring (LOCOS)

De auteurs introduceren Logit-Contribution Scoring (LOCOS), een "write-aware" detector die attention heads evalueert op basis van hun bijdrage aan de uiteindelijke antwoord-logit, in plaats van hun attention distributie.

Kernmechanisme

LOCOS ontleedt de operatie van een attention head in twee circuits:

1. QK Circuit: Bepaalt waar de head leest (attention weights $\alpha$).
2. OV Circuit: Bepaelt wat de head schrijft naar de residual stream (output projectie $W_O v$).

De methode scoort een head $(l, h)$ door de output van de OV-circuit te projecteren op de unembedding richting van de correcte antwoord-token ($u_{y_t}$).

De Drie-Stappen Procedure

1. Per-Positie Logit Bijdrage:
   Voor een decoding stap $t$ die token $y_t$ genereert, wordt de bijdrage van een bron-positie $j$ via head $(l, h)$ berekend als:
   $$\phi^{(l,h)}_{t,j} = u_{y_t}^\top o^{(l,h)}_{t,j} = \alpha^{(l,h)}_{t,j} \cdot u_{y_t}^\top W^{(l,h)}_O v^{(l,h)}_{t,j}$$
   Deze scalaire $\phi$ vangt de volledige pipeline: een head kan een hoge attention ($\alpha$) hebben maar een bijna nul bijdrage als de OV-output orthogonaal is aan de antwoord-richting. Omgekeerd kan een head met een gematigde attention een hoge bijdrage hebben als de OV-circuit de geattendde representatie transformeert naar een antwoord-gealigneerde vector.

2. Ruimtelijke Contrast:
   Om retrieval-specifieke signalen te isoleren, vergelijkt LOCOS bijdragen van needle posities (de relevante contextspan $[s_\tau, e_\tau)$) met off-needle posities.
   $$\Phi^{(l,h),+}_t = \sum_{j \in \text{needle}} \phi^{(l,h)}_{t,j}$$
   $$\Phi^{(l,h),-}_t = \frac{e_\tau - s_\tau}{N_t - (e_\tau - s_\tau)} \sum_{j \notin \text{needle}} \phi^{(l,h)}_{t,j}$$
   De herschalingsfactor zorgt ervoor dat de off-needle som vergelijkbaar is met de needle som qua lengte. De score voor een enkele stap is het verschil: $\Phi^{(l,h),+}_t - \Phi^{(l,h),-}_t$. Dit ruimtelijke contrast (needle vs. off-needle binnen een enkele stap) vervangt het temporele contrast (antwoord vs. niet-antwoord stappen) dat in eerder werk werd gebruikt.

3. Aggregatie:
   De definitieve score $S_{l,h}$ is het gemiddelde van deze ruimtelijke contrasten over alle decoding stappen waar het model een correct antwoord genereert, geaggregeerd over alle succesvolle trials in een benchmark. Negatieve scores duiden op heads waarvan de bijdragen aan de antwoord-logit voornamelijk afkomstig zijn van off-needle posities.

3. Belangrijkste Bijdragen

• Write-Aware Detectie: Het paper stelt vast dat het identificeren van non-literaal retrieval vereist dat de output van de OV-circuit wordt gemeten, niet alleen de attention weights.
• LOCOS Algoritme: Een single-forward-pass methode die heads scoort door hun geprojecteerde bijdrage aan de antwoord-logit, waarbij needle en off-needle bronnen worden gecontrasteerd.
• Causale Validatie: Door middel van mean-ablation experimenten demonstreren de auteurs dat LOCOS heads identificeert die causaal kritiek zijn voor non-literaal retrieval, terwijl attention-gebaseerde baselines hierin falen.

4. Experimentele Resultaten

De methode werd geëvalueerd op de NoLiMa benchmark (een non-literaal needle-in-a-haystack taak) over drie modelfamilies: Qwen3 (8B, 14B, 32B), Gemma-3 (12B, 27B), en OLMo-3.1 (32B).

Ablatie Prestaties (NoLiMa)

• Superieure Degradatie: Het mean-ablateren van de top-$k$ LOCOS heads veroorzaakt een steilere daling in ROUGE-L vergeleken met alle attention-gebaseerde baselines (Wu/NIAH-scored, Wu/NoLiMa-scored, en random heads).
• Qwen3-8B Case Study: Het ablateren van de top 50 LOCOS heads laat de ROUGE-L instorten van 0.401 naar 0.000. In contrast hiermee behoudt de sterkste attention-gebaseerde baseline (Wu/NIAH) een ROGUE-L van 0.292 bij dezelfde diepte.
• Specificiteit: Het ablateren van LOCOS heads degradeert de contextuele retrieval ernstig, terwijl parametrische recall (bijv. stad-land associaties) en arithmetische redenering intact blijven. Dit bevestigt dat de geïdentificeerde heads retrieval-specifiek zijn, en niet generiek belangrijk voor de modeloutput.

Literaal vs. Non-Literaal Specificiteit

• Het ablateren van top-$k$ LOCOS heads degradeert zowel non-literaal (NoLiMa) als literaal (NIAH) retrieval.
• Echter, de daling is aanzienlijk steiler op NoLiMa. Dit geeft aan dat hoewel LOCOS heads vangt die relevant zijn voor beide regimes, het uniek de subset van heads identificeert die verantwoordelijk zijn voor non-literaal synthese die door token-matching detectoren gemist wordt.

Downstream Transfer

• Op downstream long-context benchmarks (MuSiQue en BABILong) degradeert het ablateren van top-50 LOCOS heads de prestaties aanzienlijk, met name op de Qwen3 familie (bijv. MuSiQue accuratesse op Qwen3-8B daalt van 0.55 naar 0.08).
• De rangschikking van LOCOS vs. attention-gebaseerde baselines varieert per benchmark, wat suggereert dat multi-hop QA taken zowel literaal als non-literaal retrieval circuits rekruteren.

Laag Verdeling

• LOCOS scores zijn sterk geconcentreerd in late lagen voor de Qwen3 familie en Gemma-3-27B.
• De auteurs merken op dat deze concentratie beïnvloed kan zijn door de "direct-path assumption" (heads dichter bij de output hebben een directere impact op logits), maar causale patching experimenten suggeren dat het late-layer patroon geen louter artefact is.

5. Betekenis en Claims

Het paper claimt dat LOCOS de eerste mechanistische identificatie biedt van non-literaal retrieval heads.

• Beperking van Voorgaand Werk: Bestaande detectoren zijn "blind" voor non-literaal retrieval omdat ze heads belonen waar de geattendde token overeenkomt met de gegenereerde token. Dit mist heads die naar een bron attenderen en een ander antwoord synthetiseren.
• Mechanistisch Inzicht: Door de projectie van de OV-circuit op de antwoord-richting te meten, vangt LOCOS de "write" operatie die essentieel is voor synthese.
• Praktisch Nut: De geïdentificeerde heads zijn causaal kritiek voor long-context redeneren. Deze bevinding heeft implicaties voor:
  • Interpreteerbaarheid: Het bieden van een nauwkeurigerere kaart van hoe LLM's context benaderen en synthetiseren.
  • Efficiëntie: Het mogelijk maken van meer gerichte KV cache compressie door alleen de specifieke heads te behouden die verantwoordelijk zijn voor retrieval, in plaats van uniforme of attention-gebaseerde heuristieken.
  • Mitigatie van Hallucinaties: Het bieden van een mechanisme om te identificeren en in te grijpen op heads die er niet in slagen de relevante context correct te extraheren.

De auteurs blijven bescheiden over de generalisatie en merken op dat hun bevindingen specifiek zijn voor de geteste architecturen (standaard decoder-only transformers) en de NoLiMa benchmark. Ze beweren niet dat LOCOS alle retrieval mechanismen identificeert, maar wel dat het succesvol de subset isoleert die gemist wordt door token-matching criteria.