Résumé technique : Mécanismes de mélange dans la récupération d'entités en contexte

Énoncé du problème

Les modèles de langage (LM) excellent dans le raisonnement en contexte, une capacité fondamentalement dépendante de la capacité à lier des entités apparentées au sein d'un texte (par exemple, lier « Ann » à « tarte » dans « Ann aime la tarte ») et à les récupérer ensuite sur la base d'une requête (par exemple, répondre à « Qui aime la tarte ? »). Alors que des recherches antérieures ont établi que les LM utilisent un mécanisme positionnel — récupérant des entités en fonction de leur indice dans la séquence de contexte — cette vision a été largement dérivée d'études impliquant de courtes listes de groupes d'entités (par exemple, $n=2$ ou $3$).

Le problème central abordé dans ce travail est la fiabilité de ce mécanisme positionnel dans des contextes plus complexes et réalistes. À mesure que le nombre de groupes d'entités liées augmente, les LM éprouvent des difficultés à raisonner sur de longs contextes, présentant des phénomènes tels que l'effet « perdu au milieu ». Les auteurs examinent si le mécanisme positionnel se généralise à des séquences plus longues ou si les LM emploient des stratégies alternatives pour maintenir la précision de la récupération.

Méthodologie

Les auteurs utilisent des techniques d'interprétabilité mécaniste, spécifiquement des interventions d'échange, pour sonder la structure causale interne des LM.

