Résumé Technique : Démasquer la Distillation On-Policy

Énoncé du Problème

La distillation on-policy (OPD) est devenue une technique fondamentale de post-entraînement pour les modèles de raisonnement, offrant une supervision dense, token par token, pour compléter l'apprentissage par renforcement (RL) à récompense sparse. Dans ce paradigme, un modèle « enseignant » (soit un modèle externe plus grand, soit le modèle étudiant lui-même disposant d'un contexte privilégié) guide la génération de tokens par l'étudiant. Cependant, les praticiens manquent de directives fondées sur des principes concernant les choix de conception critiques :

• Sélection de l'enseignant : L'enseignant doit-il être un modèle externe plus grand ou l'étudiant lui-même (auto-distillation) ?
• Conception du contexte : Quel contexte spécifique (par exemple, traces complètes de solutions, solutions résumées, ou exemples de réponses incorrectes) doit servir de signal de supervision ?
• Granularité : La configuration optimale varie-t-elle selon la tâche, la capacité de l'étudiant, ou même les tokens individuels ?

L'évaluation actuelle repose sur des métriques de performance agrégées issues de runs d'entraînement coûteux, ce qui obscurcit les dynamiques au niveau des tokens individuels. Il reste unclear dans quelles conditions le signal de l'enseignant est bénéfique, neutre ou préjudiciable.

Méthodologie

Les auteurs introduisent un cadre diagnostique sans entraînement opérant à la plus haute résolution : par token, par question et par configuration d'enseignant. Ce cadre se compose de trois composants principaux :

1. Le Gradient Idéal par Nœud

Pour évaluer la guidance de l'enseignant, les auteurs dérivent un gradient « idéal » maximisant la probabilité pour l'étudiant d'atteindre une réponse correcte.

• Arbre de Génération : Ils décomposent le processus de génération en un arbre où les nœuds représentent les positions des tokens et les arêtes représentent les choix de tokens suivants.
• Probabilité de Succès Empirique ($\hat{P}^k_{succ}$) : En échantillonnant des déroulements (rollouts) à partir de l'étudiant, ils estiment la probabilité d'atteindre une réponse correcte étant donné la sélection du token $k$ au nœud $u$.
• Objectif Idéal : L'objectif local idéal consiste à maximiser le taux de succès attendu : $L_{ideal}(u) = \sum_k P^k_\theta P^k_{succ}$.
• Gradient Idéal : Le gradient de cet objectif par rapport aux logits de l'étudiant est dérivé comme $g^{ideal}_j = P^j_\theta (P^j_{succ} - \bar{P}_{succ})$. Cette direction pousse l'étudiant vers des tokens qui mènent au succès plus souvent que la moyenne.
• Lien avec Dr. GRPO : Les auteurs montrent que Dr. GRPO (une variante de GRPO sans normalisation de longueur) retrouve ce gradient idéal en espérance, validant ainsi son statut d'estimateur non biaisé du signal basé sur la récompense.

2. Analyse du Gradient de Distillation

Le cadre analyse les gradients produits par divers algorithmes de distillation (GKD, estimateurs sur un seul échantillon, MiniLLM). Il démontre que ces méthodes partagent une structure locale commune : $\frac{\partial L}{\partial z_j} = P^j_\theta (f_j - \bar{f})$, où $f_j$ est dérivé de la distribution de l'enseignant (par exemple, rapports logarithmiques) ou des récompenses au niveau de la trajectoire.

3. Score d'Alignement des Gradients

Pour quantifier dans quelle mesure une configuration d'enseignant spécifique approxime le signal idéal, les auteurs définissent le Score d'Alignement des Gradients :
$$\text{Align}(u) = \cos(g^{ideal}_u, g^D_u)$$
Cette similarité cosinus entre le gradient idéal et le gradient de distillation à chaque token mesure si la guidance de l'enseignant pousse vers le succès (alignement positif), l'échec (alignement négatif), ou est orthogonale/bruyante (alignement proche de zéro).

4. Estimation Évolutible via Déroulements Ciblés

Le calcul de $\hat{P}^k_{succ}$ pour tous les tokens dans de longues chaînes de raisonnement est computationnellement prohibitif. Les auteurs proposent un algorithme de déroulement ciblé :

• Fenêtres de Profondeur Exponentielle : La séquence de génération est partitionnée en fenêtres de taille croissante.
• Priorisation : Dans chaque fenêtre, les déroulements sont alloués aux tokens où l'enseignant diverge le plus fortement de l'étudiant (magnitude élevée du gradient GKD ou différence de probabilité).
• Efficacité : Cette approche évolue avec un budget choisi par l'utilisateur plutôt qu'avec la longueur de la séquence, la rendant réalisable pour des traces allant d'environ 200 tokens (BoolQ) à environ 30 000 tokens (AIME).

