はじめに

LEARNING LATENT PERMUTATIONS WITH GUMBEL-SINKHORN NETWORKSを読んだのでメモ． Permutationやmatchingの問題をsinkhorn normalizationを利用してend-to-endに解くモデルを提案． Sinkhorn iterationは以前ブログにまとめた Deep Learning on Graph Matchingの論文でも利用されていた，二重確率行列(行，もしくは列の和が1になる行列)を生成するテクニック．これにカテゴリカル分布から微分可能な形でサンプリングするテクニックであるgumbel softmax(またはconcrete distribution)を組み合わせることでpermutation matrixの潜在表現を取得可能にした．従来と異なる点として，permutationを教師ありで学習するのではなく，潜在表現として教師なしで獲得できる．

簡単に実装してみた．

Sinkhorn operator

従来，onehot表現の緩和として，温度パラメータを導入したsoftmax $\mathrm{softmax}_\tau(x)_i=\exp(x_i/\tau)/\sum_{j=1}\exp(x_j/\tau)$ が提案されていて，温度 $\tau$ が $\tau\rightarrow0$ のとき(確率がuniformの場合を除き)onehotベクトルとなる．ここではこの考え方を二重確率行列へと応用する．

そもそも二重確率行列へのは正規化はかなり難しい問題で，ここではsinkhorn normalizationというよく知られた方法を利用する．アルゴリズム自体は単純で，各要素が非負の行列に対し，行方向，列方向の正規化を繰り返すことで二重確率行列へと収束するというもの．ここでは任意の正方行列 $X$ に対し，sinkhorn normalizationを次のように定義する．

$\displaystyle S^0(X)=\exp(X),\\ S^l(X)=\mathcal{T}_c(\mathcal{T}_r(S^{l-1}(X))),\\ S(X)=\lim_{l\rightarrow\infty}S^l(X)$

$\mathcal{T}_r(X)=X\oslash(X\mathbf{1}_N\mathbf{1}^\top_N)$ と $\mathcal{T}_c(X)=X\oslash(\mathbf{1}_N\mathbf{1}^\top_NX)$ はそれぞれ，行方向，列方向の正規化を意味する．ただし， $\oslash$ を要素毎の割り算とし， $\mathbf{1}_N$ を各要素が1の $N$ 次元列ベクトルとする． $S(X)$ はSiknkhorn(1964)によって二重確率行列の集合(Birkhoff polytope)に属するが証明されており，ここではこの集合を $\mathcal{B}_N$ とする．

Permutationの問題はカテゴリを選ぶのと同様に最大化問題として書くことができ，permutation matrixの集合を $\mathcal{P}_N$ とし，行列のフロベニウス内積を $\langle A,B\rangle_F=\mathrm{trace}(A^\top B)$ とすると，次のように表現できる．

$\displaystyle M(X)=\underset{P\in\mathcal{P}_N}{\mathrm{arg}\max}\langle A,B\rangle_F$

ここで $M(\cdot)$ をmatching operatorと呼ぶ．このmatching operator $M(X)$ とsinkhorn normalizationの極限 $S(X/\tau)$ は $M(X)\approx S(X/\tau)$ の関係があるというのがこの論文の主張．すなわち，次の定理が成り立つということが主張．

Theorem 1.

For a doubly-stochastic matrix $P$ , define its entropy as $h(P)=-\sum_{i,j}P_{i,j}\log(P_{i,j})$ . Then, one has,

$\displaystyle S(X/\tau)=\underset{P\in\mathcal{B}_N}{\mathrm{arg}\max}\langle P,X\rangle_F+\tau h(P)$

Now, assume also the entries of $X$ are drawn independently from a distribution that is absolutely continuous with respect to the Lebesgue measure in $\mathbb{R}$ . Then, almost surely, the following convergence holds:

$\displaystyle M(X)=\lim_{\tau\rightarrow0^+}S(X/\tau)$

Proof

証明のための細かな定義から．まず $\mathcal{B}_N$ と $\mathcal{P}_n$ を次のように定義する．

