Sinkhorn Algorithm as a Special Case of Stochastic Mirror Descent - - PowerPoint PPT Presentation

Sinkhorn Algorithm as a Special Case of Stochastic Mirror Descent

Konstantin Mishchenko, KAUST

Problem 1

Let X ∈ Rn×n

++

Find vectors u, v ∈ Rn

+ such that

W = diag(u)Xdiag(v) is doubly stochastic

Problem 1

Let X ∈ Rn×n

++

Let $X\in \mathbb{R}^{n\times n}_{++}$

Find vectors u, v ∈ Rn

+ such that

Find vectors $u, v\in \mathbb{R}_+^n$ such that $W=\mathrm{diag}(u)X\mathrm{diag}(v)$ is doubly stochastic

n

X

i=1

Wij = 1 for any j

\sum_{i=1}^n W_{ij}=1\quad \text{for any } j

n

X

j=1

Wij = 1 for any i

W = diag(u)Xdiag(v) is doubly stochastic

Wij ≥ 0

Problem 2

\min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n C_{ij} X_{ij}\quad \text{s.t. } X\mathbf{1}=\mathbf{1}, X^\top \mathbf{1}=\mathbf{1}, X\ge

doubly stochastic

min

X2Rn×n n

X

i,j=1

CijXij s.t. X1 = 1, X>1 = 1, X ≥ 0

\min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n \left(C_{ij} X_{ij} +\gamma X_{ij}\log X_{ij}\right)\quad \text{s.t. } X\mathbf{1} =\mathbf{1}, X^\top \mathbf{1}=\mathbf{1} \min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n \mathcal{KL} (X||X^0)\quad \text{s.t. } X\mathbf{1}=\mathbf{1}, X^\top \mathbf{1}=\mathbf{1} \text{where }X^0\overset{\text{def}}{=} \exp\left(-\frac{C}{\gamma}\right) \text{ coordinate-wise} \mathcal{KL} (X||X^0) \overset{\text{def}}{=} \sum_{i,j=1}^n \Bigl(X_{ij}\log \frac{X_{ij}}{X^0_{ij}}-X_{ij}+X_{ij}^0 \Bigr)

Problem 2

\min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n C_{ij} X_{ij}\quad \text{s.t. } X\mathbf{1}=\mathbf{1}, X^\top \mathbf{1}=\mathbf{1}, X\ge

min

X2Rn×n n

X

i,j=1

CijXij s.t. X1 = 1, X>1 = 1, X ≥ 0

min

X2Rn×n n

X

i,j=1

(CijXij + γXij log Xij) s.t. X1 = 1, X>1 = 1

\min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n \left(C_{ij} X_{ij} +\gamma X_{ij}\log X_{ij}\right)\quad \text{s.t. } X\mathbf{1} =\mathbf{1}, X^\top \mathbf{1}=\mathbf{1} \min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n \mathcal{KL} (X||X^0)\quad \text{s.t. } X\mathbf{1}=\mathbf{1}, X^\top \mathbf{1}=\mathbf{1} \text{where }X^0\overset{\text{def}}{=} \exp\left(-\frac{C}{\gamma}\right) \text{ coordinate-wise} \mathcal{KL} (X||X^0) \overset{\text{def}}{=} \sum_{i,j=1}^n \Bigl(X_{ij}\log \frac{X_{ij}}{X^0_{ij}}-X_{ij}+X_{ij}^0 \Bigr)

Problem 2

\min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n C_{ij} X_{ij}\quad \text{s.t. } X\mathbf{1}=\mathbf{1}, X^\top \mathbf{1}=\mathbf{1}, X\ge

min

X2Rn×n n

X

i,j=1

CijXij s.t. X1 = 1, X>1 = 1, X ≥ 0

min

X2Rn×n n

X

i,j=1

(CijXij + γXij log Xij) s.t. X1 = 1, X>1 = 1

\min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n \left(C_{ij} X_{ij} +\gamma X_{ij}\log X_{ij}\right)\quad \text{s.t. } X\mathbf{1} =\mathbf{1}, X^\top \mathbf{1}=\mathbf{1} \min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n \mathcal{KL} (X||X^0)\quad \text{s.t. } X\mathbf{1}=\mathbf{1}, X^\top \mathbf{1}=\mathbf{1}

where X0 def = exp ✓ −C γ ◆ coordinate-wise

