🤖 AI

TRACE: Distilling Where It Matters via Token-Routed Self On-Policy Alignment

L'article propose TRACE, une méthode d'auto-distillation à routage de jetons qui applique sélectivement la divergence KL à des segments critiques marqués par des annotateurs durant l'apprentissage par renforcement avec récompenses vérifiables, atténuant ainsi le gaspillage de gradient et la fuite d'informations privilégiées pour améliorer significativement le raisonnement mathématique à long horizon et la généralisation hors distribution par rapport aux bases de référence GRPO standard et d'auto-distillation sur réponse complète.

Auteurs originaux : Jiaxuan Wang, Xuan Ouyang, Zhiyu Chen, Yulan Hu, Zheng Pan, Xin Li, Lan-Zhe Guo

Publié 2026-05-12

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Jiaxuan Wang, Xuan Ouyang, Zhiyu Chen, Yulan Hu, Zheng Pan, Xin Li, Lan-Zhe Guo

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous enseigniez à un étudiant à résoudre des problèmes mathématiques complexes. Vous avez un « Enseignant » (une IA intelligente) et un « Étudiant » (l'IA que vous essayez d'entraîner).

Par le passé, les chercheurs ont essayé une méthode appelée Distillation Autonome Sur-Politique. L'idée était simple : laisser l'Étudiant générer une réponse, puis demander à l'Enseignant d'examiner la réponse entière et de dire : « Tu as bien fait ceci, et ceci, et ceci... » jusqu'au bout. L'Étudiant tenterait alors de copier chaque mouvement de l'Enseignant.

Le Problème : L'Effet de « Sur-Enseignement »
L'article soutient que cette approche de « tout copier » est en réalité néfaste, en particulier pour les tâches de raisonnement longues et complexes.

L'Analogie : Imaginez un étudiant passant un examen de 10 pages. L'Enseignant parcourt le texte et surligne chaque mot individuel en rouge, en disant : « Fais-le exactement comme ça. »
Le Résultat : L'étudiant est submergé. Il arrête de penser par lui-même. Ses réponses deviennent plus courtes (parce qu'il a peur d'écrire trop), il commence à deviner au hasard (son « entropie » augmente), et il finit par échouer à l'examen. L'article qualifie cela de « inadéquation de granularité ». L'étudiant n'a pas besoin d'être corrigé sur chaque mot ; il a seulement besoin d'aide sur les quelques étapes critiques où il pourrait rester bloqué ou faire une erreur.

La Solution : TRACE (Alignement Routé par Jeton pour le Raisonnement Critique)
Les auteurs proposent une nouvelle méthode appelée TRACE. Au lieu que l'Enseignant corrige tout le devoir, TRACE agit comme un tuteur intelligent qui ne pointe que les moments spécifiques et cruciaux.

Voici comment TRACE fonctionne, en utilisant une analogie simple :

1. Le « Projecteur » (Routé par Jeton)

Au lieu d'éclairer toute la page, TRACE utilise un projecteur.

L'Annotateur : Un assistant (qui peut même être l'étudiant lui-même !) examine la réponse de l'Étudiant et ne marque que les « Segments Critiques ». Ce sont les quelques phrases où la logique est soit brillante (Segments Clés), soit dangereusement erronée (Segments d'Erreur).
La Règle : Le projecteur ne couvre qu'environ 25 % de la réponse. Le reste de la réponse est laissé tel quel.

2. L'Approche « Deux-Outils » (FKL vs RKL)

Une fois que le projecteur est focalisé sur une partie spécifique, TRACE utilise deux outils différents selon ce qu'il observe :

Si l'Étudiant est en difficulté (Sous-allocation) : L'Enseignant dit : « Hé, tu ignores cette étape importante ! Regarde-moi et fais-le de cette façon. » Cela s'appelle la Divergence de Kullback-Leibler Forward (FKL). Elle pousse l'étudiant à prêter attention au bon chemin.
Si l'Étudiant est sûrement dans l'erreur (Sur-confiance) : L'Étudiant dit : « Je suis sûr que c'est juste ! » mais c'est en fait faux. L'Enseignant dit : « Non, stop ! Tu es trop sûr de toi. Recule. » Cela s'appelle la Divergence de Kullback-Leibler Reverse (RKL). Elle atténue la confiance de l'étudiant dans la mauvaise direction.

