NVIDIA、X-Tokenを発表：投影誘導型クロストークナイザ知識蒸留でLlama-3.2-1BにおいてGOLDを平均3.82ポイント上回る

知識蒸留（KD）は、大規模教師モデルから小規模学生モデルへ「暗黙の知識」を転送する手法です。標準的なKDでは共通のトークナイザが必要であり、Phi-4-miniやQwen3-4Bなどの互換性のないトークナイザを持つ強力な教師を活用できません。NVIDIAの研究者は、ロジット分布ベースのクロストークナイザKD手法であるX-Tokenを導入しました。これは、補助的な訓練可能コンポーネントやアーキテクチャ変更を必要とせず、標準KD損失のドロップイン代替として機能します。

X-Tokenが解決する問題：

最初の問題はGOLDにおける「珍しいトークン失敗」です。トークナイザがテキストを異なる方法で分割すると、重要なトークンが一致しない珍しいサブセットに分類され、有害な勾配が生じます。例えば、Llamaは「201」を単一トークンとしてエンコードするのに対し、Qwen3は「2」「0」「1」に分割します。GOLDでは、すべての複数桁トークンが珍しいセットに分類され、確率が抑制されます。X-Tokenは投影行列Wを介して学生分布を教師語彙空間に投影することで、このノイズを回避します。

2番目の問題は「過度に保守的なマッチング」です。GOLDは厳密な文字列一致を使用して共通サブセットを定義するため、「Hundreds」と「Hund」などのほぼ等価なトークンペアを破棄します。X-TokenのH-KL損失は、投影行列に基づく緩和されたマッチングにより共通セットを拡大し、より多くのアライメント信号を保持します。

X-Tokenは3つのコンポーネントで構成されます：スパンアライメント、投影行列W、および2つの補完的損失（P-KLとH-KL）。スパンアライメントは動的計画法を使用して、同じテキスト部分文字列にデコードされるトークンをチャンクにグループ化します。投影行列Wは、正規化後の正確一致とマルチトークンルール（減衰重み）により構築され、学生トークンを教師トークンにマッピングします。P-KLはパーティションを完全に除去し、投影分布上で直接KLダイバージェンスを計算して有害な勾配を排除します。H-KLはGOLDのハイブリッド構造を維持しますが、共通セットを拡大し、パーティションが構造的に健全な場合に適用されます。

実験では、Llama-3.2-1Bを学生、Qwen3-4BおよびPhi-4-mini-Instructを教師として使用しました。Qwen3-4B（P-KLモード）では、X-Tokenの平均スコアは38.85で、GOLD（35.03）を3.82上回り、GSM8kは2.56から15.54に跳ね上がり、同一トークナイザKD（12.89）さえ凌ぎました。Phi-mini（H-KLモード）では、H-KLは39.18を達成し、GOLDを0.52上回りました。マルチ教師蒸留（Phi-mini + Llama-3B）は平均40.48に達しました。動的KD/CEスケーリングや静的ウェイトによるマルチ教師蒸留の設定も効果的でした。

X-Tokenは、追加の訓練可能パラメータを必要とせず、複数のベンチマークで最先端の結果を達成する、クロストークナイザKDの効果的なソリューションを提供します。この研究は、モデル選択、推論コスト、パフォーマンス向上に重要な影響を与える可能性があります。