\text{where }X^0\overset{\text{def}}{=} \exp\left(-\frac{C}{\gamma}\right) \text{ coordinate-wise}

KL(X||X0)

def

=

n

X

i,j=1

⇣ Xij log Xij X0

ij

− Xij + X0

ij

⌘

\mathcal{KL} (X||X^0) \overset{\text{def}}{=} \sum_{i,j=1}^n \Bigl(X_{ij}\log \frac{X_{ij}}{X^0_{ij}}-X_{ij}+X_{ij}^0 \Bigr)

min

X2Rn×n KL(X||X0)

s.t. X1 = 1, X>1 = 1

Problem 2

\min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n C_{ij} X_{ij}\quad \text{s.t. } X\mathbf{1}=\mathbf{1}, X^\top \mathbf{1}=\mathbf{1}, X\ge

min

X2Rn×n n

X

i,j=1

CijXij s.t. X1 = 1, X>1 = 1, X ≥ 0

min

X2Rn×n n

X

i,j=1

(CijXij + γXij log Xij) s.t. X1 = 1, X>1 = 1

\min_{X\in\mathbb{R}^{n\times n}} \sum_{i,j=1}^n \left(C_{ij} X_{ij} +\gamma X_{ij}\log X_{ij}\right)\quad \text{s.t. } X\mathbf{1} =\mathbf{1}, X^\top \mathbf{1}=\mathbf{1} \min_{X\in\mathbb{R}^{n\times n}} \mathcal{KL} (X||X^0)\quad \text{s.t. } X\mathbf{1}=\mathbf{1}, X^\top \mathbf{1}=\mathbf{1}

where X0 def = exp ✓ −C γ ◆ coordinate-wise

\text{where }X^0\overset{\text{def}}{=} \exp\left(-\frac{C}{\gamma}\right) \text{ coordinate-wise}

KL(X||X0)

def

=

n

X

i,j=1

⇣ Xij log Xij X0

ij

− Xij + X0

ij

⌘

\mathcal{KL} (X||X^0) \overset{\text{def}}{=} \sum_{i,j=1}^n \Bigl(X_{ij}\log \frac{X_{ij}}{X^0_{ij}}-X_{ij}+X_{ij}^0 \Bigr)

min

X2Rn×n KL(X||X0)

s.t. X1 = 1, X>1 = 1

Problem 3

\min_x\Bigl\{\mathcal{KL} (Ax||b) =\sum_{i}^n \Bigl(\langle a_i, x\rangle\log \frac{\langle a_i, x\rangle}{b_i}-\langle a_i, x\rangle+b_i \Bigr) \Bigr\}

x = vec(X) ∈ Rd, d = n · n

a1 = (1, 1, . . . , 1, 0, . . . , 0)

· · ·

ai = (0, . . . , 0, 1, . . . , 1 | {z }

i−th block

, 0, . . . , 0)

a_i = (0, \dotsc, 0, \underbrace{1, \dotsc, 1}_{i-\text{th block}}, 0, \dotsc, 0)t a_1 = (1, 1, \dotsc, 1, 0, \dotsc, 0) x=\mathrm{vec}(X)\in\mathbb{R}^d,\ d=n\cdot n X^\top \mathbf{1}=\mathbf{1}

X1 = 1

Problem 3

\min_x\Bigl\{\mathcal{KL} (Ax||b) =\sum_{i}^n \Bigl(\langle a_i, x\rangle\log \frac{\langle a_i, x\rangle}{b_i}-\langle a_i, x\rangle+b_i \Bigr) \Bigr\}

x = vec(X) ∈ Rd, d = n · n

a1 = (1, 1, . . . , 1, 0, . . . , 0)

· · ·

ai = (0, . . . , 0, 1, . . . , 1 | {z }

i−th block

, 0, . . . , 0)

a_i = (0, \dotsc, 0, \underbrace{1, \dotsc, 1}_{i-\text{th block}}, 0, \dotsc, 0)t a_1 = (1, 1, \dotsc, 1, 0, \dotsc, 0) x=\mathrm{vec}(X)\in\mathbb{R}^d,\ d=n\cdot n X^\top \mathbf{1}=\mathbf{1}

X1 = 1

same for X>1 = 1

A =      a1 a2 . . . am     

Problem 3

