Résumé technique : Pierres angulaires ou obstacles ? Décryptage des jetons rocheux dans la distillation sur politique

1. Énoncé du problème

La distillation sur politique (OPD) est devenue une pierre angulaire des pipelines de post-entraînement pour les grands modèles de langage (LLM), permettant un raffinement du raisonnement au niveau des jetons en alignant un modèle étudiant sur un enseignant supérieur, à l'aide de trajectoires échantillonnées à partir de la propre politique de l'étudiant. Alors que des travaux récents en apprentissage par renforcement avec récompenses vérifiables (RLVR) ont identifié qu'un petit sous-ensemble de jetons critiques (par exemple, des « jetons de bifurcation » à haute entropie) génère de manière disproportionnée des gains de raisonnement, une compréhension analogue au niveau des jetons de l'OPD reste inexplorée.

Dans l'OPD, l'objectif consiste à minimiser la divergence KL inverse par jeton entre les politiques de l'étudiant et de l'enseignant. Intuitivement, à mesure que l'entraînement converge, les jetons à forte perte (indiquant un décalage étudiant-enseignant) devraient diminuer. Cependant, des observations empiriques révèlent un paradoxe : même après que l'entraînement a atteint une saturation apparente, un sous-ensemble substantiel de jetons continue d'afficher une perte persistamment élevée. La question centrale abordée par ce travail est de savoir si ces jetons à perte élevée persistante sont des « Pierres angulaires » (Piliers) indispensables pour l'alignement des politiques ou simplement des « Obstacles » créant une redondance d'optimisation inutile.

2. Méthodologie

2.1 Définition et détection des jetons rocheux

Les auteurs définissent les jetons rocheux comme des types de jetons qui présentent un comportement à perte élevée persistant, stable à travers des contextes similaires, plutôt que des pics isolés dans de longues trajectoires. Le processus de détection comporte trois étapes :

1. Calcul du fardeau agrégé : Calcul d'un « Score rocheux » $R(v)$ pour chaque type de jeton $v$, défini comme le produit de sa perte moyenne au niveau du jeton ($\bar{\ell}_v$) et de sa fréquence ($Freq(v)$) dans les trajectoires générées par l'étudiant.
2. Filtrage de la persistance : Identification des occurrences qui restent à perte élevée avant et après l'entraînement OPD.
3. Filtrage de la cohérence contextuelle : Assurance que les occurrences à perte élevée se répètent dans des contextes locaux similaires (en utilisant une métrique de similarité de fenêtre contextuelle). Cela distingue les incohérences structurelles authentiques du bruit positionnel isolé.

La métrique finale, le Score rocheux sensible au contexte $R_{ctx}(v)$, combine le fardeau agrégé avec la cohérence contextuelle. Un seuil $K=100$ est sélectionné pour équilibrer la stabilité de la sélection (reproductibilité) et la couverture de la perte (capturant environ 60 % de la divergence KL au niveau du corpus).

2.2 Géométrie du gradient et analyse de la persistance

Pour déterminer si les jetons rocheux fournissent des signaux d'apprentissage utiles, les auteurs analysent la géométrie des gradients par jeton :

• Magnitude vs Fréquence : Comparaison de la magnitude du gradient par occurrence des jetons rocheux avec celle des jetons rares à forte KL.
• Alignement : Mesure de l'alignement cosinus des gradients des jetons rocheux avec une direction de descente globale équilibrée par fréquence ($G_{balanced}$).
• Stagnation : Suivi de la variation de la divergence KL par jeton entre les points de contrôle de l'entraînement précoce et tardif pour déterminer si la perte diminue effectivement.

2.3 Intervention causale : Expériences de mise hors service

Pour déterminer la nécessité fonctionnelle de ces jetons, les auteurs emploient une procédure de mise hors service au moment de l'inférence. Ils forment une politique $\pi_{\theta}^{\setminus v}$ en forçant le logit d'un jeton candidat $v$ à $-\infty$ à chaque étape de décodage. L'impact causal est mesuré comme la variation de la précision sur les benchmarks ($\Delta_B(v)$). Les jetons sont catégorisés comme suit :

• Piliers forts : La suppression dégrade significativement la précision.
• Obstacles forts : La suppression améliore la précision.
• Neutres : La suppression a un effet négligeable.

2.4 Distillation sélective (Gel des gradients)

Pour évaluer l'impact des jetons rocheux sur le processus d'entraînement, les auteurs introduisent un mécanisme de repondération de jetons sensible à la fenêtre. Ils gèlent les gradients des jetons rocheux et de leurs fenêtres locales associées à forte divergence (en fixant le poids $\lambda=0$) pendant l'entraînement OPD. Cette stratégie « Juste ne pas entraîner » leur permet d'isoler si ces jetons sont des ancres fonctionnelles essentielles ou des fardeaux d'optimisation improductifs.

3. Résultats clés