$\displaystyle \mathcal{B}_N=\{P\in[0,1]\in\mathbb{R}^{N,N}P1_N=1_N,P^\top1_N=1_N\},\\ \mathcal{P}_N=\{P\in\{0,1\}\in\mathbb{R}^{N,N}P1_N=1_N,P^\top1_N=1_N\}$

それぞれ離散か連続かの違い．また，assingment problemを $\mathcal{B}_N$ 上と $\mathcal{P}_N$ 上とでそれぞれ次のように定義する．

$\displaystyle M(X)\equiv\underset{P\in\mathcal{P}_N}{\mathrm{arg}\max}\langle P,X\rangle_F,\\ \tilde{M}(X)\equiv\underset{P\in\mathcal{B}_N}{\mathrm{arg}\max}\langle P,X\rangle_F$

基本的に $M(X)\subseteq\tilde{M}(X)$ が成り立つ．以下のLemma 1からLemma 3を利用してTheorem 1を証明する．各命題の証明は関連文献等が多いのでここでは省略する．

Lemma 1

$\displaystyle S(X/\tau)=\underset{P\in\mathcal{B}_N}{\mathrm{arg}\max}\langle P,X\rangle_F+\tau h(P)$

この最大化問題はユニークな解 $P_\tau$ が存在する．

Lemma 2

Suppose the entries of $X$ are drawn independently from a distribution that is absolutely continuous with respect to the Lebesgue measure in $\mathbb{R}$ . Then, almost surely, $\tilde{M}(X)=M(X)$ is a unique permutation matrix.

Lemma 3

Call $P_\tau$ the solution to the problem, $S(X/\tau)=\underset{P\in\mathcal{B}_N}{\mathrm{arg}\max}\langle P,X\rangle_F+\tau h(P)$ . Under the assumptions of Lemma 2, $P_\tau\rightarrow P_0$ when if $\tau\rightarrow0^+$ .

Theorem 1の最初の部分はLemma 1そのもの．極限での収束はLemma 3から得られる．

Sinkhorn Networks

ここではpermutationの潜在表現を学習するsinkhorn networkの最小構成について述べる．まず真の順列で並んだデータ $X$ とランダムに並べ替えられた $\tilde{X}$ のペアが与えられたとする．ここではあるpermutation行列によって $X_i=P^{-1}_{\theta,\tilde{X}_i}\tilde{X}_i+\epsilon_i$ の関係が成り立つとする．ただし， $\epsilon_i$ はノイズ項で， $P_{\theta,\tilde{X}_i}$ は $\tilde{X}_i$ と $\theta$ をパラメータとした $X_i$ から $\tilde{X}_i$ への写像．よって，以下の最小化問題を考える．

$\displaystyle f(\theta,X,\tilde{X})=\sum_{i=1}^M\|X_i-P^{-1}_{\theta,\tilde{X}_i}\tilde{X}_i\|^2$

$P_{\theta,\tilde{X}_i}$ の表現としてここではニューラルネットワークを考える．ネットワークは $\tilde{X}_i$ を入力とし， $P_{\theta,\tilde{X}_i}$ を出力する．この時 $\theta$ はニューラルネットの重みとバイアス項の集合 $\theta=\{(W_h,b_h)\}_h$ として定義される．ただし， $h$ はネットワークの層のインデックスを表す．最終的なpermuttation $P_{\theta,\tilde{X}_i}$ は割り当て問題の解として得られるが， $P_{\theta,\tilde{X}_i}=M(g(\tilde{X},\theta))$ は微分が不可能であるため，割り当て問題のソルバ $M(\cdot)$ の代わりにsinkhorn normalization $S(g(\tilde{X},\theta)/\tau)$ で置き換え，微分可能なものとして定義する．これにより，誤差逆伝播を用いてend-to-endに $\theta$ を学習することが可能となる．

Permutation equivariance

ここでは問題の性質上，ネットワークは $\tilde{X}$ のpermutation $P^\prime$ に対して次の関係を満たすことが好ましい．

$\displaystyle P_{\theta,P^\prime \tilde{X}}(P^\prime\tilde{X})=P^\prime(P_{\theta,\tilde{X}}\tilde{X})$

