最適輸送入門：Wasserstein距離で分布間の距離を測る

はじめに

2つの分布がどれだけ離れているかを測りたい場面は多いです。生成モデルの出力と本物のデータ、今月と先月の顧客分布、ドメイン適応での訓練元とテスト先。まず思い浮かぶのはKLダイバージェンスですが、KLには弱点があります。2つの分布の台（サポート）が重ならないと発散してしまい、少し離れると値が意味をなさなくなります。

今回、2つの分布を少しずつ引き離しながら、KL・JS・Wassersteinの3種類で距離を測りました。離すほどに、3つの性格の違いがはっきり出ます。KLは分布の重なりが消えると無限大へ発散してしまい、JSはある値で頭打ちになって「どれだけ離れているか」を伝えなくなります。対してWassersteinだけが、離れた分だけ大きくなりました。この記事では、最適輸送という考え方からWasserstein距離を導入し、なぜ分布の幾何を反映できるのか、そしてSinkhornアルゴリズムでどう高速に近似するのかを、実装で確かめます。

対象読者:

• 合成データや生成モデルの分布評価に関わる方
• 分布間距離としてKL/JSを使ってきて、その限界を感じている方
• 線形計画・最適化の応用としての最適輸送に関心がある方

記事のポイント:

• 最適輸送問題（Monge/Kantorovich）を線形計画として理解します
• Wasserstein距離がKL/JSと何が違うか（台が重ならなくても有限で、幾何を反映する）を確かめます
• Sinkhorn（エントロピー正則化）で高速に近似する仕組みと、精度と速度のトレードオフを見ます

---

なぜKL/JSでは足りないか

KLダイバージェンス $\mathrm{KL}(P\|Q)=\int p\log(p/q)$ は、$p>0$ なのに $q=0$ となる点があると無限大になります。台が食い違った瞬間に発散するということで、生成モデルのように出力の台が本物とずれやすい場面では、そのままでは扱いにくいです。しかも非対称で、$\mathrm{KL}(P\|Q)\neq\mathrm{KL}(Q\|P)$ です。

JSダイバージェンスは対称化して有限に収まるよう改良した指標ですが、別の問題があります。2つの分布が離れて重なりが消えると、JSは上限のlog2に飽和して動かなくなります。値が一定になるということは、そこから「もっと近づけたい」ときの勾配が消えるということでもあります。生成モデルの学習でJSベースの目的関数が不安定になりやすいのは、この飽和が一因です。

やりたいのは、台が重ならなくても有限で、しかも「どれだけ離れているか」という幾何をそのまま反映する距離です。土の山を別の場所へ運ぶコスト、という比喩からEarth Mover’s Distanceとも呼ばれるWasserstein距離がこれにあたり、輸送コストそのもので分布間の距離を定義します。

最適輸送問題

分布 $P$ を分布 $Q$ の形に「運ぶ」ことを考えます。各地点 $x_i$ から $y_j$ へどれだけの質量を運ぶかを表す輸送計画 $P_{ij}$ を決め、運ぶ距離に応じたコスト $C_{ij}$（たとえば $|x_i-y_j|$）の総和を最小化します。

$$W(P,Q) = \min_{\gamma \geq 0} \sum_{i,j} C_{ij}\, \gamma_{ij} \quad \text{s.t.} \quad \sum_j \gamma_{ij} = a_i,\ \ \sum_i \gamma_{ij} = b_j$$

制約は、各地点から運び出す量が元の質量 $a_i$ に、各地点へ運び込む量が目標の質量 $b_j$ に一致することです。これは輸送計画 $\gamma$ を変数とする線形計画そのもので、Kantorovichの緩和と呼ばれます。台が重ならなくても、質量を運ぶ距離という形で有限の値が定まり、離れれば離れるほどコストも増えていきます。これがKL/JSとの決定的な違いです。

Sinkhornアルゴリズム

厳密な線形計画は、地点数 $n$ に対して $n^2$ 個の変数を持ち、$n$ が大きいと重くなります。そこで目的関数にエントロピー項を足して問題を強凸化すると、行と列を交互にスケーリングするだけの単純な反復で解けます。これがSinkhornアルゴリズムです。

def sinkhorn(a, b, C, eps, iters=2000):
    K = np.exp(-C / eps)          # ギブスカーネル
    u = np.ones_like(a)
    for _ in range(iters):
        v = b / (K.T @ u)         # 列方向スケーリング
        u = a / (K @ v)           # 行方向スケーリング
    P = u[:, None] * K * v[None, :]
    return float(np.sum(P * C)), P

