💬 NLP

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

本論文は、「マスクされた蒸留（masked distillation）」というフレームワークを導入するものであり、これは、教師モデルの思考の連鎖（Chain-of-Thought）に基づいたフィードバックを受け取りつつ、解のトークンのみを予測するように学習させることで、大規模推論モデルの推論能力を生徒モデルのパラメータへと内面化させ、それによって精度を損なうことなく直接的な回答生成を可能にし、推論レイテンシを削減するものである。

原著者： Durgesh Kalwar, Vardhan Palod, Subbarao Kambhampati

公開日 2026-07-28

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Durgesh Kalwar, Vardhan Palod, Subbarao Kambhampati

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたの最もお気に入りのAIアシスタントが、まるで「非常に優秀だがおしゃべりな学生」のような世界を想像してみてください。あなたがトリッキーな質問を投げかけると、そのAIはただ答えを口にするのではなく、何ページものメモを書き殴り、図を描き、思考のあらゆるステップを言葉で説明してから、ようやく最終的な答えを書き留めます。この「思考を声に出す」プロセスは、技術の世界では「思考の連鎖（Chain of Thought）」と呼ばれ、AIモデルを数学や論理パズルにおいて非常に賢くしてきました。しかし、ここには落とし穴があります。そのメモのページを作成するには、膨大な時間とエネルギーが必要になるのです。それはまるで、食料品店へ行くためにタクシーを呼んでいるのに、どの通路が必要かを「考える」ためだけに、まず街の一周をドライブさせる費用まで支払っているようなものです。科学者たちは長い間、こう考えてきました。「AIに、その思考をすべて『頭の中』で行わせることはできないだろうか？」もしそれが実現すれば、AIは長いおしゃべりなメモをスキップして、即座に答えを出すことができ、膨大な時間とコンピューターのパワーを節約できるはずです。

「Masked Distillation（マスクされた蒸留）」と題されたこの論文は、まさにその問いに深く切り込んでいます。研究者たちは、大規模推論モデル（LRM）を用いて、おしゃっかいりで思考を声に出すことを好むスマートな「教師」モデルから、より小さな「生徒」モデルへと、その思考プロセスを内面化させる方法を模索しました。彼らは、生徒モデルがメモを書く工程を完全にスキップして最終的な答えだけを出力しながらも、依然として同等の賢さを維持できるかどうかを検証したかったのです。彼らはこれを「Masked Distillation」と呼ばれる巧妙なテクニックを用いてテストしました。教師がパズルを解いている間、生徒がそれを観察している様子を想像してください。教師は問題全体とその詳細な思考プロセスを見ていますが、生徒には問題だけが見せられています。目標は、生徒が教師の隠れたロジックを学習し、たとえ学習過程で教師が詳細なメモを持っていたとしても、メモを一切書かずに自力でパズルを解けるようになるかどうかを確認することでした。

チームは主に2つの実験を設定しました。第1の実験では、同じAIモデルを教師と生徒の両方に使い、単に思考のさせ方を変えてもらいました。第2の実験では、より大きく賢いモデルを教師とし、より小さく単純なモデルを生徒としました。また、「スキャフォールド（足場）」という概念も導入しました。これは自転車の補助輪のようなものです。時には、生徒がゼロから取り組むのではなく、教師のメモを少しだけ見ることで正解にたどり着けるかもしれません。彼らは、生徒が最高の成果を得るために、教師の思考のどの程度を見る必要があるのかをテストしました。

結果は、「驚き」と「そうはいかない」が入り混じったものでした。小学校レベルの算数の問題セットであるGSM8Kにおいて、生徒モデルは驚異的な成果を上げました。彼らは教師の思考を内面化することに成功したのです。完全に「マスクされた（思考を一切書き出さない）」生徒モデルは、正解率76.0%を記録しました。これは初期の能力からの大幅な向上であり、教師が平均2,398.1トークン（テキストの単位）を使用していたのに対し、わずか369.6トークンで回答を実現しました。これは、労力の6.5倍もの削減です！

