Résumé Technique : Auto-distillation Géométrique pour la Généralisation du Raisonnement

1. Énoncé du Problème

La distillation on-policy (OPD) est apparue comme une stratégie de post-entraînement pratique pour les grands modèles de langage (LLM), fournissant une supervision dense au niveau des tokens sur les trajectoires générées par le modèle étudiant lui-même. Une variante spécifique et efficace, l'auto-distillation à contexte privilégié (OPSD), conditionne l'« enseignant » (le même modèle) sur des informations supplémentaires, telles qu'une trace de solution complète ou un indice, tandis que l'étudiant doit raisonner à partir du contexte de problème d'origine seul.

Bien que cette configuration fournisse une supervision abondante, l'article identifie un mode de défaillance critique : la dérive prédictive (predictive drift). Parce que l'enseignant possède un contexte privilégié, il peut s'engager avec confiance dans des continuations qui sont évidentes compte tenu de la solution, mais qui restent incertaines ou non soutenues du point de vue de l'étudiant. Les objectifs de distillation standards (par exemple, la divergence de Kullback–Leibler) pondèrent les préférences de l'enseignant par sa propre masse de probabilité. Par conséquent, aux états où l'enseignant est très confiant mais où l'étudiant assigne une faible probabilité (faible chevauchement), le signal de gradient est fort. Au fil des mises à jour, ces tractions fortes forcent l'étudiant à adopter les continuations confiantes de l'enseignant avant que l'information propre à l'étudiant ne puisse les justifier. Cette accumulation de mises à jour discordantes mène à la dérive, où le modèle supprime l'incertitude et dégrade les performances de raisonnement hors distribution (OOD), entraînant souvent un « accord confiant sur de mauvaises réponses ».

Les approches existantes tentent d'atténuer cela en filtrant les signaux de l'enseignant peu fiables ou en remodelant les poids de divergence (par exemple, en utilisant la KL inverse ou des mécanismes de filtrage). Cependant, les auteurs soutiennent que ces méthodes ne parviennent pas à traiter le problème de fond : jusqu'où un signal de confiance devrait déplacer l'étudiant. La sélection seule ne peut pas réguler l'amplitude de la mise à jour ; un mouvement non contrôlé s'accumule en une dérive globale.

2. Méthodologie : GEOSD

Les auteurs introduisent la GEOSD (Geometric Self-Distillation), un objectif conçu pour contrer cette dérive en traitant le comportement prédictif comme un mouvement au sein d'un espace géométrique spécifique de distributions de tokens suivants. La méthode opère de deux manières complémentaires :

A. Atténuation de la Pression Locale (Perte de Hellinger)

Au lieu de la divergence KL standard, la GEOSD emploie une perte de Hellinger pour moduler la traction de l'enseignant à chaque état visité.

• Mécanisme : La distance de Hellinger est dérivée de la géométrie du simplexe de probabilité sous un plongement de racine carrée ($p \to \sqrt{p}$), où les distributions résident sur une hypersphère. Le gradient de la perte de Hellinger met à l'échelle chaque préférence de l'enseignant par la moyenne géométrique des probabilités de l'étudiant et de l'enseignant ($\sqrt{p_i q_i}$).
• Effet : Si l'enseignant assigne une probabilité élevée à un token mais que l'étudiant assigne une probabilité proche de zéro, le terme de chevauchement s'annule, et la traction du gradient est atténuée. Inversement, les préférences partagées sont transférées fortement. Cela garantit que la supervision se renforce progressivement à mesure que l'étudiant et l'enseignant entrent en accord, plutôt que de forcer l'étudiant à adopter immédiatement les continuations confiantes mais non soutenues de l'enseignant.

B. Contrôle de la Dérive Globale (Terme Proximal de Fisher–Rao)

Bien que la pondération par chevauchement limite la traction à des états individuels, elle ne permet pas d'empêcher l'accumulation de petites mises à jour raisonnables en une dérive globale importante au cours de l'entraînement.

