Résumé Technique : Le Score de Contribution aux Logits Identifie les Têtes de Récupération Non Littérales

1. Énoncé du Problème

Les grands modèles de langage (LLM) synthétisent fréquemment des réponses à partir de la signification sémantique d'un segment de contexte pertinent plutôt qu'en copiant littéralement les jetons (tokens). Bien que les travaux actuels de l'interprétabilité mécaniste aient identifié des « têtes de récupération » (retrieval heads) responsables de l'accès au contexte, les méthodes de détection actuelles sont fondamentalement limitées à la récupération littérale.

Les détecteurs existants (par exemple, Wu et al., 2025 ; Fu et al., 2025) reposent sur des observables de poids d'attention. Ils récompensent les têtes où le jeton source auquel l'attention est portée correspond au jeton de la réponse générée. Ce critère de « correspondance de jetons » identifie avec succès les têtes qui effectuent un copier-coller d'informations (récupération littérale), mais échoue à détecter les têtes qui effectuent une récupération non littérale. Dans les scénarios non littéraux, un modèle peut prêter attention à une expression source (par exemple, « Tour Eiffel ») mais synthétiser un jeton de réponse différent (par exemple, « Yuki ») via son circuit de sortie-valeur (OV). Le score basé sur l'attention capture où une tête lit, mais ignore ce qu'elle écrit, manquant ainsi le mécanisme critique par lequel la synthèse non littérale se produit.

2. Méthodologie : Score de Contribution aux Logits (LOCOS)

Les auteurs introduisent le Logit-Contribution Scoring (LOCOS), un détecteur « conscient de l'écriture » (write-aware) qui évalue les têtes d'attention en fonction de leur contribution au logit de la réponse finale, plutôt que de leur distribution d'attention.

Mécanisme Central

LOCOS décompose l'opération d'une tête d'attention en deux circuits :

1. Circuit QK : Détermine où la tête lit (poids d'attention $\alpha$).
2. Circuit OV : Détermine ce que la tête écrit dans le flux résiduel (projection de sortie $W_O v$).

La méthode évalue une tête $(l, h)$ en projetant la sortie de son circuit OV sur la direction de désincorporation (unembedding) du jeton de la réponse correcte ($u_{y_t}$).

La Procédure en Trois Étapes

1. Contribution par Position aux Logits :
   Pour une étape de décodage $t$ générant le jeton $y_t$, la contribution d'une position source $j$ à travers la tête $(l, h)$ est calculée comme suit :
   $$\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}$$
   Ce scalaire $\phi$ capture l'ensemble du pipeline : une tête peut avoir une attention élevée ($\alpha$) mais une contribution proche de zéro si sa sortie OV est orthogonale à la direction de la réponse. Inversement, une tête avec une attention modérée peut avoir une contribution élevée si son circuit OV transforme la représentation attendue en un vecteur aligné avec la réponse.

2. Contraste Spatial :
   Pour isoler les signaux spécifiques à la récupération, LOCOS compare les contributions des positions « aiguille » (needle — le segment de contexte pertinent $[s_\tau, e_\tau)$) par rapport aux positions « hors-aiguille » (off-needle).
   $$\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}$$
   Le facteur de redimensionnement garantit que la somme hors-aiguille est comparable à la somme de l'aiguille en termes de longueur. Le score pour une étape unique est la différence : $\Phi^{(l,h),+}_t - \Phi^{(l,h),-}_t$. Ce contraste spatial (aiguille vs hors-aiguille au sein d'une même étape) remplace le contraste temporel (étape de réponse vs étape de non-réponse) utilisé dans les travaux antérieurs.

3. Agrégation :
   Le score final $S_{l,h}$ est la moyenne de ces contrastes spatiaux à travers toutes les étapes de décodage où le modèle génère une réponse correcte, agrégée sur tous les essais réussis d'un benchmark. Les scores négatifs indiquent des têtes dont les contributions au logit de la réponse proviennent principalement de positions hors-aiguille.

3. Principales Contributions

• Détection Sensible à l'Écriture : L'article établit que l'identification de la récupération non littérale nécessite de mesurer la sortie du circuit OV, et non seulement les poids d'attention.
• Algorithme LOCOS : Une méthode en un seul passage direct (single-forward-pass) qui évalue les têtes par leur contribution projetée au logit de la réponse, en contrastant les sources aiguille et hors-aiguille.
• Validation Causale : À travers des expériences d'ablation par la moyenne, les auteurs démontrent que LOCOS identifie des têtes qui sont causalement critiques pour la récupération non littérale, alors que les bases de référence basées sur l'attention échouent à le faire.

