はじめに

Sinkhorn Distances: Optimal Transport with Entropic Constraintsを読んだのでメモ．

気持ち

WGANなどで有名になったWasserstein distanceは距離を計算するのに最適化問題を解かなければならず，離散の確率分布間の距離を図ろうとした際にはその次元数に対して $O(d^3\log d)$ の計算コストがかかることが知られている．この論文ではtransport matrix $P$ のエントロピーに関する正則化を加えることで問題を緩和し，高速に計算可能なSinkhorn Distanceを提案している．正則化の係数を限りなく0に近づけた際にSinkhorn DistanceはWasserstein distanceに一致する．

Optimal Transport

異なる確率分布間の距離を測る問題について考える．単体 $\Sigma_d:=\{x\in\mathbb{d}_+:x^T\mathbf{1}_d=1\}$ 上の二つの確率ベクトル $r,c$ (ニューラルネットやってる人的にはsoftmaxの出力で得られるベクトル)が与えられた時，次のpolyhedral set $U(r,c)$ を定義する．

$\displaystyle U(r,c):=\{P\in\mathbf{R}^{d\times d}_+|P\mathbf{1}_d=r,P^T\mathbf{1}_d=c\}$

$U(r,c)$ は行，列それぞれの和が $r$ と $c$ になる非負の要素からなる[tex;d\time d]行列全てを含む集合で，仮に $r,c$ がそれぞれ多項分布 $X,Y$ に従う変数だった場合 $U(r,c)$ は全ての可能な $X,Y$ の同時確率を含む集合になっている．この $U(r,c)$ に対してエントロピーと $P,Q\in U(r,c)$ 間のKL divergenceを次のように定める．

$\displaystyle h(P)=-\sum_{i,j=1}^dp_{ij}\log p_{ij},\:\mathbf{KL}(P\| Q)=\sum_{ij}p_{ij}\log\frac{p_{ij}}{q_{ij}}$

$r,c$ 間のcost matrix $M\in\mathbb{R}^{d\times d}$ が与えられた時，transport matrix (同時確率行列) $P$ を使って $r$ から $c$ へのmappingのコストをFrobenius inner product $\langle\cdot,\cdot\rangle$ を使って $\langle P,M\rangle$ と定義する．このコストを最小化する次の問題のことをoptimal transportと呼ぶ．

$\displaystyle d_M(r,c):=\min_{P\in U(r,c)}\langle P,M\rangle$

ただし，cost matrix $M$ は以下で定義される距離行列の推に属する．

$\displaystyle \mathcal{M}=\{M\in\mathbb{R}^{d\times d}_+:\forall i,j\leq d,m_{ij}=0\Leftrightarrow i=j,\forall i,j,k\leq d,m_{ij}\leq m_{ik}+m_{kj}\}$

Sinkhorn Distance

$rc^T$ のエントロピーが $h(rc^T)=h(r)+h(c)$ であることから，以下の不等式はtightである．

$\displaystyle \forall r,c\in\Sigma_d,\forall P\in U(r,c),h(P)\leq h(r)+h(c)$

また，エントロピーの凸性から以下の凸集合を導入することができる．

$\displaystyle U_\alpha(r,c):=\{P\in U(r,c)|\mathbf{KL}(P\|rc^T)\leq\alpha\}=\{P\in U(r,c)|h(P)\leq h(r)+h(c)-\alpha\}\subset U(r,c)$

この集合を用いてSinkhorn Distanceを次のように定義する．

$\displaystyle d_{M,\alpha}(r,c):=\min_{P\in U_\alpha(r,c)}\langle P,M\rangle$

Optimal transportにおいてエントロピーを考える必要性として2つ理由があり，一つは計算上の問題で，もう一つは直感的に次のようなものが挙げられる．すなわち，通常のoptimal transportでは解となる $P$ は単体の頂点に位置するようなものとなり，決定的な同時確率行列となる．それに対しエントロピーを考えることで確率的な揺れを許容した緩和を行なっていることになる．

Sinkhorn distanceは $\alpha$ によって次のことが言える．

$\alpha$ が十分に大きければsinkhorn distance $d_{M,\alpha}$ は輸送距離 $d_M$ に一致する．

これは任意の $P\in U(r,c),h(P)$ においてその下界が $\frac{1}{2}(h(r)+h(c))$ で抑えられることから言える．

