本論文では、知識蒸留プロセスに因果分析に基づく妨害識別メカニズムを組み込んだLeaFフレームワークを提案します。これにより、生徒モデルが推論中に因果的に重要な特徴に焦点を当てるよう導かれ、推論精度と汎化能力が向上します。
論文タイトル：
Learning to Focus: Causal Attention Distillation via Gradient-Guided Token Pruning
著者所属：
中国人民大学高瓴人工知能学院、清華大学計算機科学部
論文リンク：https://arxiv.org/pdf/2506.07851
コードリンク：
問題の背景
大規模言語モデル（LLMs）は自然言語処理タスクにおいて強力な文脈理解と言語生成能力を示しますが、長文推論や複雑な指示タスクにおいては依然として顕著な欠点があり、特に重要情報に焦点を当てる能力が劣っています。この注意散漫現象は、モデルの推論精度と生成品質を深刻に制約します。
この現象を体系的に研究するため、本研究ではまず教師モデルと生徒モデルの勾配感度を比較することで、入力中の妨害パターン（distracting patterns）を特定し、NuminaMath-CoTおよびAceCode-87Kデータセットで生徒モデルの性能を評価しました。
図1と図2に示すように、これらの妨害情報を単に剪定するだけで、平均精度が大幅に向上しました。数学訓練セットでは20%以上、コード訓練セットでは10%以上の向上です。
さらに、AMC_AIMEなどのより複雑なタスクでは、モデルの性能向上はGSM8Kよりも顕著であり、複雑な推論タスクにはモデルの有効な判断を妨げる誤解を招く要因がより多く含まれていることを示しています。
▲ 図1：コードタスクの精度向上
▲ 図2：数学タスクの精度向上
これらの発見は、妨害情報を排除し、モデルが重要情報に自律的に焦点を当てる能力を高めることが、大規模言語モデルの推論性能を向上させる上で重要な道であることを示しています。
この目的のために、著者はLeaFフレームワークを提案します。因果的な視点から、勾配誘導を利用して入力中の妨害要因を特定し除去することで、蒸留プロセス中に生徒モデルが重要な情報領域に焦点を当てるよう学習させ、モデルの推論性能を向上させます。
実験結果は、LeaFが数学推論とコード生成を含む複数の下流タスクで顕著な性能向上を達成したことを示しています。GSM8K、MATH、OlympiadBenchなどのデータセットでは平均精度が2.41%向上し、HumanEval+、LeetCode、LivecodeBenchなどのコードタスクでは平均2.48%向上しました。
さらに、推論プロセス中のモデルのアテンション分布はより集中し、一貫性が高まり、アテンションの可視化結果も本手法の解釈可能性をさらに検証しています。
LeaF：2段階モデリング、モデルのアテンションの因果性を向上
モデルが推論中に妨害情報に誤導されやすく、重要情報に焦点を当てにくいという問題を緩和するため、著者は因果駆動型のアテンション転移手法、すなわちLeaF（Learning to Focus）フレームワークを提案しました。このフレームワークは2つの主要な段階で構成されます。
妨害情報識別：勾配でモデルの注意バイアスを特徴付ける
第一段階は、入力中で生徒モデルを誤導するが推論自体には必要ないトークン、すなわち交絡トークン（confounding tokens）を特定することを目的とします。
具体的には、生徒が誤予測し教師が正予測したサンプルから、各入力トークンに対する両者の勾配感度を比較し、生徒モデルが推論時に注目（勾配値が大きい）するが教師モデルが注目しない（勾配値が小さい）トークンを潜在的な妨害要因として選別します。
さらに、これらのトークンを削除した後、生徒モデルと教師モデルの両方が正しい予測を生成できる場合、それらを交絡トークン（confounder tokens）と判定します。これは、生徒の推論を誤導するが、最終的に正しい答えを導き出すためには必要ない情報です。
▲ 図3. LeaFフレームワーク：勾配駆動の妨害識別と因果蒸留による推論能力の最適化
交絡トークン（confounding tokens）を特定した後、LeaFは反事実入力サンプルを構築する2つの処理方法を比較しました。
一括除去（Collective Pruning）：特定されたすべての交絡トークンを一度に直接削除する；
連続セグメント除去（Span Pruning）：より精緻な方法で、毎回1つの連続する妨害セグメントのみを削除し、より多くの意味的コンテキストを保持する。
予備実験により、Span Pruningがより安定しており、より優れた選択肢であることが証明されました。
▲ 図4. 剪定戦略の図
因果蒸留：反事実比較から焦点化戦略を学習
生徒モデルがより堅牢なアテンションパターンを学習するよう効果的に導くため、LeaFは元のサンプルと反事実サンプルを構築した後、2つの教師信号を融合した混合蒸留目的を設計しました。
標準蒸留（Standard Distillation）：生徒モデルが元の入力で教師モデルと一致するように維持する；
反事実蒸留（Counterfactual Distillation）：妨害情報が削除された後の入力においても、生徒が教師モデルと一致するように促す。
この二重蒸留メカニズムは、生徒モデルが教師モデルの出力挙動に一致するよう促すだけでなく、入力中の重要トークンに対する因果的判断能力も強化します。LeaFは、意味情報と因果的依存関係を同時にモデリングすることで、生徒モデルが表面的なパターンを模倣するだけで重要な因果関係を見落とすことを効果的に防ぎ、推論の堅牢性と汎化能力を向上させます。
さらに、LeaFは、元々入力側のみに作用していた指示レベル処理（Instruction-level Pruning）を、応答レベル処理（Response-level Pruning）にまで拡張しました。
具体的には、入力指示から妨害トークンを特定し除去するだけでなく、LeaFはモデルが生成した過去の応答を文脈入力として扱い、その中から後続の推論を誤導する可能性のあるトークンを動的に特定し削除します。
この戦略は、生成プロセス中に継続的に妨害を排除し、モデルが重要情報に焦点を当てる能力をさらに向上させ、より正確で焦点を絞ったコンテンツを生成するのに役立ちます。
▲ 図5 指示レベル処理から応答レベル処理への拡張
実験評価：重要アテンションに焦点を当て、推論性能を向上
著者は、数学推論とコード生成という2つの主要なタスクでLeaFフレームワークの有効性を体系的に評価しました。関連する実験はLlamaとQwenという2つの主要なモデルアーキテクチャと6つの評価基準をカバーし、モデルの推論能力を強化する上でのLeaFの役割を検証しました。
主要実験結果
実験結果は、LeaFがすべての数学およびコーディングベンチマークタスクで性能向上をもたらし、平均精度が標準蒸留方法と比較してそれぞれ2.41%と2.48%向上したことを示しています。特に、高難度ベンチマークであるOlympiadBenchでの改善は顕著であり、LeaFが複雑な推論におけるアテンション妨害の問題に効果的に対処できることを示しています。
▲ 図6 主要実験結果
さらに、交絡トークン（confounding token）の処理範囲を入力指示（Instruction-level）からモデル生成プロセス（Response-level）に拡張することで、モデル性能が大幅に向上しました。これは、生成段階にも推論に影響を与える妨害情報が存在することを示しており、段階的な処理戦略がモデルが重要情報に焦点を当て続けるのに役立つことを示しています。
分析実験：LeaFはいかにして推論の誤導を正確に識別し回避するか
LeaFフレームワークが妨害トークンを識別し排除する有効性を体系的に評価するため、著者はマスキング戦略、応答処理方法、閾値感度、ケーススタディを含む4つの観点から詳細な分析を行い、推論の堅牢性とモデルの焦点化能力の向上におけるその性能を包括的に検証しました。
4.1 勾配マスキング戦略分析：LeaFはいかにして妨害情報を正確に識別するか？
LeaFが採用する勾配マスキング戦略の有効性を体系的に評価するため、著者はこれを2つの一般的なマスキング方法、すなわちランダムマスキングとパープレキシティ（PPL）マスキングと比較実験しました。実験はGSM8K、MATH、OlympiadBenchで実施され、基礎的なものから複雑なものまで様々な数学タスクのシナリオをカバーしました。
▲ 図7：LeaF勾配マスキング戦略分析実験結果
実験観察：
勾配マスキングは他の戦略よりも顕著に優れている
MATHおよびOlympiadBenchのような複雑な推論タスクで最適な性能を達成し、LeaFの勾配誘導メカニズムが妨害トークンを効果的に特定できることを検証しました。
ランダムマスキング戦略は効果が不安定
GSM8KおよびOlympiadBenchでは性能低下さえ引き起こし、これは意味的ガイダンスなしにトークンを盲目的に削除することが蒸留信号を破壊する可能性を示唆しています。また、データ拡張だけではモデルの推論能力を向上させるには不十分であることをさらに強調しています。
パープレキシティマスキングは簡単なタスクでのみわずかに向上
複雑なタスク（OlympiadBenchなど）では、その効果はランダムマスキングに近く、これは生徒モデル自身のトークンへの注意が偏っている可能性があり、どの情報が本当に重要であるかを正確に判断するのが難しいことを示唆しています。教師モデルを導入して比較指導を行う必要性が浮き彫りになります。
結論：複雑な推論タスクにおいて、勾配差に基づくマスキング戦略は交絡トークン（confounder token）をより正確に識別でき、LeaFフレームワークにおける「教師-生徒勾配比較メカニズム」の有効性と合理性を検証しました。
4.2 応答レベル処理戦略：生成プロセス中の妨害情報も無視できない
LeaFは、入力指示（Instruct-level）から妨害トークンを識別するだけでなく、妨害検出の範囲をモデルの生成コンテンツ（Response-level）にまで拡張し、推論プロセス全体のアテンションバイアスをカバーします。
この目的のために、著者は3つの処理戦略を比較のために設計しました。
指示レベルのみ処理：入力テキストから妨害トークンを識別・除去するだけで、モデルの生成コンテンツは処理しない。
応答レベル2セグメント処理（2段）：生成コンテンツを前後2つのセグメントに分割し、各セグメントで個別に妨害トークンを検出・除去する。
応答レベル3セグメント処理（3段）：生成コンテンツを3つの連続するセグメントに分割し、各セグメントで個別に妨害検出と処理を行う。
▲ 図8：LeaF応答レベル処理戦略実験結果
実験観察：
応答レベル処理の導入はモデルの性能を顕著に向上させます。入力のみを処理する場合と比較して、生成プロセス中に妨害要素をさらに識別・除去することで、モデルの推論精度を効果的に高めます。これは、その後の生成コンテンツもアテンションバイアスの影響を受けやすいことを示しています。
2セグメントと3セグメントの処理効果は類似しています。より粒度の細かい3セグメント処理では明確な利点は得られず、2セグメントで応答中の妨害パターンをモデルが識別し学習するのに十分であることを示しています。過度の分割は過学習のリスクを高める可能性があります。
結論：交絡トークン（Confounder tokens）は入力指示に存在するだけでなく、モデルの生成経路にもしばしば隠されています。妨害識別メカニズムを生成段階にまで拡張し、分割粒度を適切に制御することは、長文推論タスクにおけるモデルのアテンション焦点化能力と全体的な性能を向上させるのに役立ちます。
4.3 閾値感度分析：小規模モデルは妨害に対してより脆弱であり、より積極的なフィルタリングが必要
モデルの妨害トークンに対する感度を調査するため、著者はLeaFフレームワーク内で交絡トークン（confounder tokens）を識別するために使用される閾値（threshold）が最終的な推論性能に与える影響を体系的に分析しました。
実験は、異なるモデル規模（LLaMA3.2-1BとLLaMA3.2-3B）のもと、2つのレベル（指示レベルと応答レベル）で実施されました。
▲ 図9：指示レベルしきい値感度分析（MathBench）
▲ 図10：ステップレベルしきい値感度分析（MathBench）
実験観察：
● 指示レベル（Instruct-level）：
LLaMA3.2-1Bは閾値0.10で最高の性能を示しました；
LLaMA3.2-3Bは閾値0.05で最適な性能を達成しました。
● 応答レベル（Response-level）：
LLaMA3.2-1Bは閾値0.15で最高の性能を示しました；
LLaMA3.2-3Bは閾値0.10で最高の結果を達成しました。
分析と解釈：
指示レベルであろうと生成レベルであろうと、より小さなモデル（1B）はより高い閾値でより良い効果を示します。これは、元の入力において妨害トークンに対してより敏感であり、堅牢性を確保するためにより積極的なフィルタリング戦略に依存していることを示しています。
より高い閾値は、誤導的なトークンをより効果的に識別しフィルタリングし、より良い学習結果をもたらします。一方、大規模モデル（3B）はより強力な表現力と妨害耐性を備えているため、より低い閾値で理想的な性能を達成できます。
結論：モデル規模は、妨害トークンに対する耐性に影響を与えます。小規模モデルは誤導されやすく、より積極的な妨害フィルタリングのために高い閾値を使用するのに適しています。
4.4 解釈可能性ケーススタディ：モデルは本当に「重要点に焦点を当てる」ことを学んだのか？
LeaFがモデルに因果性の高いアテンションパターンを学習させることを本当に導いているのか検証するため、著者は数学タスクで代表的な推論ケースを構築し、LeaFと標準的な知識蒸留（KD）モデルの推論チェーンにおけるアテンションの違いを比較しました。
ケースタスク：すべての方程式の根が実数であるかを判断する。
▲ 図11 ケーススタディ
LeaFモデルの性能：
モデルは「real number」、「all」、「are real」などの重要情報に焦点を当てることに成功しました；
「すべての根は実数である必要がある」という制約を明確に理解し、それに基づいて合理的な推論戦略を採用しました：
x = -1を明らかな実数根として識別しました；
判別式（Discriminant）の条件を適用し、二次因子も実数解を生成することを確認しました。
推論プロセス全体は論理的に明確で合理的な判断であり、正解を導き出すことに成功しました。
KDモデルの性能：
「すべての根は実数である必要がある」という核心条件を無視しました；
変数の符号を考慮せずにAM–GM不等式を誤用し（負の数を導入する可能性あり）、最終的な解答の誤りにつながりました。
分析まとめ：
このケースは、LeaFがモデルが重要情報を識別し、合理的な推論経路を構築する能力を助けることで、「表面的なマッチング」による推論の誤りを効果的に回避することを示しています。同時に、LeaFは精度を向上させるだけでなく、モデルの挙動の解釈可能性と合理性も向上させることを証明しています。
将来展望
本研究は、大規模言語モデルの因果的注意能力と推論の堅牢性を向上させる上でのLeaFフレームワークの有効性を検証し、注意バイアスを理解し緩和するための新たな道筋を提供します。教師-生徒間の勾配差分析と反事実蒸留メカニズムを導入することで、LeaFはモデルが妨害トークンを効果的に識別し回避するよう導き、真に重要な情報領域に焦点を当てることを学習させます。
今後も、さらに深く探求すべきいくつかの方向性があります。例えば、現在の実験は主に数学とコード推論タスクに焦点を当てていますが、言語理解、質問応答、マルチホップ推論など、より広範なタスクシナリオに拡張し、その汎用性とタスク横断的な堅牢性を検証することも、今後の研究の価値ある方向です。