Accueil › Glossaire › Retrieval Head

Retrieval Head

Une retrieval head est une tête d’attention spécialisée au sein d’un Transformer, responsable de la récupération d’informations pertinentes depuis le contexte d’entrée vers la sortie. Moins de 5 % des têtes d’attention d’un LLM sont des retrieval heads, mais elles sont indispensables à la factualité en long contexte.

Le problème qui a mené à la découverte

Les LLM modernes acceptent des fenêtres de contexte de 128K, 1M, voire plusieurs millions de tokens. Mais une question fondamentale restait ouverte : comment un modèle à base de Transformer parvient-il à récupérer une information précise enfouie à un endroit arbitraire dans un contexte aussi volumineux ? Le test Needle in a Haystack démontrait que certains modèles réussissaient cette tâche, mais le mécanisme interne restait opaque.

Wu et al. ont abordé cette question sous l’angle de l’interprétabilité mécaniste. Leur hypothèse, inspirée des travaux sur les CopyNet (Gu et al., 2016) et les Induction Heads (Olsson et al., 2022), était qu’il existe des têtes d’attention spécialisées dans la récupération d’information, de la même façon que les induction heads sont responsables de l’apprentissage en contexte.

Comment fonctionne une retrieval head

L’algorithme de copier-coller conditionnel

Une retrieval head implémente ce que les auteurs appellent un « algorithme de copier-coller conditionnel ». Concrètement, lorsque le modèle génère une réponse, la retrieval head identifie dans le contexte d’entrée les tokens pertinents par rapport à la question posée, puis redirige ces tokens vers la sortie. C’est un mécanisme de recopie sélective, conditionné par la requête.

Le « retrieval score » d’une tête est défini comme la fréquence avec laquelle cette tête effectue une opération de copier-coller lors de la réponse à des questions factuelles. Si le token auquel la tête porte le plus d’attention est identique au token en cours de génération, c’est un comportement de type retrieval. Les têtes qui présentent ce comportement de manière systématique sont classifiées comme retrieval heads.

Illustration concrète

Prenons un exemple simple. On donne au modèle un long texte contenant la phrase « The best fruit is orange » quelque part au milieu, puis on pose la question « What is the best fruit? ». Lors de la génération du token « orange » dans la réponse, les retrieval heads vont diriger leur attention vers le token « orange » dans le contexte d’entrée, effectuant ainsi le transfert d’information. Les streaming heads, elles, se focalisent sur les tokens récents et les attention sinks, sans participer activement à la récupération.

Cinq propriétés fondamentales des retrieval heads

L’étude systématique de Wu et al. sur 4 familles de modèles, 6 échelles de paramètres et 3 types de fine-tuning a révélé cinq propriétés remarquables :

Comment détecter les retrieval heads

Méthode originale (Wu et al.)

La détection repose sur un algorithme statistique qui calcule le retrieval score de chaque tête d’attention. Le processus se déroule comme suit :

Premièrement, on prépare un ensemble d’exemples de type Needle in a Haystack : un fait est inséré dans un long contexte, et on pose une question dont la réponse est ce fait. Deuxièmement, pour chaque tête d’attention, on examine si, lors de la génération du token de réponse, la tête « regarde » principalement le token correspondant dans le contexte d’entrée. Troisièmement, on agrège ce score sur l’ensemble des exemples pour obtenir un retrieval score moyen par tête. Les têtes avec un score élevé sont classifiées comme retrieval heads.

L’implémentation utilise FlashAttention pour le caching et l’attention standard pour le décodage (FlashAttention ne retourne pas la matrice d’attention, nécessaire pour l’analyse). Un seul GPU A100 80 Go suffit pour détecter les retrieval heads sur des contextes allant jusqu’à 50K tokens.

import json
import numpy as np

# Charger les scores par tête (exemple : LLaMA-2 7B 80K)
with open('./head_score/llama-2-7b-80k.json') as file:
    head_list = json.loads(file.readline())

# Calculer le score moyen et classer
head_score_list = [
    ([int(x) for x in l[0].split("-")], np.mean(l[1]))
    for l in head_list.items()
]
head_score_list = sorted(head_score_list, key=lambda x: x[1], reverse=True)

# Top 5 retrieval heads
for head, score in head_score_list[:5]:
    print(f"Couche {head[0]}, Tête {head[1]} → Score: {score:.2f}")

Sur LLaMA-2 7B 80K, les retrieval heads les plus puissantes sont situées aux couches 16 (tête 19, score 0.94) et 11 (tête 15, score 0.92). Quelques échantillons suffisent à identifier de manière stable les retrieval heads les plus fortes.

Méthode DuoAttention (Xiao et al., MIT)

