函数光滑近似（2）：softmax与argmax

本系列的第一篇函数光滑近似（1）：max 函数，我们讨论了 max 函数的光滑近似，softmax 在字面上容易误解为其光滑近似，其实后者是 argmax 的光滑近似。本篇文章讨论这个问题。

在 Attention 中，常常使用 softmax 来归一化评分函数，具体可见漫谈注意力机制：Attention的全面理解，softmax 的实质可以让我们更深入理解为什么如此操作。

softmax的数学性质

softmax函数表示如下，

$$\operatorname{softmax}(\boldsymbol{x}) = \frac{e^{ x_{i}}}{ \displaystyle \sum_{i=1}^{n}e^{ x_{i}}}$$

当然我们也可以引入一个缩放系数$\alpha$，有，

$$\operatorname{softmax}(\boldsymbol{x}) = \frac{e^{ \alpha x_{i}}}{ \displaystyle \sum_{i=1}^{n}e^{\alpha x_{i}}}$$

首先我们来谈谈softmax函数的数学性质。易知softmax具有归一性，即性质一，

$$\sum_i \frac{e^{ x_{i}}}{ \displaystyle \sum_{i=1}^{n}e^{ x_{i}}} = 1$$

可以让输入的$x_1, \dots ,x_n$具有概率意义。

我们都知道，max的近似logsumexp的梯度就是 softmax 函数，性质二，

$$\frac{\partial}{\partial x_{i}} (\frac{1}{\alpha} \log(\sum_{i=1}^{n}e^{\alpha x_{i}})) = \frac{e^{\alpha x_{i}}}{ \displaystyle \sum_{i=1}^{n}e^{\alpha x_{i}}}$$

其中$\alpha=1$​，这可以看做是softmax函数的来源之一。如果以$q(x) = \operatorname{softmax}(x)$作为模型输出，那么交叉熵损失为，

$$-\sum_i p(x_i) \log q(x_i) = \log \left(\sum_{i=1}^{n}e^{x_i} \right) - x_j$$

这里$p(x)$是one-hot形式的标签。可以看到，这个logsumexp有出现了。

性质三，加性不变性，

$$\begin{align} \operatorname{softmax}(\boldsymbol{x} + c) &= \frac{e^{ x_{i}+c}}{ \displaystyle \sum_{i=1}^{n}e^{ x_{i}+c}} \newline &= \frac{e^c e^{ x_{i}}}{ \displaystyle e^c\sum_{i=1}^{n}e^{ x_{i}}} \newline &= \operatorname{softmax}(\boldsymbol{x}) \end{align}$$

把它展开了看，

$$\begin{align} \operatorname{softmax}(c+\boldsymbol{x}) &= [\frac{e^{c+x_{1}}}{\displaystyle \sum_{k=1}^{n}e^{c+x_{k}}},\cdots,\cdots,\frac{e^{c+x_{n}}}{\displaystyle \sum_{k=1}^{n}e^{c+x_{k}}}] \newline &= [\frac{e^{c}e^{x_{1}}}{e^{c} \displaystyle \sum_{k=1}^{n}e^{x_{k}}},\cdots,\cdots,\frac{e^{c} e^{x_{n}}}{\displaystyle e^{c} \sum_{k=1}^{n}e^{x_{k}}}] \newline &= \operatorname{softmax}(\boldsymbol{x}) \end{align}$$

利用这个性质，在编程实现时一般另$c = \max(x_1, \dots, x_n)$，然后计算$\operatorname{softmax}(\boldsymbol{x} - c) $​​，这样可以避免数值溢出。

性质四，softmax的梯度，

$$\begin{align} \frac{\partial}{\partial x_i} \operatorname{softmax}(\boldsymbol{x})_{[j]} &= \frac{\partial}{\partial x_i} \frac{e^{x_j}}{\displaystyle \sum_{k=1}^{n}e^{x_k}} \newline &= \begin{cases} \frac{e^{x_j}}{\displaystyle \sum_{k=1}^{n}e^{x_k}} (1 - \frac{e^{x_j}}{\displaystyle \sum_{k=1}^{n}e^{x_k}}) , \text{if} \;i = j \\ -\frac{e^{x_j}}{\displaystyle \sum_{k=1}^{n}e^{x_k}} \frac{e^{x_i}}{\displaystyle \sum_{k=1}^{n}e^{x_k}}, \text{if} \;i \ne j \end{cases} \end{align}$$