3. L'« Indice qui s'estompe » (Décroissance)

La partie la plus importante de TRACE est que l'aide de l'Enseignant est temporaire.

L'Analogie : Imaginez que l'Enseignant donne un indice à l'étudiant pendant un test d'entraînement. Au début, l'indice est fort et clair. Mais à mesure que l'étudiant s'améliore, l'indice devient de plus en plus faible jusqu'à disparaître complètement.
Le Résultat : Après une courte période de « mise en route » (environ 40 étapes pendant l'entraînement), l'Enseignant arrête de parler entièrement. L'Étudiant est laissé à lui-même pour résoudre les problèmes en utilisant sa propre logique (apprentissage par renforcement), mais il a déjà appris les habitudes critiques grâce à la courte période de guidance. Cela empêche l'étudiant de devenir dépendant de l'Enseignant.

Pourquoi Cela Fonctionne (Les Résultats)

L'article a testé cela sur des problèmes mathématiques (comme la résolution d'équations) et des questions de connaissances générales (GPQA).

Les méthodes « Tous-Jetons » ont échoué : Lorsque d'autres méthodes tentaient de corriger toute la réponse, les performances de l'IA s'effondraient, et elle commençait à donner des réponses plus courtes et moins créatives.
TRACE a réussi : En ne corrigeant que les parties critiques puis en s'estompant, TRACE a amélioré les scores de mathématiques de l'IA de manière significative (environ 2,76 % en moyenne) par rapport à la méthode standard.
Le Bonus « Auto-Apprentissage » : Même lorsque l'IA agit comme son propre tuteur (sans l'aide d'un humain ou d'une IA plus puissante), TRACE fonctionne toujours bien. Cela prouve que la méthode ne se contente pas de « tricher » en important les connaissances d'un enseignant ultra-intelligent ; elle apprend réellement à l'IA comment mieux apprendre.

En Résumé :
TRACE est comme un entraîneur qui ne crie pas sur le joueur pour chaque erreur commise pendant le match. Au lieu de cela, l'entraîneur attend les moments critiques, donne une correction précise et spécifique, puis se retire pour laisser le joueur jouer le match par lui-même. Cela maintient le joueur concentré, confiant et capable de gérer des défis longs et difficiles sans s'épuiser.

Résumé Technique : TRACE - Distillation Ciblée via l'Alignement Auto-Politique Routé par Jetons

1. Énoncé du Problème

L'article aborde les limites de la Distillation Auto-Politique (Self-OPD) lorsqu'elle est appliquée à l'Apprentissage par Renforcement à Horizon Long avec Récompenses Vérifiables (RLVR), en particulier dans les tâches de raisonnement mathématique. Bien que les méthodes Self-OPD (par exemple, SDPO, SRPO) densifient le signal de récompense épars en utilisant un instantané de la politique comme enseignant pour guider l'élève, les auteurs identifient un mode de défaillance critique lors de l'entraînement à horizon long :

Inadéquation de Granularité : Les méthodes existantes appliquent une perte de Distillation de Connaissance (KL) sur tous les jetons d'une réponse. Cependant, l'analyse montre que ~97 % de la masse de probabilité dans les trajectoires correctes est déjà partagée entre l'élève et l'enseignant, ce qui signifie que la plupart des gradients KL sont redondants.
La « Taxe de Distillation Uniforme » : L'application de KL sur tous les jetons force le modèle à dépenser des gradients sur des positions qui n'ont pas besoin de correction. Cela conduit à :
- Fuite d'Informations Privilegiées : L'élève absorbe des corrélations présentes dans le contexte privilégié (par exemple, les traces vérifiées) qui sont absentes lors de l'inférence, provoquant une dégradation Hors Distribution (OOD).
- Hausse de l'Entropie et Effondrement du Raisonnement : Les auteurs observent un « effondrement à trois symptômes » lors de l'entraînement à horizon long : la longueur des réponses diminue drastiquement, l'entropie de l'acteur par jeton augmente (souvent >4 fois la ligne de base), et la précision de validation chute significativement.
La Question Centrale : La bonne unité de distillation n'est pas la réponse entière, mais le petit sous-ensemble de jetons de raisonnement décisifs où la distribution de l'élève a réellement besoin de correction.