1. Configuration de la tâche : L'étude utilise des tâches de raisonnement en contexte templatiques impliquant des « Groupes d'entités » (tuples d'entités remplissant des rôles spécifiques, par exemple Nom + Nourriture). Le modèle est interrogé pour récupérer une entité cible sur la base d'une entité requête au sein d'un groupe spécifique.
2. Conception contrefactuelle : Pour distinguer entre des hypothèses concurrentes, les auteurs construisent des jeux de données appariés d'entrées originales et contrefactuelles. Ces paires sont conçues de telle sorte que trois mécanismes hypothétiques — Positionnel, Lexical et Réflexif — prédiraient des sorties distinctes si une intervention d'échange était effectuée sur leurs variables intermédiaires respectives.
   • Mécanisme positionnel ($P$) : Récupère l'entité située à l'indice du groupe de la requête contrefactuelle.
   • Mécanisme lexical ($L$) : Récupère l'entité liée à l'entité requête contrefactuelle (par exemple, si la requête change de « Tim » à « Ann », il récupère ce qu'« Ann » aime).
   • Mécanisme réflexif ($R$) : Récupère l'entité cible via un pointeur direct émanant de la cible elle-même, déréférencé par la requête.
3. Protocole d'intervention : Les auteurs effectuent des interventions d'échange sur le flux résiduel du dernier jeton à des couches spécifiques ($\ell$) identifiées comme le point juste avant que la récupération ne se produise. Ils patchent le vecteur du flux résiduel de l'exécution contrefactuelle dans l'exécution originale et observent la distribution du prochain jeton résultante.
4. Modélisation causale : Un modèle causal ($M$) est développé pour simuler le comportement du LM. Ce modèle combine les trois mécanismes sous forme de mélange pondéré :
   • Le terme positionnel est modélisé comme une distribution gaussienne centrée sur l'indice de la requête, avec une variance dépendante de l'indice.
   • Les termes lexical et réflexif sont modélisés comme des distributions one-hot.
   • Le modèle est entraîné pour prédire la distribution de logits du LM en utilisant la divergence de Jensen–Shannon (JSD) comme métrique de perte.
5. Portée : Les expériences couvrent neuf modèles répartis sur trois familles (Llama, Gemma, Qwen) allant de 2 à 72 milliards de paramètres, évalués sur dix tâches de liaison distinctes.

Contributions clés

1. Identification de trois mécanismes distincts

L'article remet en question la vision prévalente selon laquelle la récupération est uniquement positionnelle. Il démontre que les LM emploient un mélange de trois mécanismes :

• Positionnel : Efficace pour les entités au début et à la fin du contexte, mais devient bruyant et diffus pour les positions intermédiaires à mesure que le nombre de groupes ($n$) augmente.
• Lexical : Récupère le contrepartie liée de l'entité requête (par exemple, « tarte » $\to$ « Ann »). Ce mécanisme fournit un signal net et one-hot.
• Réflexif : Utilise un pointeur direct de l'entité cible vers elle-même, qui est récupéré via la requête. Ceci est particulièrement crucial lorsque la cible apparaît avant la requête dans une clause (par exemple, « Tim aime le thé » $\to$ « Qui aime le thé ? »), un scénario où la copie lexicale est architecturalement difficile pour les modèles auto-régressifs.

2. Le motif de fiabilité en « U »

Grâce à des expériences d'intervention extensives, les auteurs découvrent un motif cohérent :

• Premières et dernières positions : La récupération est dominée par le mécanisme positionnel, qui est hautement concentré et fiable.
• Positions intermédiaires : Le signal positionnel se dégrade considérablement. Pour compenser, les LM s'appuient fortement sur les mécanismes lexical et réflexif, qui fournissent des signaux plus nets pour désambiguïser l'entité correcte.
• Interaction : Les mécanismes présentent une « synergie compétitive ». Lorsque les indices s'alignent, les signaux s'amplifient ; lorsqu'ils divergent, ils se suppriment mutuellement.

3. Modèle causal de la prédiction du prochain jeton

Les auteurs construisent un modèle causal qui combine ces trois mécanismes. Ce modèle atteint un accord de 95 % (mesuré par la similarité de Jensen–Shannon) avec les distributions réelles de prochains jets des LM à travers diverses tâches et tailles de modèles. En revanche, les modèles s'appuyant uniquement sur le mécanisme positionnel (la vision prévalente) obtiennent des scores significativement plus bas (~45 %), performant moins bien qu'une base uniforme dans certaines configurations.

4. Généralisation aux entrées naturalistes

L'étude valide que ces résultats tiennent dans des contextes plus naturalistes où les groupes d'entités sont entrelacés avec du texte « remplisseur » (phrases ouvertes). Bien que l'ajout de jetons de remplissage augmente le bruit et affaiblisse légèrement le mécanisme lexical, la dépendance fondamentale au mélange de trois mécanismes persiste, offrant une explication mécaniste de l'effet « perdu au milieu ».

Résultats

• Prévalence des mécanismes : Dans les positions intermédiaires de longs contextes, le mécanisme positionnel ne représente que ~20 % du comportement du modèle, tandis que les mécanismes lexical et réflexif deviennent les principaux moteurs de la récupération correcte.
• Robustesse des modèles : Le motif de mélange est cohérent à travers les familles de modèles (Llama, Gemma, Qwen) et les échelles (de 2 à 72 milliards de paramètres).
• Validation réflexive : Grâce à des conceptions contrefactuelles spécifiques où la réponse contrefactuelle n'existe pas dans le contexte original, les auteurs prouvent l'existence d'un mécanisme de pointeur réflexif distinct du jeton de réponse lui-même.
• Études d'ablation : Le retrait d'un seul mécanisme du modèle causal dégrade considérablement les performances, confirmant que les trois sont nécessaires pour une abstraction précise du comportement des LM.

Signification et affirmations

L'article prétend établir une image plus complète de la manière dont les LM lient et récupèrent des entités en contexte. En allant au-delà de la vision positionnelle simpliste, les auteurs fournissent une explication mécaniste à la fois des forces (récupération robuste via des signaux multiples) et des fragilités (bruit dans les positions intermédiaires conduisant à des erreurs « perdu au milieu ») des LLM dans des contextes longs.

Le travail suggère que le phénomène « perdu au milieu » n'est pas simplement un échec de l'attention, mais une dégradation spécifique du signal positionnel, qui n'est que partiellement compensée par les mécanismes lexical et réflexif. Les auteurs publient leur code et leurs données pour faciliter de nouvelles recherches sur les mécanismes causaux du raisonnement en contexte.