しかし、よりトリッキーな数字パズルである「Countdown」では、物語が変わりました。ここでは、頭の中ですべてを行おうとした（完全にマスクされた）生徒モデルは、実際には精度が低下し、41.7%に落ち込みました。この特定のタイプのパズルにおいては、生徒が正解するためには教師のメモをいくらか見る必要があることが判明したのです。しかし、ここで面白いことがあります。研究者たちは「スイートスポット（最適解）」を見つけ出しました。生徒に教師の思考のほんの一部（約30%、あるいは「アルファ」値0.3）を見せることで、精度が86.2%へと急上昇したのです。これは教師の精度（87.3%）に肉薄する数値であり、かつ教師よりも約1.3倍少ないトークン数で済みました。それはまるで、完璧にバランスを取るために、自転車本体ではなく、ちょうど良い補助輪を与えられたかのようでした。

論文はまた、これらの生徒モデルが未知のパズル（分布外タスク）をどの程度扱えるかについても調査しました。算数の側面では生徒モデルは好成績を収めましたが、Countdownのパズルについては結果がやや混在していました。興味深いことに、一部のケースでは、少量のスキャフォールド（足場）を用いて訓練された生徒の方が、パズルが難しくなったり変化したりした際に、教師自身よりも優れた汎化性能を示すことがありました。

一つの重要な発見は、教え方そのものが重要であるということです。生徒に教師のメモを見せて、それをそのままコピーさせるように教えようとした場合（教師あり微調整と呼ばれる手法）、それはあまりうまくいきませんでした。「Masked Distillation」の手法、つまり生徒の推測を教師の実際の思考と比較するより複雑な方法を用いる方が、思考スキルを転移させる上ではるかに効果的でした。

結論として、この論文は、すべての問題に対して思考ステップを魔法のように消去できるわけではないことを示唆しています。AIが「静かに思考」できるかどうかは、問題の種類や、そのトピックについて生徒モデルがすでにどれほど知っているかに大きく依存します。基本的な算数のような馴染みのあるタスクであれば、AIは静かに思考することを学習し、大量の時間を節約できます。より難解で未知のパズルの場合は、依然として少しの助け——つまり、メモによる小さなスキャフォールド——が必要になるかもしれません。研究者たちは、この「スキャフォールド」は調整可能な強力なツールであり、AIがいかに賢いか、そしていかに速く答えを出せるかの間の完璧なバランスを見つけることができると結論付けています。

技術要約：Masked Distillation（マスク付き蒸留）– 言語モデルにおける思考の連鎖（Chain-of-Thought）の内在化

問題提起

大規模推論モデル（LRM）は、複雑なタスク（数学、コード、プランニング）において高い性能を発揮するが、これは最終的な回答を生成する前に、長い明示的な中間トークンの連鎖（Chain-of-Thought: CoT）を生成することによって実現されている。これらのトレースは精度を向上させる一方で、推論のレイテンシ、メモリ使用量（KVキャッシュのフットプリント）、およびサービングのエネルギー消費に大きなコストをもたらし、これらはトレースの長さに比例して線形に増加する。

近年の研究では、これらの中間トレースの正当性と最終回答の正当性は因果関係的に結びついておらず、またトレースの長さが問題の複雑さと信頼性のある相関関係にあるわけでもないことが示唆されている。ここから重要な問いが生じる：これらの中間トークンに表現されている計算を言語モデルのパラメータ内に内在化させることが可能か？ つまり、推論時に明示的な推論ステップを生成することなく、直接（あるいは大幅に短いトレースで）正しい回答を出力できるようにできるだろうか。

手法：Masked Distillation（マスク付き蒸留）

著者らは、「思考する」教師モデルの推論プロセスを「思考しない」生徒モデルへと内在化させるために設計された知識蒸留フレームワークである、Masked Distillationを提案している。

