method

原来的softmax公式是通过下面的计算得到概率分布

$\text{softmax}_i(z) = \frac{e^{z_i}}{\sum_j e^{z_j}}$

现在，作者考虑直接将隐状态向一个 $K-1$ 维的单纯形投影，求得的 $p$ 即是softmax的结果

$\begin{aligned} \Delta^{K-1} &:=\{p\in \mathbb{R}^K | 1^Tp=1, p \geq 0\} \\ \text{sparsemax}(z) &:= \mathop{\arg\max}\limits_{p \in \Delta^{K-1}} ||p-z||^2 \end{aligned}$

为什么是K-1维

因为有 $1^Tp=1$ 的限制，导致自由度减1

和 $x+y+z=1$ 表示的是一个平面而不是三维空间是一个道理

如何求解

将原问题转化为一个优化问题：

$\begin{aligned} \text{min} \quad & ||p-z||^2 \\ s.t. \quad & 1^Tp=1 \\ & p \geq 0 \end{aligned}$

考虑求解其拉格朗日对偶问题：

$L(p, \lambda, \mu) = \frac{1}{2} ||p-z||^2 - \lambda^Tp + \mu(1^Tp-1)$

这里系数 $\frac{1}{2}$ 是为了方便求导后消去系数

考虑KKT条件：

$\left\{ \begin{aligned} \nabla_p L(p^*, \lambda^*, \mu^*) &= p^*-z-\lambda^*+\mu^*1 \qquad&(1)\\ 1^Tp^* &= 1 &(2)\\ p^* &\geq 0 &(3)\\ \lambda^* &\geq 0 &(4)\\ \lambda^*p^* &= 0 &(5) \end{aligned} \right.$

对于 $p_i^{\ast} > 0$，由于式(5)，此时必有 $\lambda_i^{\ast}=0$，所以由式(1)，得：

$p_i^* = z_i - \mu, \quad s.t.\quad p_i^* > 0$

由式(2)，得：

$\begin{aligned} \sum_{i \in K}p_i^* &= 0 + \sum_{i \in S(z)}p_i^* \\ 1 &= \sum_{i \in S(z)}(z_i - \mu) \\ \mu &= \frac{\sum_{i \in S(z)}z_i - 1}{|S(z)|} \end{aligned}$

其中，$S(z) = \{j \in K \, | \,\, p_j^*>0\}$

由此，我们知道了如何从 $z$ 得到 $p$，即：

$p = \text{sparsemax}(z) = [z - \mu]_+$

求出 $\mu$ 的关键，是求出多大 $|S(z)|$ 才能正好满足：