💬 NLP

Masked Distillation: Internalizing the Chain-of-Thought in Language Models

Cet article introduit la « distillation masquée », un cadre qui internalise les capacités de raisonnement des grands modèles de raisonnement dans les paramètres d'un modèle étudiant en l'entraînant à prédire uniquement les jetons de solution tout en recevant un retour conditionné par la chaîne de pensée de l'enseignant, permettant ainsi une génération de réponses directe et une réduction de la latence d'inférence sans sacrifier la précision.

Auteurs originaux : Durgesh Kalwar, Vardhan Palod, Subbarao Kambhampati

Publié 2026-07-28

📖 6 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Durgesh Kalwar, Vardhan Palod, Subbarao Kambhampati

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez un monde où votre assistant IA préféré est comme un étudiant brillant mais bavard. Lorsque vous lui posez une question difficile, il ne se contente pas de donner la réponse brusquement ; il griffonne des pages de notes, dessine des diagrammes et détaille chaque étape de son raisonnement avant de finalement écrire la solution. Ce processus de « pensée à voix haute », connu dans le monde technologique sous le nom de « Chaîne de Pensée » (Chain of Thought), a rendu les modèles d'IA incroyablement intelligents pour les problèmes de mathématiques et de logique. Cependant, il y a un hic : ces pages de notes prennent beaucoup de temps et d'énergie à générer. C'est comme payer un chauffeur de taxi pour vous emmener au supermarché, mais aussi le payer pour faire tout le tour du pâté de maisons avant de « réfléchir » à l'allée dont vous avez besoin. Les scientifiques se demandent depuis longtemps : peut-on apprendre à l'IA à faire toute cette réflexion à l'intérieur de sa tête ? Si nous y parvenions, l'IA pourrait sauter les longues notes bavardes et vous donner la réponse instantanément, économisant ainsi un temps et une puissance de calcul massifs.

Cet article, intitulé « Masked Distillation » (Distillation Masquée), plonge directement dans cette question. Les chercheurs, travaillant avec des Modèles de Raisonnement de Grande Taille (LRM), ont cherché à savoir s'ils pouvaient prendre un modèle « enseignant » intelligent qui aime réfléchir à voix haute et entraîner un modèle « étudiant » plus petit à intérioriser ce processus de réflexion. Ils voulaient voir si l'étudiant pouvait apprendre à sauter les notes entièrement pour simplement donner la réponse finale, tout en étant tout aussi intelligent. Ils ont testé cela en utilisant une technique ingénieuse appelée « Masked Distillation ». Imaginez que l'enseignant résout un puzzle pendant que l'étudiant regarde. L'enseignant voit la question complète et son propre processus de pensée détaillé, mais l'étudiant n'est présenté qu'à la question. Le but était de voir si l'étudiant pouvait apprendre la logique secrète de l'enseignant et ensuite résoudre le puzzle par lui-même sans jamais avoir besoin d'écrire les étapes, même si l'enseignant avait toutes les notes pour guider l'apprentissage.

L'équipe a mis en place deux expériences principales. Dans la première, ils ont utilisé le même modèle d'IA comme enseignant et comme étudiant, en lui demandant simplement de réfléchir différemment. Dans la seconde, ils ont utilisé un modèle plus grand et plus intelligent comme enseignant et un modèle plus petit et plus simple comme étudiant. Ils ont également introduit un concept d'« échafaudage » (scaffold). Voyez cela comme des petites roues sur un vélo. Parfois, l'étudiant peut avoir besoin de voir un peu des notes de l'enseignant (l'échafaudage) pour trouver la bonne réponse, plutôt que d'essayer de tout faire à partir de zéro. Ils ont testé de combien de la pensée de l'enseignant l'étudiant avait besoin pour obtenir les meilleurs résultats.

