Résumé technique : Distillation de données hors ligne corrigée par la distribution pour les grands modèles de langage

Énoncé du problème

L'article aborde le défi du transfert des capacités de raisonnement complexes, propres aux grands modèles de langage (LLM) « enseignants » puissants, vers des modèles « élèves » plus petits et contraints en ressources. Bien que la distillation hors ligne (l'entraînement d'un élève sur des traces de raisonnement fixes générées par un enseignant) soit efficace en échantillons et fournisse une supervision de haute qualité, elle souffre d'une dérive distributionnelle fondamentale.

Durant l'entraînement, le modèle élève est conditionné par des préfixes générés par l'enseignant. Cependant, lors de l'inférence, l'élève doit générer ses propres préfixes de manière autorégressive. À mesure que les tokens générés par l'élève s'écartent de la trajectoire de l'enseignant, il pénètre dans des régions de l'espace des tokens sous-représentées dans les données d'entraînement. Cela entraîne des erreurs cumulatives et un comportement peu fiable sur de longues trajectoires de raisonnement, une limitation que les objectifs de distillation affinés existants (par exemple, la minimisation de la divergence KL bidirectionnelle) ne résolvent pas entièrement, car ils optimisent toujours par rapport à une distribution fournie par l'enseignant.

À l'inverse, les méthodes on-policy (par exemple, GRPO, RLHF) atténuent cette dérive en s'entraînant sur des données auto-générées, mais elles entraînent des coûts de calcul élevés pour l'échantillonnage en ligne et peinent souvent avec des traces de faible qualité lors des premières étapes de l'entraînement.

Méthodologie : DISCORD

Les auteurs proposent DISCORD (Distribution-Corrected Distillation), un cadre hors ligne principiel qui conserve l'efficacité des données générées par l'enseignant tout en corrigeant le décalage distributionnel. Le mécanisme central consiste à répondérer de manière adaptative la supervision par les tokens de l'enseignant en fonction du support actuel de l'élève (sa probabilité de générer ce token).

Composants techniques clés :

1. Répondération consciente du support de l'élève :
   La méthode introduit un rapport de densité pour répondérer les trajectoires d'entraînement $\{x, y \sim \pi_T(x; c)\}$ :
   $$\frac{\bar{\pi}_\theta(y|x)}{\bar{\pi}_T(y|x;c)}$$
   où $\pi_\theta$ est l'élève, $\pi_T$ est l'enseignant, et la barre indique une opération de stop-gradient. Ce rapport privilégie les tokens de l'enseignant que l'élève est susceptible de rencontrer lors de l'inférence.

2. Transformation stabilisée :
   L'application directe du rapport de densité provoque une instabilité de l'entraînement en raison de valeurs non bornées et d'une variance élevée sur les longues séquences. Les auteurs appliquent deux stabilisations :

   • Transformation sigmoïde : Le rapport de log-densité est passé à travers une fonction sigmoïde $\sigma(\tau \cdot \log(\dots))$ pour mapper la correction dans une plage bornée $(0, 1)$, préservant l'ordre du rapport tout en empêchant l'explosion.
   • Approximation au niveau du token : Pour atténuer le biais de longueur (où les séquences plus longues accumulent des poids agrégés plus importants), la correction est appliquée au niveau du token. Pour chaque token $y_t$ dans une trajectoire, le poids est :
     $$r_t = \sigma\left(\tau \log \frac{\bar{\pi}_\theta(y_t | x; y_{<t})}{\bar{\pi}_T(y_t | x; c; y_{<t})}\right)$$
     Ce poids est appliqué comme un multiplicateur de gradient lors de l'entraînement de l'objectif de base (par exemple, SFT ou divergence KL).

3. Interprétation :
   Le poids de correction $r_t$ (avec $\tau=1$) est interprété comme la probabilité a posteriori qu'un token ait été tiré de la distribution de l'élève plutôt que de celle de l'enseignant, en supposant une a priori uniforme. Cela relie la méthode aux principes de décodage contrastif, mais les applique durant l'entraînement pour privilégier l'apprentissage dans les régions que l'élève peut habiter de manière fiable.

