AttnTrace: Contextual Attribution of Prompt Injection and Knowledge Corruption

Ce papier présente AttnTrace, une nouvelle méthode efficace et précise basée sur les poids d'attention pour retracer les sources de contexte dans les grands modèles de langage, permettant ainsi d'identifier rapidement les injections de prompts et les corruptions de connaissances avec un coût de calcul réduit par rapport aux solutions existantes.

L'essentiel

🕵️‍♂️ AttnTrace : Le détective qui trouve le coupable dans une foule de documents

Imaginez que vous avez un super-intelligent (une Intelligence Artificielle ou IA) capable de lire des milliers de pages en une seconde. C'est ce qu'on appelle un "modèle de langage à long contexte". Ce super-intelligent est utilisé pour répondre à des questions complexes en consultant une énorme bibliothèque de documents (des articles, des mémoires, des emails, etc.).

Mais voici le problème : un pirate informatique peut glisser un petit mot d'ordre secret (une "injection de prompt") au milieu de ces milliers de pages. Ce mot d'ordre dit à l'IA : "Oublie tout ce que tu as lu, réponds uniquement ceci !" (par exemple, donner une excellente note à un article scientifique nul).

L'IA obéit et produit une réponse fausse. La question est : Comment retrouver le petit mot d'ordre coupable au milieu de cette montagne de texte ?

C'est là qu'intervient AttnTrace.