\min_x\Bigl\{\mathcal{KL} (Ax||b) =\sum_{i}^{2n} \Bigl(\langle a_i, x\rangle\log \frac{\langle a_i, x\rangle}{b_i}-\langle a_i, x\rangle+b_i \Bigr) \Bigr\}

x = vec(X) ∈ Rd, d = n · n

a1 = (1, 1, . . . , 1, 0, . . . , 0)

· · ·

ai = (0, . . . , 0, 1, . . . , 1 | {z }

i−th block

, 0, . . . , 0)

a_i = (0, \dotsc, 0, \underbrace{1, \dotsc, 1}_{i-\text{th block}}, 0, \dotsc, 0)t a_1 = (1, 1, \dotsc, 1, 0, \dotsc, 0) x=\mathrm{vec}(X)\in\mathbb{R}^d,\ d=n\cdot n X^\top \mathbf{1}=\mathbf{1}

X1 = 1

same for X>1 = 1

A= \begin{pmatrix}a_1 \\ a_2 \\ \vdots \\ a_m\end{pmatrix}

min

x

n KL(Ax||b) =

2n

X

i

⇣ hai, xi log hai, xi bi hai, xi + bi ⌘o

One method to solve them all

a_i = (0, \dotsc, 0, \underbrace{1, \dotsc, 1}_{i-\text{th block}}, 0, \dotsc, 0)t a_1 = (1, 1, \dotsc, 1, 0, \dotsc, 0) x=\mathrm{vec}(X)\in\mathbb{R}^d,\ d=n\cdot n

W = diag(u)Xdiag(v) is doubly stochastic

• 1. Normalize all rows

• 2. Normalize all columns

Loop

