KL・JSダイバージェンス
KLダイバージェンス(Kullback-Leibler Divergence)は、2 つの確率分布の差異を計る尺度です(divergence は「ずれ・乖離」の意味で、値が無限大に発散することとは別の意味です)。
分布 $p$ に従うデータを、分布 $q$ に最適化した符号で符号化したときに生じる、余分な平均ビット数を表します。
KLダイバージェンスは以下の式で定義されます。
\[\begin{aligned} D_{\text{KL}}(p \parallel q) = \sum_x p(x) \log_2 \frac{p(x)}{q(x)} \quad \text{離散確率変数} \\ D_{\text{KL}}(p \parallel q) = \int p(x) \log_2 \frac{p(x)}{q(x)}\,dx \quad \text{連続確率変数} \end{aligned}\]- $p(x)=0$ の項は $0\log\frac{0}{q(x)}:=0$ と約束します
- $p(x)>0$ かつ $q(x)=0$ となる $x$ があると、$D_{\text{KL}}(p\parallel q)=\infty$ になります
- 常に $D_{\text{KL}}(p\parallel q)\ge0$ で、等号は $p=q$ のときに限ります(ギブスの不等式。情報量 を参照)
交差エントロピー $H(p, q)$ を使用して以下のように分解することができます。
\[\begin{aligned} H(p, q) = H(p) + D_{\text{KL}}(p \parallel q) \\ D_{\text{KL}}(p \parallel q) = H(p,q) - H(p) \geq 0 \end{aligned}\]つまり、$p$ の符号を使わず $q$ の符号を使ったことによる平均符号長の増加分が KLダイバージェンスです。
NOTE
本ファイルの冒頭の定義は情報理論の慣習に合わせて $\log_2$ を使っています。「深層学習 の KLダイバージェンス」の節以降は、実装に合わせて底を明示せず $\log$ と書き、自然対数を想定します。底を変えても、すべての項で同じ底を使えば式の関係は変わりません。
深層学習 の KLダイバージェンス
KLダイバージェンス $D_{\text{KL}}(p \parallel q)$ は、基準である真の分布 $p$ をモデルの分布 $q$ に置き換えることで、平均としてどれだけ余分な情報量(符号長)が必要になるかを計算する指標です。
真の分布 $p$ が固定されている 深層学習 においては、$H(p)$ は $q$(モデルのパラメータ)に依存しない定数です。したがって、交差エントロピーを最小化することは、KLダイバージェンスを最小化すること(予測分布を正解分布に近づけること)と同義になります。
\[\underset{q}{\arg\min}\ H(p,q) = \underset{q}{\arg\min}\ \bigl(H(p) + D_{\text{KL}}(p\parallel q)\bigr) = \underset{q}{\arg\min}\ D_{\text{KL}}(p\parallel q)\]- $p(x)$(真の分布):教師データ(正解ラベル)の分布
- $q(x)$(モデルの分布):モデルが出力した予測(Softmax等を通した確率)の分布
$H(p)$ は $q$ に依存しない定数なので、交差エントロピー $H(p,q)$ と KLダイバージェンス $D_{\text{KL}}(p\parallel q)$ の勾配は同じです。
\[\frac{\partial H(p,q)}{\partial z_i} = \frac{\partial D_{\text{KL}}(p\parallel q)}{\partial z_i} = q_i - p_i\]softmax に入る前のスコア(logits)$z_i$ に対する勾配は、どちらの損失でも次のようになります($p$ が one-hot でもソフトターゲットでも同様です)。
\[\frac{\partial L}{\partial z_i} = q_i - p_i\]$- \log q_t(k)$ の導出
交差エントロピー $H(p_t, q_t)$ は、情報理論において以下のように分解できます。
\[H(p_t, q_t) = H(p_t) + D_{\text{KL}}(p_t \parallel q_t)\]- $p_t$ : 教師データ(正解ラベル)の分布
- $q_t$ : モデルが出力した予測(Softmax等)の分布
- $H(p_t)$ : 教師データ自体が持つエントロピー(不確実性)
- $D_{\text{KL}}(p_t \parallel q_t)$ : 2つの分布の差異を表す KLダイバージェンス
深層学習 の目的は、予測分布 $q_t$ を正解分布 $p_t$ に近づける(=KLダイバージェンスを最小化する)ことです。
正解ラベルが one-hot ベクトルの場合、数式はシンプルになります。正解トークンのインデックスを $k$ とします。
1. 正解データのエントロピー $H(p_t)$ が 0 になる
正解が完全に固定されている(不確実性がゼロ)ため、教師データ自体のエントロピーは $0$ になります。
ここで分布の次元数は語彙数 $V$ とします。$p_t(i)=0$ の項は $0\log 0:=0$ の約束により $0$ とします。
これを分解式に代入すると、以下の関係が成り立ちます。
\[H(p_t, q_t) = 0 + D_{\text{KL}}(p_t \parallel q_t) = D_{\text{KL}}(p_t \parallel q_t)\]one-hot の場合は、交差エントロピーの値そのものが、KLダイバージェンスの値と一致します。
2. 計算が正解トークンだけの一項に絞られる
交差エントロピーの定義式に one-hot ベクトル の性質を当てはめます。
\[L_t = - \sum_{i=1}^{V} p_t(i) \log q_t(i)\]$i \neq k$ のときは $p_t(i) = 0$ となるため、正解インデックス $k$ 以外の項はすべて消滅します。
\[\begin{aligned} L_t &= - \Big( 0 \cdot \log q_t(1) + \dots + 1 \cdot \log q_t(k) + \dots + 0 \cdot \log q_t(V) \Big) \\ &= - \log q_t(k) \end{aligned}\]正解が one-hot でないケース(1 項への簡約が使えない場合)
正解ラベルが one-hot ではない( 100% の正解が 1 つとは限らない)タスクでは、$-\log q_t(k)$ の 1 項への簡約は使えません。損失は、全項の和
\[L_t = - \sum_{i=1}^{V} p_t(i) \log q_t(i) = H(p_t) + D_{\text{KL}}(p_t\parallel q_t)\]としてそのまま計算します。ここで $H(p_t)$ は $q_t$ に依存しない正の定数なので、この損失を最小化することと、$D_{\text{KL}}(p_t\parallel q_t)$ を最小化することは同じです(最小化する $q_t$ も勾配も同じです)。
NOTE
「ソフトターゲットでは KLダイバージェンスが必須」というわけではありません。交差エントロピー(全項の和)で足ります。実装で KLダイバージェンス(例:PyTorch のKLDivLoss)を使うことが多いのは、最適時に損失が $0$ になり、ターゲットとの一致度をログで見やすいなどの理由によるものです。
1. 知識蒸留(Knowledge Distillation)
- 概要: 巨大なモデル(親)の持つ知識を、軽量なモデル(子)に継承させる技術。
- ターゲット: 子モデルは、one-hot ではなく、親モデルが出力した
[犬: 0.7, 猫: 0.2, 狐: 0.1]というなだらかな確率分布(ソフトターゲット)に近づくように学習します。 - 温度 $T$: 実際の蒸留では、親と子の logits を温度 $T$ で割った softmax で軟化した分布 $p^{(T)}$, $q^{(T)}$ を使うのが一般的です。勾配の大きさを揃えるため、損失に $T^2$ を掛けます。
2. ラベルスムージング(Label Smoothing)
- 概要: 過学習を防ぐために、正解を
1.0ではなく0.9に落とし、残りの0.1を他のクラスに均等に配る(例:[0.9, 0.05, 0.05])正則化テクニック - ターゲット: スムージング後の分布 $p_t$ を正解分布として使い、損失は $-\sum_i p_t(i)\log q_t(i)$ をそのまま計算します。$p_t$ には $0$ でない値が複数含まれるので、1 項には簡約できません
- このときの $H(p_t)$ は $0$ ではなく正の定数です(上の例では約 0.394 nat)が、$q_t$ には依存しないため、最小化には影響しません
要約
- 通常の one-hot 学習: $H(p_t)=0$ なので、交差エントロピーは KLダイバージェンスと一致し、数式では正解トークンの確率を最大化する( $q_{t}(k) = 1$ に最大化すると $- \log q_t(k)$ は最小化)と同じになる
- 正解が確率分布の学習(蒸留・ラベルスムージング): 1 項への簡約は使えず、全項の和をそのまま計算する。$H(p_t)$ は定数なので、交差エントロピーの最小化と KLダイバージェンスの最小化は同じ
JSダイバージェンス(Jensen-Shannon Divergence)
KLダイバージェンスの弱点:非対称性
これまで見てきたKLダイバージェンス $D_{\text{KL}}(p\parallel q)$ には、以下の不便な性質があります。
\[D_{\text{KL}}(p\parallel q) \neq D_{\text{KL}}(q\parallel p) \quad (\text{一般に})\]「$p$ を基準に $q$ を評価する」場合と「$q$ を基準に $p$ を評価する」場合とで値が異なるため、KLダイバージェンスは距離(distance)とは呼べません(距離であれば $d(p,q)=d(q,p)$ という対称性が必要です)。三角不等式も一般に満たしません。
また、$q(x)=0$ となる点で $p(x)>0$ だと $\log\frac{p(x)}{q(x)}\to\infty$ となり、無限大になってしまう問題もあります。
向きによって、最小化したときの $q$ の振る舞いも異なります。
- $D_{\text{KL}}(p\parallel q)$ の最小化:$p(x)>0$ の場所で $q(x)$ が $0$ になると無限大になるため、$q$ は $p$ の全体を覆うように広がります
- $D_{\text{KL}}(q\parallel p)$ の最小化:$q$ は $p$ が小さい場所を避け、$p$ の 1 つの山に集中しやすくなります
JSダイバージェンスの定義
この非対称性を解消するために考案されたのがJSダイバージェンスです。$p$ と $q$ の「中間分布」
\[m(x) = \frac{1}{2}\bigl(p(x)+q(x)\bigr)\]を経由し、$p$ と $m$、$q$ と $m$ のKLダイバージェンスの平均として定義します。
\[D_{\text{JS}}(p \parallel q) = \frac{1}{2} D_{\text{KL}}(p \parallel m) + \frac{1}{2} D_{\text{KL}}(q \parallel m)\]エントロピーを使うと、次のようにも書けます。
\[D_{\text{JS}}(p \parallel q) = H(m) - \frac{1}{2}\bigl(H(p) + H(q)\bigr)\]$H$ は凹関数なので、$H(m)\ge\frac{1}{2}(H(p)+H(q))$ より $D_{\text{JS}}\ge0$ が直接分かります。
値域
- $D_{\text{JS}}(p\parallel q)=0$ となるのは $p=q$ のときに限ります
- 上限は $\log 2$(底が 2 なら $1$ ビット、自然対数なら $\ln 2$)で、$p$ と $q$ の支持(確率が正の領域)が重ならないときに達成されます
- 有界である理由:$m(x)\ge\frac{1}{2}p(x)$ より $\frac{p(x)}{m(x)}\le2$ なので、$D_{\text{KL}}(p\parallel m)\le\log2$ です。$D_{\text{KL}}(q\parallel m)$ も同様です
KLダイバージェンスに対する利点
| 性質 | KLダイバージェンス | JSダイバージェンス |
|---|---|---|
| 対称性 | $D_{\text{KL}}(p\parallel q)\neq D_{\text{KL}}(q\parallel p)$ | $D_{\text{JS}}(p\parallel q) = D_{\text{JS}}(q\parallel p)$(定義から明らか) |
| 値の範囲 | $[0,\infty]$(無限大になりうる) | $[0,\log 2]$($\log_2$ なら $[0,1]$、自然対数なら $[0,\ln2]$)に有界 |
| $q(x)=0,\ p(x)>0$ での挙動 | 無限大になる | $m(x)=\frac{1}{2}p(x)>0$ となるため有限 |
さらに $\sqrt{D_{\text{JS}}(p\parallel q)}$ は距離の公理(三角不等式を含む)を満たすことが知られており、これをJensen-Shannon距離と呼びます(Endres & Schindelin, 2003 / Österreicher & Vajda, 2003 など)。KLダイバージェンスが「距離もどき」に留まるのに対し、JSダイバージェンスの平方根は数学的に正式な距離になる、という違いがあります。
機械学習での利用例:GAN
敵対的生成ネットワーク(GAN)の元論文(Goodfellow et al., 2014)では、識別器 $D$ を最適にしたとき、目的関数が「本物のデータ分布 $p_{\text{data}}$」と「生成分布 $p_g$」の間のJSダイバージェンスで表せることが示されています(自然対数を使っており、$D_{\text{JS}}$ も自然対数で測った値です)。
\[\max_D V(D,G) = -\log 4 + 2\,D_{\text{JS}}(p_{\text{data}} \parallel p_g)\]したがって、最適な $D$ のもとで生成器 $G$ に関して最小化することは、$D_{\text{JS}}(p_{\text{data}}\parallel p_g)$ を最小化することと等価です。生成器の学習が進み $p_g \to p_{\text{data}}$ となるほど $D_{\text{JS}}\to0$ に近づく、という形で、「2つの分布の近さ」を測る尺度としてJSダイバージェンスが理論的な支柱になっています。
NOTE
これは識別器が最適であるという理想化のもとでの理論的な等価性です。実際の学習では $D$ は最適にならず、生成器の損失にも non-saturating 版などが使われます。
JSダイバージェンスの限界
$p_{\text{data}}$ と $p_g$ の支持が重ならない場合、$D_{\text{JS}}=\log 2$ で一定になります。このとき、$p_g$ のパラメータを少し動かしても値が変わらないため、勾配が $0$ になり、生成器が学習の方向を得られません。
高次元空間では、データ分布と生成分布の支持がほとんど重ならないことがよくあるため、これが GAN の学習が不安定になる原因の 1 つとされています。Wasserstein GAN は、この問題を動機の 1 つとして提案されました。
まとめ(KL・JSダイバージェンス)
| KLダイバージェンス | JSダイバージェンス | |
|---|---|---|
| 対称性 | なし | あり |
| 値域 | $[0,\infty]$ | $[0,\log2]$(有界) |
| 距離の公理 | 満たさない | 本体は満たさない(平方根は距離) |
| 主な用途 | 交差エントロピー誤差(分類の損失関数)、知識蒸留、ラベルスムージング | GANの理論的な目的関数、2分布間の対称な近さの評価 |
| 限界 | 支持が合わないと無限大 | 支持が重ならないと定数になり勾配が消える |