🤖 AI

Trajectory-Refined Distillation

Cet article introduit la Distillation par Raffinement de Trajectoire (TRD), une nouvelle méthode de distillation on-policy qui atténue le problème structurel d'« échec de préfixe » en appliquant des corrections au niveau de la trajectoire aux déploiements de l'étudiant sous la direction de l'enseignant, améliorant ainsi la précision et la couverture du raisonnement à travers divers benchmarks et échelles de modèles.

Auteurs originaux : Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

Publié 2026-06-09

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

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 enseignez à un étudiant (l'IA) comment résoudre un problème mathématique complexe. Vous avez un enseignant brillant (une IA plus avancée) qui connaît la bonne réponse.

Dans la méthode standard décrite dans l'article, appelée On-Policy Distillation (OPD), le processus fonctionne comme suit :

L'étudiant essaie de résoudre le problème par lui-même.
À mesure que l'étudiant écrit ses étapes, l'enseignant regarde et dit : « Pour ce mot spécifique que tu viens d'écrire, voici la probabilité de ce qui vient ensuite. »
L'étudiant essaie d'ajuster son cerveau pour correspondre aux probabilités de l'enseignant pour chaque mot.

Le Problème : « Prefix Failure » (Le mauvais tournant)

L'article identifie une faille majeure dans cette méthode standard appelée Prefix Failure (Échec de préfixe).

Imaginez que l'étudiant suit un chemin pour résoudre le problème. À l'étape 3, il prend un mauvais tournant. Il se retrouve sur une impasse.

Le dilemme de l'enseignant : L'enseignant voit que l'étudiant est perdu. L'enseignant sait que le bon chemin commence par « Attendez, en fait... » et prend ensuite une direction complètement différente.
La confusion : Mais l'enseignant est contraint de regarder l'emplacement actuel de l'étudiant (l'impasse). Ainsi, le conseil de l'enseignant devient un mélange confus : « Reste sur cette impasse parce que tu es déjà ici » ET « Fais demi-tour immédiatement. »
Le résultat : L'étudiant reçoit un signal brouillé. Il essaie d'apprendre de l'enseignant, mais parce que l'étudiant est coincé sur le mauvais chemin, le conseil de l'enseignant est fragmenté. L'étudiant apprend à être « bon pour être perdu » plutôt qu'à apprendre comment revenir sur le bon chemin. L'article appelle cela un « mélange bimodal » — l'enseignant essaie d'enseigner deux choses opposées à la fois, et l'étudiant est confus.

Les solutions existantes ont tenté de corriger cela en ajustant simplement la « note » (la fonction de perte) pour des mots spécifiques, mais elles n'ont pas corrigé le fait que l'étudiant était toujours coincé sur le mauvais chemin.

La Solution : Trajectory-Refined Distillation (TRD)

Les auteurs proposent une nouvelle méthode appelée Trajectory-Refined Distillation (TRD). Au lieu de simplement noter le brouillon désordonné de l'étudiant mot par mot, ils changent le flux de travail :

Le Brouillon : L'étudiant fait toujours une première tentative pour le problème (le « raw rollout » ou déploiement brut).
La Réécriture : Avant la notation, l'enseignant prend ce brouillon désordonné et le réécrit. L'enseignant corrige les mauvais tournants, rectifie la logique et produit une version propre et parfaite de la solution basée sur le point de départ de l'étudiant.
La Leçon : L'étudiant apprend alors de cette version propre et corrigée, et non du brouillon désordonné original.

L'Analogie :

L'ancienne méthode : Vous écrivez une histoire avec une incohérence dans l'intrigue au milieu. Votre éditeur lit votre texte et dit : « Pour le mot 'chat' au paragraphe 3, vous auriez dû écrire 'chien'. » Mais vous êtes toujours bloqué au paragraphe 3 en essayant de comprendre pourquoi vous avez écrit 'chat' en premier lieu. C'est déroutant.
La méthode TRD : Vous écrivez l'histoire avec l'incohérence. Votre éditeur prend votre brouillon, corrige l'incohérence, et réécrit tout le milieu de l'histoire pour que le récit soit fluide. Ensuite, vous étudiez la version réécrite par l'éditeur pour apprendre à mieux écrire des histoires la prochaine fois.

Pourquoi cela fonctionne mieux

Corrige la cause profonde : En corrigeant le « mauvais tournant » avant que l'étudiant ne tente d'apprendre de celui-ci, le conseil de l'enseignant n'est plus confus. C'est un chemin unique et clair.
Explore de nouveaux chemins : Même si l'étudiant trouve la bonne réponse dès la première tentative, l'enseignant peut lui montrer un chemin différent, plus court ou plus intelligent pour résoudre le problème. Cela apprend à l'étudiant à être plus flexible.
Efficacité : L'article a constaté que les chemins « réécrits » étaient souvent beaucoup plus courts et clairs que les tentatives originales et confuses de l'étudiant, ce qui rendait l'entraînement plus rapide.