注意到取$\log$的softmax与logsumexp的关系，

$$\log \operatorname{softmax}(\boldsymbol{x}) = x_j - \operatorname{logsumexp}(\boldsymbol{x})$$

因此，从数学性质上看，$f(x) = \exp(x)$​是将任意实数$x$​映射到非负实数的单调的、光滑的、最简单的初等函数。

softmax的导出

设有向量

$$\boldsymbol{x} = [x_{1}, \cdots,x_{N}]$$

argmax的定义如下，

$${\underset {x\in S}{\operatorname {arg\,max} }}\,f(x)=\{x\in S~:~f(s)\leq f(x){\text{ for all }}s\in S\}$$

例如，最大值位置表示如下，

$$i = \underset{i=1,\cdots,N}{\text{argmax}} \; \boldsymbol{x}$$

比如圆周率的头几个数，

$$\boldsymbol{x} = [3,1,4,1,5,9,2]$$

结果为

$$6 = \underset{i=1,\cdots,N}{\text{argmax}} \; [3,1,4,1,5,{\color{red} {9}},2]$$

表达成 one-hot 形式，即第$i$个位置为1，其他为0，

$$[0,\cdots,1,\cdots,0] = \text{one-hot}(\underset{i=1,\cdots,N}{\text{argmax}} \; \boldsymbol{x})$$

续上述例子

$$[0,0,0,0,0,1,0] = \text{one-hot}(\underset{i=1,\cdots,N}{\text{argmax}} \; \boldsymbol{x})$$

假设 $\boldsymbol{x} = [x_{1}, \cdots,x_{N}]$​ 中每个元素互不相等，这里用到上一篇推导的max的光滑近似，

$$\begin{align} [0,\cdots,1,\cdots,0] &= \text{one-hot}(\underset{i=1,\cdots,N}{\text{argmax}} \; \boldsymbol{x}) \newline &= \text{one-hot}(\underset{i=1,\cdots,N}{\text{argmax}} \; \boldsymbol{x}-\max(\boldsymbol{x})) \newline &= \text{one-hot}(\underset{i=1,\cdots,N}{\text{argmax}} \; e^{\boldsymbol{x}-\max(\boldsymbol{x})}) \newline &\approx \text{one-hot}(\underset{i=1,\cdots,N}{\text{argmax}} \; e^{\boldsymbol{x}-\log{ \displaystyle \sum_{k=1}^{N}}e^{x_{k}}}) \;\; \color{red}{\text{max的近似}}\newline &\approx \Big[\frac{e^{x_{1}}}{\displaystyle \sum_{k=1}^{N}e^{x_{k}}},\cdots,\cdots,\frac{e^{x_{N}}}{\displaystyle \sum_{k=1}^{N}e^{x_{k}}} \Big] \newline &= \text{softmax}(\boldsymbol{x}) \end{align}$$

需要说明几点：

• 引入$\big[x_i - \max(\boldsymbol{x}) \big]$使得最大值为0，使得$e^0 = 1$，对应one-hot中的1
• 引入$e^x$是考虑到$e^0=1, 0 \lt e^{x|_{x \lt 0}} \lt 1$，并拉大$[x_1, x_2, \dots, x_n]$间的距离，更好适配one-hot特点
• max不具有光滑性，被替换为其光滑近似logsumexp

理解好这三点就明白上述推导过程。

可以看到，softmax 容易让人误解成 soft 的 max，其实是 soft 的 one-hot(argmax)，它们本是同根生。在这个基础上，使用递归的思想，不难解决 softmax 问题的 topk 问题，即先把最大的计算出来，然后把其剔除，接着计算剩下最大的，以此类推…