Contributions clés

• Un cadre hors ligne principiel : DISCORD est la première méthode à corriger efficacement la dérive distributionnelle enseignant-élève en utilisant uniquement des données hors ligne fournies par l'enseignant, éliminant ainsi le besoin de déploiements coûteux d'élèves en ligne.
• Intégration flexible : Le mécanisme de répondération au niveau du token est compatible avec diverses architectures de distillation hors ligne, y compris le fine-tuning supervisé (SFT) standard et les objectifs basés sur la divergence KL.
• Techniques de stabilisation : L'introduction de la transformation sigmoïde et de l'approximation au niveau du token résout les problèmes de variance et d'instabilité inhérents aux corrections basées sur le rapport de densité pour les longues séquences.

Résultats expérimentaux

La méthode a été évaluée sur des benchmarks de raisonnement mathématique (GSM8K, MATH, MATH500) et des tâches compétitives plus difficiles tenues à l'écart (AMC, AIME, OlympiadBench, Omni-Math) en utilisant les modèles Qwen2.5 et Llama-3.2.

• Gains de performance : DISCORD surpasse systématiquement les bases hors ligne solides (SFT, ToDi, Symmetric KL) sur les tâches intra-domaine et extra-domaine (OOD). Par exemple, dans la distillation Qwen2.5-7B $\to$ 3B, DISCORD a amélioré la précision sur MATH500 de 4,60 points par rapport au SFT standard et de 2,60 points par rapport à Symmetric KL.
• Robustesse à la dérive des préfixes : Les évaluations utilisant ExAccErr (divergence distributionnelle) montrent que, tandis que la performance du SFT standard se dégrade considérablement à mesure que la longueur du préfixe augmente (les préfixes générés par l'élève s'écartent de l'enseignant), DISCORD maintient une performance stable, indiquant une réduction des erreurs cumulatives.
• Efficacité par rapport aux méthodes en ligne : DISCORD atteint des performances comparables aux méthodes de déploiement d'élèves en ligne (comme GKD) mais avec un coût de calcul significativement inférieur. Dans le cadre 7B $\to$ 3B, DISCORD n'a requis que 3,12 heures-GPU contre 49,56 heures-GPU pour GKD, sans utiliser de tokens d'élèves en ligne.
• Qualité du raisonnement : L'analyse qualitative révèle que DISCORD produit des traces plus propres, plus concises et logiquement cohérentes. Il réduit les modes d'échec courants tels que la répétition, la continuation après la réponse et l'existence de multiples réponses finales distinctes. Crucialement, dans les cas de « récupération d'échec », DISCORD préserve la structure mathématique sous-jacente de la solution de l'enseignant, même lorsqu'il n'imit pas la mise en forme de surface de l'enseignant.

Importance et affirmations

L'article affirme qu'un entraînement léger, conscient de la correction distributionnelle, peut considérablement renforcer la distillation de raisonnement hors ligne sans nécessiter d'interaction en ligne.

• Combler le fossé : Le travail démontre qu'il est possible de réduire l'écart de performance entre les méthodes de distillation hors ligne et en ligne en répondérant intelligemment les données hors ligne existantes, plutôt qu'en générant de nouvelles données.
• Généralisabilité : Bien que les expériences se concentrent sur le raisonnement mathématique, les auteurs postulent que la méthode est applicable à d'autres contextes de génération à long horizon, y compris la génération de code et le raisonnement scientifique.
• Déploiement pratique : En évitant les coûts de calcul élevés de l'échantillonnage en ligne, DISCORD offre une voie viable pour déployer des modèles de raisonnement capables dans des environnements contraints en ressources.

Les auteurs restent modestes, notant que si leur méthode améliore les bases hors ligne précédentes, les méthodes en ligne comme GKD conservent encore un avantage dans certains benchmarks spécifiques intra-domaine en raison de leur observation directe de la distribution de déploiement de l'élève. Cependant, DISCORD offre un compromis supérieur entre performance et efficacité computationnelle.