1. 为什么要引入Softmax？

在进行 $n$ 分类任务时，神经网络的最后一层输出（logits）是一个 $n$ 维向量 $\mathbf{z}\in {\mathbb{R}}^n$。我们希望将其转化为概率分布，以便选取概率最大的类别作为神经网络的最终分类结果。

直观上来讲，$\mathbf{z}$ 中哪个分量越大，相应的概率就应该越高，我们可以很容易想到使用下面的公式进行概率建模：

$$P(z_i)=\frac{z_i}{z_1+z_2+\cdots +z_n}\text{\hspace{1em}(1)}$$

显然有 ${\sum }_{i}P(z_i)=1$，但由于 $z_i\in \mathbb{R}$，$P(z_i)$ 可能为负数，这不符合概率的性质。

注意到指数函数 $e^x$ 满足单调性，非负性，导数不变性，并且可以放大输入值之间的差异，因此我们可以先让 $\mathbf{z}$ 过一遍指数函数，然后再使用公式 $(1)$ 进行概率建模，此时便得到了Softmax公式：

$${\sigma }_i=\text{Softmax}(z_i)=\frac{e^{z_i}}{{\sum }_{j=1}^n{e}^{z_j}}\text{\hspace{1em}(2)}$$

可以看出，$\text{Softmax}(\cdot )$ 函数是一个 ${\mathbb{R}}^n\to {\mathbb{R}}^n$ 的映射。

📝 为什么Softmax函数中要选取 $e^x$，而不是 $2^x,3^x$ 等函数？注意到 $(a^x{)}^{\prime }=a^x\ln a$，如果不选取 $e^x$，那么神经网络在反向传播计算梯度时会产生大量的 $\ln a$，这对于计算并不友好，而选择 $e^x$ 就可以简化掉这些繁琐的对数项。

2. Softmax的导数计算

先前我们提到过，Softmax函数接受 $n$ 个输入，并产生相应的 $n$ 个输出。对于每一个输出，我们都可以对任意一个输入进行求导，所以Softmax函数的导数实际上是一个 $n\times n$ 的雅克比矩阵：

$$\mathbf{J}={\left[\begin{array}{ccc}\frac{\partial {\sigma }_1}{\partial z_1}& \cdots & \frac{\partial {\sigma }_1}{\partial z_n}\\ ⋮& \ddots & ⋮\\ \frac{\partial {\sigma }_n}{\partial z_1}& \cdots & \frac{\partial {\sigma }_n}{\partial z_n}\end{array}\right]}_{n\times n}$$

接下来我们关注如何计算 $\frac{\partial {\sigma }_i}{\partial z_j}$，此时要分两种情况讨论。

当 $i=j$ 时，此时计算的是对角线上的元素：

$$\begin{array}{rl}\frac{\partial {\sigma }_i}{\partial z_i}& =\frac{\partial }{\partial z_i}\left(\frac{e^{z_i}}{{\sum }_{j=1}^n{e}^{z_j}}\right)=\frac{e^{z_i}\cdot {\sum }_{j=1}^n{e}^{z_j}-e^{z_i}\cdot e^{z_i}}{({\sum }_{j=1}^n{e}^{z_j}{)}^2}\\ & =\frac{e^{z_i}\cdot ({\sum }_{j=1}^n{e}^{z_j}-e^{z_i})}{({\sum }_{j=1}^n{e}^{z_j})\cdot ({\sum }_{j=1}^n{e}^{z_j})}=\frac{e^{z_i}}{{\sum }_{j=1}^n{e}^{z_j}}\cdot \left(1-\frac{e^{z_i}}{{\sum }_{j=1}^n{e}^{z_j}}\right)\\ & ={\sigma }_i(1-{\sigma }_i)\end{array}$$

当 $i\ne j$ 时，此时计算的是非对角线上的元素：

