Logit蒸留（ロジットジョウリュウ）
Logit蒸留とは、教師モデルのsoftmax出力層のlogit（対数確率）を生徒モデルが模倣する知識蒸留の最も基本的な手法で、KL-divergence損失を用いて教師の出力確率分布全体を生徒に転移する。
Logit蒸留は知識蒸留の最も直接的な手法であり、教師モデルのsoftmax出力層の生のlogit値（softmax適用前の対数確率）を生徒モデルが模倣する。教師の最終層出力z_teacherに温度Tを適用したsoftmax σ(z_teacher/T)と、生徒の対応する出力σ(z_student/T)の間のKL-divergenceを最小化することで知識を転移する。
数学的定式化
Logit蒸留の損失関数は以下のように定式化される：
- 入力: トークン列x = (x_1, ..., x_n)
- 教師の出力: p_T(y|x) = softmax(z_teacher / T) ※T: 温度パラメータ
- 生徒の出力: p_S(y|x) = softmax(z_student / T)
- 蒸留損失: L_KD = T² × KL(p_T || p_S)
- タスク損失: L_CE = CrossEntropy(p_S(y|x, T=1), y_true)
- 最終損失: L = α × L_KD + (1-α) × L_CE
T²の補正項は、温度Tで割ることにより勾配が1/T²倍にスケールダウンされるのを補償する。
Word-Level vs Sequence-Level Logit蒸留
| 手法 | 粒度 | メモリ使用量 | 精度 | 計算コスト | 適用場面 |
|---|---|---|---|---|---|
| Word-Level KD | 各トークン独立 | 低（top-K保存） | 中〜高 | 低い | 大規模LLM（70B+）の蒸留 |
| Sequence-Level KD | 文全体の確率 | 高（全系列保存） | 最高 | 非常に高い | 高品質な小型モデル生成 |
| Top-K Logit KD | 上位K個のlogit | 最低 | 中 | 最低 | メモリ制約環境 |
| On-Policy KD | 生徒の出力で蒸留 | 中 | 高 | 高い | 分布シフト対策 |
実装上の重要なポイント
Logit蒸留を実装する際の注意点は以下の通りである：
- 語彙サイズの一致: 教師と生徒でtokenizerが異なる場合、語彙マッピングが必要。Llama（32K語彙）→Phi（51K語彙）のような異なるtokenizer間ではToken Alignment技術を使用
- メモリ効率: 70Bモデルの全logitを保存すると1トークンあたり128KB（32K語彙×fp32）。バッチサイズ2048、シーケンス長2048で約537GBのディスク容量。top-256のみ保存で1/128に圧縮可能
- 温度スケジューリング: 学習初期はT=8〜10（広い探索）、後期はT=2〜4（教師の自信を重視）に段階的に下げるWarm-Cool戦略が有効
- Mixed Precision: BF16精度での蒸留がFP32と同等の結果を得られることが確認されており、A100でのスループットが約2倍に向上
Logit蒸留の最新研究（2024-2026年）
- GKD（Generalized Knowledge Distillation）: Google DeepMindが2024年に提案。生徒自身の生成結果に対して教師のlogitを適用するOn-Policy手法。従来のOff-Policy蒸留より分布シフトに強い
- MiniLLM: Microsoftが提案。逆KL-divergence KL(p_S || p_T)を使用し、生徒が低確率領域を過度に模倣することを防止。GPT-2 1.5B→120Mで従来手法比5%改善
- DistiLLM: Sunら(2024)が提案。Skewed KL-divergenceを使用し、順方向・逆方向KLの両方の利点を統合。Llama 2 13B→1.1Bで最高精度達成
- f-Distill: f-divergenceファミリーからタスクに最適な発散度を自動選択。α-divergence（α=0.5〜2.0）のグリッドサーチで最適化
主要クラウドでのLogit蒸留サポート
| プラットフォーム | 教師logitアクセス | 蒸留ツール | 備考 |
|---|---|---|---|
| Azure AI | Phi-3/Llama系のみ | Azure ML | logprobs APIで上位100まで取得可能 |
| Google Vertex AI | Gemma系のみ | Vertex Training | GKDパイプライン統合 |
| AWS SageMaker | 自前ホスト必須 | SageMaker Training | カスタムコンテナ必要 |
| Hugging Face | 全OSS対応 | transformers Trainer | KDTrainerクラス提供 |
よくある質問（FAQ）
Q1: Logit蒸留とFeature蒸留のどちらが効果的ですか？ A: Logit蒸留の方が実装が簡単で一般的に高い精度を達成する。ただし教師と生徒のアーキテクチャが大きく異なる場合（例: Transformer→RWKV）はFeature蒸留が適している。両者を組合せたMulti-Level Distillationが最も効果的だが、ハイパーパラメータ調整が複雑になる。
Q2: 商用APIモデル（GPT-4, Claude）からのLogit蒸留は可能ですか？ A: 直接的なLogit蒸留は不可能。商用APIはlogitの一部（上位5〜20トークン）しか返さず、全語彙のlogitは取得できない。代替として、APIモデルで合成データを生成し、生徒をファインチューニングする「合成データ蒸留」が実用的な手段となっている。
Q3: Logit蒸留に必要なGPUメモリはどれくらいですか？ A: 教師7B+生徒1.5Bの場合、BF16でA100 80GB×1基で実行可能。教師70B+生徒7Bの場合、A100 80GB×4基（教師をtensor parallelismで分散）が最小構成。vLLMで教師をServing化してlogitを事前計算すれば、蒸留時は生徒のみGPUに載せればよい。
まとめ
- Logit蒸留は教師のsoftmax出力分布をKL-divergenceで生徒に転移する基本手法
- 温度パラメータT=4〜8が汎用的で、Warm-Coolスケジューリングが効果的
- Word-Level KDがメモリ効率と精度のバランスに優れる
- GKD、MiniLLM等の最新手法でOn-Policy蒸留と発散度選択が進化中