Les résultats étaient un mélange de « wow » et de « pas si vite ». Sur un ensemble de données de problèmes mathématiques de niveau primaire appelé GSM8K, les modèles étudiants étaient formidables. Ils ont réussi à intérioriser la pensée de l'enseignant. L'étudiant entièrement « masqué » (celui qui n'écrivait jamais de notes) a obtenu 76,0 % de bonnes réponses, ce qui est un bond énorme par rapport à sa capacité initiale, et il l'a fait en utilisant seulement 369,6 tokens (unités de texte) en moyenne, contre 2 398,1 tokens pour l'enseignant. C'est une réduction d'effort massive de 6,5 fois !

Cependant, sur un puzzle numérique plus complexe appelé Countdown, l'histoire change. Ici, l'étudiant qui essayait de tout faire dans sa tête (entièrement masqué) a en fait été moins performant, tombant à 41,7 % de précision. Il s'est avéré que pour ce type spécifique de puzzle, l'étudiant avait besoin de voir une partie des notes de l'enseignant pour réussir. Mais voici la partie cool : les chercheurs ont trouvé un point d'équilibre. En laissant l'étudiant voir juste un petit peu de la pensée de l'enseignant (environ 30 % des notes, ou un « alpha » de 0,3), la précision de l'étudiant a grimpé en flèche pour atteindre 86,2 %. C'était presque aussi bon que l'enseignant (87,3 %), mais utilisait tout de même environ 1,3 fois moins de tokens que l'enseignant. C'était comme donner à l'étudiant juste assez de petites roues pour garder l'équilibre parfaitement sans avoir besoin de tout le vélo.

L'article a également examiné la capacité de ces étudiants à gérer de nouveaux puzzles inédits (tâches hors distribution). Sur le plan mathématique, les étudiants s'en sont bien sortis, mais sur les puzzles Countdown, les résultats étaient un peu plus mitigés. Curieusement, dans certains cas, les étudiants entraînés avec un peu d'échafaudage ont mieux généralisé que l'enseignant lui-même lorsque les puzzles devenaient plus difficiles ou changeaient légèrement.

Une découverte clé est que la méthode d'enseignement importait. Lorsqu'ils ont essayé d'enseigner à l'étudiant simplement en lui montrant les notes de l'enseignant et en lui demandant de les copier (une méthode appelée Apprentissage Supervisé par Affinement ou Supervised Fine-Tuning), cela n'a pas aussi bien fonctionné. La méthode de « Masked Distillation », qui utilise une façon plus complexe de comparer les suppositions de l'étudiant aux pensées réelles de l'enseignant, était beaucoup plus efficace pour transférer les compétences de « réflexion ».

En fin de compte, l'article suggère que nous ne pouvons pas simplement supprimer magiquement toutes les étapes de réflexion pour chaque problème. La capacité d'une IA à « réfléchir silencieusement » dépend fortement du type de problème et de ce que l'étudiant sait déjà sur ce sujet. Pour des tâches familières comme les mathématiques de base, l'IA peut apprendre à réfléchir silencieusement et gagner énormément de temps. Pour des puzzles plus difficiles et moins familiers, elle pourrait encore avoir besoin d'un petit coup de pouce — un petit échafaudage de notes — pour accomplir le travail correctement. Les chercheurs concluent que cet « échafaudage » est un outil puissant que l'on peut ajuster, permettant de trouver l'équilibre parfait entre l'intelligence de l'IA et la rapidité avec laquelle elle peut donner une réponse.

Résumé Technique : Distillation Masquée – Internaliser la Chaîne de Pensée dans les Modèles de Langage

Énoncé du Problème

Les Grands Modèles de Raisonnement (LRM) atteignent des performances élevées sur des tâches complexes (mathématiques, code, planification) en générant de longues chaînes explicites de jetons intermédiaires (Chaîne de Pensée ou CoT - Chain-of-Thought) avant de produire une réponse finale. Bien que ces traces améliorent la précision, elles engendrent des coûts significatifs en termes de latence d'inférence, d'utilisation de la mémoire (empreinte du cache KV) et d'énergie de service, qui augmentent linéairement avec la longueur de la trace.

Des découvertes récentes suggèrent que la justesse de ces traces intermédiaires n'est pas liée causalement à la justesse de la réponse finale, et que la longueur de la trace ne corrèle pas de manière fiable avec la complexité du problème. Cela soulève une question critique : La computation exprimée dans ces jetons intermédiaires peut-elle être internalisée dans les paramètres d'un modèle de langage, permettant ainsi de produire des réponses correctes directement (ou avec des traces nettement plus courtes) sans générer explicitement les étapes de raisonnement lors de l'inférence ?

Méthodologie : Distillation Masquée

Les auteurs proposent la Distillation Masquée, un cadre de distillation de connaissances conçu pour internaliser le processus de raisonnement d'un modèle enseignant « pensant » dans un modèle étudiant « non-pensant ».

Cadre Fondamental

Dans ce cadre, le modèle étudiant est entraîné pour prédire uniquement les jetons de la solution conditionnés par la question d'entrée. Le modèle enseignant, cependant, fournit un feedback basé à la fois sur la question et sur sa propre trace de CoT générée. L'objectif est de minimiser la divergence entre la distribution de l'étudiant (conditionnée uniquement par l'entrée) et la distribution de l'enseignant (conditionnée par l'entrée et le raisonnement intermédiaire).