Les Résultats

Les chercheurs ont testé cette méthode sur des compétitions mathématiques difficiles (comme l'AIME et l'HMMT) et des défis de codage. Ils ont constaté que :

La nouvelle méthode (TRD) a systématiquement battu les anciennes méthodes.
Elle était particulièrement efficace pour aider l'IA à résoudre les problèmes les plus difficiles où elle restait habituellement bloquée.
Elle fonctionnait que l'enseignant soit une IA distincte et plus grande, ou la même IA agissant comme son propre « enseignant intelligent » (en utilisant une technique appelée « information privilégiée »).

En résumé, l'article soutient que pour enseigner efficacement à une IA, vous ne devriez pas simplement noter ses erreurs mot par mot pendant qu'elle les commet. Au lieu de cela, vous devriez laisser l'enseignant corriger tout le chemin d'abord, puis laisser l'étudiant apprendre de ce voyage corrigé.

Résumé Technique : Trajectory-Refined Distillation (TRD)

1. Énoncé du Problème : Échec de Préfixe dans la Distillation On-Policy

L'article identifie une limitation structurelle fondamentale dans la Distillation On-Policy (OPD) et l'Auto-Distillation On-Policy (OPSD) pour les Grands Modèles de Langage (LLM), appelée échec de préfixe (prefix failure).

Dans l'OPD/OPSD standard, un modèle étudiant génère un déroulement (rollout) $y_o$ , et un modèle enseignant fournit une supervision dense, par jeton (token), basée sur ce déroulement spécifique. Les auteurs soutiennent que lorsque le déroulement de l'étudiant emprunte un chemin de raisonnement incorrect (un « préfixe défaillant »), le signal de supervision de l'enseignant devient structurellement compromis :

Mélange Bimodal de l'Enseignant : Lorsqu'un préfixe est incorrect, la distribution de l'enseignant devient un mélange de deux modes : l'un qui poursuit le préfixe défaillant (pour la cohérence de la séquence) et un autre qui pivote vers la solution correcte.
Gradients Fragmentés : Les fonctions de perte par jeton standard (divergence KL) évaluent la supervision le long du trajet défaillant gelé. Par conséquent, les mises à jour de gradients sont fragmentées. L'enseignant recommande de manière répétée un jeton de correction (ex: « Attendez » ou « En fait ») conditionné sur le mauvais contexte (le préxiur défaillant), plutôt que de se conditionner sur le chemin de correction en cours de déploiement.
Inefficacité des Corrections au Niveau du Jeton : Les remèdes existants (ex: écrêtage de la perte, troncature Top-K, repondération) opèrent au niveau de la perte par jeton. Ils tentent de sélectionner des jetons « informatifs » ou de supprimer des jetons « bruyants », mais échouent à traiter la cause profonde : le préfixe défaillant lui-même reste structurellement intact. Les paires de supervision délivrées à l'étudiant divergent du chemin de correction idéal, empêchant le modèle d'apprendre efficacement à se rétablir de ses erreurs.

2. Méthodologie : Trajectory-Refined Distillation (TRD)

Pour atténuer l'échec de préfixe, les auteurs proposent la Trajectory-Refined Distillation (TRD), une méthode qui déplace l'intervention du niveau du jeton vers le niveau du trajet (trajectory).

Mécanisme Central

Au lieu d'utiliser directement le déroulement brut de l'étudiant $y_o$ pour la distillation, la TRD introduit une étape de raffinement avant le calcul de la perte :

Échantillonnage : Le modèle étudiant $\pi_\theta$ génère un déroulement on-policy brut $y_o \sim \pi_\theta(\cdot | x)$ .
Raffinement : Le modèle enseignant $\pi_T$ $π_{T}$ (ou l'étudiant lui-même dans l'OPSD, conditionné par des informations privilégiées comme la vérité terrain $y^*$ $y^{*}$ ) génère un trajet raffiné $y_r$ $y_{r}$ .
- Dans l'OPD : $y_r \sim \pi_T(\cdot | x, y_o)$ .
- Dans l'OPSD : $y_r \sim \text{sg}[\pi_\theta(\cdot | x, y_o, y^*)]$ .
- Crucialement, l'enseignant est incité à réviser le déroulement de l'étudiant guidé par la solution de référence, « réécrivant » ainsi les parties erronées du préfixe tout en maintenant les schémas de raisonnement de l'étudiant là où ils sont valides.
Distillation : Le trajet raffiné $y_r$ est utilisé comme cible de supervision pour la perte standard d'appariement de distribution (KL directe avec le vocabulaire complet).

Justification Théorique

La TRD traite le problème des « gradients fragmentés » en garantissant que le contexte de supervision croît le long du chemin de correction plutôt que le long du préfixe défaillant.

