在深度学习框架中，Softmax 函数是将 logit 向量转换为概率分布的核心操作，广泛应用于分类任务的输出层。无论是卷积神经网络还是 Transformer 架构中的注意力机制，Softmax 都是不可或缺的组成部分。理解其 Jacobian 矩阵的数学推导，不仅是掌握反向传播原理的基础，更是实现高效、正确梯度计算的关键。
Softmax 函数定义与数学形式
Softmax 函数的数学定义为：对于输入向量 $x \in \mathbb {R}^K$，输出向量 $S (x) = [S_1, S_2, \ldots, S_K]$，其中每个分量表示为：
$$S_i = \frac{\exp(x_i)}{\sum_{j=1}^{K} \exp(x_j)}$$
这个看似简单的归一化指数函数，实际上创造了一个巧妙的非线性变换：它将任意实数向量映射为和为 1 的概率分布，同时保留了输入之间的相对大小关系。值得注意的是，Softmax 的输出具有明显的耦合特性 —— 每个输出分量都依赖于所有输入分量，这直接导致其 Jacobian 矩阵是一个稠密矩阵，而非对角矩阵。
Jacobian 矩阵的严格推导
要推导 Softmax 的 Jacobian 矩阵，我们需要计算输出对输入的偏导数 $\frac {\partial S_i}{\partial x_j}$。设分母 $Z = \sum_{k=1}^{K} \exp (x_k)$，则 $S_i = \exp (x_i) / Z$。根据商求导法则，分两种情况讨论：
当 $i = j$ 时，即对角线元素，使用乘积求导：
$$\frac{\partial S_i}{\partial x_i} = \frac{\exp(x_i) \cdot Z - \exp(x_i) \cdot \exp(x_i)}{Z^2} = \frac{\exp(x_i)}{Z}\left(1 - \frac{\exp(x_i)}{Z}\right) = S_i(1 - S_i)$$
当 $i \neq j$ 时，即非对角线元素：
$$\frac{\partial S_i}{\partial x_j} = \frac{0 - \exp(x_i) \cdot \exp(x_j)}{Z^2} = -\frac{\exp(x_i)}{Z} \cdot \frac{\exp(x_j)}{Z} = -S_i S_j$$
综合两种情况，Softmax 的 Jacobian 矩阵 $J$ 可以表示为：
$$J_{ij} = \frac{\partial S_i}{\partial x_j} = S_i(\delta_{ij} - S_j)$$
其中 $\delta_{ij}$ 是 Kronecker 符号，当 $i = j$ 时为 1，否则为 0。这个优雅的统一表达式揭示了 Softmax 的本质特性：对角线元素为 $S_i (1 - S_i)$，非对角线元素为 $-S_i S_j$。所有行和为零是另一个重要性质，这反映了概率分布的归一化约束。
反向传播中的梯度计算
在深度学习的反向传播算法中，我们需要计算损失函数 $L$ 对 Softmax 输入 $x$ 的梯度。假设上游梯度 $g \in \mathbb {R}^K$ 是损失对 Softmax 输出的梯度，即 $g_i = \frac {\partial L}{\partial S_i}$，那么对输入的梯度通过 Jacobian 矩阵的转置传播：
$$\frac{\partial L}{\partial x_j} = \sum_i g_i \frac{\partial S_i}{\partial x_j} = \sum_i g_i S_i(\delta_{ij} - S_j)$$
展开整理后得到闭式表达式：
$$\frac{\partial L}{\partial x_j} = S_j \left( g_j - \sum_i g_i S_i \right)$$
用向量形式表示为：
$$\nabla_x L = S \circ \left( g - (g^\top S) \mathbf{1} \right)$$
其中 $\circ$ 表示逐元素乘法，$\mathbf {1}$ 是全 1 向量。这个公式给出了无需显式构造完整 Jacobian 矩阵的梯度计算方法，是工程实现的核心依据。
与 Cross-Entropy Loss 结合的简化
当 Softmax 与交叉熵损失函数配合使用时，梯度形式会显著简化。设目标标签为 one-hot 向量 $y$，交叉熵损失为 $L = -\sum_i y_i \log S_i$。对 Softmax 输出的上游梯度为：
$$g_i = \frac{\partial L}{\partial S_i} = -\frac{y_i}{S_i}$$
对于目标类别 $c$（即 $y_c = 1$），有 $g_c = -1/S_c$；对于其他类别 $i \neq c$，有 $g_i = 0$。将这个特殊的 $g$ 代入一般梯度公式，惊人的简化发生了：
$$\frac{\partial L}{\partial x_j} = S_j \left( -\frac{\delta_{jc}}{S_c} - \sum_i g_i S_i \right) = S_j - y_j$$
这就是深度学习中著名的「Softmax 梯度 = 预测概率减去真实标签」结论。工程实践中，当使用 Cross-Entropy 损失时，直接计算 $S - y$ 而非通过完整的 Jacobian 传播，既高效又数值稳定。
工程实现的关键参数
在实际框架实现中，需要关注以下工程参数。首先是数值稳定性处理：由于指数函数在输入值较大时容易溢出，实践中采用 Log-Sum-Exp 技巧：$S_i = \exp (x_i - \max_k x_k) / \sum_j \exp (x_j - \max_k x_k)$。这个操作不影响最终结果，但显著提高了数值稳定性。
其次是梯度监控要点：计算完成后，应验证 $\sum_j \frac {\partial L}{\partial x_j} \approx 0$，这源于 Jacobian 矩阵行和为零的性质，可作为正确性检查的便捷方式。
第三是批处理维度：对于批量输入 $X \in \mathbb {R}^{B \times K}$，梯度计算沿最后一个维度（$K$ 维）独立进行，每个样本的梯度形状保持为 $B \times K$。在 PyTorch 中，可通过 F.softmax(input, dim=-1) 和自动求 grad 系统正确处理。
最后是内存与计算权衡：虽然 Jacobian 矩阵是 $K \times K$ 稠密矩阵，但通过上述闭式公式可以在 $O (K)$ 时间复杂度内完成梯度计算，无需显式存储 $K^2$ 个元素，这对大 vocabulary 场景尤为重要。
总结
Softmax 的 Jacobian 矩阵推导 $J_{ij} = S_i (\delta_{ij} - S_j)$ 揭示了深度学习中梯度流动的数学本质。掌握这一推导过程，不仅能够帮助理解反向传播的链式法则，更能在自定义层或特殊损失函数设计中正确实现梯度传播。当与 Cross-Entropy 结合时，简化的 $S - y$ 形式体现了深度学习中理论推导与工程实现的完美统一。
资料来源：ELI Bendersky 的 Softmax 导数推导文章、idlemachines.co.uk 的 Softmax Backward 实现指南。