$$\begin{array}{rl}\frac{\partial {\sigma }_i}{\partial z_j}& =\frac{\partial }{\partial z_j}\left(\frac{e^{z_i}}{{\sum }_{j=1}^n{e}^{z_j}}\right)=\frac{0-e^{z_i}\cdot e^{z_j}}{({\sum }_{j=1}^n{e}^{z_j}{)}^2}\\ & =-\frac{e^{z_i}}{{\sum }_{j=1}^n{e}^{z_j}}\cdot \frac{e^{z_j}}{{\sum }_{j=1}^n{e}^{z_j}}\\ & =-{\sigma }_i{\sigma }_j\end{array}$$

有了这些结果，我们可以尝试对雅可比矩阵进行一些变换：

$$\begin{array}{rl}\mathbf{J}& =\left[\begin{array}{cccc}{\sigma }_1(1-{\sigma }_1)& -{\sigma }_1{\sigma }_2& \cdots & -{\sigma }_1{\sigma }_2\\ -{\sigma }_2{\sigma }_1& {\sigma }_2(1-{\sigma }_2)& \cdots & -{\sigma }_2{\sigma }_n\\ ⋮& ⋮& \ddots & ⋮\\ -{\sigma }_n{\sigma }_1& -{\sigma }_n{\sigma }_2& \cdots & {\sigma }_n(1-{\sigma }_n)\end{array}\right]\\ & =\left[\begin{array}{cccc}{\sigma }_1& 0& \cdots & 0\\ 0& {\sigma }_2& \cdots & 0\\ ⋮& ⋮& \ddots & ⋮\\ 0& 0& \cdots & {\sigma }_n\end{array}\right]-\left[\begin{array}{cccc}{\sigma }_1{\sigma }_1& {\sigma }_1{\sigma }_2& \cdots & {\sigma }_1{\sigma }_n\\ {\sigma }_2{\sigma }_1& {\sigma }_2{\sigma }_2& \cdots & {\sigma }_2{\sigma }_n\\ ⋮& ⋮& \ddots & ⋮\\ {\sigma }_n{\sigma }_1& {\sigma }_n{\sigma }_2& \cdots & {\sigma }_n{\sigma }_n\end{array}\right]\\ & =\boxed{\text{diag}(\sigma )-\sigma {\sigma }^{\text{T}}}\end{array}$$

其中 $\sigma =({\sigma }_1,{\sigma }_2,\cdots ,{\sigma }_n)$ 是列向量，$\text{diag}(\sigma )$ 表示使用向量 $\sigma$ 构建对角矩阵。

我们可以利用这一结果，编写相应的代码实现Softmax及其导数计算：

import numpy as npdef softmax(z):e_z = np.exp(z)return e_z / e_z.sum()def softmax_jacobian(z):s = softmax(z).reshape(-1, 1)diag_s = np.diagflat(s)outer_s = np.dot(s, s.T)return diag_s - outer_sif __name__ == "__main__":z = np.array([1.0, 2.0, 3.0])print("Input vector z:")print(z)print("\nSoftmax of z:")s = softmax(z)print(s)print("\nJacobian matrix of the softmax function at z:")J = softmax_jacobian(z)print(J)

输出：

Input vector z:
[1. 2. 3.]Softmax of z:
[0.09003057 0.24472847 0.66524096]Jacobian matrix of the softmax function at z:
[[ 0.08192507 -0.02203304 -0.05989202][-0.02203304  0.18483645 -0.1628034 ][-0.05989202 -0.1628034   0.22269543]]

3. Softmax及其导数的一些性质

性质1：Softmax具有平移不变性，即

$$\text{Softmax}([z_1+a,z_2+a,\cdots ,z_n+a])=\text{Softmax}([z_1,z_2,\cdots ,z_n])$$

证明：这是一个十分显然的事情

$$\text{LHS}=\frac{e^{z_i+a}}{{\sum }_{j=1}^n{e}^{z_j+a}}=\frac{e^{z_i}}{{\sum }_{j=1}^n{e}^{z_j}}=\text{RHS}$$

性质2：Softmax的导数 $\mathbf{J}$ 中，任意一行或任意一列的和均为 $0$.

证明：因为 $\mathbf{J}$ 是一个对称矩阵，即 $\mathbf{J}={\mathbf{J}}^{\text{T}}$，我们只需证明 $\mathbf{J}$ 的任意一行的和为 $0$ 即可。

