KL Divergence Loss（ケーエルダイバージェンスロス）
Kullback-Leibler ダイバージェンスに基づく損失関数で、2 つの確率分布間の差異を測定する。LLM では知識蒸留（教師モデルの出力分布を生徒モデルに転写）や、RLHF における方策制約（元のモデルからの逸脱防止）に広く使われる。非対称性を持ち KL(P||Q) ≠ KL(Q||P) である点が特徴。
KL Divergence Loss とは
KL Divergence（Kullback-Leibler ダイバージェンス）は、確率分布 P と Q の間の「情報量の差」を測定する非対称な距離尺度である。P を真の分布、Q を近似分布として：
D_KL(P || Q) = Σ P(x) * log(P(x) / Q(x))
LLM の文脈では、知識蒸留（Distillation）と RLHF/DPO のアライメント制約の 2 つの主要用途がある。2024-2026 年の LLM 開発では、大規模モデルの軽量化（蒸留）とアライメント制約として不可欠な損失関数となっている。
主な用途と特徴
1. 知識蒸留（Knowledge Distillation）
教師モデル（大規模）の出力確率分布を生徒モデル（小規模）に学習させる：
L_distill = α * T^2 * D_KL(P_teacher/T || P_student/T) + (1-α) * L_CE
- T（温度）: 2-4 が標準。高温でソフトラベルの情報量増加
- α: 蒸留損失と Cross-Entropy 損失のバランス。0.5-0.9
- Llama 3.2 1B/3B: Llama 3.1 70B からの蒸留で KL Loss 使用
- Gemma 2 2B: Gemma 2 27B からの蒸留
- Phi-3-mini: GPT-4 の出力を教師として蒸留
| 蒸留ペア | 教師 | 生徒 | 性能保持率 | パラメータ削減 |
|---|---|---|---|---|
| Llama 3.1 70B → 8B | 70B | 8B | 92% | 88.6% |
| Gemma 2 27B → 2B | 27B | 2B | 85% | 92.6% |
| GPT-4 → Phi-3 | 非公開 | 3.8B | 80% | 99%+ |
| Mistral Large → 7B | 123B | 7B | 88% | 94.3% |
2. RLHF/DPO の KL 制約
アライメント時にモデルが元の SFT モデルから大きく逸脱しないよう KL ペナルティを課す：
R_total = R_reward - β * D_KL(π_θ || π_ref)
- β: KL ペナルティ係数。PPO では 0.01-0.1、DPO では 0.1-0.5
- π_ref: 参照方策（SFT モデル）
- KL 制約なしだとモデルが報酬ハッキング（reward hacking）に陥り、品質が急激に劣化
3. VAE の正則化
変分オートエンコーダーの潜在空間を標準正規分布に近づける：
L_VAE = L_reconstruction + β * D_KL(q(z|x) || p(z))
- β-VAE: β>1 で latent の解きほぐし（disentanglement）促進
- テキスト VAE: Optimus、DELLA で自然言語生成に応用
KL Divergence の性質
- 非対称性: D_KL(P||Q) ≠ D_KL(Q||P)。蒸留では教師→生徒の方向（Forward KL）が標準
- Forward KL（mode-covering）: 真の分布 P がゼロでない箇所を Q もカバーするよう学習。平均的な振る舞いを再現
- Reverse KL（mode-seeking）: Q が高確率を割り当てる箇所で P もゼロでないことを保証。鋭い分布を学習
- 非負性: D_KL(P||Q) ≥ 0、等号は P=Q のときのみ
- 三角不等式を満たさない: 距離の公理を満たさないため「ダイバージェンス」と呼ぶ
実装のポイント
- PyTorch:
F.kl_div(log_q, p, reduction='batchmean')— 入力は log 確率と確率 - 数値安定性: P(x)=0 の項は 0*log(0)=0 と定義（L'Hopital の定理）
- 温度スケーリング: logits を T で割ってからソフトマックス、損失に T^2 を掛ける（勾配スケールの補正）
- メモリ効率: 語彙サイズ 128K の LLM では教師・生徒双方の logits をメモリに保持する必要があり、チャンク処理が必要
代替・関連する距離尺度
| 距離尺度 | 対称性 | 用途 | 特徴 |
|---|---|---|---|
| KL Divergence | 非対称 | 蒸留, RLHF | 最も標準的 |
| Jensen-Shannon | 対称 | GAN | KL の対称化版 |
| Wasserstein | 対称 | W-GAN | 分布の形状を考慮 |
| Total Variation | 対称 | 理論解析 | 最大確率差 |
| f-Divergence | 非対称 | 一般化フレーム | KL, χ^2 等を統一 |
よくある質問（FAQ）
Q1: 蒸留で Forward KL と Reverse KL のどちらを使うべきですか？ A: 標準は Forward KL（D_KL(P_teacher || P_student)）。教師の出力分布全体をカバーする mode-covering な学習となり、生徒が教師の多様な出力パターンを学習できる。Reverse KL は生徒が自信のある領域に特化する mode-seeking 学習で、特定タスクへの特化が目的なら有効。
Q2: KL ペナルティの β はどう調整しますか？ A: PPO では β=0.01-0.1 が標準。β が小さいとアライメントは効くが reward hacking リスク増加、大きいと SFT モデルから動かず学習が進まない。動的 β 調整（KL 値のターゲットに応じて β を増減）が InstructGPT で提案され、実践的に有効。
Q3: KL Divergence が無限大になる場合の対処法は？ A: Q(x)=0 かつ P(x)>0 の箇所で発生。対処として label smoothing（Q に微小な一様分布を混合）、またはクリッピング（Q の最小値を 1e-8 に制限）を適用する。蒸留では温度 T を上げることで教師の出力が滑らかになり、ゼロ確率問題が緩和される。
まとめ
- KL Divergence は LLM の蒸留とアライメント制約の両方で中核的な役割
- 蒸留では温度 T=2-4、Forward KL が標準
- RLHF/DPO では β=0.01-0.5 で方策逸脱を制約
- 非対称性（Forward vs Reverse）の選択がモデルの振る舞いを決定
- 数値安定性のため label smoothing やクリッピングが必須