Récupération Idéale du Gradient : En supervisant le long de $y_r$ , le contexte de chaque jeton inclut le préfixe corrigé, permettant à l'enseignant de fournir un signal de correction cohérent et multi-étapes ( $g_{ideal}$ ) plutôt qu'un signal répétitif et mal adapté au contexte ( $g_{frag}$ ).
Exploration : Même lorsque $y_o$ est correct, la TRD expose l'étudiant à des dérivations valides alternatives suggérées par l'enseignant, élargissant la distribution des chemins de raisonnement corrects que l'étudiant apprend.

Algorithme

Le processus est résumé dans l'Algorithme 1 :

Échantillonner $y_o \sim \pi_\theta(\cdot | x)$ .
Générer $y_r \sim \pi_T(\cdot | x, y_o)$ .
Mettre à jour $\theta$ en utilisant la divergence KL entre $\pi_\theta$ et la distribution de l'enseignant sur $y_r$ .

3. Contributions Clés

Identification de l'Échec de Préfixe : L'article formalise l'« échec de préfixe » comme une cause structurelle de dégradation dans la distillation on-policy, expliquant pourquoi les interventions au niveau du jeton échouent à récupérer des erreurs de raisonnement.
Correction au Niveau du Trajet : La TRD est proposée comme la première conception d'optimisation au niveau du trajet qui cible explicitement la récupération des préfixes défaillants tout en respectant la contrainte de support on-policy, sans surcoût computationnel prohibitif.
Amélioration à Double Régime : La TRD améliore les performances dans deux régimes :
- Correction d'Erreur : Elle atténue l'échec de préfixe en corrigeant les préfixes problématiques avant la distillation.
- Exploration : Elle élargit la couverture du raisonnement en exposant les étudiants à des dérivations valides alternatives, même lorsque le déroulement original était correct.
Applicabilité à l'Auto-Distillation : La méthode s'étend naturellement à l'OPSD, où l'étudiant agit comme son propre enseignant conditionné par des informations privilégiées (vérité terrain), éliminant le besoin d'un modèle enseignant distinct.

4. Résultats Expérimentaux

Les auteurs ont évalué la TRD sur cinq benchmarks de mathématiques de compétition (AIME24/25, HMMT25, BeyondAIME, AMOBench) et des benchmarks de génération de code (HumanEval+, MBPP+, LiveCodeBench) en utilisant des modèles Qwen3 (1.7B, 4B, 8B).

Principales Découvertes

Performance Supérieure : La TRD surpasse systématiquement les bases précédentes (KL directe/inverse, avec/sans écrêtage/troncature) dans les configurations OPD et OPSD.
Gains Significatifs sur les Tâches Difficiles : Les améliorations sont plus marquées sur les benchmarks les plus complexes.
- AMOBench : La TRD a obtenu une amélioration relative d'environ 50 % du Pass@16 pour Qwen3-8B sous la configuration OPSD par rapport au modèle de base.
- AIME24/25 : La TRD a montré des gains constants, particulièrement pour les petits modèles (ex : +4,6 % sur AIME24 pour Qwen3-1.7B).
Stabilité : Contrairement à certaines bases qui se dégradent sur certains benchmarks (ex : la KL inverse qui fait chuter les performances sur AIME24), la TRD n'est jamais descendue en dessous de la performance du modèle de base.
Efficacité : Dans le cadre de l'OPSD, les trajets raffinés $y_r$ étaient nettement plus courts (longueur médiane compressée d'environ 9x) que les déroulements bruts $y_o$ , réduisant le temps de calcul de l'entraînement d'environ 60 % malgré l'étape d'échantillonnage supplémentaire.
Génération de Code : La TRD a égalé ou dépassé les bases sur HumanEval+ et MBPP+, bien que les gains sur LiveCodeBench aient été limités, suggérant que la capacité de raffinement de l'enseignant peut être limitée par la difficulté de la tâche.

5. Signification et Revendications

L'article affirme que la TRD représente un changement nécessaire dans la stratégie de post-entraînement des LLM. Les auteurs soutiennent que :

Les interventions au niveau du jeton sont insuffisantes : La nature structurelle de l'échec de préfixe nécessite des corrections au niveau du trajet.
Le support on-policy est préservé : La TRD parvient à ces corrections sans violer la contrainte on-policy (contrairement à l'utilisation d'une politique experte fixe), ce qui la rend compatible avec les paradigmes d'apprentissage efficaces en calcul.
L'auto-amélioration est viable : Le succès de la TRD dans le cadre de l'OPSD démontre que les modèles peuvent effectivement raffiner leurs propres chemins de raisonnement en utilisant des informations privilégiées (vérité terrain) sans nécessiter de modèle enseignant distinct et plus grand, offrant une voie légère vers l'auto-amélioration.

Ce travail suggère que les futures méthodes de distillation doivent tenir compte de l'intégrité structurelle du trajet de raisonnement, et non seulement de l'alignement statistique des jetons individuels.