Résumé technique : Combattre la dégradation de la fidélité de la supervision dans la distillation on-policy

Problème : Dégradation de la fidélité de la supervision (SFD)

Le papier identifie un goulot d'étranglement critique dans la distillation on-policy (OPD), un paradigme où un modèle étudiant apprend à partir de ses propres trajectoires générées en utilisant un feedback au niveau du token provenant d'un modèle enseignant. Bien que l'OPD ait été efficace pour transférer les capacités de raisonnement, les auteurs observent que ses performances se dégradent considérablement à mesure que la longueur de la génération augmente (par exemple, les performances plafonnent autour de 16k tokens et déclinent à 39k tokens).

Les auteurs attribuent cela à la dégradation de la fidélité de la supervision (Supervision Fidelity Decay - SFD). À mesure que l'étudiant génère des séquences plus longues, ses sorties dérivent progressivement de la distribution d'entraînement de l'enseignant (hors distribution ou OOD - Out-of-Distribution). Cette dérive fait que la distribution du prochain token de l'enseignant devient moins confiante et moins discriminante.

• Mécanisme : Dans la distillation reverse-KL standard, le signal d'apprentissage (avantage) repose sur la différence entre la log-probabilité de l'étudiant et la log-probabilité de l'enseignant. À mesure que la distribution de l'enseignant devient diffuse (entropie élevée) en raison de la SFD, la variance des log-probabilités de l'enseignant s'effondre.
• Conséquence : La contribution discriminative de l'enseignant au gradient s'annule. Le signal d'apprentissage se réduit effectivement à un signal « uniquement étudiant » qui renforce les modes existants de l'étudiant sans correction. Cela crée une boucle de rétroaction auto-renforçante où la dérive s'accumule, réduisant l'horizon d'apprentissage effectif et établissant un plafond sur la longueur du raisonnement.

Méthodologie : Récompense de groupe avec anticipation (Lookahead Group Reward - LGR)

Pour remédier à la SFD, les auteurs proposent la récompense de groupe avec anticipation (Lookahead Group Reward - LGR), une méthode qui maintient activement la qualité de la supervision en optimisant la confiance future de l'enseignant plutôt que la correspondance locale actuelle.

Intuition centrale

Même lorsque la supervision locale de l'enseignant à la position $t$ s'est effondrée (c'est-à-dire que la distribution est proche de l'uniformité), différents choix de tokens à $t$ peuvent conduire à des niveaux de confiance de l'enseignant significativement différents à la position suivante ($t+1$). Le papier démontre théoriquement qu'une discriminabilité un pas en avant peut survivre même lorsque la supervision locale échoue. En sélectionnant les tokens qui maximisent la confiance de l'enseignant à l'étape suivante, l'étudiant peut ralentir le taux de dérive.

Conception de l'algorithme