🧠 Le problème : Trouver une aiguille dans une botte de foin (et c'est bruyant)

Jusqu'à présent, les méthodes pour trouver ce coupable étaient soit :

1. Trop lentes : Comme essayer de lire chaque page, une par une, pour voir si elle change la réponse. Cela prend des heures.
2. Trop imprécises : Comme essayer de deviner qui a parlé dans une salle de concert bruyante. Les méthodes existantes se perdaient souvent dans le bruit.

🔍 La solution : AttnTrace, le détective qui écoute les "regards"

Les IA modernes fonctionnent avec un mécanisme appelé "Attention". Imaginez que l'IA lit un texte et que, pour chaque mot qu'elle écrit, elle regarde certains mots du texte original avec plus ou moins d'intensité. C'est comme si elle pointait un doigt vers les mots importants.

AttnTrace utilise ce "doigt pointé" (les poids d'attention) pour trouver le coupable. Mais il y a deux pièges que les anciens détectives ne pouvaient pas contourner :

1. Le bruit de fond : Parfois, l'IA regarde des mots inutiles (comme des points ou des virgules) juste par habitude. Cela brouille les pistes.
2. La dispersion : Si le pirate utilise plusieurs petits mots d'ordre différents pour tromper l'IA, l'IA se divise en plusieurs directions. Au lieu de pointer fort vers un seul coupable, elle regarde faiblement partout. C'est comme si le détective regardait dix suspects différents avec un regard vague au lieu de se concentrer sur le vrai coupable.

🛠️ Les deux super-pouvoirs d'AttnTrace

Pour résoudre ces problèmes, les chercheurs ont créé deux astuces magiques :

1. Le filtre "Top-K" (Ne regarder que les meilleurs regards)

Au lieu de compter tous les regards de l'IA (y compris ceux sur les virgules inutiles), AttnTrace ne garde que les regards les plus intenses (les "Top-K").

• L'analogie : Imaginez que vous essayez de savoir qui a crié dans une foule. Au lieu d'écouter tout le monde, vous ne vous intéressez qu'aux 5 personnes qui ont crié le plus fort. Cela élimine le bruit et vous donne une piste claire.

2. L'échantillonnage contextuel (La méthode du "Jeux de l'ombre")

C'est l'astuce la plus brillante pour contrer la dispersion.

• L'analogie : Imaginez que vous cherchez à savoir quel suspect a volé un gâteau dans une pièce remplie de 10 suspects. Si vous les regardez tous ensemble, l'IA (le voleur) peut se cacher en regardant tout le monde.
• La technique d'AttnTrace : Au lieu de regarder les 10 suspects ensemble, le détective crée 30 petites versions de la pièce. Dans chaque version, il ne met que 4 suspects au hasard.
  • Dans la version A, le vrai coupable est seul avec 3 innocents. L'IA est obligée de le regarder très fort !
  • Dans la version B, c'est un autre innocent qui est là. L'IA le regarde.
  • En répétant ce jeu 30 fois et en faisant la moyenne, le vrai coupable ressort toujours comme celui qui a reçu le plus de regards intenses, même s'il était caché dans la grande foule au début.

🚀 Les résultats : Plus rapide et plus précis

Les tests montrent qu'AttnTrace est un champion :

• Précision : Il trouve le coupable dans 95% des cas, là où les anciennes méthodes échouaient souvent.
• Vitesse : Il est environ 10 à 20 fois plus rapide que les meilleurs détectives précédents. Là où un ancien détective prenait 100 secondes pour analyser un texte, AttnTrace le fait en 10 secondes.
• Robustesse : Même si le pirate essaie de se cacher en changeant son style ou en divisant son message en plusieurs morceaux, AttnTrace le repère toujours.

🌍 Pourquoi c'est important dans la vraie vie ?

L'article donne un exemple concret : des chercheurs ont découvert que des gens inséraient des instructions cachées dans des articles scientifiques pour manipuler les IA qui écrivent les critiques de ces articles.
Grâce à AttnTrace, on peut maintenant remonter le fil : on prend la critique fausse générée par l'IA, et on peut dire : "Attendez, c'est cette phrase précise, cachée au milieu du document, qui a forcé l'IA à mentir."

En résumé

AttnTrace, c'est comme donner à un détective une paire de lunettes spéciales qui lui permettent de voir exactement qui l'IA a écouté pour prendre sa décision. Grâce à deux astuces (ignorer le bruit et tester des petits groupes de suspects), ce détective trouve le coupable beaucoup plus vite et plus sûrement que jamais auparavant, protégeant ainsi nos systèmes d'IA contre la manipulation.

Résumé technique

1. Problématique et Contexte

Les modèles de langage à grand contexte (Long-context LLMs), tels que GPT-5, Gemini-2.5-Pro ou Claude-Sonnet-4, sont de plus en plus utilisés dans des systèmes avancés comme la génération augmentée par récupération (RAG) et les agents autonomes. Ces systèmes permettent au modèle de traiter des instructions accompagnées de contextes vastes (extraits de bases de connaissances, mémoires ou Internet) pour générer des réponses pertinentes.

Cependant, ces systèmes sont vulnérables à deux types d'attaques majeures :

• L'injection de prompt (Prompt Injection) : Un attaquant insère des instructions malveillantes dans le contexte pour détourner le modèle de ses consignes initiales.
• La corruption de connaissances (Knowledge Corruption) : Un attaquant empoisonne les données de la base de connaissances pour que le modèle génère des réponses incorrectes ou spécifiques.

Le défi principal : Une fois qu'un modèle génère une réponse malveillante, il est difficile d'identifier exactement quel fragment de texte dans le contexte long a provoqué cette sortie. Les méthodes actuelles de "retraçage de contexte" (context traceback) souffrent soit d'une performance sous-optimale (précision/rappel faibles), soit d'un coût computationnel prohibitif (plusieurs centaines de secondes par échantillon).

2. Méthodologie : AttnTrace

Les auteurs proposent AttnTrace, une nouvelle méthode de traçage de contexte basée sur les poids d'attention générés par le modèle de langage lui-même. L'idée centrale est que les mécanismes d'attention des architectures Transformer capturent naturellement l'influence des tokens d'entrée sur la génération de la réponse.

Limites des approches baselines

Une approche naïve consisterait à moyenner les poids d'attention entre tous les tokens d'un texte et les tokens de sortie. Les auteurs identifient deux problèmes majeurs avec cette méthode :

1. Bruit des poids d'attention : Les poids d'attention sont souvent concentrés sur un petit nombre de tokens (phénomène de "attention sink", comme les points ou les délimiteurs), tandis que le reste du texte contient du bruit. La moyenne simple dilue l'importance des tokens critiques.
2. Dispersion des poids d'attention : Lorsque plusieurs textes malveillants distincts peuvent induire la même sortie, le modèle répartit son attention entre toutes ces sources. Cela affaiblit le signal d'attention pour chaque texte individuel, rendant leur identification difficile.

Les deux techniques clés d'AttnTrace

Pour surmonter ces limites, AttnTrace intègre deux techniques innovantes :

1. Moyenne des Top-K tokens (Top-K Averaging) :
   Au lieu de moyenner les poids d'attention sur tous les tokens d'un texte, AttnTrace sélectionne uniquement les K tokens ayant les poids d'attention les plus élevés au sein de ce texte.

   • Insight : Les textes véritablement influents contiennent quelques tokens critiques (souvent des instructions clés) qui portent les poids d'attention les plus significatifs. Ignorer les tokens "bruyants" améliore la précision du score d'attribution.

2. Échantillonnage de contexte (Context Subsampling) :
   Pour résoudre le problème de la dispersion, la méthode procède par échantillonnage aléatoire du contexte.

   • Processus : On sélectionne aléatoirement un sous-ensemble de textes du contexte (par exemple 40 %). On calcule le score d'attribution pour ce sous-ensemble. On répète ce processus plusieurs fois (B itérations) avec différents sous-ensembles.
   • Agrégation : Le score final pour chaque texte est la moyenne de ses scores à travers toutes les itérations.
   • Insight : En réduisant la présence simultanée de multiples textes malveillants dans un échantillon, on force le modèle à concentrer son attention sur les textes présents, évitant ainsi la dilution du signal. Les résultats sont ensuite agrégés pour obtenir une estimation robuste.

Analyse Théorique

Les auteurs fournissent une analyse théorique démontrant que lorsque le nombre de tokens ayant des états cachés similaires (issus de textes malveillants visant le même but) augmente, la borne supérieure du poids d'attention maximal diminue. Cela confirme mathématiquement pourquoi la dispersion se produit et pourquoi l'échantillonnage (qui réduit le nombre de textes concurrents) permet de restaurer la capacité de détection.

3. Contributions Principales

1. Proposition d'AttnTrace : Une nouvelle méthode de traçage de contexte exploitant les poids d'attention natifs des LLM, sans nécessiter de réentraînement ou d'accès aux gradients.
2. Innovations Techniques : Développement et validation théorique/empirique des techniques de "Top-K averaging" et de "Context Subsampling" pour filtrer le bruit et contrer la dispersion de l'attention.
3. Évaluation Systématique : Une comparaison exhaustive contre l'état de l'art (TracLLM, Shapley, LIME, etc.) sur divers jeux de données (HotpotQA, MuSiQue, NQ) et types d'attaques (injection de prompt, corruption de connaissances).
4. Amélioration de la Détection : Démonstration que AttnTrace peut servir de pré-filtre ("attribution-before-detection") pour améliorer les méthodes de détection d'injection de prompt existantes dans des contextes longs.
5. Étude de Cas Réelle : Application réussie pour identifier des instructions injectées dans des articles scientifiques destinés à manipuler les revues générées par IA.

4. Résultats Expérimentaux

Les résultats montrent que AttnTrace surpasse nettement les méthodes existantes :

• Précision et Rappel : Sur le jeu de données HotpotQA (attaques de corruption de connaissances), AttnTrace atteint un Précision/Rappel de 0,95/0,95, contre 0,80/0,80 pour la meilleure méthode de référence (TracLLM). Sur MuSiQue (injection de prompt), il atteint 0,99/0,81 contre 0,94/0,77 pour TracLLM.
• Efficacité Computationnelle : AttnTrace est considérablement plus rapide. Il traite un échantillon en ~10 à 20 secondes, tandis que TracLLM et les méthodes basées sur Shapley prennent plus de 100 à 1000 secondes (facteur d'accélération de 15x à 20x).
• Robustesse : AttnTrace reste efficace même face à des attaques adaptatives fortes (attaques par optimisation visant à minimiser les poids d'attention) et des attaques par fractionnement de charge (payload splitting).
• Amélioration de la Détection : En utilisant AttnTrace pour isoler les textes suspects avant de les soumettre à un détecteur (comme DataSentinel), le taux de faux positifs (FPR) est considérablement réduit pour les contextes longs.

5. Signification et Impact

• Outil de Forensique Post-Attaque : AttnTrace fournit un moyen fiable d'identifier la source racine d'une attaque dans des systèmes LLM complexes, ce qui est crucial pour l'audit de sécurité et la compréhension des défaillances.
• Efficacité et Évolutivité : Contrairement aux méthodes basées sur la perturbation (qui nécessitent des milliers de passes forward), AttnTrace est léger et scalable, le rendant applicable à des systèmes réels.
• Généralité : La méthode fonctionne sur une large gamme de modèles (Open-source comme Llama, Qwen, et fermés comme GPT-5, Gemini) et s'adapte à différents types d'attaques.
• Applications Pratiques : L'étude de cas sur la manipulation de revues académiques par IA démontre le potentiel de l'outil pour protéger l'intégrité des processus d'évaluation automatisés.

En conclusion, AttnTrace comble un vide important dans la sécurité des LLM à grand contexte en offrant une solution rapide, précise et robuste pour retracer l'origine des manipulations malveillantes, surpassant les approches actuelles à la fois en performance et en coût computationnel.