中間プローブを用いたタスク特化型知識蒸留手法

要約
この研究では、大規模言語モデルから小さなモデルへの知識蒸留において、従来の出力ベースではなく中間表現を活用する新しい手法を提案している。従来の知識蒸留では教師モデルの出力分布を訓練信号として使用するが、推論タスクにおいては正しい情報が語彙投影の過程で失われがちだった。提案手法では、凍結された教師モデルの隠れ状態に軽量プローブを訓練し、出力ロジットではなくプローブの予測を生徒モデルの監督信号として使用する。この手法をAQuA-RAT、ARC Easy/Challenge、MMLUの4つの推論ベンチマークで評価した結果、一貫した性能改善が確認され、特に限られたデータ量での改善が顕著だった。中間表現から訓練されたプローブは教師モデル自身の出力よりもクリーンなラベルを提供し、ノイズ除去効果がある。この手法はアーキテクチャの変更を必要とせず、追加の計算コストも最小限に抑えられる。
洞察・気づき
この研究は知識蒸留の根本的な課題を解決する重要な進歩を示している。従来、教師モデルの「最終出力」を信頼できる指標として扱ってきたが、実際には推論プロセスの中間段階により価値のある情報が含まれている可能性がある。これは人間の学習プロセスにも似ており、専門家の「思考過程」を学ぶことが「最終的な答え」を覚えるよりも有効であることと類似している。技術的な観点では、この手法は既存のモデルアーキテクチャを変更することなく性能向上を実現できるため、実用性が高い。また、プローブ訓練のコストが低いことから、大規模モデルの活用における計算効率の課題にも対応している。今後、このアプローチは推論タスク以外の分野にも応用できる可能性があり、より効率的なAIシステムの開発に貢献すると考えられる。