コアフレームワーク

このフレームワークにおいて、生徒モデルは入力質問に条件付けられた解のトークンのみを予測するように訓練される。一方、教師モデルは、質問とその生成したCoTトレースの両方に条件付けられたフィードバックを提供する。目的関数は、生徒の分布（入力のみに条件付けられたもの）と、教師の分布（入力および中間推論に条件付けられたもの）の間のダイバージェンスを最小化することである。

訓練目的関数には、**オンポリシーの逆KLダイバージェンス（reverse-KL divergence）**が使用される：
$L_{MD} = D_{KL}(\pi_S(\cdot | x, y_{<t}) \parallel \pi_T(\cdot | x, \text{ITs}, y_{<t}))$
ここで、 $\pi_S$ は生徒のポリシー、 $\pi_T$ は教師のポリシー、 $x$ は入力、 $y_{<t}$ は以前に生成されたトークン、 $\text{ITs}$ は教師の中間トークンを表す。

実験設定

著者らは、このフレームワークを以下の2つの設定で実装している：

自己蒸留（Self-Distillation）: 同一のモデルが、思考モードのときは教師として、非思考モードのときは生徒として機能する。
デュアルモデル設定（Dual-Model Setting）: より大規模な推論モデルが教師として機能し、別のより小さな非思考生徒モデルを監督する。

$\alpha$ -Suffix Scaffold（ $\alpha$ -サフィックス・スキャフォールド）

内部化と推論コストのトレードオフを探索するため、著者らは $\alpha$ -suffix masked distillationというバリアントを導入している。

$\alpha = 0$ (完全マスク/Fully Masked): 生徒は中間トレース全体を内在化し、解のみを出力する。
$\alpha = 1$ (非マスク/Non-Masked): 生徒は回答の前にフルトレースを再現する（標準的なCoT生成と同等）。
$0 < \alpha < 1$ : 生徒は教師のトレースの $(1-\alpha)$ の部分（プレフィックス）を内在化し、推論時には残りの $\alpha$ 割合のサフィックス（接尾辞）と解を出力するように訓練される。

これにより、中間トークンを、部分的に保持するか、あるいは完全に内在化させるべき「足場（scaffold）」として扱うことができる。

主な貢献

フレームワークの提案: CoT推論を言語モデルのパラメータに内在化させ、直接的な回答生成を可能にするMasked Distillationを導入した。
スキャフォールド分析: $\alpha$ -suffixパラメータを定義および評価し、精度と推論コストのトレードオフを調整するための制御可能な軸を確立した。
包括的な評価: 以下の2つのドメインにおいて広範な実験を実施した：
- 数学: GSM8K（イン・ディストリビューション）、MATH-500、AIME-25（アウト・オブ・ディストリビューション）。
- Countdown: 数値パズル探索タスクであり、ターゲット範囲や探索深度のバリエーションを含む。
アブレーション研究: 教師のトレースを用いた教師あり微調整（SFT）および、明示的な「ティーチング・トークン」をスキャフォールドとして用いる手法と比較した。

結果

1. 完全な内在化のタスク依存性

推論を完全に内在化できる能力（ $\alpha=0$ ）は、タスクおよび事前学習中のドメインへの生徒の事前露出度（prior exposure）に強く依存する：

GSM8K (数学): 完全な内在化は成功する。自己蒸留の設定において、生徒は76.0%の精度（教師の87.9%に対し）を達成しつつ、推論トークン数を6.5倍削減した。デュアルモデル設定では、精度はベースモデルの15.69%から45.26%に向上し、5倍のトークン削減を実現した。
Countdown: スキャフォールドなしでは、完全な内在化は失敗する。自己蒸留において、完全マスクされた生徒は**41.7%**の精度に低下した（ベースモデルの56.2%を下回る）。デュアルモデル設定では、0%から34.1%に向上したが、教師（87.3%）や非マスク版（81.1%）には遠く及ばなかった。

