Résumé technique : Distillation sur politique avec sélection de déroulements enseignants par le meilleur parmi N (BRTS)

1. Énoncé du problème

La distillation sur politique (OPD) est apparue comme une méthode d'entraînement postérieur économe en données pour les grands modèles de langage (LLM), où un modèle étudiant est supervisé sur ses propres trajectoires échantillonnées en utilisant les log-probabilités par token de l'enseignant. Bien que l'OPD atténue le biais d'exposition en alignant la supervision sur les états que l'étudiant visite réellement, les implémentations standard souffrent de signaux de supervision à forte variance.

La limitation centrale identifiée est double :

1. Contextes étudiants bruyants : La supervision est calculée sur des préfixes générés par un étudiant imparfait, qui peut dériver vers des états de raisonnement bruyants où le feedback local de l'enseignant est moins informatif.
2. Déroulements enseignants stochastiques : L'OPD standard repose généralement sur un seul déroulement enseignant par invite. Puisque l'enseignant est stochastique, une seule trajectoire échantillonnée peut être incorrecte, peu informative ou mal alignée avec le comportement de raisonnement actuel de l'étudiant. Cela se traduit par un guidage peu fiable, en particulier sur des invites difficiles où les sorties échantillonnées de l'enseignant varient considérablement en termes de justesse et de style.

Par conséquent, l'OPD standard risque de propager des signaux enseignants erronés ou de ne pas exposer l'étudiant à des chemins de raisonnement complets et corrects, limitant ainsi son efficacité sur les benchmarks de raisonnement complexes.

2. Méthodologie : Cadre BRTS

Les auteurs proposent BRTS (Best-of-N Rollout Teacher Selection), un cadre conçu pour sélectionner des trajectoires enseignants fiables avant leur utilisation pour la supervision dans la boucle OPD. BRTS enrichit la distillation standard de contexte étudiant avec une branche de supervision de contexte enseignant construite à partir d'une trajectoire enseignante sélectionnée et de haute qualité.

2.1 Sélection de trajectoires consciente de la justesse et de l'alignement

Pour chaque invite, BRTS échantillonne un pool de $N$ déroulements enseignants sans condition ($\{y_{T,i}\}_{i=1}^N$). Le processus de sélection suit une hiérarchie de priorités :

