Logit蒸留とは、教師モデルのsoftmax出力層のlogit(対数確率)を生徒モデルが模倣する知識蒸留の最も基本的な手法で、KL-divergence損失を用いて教師の出力確率分布全体を生徒に転移する。
Logit蒸留は知識蒸留の最も直接的な手法であり、教師モデルのsoftmax出力層の生のlogit値(softmax適用前の対数確率)を生徒モデルが模倣する。教師の最終層出力z_teacherに温度Tを適用したsoftmax σ(z_teacher/T)と、生徒の対応する出力σ(z_student/T)の間のKL-divergenceを最小化することで知識を転移する。
Logit蒸留の損失関数は以下のように定式化される:
T²の補正項は、温度Tで割ることにより勾配が1/T²倍にスケールダウンされるのを補償する。
| 手法 | 粒度 | メモリ使用量 | 精度 | 計算コスト | 適用場面 |
|---|---|---|---|---|---|
| Word-Level KD | 各トークン独立 | 低(top-K保存) | 中〜高 | 低い | 大規模LLM(70B+)の蒸留 |
| Sequence-Level KD | 文全体の確率 | 高(全系列保存) | 最高 | 非常に高い | 高品質な小型モデル生成 |
| Top-K Logit KD | 上位K個のlogit | 最低 | 中 | 最低 | メモリ制約環境 |
| On-Policy KD | 生徒の出力で蒸留 |
| 中 |
| 高 |
| 高い |
| 分布シフト対策 |
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クラス提供 |
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に載せればよい。