Résumé Technique : Tous les Désaccords ne sont pas Apprenables : Enseignabilité des Tokens en Distillation On-Policy

Énoncé du Problème

La distillation on-policy (OPD) entraîne un modèle de langage étudiant sur ses propres déroulements en utilisant une supervision au niveau des tokens provenant d'une politique enseignante. Bien que les méthodes OPD sélectives récentes tentent d'améliorer l'efficacité en priorisant des tokens « saillants » — tels que ceux présentant une forte entropie ou un fort désaccord enseignant-étudiant (divergence KL) — ces approches reposent sur l'hypothèse qu'un désaccord élevé équivaut à une grande valeur d'apprentissage.

L'article identifie un défaut critique dans cette hypothèse : le désaccord KL brut est un proxy grossier pour l'apprentissage. Il confond deux scénarios qualitativement différents :

1. Désaccord Apprenable : L'enseignant attribue une masse de probabilité corrective à des tokens déjà présents dans le support prédictif actuel (top-$K$) de l'étudiant.
2. Désaccord Incompatible : L'enseignant attribue une masse à des tokens bien en dehors du support actuel de l'étudiant. Bien que cela crée un écart KL important, l'étudiant ne peut pas facilement absorber cette correction, conduisant à des effets d'apprentissage local faibles ou négatifs.

Le problème central réside dans la distinction entre quels signaux enseignants au niveau des tokens sont réellement apprenables (c'est-à-dire capables d'induire une amélioration locale) et ceux qui sont simplement saillants ou incompatibles.

Méthodologie

1. Diagnostic à Contexte Fixe

Pour isoler la valeur d'apprentissage de tokens spécifiques du bruit des déroulements et des changements de contexte en aval, les auteurs introduisent un diagnostic à contexte fixe.

• Protocole : Les préfixes générés par l'étudiant sont collectés et figés en tant que banque de contextes. Les modèles étudiants initial (avant mise à jour) et entraîné sont re-notés par rapport à la même distribution enseignante sur ces contextes fixes.
• Métrique : Le Gain de Token à Contexte Fixe ($G^{\text{fix}}_{i,t}$) est défini comme la réduction de la divergence KL entre l'enseignant et l'étudiant après l'entraînement, mesurée sur le même préfixe.
• Objectif : Cela permet aux auteurs de corréler les signaux pré-mise à jour (comme le KL brut) avec une amélioration locale réelle, en contrôlant la variance d'échantillonnage.

2. Formalisation de l'Enseignabilité des Tokens

En utilisant ce diagnostic, les auteurs décomposent le désaccord en deux composantes basées sur la compatibilité du support local :

• Support Local ($S^S_t(K)$) : L'ensemble des $K$ meilleurs tokens de l'étudiant.
• Masse de Compatibilité ($C_t$) : La masse de probabilité totale que l'enseignant attribue au support local de l'étudiant $S^S_t(K)$.
• Décomposition :
  • Désaccord Apprenable ($D^L_t$) : Désaccord élevé où la correction de l'enseignant est alignée avec le support de l'étudiant (Forte $C_t$).
  • Désaccord Incompatible ($D^I_t$) : Désaccord élevé où la correction de l'enseignant tombe en dehors du support de l'étudiant (Faible $C_t$).

Le Score d'Enseignabilité est défini comme le produit du désaccord normalisé et de la compatibilité :
$$s^{\text{teach}}_t = \tilde{D}_t \cdot \tilde{C}_t$$
où $\tilde{D}_t$ est le désaccord normalisé et $\tilde{C}_t$ est la masse de compatibilité normalisée.

3. OPD Sensible à l'Enseignabilité (TA-OPD)

Sur la base du constat que $D^L_t$ prédit mieux le gain à contexte fixe que le KL brut, les auteurs proposent TA-OPD, une méthode légère de sélection de positions de tokens :

• Sélection : Pour un budget donné (par exemple, 5 % ou 10 % des tokens), TA-OPD sélectionne les $N$ meilleures positions ayant les scores d'enseignabilité les plus élevés ($s^{\text{teach}}_t$).
• Implémentation : Elle ne nécessite que les probabilités de tokens de l'enseignant et de l'étudiant (spécifiquement les log-probabilités top-$K$) disponibles lors de l'OPD standard. Elle ne requiert pas de modèles de récompense, de vérificateurs ou d'étiquettes supplémentaires.
• Contraste : Contrairement aux méthodes précédentes qui sélectionnent en fonction de l'entropie ou de la divergence brute (TIP), TA-OPD filtre explicitement les corrections « alignées avec le support ».