2. Méthodologie : TRACE

Les auteurs proposent TRACE (Token-Routed Alignment for Critical rEasoning), une méthode qui achemine les signaux de distillation uniquement vers des plages de jetons spécifiques plutôt que de les appliquer uniformément.

Composants Principaux

Self-OPD Routé par Jetons : Au lieu d'une perte KL globale, TRACE utilise un annotateur privilégié ( $\pi_A$ ) pour identifier des plages spécifiques dans une trajectoire :
- Plages Clés ( $K_y$ ) : Étapes de raisonnement décisives dans les trajectoires correctes (Récompense=1).
- Plages d'Erreur ( $E_y$ ) : Erreurs de raisonnement à l'origine des trajectoires incorrectes (Récompense=0).
- Non-Plages ( $N_y$ ) : Tous les jetons restants.
- Contrainte : La couverture totale des plages est limitée à $\alpha = 0,25$ (25 % des jetons) pour limiter l'exposition aux informations privilégiées.
Routage Diagnostique Grossier :
- L'annotateur émet une étiquette de type grossière (par exemple, case_split, missed_case) mais pas le texte de la plage ni le contenu complet de la trajectoire.
- L'Enseignant ( $\pi_T$ ) reçoit le problème et cette étiquette grossière comme préfixe diagnostique privé. Il calcule les logits de manière causale sur le préfixe généré par l'élève, mais ne voit pas les emplacements ou le contenu spécifiques des plages.
- L'Élève ( $\pi_S$ ) est entraîné en utilisant une perte routée où le terme KL est appliqué uniquement aux plages masquées.
Routage Asymétrique de Divergence :
- KL Direct (FKL) sur les Plages Clés : Appliqué aux trajectoires correctes. Le FKL ( $KL(\pi_T \parallel \pi_S)$ ) fournit une « poussée » aux jetons que l'élève sous-alloque, indépendamment de la masse de probabilité actuelle de l'élève.
- KL Inverse (RKL) sur les Plages d'Erreur (Optionnel) : Appliqué aux trajectoires incorrectes. Le RKL ( $KL(\pi_S \parallel \pi_T)$ ) supprime les jetons auxquels l'élève est trop confiant mais que l'enseignant désapprouve.
- Aucun KL sur les Non-Plages : Ces jetons reposent uniquement sur la perte GRPO standard (Optimisation de Politique Relative par Groupe).
Calendrier d'Exposition Finie :
- Le poids KL ( $\lambda_k$ ) est lissé vers zéro après un court échauffement (par exemple, 40 étapes).
- Cela garantit que l'exposition cumulative aux gradients privilégiés reste finie sur l'horizon d'entraînement, empêchant l'accumulation à long terme de corrélations dégradant la distribution OOD.

3. Analyse Théorique

L'article fournit une justification théorique de la stratégie de routage :

Pression Asymétrique Ponctuelle (Corollaire 2) :
- Le FKL est efficace pour les jetons sous-alloqués (où la probabilité de l'élève $p \ll$ probabilité de l'enseignant $q$ ) car la pression du gradient s'échelonne avec l'écart $(q-p)$ et ne s'annule pas lorsque $p \to 0$ .
- Le RKL est efficace pour les jetons confiants-mais-erronés (où $p \gg q$ ) car le gradient s'échelonne avec la sur-confiance de l'élève ( $p$ ).
Borne d'Exposition Cumulative (Proposition 3) : En combinant le masquage des plages (limitant la bande passante à $\alpha$ ) et la décroissance KL (limitant la durée), TRACE borne l'exposition cumulative aux gradients privilégiés à $O(\alpha \Lambda^2)$ , empêchant la croissance illimitée observée dans les lignes de base KL persistantes sur tous les jetons.
Signal vs Fuite : La méthode maximise le signal d'alignement sur les plages clés tout en minimisant le risque de fuite d'informations privilégiées provoquant des corrélations fallacieuses.

