ロジット蒸留（ロジットジョウリュウ）

ロジット蒸留とは

ロジット蒸留（Logit Distillation / Logit-based Knowledge Distillation）は、教師モデルのソフトマックス関数適用前の生の出力値（ロジット）を直接使って生徒モデルを学習させる蒸留手法です。Hinton の元論文ではソフトマックス後の確率分布を使いましたが、ロジットを直接使うことで温度パラメータの調整が不要になり、より豊富な情報を効率的に転移できるケースがあります。

ロジットとソフトターゲットの違い

ソフトマックス関数は入力ロジットを確率分布に変換しますが、この変換で情報が失われます。特に、非常に小さな確率のクラス（ロジット値が大きな負の値）の情報は、ソフトマックス後にはほぼ0に潰されてしまいます。

主要な手法

標準ロジットマッチング

教師と生徒のロジットベクトル間のMSE（平均二乗誤差）を最小化します。実装が最もシンプルで、温度パラメータの調整が不要です。

Scaled Logit Distillation

ロジットのスケール（絶対値の大きさ）を正規化してからマッチングする手法です。教師と生徒のモデルサイズが大きく異なる場合、ロジットのスケールも異なるため、正規化が精度向上に寄与します。

DKD（Decoupled Knowledge Distillation）

ロジットを「ターゲットクラスのロジット（正解に関する知識）」と「非ターゲットクラスのロジット（クラス間関係の知識）」に分離して、それぞれ独立に蒸留する手法です。2022年のCVPR論文で提案され、従来のKD損失を大幅に上回る性能を示しました。

LLM での適用

LLM のロジットは語彙サイズ（32,000〜128,000次元）の巨大なベクトルです。全ロジットを保存・転送するコストが膨大なため、以下の工夫が用いられます。

Top-K ロジット蒸留

教師モデルの上位K個（K=50-100）のロジットのみを使用します。重要度の低い語彙のロジットは無視することで、ストレージとメモリを大幅に節約します。

トークンレベル蒸留

LLM の自己回帰生成では、各トークン位置で語彙サイズ分のロジットが出力されます。全位置・全語彙のロジットを保存すると膨大なストレージが必要になるため、サンプリング戦略やストリーミング蒸留が用いられます。

ソフトターゲット蒸留との比較実験

ロジットベースの手法は特にクラス数が多いタスクで優位性を発揮します。

実装上の注意点

数値安定性

ロジットの絶対値が非常に大きくなる場合（100以上）、MSE損失が爆発する可能性があります。ロジットの標準化（z-score normalization）やクリッピングで安定化を図ります。

メモリ効率

LLM の場合、語彙サイズ128Kのロジットを FP32 で保持すると 1トークンあたり512KBです。FP16やBFloat16で保存するか、Top-K フィルタリングでメモリを削減します。

FAQ

Q1: ロジット蒸留とソフトターゲット蒸留はどちらが良いですか？

A1: タスクとモデルサイズに依存しますが、ロジット蒸留（特にDKD）が多くのベンチマークで優位です。ロジット蒸留は温度パラメータの調整が不要で再現性が高く、特にクラス数が多いタスク（ImageNet、LLMの語彙予測）で差が大きくなります。

Q2: ロジットにアクセスできないAPIモデル（GPT-4等）でもロジット蒸留は可能ですか？

A2: いいえ。ロジット蒸留には教師モデルの生ロジットベクトルが必要です。APIモデルは通常、最終的なテキスト出力かTop-K log probability のみを返すため、フルロジット蒸留は不可能です。APIモデルからの蒸留にはResponse-based手法を使います。

Q3: DKDのデカップリングとは具体的に何を分離していますか？

A3: DKDはKLダイバージェンスを2つの独立な項に分解します。TCKD（Target Class KD）は正解クラスの確率のみに注目し、NCKD（Non-Target Class KD）は不正解クラス間の関係性に注目します。それぞれに独立のハイパーパラメータ（α, β）を設定できるため、タスクに応じた最適な蒸留が可能です。