Softmax 输出及其反向传播推导

Softmax 是在神经网络的中应用极其普遍的一种操作, 搞清楚 softmax 极其对对应的 cross entropy loss function 的意义及数学推导过程可以有好的帮助理解源码. 在这篇文章中, 详细地记录了 softmax, cross entropy loss, 和其对应的梯度的推导过程, 最后, 附加了 MXNet 中的梯度计算代码, 可以帮助直观的理解开源代码中梯度的计算过程.

Softmax Function

首先给出 Softmax function 的定义
\[ y_c=\zeta(\textbf{z})_c = \dfrac{e^{z_c}}{\sum_{d=1}^C{e^{z_d}}} \text{ for } c=1, \dots, C \]
其中, \(\textbf{z}\) 是输入, C 维的向量, \(\textbf{y}\) 是输出, C 维向量. 而且, \(y\) 的取值范围是 \(\lbrack 0,1 \rbrack\). 这样, C 个类别的概率分别是
\[ \left\lbrack\matrix{P(t=1\vert\textbf{z})\cr \vdots \cr P(t=C\vert\textbf{z})}\right\rbrack=\left\lbrack\matrix{\zeta(\textbf{z}_1)\cr \vdots \cr \zeta(\textbf{z}_C)}\right\rbrack = \dfrac{1}{\sum_{d=1}^C{e^{z_d}}}\left\lbrack\matrix{e^{z_1} \cr \vdots \cr e^{e^{z_C}} }\right\rbrack \]

Softmax 函数的导数

先给出 softmax 的导数, 后面会用到. 首先, 为了方便, 定义符号 \(\sum_C=\sum_{d=1}^C{e^{z_d}}\text{ for } c=1 \dots C\), 所以, 给定一个feature vector, 输出类别 \(c\) 的概率是 \(y_c=\dfrac{e^{z_c}}{\sum_c}\).
\[\begin{array}{l}if\;i = j:\;\dfrac{{\partial {y_i}}}{{\partial {z_j}}} = \dfrac{{\partial {y_i}}}{{\partial {z_i}}} = \dfrac{{\partial \dfrac{{{e^{{z_i}}}}}{{\sum\nolimits_C {} }}}}{{\partial {z_i}}} = \dfrac{{{e^{{z_i}}} \cdot \sum\nolimits_C {} - {e^{{z_i}}} \cdot {e^{{z_i}}}}}{{{{\left( {\sum\nolimits_C {} } \right)}^2}}} = \dfrac{{{e^{{z_i}}}\left( {\sum\nolimits_C {} - {e^{{z_i}}}} \right)}}{{\sum\nolimits_C {} \cdot \sum\nolimits_C {} }}{\rm{ = }}\dfrac{{{e^{{z_i}}}}}{{\sum\nolimits_C {} }} \cdot \dfrac{{\sum\nolimits_C {} - {e^{{z_i}}}}}{{\sum\nolimits_C {} }} = {y_i} \cdot \left( {1 - {y_i}} \right)\\if\;i \ne j:\;\dfrac{{\partial {y_i}}}{{\partial {z_j}}} = \dfrac{{\partial \dfrac{{{e^{{z_i}}}}}{{\sum\nolimits_C {} }}}}{{\partial {z_j}}} = \dfrac{{0 \cdot \sum\nolimits_C {} - {e^{{z_i}}} \cdot {e^{{z_j}}}}}{{{{\left( {\sum\nolimits_C {} } \right)}^2}}} = - \dfrac{{{e^{{z_i}}}}}{{\sum\nolimits_C {} }} \cdot \dfrac{{{e^{{z_j}}}}}{{\sum\nolimits_C {} }} = - {y_i} \cdot {y_j}\end{array}\]

损失函数 cross entropy loss function

目标函数本质上是通过最大似然推倒出来的, cross entropy 不过是个高大上的名词.
\[ \underset{\theta}{\arg\max} \cal{L}(\theta \vert \textbf{t,z}) \]
把上述似然函数写成联合分布的形式:
\[ P(\textbf{t,x}\vert \theta)=P(\textbf{t}\vert \textbf{z}, \theta)P(\textbf{z}\vert \theta) \]

\(\textbf{z}\) 的分布对于结果没有影响, 只需要关心 \(P(\textbf{t}\vert \textbf{z}, \theta)\). 对于固定的 \(\theta\), 可以记成\(P(\textbf{t}\vert \textbf{z})\), 又因为各个 \(t_i\) 是独立的, 只有一个 \(t_i\) 是激活的, 所以,
\[ P(\textbf{t}\vert\textbf{z})=\prod_{i=c}^{C}P\left( t_c \vert \textbf{z}\right)^{t_c}=\prod_{i=c}^C\zeta(\textbf{z})_c^{t_c}=\prod_{i=c}^{C}y_c^{t_c} \]
maximizing 上面的式子, 相当于 minimizing 上述式子去对数, 再取负数, 最终得到的优化目标是:
\[ -\log \prod_{i=c}^{C}y_c^{t_c}=-\sum_{i=c}^{C}t_c \cdot \log(y_c) \]
一个 batch 的 cross-entropy error function 是
\[ \xi(X, Y) = \sum_{i=1}^n \xi(\textbf{t}_i, \textbf{y}_i) = - \sum_{i=1}^n \sum_{i=c}^C t_{ic} \cdot \log(y_{ic}) \]
当且仅当样本 \(i\)属于类别 \(c\) 的时候, \(t_{ic}=1\)

corss entropy loss function 的梯度

\[ \begin{align} \dfrac{\partial\xi}{\partial z_i} &= - \sum_{j=1}^C \dfrac{\partial t_j \log(y_j)}{\partial z_i} \\ &= - \sum_{j=1}^C t_j \dfrac{\partial \log(y_j)}{\partial z_i} \\ &= - \sum_{j=1}^C t_j \dfrac{1}{y_j} \dfrac{\partial y_j}{\partial z_i} \\ &= - \dfrac{t_i}{y_i} \dfrac{\partial y_i}{\partial z_i} - \sum_{j \neq i}^C \dfrac{t_j}{y_j} \dfrac{\partial y_j}{\partial z_i} \\ &= - \dfrac{t_i}{y_i} y_i(1-y_i) - \sum_{j \neq i}^{C} \dfrac{t_j}{y_j}(-y_jy_j) \\ &= -t_i + t_iy_i + \sum_{j \neq i}^{C} t_jy_i \\ &= -t_i + \sum_{j=1}^C t_jy_i \\ &= -t_i + y_i \sum_{j=1}^C t_j \\ &= y_i - t_i \end{align} \]

源码分析

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
template<typename DType>
inline void SoftmaxGrad(Tensor<cpu, 2, DType> dst,
const Tensor<cpu, 2, DType> &src,
const Tensor<cpu, 1, DType> &label) {
for (index_t y = 0; y < dst.size(0); ++y) {
const index_t k = static_cast<int>(label[y]);
for (index_t x = 0; x < dst.size(1); ++x) {
if (x == k) {
dst[y][k] = src[y][k] - 1.0f;
} else {
dst[y][x] = src[y][x];
}
}
}
}