Résultats Clés

Le cadre a été appliqué aux modèles Qwen3-0.6B et Qwen3-1.7B sur les jeux de données BoolQ, MMLU et AIME 2025 avec 8 configurations d'enseignant différentes.

1. La Distillation Aide Davantage sur les Chemins Incorrects

La découverte la plus cohérente est que la guidance par distillation présente un alignement significativement plus élevé sur les déroulements incorrects que sur les corrects.

• Lorsque l'étudiant est sur un chemin menant à l'échec, l'enseignant pousse de manière fiable vers des tokens menant au succès.
• Lorsque l'étudiant est déjà sur un chemin correct, le signal de l'enseignant devient bruyant et faiblement aligné avec l'idéal, car l'étudiant performe déjà bien.
• Ce schéma se maintient à travers tous les paramètres, les échelles de modèles et les jeux de données.

2. Interaction entre la Conception du Contexte et la Capacité de l'Étudiant

Le contexte de distillation optimal n'est pas universel ; il dépend conjointement de la capacité de l'étudiant et de la difficulté de la tâche :

• Petits Modèles (0.6B) : Bénéficient le plus de l'auto-distillation avec des démonstrations correctes uniquement, dans leur propre style de raisonnement. Ils peinent à décoder les schémas de raisonnement d'enseignants externes beaucoup plus grands. Des contextes résumés peuvent nuire aux performances si l'étudiant a besoin de traces complètes étape par étape.
• Modèles Plus Grands (1.7B) : Peuvent exploiter efficacement des enseignants externes (par exemple, Qwen3-8B) et des contextes résumés, ce qui double presque l'alignement par rapport aux traces brutes.
• Démonstrations Incorrectes : L'inclusion d'exemples incorrects (Self-1C1W) nuit généralement aux performances sur les tâches de raisonnement court (BoolQ, MMLU) en introduisant du bruit. Cependant, sur des problèmes mathématiques difficiles (AIME), les démonstrations incorrectes fournissent des signaux contrastifs utiles, devenant la meilleure configuration d'enseignant pour les questions difficiles.

3. Aucune Recette Universelle

Il n'existe ni enseignant unique ni configuration de contexte qui maximise l'alignement sur toutes les tâches.

• Dépendance à la Tâche : Les enseignants externes surpassent l'auto-distillation pour le modèle 1.7B sur BoolQ, mais pas sur MMLU.
• Dépendance à la Difficulté : Sur les mathématiques difficiles, l'inclusion d'une démonstration incorrecte améliore l'alignement, alors qu'elle le dégrade sur des tâches plus simples.
• Limites Prédictives : Bien que la divergence entre les distributions de l'étudiant et de l'enseignant (KL, JS) corrèle positivement avec l'alignement, la corrélation est faible ($|\rho| < 0.05$). La divergence peut servir de filtre peu coûteux de condition nécessaire, mais n'est pas un prédicteur fiable de savoir si un désaccord spécifique est utile.

Signification et Revendications

L'article revendique fournir la première méthodologie rigoureuse et sans entraînement pour évaluer la qualité de la guidance de l'enseignant au niveau du token. Ses contributions principales sont :

1. Cadre Diagnostique : Un pipeline évolutif pour calculer l'alignement des gradients hors ligne, révélant que les signaux des enseignants sont un mélange hétérogène de contributions bénéfiques, neutres et préjudiciables.
2. Insight Mécanistique : La découverte que la distillation est plus efficace sur les déroulements échoués suggère que les pipelines d'entraînement actuels, qui traitent tous les tokens de manière égale, peuvent gaspiller le budget de gradient sur des tokens où l'étudiant est déjà correct.
3. Nuance Contextuelle : La démonstration que la « compréhensibilité » est un facteur clé — les signaux ne sont utiles que si l'étudiant peut les décoder — explique pourquoi les petits modèles échouent à apprendre d'enseignants externes beaucoup plus grands et pourquoi la conception du contexte (résumé vs traces brutes) doit être adaptée à la capacité du modèle.
4. Directions Futures : Les résultats motivent des analyses diagnostiques par tâche et par token et suggèrent des algorithmes d'entraînement futurs qui pourraient pondérer la perte de distillation en fonction de l'alignement (par exemple, en mettant l'accent sur les trajectoires échouées) ou utiliser la divergence comme mécanisme de contrôle.

Les auteurs concluent que si la distillation on-policy est un outil puissant, son efficacité est hautement conditionnelle, et la configuration « optimale » évolue dynamiquement en fonction de l'état de l'étudiant, de la difficulté de la tâche et du token spécifique en cours de génération.