置換矩陣也能求導優化

置換矩陣也能求導優化

本文是對論文 Learning Latent Permutations with Gumbel-Sinkhorn Networks的閱讀筆記。

很多時候我們都希望學習一個置換矩陣(permutation matrix)，用來找到一個合適的排序，或者解決一個指派問題，就是找到一個最優的分配策略，他可以用匈牙利算法在多項式時間內解決，然後這個問題是不可微的，也就不能放在神經網絡中。他可以形式化的寫作：

$M(X)=\underset{P\in \mathcal{P}_{N}}{\arg\max} \langle P,X\rangle _{F} =\underset{P\in \mathcal{P}_{N}}{\arg\max} tr\left( P^{T} X\right)$

其中X是NxN維矩陣表示每個指派的收益，P是置換矩陣。$\displaystyle \langle \cdot \rangle _{F}$是Frobenius內積。

指派問題與softmax的聯繫

其實仔細想想，置換矩陣不就是相當於每一行每一列都是一個one-hot嗎，而one-hot一個著名的弱化的例子就是softmax了，我們知道當softmax的溫度趨於0的時候，softmax會變成one-hot：

$softmax_{\tau }( x)_{i} =\frac{\exp( x_{i} /\tau )}{\sum _{j}\exp( x_{j} /\tau )}$

那麼他跟指派問題有什麼關係呢？其實我們可以定義一個對行列不停地分別除以每一行的和，以及除以每一列的和的操作，這個操作稱爲Sinkhorn operator S(X) 可以定義如下：

$\begin{aligned} S^{0} (X) & =\exp (X)\\ S^{l} (X) & =\mathcal{T}_{c}\left(\mathcal{T}_{r}\left( S^{l-1} (X)\right)\right)\\ S(X) & =\lim _{l\rightarrow \infty } S^{l} (X) \end{aligned}$

其中$\displaystyle \mathcal{T}_{r} (X)=X\oslash \left( X\mathbf{1}_{N}\mathbf{1}^{\top }_{N}\right) ,\mathcal{T}_{c} (X)=X\oslash \left(\mathbf{1}_{N}\mathbf{1}^{\top }_{N} X\right)$，可以證明S(X)一定會收斂到一個叫Birkhoff polytope的空間上，記爲

$\mathcal{B}_{N} =\left\{P\in [0,1]\in \mathbb{R}^{N,N} P1_{N} =1_{N} ,P^{\top } 1_{N} =1_{N}\right\}$

可以證明$\displaystyle \mathcal{B}_{N}$是包含所有置換矩陣的。於是，指派問題與softmax的聯繫可以用下面這個定理聯繫起來：

定理1： 定義隨機矩陣P的熵爲$\displaystyle h(P)=-\sum _{i.j} p_{ij}\log p_{ij}$，則

$S(X/\tau )=\underset{P\in \mathcal{B}_{N}}{\arg\max} \langle P,X\rangle _{F} +\tau h(P)$

如果X的產生是獨立的，那麼一定有

$M(X)=\lim _{\tau \rightarrow 0^{+}} S(X/\tau )$

現在證明第一個，考慮如下拉格朗日有約束的優化：

$\mathcal{L} (\alpha ,\beta ,P)=\langle P,X\rangle _{F} +\tau h(P)+\alpha ^{\top }( P1_{N} -1_{N}) +\beta ^{\top }\left( P^{\top } 1_{N} -1_{N}\right)$

令$\partial \mathcal{L} /\partial P=0$ 一定有對於每個$i,j$

$p^{i,j}_{\tau } =\exp( \alpha _{i} /\tau -1/2)\exp( X_{i,j} /\tau )\exp( \beta _{j} /\tau -1/2)$

也就是說對於任意正數的對角矩陣$D_{1} ,D_{2} ,$都有$P_{\tau } =D_{1}\exp( X_{i,j} /\tau ) D_{2}$ 根據Sinkhorn’s theorem, 一定有$S(X/\tau )=P_{\tau }$

與softmax的聯繫

實際上，以上的推導只是基於一個softmax的原理簡單的推廣:

$\arg\max_{i} x_{i} =\lim _{\tau \rightarrow 0^{+}}\operatorname{softmax} (x/\tau )$

即是softmax所近似的argmax值，實際上，這個argmax可以表示成一個關於one-hot向量的優化問題，即尋找一個one-hot向量，使得他跟x的內積最大：

$\arg\max_{i} x_{i} =\arg\max_{e} \langle e,x\rangle$