$$\sum _{j=1}^n{\mathbf{J}}_{ij}=\sum _{j\ne i}(-{\sigma }_i{\sigma }_j)+{\sigma }_i(1-{\sigma }_i)=-\sum _{j\ne i}{\sigma }_i{\sigma }_j-{\sigma }_i{\sigma }_i+{\sigma }_i=-{\sigma }_i+{\sigma }_i=0$$

由以上结果可知 $J1=0$，这说明 $1$ 是 $\mathbf{J}$ 的一个特征向量，对应的特征值为 $0$。

性质3：$\mathbf{J}$ 为半正定矩阵且 $\text{rank}(\mathbf{J})=n-1$。

证明：对于任意向量 $\mathbf{x}\in {\mathbb{R}}^n$，考虑二次型：

$${\mathbf{x}}^{\mathrm{T}}\mathbf{Jx}={\mathbf{x}}^{\mathrm{T}}\left(\text{diag}(\sigma )-\sigma {\sigma }^{\mathrm{T}}\right)\mathbf{x}=\sum _{i=1}^n{\sigma }_i{x}_i^2-{\left(\sum _{i=1}^n{\sigma }_i{x}_i\right)}^2$$

将 ${\sigma }_i$ 视为概率分布，定义随机变量 $X$，其取值为 $x_i$，概率为 ${\sigma }_i$。则：

$${\mathbf{x}}^{\mathrm{T}}\mathbf{Jx}=\mathbb{E}[X^2]-(\mathbb{E}[X]{)}^2=\mathrm{Var}(X)\ge 0$$

因此 $\mathbf{J}$ 是半正定矩阵。

当 $\mathrm{Var}(X)>0$ 时，意味着随机变量 $X$ 至少有两个不同的取值，也就是说 $\mathbf{x}$ 不是一个常数向量，且 $\Vert \mathbf{x}{\Vert }^2>0$，从而

$$0<\mathrm{Var}(X)={\mathbf{x}}^{\mathrm{T}}\mathbf{Jx}=\lambda \Vert \mathbf{x}{\Vert }^2\Rightarrow \lambda >0$$

结合性质2可知 $\mathbf{J}$ 只有一个零特征值，剩余的 $n-1$ 个特征值均为正数，从而 $\text{rank}(\mathbf{J})=n-1$。

4. 交叉熵损失的梯度计算

设Softmax的输入为 $\mathbf{z}$，输出为 $\sigma$，标签为 $\mathbf{y}$（one-hot向量），那么交叉熵损失定义为：

$$\mathcal{L}=-\sum _{i=1}^n{y}_i\cdot \log {\sigma }_i=-{\mathbf{y}}^{\text{T}}\log \sigma$$

于是

$$\begin{array}{r}\frac{\partial \mathcal{L}}{\partial \mathbf{z}}={\left(\frac{\partial \sigma }{\partial \mathbf{z}}\right)}^{\text{T}}\cdot \frac{\partial \mathcal{L}}{\partial \sigma }=(\text{diag}(\sigma )-\sigma {\sigma }^{\mathrm{T}})\cdot \left(-\frac{\mathbf{y}}{\sigma }\right)=\sigma -\mathbf{y}\end{array}$$

这说明损失函数对输入 $\mathbf{z}$ 的梯度可以简洁地表示为Softmax的输出与真实标签之间的差值。

5. Softmax的各种变体

在这一节中，我们会使用C++来实现Softmax的各种变体。

5.1 Naive Softmax

即最原始的softmax。总共需要两次遍历，第一次遍历是用来计算分母，第二次遍历是用来计算每一个输出。

std::vector<double> softmax(const std::vector<double>& input) {std::vector<double> output(input.size());double sum = 0.0;for (size_t i = 0; i < input.size(); i++) {output[i] = std::exp(input[i]);sum += output[i];}for (size_t i = 0; i < output.size(); i++) {output[i] /= sum;}return output;
}

5.2 Safe Softmax

在实际计算中，$e^x$ 很容易发生上溢问题。例如，在 Python 中，当 $x\ge 710$ 时，计算 $e^x$ 就会导致数值溢出。此外，在深度学习模型中，fp16 精度广泛用于加速训练，但它的最大表示范围仅为 $65536$。这意味着，当 $x\ge 11$ 时，$e^x$ 的值将超过 fp16 能够表示的数值范围，进而导致上溢错误。