DuoAttention (ICLR 2025) propose une approche alternative, basée sur l’optimisation plutôt que la statistique. Au lieu de mesurer le comportement de copier-coller, DuoAttention évalue directement l’impact de la compression du KV cache de chaque tête sur la qualité de la sortie. Les têtes dont la compression dégrade significativement la sortie sont identifiées comme retrieval heads. Cette méthode est plus robuste car elle mesure la déviation de sortie plutôt qu’un motif attentionnel spécifique.

Retrieval Heads vs Streaming Heads

La découverte des retrieval heads a naturellement conduit à caractériser les autres têtes d’attention. DuoAttention formalise cette dichotomie :

La proportion de retrieval heads varie selon l’architecture. Les modèles à Multi-Head Attention (MHA) comme LLaMA-2 ont environ 25 % de retrieval heads (par le critère DuoAttention). Les modèles à Grouped-Query Attention (GQA) comme LLaMA-3 en ont environ 50 %, car le partage de clés/valeurs entre têtes rend chaque « groupe » plus critique.

Applications pratiques

Optimisation du KV cache

L’application la plus immédiate des retrieval heads est l’optimisation du KV cache, le principal goulot d’étranglement mémoire pour l’inférence en long contexte. Pour LLaMA-2 7B, le KV cache de 100K tokens occupe plus de 50 Go de mémoire GPU. En compressant le KV cache des streaming heads (qui n’ont besoin que des tokens récents) tout en conservant le cache complet des retrieval heads, DuoAttention réduit la mémoire jusqu’à 2,55× pour les modèles MHA et 1,67× pour les modèles GQA, avec une perte de précision minimale sur le test NIAH et sur LongBench.

Résultat spectaculaire : combiné avec la quantification (poids en 4 bits, KV cache en 4 bits), DuoAttention permet à LLaMA-3 8B de décoder un contexte de 3,3 millions de tokens sur un seul GPU A100 80 Go.

Réduction des hallucinations

L’étude montre que l’activation des retrieval heads est corrélée avec la factualité de la sortie. Quand les retrieval heads s’activent correctement (elles « trouvent » le bon passage), la réponse est factuelle. Quand elles échouent, le modèle hallucine. Cette observation ouvre la voie à des systèmes de détection d’hallucinations basés sur le monitoring de l’activité des retrieval heads pendant la génération.

Impact sur le raisonnement Chain-of-Thought

Les retrieval heads influencent fortement le raisonnement Chain-of-Thought (CoT). En CoT, le modèle doit fréquemment se référer à la question initiale et au contexte précédemment généré pour construire son raisonnement étape par étape. Ce va-et-vient constant repose sur les retrieval heads. Les tâches où le modèle génère directement une réponse à partir de ses connaissances entraînées (sans besoin de se référer au contexte) sont beaucoup moins affectées par la désactivation des retrieval heads.

Extension : QRHead (Query-Focused Retrieval Head)

Zhang et al. (2025) ont introduit QRHead, une amélioration du concept de retrieval head. Au lieu de mesurer le comportement de copier-coller générique, QRHead agrège les scores d’attention par rapport à la requête spécifique de l’utilisateur, en utilisant quelques exemples de tâches réelles (QA long contexte). Le retriever QRRetriever qui en découle sélectionne les passages les plus pertinents en fonction de la masse attentionnelle des QRHeads. Sur les tâches de raisonnement multi-hop (LongMemEval, CLIPPER), cette approche améliore les performances de plus de 10 % par rapport au contexte complet.

Place dans l’écosystème

Limites de la recherche actuelle

La découverte des retrieval heads est une avancée significative, mais plusieurs limites subsistent. Les études ont porté principalement sur des modèles open-weight (LLaMA, Mistral, Qwen, Yi). Les modèles propriétaires (GPT-5.4, Claude Opus 4.6, Gemini 3.1 Pro) ne sont pas analysables avec cette méthode car ils ne donnent pas accès aux matrices d’attention internes. La relation exacte entre retrieval heads et le phénomène Lost in the Middle n’est pas entièrement clarifiée : les retrieval heads expliquent le « quoi » (quelles têtes récupèrent l’info) mais pas complètement le « pourquoi » du biais positionnel en U. Enfin, les retrieval heads expliquent la récupération factuelle, mais pas les capacités de raisonnement, de synthèse ou de compréhension globale d’un long contexte.

Verdict

La découverte des retrieval heads est l’une des contributions les plus importantes en interprétabilité mécaniste des LLM de ces dernières années. Pour la première fois, on identifie un sous-réseau concret et spécialisé qui implémente un algorithme fonctionnel (copier-coller conditionnel) au sein d’un Transformer. Les implications sont directes : compression intelligente du KV cache (DuoAttention réduit la mémoire de 2,5×), compréhension de la factualité, et détection potentielle d’hallucinations. Si vous travaillez sur l’optimisation d’inférence en long contexte ou sur la fiabilité des LLM, c’est un article incontournable.