跟上面的證明差不多，我們可以寫出e的帶約束的拉格朗日公式，並對e求導，令它等於0，可以得到e的形式就是一個softmax公式，於是

$\operatorname{softmax} (x/\tau )\equiv \frac{\exp (x/\tau )}{\sum _{i=1}\exp( x_{i} /\tau )} =\underset{e\in \mathcal{S}_{n}}{\arg\max} \langle e,x\rangle +\tau h(e)$

可以簡單證明一下：

$\begin{aligned} \frac{\partial }{\partial e}\left( \langle e,x\rangle +\tau h(e)+\alpha ^{\top }( e1_{N} -1_{N})\right) & =x+\tau \log( e) +\tau +\alpha ^{\top } =0\\ \Longrightarrow & e=\exp( -x/\tau )\exp\left( -1-\alpha ^{\top } /\tau \right) \end{aligned}$

又因爲$\displaystyle \sum _{i} e_{i} =1$,所以一定有$\displaystyle \exp\left( -1-\alpha ^{\top } /\tau \right) =1/\sum _{i=1}\exp( x_{i} /\tau )$.

於是我們理解permutation matrix本質上可以弱化爲softmax組成的矩陣，因此我們完全可以利用gumbel分佈對softmax的重參數化能力來重參數化這個置換矩陣，從而採用VAE來求解一個隱變量時置換矩陣的生成模型，而這就是這篇論文所做的貢獻了。

關於gumbel的重參數可以參考我前面的一篇文章：

帶你認識神奇的Gumbel trick
https://blog.csdn.net/a358463121/article/details/80820878

對softmax的進一步思考

當我們把softmax寫成下面這一條公式的時候，事情就變得越發有趣：
$\operatorname{softmax} (x/\tau )\equiv \frac{\exp (x/\tau )}{\sum _{i=1}\exp( x_{i} /\tau )} =\underset{e\in \mathcal{S}_{n}}{\arg\max} \langle e,x\rangle +\tau h(e)$
我們發現，softmax的指數項的來源，其實是那個$h(e)$產生，因爲裏面有個log，爲了得到偏導數等於0的e，這個log要取個指數項消掉。然而對於熵$h(x)$，其實除了香農熵的經典定義外，還有很多其他的定義的，這是否意味着，softmax還擁有其他的變種？

答案是肯定的，在論文[3]中說明了，只要將香農熵換成，Gini entropy：$\mathrm{H}^{\mathrm{G}}(\boldsymbol{p}):=\frac{1}{2} \sum_{j} p_{j}\left(1-p_{j}\right)$，就能推導出全新的sparse softmax，這個softmax具有sparse的性質：

$\text { sparsemax }(\boldsymbol{x}):=\underset{\boldsymbol{p} \in \Delta^{d}}{\operatorname{argmin}}\|\boldsymbol{p}-\boldsymbol{x}\|^{2}$

實際上，這個sparse softmax是在論文 [2] 中提出的，有了這個思路，其實我們能夠推導出各種各樣的變種，論文就提出了用Tsallis α-entropies:

$\mathrm{H}_{\alpha}^{\mathrm{T}}(\boldsymbol{p}):=\left\{\begin{array}{ll} \frac{1}{\alpha(\alpha-1)} \sum_{j}\left(p_{j}-p_{j}^{\alpha}\right), & \alpha \neq 1 \\ \mathrm{H}^{\mathrm{S}}(\boldsymbol{p}), & \alpha=1 \end{array}\right.$

於是就推導出了entmax
$\alpha \text { -entmax }(\boldsymbol{z}):=\underset{\boldsymbol{p} \in \Delta^{d}}{\operatorname{argmax}} \boldsymbol{p}^{\top} \boldsymbol{z}+\mathrm{H}_{\alpha}^{\mathrm{T}}(\boldsymbol{p})$

這部分有興趣的各位可以去看看[3].

參考文獻

[1] Mena G, Belanger D, Linderman S et.al. Learning Latent Permutations with Gumbel-Sinkhorn Networks[J]. Iclr 2018, 2018(2011): 1–14.
[2] Martins A F T, Astudillo R F. From softmax to sparsemax: A sparse model of attention and multi-label classification[J]. 33rd International Conference on Machine Learning, ICML 2016, 2016, 4: 2432–2443.
[3] Peters B, Niculae V, Martins A F T. Sparse Sequence-to-Sequence Models[J]. 2019.