Sinkhorn Distance の計算方法を最適化問題から見て、Python でその動作を確認する。

Sinkhorn Distance は最適輸送 (Optimal Transport) にエントロピー正則化を適用したものである。
そこで、まず最適輸送問題について具体例とともに確認する。

参考論文

Sinkhorn Distances: Lightspeed Computation of Optimal Transport

同論文の解説記事

Sinkhorn Distances: Optimal Transport with Entropic Constraintsを読んだのでメモ - 機械学習とかコンピュータビジョンとか

\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 の収束性などについては、以下の論文や記事が詳しい。

この反復の問題点として、計算の不安定さが挙げられる。
$\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 についての解説記事などが参考になる。

logsumexpとは｜Numpy・PyTorchによる実装例も解説！

具体例

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$ に近づいているのがわかる。