鉴于Softmax的平移不变性以及 $e^x$ 对于负输入的计算会更加精确，我们可以通过将所有输入减去其最大值来缓解上溢问题，即计算 $\text{Softmax}(\mathbf{z}-\max (\mathbf{z}))$。这一操作不仅可以避免数值溢出，还能提升计算的稳定性和精度，并且不影响Softmax的输出结果。

大部分的深度学习框架都采用了Safe Softmax，但它的缺点是需要进行三次遍历，第一次用来统计最大值，第二次用来计算分母，第三次用来计算每个输出。

std::vector<double> safe_softmax(const std::vector<double>& input) std::vector<double> safe_softmax(const std::vector<double>& input) {std::vector<double> output(input.size());double max_val = input[0];for (size_t i = 1; i < input.size(); i++) {max_val = std::max(max_val, input[i]);}double sum = 0.0;for (size_t i = 0; i < input.size(); i++) {output[i] = std::exp(input[i] - max_val);sum += output[i];}for (size_t i = 0; i < output.size(); i++) {output[i] /= sum;}return output;
}

5.3 Online Softmax

Safe Softmax需要三次遍历，那有没有可能将其缩减至两次遍历呢？

定义 $m_k={\max }_{i=1}^k{z}_i$，那么有 $m_n=\max (\mathbf{z})$。定义 $d_k={\sum }_{i=1}^k{e}^{z_i-m_n}$，那么 $d_n$ 就是Safe Softmax中的分母。

回顾Safe Softmax的计算过程，为了得到 $d_n$，我们需要从 $d_1$ 计算到 $d_n$，但无论是哪一个 $d_k$，都会涉及到 $m_n$，这意味着在此之前必须先通过一轮循环得到 $m_n$。

能否把计算 $m_n$ 的这一轮循环砍掉呢？一个直观的想法就是在计算 $d_1$ 到 $d_n$ 的过程中同时去计算 $m_1$ 到 $m_n$。但此时计算 $d_k$ 时，我们只有 $m_k$，并没有 $m_n$，该怎么办呢？

先将计就计试一下，令 $d_k^{\prime }={\sum }_{i=1}^k{e}^{z_i-m_k}$，容易发现

$$d_n^{\prime }=\sum _{i=1}^n{e}^{z_i-m_n}=d_n$$

这说明我们可以用 $d_n^{\prime }$ 作为Safe Softmax的分母。接下来只需要关注如何去计算 $d_k^{\prime }$。注意到

$$\begin{array}{rl}{d}_k^{\prime }& =\sum _{i=1}^k{e}^{z_i-m_k}\\ & =\sum _{i=1}^{k-1}{e}^{z_i-m_k}+e^{z_k-m_k}\\ & =\left(\sum _{i=1}^{k-1}{e}^{z_i-m_{k-1}}\right)\cdot e^{m_{k-1}-m_k}+e^{z_k-m_k}\\ & =d_{k-1}^{\prime }\cdot e^{m_{k-1}-m_k}+e^{z_k-m_k}\end{array}$$

且 $d_1^{\prime }=e^{z_1-m_1}$，于是我们可以用这个递推式来计算得到 $d_n^{\prime }$。

std::vector<double> online_softmax(const std::vector<double>& input) {std::vector<double> output(input.size());double pre_max_val = input[0];double sum = std::exp(input[0] - pre_max_val);for (size_t i = 1; i < input.size(); i++) {double cur_max_val = std::max(pre_max_val, input[i]);sum = sum * std::exp(pre_max_val - cur_max_val) + std::exp(input[i] - cur_max_val);pre_max_val = cur_max_val;}for (size_t i = 0; i < input.size(); i++) {output[i] = std::exp(input[i] - pre_max_val) / sum;}return output;
}

---

Ref

[1] https://zhuanlan.zhihu.com/p/638788074
[2] https://courses.cs.washington.edu/courses/cse599m/23sp/notes/flashattn.pdf
[3] https://arxiv.org/pdf/1805.02867v2