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} \]

源码分析

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];
      }
    }
  }
}