Escaping the KL Agreement Trap in On-Policy Distillation

Résumé Technique : Échapper au piège de l'accord KL dans la distillation on-policy

1. Définition du problème

La distillation on-policy (OPD) est devenue une méthode puissante pour transférer les capacités de grands modèles enseignants vers des modèles étudiants plus petits en fournissant une supervision dense, au niveau du jeton (token), sur les déploiments générés par l'étudiant. Cependant, cet article identifie un mode de défaillance critique appelé le piège de l'accord de faible KL (low-KL agreement trap).

Dans une OPD standard, une faible divergence de Kullback-Leibler (KL) inverse entre l'étudiant et l'enseignant est généralement interprétée comme un signe de progrès ou d'alignement. Les auteurs observent qu'au fur et à mesure de la progression d'un déploiement de l'étudiant, celui-ci peut dériver vers un « préfixe corrompu » (par exemple, une erreur de raisonnement précoce ou un schéma de dégénérescence répétitif). Conditionné par cet état dégradé, le modèle enseignant peut produire des continuations qui sont localement plausibles mais non informatives, entraînant une faible divergence KL. Cela crée un accord dégénéré : l'enseignant et l'étudiant sont d'accord, mais le signal est faible car l'enseignant se contente de suivre l'erreur de l'étudiant plutôt que de fournir une guidance corrective.

L'article soutient que les méthodes d'OPD existantes traitent l'intégralité du déploiement comme une supervision utile. Par conséquent, elles gaspillent de la puissance de calcul sur les jetons suivant l'entrée dans le piège, là où le signal de supervision est dilué ou nuisible, détournant potentiellement le processus d'optimisation.

2. Méthodologie : KAT (KL Agreement Trap Termination)

Pour remédier à cela, les auteurs proposent KAT, une règle de terminaison de déploiement en ligne conçue pour détecter et filtrer ces régions d'accord dégénéré sans modifier l'objectif d'OPD sous-jacent ni nécessiter de modèles auxiliaires.

Mécanisme central :
KAT réutilise la divergence KL inverse ($d_t$) déjà calculée pendant le processus d'OPD. Au lieu de seuiller un seul jeton KL, KAT surveille une statistique de fenêtre glissante ($z_t$) de la moyenne de la KL sur une fenêtre de $W$ jetons.

Seuillage adaptatif :
Reconnaissant que l'échelle de la divergence KL change tout au long de l'entraînement, KAT emploie un seuil dynamique basé sur les données ($s_k$) :

• Phase de chauffe (Warmup) : Pendant les $K$ premières étapes de l'optimiseur, KAT collecte les valeurs minimales de la fenêtre glissante de la KL ($m(\hat{y})$) des déploiements terminés dans un tampon FIFO.
• Calibration : Après la chauffe, le seuil d'arrêt est fixé au quantile $(1-\eta)$ des minima récents du tampon. Cela garantit que la règle s'adapte à l'évolution de la politique de l'étudiant et à l'échelle de la KL.

Logique de terminaison :
Un déploiement est interrompu en ligne lorsque la statistique de la fenêtre glissante de la KL reste inférieure au seuil adaptatif pendant $T$ fenêtres consécutives. Crucialement, KAT inclut une longueur d'exemption ($L_0$) pour éviter une terminaison prématurée durant la phase initiale d'un déploiement où la KL chute naturellement. Une fois le déclenchement confirmé, le déploiement est tronqué au début de la première fenêtre de faible KL, et aucun jeton supplémentaire n'est généré ou utilisé pour la supervision.

3. Contributions clés

L'article apporte trois contributions principales :

1. Identification des pièges d'accord de faible KL : Les auteurs caractérisent un mode de défaillance spécifique dans l'OPD où une faible KL persistante indique un « accord dégénéré » sur des préfixes corrompus plutôt qu'un alignement bénéfique.
2. Analyse de la qualité de la supervision : À travers une analyse de la géométrie du gradient (projection des gradients sur les sous-espaces de mise à jour principaux) et une visualisation du vocabulaire, l'article démontre que :
   • Les gradients dans les phases d'« accord » et de « post-accord » s'alignent mal avec les directions principales de mise à jour des paramètres.
   • Le vocabulaire de l'enseignant passe de jetons liés au raisonnement à des fragments de surface locaux ou des continuations génériques dans ces régions.
   • La qualité de la supervision ne se rétablit pas même si la KL augmente légèrement après le piège.
3. Proposition de KAT : Une règle de terminaison adaptative et en ligne qui filtre les suffixes non informatifs. Elle améliore l'efficacité et la précision de l'entraînement en concentrant les mises à jour sur les préfixes informatifs pré-piège, ne nécessitant qu'une gestion de données en $O(1)$ par jeton.

4. Résultats expérimentaux

Les auteurs ont évalué KAT sur quatre benchmarks de raisonnement mathématique (AMC, MATH500, MinervaMath, AIME24) en utilisant deux échelles d'étudiants (Qwen3-0.7B/1.6B et Qwen3-4B) avec un enseignant Qwen3-8B.

Gains de performance :

• Précision : KAT a amélioré la précision moyenne avg@k de 2,66 % et pass@k de 3,43 % par rapport à l'OPD standard.
• Efficacité : La longueur moyenne des déploiements a été réduite de 59,73 % (passant d'environ 1480 à environ 596 jetons).
• Coût d'entraînement : Sur l'étudiant 4B, KAT a réduit le temps d'entraînement de 2,4x (26,28h à 11,00h) et l'utilisation de calcul de 59,4 % par rapport à l'OPD standard.

Analyse comparative :

• vs Préfixe fixe : KAT a surpassé la troncature à longueur fixe (ex: 1024 ou 2048 jetons), démontrant que la qualité du point de troncature importe plus que la longueur elle-même. Les préfixes fixes coupent souvent un raisonnement utile ou conservent des suffixes inutiles.
• vs Terminaison aléatoire : KAT a nettement mieux performé que la troncature aléatoire, confirmant que le signal de faible KL est un indicateur spécifique de supervision faible plutôt qu'une simple métrique de longueur.
• Ablation : Supprimer les étapes d'« accord » et de « post-accord » de l'entraînement (comme le fait KAT) a produit une précision plus élevée que de les inclure, validant que ces régions fournissent des signaux faibles ou nuisibles.

5. Signification et revendications

L'article affirme que KAT offre une approche plus efficace et plus performante de la distillation on-policy en repensant fondamentalement la manière dont les signaux de supervision sont utilisés. En identifiant que une faible KL n'est pas toujours bénigne, la méthode empêche la dilution des signaux d'apprentissage par des jetons non informatifs.

Les auteurs soulignent que KAT atteint ces gains sans changer l'objectif d'OPD ni introduire de nouvelles fonctions de perte ou de modèles auxiliaires. Il tire parti du calcul existant de la KL inverse pour prendre des décisions en ligne et adaptatives. L'article conclut que KAT rend l'OPD à la fois plus efficace (meilleure précision) et plus performante (calcul réduit), particulièrement dans les tâches de raisonnement mathématique où les chemins de raisonnement peuvent facilement dégénérer en états irrécupérables.

Limites notées :
Les auteurs notent modestement que ce travail se concentre sur le raisonnement mathématique. Le comportement de la règle de terminaison dans des domaines plus larges comme la génération de code, le suivi d'instructions ouvert ou les interactions multi-tours reste à étudier. De plus, bien que les hyperparamètres soient intuitifs, une orientation plus systématique pour leur réglage dans divers contextes pourrait améliorer l'utilisabilité.