{

X^{k+1}=\arg\min_{X\mathbf{1}=\mathbf{1}}\left\{\mathcal{KL}(X||X^k) \right\} X^{k+2}=\arg\min_{X^\top\mathbf{1}=\mathbf{1}}\left\{\mathcal{KL}(X||X^{k+1}) \right\} f_1(x)=\mathcal{KL}(A_{\text{rows}}x||\mathbf{1}),\ f_2(x)=\mathcal{KL}(A_{\text{cols}}x||\mathbf{1})

One method to solve them all

a_i = (0, \dotsc, 0, \underbrace{1, \dotsc, 1}_{i-\text{th block}}, 0, \dotsc, 0)t a_1 = (1, 1, \dotsc, 1, 0, \dotsc, 0) x=\mathrm{vec}(X)\in\mathbb{R}^d,\ d=n\cdot n

• 1. Normalize all rows

• 2. Normalize all columns

Loop

{

Xk+1 = arg min

X1=1

• KL(X||Xk)

X^{k+1}=\arg\min_{X\mathbf{1}=\mathbf{1}}\left\{\mathcal{KL}(X||X^k) \right\}

Xk+2 = arg min

X>1=1

• KL(X||Xk+1)

{

Loop

X^{k+2}=\arg\min_{X^\top\mathbf{1}=\mathbf{1}}\left\{\mathcal{KL}(X||X^{k+1}) \right\} f_1(x)=\mathcal{KL}(A_{\text{rows}}x||\mathbf{1}),\ f_2(x)=\mathcal{KL}(A_{\text{cols}}x||\mathbf{1})

min

X2Rn×n KL(X||X0)

s.t. X1 = 1, X>1 = 1

One method to solve them all

a_i = (0, \dotsc, 0, \underbrace{1, \dotsc, 1}_{i-\text{th block}}, 0, \dotsc, 0)t a_1 = (1, 1, \dotsc, 1, 0, \dotsc, 0) x=\mathrm{vec}(X)\in\mathbb{R}^d,\ d=n\cdot n

• 1. Normalize all rows

• 2. Normalize all columns

Loop

{

Xk+1 = arg min

X1=1

• KL(X||Xk)

X^{k+1}=\arg\min_{X\mathbf{1}=\mathbf{1}}\left\{\mathcal{KL}(X||X^k) \right\}

Xk+2 = arg min

X>1=1

• KL(X||Xk+1)

{

Loop

X^{k+2}=\arg\min_{X^\top\mathbf{1}=\mathbf{1}}\left\{\mathcal{KL}(X||X^{k+1}) \right\}

f1(x) = KL(Arowsx||1), f2(x) = KL(Acolsx||1)

f_1(x)=\mathcal{KL}(A_{\text{rows}}x||\mathbf{1}),\ f_2(x)=\mathcal{KL}(A_{\text{cols}}x||\mathbf{1})

min

x

n KL(Ax||b) =

2n

X

i

⇣ hai, xi log hai, xi bi hai, xi + bi ⌘o

\log (x^{k+1}) = \log (x^k) - \eta \nabla f_i(x^k),\ i\sim U(\{1, 2\})

log(xk+1) = log(xk) ηrfi(xk), i ⇠ U({1, 2})

Intuition

\min_x \frac{1}{2}\|x-x^0\|^2\ \text{s.t. } Ax=b

min

x

1 2kx x0k2 s.t. Ax = b

xk+1 = Πhai,xi=bi(xk), i ∼ U({1, . . . , m})

x^k \in x^0 + \mathrm{Range}(A^\top) x^{k+1}=\Pi_{\langle a_i, x\rangle = b_i}(x^k),\ i\sim U(\{1,\dotsc, m\})

Kaczmarz Algorithm

x^{k+1}=x^k+\frac{b_i-\langle a_i, x\rangle}{\|a_i\|^2}a_i

Intuition

\min_x \frac{1}{2}\|x-x^0\|^2\ \text{s.t. } Ax=b

min

x

1 2kx x0k2 s.t. Ax = b

xk+1 = Πhai,xi=bi(xk), i ∼ U({1, . . . , m})

xk ∈ x0 + Range(A>)

x^k \in x^0 + \mathrm{Range}(A^\top) x^{k+1}=\Pi_{\langle a_i, x\rangle = b_i}(x^k),\ i\sim U(\{1,\dotsc, m\})

Kaczmarz Algorithm

x^{k+1}=x^k+\frac{b_i-\langle a_i, x\rangle}{\|a_i\|^2}a_i

xk+1 = xk + bi hai, xi kaik2 ai

x^{k+1}=x^k-\eta \nabla f_i(x^k),\quad f_i(x) = \frac{1}{2\|a_i\|^2}(\langle a_i, x\rangle - b_i)^2

xk+1 = xk ηrfi(xk), fi(x) = 1 2kaik2 (hai, xi bi)2

Sinkhorn Algorithm

a_i = (0, \dotsc, 0, \underbrace{1, \dotsc, 1}_{i-\text{th block}}, 0, \dotsc, 0)t a_1 = (1, 1, \dotsc, 1, 0, \dotsc, 0) x=\mathrm{vec}(X)\in\mathbb{R}^d,\ d=n\cdot n

f1(x) = KL(Arowsx||1), f2(x) = KL(Acolsx||1)

min

x

n KL(Ax||b) =

2n

X

i

⇣ hai, xi log hai, xi bi hai, xi + bi ⌘o

\log x^k \in \log x^0 + \textrm{Range}(A^\top)

log(xk+1) = log(xk) ηrfi(xk), i ⇠ U({1, 2})

Sinkhorn Algorithm

a_i = (0, \dotsc, 0, \underbrace{1, \dotsc, 1}_{i-\text{th block}}, 0, \dotsc, 0)t a_1 = (1, 1, \dotsc, 1, 0, \dotsc, 0) x=\mathrm{vec}(X)\in\mathbb{R}^d,\ d=n\cdot n

f1(x) = KL(Arowsx||1), f2(x) = KL(Acolsx||1)

min

x

n KL(Ax||b) =

2n

X

i

⇣ hai, xi log hai, xi bi hai, xi + bi ⌘o

log xk ∈ log x0 + Range(A>)

\log x^k \in \log x^0 + \textrm{Range}(A^\top)

Arows, Acols ∈ {0, 1}n×n

Xk+1 = diag(uk)X0diag(vk)

X^{k+1}=\mathrm{diag}(u^k)X^0\mathrm{diag}(v^k)

log(xk+1) = log(xk) ηrfi(xk), i ⇠ U({1, 2})

Reference

• R. Sinkhorn, Diagonal equivalence to matrices with

prescribed row and column sums, 1967

• K. M., Sinkhorn Algorithm as a Special

Case of Stochastic Mirror Descent, 2019 arXiv:1909.06918 NeurIPS workshop on Optimal Transport and its Application to Machine Learning Canada, December 2019 https://sites.google.com/view/otml2019/