置换矩阵也能求导优化
本文是对论文 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.