Résumé technique : Repenser la supervision des récompenses : l'auto-distillation conditionnée par des rubriques

Énoncé du problème

Le post-entraînement des modèles de langage de raisonnement (LLM) repose actuellement fortement sur deux paradigmes, qui souffrent tous deux de limitations significatives dans la fourniture d'un feedback fin :

1. Apprentissage par renforcement avec récompenses vérifiables (RLVR) : Les méthodes comme le Group Relative Policy Optimization (GRPO) compressent généralement le feedback évaluatif en un signal scalaire clairsemé (par exemple, une correction binaire ou un score unique). Bien qu'efficace pour les tâches vérifiables, cette approche occulte le pourquoi une trajectoire a réussi ou échoué, créant un « goulot d'étranglement de l'attribution de crédit » où le modèle ne peut pas distinguer les étapes de raisonnement intermédiaires correctes des incorrectes.
2. Auto-distillation on-policy (OPSD) : Ces méthodes remplacent les récompenses clairsemées par un guidage dense au niveau du jeton (token) provenant d'un modèle enseignant. Cependant, les approches existantes conditionnent l'enseignant sur une trajectoire de référence privilégiée unique (une solution d'or spécifique). Cela lie la supervision à un chemin de raisonnement particulier, risquant de sur-spécifier la manière dont une réponse doit être produite et ignorant l'espace des autres dérivations valides. De plus, si l'étudiant s'écarte légèrement de ce chemin unique, l'enseignant peut encourager une révision globale plutôt qu'une correction localisée.

De plus, bien que les rubriques (critères d'évaluation structurés) offrent une représentation de la qualité plus riche et plus interprétable qu'un score scalaire, les travaux antérieurs ont largement écarté cette structure lors de l'optimisation en agrégeant les jugements de la rubrique en une récompense scalaire unique.

Méthodologie : Auto-distillation conditionnée par des rubriques (RCSD)

Les auteurs proposent l'Auto-distillation conditionnée par des rubriques (RCSD), un cadre qui traite les rubriques comme une supervision structurée privilégiée du côté de l'enseignant pour l'auto-distillation on-policy. Au lieu de réduire le feedback de la rubrique en un scalaire, la RCSD conditionne le modèle enseignant sur les informations de la rubrique au niveau des critères afin de fournir un guidage dense au niveau du jeton sur les propres trajectoires échantillonnées de l'étudiant.

Le cadre fonctionne via un pipeline en deux étapes :

Étape I : Apprentissage d'un générateur de rubriques

Pour remédier au coût élevé de l'obtention de rubriques spécifiques à chaque instance, les auteurs entraînent un générateur de rubriques dédié.

• Politique Enseignante (Teacher Policy) : Reçoit la question ($x$) et une réponse de référence ($y^\star$) comme entrées privilégiées pour inférer les critères d'évaluation appropriés.
• Politique Étudiante (Student Policy) : Observe uniquement la question ($x$) et apprend à générer une rubrique ($r$) qui approxime les critères de l'enseignant.
• Objectif : L'étudiant est entraîné via une distillation on-policy pour correspondre à la distribution de jetons de l'enseignant le long de la propre trajectoire de rubrique échantillonnée par l'étudiant. Cela permet d'amortir la génération de critères spécifiques à l'instance en un modèle réutilisable.

Étape II : Raisonnement conditionné par la rubrique

Le générateur de rubriques entraîné est utilisé pour guider l'entraînement du modèle de raisonnement.

• Politique Enseignante : Reçoit la question ($x$) et la rubrique générée ($r$) comme contexte privilégié.
• Politique Étudiante : Observe uniquement la question ($x$) et génère une trajectoire de raisonnement ($\hat{y}$).
• Objectif : L'étudiant minimise la divergence de Kullback-Leibler (KL) entre sa distribution et la distribution de l'enseignant, où l'enseignant est conditionné par la rubrique.
  $$\mathcal{L}_{reason} = \mathbb{E}_{\hat{y} \sim p_S} \left[ \sum_{t=1}^{|\hat{y}|} D_{KL} \left( p_T(\cdot \mid x, r, \hat{y}_{<t}) \parallel p_S(\cdot \mid x, \hat{y}_{<t}) \right) \right]$$
• Conception clé : En se conditionnant sur la rubrique plutôt que sur une seule réponse de référence, l'enseignant fournit un guidage conscient des critères. Cela permet au signal d'optimisation de distinguer les différentes dimensions de la correction partielle (par exemple, une étape correcte suivie d'une erreur locale) plutôt que de les compresser dans un score unique ou un chemin de référence unique. Les auteurs utilisent la divergence KL directe (Forward KL) pour encourager l'étudiant à couvrir le support de la distribution de l'enseignant conditionnée par la rubrique, préservant ainsi la diversité à travers les chemins de raisonnement valides.

Contributions clés

1. Nouvelle interface de supervision : La RCSD reformule les rubriques non pas simplement comme des outils de notation, mais comme des interfaces de supervision structurées qui façonnent l'apprentissage au niveau du jeton pendant l'optimisation.
2. Pipeline en deux étapes : Les auteurs introduisent une méthode évolutive pour apprendre automatiquement des rubriques spécifiques aux tâches (Étape I) puis les utilisent pour guider le raisonnement (Étape II), éliminant la dépendance à l'annotation manuelle de rubriques ou aux récompenses coûteuses de type "LLM-as-a-Judge" pendant la phase de distillation.
3. Apprentissage on-policy conscient des critères : Contrairement à l'OPSD standard qui se conditionne sur une trajectoire unique, la RCSD se conditionne sur des critères, permettant une attribution de crédit fine qui pénalise les erreurs locales spécifiques tout en préservant les étapes correctes, même sur les propres trajectoires échantillonnées de l'étudiant.

Résultats expérimentaux

La méthode a été évaluée sur une suite diversifiée de benchmarks de raisonnement scientifique, incluant des tâches vérifiables (GPQA-DIAMOND, SciBench, PIQA) et des tâches ouvertes non vérifiables (ResearchQA, RubricHub).

• Performance : La RCSD a obtenu le meilleur score moyen global (70,6) sur l'ensemble des benchmarks. Elle a surpassé le modèle de base Qwen3-8B de 4,7 points, dépassant GRPO de 1,0 point et OPSD de 0,9 point en moyenne.
• Spécificité du domaine : Les gains ont été plus prononcés sur les tâches de raisonnement scientifique et basées sur des rubriques (par exemple, +8,2 points sur ResearchQA, +4,9 points sur RubricHub), là où les récompenses scalaires échouent souvent à capturer la qualité de la réponse.
• Généralisation : La méthode a démontré une robustesse de généralisation vers des benchmarks médicaux hors domaine (MedMCQA, PubMedQA) sans oubli catastrophique, suggérant que les critères appris capturent des structures de raisonnement générales.
• Études d'ablation :
  • Type de perte : La divergence KL directe (Forward KL) s'est avérée être l'objectif le plus efficace, produisant la performance moyenne la plus élevée et maintenant une entropie (diversité) de l'étudiant plus élevée par rapport à la Reverse KL ou à la divergence de Jensen-Shannon.
  • Source de la rubrique : L'utilisation de rubriques générées par le générateur de l'Étape I a donné des résultats compétitifs avec l'utilisation de rubriques de référence de vérité terrain, validant la scalabilité de l'approche en deux étapes.
  • Robustesse : La méthode est restée efficace même avec des entrées de rubriques dégradées (génériques, bruitées ou réduites), bien que les rubriques spécifiques à l'instance fournissent la supervision la plus forte.

Signification et affirmations

L'article affirme que la RCSD offre une approche sans vérificateur pour le post-entraînement de modèles de raisonnement ouverts. En préservant la structure au niveau des critères pendant l'optimisation et en convertissant les rubriques en un guidage dense au niveau du jeton, la méthode répond au goulot d'étranglement de l'attribution de crédit inhérent à l'optimisation par récompense scalaire et au biais lié au chemin spécifique de la distillation.

Les auteurs soutiennent que la RCSD est particulièrement significative pour les tâches où les réponses de haute qualité ne peuvent pas être pleinement capturées par une vérification automatique (réponses par correspondance exacte ou tests unitaires) ou par des récompenses de résultat scalaires. En permettant au modèle d'apprendre ce que les propriétés d'une réponse doivent satisfaire (via les rubriques) plutôt que comment imiter un chemin de solution spécifique, la RCSD facilite des améliorations de raisonnement plus flexibles et robustes. Le travail suggère qu'un feedback structuré et multidimensionnel est un signal de supervision supérieur pour les tâches de raisonnement complexes par rapport aux récompenses éparses et à la distillation sur une trajectoire unique.