到此，你可能意识到，为了平滑$\text{one-hot}(\underset{i=1,\cdots,N}{\text{argmax}} \; \boldsymbol{x})$，弄出了$\text{softmax}(\boldsymbol{x})$有点吃力不讨好，为什么？因为$\text{softmax}(\boldsymbol{x})$需要求和，当向量维度很大时，还真不是简单的事情。比如，层次化 Softmax，重要性采样和噪声对比估计来解决。但是当维度不太大时，softmax的影响较少。后续会谈到。

softmax在深度学习场景中的解释

对于一个$n$分类问题，模型输出 softmax 分布对应的是类别的概率，

$$(p_{1}, \dots, p_{n}) = \frac{(e^{x_{1}},\dots,e^{x_{n}})}{\displaystyle \sum_{i=1}^{n}e^{x_{i}}}$$

$p_i$表示类别为$i$的概率。根据以上推导可知，模型本来是可以直接输出one-hot(argmax)，直接指定样本的类别$i$​，但是one-hot(argmax)太hard了，不具有光滑性，不利于梯度优化，为此把输出改为其光滑版本，即softmax。这可以看做是多分类问题的输出使用softmax的基本原因，其他的解释都是使用softmax后的side effect。

在注意力机制中，使用 softmax 函数对评分函数输出进行归一化，假设有一向量序列 $\boldsymbol{X} = [\boldsymbol{x}_{1}, \dots, \boldsymbol{x}_{n}] \in \mathbb{R}^{n \times d}$​，获得注意力分布，

$$\begin{align} p(z=i|\boldsymbol{X},\boldsymbol{q}) &= \operatorname{softmax}(\alpha_{1},\dots,\alpha_{n}) \newline &= \frac{\exp(\alpha_{i})}{\displaystyle{\sum_{i=1}^{n}\exp(\alpha_{i})}} \newline &= \frac{\exp(s(\boldsymbol{q}, \boldsymbol{x}_{i}))}{\displaystyle{\sum_{i=1}^{n}\exp(s(\boldsymbol{q}, \boldsymbol{x}_{i}))}} \end{align}$$

$s(\boldsymbol{q}, \boldsymbol{x}_{i})$​​​为评分函数。然后加权平均，即注意力分布下的均值，获得注意力的输出，

$$\operatorname{Attention}(\boldsymbol{X},\boldsymbol{q}) = \sum_{i=1}^{n}p(z=i|\boldsymbol{X},\boldsymbol{q})\boldsymbol{x}_{i}$$

由于softmax是one-hot(argmax)的光滑化，可以这样解释注意力机制。本来模型可以在给定查询向量$\boldsymbol{q}$情况下，选择输出$s(\boldsymbol{q}, \boldsymbol{x}_{i})$最大值所对应的向量$\boldsymbol{x}_i$，即

$$p(z=i|\boldsymbol{X},\boldsymbol{q}) = \text{one-hot}\Big(\underset{i=1,\cdots,N}{\text{argmax}} \; s(\boldsymbol{q}, \boldsymbol{x}_{i}) \Big)$$

但是one-hot(argmax)太hard了，不具有光滑性，不利于梯度优化，为此把输出改为其光滑版本，即softmax。这可以看做是注意力机制中，对评分函数进行softmax归一化的基本原因，其他的解释都是使用softmax后的side effect。

softmax还可以近似计算数据集$R$的正确率，

$$\begin{align} \operatorname{acc}(R) &= \frac{1}{|R|} \sum_{C_k \in R} \operatorname{one-hot}(C_k) \cdot \operatorname{one-hot}(\arg\max(\boldsymbol{x})) \newline &\approx \frac{1}{|R|} \sum_{C_k \in R} \operatorname{one-hot}(C_k) \cdot \operatorname{softmax}(\boldsymbol{x}) \end{align}$$

argmax的导出

softmax 可以很自然地处理 argmax 的平滑问题，

