最適輸送 — 分布を運ぶ数学

2つの確率分布の距離を「砂を運ぶ最小コスト」として測るのが最適輸送。Wasserstein距離の定義から、実用の主役であるSinkhornアルゴリズム、WGAN・FID・Flow Matchingへの接続、そして現場で踏む落とし穴までを比喩とnumpyコードで積み上げる。

Sinkhorn Distances: Lightspeed Computation of Optimal Transport

分布どうしの距離で、素直に困る場面

機械学習の中身は、多くの場合「モデルの分布を、データの分布に近づける」作業です。近づけるには、まず「いまどれくらい離れているか」を1つの数字にしなければ始まりません。その定番がKL情報量ですが、KLには決定的に苦手な場面があります。2つの分布が重なっていないときです。

本物のデータは $x=0$ にだけ立つ細い針、モデルは $x=\theta$ にだけ立つ細い針だとしましょう。$\theta$ が $0.01$ でも $100$ でも、「重なりがゼロ」という一点においてKLは無限大です（Jensen–Shannonダイバージェンスなら $\log 2$ で一定）。つまり $\theta$ を動かしても値が変わらない＝勾配がゼロ。人間の感覚では $\theta=0.01$ の方が明らかに近いのに、その差が数字に出ません。直感が使っているのは「針が立っている場所」ですが、KLは各点の確率を縦に比べるだけで、横方向の距離を見ていないからです。

最適輸送（Optimal Transport）はここを埋めます。片方の分布を、もう片方の形になるまで「運んだ」ときの最小コストを距離と定義する。運ぶという言葉どおり、横方向の距離が値の中に入ります。

比喩: 引っ越し業者の見積もり

あなたは土建屋で、いくつかの土場に砂が山積みになっていて、いくつかの現場がそれぞれ決まった量の砂を欲しがっています。運賃は「運ぶ量 × 運ぶ距離」。全部を運び終えるやり方は無数にありますが、総額が最小になる運び方がひとつあり、その最小額が砂の配置と現場の需要のあいだの「距離」です。

これは比喩であると同時に文字通りの起源でもあります。1781年、ガスパール・モンジュが「土砂の切り取りと盛り土」の問題として定式化したのが最適輸送の始まりでした。

直感: 決めるのは「対応」ではなく「輸送表」

モンジュの定式化には弱点がありました。「土場 $i$ の砂は丸ごと現場 $\sigma(i)$ へ」という1対1の対応を探すので、土場が3つ・現場が2つならそんな対応は作れません。1942年にレオニート・カントロヴィチが与えた解決は身も蓋もないほど単純で、砂を分割してよいことにする。すると決めるべきものは次のような表になります。

表 $P$ の $(i,j)$ 成分が「土場 $i$ から現場 $j$ へ運ぶ量」です。制約は3つだけ。すべて非負、行の合計が各土場の在庫、列の合計が各現場の需要。この条件を満たす表の集合を輸送多面体と呼びます。

総運賃は「運ぶ量 × 単価」を全マス足したものなので、目的関数も制約も表の成分について線形です。つまり最適輸送は線形計画そのものになります。分割を許した瞬間に解の存在が保証され、既知の道具箱が丸ごと使えるようになった——これがカントロヴィチの緩和が革命的だった理由です。

仕組み: Wasserstein距離の定義

以上を1本の式にまとめます。

$\mu,\nu$ が比べたい2つの分布、$x_i$ と $y_j$ がそれぞれの点の位置、$a_i$ と $b_j$ がそこに乗っている確率（＝砂の量）です。$U(a,b)$ はさきほどの輸送多面体、$d$ は地面の上の距離（普通はユークリッド距離）。$p$ は運賃の効き方で、$p=1$ なら距離に比例、$p=2$ なら距離の2乗、つまり遠くへ運ぶことを強く嫌います。外側の $1/p$ 乗は単位を「長さ」に戻すためです。

平たく言えばこの式は、「運び方の表をすべて考え、総運賃が最小のものを選ぶ。その最小運賃が距離である」と言っているだけです。$p=1$ の場合を特にEarth Mover's Distance（土を動かす距離）と呼びます。名前がそのまま説明になっています。

総当たりは、なぜ無理なのか

素朴には「全部の運び方を試せばいい」と思えます。しかし $n$ 個の点を $n$ 個へ1対1で割り当てるだけでも $n!$ 通りあり、$n=20$ で約 $2.4 \times 10^{18}$ 通りです。

線形計画として厳密に解く専用アルゴリズム（ネットワークシンプレックス法、オークション法など）を使っても、点が $n$ 個ならおおむね $n^3$ に対数がつく程度の手間がかかります。数千点までなら現実的ですが、学習の1ステップごとにミニバッチ同士を比べたい生成モデルには重すぎます。

1次元だけは、並べ替えるだけで解ける

ところが例外がひとつあります。点が数直線の上に並んでいる場合、最適な運び方は考えるまでもなく、小さい方から順に対応させるのが常に最適です。2本の運搬経路が交差（追い越し）していたら、交差をほどいた方が必ず総距離は短くなるからです。

$$W_p^p(\mu,\nu)=\int_0^1 \left|F^{-1}(u)-G^{-1}(u)\right|^p du$$

$F^{-1}$ と $G^{-1}$ は2つの分布の分位点関数（下から $u$ 割の位置にある値）です。式は「同じ分位点どうしを突き合わせて差を測り、平均する」と言っています。同じ個数のサンプルなら、両方をソートして順に引き算するだけで $O(n\log n)$。Pythonなら scipy.stats.wasserstein_distance が1行でやってくれます。

この「1次元だけは激安」という事実は後で効きます。高次元の点群をランダムな向きに何本も射影し、各方向の1次元Wasserstein距離を平均する——これが実用されている近似、Sliced Wasserstein距離です。