（機械学習）確率分布間の距離・ダイバージェンス

機械学習や深層学習では、2つの確率分布間の差異を表すために、様々な距離・ダイバージェンス・情報量が使われている。様々な尺度があるが、ここではそれらの中から個人的に見かけたことのある尺度に絞って紹介する。
（エントロピー（情報量）周りは次回予定）

※以降の定義で連続確率と離散確率の2パターンがある場合、離散確率版の式のみを提示する

目次

• KLダイバージェンス (KLD: Kullback-Leibler divergence)
• JSダイバージェンス (JSD: Jensen-Shannon divergence)
• Hellinger距離 (Hellinger distance)
• バタチャリヤ距離 (Bhattacharyya distance)
• Wasserstein距離 (EMD: Earth mover’s distance)
• MMD (Maximum Mean Discrepancy)
• マハラノビス距離 (Mahalanobis distance)

KLダイバージェンス (KLD: Kullback-Leibler divergence)

\begin{align}
D_{KL}(P\parallel Q) &= \sum_iP(i)\log \frac{P(i)}{Q(i)} \quad ∈ [0,∞)\\
 &=H(P,Q)-H(P) \\ % \quad⋯\ H(P,Q)\text{はcross entropy} \\
 & \ \ \ ※H(P,Q)\text{はcross entropy}、H(P)\text{はentropy}
\end{align}

確率分布Pから見た確率分布Qの非重なり度合いを示し、以下の特徴を持つ。

• 非対称性（PとQを入れ替えて計算すると異なる値になる）等の性質があり、距離の定義を満たさない。
  • → 距離でなくダイバージェンスと呼ばれる。
• 重なり度合いのみを見るため、分布間のユークリッド距離的な離れ具合は考慮されない。
• $P(i)>0\ かつ\ Q(i)=0$ なケースは上手く扱えない。
  （極小値εで0除算を避けても、数値不安定化は避けられない）
• 単位は情報量（bits, nats等）

深層学習周りでよく使われており、以下のような例がある。

• （例1）DNNのクラス分類タスク
  • DNNの推論分布と正解分布（one-hot）間のクロスエントロピーがロスとして用いられる。このロスの最小化はKLD最小化と同等である。
    （正解分布PのエントロピーH(P)は固定値なので）
• （例2）生成モデル（GANやDiffusion model等）の生成タスク
  • 究極的には、学習データ分布に対する生成データ分布の尤度最大化を目指している。これは学習データ分布と生成データ分布間のKLD最小化を意味する。

Forward KLと Reverse KL

PとQを入れ替えると異なる結果になると説明したが、それらは以下のように区別されることがある。

• $D_{KL}(P\parallel Q)=\text{Forward KL}$
• $D_{KL}(Q\parallel P)=\text{Reverse KL}$

特に生成モデルにおいては、これらの違いが生成データの多様性などに影響してくる。
（こちらでもう少し詳細な解説がされている）

JSダイバージェンス (JSD: Jensen-Shannon divergence)

\begin{align}
D_{JS}(P \parallel Q) &= \frac{1}{2} D_{KL}(P \parallel \frac{P+Q}{2}) + \frac{1}{2} D_{KL}(Q \parallel \frac{P+Q}{2}) \quad ∈ [0,\log(2)] \\
 &= H(\frac{P+Q}{2}) - \frac{H(P)+H(Q)}{2} \\
 % &= I(X,Z) \\
 & ※\frac{P+Q}{2}は混合確率分布
\end{align}

KLDを対称化したような尺度になっている。また、混合確率分布を取ることでKLDより数値安定性が向上している。

実例としては、「生成モデルのGAN学習ではJSDの最小化問題を解いている」というのが有名かもしれない。（真の目的であるKLD最小化の近似としてだが）

（参考）二乗根を取ると距離になる： $\text{JS distance} = \sqrt{D_{JS}(P \parallel Q)}$

Hellinger距離 (Hellinger distance)

\begin{align}
D_H(P,Q) &= \sqrt{\frac{\sum{ \left( \sqrt{P(i)}-\sqrt{Q(i)} \right) ^2}}{2}} \quad ∈ [0,1] \\
&=\sqrt{1-\text{BC}(P,Q)} \qquad \text{where} \quad \text{BC}(P,Q)=\sum{\sqrt{P(i)Q(i)}} \\
& \quad※\text{BC}はバタチャリヤ係数
\end{align}

KLDとHellinger距離はf-divergenceと呼ばれる尺度で一般化でき、KLDの亜種と言える。KLDはlogの差（log P/Q=log P - log Q）をベースにしているが、Hellinger距離は二乗根の差をベースにしているという点が異なる。
このようにlogの代わりに二乗根にすることで0除算リスク（log 0リスク）が無くなるため、Qが0確率を含んでいても問題なく扱える。（数値不安定化リスクなし）
あと、名前から分かるように対称性などを有した距離である。

バタチャリヤ距離 (Bhattacharyya distance)