2. $\alpha$ -Suffix Scaffoldingの有効性

少量のスキャフォールド（トレースの一部を出力すること）を提供することで、特に完全な内在化が失敗するタスクにおいて、性能が大幅に回復する：

Countdown (自己蒸留): $\alpha = 0.3$ に設定することで、精度は41.7%から86.2%に上昇し、実質的に教師（87.3%）に匹敵する結果となった。これは非マスク版と比較して、推論トークン数を約1.3倍削減している。
Countdown (デュアルモデル): $\alpha = 0.5$ のバリアントは**84.2%**の精度を達成し、非マスクのベースライン（81.1%）を上回りつつ、より少ないトークンを使用しており、パレート最適点を示した。
GSM8K: 小規模なスキャフォールド（ $\alpha=0.3$ ）は、精度の差の大部分を回復させた（83.2% vs 87.9%）。これは非マスク版の約半分のトークンコストであった。

3. アウト・オブ・ディストリビューション (OOD) への汎化

ターゲット範囲のシフト: 様々なスキャフォールド・パラダイムで訓練されたモデルは、ターゲットとなる数値範囲が変化したOODタスクに対してクリーンに転移した。
探索深度のシフト: Countdownの自己蒸留設定において、スキャフォールド訓練された生徒（特に $\alpha=0.5$ ）は、探索深度が増大した場合のOODシナリオにおいて、非マスク版よりも優れた汎化性能を示した（32.8% vs 25.7%）。
デュアルモデルの限界: より困難な数学ベンチマーク（MATH-500, AIME-25）におけるデュアルモデル設定では、教師と生徒の容量差（capacity gap）が大きすぎた。その結果、完全な内在化もスキャフォールドもOOD性能を回復できず、すべてのバリアントがベースとなる生徒の底辺値へと崩壊した。

4. 訓練目的関数

Reverse-KL vs. SFT: 本論文では、教師のトレースを用いた**教師あり微調整（SFT）**は、スキャフォールドを効果的に活用できないことが示されている。Countdownにおいて、スキャフォールドを用いたSFTは約38%の精度であったのに対し、同じスキャフォールドを用いたReverse-KLによるMasked Distillationは約67%に達した。
ティーチング・トークン: 明示的な「ティーチング・トークン」をスキャフォールドとして提供することは、完全マスクよりも性能を向上させたが、 $\alpha$ -suffixアプローチほど効果的ではなかった。

意義と主張

本論文は、**サフィックス・スキャフォールディング（suffix scaffolding）**が、精度と推論効率のバランスを取るための制御可能なメカニズムを提供すると主張している。最適な動作点は、特定のタスクおよび教師と生徒の間の容量差によって決定される。

内在化は可能だが普遍的ではない: モデルは、自身が既知のドメイン（GSM8Kなど）については推論を内在化できるが、新しい探索タスク（Countdownなど）においては、明示的なトレースを保持せずに進めることは困難である。
コスト・パフォーマンスのトレードオフ: 少量の明示的な推論（例： $\alpha=0.3$ ）を提供することで、フルCoT生成の性能にほぼ限りなく近づきつつ、大幅な推論コストの削減（トークン数の1.3倍から2倍の削減）が可能となる。
汎化: スキャフォールディングは、特定のOODシナリオ（自己蒸留における探索深度のシフトなど）において汎化性能を向上させることができる。これは、部分的なトレースを保持することが、特定の文脈において、完全に内在化されたモデルやフルトレースを用いるモデルよりも、複雑な探索空間をナビゲートするのに役立つことを示唆している。

著者らは、完全な内在化は効率化のための魅力的な目標であるが、 $\alpha$ -suffixアプローチは、高い性能を維持しながらレイテンシを低減できる、実用的で調整可能な代替案であると結論付けている。