4. Contributions Clés

Diagnostic : Identification d'une « taxe de distillation uniforme » dans le Self-OPD sur tous les jetons, reliant la hausse de l'entropie, la suppression épistémique et la fuite de gradients privilégiés à une inadéquation de granularité.
Méthode et Théorie : Introduction de TRACE, qui achemine $\{FKL, RKL, \emptyset\}$ en fonction de la classe de jeton (Clé, Erreur, Non-plage) et décroît le canal privilégié. La théorie explique pourquoi le FKL domine pour les plages clés dans les régimes à base forte et pourquoi le coin optimal bascule vers le RKL pour les plages d'erreur dans les régimes à base plus faible.
Preuves Empiriques :
- Performance : Sur Qwen3-8B, TRACE améliore la moyenne sur 5 benchmarks de 2,76 points de pourcentage par rapport à GRPO.
- Stabilité : Contrairement à SDPO/SRPO, TRACE préserve les performances OOD du modèle de base sur GPQA-Diamond, là où les lignes de base se dégradent.
- Auto-Annotation : La méthode fonctionne avec un auto-annotateur en ligne (la politique de l'élève elle-même), récupérant ~69 % du gain d'un annotateur API externe puissant, prouvant que les gains proviennent du mécanisme de routage et non simplement des capacités externes.
- Adaptation Trans-Échelle : Sur une base plus faible (Qwen3-1.7B), le coin de routage dominant s'inverse naturellement de FKL-sur-Clé vers RKL-sur-Erreur, correspondant aux prédictions théoriques.

5. Résultats

Qwen3-8B (Base Forte) : TRACE (FKL sur les Plages Clés) atteint une précision moyenne de 81,51 % sur 5 benchmarks (MATH-500, AIME 24/25, AMC 23, GPQA-Diamond), surpassant GRPO (78,75 %) et la ligne de base Self-OPD la plus forte précédente RLSD (78,99 %). Crucialement, elle maintient le score du modèle de base sur GPQA-Diamond (58,33 % contre 58,27 % pour la base), alors que GRPO et les lignes de base sur tous les jetons chutent significativement.
Qwen3-1.7B (Base Plus Faible) : TRACE (RKL sur les Plages d'Erreur) atteint 60,16 % en moyenne, la seule méthode entraînée à dépasser le modèle de base (58,77 %). Les lignes de base sur tous les jetons (SDPO/SRPO) s'effondrent, chutant de 11 à 19 points en dessous de GRPO.
Études d'Ablation :
- Localisation du Masque : Remplacer les masques de plages spécifiques de TRACE par des masques aléatoires de 25 % ou un KL sur tous les jetons entraîne des baisses de performance, confirmant que quels jetons sont sélectionnés importe plus que la quantité.
- Qualité de l'Annotateur : L'utilisation de l'élève en cours d'entraînement actif comme son propre annotateur génère +1,90 point, démontrant que la méthode ne dépend pas de l'importation de capacités de raisonnement « surhumaines » externes.

6. Signification et Revendications

L'article revendique que TRACE offre une solution géométrique de perte à l'instabilité du Self-OPD à horizon long, plutôt qu'une correction spécifique au domaine. Sa signification réside dans :

Découplage de la Distillation et de l'Instabilité : En restreignant le signal « privilégié » aux plages critiques et en le décroissant, TRACE évite la hausse de l'entropie et l'effondrement du raisonnement observés dans SDPO et SRPO.
Routage Adaptatif : Il démontre que la direction optimale de distillation (FKL vs RKL) n'est pas un hyperparamètre fixe mais dépend du profil d'erreur du modèle de base (sous-allocation vs sur-confiance).
Faisabilité Pratique : Il prouve qu'une distillation de haute qualité peut être atteinte sans enseignants externes ni rubriques coûteuses par requête, reposant plutôt sur une annotation diagnostique grossière et légère pouvant être auto-générée.

Les auteurs soulignent que TRACE est un « cadre routé par coin » plutôt qu'une recette fixe, où l'action de routage spécifique (FKL vs RKL) doit être choisie en fonction du mode d'erreur résiduelle du modèle de base.