Résumé Technique : OISD – Auto-distillation Interne On-Policy des Modèles de Langage

1. Énoncé du Problème

Les récentes approches de post-entraînement par apprentissage par renforcement (RL) pour les grands modèles de langage (LLM), en particulier celles utilisant l'optimisation de politique relative par groupe (GRPO) avec des récompenses vérifiables, ont considérablement fait progresser le raisonnement mathématique. Cependant, ces méthodes sont principalement centrées sur la sortie finale. Elles optimisent la politique d'action finale à l'aide de récompenses clairsemées au niveau des résultats (par exemple, la justesse de la réponse finale), tout en négligeant largement les signaux prédictifs, les processus de raisonnement et les représentations internes encodés dans les couches intermédiaires.

Bien que la distillation de connaissances (KD) offre un mécanisme pour fournir une supervision dense au niveau des tokens, les paradigmes existants souffrent de limitations spécifiques :

• Distillation Off-Policy : Crée un décalage de distribution car l'élève est entraîné sur des trajectoires générées par un enseignant fixe, et non sur ses propres générations.
• Distillation On-Policy : Réduit ce décalage mais repose généralement sur des modèles enseignants externes.
• Auto-distillation On-Policy : Les méthodes récentes éliminent les enseignants externes mais dépendent d'informations privilégiées injectées de l'extérieur (par exemple, des invites privilégiées, des solutions de vérité terrain ou des retours d'information auxiliaires) pour superviser le modèle.

Il existe un vide dans l'exploitation des propres calculs internes du modèle comme source de supervision, sans nécessiter de données privilégiées externes ni de multiples politiques d'action. De plus, des recherches antérieures indiquent que les représentations de raisonnement évoluent systématiquement à travers les couches (des prédictions de candidats larges dans les premières couches aux décisions convergées dans les couches profondes), suggérant que les couches intermédiaires pourraient être guidées par les couches profondes pour améliorer la cohérence du raisonnement.

2. Méthodologie : Cadre OISD

Les auteurs proposent OISD (Auto-distillation Interne On-Policy), un cadre qui traite la couche finale du modèle comme un « enseignant interne détaché » et une couche intermédiaire sélectionnée comme un « élève interne auxiliaire ». Cela se produit au cours d'un seul passage avant (forward pass) durant le processus de déploiement (rollout) et d'optimisation GRPO.

Architecture Principale

• Politique d'Action Unique : La couche finale ($L$) reste la seule politique pour générer les déploiements et optimiser l'objectif GRPO.
• Enseignant-Élève Interne : Au cours du même passage avant, l'état caché d'une couche intermédiaire sélectionnée ($l < L$) est projeté dans l'espace des tokens à l'aide de la matrice de décodage (suivant le cadre de la logit lens). Cela crée une distribution « élève » provisoire dérivée du même contexte que l'« enseignant » (couche finale).
• Aucune Privilège Externe : Les signaux de supervision sont dérivés entièrement des états internes du modèle sur le même déploiement on-policy, ne nécessitant aucune solution de vérité terrain externe ni aucune invite privilégiée.

Objectifs d'Alignement

OISD emploie deux mécanismes d'alignement interne complémentaires pour transférer les signaux prédictifs de la couche finale vers la couche intermédiaire :

1. Alignement des Logits (« Comment Penser ») :

   • Objectif : Transférer les comportements de raisonnement de haut niveau et les croyances prédictives.
   • Mécanisme : La divergence de Jensen–Shannon (JS) est calculée entre les distributions de tokens mises à l'échelle par température de la couche intermédiaire ($p^l_t$) et de la couche finale ($q^L_t$).
   • Pondération : La perte est pondérée par l'avantage GRPO signé ($A_t$). Les tokens à fort avantage attirent la couche intermédiaire vers la distribution de la couche finale, tandis que les tokens à faible avantage la repoussent pour éviter de renforcer les continuations médiocres.
   • Formule : $\mathcal{L}^{(l)}_{\text{think}} = \mathbb{E}_t [ A_t \cdot \text{JS}(p^{l,\tau}_t, \text{sg}(q^{L,\tau}_t)) ]$

2. Alignement de l'Attention (« Où Regarder ») :

   • Objectif : Imposer des motifs d'attention cohérents et un comportement de routage des preuves.
   • Mécanisme : La divergence JS est calculée entre les poids d'attention de la couche intermédiaire et de la couche finale sur un ensemble échantillonné de clés causales (combinant des fenêtres locales et des positions espacées globales pour gérer la mémoire).
   • Pondération : À l'instar de l'alignement des logits, la perte est pondérée par l'avantage signé. Les avantages positifs encouragent la couche intermédiaire à prêter attention aux mêmes preuves que la couche finale ; les avantages négatifs réduisent l'accord avec les motifs à faible récompense.
   • Formule : $\mathcal{L}^{(l)}_{\text{attn}} = \mathbb{E}_t [ A_t \cdot \frac{1}{H} \sum_{h=1}^H \text{JS}(\tilde{a}^l_{t,h}, \text{sg}(\tilde{a}^L_{t,h})) ]$

Objectif d'Entraînement

La perte totale combine la perte GRPO standard avec les deux pertes d'alignement interne :
$$\mathcal{L}_{\text{OISD}} = \mathcal{L}_{\text{GRPO}} + \lambda_{\text{think}} \mathcal{L}_{\text{think}} + \lambda_{\text{attn}} \mathcal{L}_{\text{attn}}$$
où $\lambda_{\text{think}}$ et $\lambda_{\text{attn}}$ contrôlent l'intensité de la distillation.

3. Contributions Clés

1. Nouveau Paradigme : Introduction de l'Auto-distillation Interne On-Policy, un paradigme de RL de raisonnement qui exploite les propres calculs internes du modèle pour la supervision tout en maintenant une politique d'action unique.
2. Cadre OISD : Une implémentation spécifique utilisant la couche finale comme enseignant détaché et une couche intermédiaire comme élève, transférant des signaux prédictifs on-policy sans informations privilégiées externes.
3. Mécanismes d'Alignement Dual : Développement de deux objectifs complémentaires :
   • Alignement des Logits : Supervise le processus de raisonnement (« comment penser »).
   • Alignement de l'Attention : Supervise le routage des preuves (« où regarder »).
4. Validation Empirique : Démonstration d'améliorations substantielles et cohérentes par rapport à des bases de RL de raisonnement solides (incluant GRPO et BuPO) sur quatre benchmarks de raisonnement mathématique.

4. Résultats Expérimentaux

Les auteurs ont évalué OISD sur les modèles Qwen3-4B, Qwen3-8B et OctoThinker à travers quatre benchmarks : AMC23, MATH500, AIME24 et AIME25.

• Gains de Performance :
  • Sur Qwen3-4B, OISD a obtenu un score moyen de 64,75, surpassant GRPO (55,08) de 9,67 points et BuPO (58,51) de 6,24 points.
  • Les améliorations ont été particulièrement significatives sur les benchmarks AIME (par exemple, +14,31 points par rapport à GRPO sur AIME24).
  • Sur Qwen3-8B, OISD a obtenu la meilleure moyenne globale de 66,72, surpassant GRPO (64,23) et BuPO (66,36).
• Analyse Pass@K : OISD a constamment atteint des taux Pass@K plus élevés pour toutes les valeurs de $K$ par rapport à GRPO, indiquant une amélioration de la qualité globale de l'ensemble des solutions échantillonnées, et non seulement de la précision moyenne.
• Études d'Ablation :
  • L'alignement des Logits et l'alignement de l'Attention ont tous deux contribué positivement. L'alignement des Logits a montré le gain le plus important sur AIME24, tandis que l'alignement de l'Attention a légèrement mieux performé en moyenne.
  • La combinaison des deux a produit les meilleurs résultats, confirmant leur nature complémentaire.
• Comparaison avec OPSD : OISD (64,75) a nettement surpassé OPSD (56,63), une méthode d'auto-distillation on-policy de l'état de l'art qui repose sur des solutions de vérité terrain privilégiées. Cela met en évidence l'efficacité de la supervision purement interne d'OISD.
• Analyse Qualitative :
  • Dynamique d'Entraînement : OISD a montré une amélioration régulière de la récompense et une augmentation initiale de l'entropie (suggérant une exploration plus large) avant de se stabiliser.
  • Visualisation de l'Attention : Les couches intermédiaires dans OISD se sont progressivement alignées sur l'attention de la couche finale envers les tokens de preuve critiques (par exemple, les extrémités, le supremum) au cours de l'entraînement.
  • Visualisation Logit-Lens : OISD a présenté un alignement plus précoce entre les prédictions intermédiaires et finales sur les tokens critiques pour le raisonnement par rapport à GRPO, qui produisait souvent des prédictions intermédiaires fragmentées.

5. Signification et Revendications

L'article revendique que OISD initie l'investigation de l'auto-distillation interne on-policy. Sa signification principale réside dans :

• Élimination de la Dépendance Externe : Il fournit une supervision interne dense sans nécessiter d'enseignants externes, d'invites privilégiées ou de solutions de vérité terrain durant l'entraînement.
• Cohérence de la Politique : Il préserve une politique d'action unique pour la génération de déploiements et l'optimisation, évitant les problèmes de décalage de distribution des méthodes off-policy.
• Raisonnement Amélioré : En alignant les représentations intermédiaires avec le raisonnement et les motifs d'attention de la couche finale, OISD encourage le modèle à apprendre des représentations intermédiaires plus fortes et plus globalement cohérentes.

Les auteurs notent des limitations, notamment le focus sur le raisonnement mathématique avec des réponses vérifiables et l'absence d'une étude approfondie sur la sélection optimale des couches. Ils suggèrent que les travaux futurs devraient explorer la généralisation à des tâches de raisonnement plus complexes et la supervision interne multi-couches.