Sinkhorn Distance の計算方法を最適化問題から見て、Python でその動作を確認する。
Sinkhorn Distance は最適輸送 (Optimal Transport) にエントロピー正則化を適用したものである。
そこで、まず最適輸送問題について具体例とともに確認する。
参考論文
同論文の解説記事
\newcommand{\vector}[1]{\mathrm{\mathbf #1}}
\newcommand{\subjectto}{\mathop{\mathrm{subject\ to}}}
\newcommand{\where}{\mathop{\mathrm{where}}}
\newcommand{\minimize}{\mathop{\mathrm{minimize}}}
\newcommand{\diag}{\mathop{\mathrm{diag}}}
最適輸送問題
最適輸送問題は以下の最適化問題として与えられる。
\begin{align}
\minimize_{\vector P\in U(\vector r, \vector c)} \quad & \left<\vector P,\vector M\right> \\
\where \quad & U(\vector r, \vector c) = \left\{P \in \mathbb{R}^{m \times n}_+ \mid \vector P \vector 1_n = \vector r,\ \vector P^\top \vector 1_m = \vector c \right\}
\end{align}
ここで、$\vector M \in \mathbb{R}^{m \times n}$ は確率ベクトル $\vector r \in \Sigma_m$ から $\vector c \in \Sigma_n$ へのコスト行列、決定変数 $\vector P$ は $\vector r$ から $\vector c$ への輸送確率行列となっている ($\Sigma_d:={\vector x\in\mathbb{R}^d \mid \vector x ^\top \vector 1_d = 1}$)。
この問題の制約集合を書き下すと、以下の線形計画問題となり、線形計画のアルゴリズム (シンプレックス法とか内点法とか) を用いて問題を解くことができる。
\begin{align}
\minimize_{\vector P\in\mathbb{R}^{m \times n}} \quad & \left<\vector P,\vector M\right> \\
\subjectto \quad & \vector P \vector 1_n = \vector r,\ \vector P^\top \vector 1_m = \vector c,\ \vector P \geq 0
\end{align}
この最適化問題の最適値 $d_{\vector M}(\vector r, \vector c) := \min_{\vector P \in U(\vector r, \vector c)}\left<\vector P,\vector M\right>$ は Earth Mover's Distance (EMD) と呼ばれている (Wikipedia)。
以降 $\vector r > 0, \vector c > 0$ とする。($r_i$ = 0 のとき $P_{ij}=0$ となるため、$\vector c$ も同様)
具体例
最適輸送の具体例として、画像の色変換を適当に実装してみる。
まず必要なライブラリをインポートする。
import cv2
import numpy as np
from scipy.optimize import linprog
適当な画像を２枚用意し、Lab 変換を施し、a と b のヒストグラムを作成する (上式の $\vector r, \vector c$ にあたる)。
max_val = 256 * 256
# Image A
bird_blue = cv2.imread('../data/25193588_m.jpg')
lab_blue = cv2.cvtColor(bird_blue, cv2.COLOR_BGR2Lab).astype(int)
ab_blue = lab_blue[..., 1] * 256 + lab_blue[..., 2]
hist_blue = np.bincount(ab_blue.ravel(), minlength=max_val)
nonzero_blue = np.nonzero(hist_blue)[0]
hist_blue = hist_blue[nonzero_blue]
hist_blue = hist_blue / hist_blue.sum()
# Image B
bird_green = cv2.imread('../data/25421259_m.jpg')
lab_green = cv2.cvtColor(bird_green, cv2.COLOR_BGR2Lab).astype(int)
ab_green = lab_green[..., 1] * 256 + lab_green[..., 2]
hist_green = np.bincount(ab_green.ravel(), minlength=max_val)
nonzero_green = np.nonzero(hist_green)[0]
hist_green = hist_green[nonzero_green]
hist_green = hist_green / hist_green.sum()
コスト行列 $\vector M$ を作成する。
cost_matrix = np.abs(nonzero_blue.reshape(-1, 1) - nonzero_green.reshape(1, -1))
cost_matrix = (cost_matrix // 256 + cost_matrix % 256) / 255.
最後に Scipy の線形計画関数 scipy.optimize.linprog で使用できるように array を整え、線形計画問題を解く。
n_row, n_col = cost_matrix.shape
prob = np.hstack([hist_blue, hist_green])
rowsum_coef = np.tile(np.identity(n_col), n_row)
colsum_coef = np.repeat(np.identity(n_row), n_col, axis=1)
coef = np.vstack([colsum_coef, rowsum_coef])
res = linprog(cost_matrix.ravel(), A_eq=coef, b_eq=prob, bounds=(0, None))
optim_transport = res.x.reshape(n_row, n_col)
optim_distance = res.fun
と書いたが、行列のサイズが大きくなりすぎるため、メモリに乗らない。
scipy.sparse を使うとなんとかなるかもしれないが、ここではPOT: Python Optimal Transport という最適輸送ライブラリを使用する。
POT は pip install POT でインストールできる。
POT を用いると、以下で最適輸送 (EMD) を計算できる。
import ot
optim_transport = ot.emd(hist_blue, hist_green, cost_matrix)
ここでは輸送問題を解く際に、Displacement interpolation using Lagrangian mass transport | Proceedings of the 2011 SIGGRAPH Asia Conference のアルゴリズムを使用しているらしい。
求まった最適解に対し、例えば argmax を取ることで輸送先を一意に決定し、画像の色を変換することができる。
lab_transported = lab_blue.copy()
transport_blue = np.arange(max_val)
transport_blue[nonzero_blue] = nonzero_green[np.argmax(optim_transport, axis=1)]
ab_transported = transport_blue[ab_blue]
lab_transported[..., 1] = ab_transported // 256
lab_transported[..., 2] = ab_transported % 256
transported = cv2.cvtColor(lab_transported.astype(np.uint8), cv2.COLOR_Lab2BGR)
Sinkhorn Distance
Earth Mover's Distance の最適化問題にエントロピーの制約
$$h(\vector P) \geq h(\vector r) + h(\vector c) - \alpha$$
を追加した次の問題を考える。(エントロピー: $h(\vector P) = -\sum_{ij}P_{ij}\log P_{ij}$)
\begin{align}
\minimize_{\vector P\in U_\alpha(\vector r, \vector c)} \quad & \left<\vector P,\vector M\right> \\
\where \quad & U_\alpha(\vector r, \vector c) = \left\{\vector P \in \mathbb{R}^{m \times n}_+ \mid \vector P \vector 1_n = \vector r,\ \vector P^\top \vector 1_m = \vector c, h(\vector P) \geq h(\vector r) + h(\vector c) - \alpha \right\}
\end{align}
この問題の最適値を Sinkhorn Distance という。制約式中の $\alpha$ の値に対応した $\lambda > 0$ を使用することで、以下の等価な問題が導かれる。
\begin{align}
\minimize_{\vector P\in\mathbb{R}^{m \times n}} \quad & \left<\vector P,\vector M\right> - \frac{1}{\lambda} h(\vector P)\\
\subjectto \quad & \vector P \vector 1_n = \vector r,\ \vector P^\top \vector 1_m = \vector c,\ \vector P \geq 0
\end{align}
エントロピーの狭義凹性より、この問題は狭義凸最適化問題である。また、制約条件はいずれも線形制約であるため、KKT 条件より大域的最適解を求められる。この問題のラグランジアンは次式で与えられる。
\begin{align}
\mathcal{L} \left(\vector P, \vector{x}, \vector{y}, \vector{Z}\right) &= \left<\vector P,\vector M\right> - \frac{1}{\lambda} h(\vector P) + \vector x^\top \left(\vector P \vector 1_n - \vector r\right) + \vector y^\top \left(\vector P^\top \vector 1_m - \vector c\right) - \left<\vector Z, \vector P\right> \\
&= \left<\vector P, \vector M + \frac{1}{\lambda}\log\vector P + \vector x \vector 1_n^\top + \vector 1_m \vector y^\top - \vector Z\right> - \vector x^\top \vector r - \vector y^\top \vector c
\end{align}
これより、KKT 条件は
\frac{\partial \mathcal{L}}{\partial P_{ij}} = M_{ij} + \frac{1}{\lambda}\left(\log {P_{ij}} + 1\right) + x_i + y_j - Z_{ij} = 0 \\
\vector P \vector 1_n = \vector r \\
\vector P^\top \vector 1_m = \vector c \\
\vector P \geq 0,\ \vector Z \geq 0,\ P_{ij}Z_{ij} = 0
となる。$P_{ij} > 0$ のとき、$P_{ij}Z_{ij} = 0$ より $Z_{ij} = 0$ となる。
また、$\frac{\partial \mathcal{L}}{\partial P_{ij}} = 0$ より、
P_{ij} = \exp\left(-\lambda\left(M_{ij} + x_i + y_j - Z_{ij}\right) - 1\right) = \exp\left(-\lambda x_i - \frac12\right) \exp\left(-\lambda M_{ij} + \lambda Z_{ij}\right) \exp\left(-\lambda y_j - \frac12\right) > 0
となるため、$Z_{ij} = 0$ である。
ここで、
\vector u := \exp\left(-\lambda \vector x - \frac12\right),\ \vector v := \exp\left(-\lambda \vector y - \frac12\right),\ \vector K := \exp\left(-\lambda \vector M\right)
とすると、KKT 条件を満たす $\vector P$ は
\vector P = \diag(\vector u)\vector K \diag(\vector v)
とかける。
以上をまとめると、Sinkhorn Distance の最適化問題の最適解は正の値をとる行列であり ($P_{ij} > 0$)、その行和、列和は確率ベクトルであり ($\vector P \vector 1_n = \vector r, \vector P^\top \vector 1_m = \vector c, \sum_i r_i = \sum_j c_j = 1$)、正の値をとる対角行列を用いて $\vector P = \diag(\vector u)\vector K \diag(\vector v)$ と分解される。
これらの条件を満たす行列 $\vector P$ は、以下の Sinkhorn Iteration によって生成できる。
\begin{align}
\vector u^{k+1} &= \frac {\vector r} {\vector K \vector v^{k}} \\
\vector v^{k+1} &= \frac {\vector c} {\vector K^\top \vector u^{k+1}}
\end{align}
ここでの除算はベクトルの要素ごとの除算である。
この計算では、$\diag\left(\vector u^k\right)\vector K \diag\left(\vector v^k \right)$ の行和、列和がそれぞれ $\vector r, \vector c$ となるように $k+1$ 回目の更新を行なっている。
Sinkhorn Iteration の収束性などについては、以下の論文や記事が詳しい。
- Concerning nonnegative matrices and doubly stochastic matrices
- On the scaling of multidimensional matrices
- 輸送問題の解を微分する - Qiita
この反復の問題点として、計算の不安定さが挙げられる。
$\vector K = \exp\left(-\lambda \vector M\right)$ と指数関数を用いて行列 $\vector K$ を計算している都合上、正則化パラメータ $\lambda$ の値によっては $\vector K$ の値が非常に大きく (小さく) なり、overflow を起こしたり 0 に落ちる可能性がある。
そこで、対数関数を $\vector u, \vector v$ の更新式に作用させる手法が Interpolating between Optimal Transport and MMD using Sinkhorn Divergences などで提案されている。
\begin{align}
\log u^{k+1}_i &= \log r_i - \log\left(\sum_{j=1}^n \exp \left(- \lambda M_{ij} + \log v^{k}_j \right)\right) \\
\log v^{k+1}_j &= \log c_j - \log\left(\sum_{i=1}^m \exp \left(- \lambda M_{ij} + \log u^{k+1}_i \right)\right)
\end{align}
なぜこの更新式が overflow を回避できるかについては、以下の logsumexp についての解説記事などが参考になる。
具体例
EMD の具体例で使用した値 (hist_blue, hist_green, cost_matrix) をそのまま再利用する。
Sinkhorn Iteration と対数関数を用いた安定化法は次のように実装できる。
import numpy as np
from scipy.special import logsumexp
def sinkhorn(r, c, cost_matrix, lambd=1.0, num_iters=100):
K = np.exp(-cost_matrix * lambd)
u = np.ones_like(r)
v = np.ones_like(c)
for _ in range(num_iters):
u = r / (K @ v)
v = c / (K.T @ u)
P = u.reshape(-1, 1) * (K * v.reshape(1, -1))
return P
def sinkhorn_log(r, c, cost_matrix, lambd=1.0, num_iters=100):
K = -cost_matrix * lambd
n_row, n_col = K.shape
logu = np.zeros_like(r)
logv = np.zeros_like(c)
logr = np.log(r)
logc = np.log(c)
for _ in range(num_iters):
logu = logr - logsumexp(K + logv[None, :], axis=1)
logv = logc - logsumexp(K + logu[:, None], axis=0)
P = np.exp(logu.reshape(-1, 1) + K + logv.reshape(1, -1))
return P
例えば、$\lambda = 1000$ として最適化問題を解くと、対数関数による安定化がない場合は warning を吐く。
optim_transport = sinkhorn(hist_blue, hist_green, cost_matrix, lambd=1000, num_iters=100)
# RuntimeWarning: divide by zero encountered in divide
# u = r / (K @ v)
# RuntimeWarning: overflow encountered in divide
# u = r / (K @ v)
# RuntimeWarning: invalid value encountered in matmul
# v = c / (K.T @ u)
optim_transport = sinkhorn_log(hist_blue, hist_green, cost_matrix, lambd=1000, num_iters=100)
Sinkhorn distance も同様に POT に実装がある。
正則化パラメータ reg は上記の実装の lambd とは逆数になっている。
また、対数による安定化は method='sinkhorn_log' で指定できる。
optim_transport = ot.sinkhorn(hist_blue, hist_green, cost_matrix, reg=0.001, numItermax=100, method='sinkhorn_log')
求まった最適解に対し、例えば argmax を取ることで輸送先を一意に決定し、画像の色を変換することができる。
lab_transported = lab_blue.copy()
transport_blue = np.arange(max_val)
transport_blue[nonzero_blue] = nonzero_green[np.argmax(optim_transport, axis=1)]
ab_transported = transport_blue[ab_blue]
lab_transported[..., 1] = ab_transported // 256
lab_transported[..., 2] = ab_transported % 256
transported = cv2.cvtColor(lab_transported.astype(np.uint8), cv2.COLOR_Lab2BGR)
まとめと雑感
もともと Sinkhorn Distance は、その論文のタイトル「Sinkhorn Distances:
Lightspeed Computation of Optimal Transport」 にあるように最適輸送を高速に計算する方法として提案された。
しかし、最適輸送のアルゴリズムの発展により、現在 (2022/12/31) の POT ライブラリでは ot.emd のほうが ot.sinkhorn より高速に動作している。
その一方で、Sinkhorn Distance は「微分可能」な距離指標として深層学習へ応用されており、EMD と Sinkhorn Distance にはそれぞれ利点があるように思える。
余談: エントロピー正則化について
エントロピー正則化の効果を理解するため、以下の最適化問題を考える。
\begin{align}
\minimize_{x \in \mathbb R} \quad & f(x) := -ax - b(1 - x) + \tau \left(x \log x + (1 - x) \log(1 - x) \right) \\
\subjectto \quad & 0 \leq x \leq 1 \\
\where \quad & a, b \in \mathbb R, \tau > 0
\end{align}
これはロジット $[a, b]$ が得られたときのクラス確率 $x$ を求める最適化問題と解釈できる。
この問題のラグランジアンは以下の通り。
\mathcal L(x, \alpha, \beta) = -ax - b(1 - x) + \tau \left(x \log x + (1 - x) \log(1 - x) \right) - \alpha x + \beta (x - 1)
よって、KKT 条件は次のようになる。
\begin{align}
\frac{\partial \mathcal L}{\partial x} = -a + b + \tau \log \frac{x}{1 - x} - \alpha + \beta = 0 \\
x \geq 0, \alpha \geq 0, \alpha x = 0 \\
x \leq 1, \beta \geq 0, \beta (1 - x) = 0
\end{align}
よって、KKT 条件を満たす $x$ は次式で表される。
\begin{align}
\frac x {1 - x} &= \exp \left( \frac 1 \tau ( a - b + \alpha - \beta) \right) \\
x &= \frac{\exp \left( \frac 1 \tau ( a - b + \alpha - \beta) \right)}{1 + \exp \left( \frac 1 \tau ( a - b + \alpha - \beta) \right)}
\end{align}
これより、$0 < x < 1$ であるので、$\alpha = \beta = 0$ となる。
よって、最適解は以下の温度付きソフトマックス関数となる。
x = \frac{\exp \left( \frac 1 \tau (a - b) \right)}{1 + \exp \left( \frac 1 \tau (a - b) \right)} = \frac{\exp \left( \frac 1 \tau a \right)}{\exp \left( \frac 1 \tau a \right) + \exp \left( \frac 1 \tau b \right)}
実際に目的関数をプロットするとこんな感じ。($\tau=0,0.5,1.0, a=0.5, b=-0.5$)
$\tau$ が大きくなるにつれて、最適解がエントロピー最大の点 $x=0.5$ に近づいているのがわかる。