L'objectif d'entraînement utilise la divergence KL inverse on-policy :
$L_{MD} = D_{KL}(\pi_S(\cdot | x, y_{<t}) \parallel \pi_T(\cdot | x, \text{ITs}, y_{<t}))$
Où $\pi_S$ est la politique de l'étudiant, $\pi_T$ est la politique de l'enseignant, $x$ est l'entrée, $y_{<t}$ sont les jetons précédemment générés, et $\text{ITs}$ sont les jetons intermédiaires de l'enseignant.

Paramètres Expérimentaux

Les auteurs instaurent ce cadre dans deux configurations :

Auto-distillation : Le même modèle sert d'enseignant (en mode pensée) et d'étudiant (en mode non-pensée).
Configuration à Double Modèle : Un modèle de raisonnement plus grand sert d'enseignant, supervisant un étudiant non-pensant distinct et plus petit.

L'Échafaudage par Suffixe $\alpha$

Pour explorer le compromis entre l'internalisation complète et le coût d'inférence, les auteurs introduisent une variante de la distillation masquée par suffixe $\alpha$ :

$\alpha = 0$ (Entièrement Masqué) : L'étudiant internalise l'intégralité de la trace intermédiaire et émet uniquement la solution.
$\alpha = 1$ (Non-Masqué) : L'étudiant reproduit la trace complète avant la réponse (équivalent à une génération de CoT standard).
$0 < \alpha < 1$ : L'étudiant internalise le préfixe de $(1-\alpha)$ de la trace de l'enseignant et est entraîné à émettre le suffixe restant de fraction $\alpha$ de la trace lors de l'inférence, aux côtés de la solution.

Cela permet aux auteurs de traiter les jetons intermédiaires comme un « échafaudage » qui peut être partiellement conservé ou entièrement internalisé.

Contributions Clés

Cadre Proposé : Introduction de la Distillation Masquée pour internaliser le raisonnement CoT dans les paramètres du modèle, permettant la génération directe de réponses.
Analyse de l'Échafaudage : Définition et évaluation du paramètre de suffixe $\alpha$ , établissant un axe contrôlé pour ajuster le compromis précision/coût d'inférence.
Évaluation Exhaustive : Réalisation d'expérimentations approfondies sur deux domaines :
- Mathématiques : GSM8K (In-Distribution), MATH-500, et AIME-25 (Out-of-Distribution).
- Countdown : Un jeu de recherche de type puzzle numérique avec des variations de plage cible et de profondeur de recherche (splits OOD).
Études d'Ablation : Comparaison de la Distillation Masquée (utilisant la KL inverse) par rapport au Fine-Tuning Supervisé (SFT) sur les traces de l'enseignant, ainsi qu'à l'utilisation de « jetons d'enseignement » explicites comme échafaudages.

Résultats

1. Dépendance de l'Internalisation Complète vis-à-vis de la Tâche

La capacité d'internaliser totalement le raisonnement ( $\alpha=0$ ) dépend fortement de la tâche et de l'exposition préalable de l'étudiant au domaine lors du pré-entraînement :

GSM8K (Mathématiques) : L'internalisation complète est un succès. Dans le cadre de l'auto-distillation, l'étudiant a atteint 76,0 % de précision (contre 87,9 % pour l'enseignant) tout en réduisant les jetons d'inférence par 6,5×. Dans la configuration à double modèle, la précision est passée de 15,69 % (base) à 45,26 % avec une réduction de jetons de 5×.
Countdown : L'internalisation complète échoue sans échafaudage. En auto-distillation, l'étudiant entièrement masqué est tombé à 41,7 % de précision (en dessous des 56,2 % du modèle de base). Dans la configuration à double modèle, il est passé de 0 % à 34,1 %, mais reste bien en dessous de l'enseignant (87,3 %) et de la variante non-masquée (81,1 %).