正則化の強さ $\varepsilon$ が挙動を決めます。小さいほど厳密解に近づきますが、$K=\exp(-C/\varepsilon)$ が数値的に潰れて不安定になりやすいです。大きいほど計算は安定して速いですが、エントロピー正則化のバイアスで真の距離より大きめに出ます。

距離を実装して比べる

前半は、2つのガウス分布（標準偏差0.7、等分散）を0から3まで引き離しながら、KL・JS・W1をグリッド上で計算します。後半は、離散化した2つの分布に対して、scipy.optimize.linprogで厳密な最適輸送を解き、Sinkhornと突き合わせます。乱数シードは42に固定しました。

分布が離れるにつれて、3つの距離の性格の違いがはっきり出ます。

分布同士の重なりが消えていく様子を並べると、JSが飽和し、KLが発散する理由が直感的に分かります。ずれが3のとき、2つのガウスはほとんど重なっていません。

数値で並べると次のようになります。W1がずれと一致して線形に伸びるのに対し、JSはlog2（0.693）に漸近して差を伝えなくなります。

ずれ	KL	JS	Wasserstein-1
0	0.00	0.000	0.00
1	1.02	0.205	1.00
2	4.08	0.505	2.00
3	9.18	0.650	3.00

後半、離散分布（平均0と2）での厳密なWasserstein距離は2.0000でした。Sinkhornを $\varepsilon$ を変えて走らせると、精度と速度のトレードオフがそのまま数字に出ます。

手法	コスト	厳密解との誤差	時間
厳密LP	2.0000	—	40.0 ms
Sinkhorn ε=0.02	2.0000	0.0000	12.2 ms
Sinkhorn ε=0.1	2.0002	0.0002	8.2 ms
Sinkhorn ε=0.3	2.0037	0.0037	7.1 ms
Sinkhorn ε=1.0	2.0200	0.0200	7.4 ms

三つの距離の使いどころ

Wasserstein距離だけが、ずれに対して線形に、一定の傾きで増えました。傾きが一定ということは、最適化の勾配情報として使えます。JSが0.65で飽和して勾配を失うのとは対照的で、Wasserstein GANがJSベースのGANより学習が安定しやすい理由がここにあります。KLは発散気味に膨らみ、値そのものは大きいですが解釈しづらいです。

Sinkhornは $\varepsilon$ を小さくすれば厳密解にほぼ一致しました（ε=0.02で誤差0.0000）。$\varepsilon$ を大きくすると誤差が0.02まで増える一方、反復は安定します。エントロピー正則化は必ず真の距離を上振れさせるバイアスを持つので、厳密さが要るなら $\varepsilon$ を下げ、規模が大きく速度が要るなら少し上げる、という使い分けになります。バイアスが気になる用途では、Sinkhorn divergenceで補正する手もあります。

速度の面はどうでしょうか。今回の $n=60$ では厳密LPも40msで済み、差はまだ小さいです。効いてくるのは規模で、厳密LPは変数が $n^2$ に膨らむため、地点数が数百を超えると急速に重くなります。Sinkhornは1反復あたり行列積だけなのでスケールしやすく、GPUにも載ります。実務で分布の忠実度を大量に測るなら、Sinkhornを既定に置くのが現実的です。

実務で気をつける点は二つあります。標本からWassersteinを推定する場合、高次元では次元の呪いでサンプル効率が悪く、必要標本数が急増します。コスト行列の設計（どの距離計量を使うか）が結果を左右する点も見落とされやすいです。合成データの品質評価で分布の類似度を指標化するときは、この2点、つまり次元と計量の選び方を最初に押さえておくと、指標が安定します。

まとめ

分布間の距離は、輸送コストとして定義すると台が重ならなくても有限に収まり、離れ具合という幾何を反映します。今回の実験では、ずれ3に対してWasserstein-1が3.00と線形に伸びる一方、JSはlog2で飽和し、KLは9.18まで膨らみました。生成モデルの評価、ドメイン適応、合成データの忠実度指標といった場面で、まずはSinkhornを小さめの $\varepsilon$ から試し、必要に応じて厳密LPと突き合わせて $\varepsilon$ を詰めていくとよいでしょう。次の一歩としては、実データの2群でWasserstein距離を測り、KLで見ていたときと結論が変わるかを確かめてみることをお勧めします。