最適輸送とWasserstein距離をわかりやすく:Kantorovich双対とSinkhornまで

2つの分布が「どれくらい違うか」を測りたい場面は山ほどあります。たとえば、生成モデルが作った画像の分布と本物の画像の分布。あるいは、別の工場で集めたセンサーデータと自分の工場のデータ。これらの「隔たり」をひとつの数字で表せれば、モデルを学習させたり、データを揃えたりできます。

ところが、よく使われるKLダイバージェンスには弱点があります。2つの分布の台(値を取りうる範囲)が重なっていないと、KLは無限大に発散してしまうのです。発散した瞬間、勾配は意味を失い、学習は止まります。これは「全然違う2つの分布」を扱う生成モデルでは致命的でした。

そこで登場するのが最適輸送(Optimal Transport, OT)と、そこから定義されるWasserstein距離です。これは「片方の分布を、もう片方の形になるよう、最小のコストで運ぶにはいくらかかるか」という、とても物理的な発想に基づきます。台が重ならなくても有限の値を返し、しかも分布をずらすほど滑らかに値が変わる。だから勾配が死にません。

この性質のおかげで、最適輸送は次のような場面で主役になっています。

  • 生成モデル:Wasserstein GAN(WGAN)は、まさにこのWasserstein距離を最小化して画像を生成します。拡散モデルやflow matchingも、最適輸送の視点で「最短経路で分布を移す」と捉え直せます。
  • ドメイン適応:訓練データと本番データの分布がずれているとき、最適輸送で片方をもう片方に「寄せる」ことで、モデルを作り直さずに精度を保てます。

本記事の内容

  • 最適輸送の直感(砂山を運ぶ問題)と、Monge・Kantorovichの2つの定式化
  • Wasserstein距離の定義、1次元での閉形式、Kantorovich-Rubinstein双対
  • KLとの決定的な違い(台が重ならないときの挙動)
  • 計算法:線形計画とエントロピー正則化(Sinkhornアルゴリズム)をPythonで自前実装
  • 応用:WGAN・ドメイン適応・Wasserstein重心と補間

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

最適輸送とは:砂山を運ぶ問題

まず、数式をいったん忘れて、土木工事を想像してください。ここに砂山がひとつあります。これを少し離れた場所に、別の形の砂山として作り直したい。スコップで砂をすくい、運び、目標の場所に積み直します。

このとき、コストは「運ぶ砂の量 × 運ぶ距離」で決まります。近くに運べば安く、遠くに運べば高い。たくさん運べばその分高い。では、目標の形を完成させる運び方のうち、総コストが最小になる運び方は何でしょうか。これが最適輸送の問題です。

最適輸送の概念:砂山を別の形に運ぶ

上の図では、左にある青い砂山 $\mu$ を、右のオレンジの形 $\nu$ に作り変えています。矢印は「どこからどこへ砂を運ぶか」を表します。直感的に、すぐ隣に運べる砂はそのまま使い、足りない右側へは左から運んでくる、という配分が安く済みそうです。この「最も安い運び方」を求めるのが本記事の主題です。

数学では、砂山を確率分布に置き換えます。砂の総量を1に正規化すれば、砂山は確率分布 $\mu$、目標の形は確率分布 $\nu$ とみなせます。すると「2つの分布の隔たり」を、運搬コストの最小値で測れるようになります。ここから、その最小値を厳密に定義していきましょう。

Monge問題:決定的な輸送写像

最初にこの問題を定式化したのは、18世紀のフランスの数学者モンジュ(Gaspard Monge)でした。彼の考え方はシンプルです。「ある場所 $x$ にある砂は、すべて1か所 $T(x)$ に運ぶ」と決めてしまうのです。この $T$ を輸送写像と呼びます。

イメージとしては、砂粒ひとつひとつに「行き先」が一意に割り当てられている状態です。位置 $x$ の砂はすべて $T(x)$ へ。途中で分割して別々の場所に送ることは許しません。

輸送写像 $T$ が満たすべき条件は、「運んだ結果が目標の分布 $\nu$ になる」ことです。これは押し出し測度という言葉で $T_\# \mu = \nu$ と書きます。要するに、$\mu$ に従う砂を $T$ で動かしたら、ちょうど $\nu$ の形になる、という制約です。

Monge問題は、この制約のもとでコストを最小化します。コスト関数を $c(x, y)$(位置 $x$ から $y$ へ運ぶ単価)とすると、