これはネットワークの出力とpermutation行列に $g(\tilde{X}_i,\theta)=P_i$ という関係（すなわち入力の $i$ 番目の要素が出力の $i$ 行目にのみ影響を与える）を満たすようにすればよい．なので，ネットワークは与えられた集合 $\tilde{X}$ の要素を独立に扱い対応するpermutation行列の各行を出力し，全ての行をstackすることで $P_{\theta,\tilde{X}_i}$ を作る．より具体的には論文のFigure 1を参照．

Probabilistic aspects

Sinkhorn networkの確率モデルとしての見方を与える．ポテンシャル関数 $X:\mathcal{Y}\rightarrow\mathbb{R}$ での離散集合 $\mathcal{Y}$ 上の分布は

$\displaystyle p(y|X)\propto\exp(X(y))\mathbf{1}_{y\in\mathcal{Y}}$

として表現される． Gumbel max trickによりこの分布からのサンプルはi.i.d.のガンベルノイズ $\gamma(y)$ を用いて $\mathrm{arg}\max_{y\in\mathcal{Y}}\{X(y)+\gamma(y)\}\sim p(\cdot|X)$ として計算できる．

一方で，今回の問題では集合の大きさが $|\mathcal{Y}|=N!$ となっているため扱いにくい．なので $\mathcal{Y}=\prod_{i=1}^N\mathcal{Y}_i$ のように分解し，さらにrank-oneのノイズ $\gamma(y)=\sum_{i=1}^N\gamma_i(y_i)$ を使うことで扱いやすい形に変える．

permutation $P$ は $X$ をパラメータとするlinear potential $X(P)=\langle X,P\rangle_F$ のGumbel-Matching分布 $\mathcal{G.M.}(X)$ に従うものとする． $\epsilon$ がi.i.d.のGumbel noiseとした時 $M(X+\epsilon)\sim\mathcal{G.M.}(X)$ とすることができる．ただ，Gumbel-Matching distributionからのサンプリングは微分不可能であるため，ニューラルネットとの相性が悪い．ここでは新たにパラメータ $X$ と温度 $\tau$ を持つGumbel-Sinkhorn distributionを考える．すなわち $P\sim\mathcal{G.S.}(X,\tau),P=S( (X+\epsilon)/\tau)$ ．

カテゴリカル分布のGumbel-softmax等による連続緩和と違い，Gumbel-MatchingとGumbel-Sinkhorn distributionは扱いにくい．一方で近年提案されたlikelihood-free methodsを使えば推定することは不可能ではない．観測 $Y$ と潜在変数 $Z=\{P,W\}$ を考える．ただし， $P$ はpermutationで $W$ はその他の変数とする．この時事後分布 $p(\{p,W\}|Y)$ を以下のELBOの最大化によって近似する方法を考える．

$\displaystyle p(y)\geq E_{q(Z|Y)}(\log p(Y|Z))-KL(q(Z|Y)\| p(Z))$

ここでは平均場近似により事前分布と変分事後分布が $q(\{P,W\}|Y)=q(P)q(W),p(P,W)=p(P)p(W)$ のように分解可能であることを仮定する．これにより $W$ を無視することができ一般性を失うことなく $Z=P$ であることを仮定できる．

ここでは $P$ はGumbel-Matching分布に従うが，微分可能にするためGumbel-Sinkhorn分布 $\mathcal{G.S.}(X,\tau)$ へ置き換える．この時，事前分布は $\mathcal{G.S.}(X=0,\tau_{prior}$ ，変分事後分布は $\mathcal{G.S.}(X,\tau)$ とする．

$\mathcal{G.S.}$ の陽な表現がないため，ELBOのKL項を計算することができない．そこでreparameterizationにより $S((X+\epsilon)/\tau)\sim\mathcal{G.S.}(X,\tau)$ ， $S(\epsilon/\tau_{prior})\sim\mathcal{G.S.}(X=0,\tau_{prior})$ とし，KL項を $KL((X+\epsilon)/\tau)\|\epsilon/\tau_{prior})$ として計算をする．このKLの具体的計算はAppendixを参照．