$$\begin{align} \text{argmax}(\boldsymbol{x}) &= \sum_{i=1}^{n}i \times \text{one-hot}(\text{argmax}(\boldsymbol{x}))_{i} \newline &\approx \sum_{i=1}^{n} i \times \text{softmax}(\boldsymbol{x}) \newline &= \frac{\displaystyle \sum_{i=1}^{n}i \times e^{x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{x_{i}}} \end{align}$$

所以，到这里，我们已经区分了 argmax、softmax、soft 的 max 的区别了吧。使用递归的思想，不难解决 softkmax 问题。

argmax的光滑近似容易让人联想到基于softmax的加权平均。假设有序列$[x_1, \dots, x_n]$，那么$x_i$对应的权重由softmax获得，

$$w_{i} = \frac{e^{x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{x_{i}}}$$

对序列加权平均有，

$$\begin{align} \bar{x} &= \sum_{i=1}^{n} w_{i} x_{i} \\ &= \sum_{i=1}^{n} \frac{e^{x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{x_{i}}} x_{i} \\ &= \frac{\displaystyle \sum_{i=1}^{n}x_{i} e^{x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{x_{i}}} \end{align}$$

这个结果像什么？这个结论也是$\max(\boldsymbol{x})$的光滑近似的一种形式，

$$\begin{align} \max(\boldsymbol{x}) &= \boldsymbol{x}_{\arg\max(\boldsymbol{x})} \newline &= \boldsymbol{x} \otimes \operatorname{one-hot}[\arg \max(\boldsymbol{x})] \newline &= \sum_{i=1}^{n} x_i \times \operatorname{one-hot}[\arg \max(\boldsymbol{x})]_{[i]} \newline &\approx \sum_{i=1}^{n} \frac{e^{x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{x_{i}}} x_{i} \\ &= \frac{\displaystyle \sum_{i=1}^{n}x_{i} e^{x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{x_{i}}} \end{align}$$

如果考虑参数有，

$$\max(\boldsymbol{x};\alpha) \approx \frac{\displaystyle \sum_{i=1}^{n}x_{i} e^{\alpha x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{\alpha x_{i}}}$$

有三个关键点：

• 如果$\alpha = 0$，则上式右侧为均值
• 如果$\alpha = +\infty$，则上式右侧为$\max(\boldsymbol{x})$
• 如果$\alpha = -\infty$，则上式右侧为$\min(\boldsymbol{x})$

取$x_1 = 0, x_2 = x, \alpha=1$​有特例一，

$$\max(0, x) \approx \frac{x e^x}{1 + e^x} = x \sigma(x)$$

取$x_1 = -x, x_2 = x$，有特例二，

$$\begin{align} |x| &= \max\{-x, x \} \newline &\approx \frac{x(e^{\alpha x} - e^{-\alpha x})}{e^{\alpha x} + e^{-\alpha x}} \newline &= x \tanh(\alpha x) \end{align}$$

实现

numpy实现，

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    x = np.exp(x)
    x = x / x.sum(axis=axis, keepdims=True)
    return x

我们接着看看 tensorflow 中的实现，

def softmax(x, axis=-1):
    """Tensorflow实现的softmax"""
    x = x - tf.reduce_max(x, axis, keepdims=True)
    x = tf.exp(x)
    return x / tf.reduce_sum(x, axis, keepdims=True)

总结

通过以上的推导，我们知道，

$$\operatorname{softmax}(\boldsymbol{x}) \approx \operatorname{one-hot}(\underset{i=1,\cdots,N}{\operatorname{argmax}} \; \boldsymbol{x})$$

对于$\arg \max$则有光滑表示，

$$\begin{align} \text{argmax}(\boldsymbol{x}) &= \sum_{i=1}^{n}i \times \text{one-hot}(\text{argmax}(\boldsymbol{x}))_{i} \newline &\approx \sum_{i=1}^{n} i \times \text{softmax}(\boldsymbol{x}) \end{align}$$

注意力机制中，对评分函数进行softmax归一化的解释。对多分类问题中，使用softmax输出作为类别概率分布的解释。

参考

[1] https://en.wikipedia.org/wiki/Heaviside_step_function

[2] https://www.elen.ucl.ac.be/Proceedings/esann/esannpdf/es2014-153.pdf

转载请包括本文地址：https://allenwind.github.io/blog/9905
更多文章请参考：https://allenwind.github.io/blog/archives/