$\alpha=0$ の時， $M$ がユークリッド距離行列（ $m_{ij}=\|\varphi_i-\varphi_j\|_2^2\|）であれば[tex:d_{M,0}=r^TMc$ は負定値カーネルとなる．

これらのextremeな場合に限らず全ての可能な $\alpha$ に対しsinkhorn distanceは可換かつ三角不等式を満たす．ただし $\alpha$ が十分小さい場合において $h(r)\gt 0$ であるような任意の $r$ に対し $d_{M,\alpha}(r,r)\gt 0$ であるため，sinkhorn distanceは距離の公理（ $d(x,y)=0\Leftrightarrow x=y$ ）を満たさない．ただしこれは $d_{M,\alpha}$ に $\mathbf{1}_{r\neq c}$ をかけることで満たすことができる．これをTheorem 1とする．

Theorem 1.

全ての $\alpha\geq 0,M\in\mathcal{M}$ において, $d_{M,\alpha}$ は対称かつ三角不等式を満たす．また， $(r,c)\mapsto\mathbf{1}_{r\neq c}d_{M,\alpha}(r,c)$ は距離の公理を全て満たす．

Theorem 1の証明には以下のLemma 1が重要となる．

Lemma 1 (Gluing Lemma With Entropic Constraint)

$\alpha\geq 0,:x,y,z\in\Sigma_d,:P\in U_\alpha(x,y),:Q\in U_\alpha(y,z)$ とする． $S$ を $s_{ik}:=\sum_j\frac{p_{ij}q_{jk}}{y_j}$ で定義される $d\times d$ 行列とすると $S$ は $S\in U_\alpha(x,z)$ となる．

実際にTheorem 1の証明をする． $d_{M,\alpha}$ の対称性は $M$ の対称性から直接いえる． $x,y,z\in\Sigma_d$ とし， $P\in U_\alpha(x,y),Q\in U_\alpha(y,x)$ を $d_{M,\alpha}(x,y),d_{M,\alpha}(y,z)$ の二つの輸送問題の解とする．Lemma 1で与えられた $S$ を用いて以下のように三角不等式を示す．

$\displaystyle d_{M,\alpha}(x,y)=\min_{P\in U_\alpha(x,y)}\langle P,M\rangle\leq\langle S,M\rangle=\sum_{ik}m_{ik}\sum_j\frac{p_{ij}q_{jk}}{y_j}\leq\sum_{ijk}(m_{ij}+m_{jk})\frac{p_{ij}q_{jk}}{y_j}\\ =\sum_{ijk}m_{ij}{p_{ij}q_{jk}}{y_j}+m_{jk}{p_{ij}q_{jk}}{y_j}=\sum_{ij}m_{ij}p_{ij}\sum_k\frac{q_{jk}}{y_j}+\sum_{jk}m_{jk}q_{jk}\sum_i\frac{p_{ij}}{y_j}\\ =\sum_{ij}m_{ij}p_{ij}+\sum_{jk}m_{jk}q_{jk}=d_{M,\alpha}(x,y)+d_{M,\alpha}(y,z)$

$P,Q$ は同時確率行列であるため $\sum_k\frac{q_{jk}}{y_j}=1,\sum_i\frac{p_{ij}}{y_j}=1$ が成り立つことを使っている．ちなみに論文では $\min_{P\in U_\alpha(x,z)}$ となっていたがおそらくタイポ．

Computing Regularized Transport with Sinkhorn's Algorithm

ここではsinkhorn distanceのエントロピー拘束に対するラグランジュ乗数を考える．

$\displaystyle \mathrm{For}:\lambda\gt 0,d^\lambda_M(r,c):=\langle P^\lambda,M\rangle,:\mathrm{where}:P^\lambda=\underset{P\in U(r,c)}{\mathrm{arg}\min}\langle P,M\rangle-\frac{1}{\lambda}h(P)$

双対定理から $d_{M,\alpha}(r,c)=d^\lambda_M(r,c)$ が得られ，この $d^\lambda_M$ をdual-Sinkhorn divergenceと呼ぶ． $d^\lambda_M$ は元の問題 $d_M$ に比べて計算が非常に楽という利点がある．Dual-Sinkhorn divergenceを利用して以下の命題が導かれる．

Lemma 2.

$\lambda\gt 0$ において解 $P^\lambda$ は一意に求まり， $P^\lambda=\mathrm{diag}(u)K\mathrm{diag}(v)$ の形で表現される．ただし， $u,v$ は非負の $\mathbb{R}^d$ ベクトルで $K:=e^{-\lambda M}$ は $-\lambda M$ のelement-wise exponential．

この命題は以下のように $P\mathbf{1}_d=r,P^T\mathbf{1}_d=c$ という二つの拘束に対する双対変数 $\alpha,\beta\in\mathbb{R}^d$ を使ってdual-sinkhorn divergenceを書くことで導ける．

$\displaystyle \mathcal{L}(P,\alpha,\beta)=\sum_{ij}\frac{1}{\lambda}p_{ij}\log p_{ij}+p_{ij}m_{ij}+\alpha^T(P\mathbf{1}_d-r)+\beta^T(P^T\mathbf{1}_d-c)$

$\partial\mathcal{L}/\partial p_{ij}=0$ から $p_{ij}=e^{-1/2-\lambda\alpha_i}e^{-\lambda m_{ij}}e^{-1/2-\lambda\beta_j}$ となる．Sinkhorn's theoremから $U(r,c)$ に属する $\mathrm{diag}(u)K\mathrm{diag}(v)$ が $u,v\geq\mathbf{0}_d$ において一意に存在する事が言える．よって $P^\lambda$ はSinkhorn's fixed point iterationによって， $(u,v)\leftarrow(r./Kv,c./K'u)$ と計算可能．

以上から $K,r,c$ が与えられれば十分な回数更新を繰り返す事で $P^\lambda$ を計算できる．更新の式は $u\leftarrow r./K(c./K'u)$ として一つの式で表現できる．ここの詳しい計算手順は論文のAlgorithm 1にまとまっているので要参照．

まとめ

一部証明を厳密に追えてないけど，色々な論文で使われているのがよくわかる．実験で一部気になるのはFigure 2の結果でEMDよりtest errorが低いのがよくわからない．実験は読んでないから細かい実験条件はわからないが直感的には厳密に解くEMDの方が良さそうだが，同時確率行列への緩和がいい方向に効いたという事なのか．何はともあれ簡単に実装してみたい．

機械学習とかコンピュータビジョンとか

CVやMLに関する勉強のメモ書き。

Sinkhorn Distances: Optimal Transport with Entropic Constraintsを読んだのでメモ