Contributions Clés

1. Diagnostic à Contexte Fixe : Une méthode pour mesurer la réduction de KL dans le même contexte, reliant les signaux de tokens pré-mise à jour à l'amélioration locale de l'OPD sans bruit de déroulement confondant.
2. Formalisation de l'Enseignabilité des Tokens : La démonstration que le désaccord KL brut confond les signaux apprenables et incompatibles. L'article prouve que le désaccord aligné avec le support est le prédicteur fiable d'un apprentissage utile.
3. Algorithme TA-OPD : Une méthode de sélection pratique et légère qui applique la perte OPD uniquement aux positions à haute enseignabilité, améliorant l'efficacité sans outils externes.
4. Validation Empirique : Des expériences extensives montrant que la sélection basée sur l'enseignabilité produit de meilleures performances en aval que la sélection basée sur la saillance, même avec des budgets de tokens considérablement réduits.

Résultats Expérimentaux

Les auteurs ont évalué TA-OPD sur quatre configurations enseignant-étudiant impliquant les modèles Qwen3 (4B, 8B, 14B) et Qwen2.5 (3B), en utilisant des benchmarks tels que AIME, GPQA-Diamond, HumanEval et MATH-500.

• Performance vs OPD Complet : Dans de nombreuses configurations, TA-OPD a égalé ou surpassé l'OPD sur tous les tokens tout en n'utilisant que 5 % des tokens. Par exemple, dans la configuration Qwen3-4B $\to$ 1.7B, TA-OPD a obtenu un score moyen de 44,89 contre 42,37 pour l'OPD complet, malgré l'utilisation d'un budget de 10 % (et surpassant l'OPD complet sur des métriques spécifiques).
• Comparaison avec les Lignes de Base : TA-OPD a constamment surpassé les lignes de base basées sur l'entropie, la divergence et TIP (Importance des Tokens dans l'OPD).
  • Dans la configuration à backbone croisé (DeepSeek-R1-Distill-Qwen-14B $\to$ Qwen2.5-3B), TA-OPD a amélioré le score moyen à 30,62, surpassant significativement le modèle de base (29,98) et l'OPD complet (28,76).
• Sensibilité au Budget : Les résultats indiquent que la performance de l'OPD est limitée par la qualité des tokens plutôt que purement par le budget. Augmenter le budget de tokens au-delà de 10 % n'a souvent pas produit de gains supplémentaires et a parfois dégradé les performances, renforçant l'idée que toute supervision supplémentaire n'est pas bénéfique.
• Robustesse : Le « écart de gain à fort support » positif (la différence de performance entre les tokens à haute et basse enseignabilité) est resté cohérent à travers différents contextes (mathématiques, GSM8K-COT), échelles d'enseignants (8B, 14B) et tailles de support ($K=8, 16, 32$).

Importance et Revendications

L'article reframe la distillation on-policy sélective non pas comme un problème de sélection de tokens « saillants » ou « incertains », mais comme un problème d'identification de signaux enseignants localement apprenables.

• Revendication Centrale : Les gains de l'OPD dépendent de la compatibilité entre la correction de l'enseignant et le support prédictif actuel de l'étudiant. Une forte divergence KL seule est insuffisante ; l'enseignant doit attribuer une masse à des tokens que l'étudiant considère déjà comme plausibles.
• Impact Pratique : TA-OPD offre un moyen de réduire drastiquement le coût computationnel de la distillation (en filtrant les tokens incompatibles) tout en maintenant ou améliorant les performances du modèle. Cela est réalisé sans la surcharge des modèles de récompense ou des vérificateurs, le rendant applicable aux pipelines d'entraînement standards.
• Limites : Les auteurs notent que leur analyse se concentre sur le raisonnement riche en mathématiques et les modèles de la famille Qwen. Bien que le diagnostic soit agnostique au modèle, une validation plus large sur le dialogue, le code et des backbones non-Qwen est nécessaire pour confirmer la généralisabilité. De plus, le « budget de tokens » fait référence à la densité de supervision, et pas nécessairement à une accélération du temps réel, car la méthode sélectionne des positions plutôt que de réduire le calcul du transformateur.

En conclusion, ce travail suggère qu'une distillation efficace nécessite une stratégie de supervision consciente de la compatibilité, où l'étudiant apprend le mieux à partir de signaux enseignants qui sont à la fois correctifs et localement accessibles.