3.1 Identité et omniprésence

• Composition : Les jetons rocheux sont principalement des échafaudages syntaxiques et structurels, incluant les délimiteurs LaTeX/mathématiques, les structures Markdown/espaces blancs, les marqueurs de discours (par exemple, « Donc », « Attendez ») et les chiffres. Ils marquent les frontières du raisonnement plutôt que des mots porteurs de contenu.
• Prévalence : Bien qu'ils ne représentent qu'environ 6 % du vocabulaire, les jetons rocheux comptent pour environ 18 % de la fréquence totale des jetons dans les sorties générées.

3.2 Le paradoxe du gradient

• Alignement : Contrairement à l'hypothèse selon laquelle ils seraient du bruit, les jetons rocheux présentent un alignement directionnel positif non trivial avec l'objectif d'optimisation global. Ils fournissent des gradients qui pointent vers la bonne voie d'optimisation.
• Stagnation : Malgré cet alignement, les jetons rocheux restent stagnants tout au long de l'entraînement. Leur divergence KL ne diminue pas de manière significative (médiane $\Delta KL \approx 0$), contrairement aux jetons rares à forte KL qui montrent une réduction substantielle.
• Cause : Les auteurs émettent l'hypothèse que les optimiseurs adaptatifs (comme Adam) suppriment le taux d'apprentissage effectif pour ces jetons omniprésents en raison d'un déséquilibre de classe à queue lourde, piégeant l'étudiant dans des habitudes structurelles préexistantes malgré la présence de signaux directionnels corrects.

3.3 Nécessité causale (Piliers vs Obstacles)

• Rareté des piliers : Seule une infime fraction (1,5 % à 3,5 %) des jetons rocheux agit comme des « Piliers forts » dont la suppression dégrade les performances de raisonnement.
• Absence d'obstacles : Aucun « Obstacle fort » n'a été trouvé ; la suppression des jetons rocheux n'a jamais amélioré les performances.
• Orthogonalité : L'importance causale d'un jeton (le fait d'être un pilier) est orthogonale aux métriques standard telles que l'entropie, la fréquence ou la KL résiduelle. Les taxonomies d'importance existantes échouent à prédire quels jetons sont causalement nécessaires.
• Zone tampon : Le modèle étudiant semble maintenir une « zone tampon » de résidus persistants pour protéger quelques ancres sémantiques critiques, plutôt que de s'appuyer sur l'ensemble des jetons à perte élevée.

3.4 Efficacité de l'entraînement

• Gel sélectif : En gelant les gradients des jetons rocheux (et de leurs fenêtres locales) pendant l'entraînement, les auteurs ont obtenu une augmentation de 1,4 à 1,7 fois de la vitesse d'entraînement en temps réel.
• Impact sur les performances : Cette accélération a entraîné une perte négligeable dans les performances de raisonnement en aval (mesurées sur les benchmarks AIME, HMMT et MATH). Cela suggère qu'une partie substantielle de la bande passante d'optimisation consacrée aux jetons rocheux est redondante.

4. Signification et affirmations

L'article remet en question la sagesse conventionnelle selon laquelle l'alignement de chaque jeton à perte élevée est indispensable pour l'apprentissage dans la distillation sur politique. Les auteurs affirment :

1. Identification d'une dynamique cachée : Ils identifient les « jetons rocheux » comme un phénomène de perte élevée persistant qui révèle un aspect caché des dynamiques d'optimisation de l'OPD, distinct des « jetons de bifurcation » observés dans le RLVR.
2. Inefficacité des pipelines actuels : Une partie substantielle des signaux d'entraînement dans l'OPD fournit des contributions nettement inférieures à l'attendu aux performances de raisonnement ultime du modèle.
3. Frontière efficacité-performance : Le travail démontre que contourner stratégiquement ces « obstacles » (via le gel des gradients) peut considérablement rationaliser le processus d'alignement. Cela offre un paradigme plus efficace pour la distillation de modèles à grande échelle sans compromettre l'intégrité du raisonnement.

L'article conclut que si les jetons rocheux ne sont pas entièrement jetables (car un petit sous-ensemble agit comme des piliers nécessaires), une allocation non uniforme des gradients peut efficacement contourner la redondance computationnelle, optimisant ainsi la frontière efficacité-performance dans l'alignement sur politique.

5. Limites

Les auteurs notent que leurs évaluations empiriques reposent principalement sur le raisonnement mathématique compétitif, et que le rôle des jetons rocheux peut différer dans la génération ouverte ou les tâches de codage. De plus, l'étude utilise un appariage spécifique de distillation de 30B vers 4B, laissant la généralisabilité à travers différentes familles d'architectures ou échelles comme une question ouverte. Enfin, la stratégie d'atténuation emploie un masque binaire strict de gel des gradients, que les auteurs suggèrent pourrait être affiné dans des travaux futurs avec des schémas de repondération dynamique.