\begin{align}
D_B &= -\ln{\text{BC}(P,Q)} \quad ∈ [0,∞) \qquad \text{where}  \quad \text{BC}(P,Q)=\sum{\sqrt{P(i)Q(i)}} \\
& \quad※\text{BC}はバタチャリヤ係数
\end{align}

外側でlog、内側で二乗根を取るため、数値安定性はKLDとHellinger距離の中間ぐらい。
そして、こんな名前なのに距離でない。（対称性はあるけど三角不等式を満たさない）

参考までに、KLDとHellinger距離とは以下の関係性を持つ。
$$D_{KL} \ge 2 D_B \ge 2D_H^2$$

Wasserstein距離 (EMD: Earth mover’s distance)

\begin{align}
D_{W1}(P,Q)&=\text{EMD}(P,Q) \\
&= \inf _{γ∈\text{Π}(P,Q)} \mathbb E_{(x,y) \sim γ} [d(x,y)] \\
&= \min_{f∈F} \sum d_{ij}f_{ij} \quad ∈ [0,∞)\\
&※i,jは位置、d_{ij}はi・j間の距離、f_{ij}はi・j間で移送する荷物量\\
&※\min_{f∈F}は最小の総移送量のケース\\
&※P,Qが確率分布であるため、分母\text{（＝1）}は省略
\end{align}

Wasserstein距離は、確率分布のグラフを荷物の山と見立てたときの、PからQへの最小の総移送量（$\sum$荷物量×移送距離）として解釈される。これまで説明してきたすべての尺度は分布の重なり度合いしか見ないため、分布間がユークリッド距離的な意味合いで「どれぐらい離れているか」を表現できなかったが、Wasserstein距離はそれを表現できる。

下記の図では完全に重複しない3つの一様分布を描画している。そして、左の確率分布P（青）を基準として、確率分布Q1（赤）と確率分布Q2（緑）へのWasserstein距離・Hellinger距離を表記した。この2つの距離を比較すると、Hellinger距離はx軸上の距離が異なるどちらの分布に対しても1.0と同じ値を取る一方で、Wasserstein距離はQ1に対して10.0、Q2に対して20.0となっており、x軸の離れ具合を表せていることがわかる。

p-Wasserstein距離

Wasserstein距離は任意のp次モーメントを選ぶことができ、以下のように一般化される。

D_{Wp}(P,Q)= \min_{f∈F} \left(\sum d_{ij}^p f_{ij}\right)^{\frac{1}{p}}

特に2-Wasserstein距離は、GANの発展形であるWasserstein GANや生成タスクの重要な評価指標であるFID (Fréchet inception distance)で採用されている。
（Fréchet-Wasserstein距離や確率分布間のFréchet distanceと呼ばれることもある）

MMD (Maximum Mean Discrepancy)

\begin{align}
\text{MMD}^2_{emp}(P, Q) &= \frac{1}{n^2} \sum\sum k(x_i,x_j) + \frac{1}{m^2} \sum\sum k(y_i,y_j) - \frac{2}{nm} \sum\sum k(x_i, y_j) \quad ∈ [0,∞)\\
& ※ ここでは、kは\text{RBF}カーネル: k(x,y)=\exp⁡(−γ∥x−y∥^2) \\
& ※ x\sim P,\ y\sim Q\ だが、ここでは経験分布からのサンプルとして扱っている
\end{align}

MMDはカーネルを使った尺度であり、以下の特徴を持つ。

• カーネルを使うので高次元への耐性がある
• カーネル選びが必要になる
• Wasserstein距離と同様に分布間が「どれぐらい離れているか」を表すことができる
• 名前に"distance"は入ってないがMMDは距離

マハラノビス距離 (Mahalanobis distance)

\begin{align}
\text{（点・点）}\quad D_M(\boldsymbol{a},\boldsymbol{b};P)&=\sqrt{(\boldsymbol{a}-\boldsymbol{b})^T Σ_P^{-1} (\boldsymbol{a}-\boldsymbol{b})} \\
\text{（点・確率分布）}\ D_M(\boldsymbol{a};P)&=\sqrt{(\boldsymbol{a}-\boldsymbol{μ}_P)^T Σ_P^{-1} (\boldsymbol{a}-\boldsymbol{μ}_P)} \\
\text{（確率分布・確率分布）} D_M(P,Q)&=\sqrt{(\boldsymbol{μ}_P-\boldsymbol{μ}_Q)^T \left(\frac{Σ_P+Σ_Q}{2}\right)^{-1} (\boldsymbol{μ}_P-\boldsymbol{μ}_Q)} \quad ∈ [0,∞)\\
\end{align}

マハラノビス距離は確率分布の共分散行列で正規化＋相関補正されたユークリッド距離により定義されている。マハラノビス距離は点・点または点・確率分布間の距離だが、上式3個目のように確率分布・確率分布間の距離に応用できなくもない。これも分布間が「どれぐらい離れているか」を表せる。
（正規分布にある程度近い分布を想定できない場合は適さないが）