Trust-Region Behavior Blending for On-Policy Distillation

Résumé Technique : Mélange de Comportement par Région de Confiance pour la Distillation On-Policy

Énoncé du Problème

La distillation on-policy (OPD) traite le décalage de préfixe inhérent à la distillation hors-ligne en entraînant un modèle étudiant sur des préfixes échantillonnés à partir de sa propre politique actuelle tout en s'alignant sur un enseignant plus fort. Cependant, cette approche introduit une fragilité critique lors des premières étapes de l'entraînement : les modèles étudiants faibles génèrent souvent des préfixes de faible qualité. Lorsque la distribution des trajectoires de l'étudiant est médiocre, la supervision de l'enseignant est appliquée à des préfixes de faible qualité ou « inutilisables », limitant l'efficacité du signal d'apprentissage.

Les tentatives précédentes pour stabiliser ce processus consistent souvent à déplacer la collecte hors-ligne (par exemple, en injectant des jetons de l'enseignant ou en modifiant les distributions cibles). Or, cela peut perturber la distribution d'entraînement cible. À l'inverse, les déploiements (rollouts) purement étudiants préservent la distribution mais souffrent d'une qualité initiale médiocre. Le défi central est d'améliorer la qualité des premiers déploiements sans déplacer de façon permanente le processus de collecte hors de la politique (off-policy) ni altérer l'objectif fondamental de distillation par préfixe.

Méthodologie : Mélange de Comportement par Région de Confiance (TRB)

Les auteurs proposent le Mélange de Comportement par Région de Confiance (TRB), une méthode de chauffe (warmup) conçue pour contrôler la politique de comportement lors de la phase initiale de collecte de déploiements. Le TRB repose sur le principe de trouver une politique de comportement aussi proche que possible de l'enseignant, tout en étant strictement contrainte dans une « région de confiance » centrée sur la politique actuelle de l'étudiant.

1. Optimisation par Préfixe
À chaque préfixe de génération $h$, le TRB définit une politique de comportement $\mu^*(\cdot | h)$ en résolvant un problème d'optimisation contraint. L'objectif est de minimiser la divergence de Kullback-Leibler (KL) par rapport à la politique de l'enseignant $\pi_T$, sous la contrainte d'une divergence KL par rapport à la politique de l'étudiant $\pi_S$ :

$$\mu^*(\cdot | h) = \arg \min_{\mu} D_{KL}(\mu \parallel \pi_T) \quad \text{s.t.} \quad D_{KL}(\mu \parallel \pi_S) \le \epsilon$$

Ici, $\epsilon$ représente la déviation locale autorisée par rapport à l'étudiant. Cette formulation garantit que la distribution d'échantillonnage est attirée vers les jetons soutenus par l'enseignant, mais reste bornée dans une région de confiance centrée sur l'étudiant.

2. Solution à Forme Fermée
L'optimisation produit une solution à forme fermée où la politique de comportement est une interpolation géométrique entre les politiques de l'étudiant et de l'enseignant :
$$\mu_\beta(a | h) = \frac{\pi_S(a | h)^{1-\beta} \pi_T(a | h)^\beta}{Z_\beta(h)}$$
Le coefficient $\beta^*(h)$ est la plus grande valeur réalisable dans $[0, 1]$ telle que la contrainte KL soit satisfaite. Si l'enseignant est déjà dans la région de confiance, $\mu^*$ devient l'enseignant ; si $\epsilon=0$, elle revient à l'étudiant.

3. Chauffe avec Recul (Annealed Warmup)
Le TRB est appliqué uniquement durant une phase de chauffe initiale. Le budget KL $\epsilon$ est réduit linéairement vers zéro sur un horizon $K$ :
$$\epsilon_k = \epsilon_0 \left(1 - \frac{k}{K}\right)$$
Cela permet une forte guidance de l'enseignant au début de l'entraînement pour améliorer la qualité des préfixes, tout en revenant progressivement à un échantillonnage de déploiement pur de l'étudiant à mesure que celui-ci s'améliore. Crucialement, la perte de distillation reverse-KL par préfixe reste inchangée ; seule la politique de collecte de données est modifiée.

Principales Contributions

• Une Nouvelle Stratégie de Chauffe : Le TRB introduit une méthode pour contrôler le comportement des déploiements précoces en optimisant une politique guidée par l'enseignant sous une contrainte explicite de KL centrée sur l'étudiant, plutôt qu'en injectant simplement des jetons ou en modifiant la fonction de perte.
• Préservation des Propriétés On-Policy : Contrairement aux méthodes qui mélangent de façon persistante les politiques de l'enseignant et de l'étudiant ou modifient la distribution cible, le TRB garantit que l'entraînement revient finalement à des déploiements purement on-policy, préservant ainsi la distribution cible.
• Formulation Théorique : L'article fournit une solution à forme fermée pour la politique de comportement qui équilibre la guidance de l'enseignant et la stabilité de l'étudiant, justifiée par la monotonie de la divergence KL par rapport au coefficient d'interpolation.

Résultats Expérimentaux

Les auteurs ont évalué le TRB à travers deux contextes de distillation de raisonnement mathématique :

1. Qwen3-1.7B-Base distillé à partir de Q�wen3-8B.
2. Qwen3-0.6B-Base distillé à partir de Qwen3-4B.

Performance :

• Le TRB a obtenu la meilleure performance moyenne parmi toutes les méthodes comparées dans les deux contextes, surpassant l'OPD vanilla, le mélange à $\epsilon$ fixe, Veto, SKD (injection intercalée de l'enseignant), la chauffe par température et la chauffe par SFT.
• Dans le contexte Qwen3-1.7B, le TRB a atteint un pass@1 moyen de 33,2 %, contre 32,3 % pour l'OPD vanilla.
• Dans le contexte Qwen3-0.6B, le TRB a atteint 44,4 %, contre 44,0 % pour l'OPD vanilla.

