置换矩阵也能求导优化
置换矩阵也能求导优化
本文是对论文 Learning Latent Permutations with Gumbel-Sinkhorn Networks的阅读笔记。
很多时候我们都希望学习一个置换矩阵(permutation matrix)，用来找到一个合适的排序，或者解决一个指派问题，就是找到一个最优的分配策略，他可以用匈牙利算法在多项式时间内解决，然后这个问题是不可微的，也就不能放在神经网络中。他可以形式化的写作：
M ( X ) = arg max P ∈ P N ⟨ P , X ⟩ F = arg max P ∈ P N t r ( P T X ) 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) M(X)=P∈PNargmax⟨P,X⟩F=P∈PNargmaxtr(PTX)
其中X是NxN维矩阵表示每个指派的收益，P是置换矩阵。 ⟨ ⋅ ⟩ F \displaystyle \langle \cdot \rangle _{F} ⟨⋅⟩F是Frobenius内积。
指派问题与softmax的联系
其实仔细想想，置换矩阵不就是相当于每一行每一列都是一个one-hot吗，而one-hot一个著名的弱化的例子就是softmax了，我们知道当softmax的温度趋于0的时候，softmax会变成one-hot：
s o f t m a x τ ( x ) i = exp ( x i / τ ) ∑ j exp ( x j / τ ) softmax_{\tau }( x)_{i} =\frac{\exp( x_{i} /\tau )}{\sum _{j}\exp( x_{j} /\tau )} softmaxτ(x)i=∑jexp(xj/τ)exp(xi/τ)
那么他跟指派问题有什么关系呢？其实我们可以定义一个对行列不停地分别除以每一行的和，以及除以每一列的和的操作，这个操作称为Sinkhorn operator S(X) 可以定义如下：
S 0 ( X ) = exp ( X ) S l ( X ) = T c ( T r ( S l − 1 ( X ) ) ) S ( X ) = lim l → ∞ S l ( 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} S0(X)Sl(X)S(X)=exp(X)=Tc(Tr(Sl−1(X)))=l→∞limSl(X)
其中 T r ( X ) = X ⊘ ( X 1 N 1 N ⊤ ) , T c ( X ) = X ⊘ ( 1 N 1 N ⊤ X ) \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) Tr(X)=X⊘(X1N1N⊤),Tc(X)=X⊘(1N1N⊤X)，可以证明S(X)一定会收敛到一个叫Birkhoff polytope的空间上，记为
B N = { P ∈ [ 0 , 1 ] ∈ R N , N P 1 N = 1 N , P ⊤ 1 N = 1 N } \mathcal{B}_{N} =\left\{P\in [0,1]\in \mathbb{R}^{N,N} P1_{N} =1_{N} ,P^{\top } 1_{N} =1_{N}\right\} BN={ P∈[0,1]∈RN,NP1N=1N,P⊤1N=1N}
可以证明 B N \displaystyle \mathcal{B}_{N} BN是包含所有置换矩阵的。于是，指派问题与softmax的联系可以用下面这个定理联系起来：
定理1： 定义随机矩阵P的熵为 h ( P ) = − ∑ i . j p i j log p i j \displaystyle h(P)=-\sum _{i.j} p_{ij}\log p_{ij} h(P)=−i.j∑pijlogpij，则
S ( X / τ ) = arg max P ∈ B N ⟨ P , X ⟩ F + τ h ( P ) S(X/\tau )=\underset{P\in \mathcal{B}_{N}}{\arg\max} \langle P,X\rangle _{F} +\tau h(P) S(X/τ)=P∈BNargmax⟨P,X⟩F+τh(P)
如果X的产生是独立的，那么一定有
M ( X ) = lim τ → 0 + S ( X / τ ) M(X)=\lim _{\tau \rightarrow 0^{+}} S(X/\tau ) M(X)=τ→0+limS(X/τ)
现在证明第一个，考虑如下拉格朗日有约束的优化：
L ( α , β , P ) = ⟨ P , X ⟩ F + τ h ( P ) + α ⊤ ( P 1 N − 1 N ) + β ⊤ ( P ⊤ 1 N − 1 N ) \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) L(α,β,P)=⟨P,X⟩F+τh(P)+α⊤(P1N−1N)+β⊤(P⊤1N−1N)
令 ∂ L / ∂ P = 0 \partial \mathcal{L} /\partial P=0 ∂L/∂P=0 一定有对于每个 i , j i,j i,j
p τ i , j = exp ( α i / τ − 1 / 2 ) exp ( X i , j / τ ) exp ( β j / τ − 1 / 2 ) p^{i,j}_{\tau } =\exp( \alpha _{i} /\tau -1/2)\exp( X_{i,j} /\tau )\exp( \beta _{j} /\tau -1/2) pτi,j=exp(αi/τ−1/2)exp(Xi,j/τ)exp(βj/τ−1/2)
也就是说对于任意正数的对角矩阵 D 1 , D 2 , D_{1} ,D_{2} , D1,D2,都有 P τ = D 1 exp ( X i , j / τ ) D 2 P_{\tau } =D_{1}\exp( X_{i,j} /\tau ) D_{2} Pτ=D1exp(Xi,j/τ)D2 根据Sinkhorn’s theorem, 一定有 S ( X / τ ) = P τ S(X/\tau )=P_{\tau } S(X/τ)=Pτ
与softmax的联系
实际上，以上的推导只是基于一个softmax的原理简单的推广:
arg max i x i = lim τ → 0 + softmax ( x / τ ) \arg\max_{i} x_{i} =\lim _{\tau \rightarrow 0^{+}}\operatorname{softmax} (x/\tau ) argimaxxi=τ→0+limsoftmax(x/τ)
即是softmax所近似的argmax值，实际上，这个argmax可以表示成一个关于one-hot向量的优化问题，即寻找一个one-hot向量，使得他跟x的内积最大：
arg max i x i = arg max e ⟨ e , x ⟩ \arg\max_{i} x_{i} =\arg\max_{e} \langle e,x\rangle argimaxxi=argemax⟨e,x⟩
跟上面的证明差不多，我们可以写出e的带约束的拉格朗日公式，并对e求导，令它等于0，可以得到e的形式就是一个softmax公式，于是
softmax ( x / τ ) ≡ exp ( x / τ ) ∑ i = 1 exp ( x i / τ ) = arg max e ∈ S n ⟨ e , x ⟩ + τ h ( e ) \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(x/τ)≡∑i=1exp(xi/τ)exp(x/τ)=e∈Snargmax⟨e,x⟩+τh(e)
可以简单证明一下：
∂ ∂ e ( ⟨ e , x ⟩ + τ h ( e ) + α ⊤ ( e 1 N − 1 N ) ) = x + τ log ( e ) + τ + α ⊤ = 0 ⟹ e = exp ( − x / τ ) exp ( − 1 − α ⊤ / τ ) \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} ∂e∂(⟨e,x⟩+τh(e)+α⊤(e1N−1N))⟹=x+τlog(e)+τ+α⊤=0e=exp(−x/τ)exp(−1−α⊤/τ)
又因为 ∑ i e i = 1 \displaystyle \sum _{i} e_{i} =1 i∑ei=1,所以一定有 exp ( − 1 − α ⊤ / τ ) = 1 / ∑ i = 1 exp ( x i / τ ) \displaystyle \exp\left( -1-\alpha ^{\top } /\tau \right) =1/\sum _{i=1}\exp( x_{i} /\tau ) exp(−1−α⊤/τ)=1/i=1∑exp(xi/τ).
于是我们理解permutation matrix本质上可以弱化为softmax组成的矩阵，因此我们完全可以利用gumbel分布对softmax的重参数化能力来重参数化这个置换矩阵，从而采用VAE来求解一个隐变量时置换矩阵的生成模型，而这就是这篇论文所做的贡献了。
关于gumbel的重参数可以参考我前面的一篇文章：
带你认识神奇的Gumbel trick
https://blog.csdn.net/a358463121/article/details/80820878
对softmax的进一步思考
当我们把softmax写成下面这一条公式的时候，事情就变得越发有趣：
softmax ( x / τ ) ≡ exp ( x / τ ) ∑ i = 1 exp ( x i / τ ) = arg max e ∈ S n ⟨ e , x ⟩ + τ h ( e ) \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(x/τ)≡∑i=1exp(xi/τ)exp(x/τ)=e∈Snargmax⟨e,x⟩+τh(e)
我们发现，softmax的指数项的来源，其实是那个 h ( e ) h(e) h(e)产生，因为里面有个log，为了得到偏导数等于0的e，这个log要取个指数项消掉。然而对于熵 h ( x ) h(x) h(x)，其实除了香农熵的经典定义外，还有很多其他的定义的，这是否意味着，softmax还拥有其他的变种？
答案是肯定的，在论文[3]中说明了，只要将香农熵换成，Gini entropy： H G ( p ) : = 1 2 ∑ j p j ( 1 − p j ) \mathrm{H}^{\mathrm{G}}(\boldsymbol{p}):=\frac{1}{2} \sum_{j} p_{j}\left(1-p_{j}\right) HG(p):=21∑jpj(1−pj)，就能推导出全新的sparse softmax，这个softmax具有sparse的性质：
sparsemax ( x ) : = argmin p ∈ Δ d ∥ p − x ∥ 2 \text { sparsemax }(\boldsymbol{x}):=\underset{\boldsymbol{p} \in \Delta^{d}}{\operatorname{argmin}}\|\boldsymbol{p}-\boldsymbol{x}\|^{2} sparsemax (x):=p∈Δdargmin∥p−x∥2
实际上，这个sparse softmax是在论文 [2] 中提出的，有了这个思路，其实我们能够推导出各种各样的变种，论文就提出了用Tsallis α-entropies:
H α T ( p ) : = { 1 α ( α − 1 ) ∑ j ( p j − p j α ) , α ≠ 1 H S ( p ) , α = 1 \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. HαT(p):={ α(α−1)1∑j(pj−pjα),HS(p),α=1α=1
于是就推导出了entmax
α -entmax ( z ) : = argmax p ∈ Δ d p ⊤ z + H α T ( p ) \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}) α -entmax (z):=p∈Δdargmaxp⊤z+HαT(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.