1. Récompense de confiance : À chaque position $t$, la méthode évalue la probabilité maximale du prochain token de l'enseignant ($\max_v \pi_T(v|x_{<t}, x_t)$) pour les $K$ meilleurs candidats de l'étudiant. Cette valeur sert de signal de récompense indiquant comment bien un candidat préserve la capacité de l'enseignant à fournir une supervision de haute qualité à l'étape suivante.
2. Normalisation de groupe : Pour gérer la variance des niveaux de confiance absolue à travers différentes tâches et positions, les récompenses brutes sont normalisées au sein du groupe des $K$ meilleurs candidats (similairement à GRPO). Cela conserve le classement relatif des candidats tout en supprimant le bruit d'échelle absolue.
3. Activation déclenchée par l'entropie : Calculer la récompense d'anticipation pour tous les tokens est coûteux en termes de calcul. Les auteurs observent que la SFD et le besoin de guidance par anticipation sont les plus critiques aux positions de haute entropie (où l'étudiant est incertain). Par conséquent, la récompense de confiance n'est calculée que pour les positions où l'entropie de génération de l'étudiant dépasse un seuil $\tau$. Cela réduit la charge de calcul de 4 à 7 fois.
4. Mécanisme d'attention en arbre (Tree-Attention) : Pour évaluer efficacement $K$ candidats pour plusieurs positions sans réexécuter le prefill de l'enseignant pour chaque candidat, les auteurs emploient un mécanisme d'attention en arbre. Celui-ci construit une séquence étendue où la branche principale est la trajectoire de l'étudiant et les tokens candidats sont ajoutés comme des branches. Un masque d'attention parcimonieux permet à l'enseignant de traiter tous les candidats en un seul passage avant (forward pass), atteignant une accélération théorique allant jusqu'à 1000$\times$ par rapport à une implémentation naïve.

La fonction de perte finale combine l'avantage reverse-KL standard avec la récompense de confiance par anticipation, pondérée par un hyperparamètre $\gamma$.

Contributions clés

1. Identification de la SFD : Le papier caractérise formellement la dégradation de la fidélité de la supervision comme un mode de défaillance fondamental de la distillation on-policy reverse-KL, prouvant qu'elle réduit le gradient en une boucle de dérive auto-renforçante qui limite la longueur du raisonnement.
2. Récompense de groupe avec anticipation (LGR) : Les auteurs proposent un remède fondé sur l'anticipation d'un pas. Au lieu de tenter de correspondre à une distribution dégradée, LGR sélectionne les tokens qui maximisent la confiance future de l'enseignant, dirigeant ainsi activement l'étudiant vers des états où la supervision reste fiable.
3. Implémentation efficace : La conception d'un mécanisme d'attention en arbre déclenché par l'entropie rend le processus d'anticipation, pourtant gourmand en calcul, pratique, permettant l'entraînement sur de longs contextes (jusqu'à 39k tokens) sans coûts prohibitifs.

Résultats expérimentaux

La méthode a été évaluée sur six benchmarks couvrant le raisonnement mathématique (AIME-24, AIME-25, AIME-26, HMMT-25, HMMT-26) et la génération de code (LiveCodeBench v6).

• Gains de performance : Pour un modèle étudiant de 7B, LGR a amélioré le score mean@8 de 2,57 points par rapport à l'OPD standard (9k tokens) en moyenne.
• Mise à l'échelle du contexte long : Les bénéfices de LGR augmentent avec la longueur de la génération. Sur le benchmark AIME-26 à 39k tokens, LGR a obtenu un gain de +4,92 points par rapport à l'OPD. En revanche, les performances de l'OPD standard se sont considérablement dégradées à cette longueur.
• Comparaison avec les baselines : LGR a surpassé d'autres méthodes de distillation, y compris JSD, REOPOLD et les variantes standard de l'OPD. Notamment, les méthodes qui ne font que stabiliser l'entraînement (comme REOPOLD) sans traiter la dégradation de la qualité de la supervision n'ont pas réussi à égaler les performances de LGR.
• Études d'ablation :
  • Max-p vs Entropie : L'utilisation de la probabilité maximale de l'enseignant (Max-p) comme métrique de confiance a surpassé l'utilisation de l'entropie ou de la perplexité, ce qui concorde avec la découverte théorique que la Max-p est un proxy direct de la proximité avec la variété (manifold) in-distribution de l'enseignant.
  • Sensibilité au poids : Le poids de la récompense de confiance $\gamma$ a montré une performance optimale dans la plage de 1,0 à 1,5. Des poids excessivement élevés ont conduit à un détournement de récompense (reward hacking) sous forme de séquences répétitives.

Signification et affirmations

Le papier affirme que LGR fournit une solution structurellement différente au problème du raisonnement à long horizon dans la distillation. Contrairement aux stratégies qui évitent passivement les régions de faible supervision (ex: troncature) ou qui stabilisent simplement l'optimisation, LGR dirige activement l'étudiant vers des trajectoires où l'enseignant maintient une haute fidélité de supervision.

Les auteurs soulignent que leur travail traite la cause profonde de la dégradation des performances dans les chaînes de raisonnement longues : l'effet cumulatif de la dérive de l'étudiant sur la confiance de l'enseignant. En opérationnalisant l'idée qu'une « plus grande confiance au prochain pas implique un signal futur plus discriminant », LGR brise le cycle vicieux de la dégradation de la supervision, permettant aux modèles de passer efficacement à des chaînes de raisonnement beaucoup plus longues (jusqu'à 39k tokens) que ce qui était possible auparavant avec la distillation on-policy standard.