4. Résultats Expérimentaux

La méthode a été évaluée sur le benchmark NoLiMa (une tâche de recherche d'aiguille dans une botte de foin non littérale) à travers trois familles de modèles : Qwen3 (8B, 14B, 32B), Gemma-3 (12B, 27B) et OLMo-3.1 (32B).

Performance de l'Ablation (NoLiMa)

• Dégradation Supérieure : L'ablation de la top-$k$ des têtes LOCOS provoque un déclin plus marqué du ROUGE-L par rapport à toutes les bases de référence basées sur l'attention (Wu/NIAH-scored, Wu/NoLiMa-scored, et têtes aléatoires).
• Étude de Cas Qwen3-8B : L'ablation des 50 meilleures têtes LOCOS fait s'effondrer le ROUGE-L de 0,401 à 0,000. En revanche, la base de référence d'attention la plus forte (Wu/NIAH) conserve un ROUGE-L de 0,292 à la même profondeur.
• Spécificité : L'ablation des têtes LOCOS dégrade sévèrement la récupération contextuelle tout en laissant intactes la récupération paramétrique (ex: associations ville-pays) et le raisonnement arithmétique. Cela confirme que les têtes identifiées sont spécifiques à la récupération, et non génériquement importantes pour la sortie du modèle.

Spécificité Littérale vs Non Littérale

• L'ablation des top-$k$ têtes LOCOS dégrade à la fois la récupération non littérale (NoLiMa) et littérale (NIAH).
• Cependant, la chute est nettement plus forte sur NoLiMa. Cela indique que si LOCOS capture des têtes pertinentes pour les deux régimes, elle identifie de manière unique le sous-ensemble de têtes responsables de la synthèse non littérale que les détecteurs de correspondance de jetons manquent.

Transfert en Aval

• Sur les benchmarks de contexte long en aval (MuSiQue et BABILong), l'ablation des top-50 têtes LOCOS dégrade significativement les performances, particulièrement sur la famille Qwen3 (par exemple, la précision sur MuSiQue pour Qwen3-8B chute de 0,55 à 0,08).
• Le classement de LOCOS par rapport aux bases de référence basées sur l'attention varie selon le benchmark, suggérant que les tâches de QA multi-étapes recrutent à la fois des circuits de récupération littérale et non littérale.

Distribution des Couches

• Les scores LOCOS sont fortement concentrés dans les couches tardives pour la famille Qwen3 et Gemma-3-27B.
• Les auteurs notent que cette concentration peut être influencée par l'« hypothèse du chemin direct » (les têtes proches de la sortie ont un impact plus direct sur les logits), mais des expériences de patch causal suggèrent que le motif des couches tardives n'est pas un simple artefact.

5. Signification et Revendications

L'article affirme que LOCOS fournit la première identification mécaniste des têtes de récupération non littérale.

• Limitation des Travaux Antérieurs : Les détecteurs existants sont « aveugles » à la récupération non littérale car ils récompensent les têtes où le jeton attendu correspond au jeton généré. Cela manque les têtes qui prêtent attention à une source et synthétisent une réponse différente.
• Aperçu Mécaniste : En mesurant la projection du circuit OV sur la direction de la réponse, LOCOS capture l'opération d'« écriture » essentielle à la synthèse.
• Utilité Pratique : Les têtes identifiées sont causalement critiques pour le raisonnement en contexte long. Cette découverte a des implications pour :
  • L'interprétabilité : Fournir une carte plus précise de la manière dont les LLM accèdent au contexte et le synthétisent.
  • L'Efficacité : Permettre une compression plus ciblée du cache KV en préservant uniquement les têtes spécifiques responsables de la récupération, plutôt que d'utiliser des heuristiques uniformes ou basées sur l'attention.
  • Atténuation des Hallucinations : Offrir un mécanisme pour identifier et intervenir sur les têtes qui échouent à récupérer correctement le contexte pertinent.

Les auteurs restent modestes quant à la généralisation, notant que leurs conclusions sont spécifiques aux architectures testées (transformers standard decoder-only) et au benchmark NoLiMa. Ils ne prétendent pas que LOCOS identifie tous les mécanismes de récupération, mais plutôt qu'elle isole avec succès le sous-ensemble manqué par les critères de correspondance de jetons.