置换矩阵也能求导优化

置换矩阵也能求导优化

本文是对论文 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.