2. Efficacité de l'Échafaudage par Suffixe $\alpha$

Fournir un faible amount d'échafaudage (émission d'une fraction de la trace) récupère significativement la performance, particulièrement sur les tâches où l'internalisation complète échoue :

Countdown (Auto-distillation) : Le réglage $\alpha = 0,3$ a porté la précision de 41,7 % à 86,2 %, égalant essentiellement l'enseignant (87,3 %) tout en réduisant les jetons d'inférence d'environ 1,3× par rapport à la variante non-masquée.
Countdown (Double Modèle) : La variante $\alpha = 0,5$ a atteint 84,2 % de précision, surpassant la ligne de base non-masquée (81,1 %) tout en utilisant moins de jetons, établissant un point Pareto-optimal.
GSM8K : De petits échafaudages ( $\alpha=0,3$ ) ont récupéré la majeure partie de l'écart de précision (83,2 % contre 87,9 %) avec environ la moitié du coût en jetons de la variante non-masquée.

3. Généralisation vers l'Out-of-Distribution (OOD)

Décalage de Plage Cible (Target-Range Shift) : Les modèles entraînés avec divers paradigmes d'échafaudage ont transféré proprement vers des tâches OOD où la plage de nombres cibles changeait.
Décalage de Profondeur de Recherche (Search-Depth Shift) : Dans le cadre de l'auto-distillation sur Countdown, les étudiants entraînés avec échafaudage (spécifiquement $\alpha=0,5$ ) ont mieux généralisé sous une profondeur de recherche accrue (32,8 % contre 25,7 %).
Limitations du Double Modèle : Dans la configuration à double modèle sur des benchmarks mathématiques plus difficiles (MATH-500, AIME-25), l'écart de capacité entre l'enseignant et l'étudiant était trop important ; ni l'internalisation complète ni l'échafaudage n'ont permis de récupérer la performance OOD, toutes les variantes s'effondrant au niveau plancher du modèle étudiant de base.

4. Objectifs d'Entraînement

Reverse-KL vs SFT : L'article conclut que le Fine-Tuning Supervisé (SFT) sur les traces de l'enseignant échoue à exploiter efficacement l'échafaudage. Sur Countdown, le SFT avec échafaudage a produit environ 38 % de précision, tandis que la Distillation Masquée par Reverse-KL avec les mêmes échafaudages a atteint environ 67 %.
Jetons d'Enseignement : Fournir des « jetons d'enseignement » explicites comme échafaudages a amélioré la performance par rapport au masquage complet, mais s'est avéré moins efficace que l'approche par suffixe $\alpha$ .

Signification et Revendications

L'article affirme que l'échafaudage par suffixe fournit un mécanisme contrôlable pour équilibrer la précision et l'efficacité de l'inférence. Le point de fonctionnement optimal est déterminé par la tâche spécifique et l'écart de capacité entre l'enseignant et l'étudiant.

L'internalisation est possible mais pas universelle : Les modèles peuvent internaliser le raisonnement pour des tâches dont ils sont familiers (comme GSM8K) mais peinent face à des tâches de recherche inédites (comme Countdown) sans conserver une trace explicite.
Compromis Coût-Performance : Un faible amount de raisonnement explicite (ex: $\alpha=0,3$ ) peut récupérer la quasi-totalité de la performance de la génération de CoT complète tout en offrant des économies substantielles de coût d'inférence (réduction de 1,3× à 2× des jetons).
Généralisation : L'échafaudage peut améliorer la généralisation dans certains scénarios OOD (ex: décalages de profondeur de recherche en auto-distillation), suggérant que conserver une trace partielle aide le modèle à naviguer dans des espaces de recherche complexes mieux qu'un modèle entièrement internalisé ou qu'une trace complète dans certains contextes.

Les auteurs concluent que, bien que l'internalisation complète soit un objectif attractif pour l'efficacité, l'approche par suffixe $\alpha$ offre une alternative pratique et ajustable qui maintient une haute performance avec une latence réduite.