$$ \begin{equation} \inf_{T:\, T_\# \mu = \nu} \int c\bigl(x, T(x)\bigr)\, d\mu(x) \end{equation} $$

を解きます。$c(x,y) = \|x – y\|$ なら距離に比例するコスト、$c(x,y) = \|x-y\|^2$ なら距離の2乗のコストです。

一見すると自然な定式化ですが、Monge問題にはやっかいな弱点があります。輸送写像 $T$ は「1点 → 1点」の決定的な対応なので、質量を分割できません。たとえば、1点に集中した分布(ディラックのデルタ)を、2点に分かれた分布へ運びたいとき、1点を2つに「割って」送る写像は存在しません。このように、$T$ が存在しない、あるいは存在しても最小化問題が非凸で扱いにくい、という困りごとが起きます。

この弱点を解消したのが、次に説明するKantorovichの緩和です。「質量を分割してもよい」と許すだけで、問題が一気に素直になります。

Kantorovich緩和:輸送計画と線形計画

20世紀のソ連の数学者カントロビッチ(Leonid Kantorovich)は、Monge問題の「1点は1点へ」という縛りを外しました。「位置 $x$ の砂を、複数の行き先に分けて運んでよい」と認めたのです。

このとき主役になるのは写像 $T$ ではなく、輸送計画 $\gamma$ です。$\gamma(x, y)$ は「位置 $x$ から位置 $y$ へ運ぶ砂の量」を表します。$x$ と $y$ の両方を引数に取る点がポイントで、これは $x$ と $y$ の結合分布(同時分布)に他なりません。

Monge写像とKantorovich計画の違い

左のMongeは、各点(青)が1点(オレンジ)へ細い矢印で結ばれています。1対1の決定的な対応です。一方、右のKantorovichは、1つの青い点から複数のオレンジへ、太さの違う矢印が伸びています。矢印の太さが運ぶ量で、質量を好きな割合で分割できることを示します。この「分割の自由」が、解の存在を保証します。

輸送計画 $\gamma$ が満たすべき条件は2つです。これを周辺条件と呼びます。

$$ \begin{equation} \int \gamma(x, y)\, dy = \mu(x), \qquad \int \gamma(x, y)\, dx = \nu(y) \end{equation} $$

最初の式は「位置 $x$ から運び出す総量は、そこにある砂 $\mu(x)$ に等しい」。2番目は「位置 $y$ へ運び込む総量は、目標 $\nu(y)$ に等しい」。つまり、$\gamma$ の2つの周辺分布がそれぞれ $\mu$ と $\nu$ になる、という条件です。この条件を満たす $\gamma$ の集合を $\Pi(\mu, \nu)$(カップリングの集合)と書きます。

Kantorovich問題は、この集合の上でコストを最小化します。

$$ \begin{equation} \mathrm{OT}_c(\mu, \nu) = \inf_{\gamma \in \Pi(\mu, \nu)} \int c(x, y)\, d\gamma(x, y) \end{equation} $$

ここで重要なのは、$\gamma$ について線形だという点です。目的関数も制約も $\gamma$ の1次式。つまりこれは線形計画問題(LP)になります。線形計画は凸問題で、必ず最小解が存在し、効率的に解けます。Mongeの非凸な難しさが、緩和ひとつで消えました。

離散の場合は、もっとわかりやすくなります。$\mu$ が $n$ 点、$\nu$ が $m$ 点に質量を持つとき、輸送計画は $n \times m$ の行列 $P$ になります。$P_{ij}$ は「点 $i$ から点 $j$ へ運ぶ量」。コストも行列 $C$($C_{ij} = c(x_i, y_j)$)で表せて、最小化したいのは $\langle P, C \rangle = \sum_{ij} P_{ij} C_{ij}$ です。周辺条件は「各行の和が $a_i$、各列の和が $b_j$」になります。

ここまでで「最適輸送=コスト最小の輸送計画を探す線形計画」という枠組みができました。次は、このコストの最小値から距離を定義します。

Wasserstein距離の定義

最適輸送のコスト最小値を使うと、分布どうしの距離が定義できます。コスト関数として「距離の $p$ 乗」 $c(x,y) = \|x – y\|^p$ を選び、最後に $p$ 乗根を取ります。

$$ \begin{equation} W_p(\mu, \nu) = \left( \inf_{\gamma \in \Pi(\mu, \nu)} \int \|x – y\|^p\, d\gamma(x, y) \right)^{1/p} \end{equation} $$

これが$p$次のWasserstein距離です。$p$ 乗根を取るのは、ふつうの距離(ユークリッド距離が $\sqrt{\sum (x_i – y_i)^2}$ であるのと同じ理屈)と単位を揃えるためです。

なぜこれが「距離」と呼べるのでしょうか。実は $W_p$ は、距離の3条件(非負性、対称性、三角不等式)をすべて満たします。$\mu = \nu$ のとき、その場で動かさない計画(対角に質量を置く)がコスト0なので $W_p = 0$。逆も成り立ちます。対称性は $\|x-y\| = \|y-x\|$ から明らか。三角不等式も「2段階で運ぶより直接運ぶ方が安い」という直感から証明できます。だから $W_p$ は分布空間における正真正銘の距離なのです。

実務で特によく使うのは $p=1$(コストが距離そのもの)と $p=2$(コストが距離の2乗)です。$W_1$ はEarth Mover’s Distance(土を動かす距離)とも呼ばれ、まさに砂山の運搬コストです。$W_2$ は滑らかな性質を持ち、Wasserstein重心や流れの理論で好まれます。

定義はできましたが、一般には $\inf$ を計算するのが大変です。ところが、1次元ならこの最小化が手計算できる美しい閉形式になります。次の節で見てみましょう。

1次元の閉形式:分位点関数の差

1次元(実数直線上)の分布に限ると、最適輸送は驚くほど単純になります。結論から言うと、$W_1$ は累積分布関数(CDF)の差の面積、あるいは同じことですが分位点関数(CDFの逆関数)の差の積分で書けます。

$$ \begin{equation} W_1(\mu, \nu) = \int_{-\infty}^{\infty} \bigl| F_\mu(x) – F_\nu(x) \bigr|\, dx = \int_0^1 \bigl| F_\mu^{-1}(q) – F_\nu^{-1}(q) \bigr|\, dq \end{equation} $$

なぜこうなるか、直感で説明します。1次元では「砂を運ぶとき、順番を入れ替えると損」という事実があります。左から $q$ パーセンタイルの砂は、運んだ先でも $q$ パーセンタイルの位置に来るのが最適です。だから「同じ分位どうしを対応させる」のが最適輸送になります。$F_\mu^{-1}(q)$ は $\mu$ の $q$ 分位点、$F_\nu^{-1}(q)$ は $\nu$ の $q$ 分位点なので、その差を $q$ について足し合わせれば総運搬コストになる、というわけです。

1次元Wasserstein距離は分位点関数の差

左の図で、青と橙はそれぞれ $\mu, \nu$ の分位点関数です。横軸は分位 $q \in [0,1]$、縦軸は値。2本の線で挟まれた紫の面積が、ちょうど $W_1$ になります。右はもとの2つのヒストグラムで、平均も分散も違う2分布を比べています。「分位点の差を全部足す」という操作が、形のずれをまるごと拾っていることがわかります。

$p=2$ でも同じ発想が使えて、$W_2^2 = \int_0^1 |F_\mu^{-1}(q) – F_\nu^{-1}(q)|^2\, dq$ となります。実装上は、両方の分布から同じ個数のサンプルを取り、それぞれソートして対応づけ、差を取るだけで $W_p$ が計算できます。これは後のPython節で確かめます。

1次元の閉形式は、計算が軽いので「スライス法」(高次元データを多数の方向に射影して1次元Wを平均する近似)にも応用されます。ただし2次元以上では一般に閉形式がありません。そこで効くのが、次に説明する双対表現です。

Kantorovich-Rubinstein双対:1-Lipschitz証人

線形計画には必ず双対問題があります。$W_1$ の場合、その双対はとても美しく、しかも機械学習で直接使われています。それがKantorovich-Rubinstein双対です。

$$ \begin{equation} W_1(\mu, \nu) = \sup_{\|f\|_{\mathrm{Lip}} \le 1} \left( \mathbb{E}_{x \sim \mu}[f(x)] – \mathbb{E}_{y \sim \nu}[f(y)] \right) \end{equation} $$

ここで $\|f\|_{\mathrm{Lip}} \le 1$ は「$f$ が1-Lipschitz」、つまり $|f(x) – f(x’)| \le \|x – x’\|$ を満たす関数という意味です。傾きの絶対値が1を超えないなだらかな関数、と思ってください。

この式の読み方を説明します。私たちは関数 $f$ を1つ選び、$\mu$ の上での平均値から $\nu$ の上での平均値を引きます。$f$ を上手に選ぶと、この差を大きくできます。その差を最大にする $f$ こそが、2分布の違いを最もよく「証言」する関数です。だから $f$ を証人関数(witness function、批評家とも)と呼びます。差の最大値が $W_1$ そのものになります。

なぜ1-Lipschitzの縛りが必要なのでしょうか。もし $f$ をいくらでも急峻にできたら、$\mu$ の所で巨大な値、$\nu$ の所でうんと小さい値を取らせて差を無限大にできてしまいます。傾きを1以下に抑えることで、「分布のずれた距離に見合った分だけ」差が出るように制限しているのです。これがWasserstein距離が距離に比例して滑らかに変化する理由でもあります。

Kantorovich双対の1-Lipschitz証人関数

上のパネルは2つのCDF($F_\mu, F_\nu$)と、それに挟まれた面積(これが $W_1$)です。下のパネルが最適な証人関数 $f$ で、傾きの絶対値がちょうど1(区分的に $\pm 1$)になっています。$F_\nu$ より $F_\mu$ が大きい区間では $f$ の傾きが $-1$、逆の区間では $+1$。つまり証人は「どちら向きに分布がずれているか」を傾きの符号で語っているわけです。

この双対こそ、WGANのcritic(批評家)の正体です。WGANでは、ニューラルネットを1-Lipschitzになるよう制約(重みクリッピングや勾配ペナルティ、あるいはスペクトル正規化)したうえで $\mathbb{E}_{\mu}[f] – \mathbb{E}_{\nu}[f]$ を最大化します。これが本物分布と生成分布の $W_1$ を推定し、生成器はその $W_1$ を小さくするよう学習します。Lipschitz制約の具体的なやり方は別記事に譲りますが、ここで「なぜcriticは1-Lipschitzでなければならないか」がはっきりしたはずです。

双対表現は理論的に強力ですが、本物の威力は「KLでは扱えない状況を扱える」点にあります。次に、KLとの決定的な違いを見ます。

KLとの対比:重ならない分布でも勾配が生きる

ここまで何度か「KLは台が重ならないと発散する」と述べてきました。これを具体例で確かめましょう。

KLダイバージェンス $\mathrm{KL}(P \| Q) = \sum_x P(x) \log \frac{P(x)}{Q(x)}$ を思い出してください。もし $P(x) > 0$ なのに $Q(x) = 0$ となる点が1つでもあると、$\log \frac{P(x)}{0} = +\infty$ となり、和全体が発散します。つまり$P$ が持つ質量を、$Q$ がゼロの場所に置いた瞬間にKLは無限大です。

これは生成モデルで頻繁に起こります。本物の画像は高次元空間の薄い「多様体」の上にしか存在せず、学習初期の生成器が作る画像はまったく別の場所にいます。2つの台がきれいに重なることはまずありません。すると $\mathrm{KL} = \infty$、勾配は定義できず、学習が進みません。

Wasserstein距離とKLダイバージェンスの比較

左の図は、細い2つの分布を距離 $\theta$ だけ離した状況です。$\theta > 0$ である限り台は重なりません。右の図がポイントです。$W_1$(緑)は $\theta$ に正比例して滑らかに増え、原点で勾配を持ちます。一方、KL(橙の破線)は $\theta > 0$ のあいだずっと $+\infty$ に貼り付いたまま。$\theta$ を変えてもKLは動かないので、「どちらにどれだけ近づければよいか」という方向の情報がゼロです。勾配が完全に死んでいます。

これがWasserstein距離が生成モデルにもたらした革命の核心です。台が重ならなくても有限で、しかも滑らかな勾配を持つ。だから生成器は「あと少し右に寄せればコストが下がる」という連続的な手がかりを得られます。WGANが学習を安定させられたのは、まさにこの性質のおかげです。

理論の美しさはわかりました。では、高次元の実データで実際に $W$ を計算するにはどうすればよいでしょうか。次は計算アルゴリズムに入ります。

計算法1:厳密な線形計画

離散分布どうしの最適輸送は、前述のとおり線形計画です。$n$ 点の $\mu$(質量ベクトル $a$)と $m$ 点の $\nu$(質量ベクトル $b$)、コスト行列 $C \in \mathbb{R}^{n \times m}$ が与えられたとき、

$$ \begin{equation} \min_{P \ge 0}\ \langle P, C \rangle \quad \text{s.t.}\quad P \mathbf{1}_m = a,\ \ P^\top \mathbf{1}_n = b \end{equation} $$

を解きます。$P \mathbf{1}_m = a$ は「各行の和が $a$」(運び出す量の制約)、$P^\top \mathbf{1}_n = b$ は「各列の和が $b$」(運び込む量の制約)です。変数は $P$ の $n \times m$ 個。これを汎用の線形計画ソルバー(scipyのlinprogなど)に渡せば厳密解が得られます。

最適輸送計画行列のヒートマップ

上は、$\mu$ が左寄り・$\nu$ が右寄りという1次元配置で厳密に解いた輸送計画 $P$ のヒートマップです。明るいほど多く運びます。質量が対角に近い帯に集中しているのがわかります。これは「近い位置どうしを対応させる」最適輸送の性質そのものです。順序を入れ替えて遠くへ運ぶと損なので、解は自然に「順番を保つ」配分になります。

厳密な線形計画は正確ですが、変数の数が $n \times m$ なので、点の数が増えると一気に重くなります。たとえば $n = m = 10000$ なら変数は1億個。これでは大規模な機械学習に使えません。そこで考案されたのが、エントロピー正則化による高速近似です。

計算法2:エントロピー正則化とSinkhornアルゴリズム

エントロピー正則化のアイデアはこうです。目的関数に「輸送計画 $P$ のエントロピー」を罰則として足します。

$$ \begin{equation} \min_{P \in \Pi(a,b)}\ \langle P, C \rangle – \varepsilon\, H(P), \qquad H(P) = -\sum_{ij} P_{ij}\bigl(\log P_{ij} – 1\bigr) \end{equation} $$

$\varepsilon > 0$ は正則化の強さです。$H(P)$ は $P$ の「ばらけ具合」で、これを大きくする(=罰則を引く)方向に効くので、解は少しぼやけて滑らかになります。なぜわざわざぼかすのか。それは、この正則化を入れると解が閉じた形のスケーリング反復で求まり、しかもGPUで超高速・微分可能になるからです。

具体的に解いてみます。エントロピー項を入れた問題のラグランジュ条件を整理すると、最適な $P$ は必ず次の形になります。

$$ \begin{equation} P_{ij} = u_i\, K_{ij}\, v_j, \qquad K_{ij} = \exp\!\left(-\frac{C_{ij}}{\varepsilon}\right) \end{equation} $$

ここで $K$ はギブスカーネル(コストを指数で潰した行列)、$u, v$ は2つの未知スケーリングベクトルです。あとは周辺条件 $P\mathbf{1} = a$ と $P^\top \mathbf{1} = b$ を満たすように $u, v$ を決めるだけ。これを代入すると、

$$ \begin{equation} u = \frac{a}{K v}, \qquad v = \frac{b}{K^\top u} \end{equation} $$

という連立になります(割り算は要素ごと)。1番目の式は「行和を $a$ に合わせるように $u$ を更新」、2番目は「列和を $b$ に合わせるように $v$ を更新」。この2つを交互に繰り返すだけで $u, v$ が収束します。これがSinkhornアルゴリズムです。やっていることは、行列 $K$ を行方向・列方向に交互にスケーリングして、周辺が $a, b$ になるよう整える操作に他なりません。

Sinkhorn反復の収束

左の図は、Sinkhorn反復のたびに「周辺制約の誤差」がどれだけ残るかを対数軸で示しています。誤差はほぼ直線的に下がっており、これは指数的な収束を意味します。数回〜数十回でほぼ制約を満たします。右の図は輸送コスト $\langle P, C \rangle$ が一定値に落ち着く様子です。交互スケーリングだけでこれほど速く収束するのが、Sinkhornが実務で愛される理由です。

では、正則化の強さ $\varepsilon$ は結果にどう効くのでしょうか。これは品質と計算のトレードオフを握る重要なパラメータです。

エントロピー正則化の強さと輸送計画

3枚のヒートマップは、$\varepsilon$ を大・中・小と変えたときの輸送計画です。$\varepsilon$ が大きい(左)と、計画は大きくぼやけて広がり、コストも高めです。$\varepsilon$ を小さくする(右)と、計画はくっきり対角に集中し、コストは厳密な線形計画の値に近づきます。つまり$\varepsilon \to 0$ で正則化なしの厳密解に収束します。一方で $\varepsilon$ を小さくしすぎると $\exp(-C/\varepsilon)$ がアンダーフローして数値不安定になるので、実用では対数領域での安定化版を使うか、適度な $\varepsilon$ を選びます。後のPython実装で、$\varepsilon$ を弱めると厳密解に近づくことを数値で確かめます。

ここまでで計算法が出そろいました。最後に、この最適輸送が実際にどんな場面で使われているかを見ていきましょう。

応用:WGAN・ドメイン適応・Wasserstein補間

Wasserstein GAN(WGAN)

すでに双対の節で触れたとおり、WGANはKantorovich-Rubinstein双対を直接使います。critic(批評家)$f$ を1-Lipschitzに制約しながら $\mathbb{E}_{\text{本物}}[f] – \mathbb{E}_{\text{生成}}[f]$ を最大化し、これを $W_1$ の推定値とします。生成器はこの $W_1$ を小さくするよう学習します。台が重ならない初期段階でも勾配が生きるので、従来のGANより学習が安定します。詳しい損失設計やLipschitz制約の実装は別記事を参照してください。

ドメイン適応

訓練したデータ(ソース)と本番のデータ(ターゲット)で分布がずれていると、モデルの精度は落ちます。最適輸送は、この2つの分布のあいだの最適な対応 $P$ を求め、ソースの各点をターゲット側へ「重心マッピング」で寄せることで分布を揃えます。各ソース点を、輸送計画の重み付き平均 $\hat{x}_i = \sum_j P_{ij} y_j / \sum_j P_{ij}$ で移すのが基本形です。

最適輸送によるドメイン適応

左は適応前で、青いソースとオレンジのターゲットが離れています。右は最適輸送でソースを移した後で、緑の点群がオレンジにぴたりと重なっています。分布が揃えば、ソースで学習した分類器をそのままターゲットに使っても精度が保てます。モデルを再学習せずにドメインのギャップを埋められるのが利点です。

Wasserstein重心と補間

最後は、最適輸送が生み出す「分布の補間」です。2つの分布 $\mu_0, \mu_1$ のあいだを、パラメータ $t \in [0,1]$ で滑らかにつなぐ経路を考えます。Wasserstein空間でこの中間点を取ると、山が「移動」しながら変形する自然な補間が得られます。これを変位補間(displacement interpolation)と呼びます。

Wasserstein補間と線形補間の比較

上はWasserstein補間で、1つの山が左から右へなめらかに移動していきます。下は単純な線形補間(密度の重み付き平均)で、途中で2つの山に分裂してしまいます。物理的に「砂山が動く」イメージに合うのは明らかに上です。この変位補間の考え方は、拡散モデルやflow matchingにつながります。これらは「ノイズ分布からデータ分布へ、最短経路(最適輸送に近い経路)で連続的に移す」流れを学習する手法で、最適輸送はその理論的な背骨になっています。さらに、複数分布の「平均」にあたるWasserstein重心は、形状の平均化やクラスタリングにも応用されます。

応用の広がりが見えたところで、これらを支える計算を実際にコードで動かしてみましょう。

Pythonでの実装

ここからは手を動かします。確認したいのは次の3点です。(1) 1次元Wassersteinがソート対応で正しく計算できること、(2) Sinkhornを自前実装し、$\varepsilon$ を弱めると厳密な線形計画の解に近づくこと、(3) 台が重ならない分布でKLが発散する一方Wが有限であること。外部の最適輸送ライブラリには頼らず、すべて自分で書きます。

1次元Wassersteinをソートで計算する

1次元では「両分布をソートして対応づける」だけで $W_p$ が求まります。平均0と平均2の正規分布(分散は同じ)なら、理論上 $W_1 = W_2 = 2$ になるはずです。

import numpy as np

rng = np.random.default_rng(42)
xm = rng.normal(0.0, 1.0, 100000)   # 平均0
xn = rng.normal(2.0, 1.0, 100000)   # 平均2(2だけ平行移動)

# 同サンプル数なので、ソートして同分位どうしを対応づける
sm = np.sort(xm)
sn = np.sort(xn)
W1 = np.mean(np.abs(sm - sn))
W2 = np.sqrt(np.mean((sm - sn) ** 2))
print(f"1次元 W1 = {W1:.4f}(理論値 2.0)")
print(f"1次元 W2 = {W2:.4f}(理論値 2.0)")

このコードの出力は W1 = 2.0070W2 = 2.0070 となります。どちらも理論値2.0にほぼ一致しました。平行移動しただけの分布では、分散が同じなら全分位が一様に2だけずれるので、$W_1$ も $W_2$ も平行移動量に等しくなります。サンプル数を増やせばさらに2.0に近づきます。ソート対応という素朴な操作で正確に計算できることが確認できました。

次は2次元以上でも使えるSinkhornを実装します。

Sinkhornアルゴリズムの自前実装

交互スケーリングをそのままコードにします。K = exp(-C/eps) を作り、uv を交互更新するだけです。

import numpy as np

def sinkhorn(a, b, C, eps, n_iter=5000):
    """エントロピー正則化最適輸送をSinkhorn(交互スケーリング)で解く。
    a, b: 周辺(確率ベクトル), C: コスト行列, eps: 正則化強さ。
    戻り値: 輸送計画P, 輸送コスト<P,C>。"""
    K = np.exp(-C / eps)
    u = np.ones_like(a)
    v = np.ones_like(b)
    for _ in range(n_iter):
        u = a / (K @ v + 1e-300)      # 行和をaに合わせる
        v = b / (K.T @ u + 1e-300)    # 列和をbに合わせる
    P = u[:, None] * K * v[None, :]
    return P, float(np.sum(P * C))

更新式 u = a / (K @ v)v = b / (K.T @ u) が、まさに本文で導いたスケーリング式です。微小値 1e-300 はゼロ割り防止です。これだけでエントロピー正則化された最適輸送が解けます。次に、この近似解が $\varepsilon \to 0$ で厳密解に迫ることを確かめます。

Sinkhornが厳密解に近づくことの確認

scipyのlinprogで厳密な線形計画を解き、Sinkhornの結果と $\varepsilon$ を変えながら比べます。

import numpy as np
from scipy.optimize import linprog

def lp_ot(a, b, C):
    """厳密な線形計画で最適輸送を解く。"""
    n, m = C.shape
    cvec = C.reshape(-1)
    A_eq = []
    for i in range(n):                       # 各行の和 = a_i
        r = np.zeros((n, m)); r[i, :] = 1.0; A_eq.append(r.reshape(-1))
    for j in range(m):                       # 各列の和 = b_j
        r = np.zeros((n, m)); r[:, j] = 1.0; A_eq.append(r.reshape(-1))
    A_eq = np.array(A_eq)
    b_eq = np.concatenate([a, b])
    res = linprog(cvec, A_eq=A_eq, b_eq=b_eq,
                  bounds=[(0, None)] * (n * m), method="highs")
    return res.x.reshape(n, m), float(np.sum(res.x.reshape(n, m) * C))

# 1次元配置: muは0.3付近, nuは0.7付近に質量
xs = np.linspace(0, 1, 30)
ys = np.linspace(0, 1, 30)
a = np.exp(-((xs - 0.3) ** 2) / 0.02); a /= a.sum()
b = np.exp(-((ys - 0.7) ** 2) / 0.02); b /= b.sum()
C = (xs[:, None] - ys[None, :]) ** 2     # 2次コスト

P_lp, cost_lp = lp_ot(a, b, C)
print(f"厳密LP コスト = {cost_lp:.5f}")
for eps in [0.3, 0.05, 0.01, 0.005, 0.001]:
    _, cost_sk = sinkhorn(a, b, C, eps)
    print(f"  Sinkhorn eps={eps:<6} コスト={cost_sk:.5f}"
          f"  (LP比 +{100*(cost_sk-cost_lp)/cost_lp:.2f}%)")

出力は次のようになります。

厳密LP コスト = 0.15989
  Sinkhorn eps=0.3    コスト=0.17814  (LP比 +11.41%)
  Sinkhorn eps=0.05   コスト=0.17253  (LP比 +7.90%)
  Sinkhorn eps=0.01   コスト=0.16398  (LP比 +2.56%)
  Sinkhorn eps=0.005  コスト=0.16195  (LP比 +1.28%)
  Sinkhorn eps=0.001  コスト=0.16010  (LP比 +0.13%)

$\varepsilon$ を $0.3 \to 0.001$ と小さくするにつれ、Sinkhornのコストが厳密LPの値 $0.15989$ にどんどん近づいています。$\varepsilon = 0.3$ では11.4%も上振れしていたのが、$\varepsilon = 0.001$ ではわずか0.13%の差です。これは本文で述べた「$\varepsilon \to 0$ で正則化なしの厳密解に収束する」という性質を、数値で裏づけています。エントロピー罰則の分だけコストが上振れし、罰則を弱めるほど真の最小コストに迫る、という関係がはっきり見えます。

最後に、Wasserstein距離の真骨頂である「KL発散 vs W有限」を確認します。

台が重ならない分布でのKL発散とW有限

平行移動した、台がまったく重ならない2つの離散分布を作り、KLとWを比べます。

import numpy as np

# 台が重ならない2分布: Pは0付近, Qは2付近
supp_p = np.array([0.0, 0.1, 0.2])
supp_q = np.array([2.0, 2.1, 2.2])
p = np.full(3, 1 / 3)
q = np.full(3, 1 / 3)

# 共通グリッドに乗せてKLを計算
grid = np.union1d(supp_p, supp_q)
P = np.array([p[list(supp_p).index(g)] if g in supp_p else 0.0 for g in grid])
Q = np.array([q[list(supp_q).index(g)] if g in supp_q else 0.0 for g in grid])
mask = P > 0
kl = np.sum(P[mask] * np.log(P[mask] / np.where(Q[mask] > 0, Q[mask], 1e-300)))
print(f"KL(P||Q) = {kl:.1f}  (Qが0の所にPの質量 → 実質的に発散)")

# 同じ2分布のW1(厳密LP)
C = np.abs(supp_p[:, None] - supp_q[None, :])
_, w1 = lp_ot(p, q, C)
print(f"W1(P, Q) = {w1:.4f}  (平行移動量2.0付近で有限)")

出力は KL(P||Q) = 689.7(ゼロ割りを 1e-300 で抑えた結果の巨大な値=実質的に発散)と W1(P, Q) = 2.0000 です。KLは台が重ならないために発散してしまい、距離としてまったく機能しません。一方 $W_1$ は2.0という有限の値を返し、しかもこれは2分布の平行移動量に一致します。KLが「無限大」としか言えない状況で、Wは「2.0だけずれている」と定量的に答えられる。 これがWasserstein距離が生成モデルやドメイン適応で選ばれる、決定的な理由です。

補遺:飛ばした証明をすべて埋める

この補遺は読み飛ばしても大丈夫です。 ここから先は、本文で証明なしに使った主張(双対性・1-Lipschitz帰着・Sinkhornの収束など)を一行ずつ埋める部分です。前半までで「分布を最小コストで運ぶ距離」というイメージは押さえられているので、結論を使えれば十分という方は次の節へ進んでください。証明を自分の手で追いたい方、論文の付録を読めるようになりたい方のための部分です。

ここまでは結果と直感を優先して、いくつかの主張を証明なしで使ってきました。「双対がなぜ成り立つのか」「なぜ $g=-f$ で1-Lipschitzに帰着するのか」「1次元の閉形式はどこから来るのか」「Sinkhornの形 $P=\mathrm{diag}(u)K\mathrm{diag}(v)$ はどう導かれ、なぜ収束するのか」。この補遺では、これら4つを1行も飛ばさずに埋めます。読者は「数学はできるがこの分野は初見」を想定し、記号は初出のたびに一言添えます。記述は Peyré–Cuturi Computational Optimal Transport、Villani Optimal Transport: Old and New、Santambrogio Optimal Transport for Applied Mathematicians に準拠します。

以下では離散の有限次元で進めます。$\mu$ は $n$ 点 $\{x_i\}$ 上の質量ベクトル $a\in\mathbb{R}^n_{\ge0}$($\sum_i a_i=1$)、$\nu$ は $m$ 点 $\{y_j\}$ 上の $b\in\mathbb{R}^m_{\ge0}$($\sum_j b_j=1$)です。コストは $C\in\mathbb{R}^{n\times m}$、$C_{ij}=c(x_i,y_j)$。記号 $\langle A,B\rangle=\sum_{ij}A_{ij}B_{ij}$ は行列のフロベニウス内積(成分どうしの積の総和)、$\mathbf{1}_m$ は全成分1の $m$ 次元ベクトルです。

補遺1:Kantorovich問題のLP双対を全行で導く

主問題(プライマル)はすでに見た線形計画です。変数は輸送計画行列 $P\in\mathbb{R}^{n\times m}$。

$$ \begin{equation} \min_{P\ge 0}\ \langle C,P\rangle \quad\text{s.t.}\quad P\mathbf{1}_m=a,\quad P^\top\mathbf{1}_n=b. \end{equation} $$

$P\mathbf{1}_m=a$ は「各行 $i$ の和が $a_i$」($x_i$ から運び出す総量)、$P^\top\mathbf{1}_n=b$ は「各列 $j$ の和が $b_j$」($y_j$ へ運び込む総量)。記号 $\min$ は最小値(達成されること—後述—を含意して $\inf$ でなく $\min$ と書きます)。

ステップ1:ラグランジアンを立てる。 等式制約 $a-P\mathbf{1}_m=0$ と $b-P^\top\mathbf{1}_n=0$ に、それぞれラグランジュ乗数ベクトル $f\in\mathbb{R}^n$(行制約に対応)と $g\in\mathbb{R}^m$(列制約に対応)を割り当てます。乗数 $f,g$ は符号制約なしの自由変数です(等式制約だから)。さらに不等式 $P\ge0$ に対しては「$P\ge0$ の集合上で $P$ について最小化する」形でそのまま扱います。ラグランジアンは

$$ \begin{equation} \mathcal{L}(P,f,g)=\langle C,P\rangle+\langle f,\,a-P\mathbf{1}_m\rangle+\langle g,\,b-P^\top\mathbf{1}_n\rangle. \end{equation} $$

ここで $\langle f,a\rangle=\sum_i f_i a_i$ はベクトルの内積です。

ステップ2:$P$ の項でまとめる。 第2・第3項を成分で書き下します。$\langle f,P\mathbf{1}_m\rangle=\sum_i f_i\sum_j P_{ij}=\sum_{ij}f_i P_{ij}$、同様に $\langle g,P^\top\mathbf{1}_n\rangle=\sum_j g_j\sum_i P_{ij}=\sum_{ij}g_j P_{ij}$ です。これを代入すると、

$$ \begin{equation} \mathcal{L}(P,f,g) =\langle f,a\rangle+\langle g,b\rangle +\sum_{ij}\bigl(C_{ij}-f_i-g_j\bigr)P_{ij}. \end{equation} $$

つまり $\langle f,a\rangle+\langle g,b\rangle$ は $P$ に依存しない定数項、残りは $P_{ij}$ に係数 $(C_{ij}-f_i-g_j)$ が掛かった線形項に整理できました。

ステップ3:双対関数($P\ge0$ 上での下限)を求める。 双対関数 $h(f,g)=\inf_{P\ge0}\mathcal{L}(P,f,g)$ を計算します。線形項 $\sum_{ij}(C_{ij}-f_i-g_j)P_{ij}$ を $P\ge0$ の範囲で最小化することを考えます。各成分は独立なので、係数の符号で場合分けします。

  • もし、ある $(i,j)$ で係数が負、すなわち $C_{ij}-f_i-g_j<0$ なら、その $P_{ij}$ を $+\infty$ に飛ばすと項は $-\infty$。よって下限は $-\infty$。
  • すべての $(i,j)$ で係数が非負、すなわち $C_{ij}-f_i-g_j\ge0$ なら、$P_{ij}\ge0$ のもとで各項の最小は $P_{ij}=0$ のとき $0$。よって線形項の最小値は $0$。

まとめると、

$$ \begin{equation} h(f,g)= \begin{cases} \langle f,a\rangle+\langle g,b\rangle, & f_i+g_j\le C_{ij}\ (\forall i,j),\\[4pt] -\infty, & \text{それ以外}. \end{cases} \end{equation} $$

下限が $-\infty$ になる $(f,g)$ は双対の最大化で自動的に捨てられるので、実質的に制約 $f_i+g_j\le C_{ij}$ を課したことになります。

ステップ4:双対問題が立つ。 双対は $h(f,g)$ の最大化なので、

$$ \begin{equation} \max_{f,g}\ \langle f,a\rangle+\langle g,b\rangle \quad\text{s.t.}\quad f_i+g_j\le C_{ij}\ (\forall i,j). \end{equation} $$

これがKantorovich双対です。本文の連続版 $\sup_{f,g}\ \mathbb{E}_\mu[f]+\mathbb{E}_\nu[g]$ s.t. $f(x)+g(y)\le c(x,y)$ と一致します($\langle f,a\rangle=\sum_i f_i a_i=\mathbb{E}_\mu[f]$ だから)。$f,g$ をKantorovichポテンシャルと呼びます。

ステップ5:弱双対(主 $\ge$ 双対)を乗数から直接示す。 任意の許容プライマル $P$($P\ge0,\ P\mathbf{1}_m=a,\ P^\top\mathbf{1}_n=b$)と、任意の許容デュアル $(f,g)$($f_i+g_j\le C_{ij}$)を取ります。制約 $f_i+g_j\le C_{ij}$ の両辺に $P_{ij}\ge0$ を掛けて(不等式の向きは保たれる)総和すると、

$$ \begin{equation} \sum_{ij}(f_i+g_j)P_{ij}\ \le\ \sum_{ij}C_{ij}P_{ij}=\langle C,P\rangle. \end{equation} $$

左辺を周辺条件で書き換えます。$\sum_{ij}f_i P_{ij}=\sum_i f_i\sum_j P_{ij}=\sum_i f_i a_i=\langle f,a\rangle$、同様に $\sum_{ij}g_j P_{ij}=\langle g,b\rangle$。よって

$$ \begin{equation} \langle f,a\rangle+\langle g,b\rangle\ \le\ \langle C,P\rangle. \end{equation} $$

これがすべての許容ペアで成り立つので、左辺の最大(双対最適値)は右辺の最小(主最適値)以下です。すなわち 双対 $\le$ 主(弱双対)。直感は明快で、ポテンシャル差 $f_i+g_j$ は「その経路に許される最大運賃 $C_{ij}$」を超えられないため、総収入は総コストを上回れない、ということです。

ステップ6:強双対(等号)の成立条件。 ここで主問題は有限個の線形不等式・等式で定義される多面体(feasible set は空でない:例えば独立カップリング $P_{ij}=a_i b_j$ が常に許容)上の有界($\langle C,P\rangle\ge \min_{ij}C_{ij}>-\infty$)な線形計画です。有限LPで主問題が許容かつ有界なら強双対が成り立つ(LP双対定理)ので、

$$ \begin{equation} \min_{P}\langle C,P\rangle=\max_{f,g}\ \langle f,a\rangle+\langle g,b\rangle. \end{equation} $$

一般の連続分布(Polish空間上のRadon測度)でも、コスト $c$ が下半連続かつ下に有界なら強双対は成立します(Kantorovich双対定理、Villani Theorem 5.10)。よって本文で「双対の最大値が $W$ そのもの」と書いたのは、この強双対に支えられています。さらに相補性条件として、最適 $P^\star$ と最適 $(f^\star,g^\star)$ は「$P^\star_{ij}>0\ \Rightarrow\ f^\star_i+g^\star_j=C_{ij}$」を満たします。ステップ5の不等式 $\sum_{ij}(C_{ij}-f_i-g_j)P_{ij}\ge0$ が等号で閉じる(双対ギャップ0)には、係数が正の $(i,j)$ で $P_{ij}=0$ でなければならないからです。これは「実際に質量を運ぶ経路ではポテンシャル差がコストにちょうど一致する(=その辺が tight)」ことを意味します。

補遺2:Kantorovich–Rubinstein双対($W_1$)を導く

コストが距離そのもの、$c(x,y)=\|x-y\|$ の場合を扱います(これが $W_1$)。補遺1の双対は $\sup_{f,g}\ \mathbb{E}_\mu[f]+\mathbb{E}_\nu[g]$ s.t. $f(x)+g(y)\le\|x-y\|$ でした。これを「$f$ ひとつ、しかも1-Lipschitz」に帰着させます。鍵は$c$-変換という最適化のテクニックです。

ステップ1:$g$ を最善に選ぶ($c$-変換)。 $f$ を固定したとき、制約 $g(y)\le\|x-y\|-f(x)$ をすべての $x$ について満たしつつ、目的の $\mathbb{E}_\nu[g]=\sum_j g_j b_j$ を最大化したい。$b_j\ge0$ なので各 $g_j$ はできるだけ大きく取るのが得。$g_j$ が満たすべき上界は全 $x_i$ にわたる下限なので、最大値は

$$ \begin{equation} g^c(y)=\inf_{x}\bigl(\|x-y\|-f(x)\bigr). \end{equation} $$

これを $f$ の $c$-変換と呼びます(記号 $\inf$ は下限)。$g$ をこの $g^c$ に置き換えても制約は破れず目的は増えるだけなので、双対は $\sup_f\ \mathbb{E}_\mu[f]+\mathbb{E}_\nu[g^c]$ に縮約できます。

ステップ2:$g^c=-f$ となるのは $f$ が1-Lipschitzのとき。 まず $g^c$ がどんな関数かを調べます。三角不等式 $\|x-y\|\ge\bigl|\,\|x-z\|-\|z-y\|\,\bigr|$ を使うと、$g^c$ は自動的に1-Lipschitzになりますが、ここでは目的の値を直接評価します。

$y$ を固定し、$g^c(y)=\inf_x(\|x-y\|-f(x))$ を考えます。候補のひとつ $x=y$ を代入すると $\|y-y\|-f(y)=-f(y)$。下限はこれ以下なので、

$$ \begin{equation} g^c(y)\ \le\ -f(y). \end{equation} $$

一方、もし $f$ が1-Lipschitz、すなわち $f(y)-f(x)\le\|x-y\|$(同値:$|f(x)-f(y)|\le\|x-y\|$)なら、移項して $\|x-y\|-f(x)\ge -f(y)$ がすべての $x$ で成り立ちます。下限を取っても向きは保たれるので $g^c(y)\ge-f(y)$。2つの不等式を合わせて、

$$ \begin{equation} f\ \text{が1-Lipschitz}\ \Longrightarrow\ g^c(y)=-f(y). \end{equation} $$

逆に、$g^c(y)=-f(y)$ が全 $y$ で成り立つなら、$\|x-y\|-f(x)\ge-f(y)$ すなわち $f(y)-f(x)\le\|x-y\|$ が全 $x,y$ で成り立ち、$f$ は1-Lipschitzです。さらに、1-Lipschitzでない $f$ を使っても、目的 $\mathbb{E}_\mu[f]+\mathbb{E}_\nu[g^c]$ は1-Lipschitz関数で置き換えて損をしません。正確な機構はこうです。$c$-変換を2回施した $f^{cc}:=(f^c)^c$ は必ず1-Lipschitzで、しかも $f^{cc}\ge f$ が成り立ちます。さらに $(f^{cc})^c=f^c$ なので、$f$ を $f^{cc}$ に置き換えても相方 $g^c=f^c$ は変わらず、$\mathbb{E}_\nu[g^c]$ は不変です。一方 $f^{cc}\ge f$ かつ $\mu\ge0$ なので $\mathbb{E}_\mu[f^{cc}]\ge\mathbb{E}_\mu[f]$。つまり置き換えで目的関数は決して下がりません。したがって最適化は1-Lipschitz関数 $f$ の上だけで行えば十分です。

ステップ3:$W_1$ の双対公式。 $g=-f$、$f$ が1-Lipschitz を代入すると、

$$ \begin{equation} W_1(\mu,\nu)=\sup_{\|f\|_{\mathrm{Lip}}\le1}\Bigl(\mathbb{E}_{x\sim\mu}[f(x)]-\mathbb{E}_{y\sim\nu}[f(y)]\Bigr). \end{equation} $$

ここで $\|f\|_{\mathrm{Lip}}=\sup_{x\ne x’}\frac{|f(x)-f(x’)|}{\|x-x’\|}$ はLipschitz定数(傾きの上限)。これが本文の式(Kantorovich–Rubinstein双対)です。$\mathbb{E}_\nu[g]=\mathbb{E}_\nu[-f]=-\mathbb{E}_\nu[f]$ から符号が引き算になりました。

WGAN criticの正体。 この $f$ こそ WGAN の critic(批評家)です。WGANは $f$ をニューラルネットで表し、1-Lipschitz制約のもとで $\mathbb{E}_{\text{本物}}[f]-\mathbb{E}_{\text{生成}}[f]$ を最大化して $W_1$ を推定します。「なぜ critic は1-Lipschitzでなければならないか」は、いま示したとおり1-Lipschitz性こそが $c$-変換で $g=-f$ に帰着するための必要十分条件だから、というのが答えです。Lipschitz制約の実装(スペクトル正規化など)は次の記事で詳説しています。

Lipschitz連続性とスペクトルノルム:1-Lipschitz証人とWGAN
1-Lipschitz制約をスペクトルノルムで課す方法を、WGAN criticの実装まで含めて解説します。

補遺3:1次元Wassersteinの閉形式を証明する

主張は、実数直線上の分布について

$$ \begin{equation} W_p^p(\mu,\nu)=\int_0^1\bigl|F_\mu^{-1}(u)-F_\nu^{-1}(u)\bigr|^p\,du \end{equation} $$

でした。$F_\mu(x)=\mu((-\infty,x])$ は累積分布関数(CDF)、$F_\mu^{-1}(u)=\inf\{x:F_\mu(x)\ge u\}$ は分位点関数(CDFの一般化逆関数。記号 $\inf$ は「$u$ パーセンタイルを与える最小の $x$」の意)。証明の核は「1次元では単調な対応づけ(順序を保つ輸送)が最適」という事実です。これを離散の rearrangement(並べ替え)不等式で厳密に示します。

ステップ1:同サイズ離散への帰着。 分位点を等間隔に $N$ 個サンプリングして近似します。$u_k=(k-\tfrac12)/N$($k=1,\dots,N$)とし、$\alpha_k=F_\mu^{-1}(u_k)$、$\beta_k=F_\nu^{-1}(u_k)$ を取ります。CDFの逆関数の定義から、$\alpha_1\le\alpha_2\le\cdots\le\alpha_N$ と $\beta_1\le\cdots\le\beta_N$ はどちらも昇順です。$N\to\infty$ でこの離散近似の輸送コストは右辺の積分に収束する(リーマン和)ので、離散で「ソート対応が最適」を示せば十分です。各点に等質量 $1/N$ を置いた経験分布どうしの $W_p^p$ を考えます。

ステップ2:同サイズ点群の最適輸送は順列で書ける。 等質量 $1/N$ の $N$ 点を $N$ 点へ運ぶ最適輸送計画 $P$ を考えます。$W_p^p$ の LP の最適解は多面体の頂点に取れ、Birkhoff–von Neumann の定理より、二重確率行列(行和・列和が一定)の頂点は置換行列です。よって最適計画は「$\alpha_k$ を $\beta_{\sigma(k)}$ へ送る」ある順列 $\sigma$ で表せ、コストは

$$ \begin{equation} \frac1N\sum_{k=1}^N\bigl|\alpha_k-\beta_{\sigma(k)}\bigr|^p \end{equation} $$

の $\sigma$ 上での最小化に帰着します。

ステップ3:恒等順列が最適(rearrangement 不等式)。 $\beta$ が昇順のとき、上の和を最小化するのは恒等 $\sigma(k)=k$、すなわちソートどうしの対応であることを示します。背理法ふうに、ある最適 $\sigma$ に「逆転対」があると仮定します。逆転対とは、$k\sigma(l)$ となる組です。このとき2点 $\alpha_k\le\alpha_l$ が、降順に並んだ $\beta_{\sigma(k)}\ge\beta_{\sigma(l)}$ へ割り当てられています。この2つの行き先を入れ替える($\alpha_k\to\beta_{\sigma(l)}$、$\alpha_l\to\beta_{\sigma(k)}$)と、コスト差は

$$ \begin{equation} \Delta=\Bigl(|\alpha_k-\beta_{\sigma(l)}|^p+|\alpha_l-\beta_{\sigma(k)}|^p\Bigr)-\Bigl(|\alpha_k-\beta_{\sigma(k)}|^p+|\alpha_l-\beta_{\sigma(l)}|^p\Bigr). \end{equation} $$

ここで4つの実数 $\alpha_k\le\alpha_l$、$\beta_{\sigma(l)}\le\beta_{\sigma(k)}$ について、関数 $t\mapsto t^p$($p\ge1$)の凸性から導かれる並べ替え不等式「小さいものは小さいものへ対応させる方が $|\cdot|^p$ の和は小さい」により $\Delta\le0$。つまり交換しても悪化せず、逆転を1つ減らせます。これを繰り返すと有限回で逆転対が消え、恒等順列に到達し、その都度コストは増えません。ゆえに恒等順列(ソート対応)が最適です。

補足($\Delta\le0$ の一行証明)。$a\le a’$、$d\le d’$ なる実数で、$\phi(t)=t^p$ が凸かつ単調なら $|a-d’|^p+|a’-d|^p\le|a-d|^p+|a’-d’|^p$ が成り立ちます。これは「クロスさせるより平行に対応させる方が距離コストが小さい」という1次元特有の性質で、$p=1$ なら絶対値の三角関係から、$p>1$ なら $\phi$ の凸性(差分の単調性)から従います。

ステップ4:積分形へ戻す。 ソート対応のコストは $\frac1N\sum_k|\alpha_k-\beta_k|^p=\frac1N\sum_k|F_\mu^{-1}(u_k)-F_\nu^{-1}(u_k)|^p$。$N\to\infty$ でこれは $\int_0^1|F_\mu^{-1}(u)-F_\nu^{-1}(u)|^p\,du$ に収束します。よって

$$ \begin{equation} W_p^p(\mu,\nu)=\int_0^1\bigl|F_\mu^{-1}(u)-F_\nu^{-1}(u)\bigr|^p\,du. \end{equation} $$

$p=1$ の場合は、変数変換 $u=F(x)$ を用いると $\int_0^1|F_\mu^{-1}-F_\nu^{-1}|\,du=\int_{-\infty}^\infty|F_\mu(x)-F_\nu(x)|\,dx$(CDF差の面積)とも書けます(本文の式)。この証明は本文のPython実装と完全に整合します:「両分布から同数サンプルを取り、それぞれソートして差を取る」操作は、まさにステップ2–3の「ソート対応が最適」をそのまま実行しているのです。実測で $W_1=W_2=2.0070$(理論値2.0)が得られたのは、平行移動した正規分布では全分位が一様に2ずれ、ソート対応がそれを正しく拾うからです。

補遺4:エントロピー正則化の最適解とSinkhornの収束

最後に、正則化問題

$$ \begin{equation} \min_{P\in\Pi(a,b)}\ \langle C,P\rangle-\varepsilon H(P), \qquad H(P)=-\sum_{ij}P_{ij}\bigl(\log P_{ij}-1\bigr) \end{equation} $$

の最適解が $P=\mathrm{diag}(u)K\mathrm{diag}(v)$($K=e^{-C/\varepsilon}$)になること、そしてSinkhorn反復が収束することを示します。$\Pi(a,b)$ は周辺が $a,b$ のカップリング集合、$\varepsilon>0$。本文の符号規約(罰則として $-\varepsilon H$ を引く)に合わせます。

ステップ1:目的関数を整理。 $H$ の定義を代入すると、$-\varepsilon H(P)=\varepsilon\sum_{ij}P_{ij}(\log P_{ij}-1)$。よって目的は

$$ \begin{equation} E(P)=\sum_{ij}C_{ij}P_{ij}+\varepsilon\sum_{ij}P_{ij}\bigl(\log P_{ij}-1\bigr). \end{equation} $$

各項は $P_{ij}\mapsto P_{ij}\log P_{ij}$ を含み、これは狭義凸(2階微分 $1/P_{ij}>0$)。線形項を足しても凸性は保たれるので、$E$ は $P$ について狭義凸。さらに $\Pi(a,b)$ はコンパクト凸集合。よって最小解は一意に存在します。最適点は内部($P_{ij}>0$)に来ます($P_{ij}\to0^+$ で $\partial E/\partial P_{ij}=C_{ij}+\varepsilon\log P_{ij}\to-\infty$、つまり境界では下げ続けられるため、最小は内部)。これでLagrange法(等式制約のみ)が使えます。

ステップ2:Lagrangianと1階条件。 周辺制約 $\sum_j P_{ij}=a_i$(乗数 $f_i$)、$\sum_i P_{ij}=b_j$(乗数 $g_j$)を付けて、

$$ \begin{equation} \mathcal{L}(P,f,g)=\sum_{ij}\Bigl(C_{ij}P_{ij}+\varepsilon P_{ij}(\log P_{ij}-1)\Bigr) -\sum_i f_i\Bigl(\sum_j P_{ij}-a_i\Bigr)-\sum_j g_j\Bigl(\sum_i P_{ij}-b_j\Bigr). \end{equation} $$

$P_{ij}$ で偏微分して0と置きます。$\frac{\partial}{\partial P_{ij}}\bigl[P_{ij}(\log P_{ij}-1)\bigr]=\log P_{ij}$(積の微分:$(\log P_{ij}-1)+P_{ij}\cdot\frac1{P_{ij}}=\log P_{ij}$)なので、

$$ \begin{equation} \frac{\partial\mathcal{L}}{\partial P_{ij}}=C_{ij}+\varepsilon\log P_{ij}-f_i-g_j=0. \end{equation} $$

ステップ3:解の形を取り出す。 $\log P_{ij}$ について解くと $\log P_{ij}=\dfrac{f_i+g_j-C_{ij}}{\varepsilon}$。両辺を指数に乗せ、指数の和を積に分解します:

$$ \begin{equation} P_{ij}=\exp\!\Bigl(\frac{f_i}{\varepsilon}\Bigr)\cdot\exp\!\Bigl(-\frac{C_{ij}}{\varepsilon}\Bigr)\cdot\exp\!\Bigl(\frac{g_j}{\varepsilon}\Bigr). \end{equation} $$

ここで $u_i=e^{f_i/\varepsilon}$、$v_j=e^{g_j/\varepsilon}$、$K_{ij}=e^{-C_{ij}/\varepsilon}$ と置けば、

$$ \begin{equation} P_{ij}=u_i\,K_{ij}\,v_j,\qquad\text{すなわち}\quad P=\mathrm{diag}(u)\,K\,\mathrm{diag}(v). \end{equation} $$

これが本文の式 $P_{ij}=u_iK_{ij}v_j$ です。$u,v$ は正のスケーリングベクトル、$K$ はギブスカーネル。Lagrangeの1階条件から、解は必ずこの「対角スケーリング2枚で挟んだ形」になることが導けました。

ステップ4:周辺条件 → Sinkhorn更新式。 あとは $u,v$ を周辺条件で決めます。行和制約 $P\mathbf{1}_m=a$ を成分で書くと、$\sum_j u_i K_{ij}v_j=u_i\sum_j K_{ij}v_j=u_i\,(Kv)_i=a_i$。よって

$$ \begin{equation} u_i=\frac{a_i}{(Kv)_i},\qquad\text{ベクトルで}\quad u=\frac{a}{Kv}\ (\text{要素ごとの割り算}). \end{equation} $$

同様に列和制約 $P^\top\mathbf{1}_n=b$ から $v_j=b_j/(K^\top u)_j$、すなわち $v=b/(K^\top u)$。この2式を交互に適用するのがSinkhorn反復です。本文の自前実装 u = a/(K@v)v = b/(K.T@u) はこの導出そのものです。

ステップ5:収束の証明(Hilbert射影距離での縮小写像)。 Sinkhornが収束する理由を述べます。1回の更新 $v\mapsto v’=b/(K^\top(a/(Kv)))$ を考え、正のベクトルどうしの距離をHilbert射影距離

$$ \begin{equation} d_H(x,y)=\log\frac{\max_i (x_i/y_i)}{\min_i (x_i/y_i)} \end{equation} $$

で測ります(スケール不変な距離:$x$ と $cx$ の距離は0、正の射影空間 cone 上の距離)。鍵となるのがBirkhoff–Hopfの定理です。これは「正の要素を持つ線形写像 $x\mapsto Kx$ は、Hilbert射影距離に関して縮小写像であり、その縮小係数(Birkhoff contraction ratio)は

$$ \begin{equation} \kappa(K)=\frac{\sqrt{\eta(K)}-1}{\sqrt{\eta(K)}+1}<1, \qquad \eta(K)=\max_{i,j,k,l}\frac{K_{ik}K_{jl}}{K_{jk}K_{il}} \end{equation} $$

で与えられる」というものです($\eta$ は $K$ の射影直径に相当)。$K=e^{-C/\varepsilon}$ は全成分正なので $\eta(K)<\infty$、ゆえに $\kappa(K)<1$。要素ごとの割り算($x\mapsto a/x$)は $d_H$ を保つ(等長:比の最大最小が入れ替わるだけ)ので、Sinkhornの1ステップ(行スケーリング+列スケーリング)の合成は $d_H$ に関して係数 $\kappa(K)^2<1$ の縮小写像になります。バナッハの不動点定理より、反復は唯一の不動点に幾何(指数)収束します。

これは本文の図(ot07_sinkhorn_conv)で「周辺制約の誤差が対数軸で直線的に下がる=指数収束」と観察したことの理論的裏づけです。縮小係数 $\kappa(K)^2$ が収束率を決め、$\varepsilon$ が小さいほど $K$ のコントラストが強く $\eta$ が大きく($\kappa\to1$)なるので、収束は遅くなります。

ステップ6:$\varepsilon\to0$ で厳密LP解へ。 最後に「罰則を弱めると厳密解に戻る」ことの直感を、実測と接続します。エントロピー項を $\varepsilon$ 倍で足しているので、$\varepsilon\to0$ で目的は元のLP $\langle C,P\rangle$ に各点収束します。狭義凸な近似問題の最小解 $P_\varepsilon$ は、$\varepsilon\to0$ で元LPの最適解集合のうち最大エントロピー解(複数最適解があるとき、最もばらけたもの)に収束することが知られています(Γ収束の議論、Peyré–Cuturi 4.2節)。コストの上振れは、エントロピー罰則の分だけ最小コストから持ち上がる量で、$\varepsilon$ に比例して縮みます。本文の実測

$$ \begin{equation} \varepsilon:0.3\to0.001,\qquad \text{LP比}\ +11.41\%\to+0.13\% \end{equation} $$

は、まさにこの収束を数値で示したものです。$\varepsilon=0.3$ では罰則が強く解が大きくぼやけてコストが11.4%上振れし、$\varepsilon=0.001$ ではほぼ厳密LP値 $0.15989$ に一致(差0.13%)しました。理論($\varepsilon\to0$ で厳密解へ)と実装(弱めるほどLP値に迫る)が、ここで完全に符合します。

以上で、本文で結果だけ述べた4つの主張—LP双対、Kantorovich–Rubinstein双対、1次元閉形式、Sinkhornの形と収束—をすべて証明できました。これで本記事の理論は自己完結します。

まとめ

本記事では、最適輸送とWasserstein距離を、直感から実装まで通して解説しました。

  • 最適輸送は「片方の分布を、もう片方の形に最小コストで運ぶ」問題。Mongeの決定的写像は質量を分割できず扱いにくいが、Kantorovichの輸送計画 $\gamma$(結合分布)への緩和で線形計画になり、解が必ず存在する。
  • Wasserstein距離 $W_p$ は最小輸送コストの $p$ 乗根で定義される、分布空間の正真正銘の距離。1次元なら分位点関数の差という閉形式で計算でき、ソート対応だけで求まる。
  • Kantorovich-Rubinstein双対は $W_1$ を1-Lipschitz証人関数の期待値差の最大化で表す。これがWGANのcriticの正体で、なぜLipschitz制約が必要かを説明する。
  • KLとの違いが核心。台が重ならないとKLは発散して勾配が死ぬが、Wは有限で滑らかな勾配を持つ。実測でもKLは実質発散、$W_1$ は平行移動量2.0を正しく返した。
  • 計算は厳密な線形計画か、エントロピー正則化+Sinkhornアルゴリズム。Sinkhornは交互スケーリングだけで高速・微分可能で、$\varepsilon \to 0$ で厳密解に収束することを数値で確認した。
  • 応用はWGAN、ドメイン適応(重心マッピングで分布を寄せる)、Wasserstein重心と変位補間(拡散・flow matchingの理論的背骨)。

最適輸送は、生成モデルから領域適応まで、現代の機械学習を横断する共通言語になっています。次のステップとして、双対表現を実際の画像生成に応用する流れと、別の分布間距離の体系を押さえておくとよいでしょう。

次のステップとして、以下の記事も参考にしてください。