Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

Résumé Technique : Ponter les trajectoires de raisonnement dans la distillation on-policy via le guidage vers le futur proche

Énoncé du problème
La distillation on-policy (OPD) est une pierre angulaire de l'entraînement post-apprentissage moderne pour les grands modèles de langage (LLM), conçue pour améliorer le raisonnement en entraînant un modèle étudiant sur des trajectoires échantillonnées à partir de sa propre politique sous la supervision d'un enseignant. L'OPD standard fonctionne selon un mécanisme de « échantillonnage de trajectoire mais apprentissage par jeton » : elle identifie les écarts via des jetons à perte élevée et les répare en utilisant une correction de divergence de KL inverse locale.

Les auteurs identifient un décalage fondamental dans cette approche, arguant que la supervision au niveau du jeton est structurellement désalignée avec les exigences de la correction de raisonnement au niveau de la trajectoire. Ce désalignement se manifeste par deux modes d'échec critiques :

1. Fausses alertes bruyantes : Les jetons à perte élevée ne correspondent pas toujours à de véritables divergences de raisonnement. Une analyse empirique révèle qu'environ 30 % des jetons à perte élevée tombent dans un régime de faible divergence, représentant des écarts de forme de surface (par exemple, des préférences stylistiques comme « et » versus « puis ») plutôt que des embranchements logiques. Traiter ces jetons comme des erreurs critiques fournit une supervision bruitée qui peut activement interférer avec la correction du raisonnement.
2. Le piège de l'apprentissage jeton par jeton : Même lorsque de véritables points de divergence sont identifiés, la correction de la divergence KL inverse par jeton unique est souvent insuffisante pour rediriger l'ensemble du chemin de raisonnement. Les échecs de raisonnement évoluent généralement comme des dérives distributionnelles à court horizon. Corriger un jeton localement ne garantit pas que le modèle s'alignera sur la future trajectoire de l'enseignant ; le modèle dérive souvent de nouveau vers des branches erronées immédiatement après la correction locale.

Méthodologie : OPD sensible à la trajectoire (TOPD)
Pour remédier à ces limitations, les auteurs proposent la TOPD (Trajectory-aware OPD), qui comble le fossé entre la correction de jeton locale et l'évolution de la trajectoire globale en incorporant des informations de trajectoire du futur proche. La méthode se compose de deux étapes principales :

1. Détection de la divergence :

   • Pour les positions candidates à perte élevée, le système revient au jeton précédent pour établir un préfixe commun.
   • L'enseignant et l'étudiant génèrent tous deux des continuations à fenêtre courte (longueur $K=50$) à partir de ce préfixe.
   • Le système mesure la disparité de niveau de trajectoire en utilisant la distance de transport optimal (OT) entre les continuations de l'enseignant et de l'étudiant.
   • Les véritables points de divergence sont identifiés comme des positions présentant à la fois une perte élevée au niveau du jeton et une divergence OT élevée sur la fenêtre courte. Cela filtre les écarts de surface bénins.

2. Injection de signal de l'enseignant au niveau de la trajectoire :

   • Au lieu d'optimiser uniquement le jeton actuel, la TOPD étend la supervision à l'ensemble de la trajectoire de la fenêtre courte.
   • En utilisant le plan de transport OT ($\gamma$), la méthode construit une « cible douce sensible à la trajectoire » pour chaque position future de l'étudiant. Cette cible aligne les futurs jetons de l'étudiant avec la trajectoire de l'enseignant sur la base de la cartographie de transport optimal.
   • L'étudiant est entraîné à correspondre à ces cibles douces alignées par OT à travers la fenêtre, apprenant ainsi efficacement comment diriger son chemin de raisonnement futur vers le chemin de l'enseignant, et non pas seulement comment corriger le jeton actuel.
   • L'objectif final combine la perte OPD standard avec un terme de perte de trajectoire ($L_{TOPD} = L_{OPD} + \lambda L_{traj}$).

Contributions clés

• Identification du décalage : Ce travail établit que la supervision au niveau du jeton est un « proxy grossier » pour la correction du raisonnement au niveau de la trajectoire, échouant à distinguer les véritables embranchements logiques des variations de surface bénignes.
• Caractérisation empirique des modes d'échec : Les auteurs fournissent des preuves qu'environ 30 % des jetons à perte élevée sont des fausses alertes de faible divergence et démontrent le « piège de l'apprentissage jeton par jeton » où les corrections locales échouent à restaurer la trajectoire de raisonnement globale.
• Principe de guidage au niveau de la trajectoire : Le papier introduit un principe où l'information de la trajectoire du futur proche est distribuée sur plusieurs jetons. En utilisant l'alignement OT sur fenêtre courte, la TOPD injecte directement les écarts de chemin enseignant-étudiant dans la fonction de perte, permettant à l'étudiant d'apprendre des changements de trajectoire plutôt que des éditions de jetons isolés.

Résultats expérimentaux
Des expériences ont été menées avec des modèles Qwen3 (distillant un enseignant de 30B vers un étudiant de 4B) sur des benchmarks mathématiques compétitifs (AIME24, AIME25, HMMT25-Feb).

• Gains de performance : L'OPD standard a amélioré la précision moyenne du modèle de démarrage à chaud de 38,9 % à 47,8 %. La TOPD a augmenté cela à 52,2 %.
• Benchmarks spécifiques : La TOPD a amélioré la précision de l'AIME24 de 60,0 % à 63,3 % et de l'AIME25 de 46,7 % à 53,3 %.
• Études d'ablation :
  • La suppression des jetons à haute perte et faible OT (fausses alertes) dans l'OPD standard a produit un gain modeste (+0,4 %), confirmant que ces jetons sont bruités.
  • La suppression des jetons à haute perte et haute OT (véritables divergences) a provoqué une chute significative de la performance, confirmant leur nécessité.
  • L'analyse de la correction de trajectoire locale a montré que si l'OPD standard réduisait partiellement la distance OT de la fenêtre courte (de 0,538 à 0,350), la TOPD a obtenu une réduction substantielle (à 0,204), démontrant une capacité supérieure à rediriger les trajectoires futures proches.
• Efficacité : La TOPD introduit un surcoût de calcul modeste, le temps d'entraînement étant de 1,41 $\times$ celui de l'OPD standard, car l'OT est calculé uniquement sur les fenêtres sélectionnées à perte élevée.

Signification et affirmations
Le papier affirme que les échecs de raisonnement dans l'OPD sont fondamentalement des phénomènes de niveau de trajectoire plutôt que des décalages de jetons isolés. En passant d'une édition de jeton ponctuelle à un alignement global sensible aux transitions, la TOPD permet un alignement du raisonnement plus efficace. Les auteurs positionnent ce travail comme une étape vers des objectifs de distillation plus sensibles à la trajectoire, arguant qu'une correction de raisonnement réussie nécessite que l'étudiant apprenne non seulement quel jeton corriger, mais aussi comment naviguer dans le chemin de raisonnement subséquent pour rester cohérent avec l'enseignant.

Limitations
Les auteurs reconnaissent que la TOPD entraîne des coûts de calcul supplémentaires dus aux calculs d'OT, se concentre principalement sur la divergence à horizon court (pouvant manquer des erreurs de contexte long) et est actuellement validée principalement sur des benchmarks de raisonnement mathématique plutôt que sur des tâches de code ou de langage ouvert. De plus, la distance OT mesure la proximité des trajectoires mais peut ne pas capturer l'équivalence sémantique là où des chemins différents mènent à la même réponse correcte.