https://github.com/dfdazac/wassdistance/tree/master
前置知识
Computational optimal transport学习
具体看到熵对偶的坐标上升那就行
L
C
ε
(
a
,
b
)
=
def.
min
P
∈
U
(
a
,
b
)
⟨
P
,
C
⟩
−
ε
H
(
P
)
\mathrm{L}_{\mathbf{C}}^{\varepsilon}(\mathbf{a}, \mathbf{b}) \stackrel{\text { def. }}{=} \min _{\mathbf{P} \in \mathbf{U}(\mathbf{a}, \mathbf{b})}\langle\mathbf{P}, \mathbf{C}\rangle-\varepsilon \mathbf{H}(\mathbf{P})
LCε(a,b)= def. P∈U(a,b)min⟨P,C⟩−εH(P)
U
(
a
,
b
)
=
def.
{
P
∈
R
+
n
×
m
:
P
1
m
=
a
and
P
T
1
n
=
b
}
\mathbf{U}(\mathbf{a}, \mathbf{b}) \stackrel{\text { def. }}{=}\left\{\mathbf{P} \in \mathbb{R}_{+}^{n \times m}: \mathbf{P} \mathbf{1}_m=\mathbf{a} \quad \text { and } \quad \mathbf{P}^{\mathrm{T}} \mathbf{1}_n=\mathbf{b}\right\}
U(a,b)= def. {P∈R+n×m:P1m=a and PT1n=b}
对偶
L
C
ε
(
a
,
b
)
=
max
f
∈
R
n
,
g
∈
R
m
⟨
f
,
a
⟩
+
⟨
g
,
b
⟩
−
ε
⟨
e
f
/
ε
,
K
e
g
/
ε
⟩
\mathrm{L}_{\mathbf{C}}^{\varepsilon}(\mathbf{a}, \mathbf{b})=\max _{\mathbf{f} \in \mathbb{R}^n, \mathbf{g} \in \mathbb{R}^m}\langle\mathbf{f}, \mathbf{a}\rangle+\langle\mathbf{g}, \mathbf{b}\rangle-\varepsilon\left\langle e^{\mathbf{f} / \varepsilon}, \mathbf{K} e^{\mathbf{g} / \varepsilon}\right\rangle
LCε(a,b)=f∈Rn,g∈Rmmax⟨f,a⟩+⟨g,b⟩−ε⟨ef/ε,Keg/ε⟩
(
u
,
v
)
=
(
e
f
/
ε
,
e
g
/
ε
)
(\mathbf{u}, \mathbf{v})=\left(e^{\mathbf{f} / \varepsilon}, e^{\mathbf{g} / \varepsilon}\right)
(u,v)=(ef/ε,eg/ε)
P = d i a g ( u ) K d i a g ( v ) , K = e x p ( − C ϵ ) \mathbf{P}=\rm{diag}\left(\mathbf{u}\right)\mathbf{K}\rm{diag}\left(\mathbf{v}\right),\quad \mathbf{K}=exp\left(-\frac{C}{\epsilon}\right) P=diag(u)Kdiag(v),K=exp(−ϵC)
坐标上升
f
(
ℓ
+
1
)
=
ε
log
a
−
ε
log
(
K
e
g
(
ℓ
)
/
ε
)
,
g
(
ℓ
+
1
)
=
ε
log
b
−
ε
log
(
K
T
e
f
(
ℓ
+
1
)
/
ε
)
.
\begin{aligned} \mathbf{f}^{(\ell+1)} & =\varepsilon \log \mathbf{a}-\varepsilon \log \left(\mathbf{K} e^{\mathbf{g}^{(\ell)} / \varepsilon}\right), \\ \mathbf{g}^{(\ell+1)} & =\varepsilon \log \mathbf{b}-\varepsilon \log \left(\mathbf{K}^{\mathrm{T}} e^{\mathbf{f}^{(\ell+1)} / \varepsilon}\right) . \end{aligned}
f(ℓ+1)g(ℓ+1)=εloga−εlog(Keg(ℓ)/ε),=εlogb−εlog(KTef(ℓ+1)/ε).
代码中有一些变化
考虑 C ∈ R n × m , f ∈ R n , g ∈ R m \mathbf{C}\in\mathbb{R}^{n\times m}, \mathbf{f}\in\mathbb{R}^n, \mathbf{g}\in\mathbb{R}^m C∈Rn×m,f∈Rn,g∈Rm
log
(
K
e
g
/
ε
)
=
log
(
[
∑
j
e
−
C
i
,
j
−
g
j
ε
]
i
)
=
log
(
[
∑
j
e
−
C
i
,
j
−
g
j
ε
e
f
i
ε
e
−
f
i
ε
]
i
)
=
log
(
[
∑
j
e
−
C
i
,
j
−
f
i
−
g
j
ε
]
i
⊙
e
−
f
ε
)
=
log
(
[
∑
j
e
−
C
i
,
j
−
f
i
−
g
j
ε
]
i
)
−
f
ε
=
logsumexp
(
−
C
−
f
T
−
g
ε
,
d
i
m
=
−
1
)
−
f
ε
\begin{aligned} &\log \left(\mathbf{K} e^{\mathbf{g} / \varepsilon}\right)\\ =&\log\left(\left[\sum_{j}e^{-\frac{C_{i,j}-g_j}{\varepsilon}}\right]_i\right)\\ =&\log\left(\left[\sum_{j}e^{-\frac{C_{i,j}-g_j}{\varepsilon}}e^{\frac{f_i}{\varepsilon}}e^{-\frac{f_i}{\varepsilon}}\right]_i\right)\\ =&\log\left(\left[\sum_{j}e^{-\frac{C_{i,j}-f_i-g_j}{\varepsilon}}\right]_i\odot e^{-\frac{\mathbf{f}}{\varepsilon}}\right)\\ =&\log\left(\left[\sum_{j}e^{-\frac{C_{i,j}-f_i-g_j}{\varepsilon}}\right]_i\right)-\frac{\mathbf{f}}{\varepsilon}\\ =&\operatorname{logsumexp}\left(-\frac{\mathbf{C}-\mathbf{f}^T-\mathbf{g}}{\varepsilon},dim=-1\right)-\frac{\mathbf{f}}{\varepsilon}\\ \end{aligned}
=====log(Keg/ε)log
[j∑e−εCi,j−gj]i
log
[j∑e−εCi,j−gjeεfie−εfi]i
log
[j∑e−εCi,j−fi−gj]i⊙e−εf
log
[j∑e−εCi,j−fi−gj]i
−εflogsumexp(−εC−fT−g,dim=−1)−εf
其中最后一步，向量和矩阵相加涉及广播机制
log
(
K
T
e
f
/
ε
)
=
log
(
[
∑
i
e
−
C
i
,
j
−
f
i
ε
]
j
)
=
log
(
[
∑
i
e
−
C
i
,
j
−
f
i
ε
e
g
j
ε
e
−
g
j
ε
]
j
)
=
log
(
[
∑
i
e
−
C
i
,
j
−
f
i
−
g
j
ε
]
j
⊙
e
−
g
ε
)
=
log
(
[
∑
i
e
−
C
i
,
j
−
f
i
−
g
j
ε
]
j
)
−
g
ε
=
logsumexp
(
−
C
−
f
T
−
g
ε
,
d
i
m
=
−
2
)
−
g
ε
=
logsumexp
(
−
(
C
−
f
T
−
g
)
T
ε
,
d
i
m
=
−
1
)
−
g
ε
\begin{aligned} &\log \left(\mathbf{K}^{\mathrm{T}} e^{\mathbf{f} / \varepsilon}\right)\\ =&\log\left(\left[\sum_{i}e^{-\frac{C_{i,j}-f_i}{\varepsilon}}\right]_j\right)\\ =&\log\left(\left[\sum_{i}e^{-\frac{C_{i,j}-f_i}{\varepsilon}}e^{\frac{g_j}{\varepsilon}}e^{-\frac{g_j}{\varepsilon}}\right]_j\right)\\ =&\log\left(\left[\sum_{i}e^{-\frac{C_{i,j}-f_i-g_j}{\varepsilon}}\right]_j\odot e^{-\frac{\mathbf{g}}{\varepsilon}}\right)\\ =&\log\left(\left[\sum_{i}e^{-\frac{C_{i,j}-f_i-g_j}{\varepsilon}}\right]_j\right)-\frac{\mathbf{g}}{\varepsilon}\\ =&\operatorname{logsumexp}\left(-\frac{\mathbf{C}-\mathbf{f}^T-\mathbf{g}}{\varepsilon},dim=-2\right)-\frac{\mathbf{g}}{\varepsilon}\\ =&\operatorname{logsumexp}\left(-\frac{\left(\mathbf{C}-\mathbf{f}^T-\mathbf{g}\right)^T}{\varepsilon},dim=-1\right)-\frac{\mathbf{g}}{\varepsilon}\\ \end{aligned}
======log(KTef/ε)log
[i∑e−εCi,j−fi]j
log
[i∑e−εCi,j−fieεgje−εgj]j
log
[i∑e−εCi,j−fi−gj]j⊙e−εg
log
[i∑e−εCi,j−fi−gj]j
−εglogsumexp(−εC−fT−g,dim=−2)−εglogsumexp(−ε(C−fT−g)T,dim=−1)−εg