• Mécanisme : La GEOSD ajoute un terme proximal qui pénalise la distance entre la distribution actuelle de l'étudiant et un point de contrôle récent ($\theta_{ckpt}$). Cette distance est mesurée à l'aide de la distance de Fisher–Rao (la longueur d'arc sur l'hypersphère), qui suit le déplacement accumulé dans la géométrie des distributions prédictives.
• Effet : Cela agit comme une « région de confiance mobile », amortissant les déplacements abrupts tout en permettant un mouvement constant lorsque le signal de l'enseignant est cohérent. Contraoth l'ancrage au modèle initial (qui interdirait l'adaptation), l'ancrage à un point de contrôle récent permet au modèle de s'adapter tout en l'empêchant de dériver trop loin de son comportement récent et stable.

C. Mise à jour par Gradient Naturel

La perte de Hellinger et le terme proximal de Fisher–Rao sont tous deux des distances dans la même géométrie de distributions de tokens suivants. Pour optimiser cet objectif, la GEOSD utilise une mise à jour par gradient naturel.

• Mécanisme : Au lieu de mettre à jour les paramètres dans l'espace euclidien, la mise à jour est préconditionnée par la Matrice d'Information de Fisher (approximée via K-FAC pour l'efficacité).
• Effet : Cela garantit que chaque étape effectue le plus petit changement possible aux prédictions de l'étudiant pour améliorer l'objectif, plutôt que de faire un pas de taille fixe dans l'espace des paramètres. Le terme proximal agit efficacement comme un « rappel vers le point de contrôle » dans cette géométrie, ramenant la politique vers le point de contrôle en fonction de la saillance comportementale des directions de paramètres.

3. Contributions Clés

1. Objectif GEOSD : Un cadre d'auto-distillation qui pondère les préférences de l'enseignant par le chevauchement enseignant-étudiant (via la perte de Hellinger) et pénalise la dérive prédictive par rapport à un point de contrôle récent (via la distance de Fisher–Rao), optimisé via des gradients naturels dans une géométrie unifiée.
2. Amélioration Empirique : La GEOSD préserve les gains ID (en distribution) de l'auto-distillation tout en améliorant significativement la performance de raisonnement OOD (hors distribution). Sur trois familles de modèles (Qwen3, Olmo, DeepSeek-R1) et cinq échelles (de 1.7B à 32B), elle améliore la précision OOD moyenne de 5,7 à 8,6 points par rapport au modèle de base, alors que les objectifs KL standards dégradent la performance OOD.
3. Analyse Mécaniste : L'article caractérise pourquoi l'appariement standard échoue en OOD : il concentre la masse de probabilité sur le meilleur choix de l'enseignant en « drainant » les alternatives lors des états à haute entropie, menant à un accord confiant sur de mauvures réponses. La GEOSD maintient la dominance du meilleur choix tout en gardant les continuations alternatives « à portée », empêant l'effondrement prématuré.

4. Résultats Expérimentaux

Les auteurs ont évalué la GEOSD sur des benchmarks de raisonnement mathématique (DAPO-Math pour l'entraînement ; AIME 2024/2025, AMC 2023, MATH-500 pour les tests OOD).

• Performance : La GEOSD a obtenu la meilleure précision OOD sur tous les benchmarks et familles de modèles testés. Par exemple, sur Qwen3-8B, elle a amélioré la précision OOD de +8,6 points par rapport au modèle de base, contre une dégradation de -8,1 points pour la Forward KL.
• Robustesse : Les gains se sont maintenus à travers diverses quantités de contexte privilégié (des préfixes courts aux solutions complètes) et à travers différentes échelles de modèles (1.7B à 32B).
• Ablation : Le retrait de la pondération par chevauchement de Hellinger, du terme proximal de Fisher–Rao ou de la mise à jour par gradient naturel a entraîné des baisses significatives de la performance OOD, confirmant que les trois composants sont nécessaires.
• Faux Consensus : L'analyse a montré que les méthodes standards (comme la Forward KL) présentent des taux élevés de « faux consensus » (des échantillons s'accordant fortement sur de mauvaises réponses) dans les états à haute entropie. La GEOSD a réduit ce taux en dessous de celui du modèle de base, maintenant une auto-cohérence informative.

5. Signification et Revendications

L'article affirme que la dégradation du raisonnement OOD dans l'auto-distillation à contexte privilégié n'est pas simplement le résultat de l'utilisation de signaux peu fiables, mais une conséquence de la manière dont ces signaux déplacent le modèle. Les objectifs standards permettent à la confiance de l'enseignant de forcer l'étudiant dans un mode étroit prématurément, drainant la masse de probabilité des alternatives viables.

La GEOSD démontre qu'une distillation de raisonnement efficace dépend non seulement de quels signaux de l'enseignant faut-il suivre, mais aussi de jusqu'où ces signaux sont autorisés à déplacer l'étudiant. En régulant ce mouvement géométriquement — en atténuant les tractions là où l'étudiant manque de soutien et en ancrant la trajectoire à un état stable récent — la GEOSD atteint un compromis supérieur entre la performance en distribution et la généralisation hors distribution. Les auteurs concluent que le contrôle de la dérive prédictive est essentiel pour l'échelle de l'auto-distillation dans les tâches de raisonnement.

Limites Notées : Les auteurs reconnaissent que leurs conclusions sont actuellement établies uniquement pour le raisonnement mathématique. L'applicabilité de la pondération de chevauchement et du contrôle de la dérive de Fisher–Rao à d'autres domaines (ex: code, raisonnement scientifique) où les états à haute entropie peuvent avoir des structures différentes reste à tester. De plus, la méthode suppose que la politique de base est largement compétente ; pour des modèles de base très faibles, le mécanisme pourrait supprimer des signaux utiles qui nécessitent des changements comportementaux importants.