Résumé Technique : Pilotage de l'Attention Guidé par les Variétés (MAGS)

Énoncé du Problème

Les modèles de langage de grande taille (LLM) possèdent souvent les connaissances sous-jacentes nécessaires à des tâches de raisonnement complexes, mais échouent fréquemment à exprimer cette capacité de manière fiable lors de la génération. Bien que l'échantillonnage répété d'un même modèle et d'un même prompt puisse produire à la fois des solutions correctes et incorrectes, les annotations au niveau du processus indiquent que les erreurs surviennent généralement à des étapes intermédiaires du raisonnement plutôt que par un manque de connaissances terminal.

Les méthodes existantes de pilotage des activations au moment de l'inférence tentent de corriger ces comportements en appliquant des vecteurs de correction fixes et précalculés au flux résiduel ou aux sorties de l'attention. Cependant, ces approches souffrent d'une inadéquation structurelle critique avec les tâches de raisonnement :

1. Perturbation Indiscriminée : Elles appliquent des corrections de manière uniforme, indépendamment de l'état actuel du modèle. Cela risque de corrompre des étapes déjà correctes tout en échouant à intercepter des erreurs qui ne se sont pas encore manifestées.
2. Nature Statique : Elles manquent de mécanismes pour détecter quand une trajectoire de raisonnement dérive vers un état d'erreur, traitant toutes les étapes de génération comme ayant un besoin égal de correction.

L'hypothèse centrale de ce travail est que les erreurs de raisonnement se manifestent par une dérive structurée dans un sous-espace de faible dimension des espaces de sortie des têtes d'attention individuelles. Les trajectoires correctes et incorrectes occupent des régions géométriquement séparables, et la transition d'un comportement correct à un comportement incorrect suit une direction spécifique de faible rang.

Méthodologie : Pilotage de l'Attention Guidé par les Variétés (MAGS)

MAGS est une intervention adaptative et consciente de la trajectoire au moment de l'inférence qui pilote les sorties des têtes d'attention uniquement lorsqu'une déviation par rapport à une « variété de justesse » est détectée. La méthode opère en trois phases :

1. Construction de la Variété d'Erreur par Contraste

Pendant une phase d'entraînement, le système analyse des paires contrastives de traces de raisonnement (même problème, une solution correcte, une solution incorrecte) pour apprendre la structure géométrique des erreurs.

• Vecteurs de Différence par Tête : Pour chaque tête d'attention $(l, h)$, le système calcule l'activation moyenne pour les traces correctes ($\mu_c$) et les traces incorrectes ($\mu_e$). Le vecteur de différence $\delta = \mu_e - \mu_c$ isole le décalage directionnel attribuable aux erreurs, annulant les spécificités du problème.
• Apprentissage du Sous-espace : En empilant ces vecteurs de différence sur plusieurs problèmes et en effectuant une Analyse en Composantes Principales (ACP), MAGS extrait une base de faible dimension $B$ (les $k$ premiers vecteurs singuliers). Cette base engendre le « sous-espace d'erreur » — les directions dans l'espace de sortie de la tête où la déviation correcte-incorrecte présente la plus grande variance.
• Référence Globale : Un centroïde global $\mu_c$ est calculé à partir de toutes les traces correctes pour servir de point de référence pendant l'inférence.

2. Détection d'Erreur par Proximité

Pendant l'inférence, MAGS surveille la sortie des têtes d'attention sélectionnées en temps réel.

• Score de Proximité : Pour chaque étape de token $t$, le système calcule la norme carrée de la projection de la sortie centrée de la tête sur le sous-espace d'erreur appris :
  $$d_t = \| B(a_t - \mu_c) \|^2$$
  Un score élevé indique que l'activation de la tête a dérivé significativement vers la variété d'erreur.
• Déclenchement Dynamique : Un seuil par tête $\tau$ est calibré sur les traces correctes (par exemple, fixé à un percentile élevé des scores provenant de trajectoires correctes). Une correction est appliquée uniquement si $d_t > \tau$. Cela garantit que les étapes déjà sur la trajectoire correcte restent inchangées.

3. Correction Ciblée

Lorsqu'un déclencheur s'active, MAGS applique une correction in situ à la sortie de la tête d'attention avant qu'elle ne soit transmise à la projection de sortie de la couche :
$$\tilde{a}_t = a_t - \alpha B^\top B (a_t - \mu_c)$$
où $\alpha \in (0, 1]$ est un hyperparamètre de force de pilotage.