1. Priorité 1 : Justesse. Le système filtre les déroulements où la réponse finale extraite correspond à la vérité terrain ($y^*$).
2. Priorité 2 : Alignement étudiant. Si plusieurs déroulements corrects existent, BRTS sélectionne celui présentant la plus forte superposition au niveau des tokens avec l'ensemble des $K$ meilleurs candidats de l'étudiant. Cela garantit que la trajectoire sélectionnée est compatible avec la distribution de raisonnement actuelle de l'étudiant.
3. Récupération guidée par la vérité terrain (Niveau 2). Si tous les échantillons enseignants sans condition échouent (c'est-à-dire qu'aucun n'est correct), BRTS invoque une étape de récupération. Il construit une invite modifiée qui injecte silencieusement la réponse de vérité terrain comme signal de validation tout en instruisant l'enseignant à produire une dérivation naturelle. Un nouveau déroulement est échantillonné ; s'il est correct, il est sélectionné. Si cela échoue également, le système revient au déroulement sans condition présentant la plus forte superposition avec l'étudiant.

2.2 Supervision de contexte enseignant

L'objectif d'entraînement combine deux composantes de perte :

• Perte de contexte étudiant ($L_{stu-ctx}$) : L'objectif OPD standard, minimisant la divergence KL inverse entre les distributions de l'étudiant et de l'enseignant conditionnées par le préfixe généré par l'étudiant. Cela corrige l'étudiant sur les états qu'il visite réellement.
• Perte de contexte enseignant ($L_{tea-ctx}$) : Une perte auxiliaire calculée sur la trajectoire enseignante sélectionnée ($y'$). Ici, les distributions de l'étudiant et de l'enseignant sont toutes deux conditionnées par le préfixe de l'enseignant. Cela expose l'étudiant à un chemin de raisonnement cohérent et correct qui peut se situer en dehors de la région de haute probabilité actuelle de l'étudiant.

La perte totale est définie comme suit :
$$L_{total} = L_{stu-ctx} + \lambda L_{tea-ctx}$$
où $\lambda$ est un coefficient de pondération (fixé empiriquement à 10).

2.3 Direction des $K$ meilleurs candidats

Pour implémenter ces pertes efficacement, BRTS utilise une agrégation des $K$ meilleurs :

• Pour la branche de contexte étudiant, l'ensemble des tokens candidats est dérivé des $K$ meilleurs tokens de l'étudiant sous le préfixe étudiant.
• Pour la branche de contexte enseignant, l'ensemble des candidats est dérivé des $K$ meilleurs tokens de l'enseignant sous le préfixe enseignant sélectionné. Cela permet à la branche auxiliaire de superviser les tokens préférés par l'enseignant que l'étudiant pourrait ne pas considérer actuellement, tout en maintenant un ensemble de candidats compact.

3. Contributions clés

Le papier résume ses contributions comme suit :

1. Identification de la supervision à forte variance : Les auteurs soulignent que la supervision de contexte étudiant standard dans l'OPD est limitée par des préfixes étudiants bruyants et des déroulements enseignants stochastiques, rendant le signal enseignant peu fiable sur les invites difficiles.
2. Introduction de BRTS : Un cadre léger pour sélectionner et distiller des trajectoires enseignants fiables. Il emploie une règle de sélection « justesse en premier, alignement en second » et inclut une solution de repli guidée par la vérité terrain pour les échantillons sans condition échoués.
3. Amélioration empirique : BRTS surpasse nettement l'OPD standard sur les benchmarks de raisonnement complexes. Le papier fournit des analyses diagnostiques démontrant comment la qualité des candidats enseignants, la couverture de récupération et les choix de modèles de base influencent les performances.

4. Résultats expérimentaux

Des expériences ont été menées sur AIME 2024, AIME 2025 et AMC 2023 en utilisant un modèle étudiant de 1,5B (DeepSeek-1.5B) et diverses configurations d'enseignants (JustRL-1.5B et DeepSeek-R1-Distill-Qwen-7B).

• Gains de performance : BRTS améliore constamment la ligne de base étudiant seul. Les gains sont les plus marqués sur le jeu de données AIME 2025 plus difficile. Par exemple, sur AIME 2024, l'utilisation de quatre candidats enseignants a atteint une précision moyenne de 0,4000 contre 0,3917 pour la ligne de base.
• Impact de la récupération de niveau 2 : L'étape de récupération guidée par la vérité terrain (Niveau 2) apporte des améliorations substantielles sur les invites difficiles où l'échantillonnage sans condition échoue. Sur AIME 2025, l'ajout du Niveau 2 a augmenté la précision du vote majoritaire de 0,2748 (Niveau 1 uniquement) à 0,3133.
• Crêtes d'apprentissage précoce : BRTS atteint un pic de performance plus élevé lors des premières étapes d'entraînement par rapport aux lignes de base, suggérant une convergence plus rapide et un meilleur compromis calcul-performance.
• Robustesse : La méthode reste efficace lors du remplacement des modèles de base enseignants (par exemple, utiliser un enseignant de 7B pour un étudiant de 1,5B) et n'est pas fragile face à de légères perturbations d'invite, qui aident en réalité à décorréler les échantillons enseignants.
• Coût computationnel : La surcharge est modeste ; augmenter le pool de candidats enseignants de 1 à 2 ajoute environ 59 secondes par étape d'entraînement sur 8×B200 GPU, ce qui est considéré comme acceptable compte tenu des gains de performance.

5. Importance et affirmations

Le papier affirme que BRTS adresse un goulot d'étranglement critique dans la distillation sur politique : l'irrémédiable fiabilité des signaux enseignants stochastiques sur les tâches difficiles. En structurant la supervision enseignante pour qu'elle soit à la fois consciente de la justesse (filtrage des erreurs) et consciente de l'alignement (correspondance avec la distribution de l'étudiant), BRTS transforme les déroulements enseignants stochastiques en une supervision structurée et de haute qualité.

Les auteurs positionnent BRTS non pas comme un remplacement de l'OPD, mais comme un enrichissement nécessaire permettant aux étudiants d'apprendre à partir de chemins de raisonnement complets et fiables, sans encourir l'oubli catastrophique souvent associé au fine-tuning supervisé standard ni la forte variance de l'OPD standard. La méthode est présentée comme une extension évolutive et légère qui améliore les capacités de raisonnement, en particulier dans les domaines où des réponses de vérité terrain sont disponibles pour vérification.