Analyse :

• Dynamique de Début d'Entraînement : L'analyse des trajectoires (Figure 2) montre que bien que certaines méthodes (comme SKD) progressent plus rapidement initialement, elles ne parviennent pas nécessairement à maintenir la performance finale la plus élevée. Le TRB offre un « démarrage rapide » qui conduit à des résultats finaux supérieurs.
• Entropie et Qualité des Préfixes : Pendant la phase de chauffe, le TRB réduit l'entropie moyenne des jetons de l'enseignant sur les préfixes visités, indiquant une guidance de meilleure qualité. Les sondes de l'étape 0 (Figure 4) démontrent que les préfixes générés sous TRB présentent des taux de réussite plus élevés lorsqu'ils sont poursuivis soit par l'enseignant, soit par l'étudiant, comparativement aux préfixes de l'OPD vanilla, suggérant que le TRB déplace avec succès la distribution d'état initiale vers des trajectoires plus apprenables.

Signification et Revendications

L'article affirme que le TRB traite efficacement la fragilité de la distillation on-policy précoce en améliorant la qualité de la distribution de données initiale sans dévier de façon permanente de l'objectif on-policy.

• Portée Modeste : Les auteurs précisent explicitement que leur étude est limitée à deux contextes de raisonnement mathématique avec des paires Qwen3-Base et un protocole d'évaluation basé sur la correction. Ils ne prétendent pas que les calendriers de chauffe spécifiques se transfèrent inchangés à d'autres domaines ou écarts enseignant-étudiant.
• Compromis de Coût : L'article reconnaît que le TRB augmente les coûts d'entraînement pendant la phase de chauffe en raison de l'exigence de décodage en ligne de l'enseignant et de la co-résidence étudiant-enseignant. Cependant, ils notent que ces coûts sont temporaires, car l'entraînement revient au profil d'exécution standard et plus efficace de l'OPD après l'horizon de chauffe.
• Intuition Fondamentale : Les résultats suggèrent que le comportement hors-ligne guidé par l'enseignant est particulièrement bénéfique lorsque les préfixes visités par l'étudiant sont significativement désalignés avec l'enseignant. Une fois ce décalage résolu, la guidance hors-ligne continue peut offrir des rendements décroissants, ce qui soutient la stratégie d'atténuation de la politique de comportement vers l'étudiant.