Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis

1. Problématique et Contexte

L'apprentissage en contexte (In-Context Learning ou ICL) permet aux grands modèles de langage (LLM) de résoudre de nouvelles tâches à partir de démonstrations fournies dans l'entrée, sans mise à jour des poids du modèle. Bien que ce phénomène soit bien documenté, son mécanisme sous-jacent reste partiellement incompris.

La recherche actuelle est dominée par deux paradigmes qui, jusqu'à présent, n'ont pas été unifiés :

1. Le paradigme introspectif (mécaniste) : Il identifie des composants spécifiques (comme les "Induction Heads" ou IH) responsables de l'ICL via des analyses de circuits et des ablations. Cependant, il ne explique pas comment ces composants réalisent les différentes fonctions de l'ICL.
2. Le paradigme holistique (fonctionnel) : Il décompose l'ICL en deux composantes fonctionnelles distinctes : la Reconnaissance de Tâche (Task Recognition - TR), qui identifie l'espace des étiquettes possibles, et l'Apprentissage de Tâche (Task Learning - TL), qui apprend la mapping spécifique entre les textes et les étiquettes. Ce paradigme ne parvient pas cependant à localiser ces fonctions sur des composants concrets du modèle.

Le problème central est donc de réconcilier ces deux perspectives : identifier précisément quelles têtes d'attention (attention heads) sont responsables de la TR et de la TL, et comprendre comment elles interagissent géométriquement pour produire l'ICL.

2. Méthodologie : TSLA et Analyse Géométrique

Les auteurs proposent un cadre unifié basé sur l'analyse des têtes d'attention, introduisant une nouvelle méthode appelée Task Subspace Logit Attribution (TSLA).

A. Identification des Têtes TR et TL via TSLA

Contrairement aux méthodes précédentes (comme l'attribution directe des logits) qui peuvent être biaisées par la surface des tokens ou ignorer la compétition entre étiquettes, TSLA utilise la géométrie de l'espace des sous-espaces (subspace) :

• Pour la Reconnaissance de Tâche (TR) : Le score TR mesure la norme de la projection de la sortie d'une tête d'attention sur le sous-espace engendré par les vecteurs de déembedding (unembeddings) des étiquettes de démonstration. Une tête est considérée comme une "tête TR" si sa sortie est fortement concentrée dans ce sous-espace, indiquant qu'elle aide le modèle à reconnaître l'espace sémantique des étiquettes.
• Pour l'Apprentissage de Tâche (TL) : Le score TL mesure la composante de la sortie de la tête (après projection sur le sous-espace de la tâche) qui s'aligne avec la direction moyenne de différence entre l'étiquette correcte et les étiquettes incorrectes. Cela isole la capacité de la tête à discriminer la bonne étiquette au sein de l'espace reconnu.

B. Validation Expérimentale

Les auteurs valident leur approche à travers plusieurs axes :

• Corrélations et Chevauchements : Analyse de la similarité entre les têtes TR/TL et les Induction Heads (IH) connues.
• Ablations Fines : Suppression sélective des têtes TR et TL pour observer l'impact sur le ratio de reconnaissance de tâche (TR ratio) et la précision globale.
• Perturbations d'Entrée : Modification des démonstrations (brouillage du texte vs changement des étiquettes) pour tester l'indépendance causale des deux mécanismes.
• Steering (Pilotage) et Analyse Géométrique : Injection des sorties des têtes dans le flux résiduel des entrées "zero-shot" pour observer les changements géométriques dans les états cachés (alignement sur le sous-espace vs rotation vers l'étiquette correcte).

3. Contributions Clés

1. Définition et Identification via TSLA : Développement d'une méthode théorique pour localiser les têtes d'attention spécialisées dans la TR et la TL, dépassant les limites des méthodes d'attribution directe (DLA).
2. Séparabilité Mécaniste : Démonstration par ablation que la TR et la TL sont contrôlées par des ensembles de têtes distincts et séparables.
3. Interprétation Géométrique Unifiée :
   • Les têtes TR alignent les états cachés avec le sous-espace de la tâche (reconnaissance de l'espace des étiquettes).
   • Les têtes TL effectuent une rotation à l'intérieur de ce sous-espace pour orienter l'état caché vers le vecteur de déembedding de l'étiquette correcte.
4. Réconciliation des Concepts Existants :
   • Les Induction Heads (IH) sont révélées comme un sous-ensemble des têtes TR, expliquant leur rôle crucial dans la reconnaissance de l'espace des étiquettes plutôt que dans la sélection spécifique de l'étiquette.
   • Les Task Vectors sont expliqués comme des agrégats des sorties de ces têtes, dont l'efficacité dépend de la tâche (classification vs génération).

4. Résultats Principaux

• Spécialisation des Têtes : Les têtes TR et TL sont hautement spécialisées. L'ablation des têtes TR fait chuter le "TR ratio" (la proportion de prédictions dans l'ensemble des étiquettes valides) de ~100% à ~20%, effondrant la précision. L'ablation des têtes TL réduit la précision mais laisse le TR ratio largement intact (le modèle devine au hasard parmi les étiquettes valides).
• Relation avec les Induction Heads (IH) : Il existe un fort chevauchement entre les têtes TR et les IH. Les IH agissent principalement en renforçant la reconnaissance de l'espace des étiquettes (TR).
• Distribution en Couches : Les têtes TR apparaissent significativement plus profondément dans le réseau que les têtes TL et les IH. Cela suggère un ordre de traitement où l'apprentissage des mappings (TL) précède la reconnaissance formelle de l'espace de la tâche (TR) dans le flux de calcul.
• Analyse Géométrique :
  • L'ajout des sorties des têtes TR aux états cachés augmente l'alignement avec le sous-espace de la tâche.
  • L'ajout des sorties des têtes TL augmente la différence de logits (logit difference) entre la bonne et les mauvaises étiquettes, effectuant une rotation discriminante.
• Dépendance à la Tâche (Steering) :
  • Pour les tâches de classification (espace d'étiquettes fixe), les vecteurs de tâche basés sur les têtes TR sont essentiels pour restaurer les performances zero-shot.
  • Pour les tâches de génération (espace d'étiquettes infini, ex: rédaction de critiques), les vecteurs basés sur les têtes TL sont plus efficaces, car ils capturent la mapping texte-étiquette complexe plutôt que la simple reconnaissance d'un ensemble fermé.

5. Signification et Impact

Ce travail fournit une explication unifiée et interprétable de l'ICL. Il démontre que l'ICL n'est pas un phénomène monolithique, mais le résultat de deux mécanismes neuronaux distincts et séquentiels :

1. La Reconnaissance de Tâche (via les têtes TR/IH) qui prépare le modèle en le plaçant dans le bon sous-espace sémantique.
2. L'Apprentissage de Tâche (via les têtes TL) qui affine la prédiction en naviguant à l'intérieur de ce sous-espace.

Cette découverte permet de mieux comprendre pourquoi les modèles échouent en zero-shot (manque de reconnaissance de tâche) et comment les interventions mécanistiques (comme l'ajout de vecteurs de tâche) peuvent restaurer ces capacités. Elle offre également un cadre pour analyser la robustesse et la généralisation des LLMs sur des tâches complexes et composées.