• Interprétation Géométrique : Cette opération projette la déviation $(a_t - \mu_c)$ sur le sous-espace d'erreur et la soustrait, éliminant efficacement la composante d'erreur tout en préservant toutes les informations dans le complément orthogonal (les $(d_h - k)$ dimensions non liées à la variété d'erreur).
• Préservation de l'Information : La méthode garantit que les informations non liées à la variété d'erreur sont préservées, évitant la « sur-correction » observée dans les méthodes statiques.

Extension Multi-Objectifs (MAGSu)

Pour gérer des objectifs conflictuels (par exemple, validité syntaxique vs affinité de liaison chimique), les auteurs proposent MAGSu, qui apprend des variétés indépendantes pour chaque objectif et applique des corrections via des ensembles disjoints de têtes, empêchant ainsi les interférences entre les objectifs de pilotage.

Contributions Clés

1. Validation de l'Hypothèse : L'article fournit des preuves empiriques que les erreurs de raisonnement se manifestent par une dérive structurée dans des sous-espaces de faible dimension de têtes d'attention spécifiques. Des expériences diagnostiques montrent que les trajectoires correctes et incorrectes sont hautement séparables par ces sous-espaces, les signaux d'erreur étant rares et concentrés dans des têtes spécifiques.
2. Mécanisme d'Intervention Adaptative : MAGS introduit un déclencheur dynamique basé sur la proximité qui applique des corrections uniquement lorsqu'une tête dérive dans le sous-espace d'erreur. Cela contraste avec les méthodes statiques qui appliquent des vecteurs fixes indépendamment de l'état d'activation actuel.
3. Supériorité Empirique : La méthode surpasse systématiquement les bases de référence non pilotées et les approches de pilotage statique existantes (ITI, Pilotage Angulaire) sur divers benchmarks.

Résultats Expérimentaux

Les auteurs ont évalué MAGS sur trois familles de modèles (Llama-3.1-8B-Instruct, Gemma-4-E4b-it et GPT-OSS-20B) dans trois domaines :

• Raisonnement Mathématique (MATH-500, GSM8K) : MAGS a obtenu les gains les plus importants sur MATH-500, améliorant la précision de +5,2 % par rapport à la base de référence Llama non pilotée et de +3,4 % par rapport à la base de référence Gemma non pilotée. L'amélioration était la plus prononcée dans les tâches comportant des dérivations multi-étapes où l'accumulation d'erreurs est probable.
• Génération de Code (HumanEval, MBPP) : MAGS a amélioré HumanEval pass@1 d'environ 4 points par rapport aux bases de référence non pilotées. Notamment, les méthodes statiques (ITI, Pilotage Angulaire) ont dégradé les performances sur la génération de code, probablement en raison de la perturbation de la cohérence syntaxique, tandis que MAGS a évité cela en supprimant les corrections sur les étapes correctes.
• Génération Moléculaire (SMILES) : Dans une tâche multi-objectifs nécessitant à la fois une validité syntaxique et une forte affinité de liaison, MAGSu a réussi à améliorer simultanément les deux métriques. En revanche, ITI a amélioré la validité mais a considérablement dégradé l'affinité de liaison, démontrant l'interférence causée par des objectifs entremêlés dans le pilotage statique.
• Fluidité et Surcharge : MAGS a maintenu des niveaux de perplexité comparables aux bases de référence non pilotées, indiquant qu'il ne déforme pas la distribution de génération. La surcharge computationnelle est négligeable, dominée par des opérations $O(K \cdot k \cdot d_h)$ par étape, où $K$ est le nombre de têtes surveillées.

Signification et Revendications

L'article revendique que MAGS reformule le pilotage des activations comme un problème de détection de dérive plutôt que comme une tâche de correction constante. En identifiant une base de faible rang qui prédit le début d'une erreur, MAGS fournit une structure géométrique interprétable que les méthodes basées sur l'amplitude scalaire ne possèdent pas.

Les auteurs affirment que :

1. Les Variétés de Justesse sont Générales : L'existence de sous-espaces de faible dimension séparant les comportements corrects et incorrects semble être une caractéristique générale de la géométrie de l'attention des LLM, applicable au raisonnement, à la programmation et à la génération scientifique.
2. La Conscience de la Trajectoire est Critique : La correction indiscriminée perturbe les étapes de raisonnement valides. La capacité d'intervenir uniquement lorsqu'une déviation est détectée est le principal moteur de la performance supérieure de MAGS et de sa robustesse face à la sensibilité aux hyperparamètres.
3. Les Contraintes Géométriques Préviennent la Sur-correction : En projetant les activations de retour sur le complément du sous-espace d'erreur, MAGS préserve les capacités intrinsèques du modèle tout en éliminant uniquement les composantes spécifiques à l'erreur.

L'étude conclut que bien que MAGS nécessite un ensemble de traces contrastives pour l'ajustement de la variété, il offre une solution robuste et à faible surcharge pour améliorer la cohérence du raisonnement dans les LLM, surpassant les interventions statiques qui ne tiennent pas compte de la nature dynamique de la trajectoire de génération.