Softmax 定义

对向量 $\mathbf{x} = (x_1, \dots, x_n)$，softmax 定义为： $$ y(\mathbf{x}) = \frac{e^{x_i}}{\sum_{k=1}^n e^{x_k}}, \quad i = 1, \dots, n $$ 记： $$ S = \sum_{k=1}^n e^{x_k} $$

Softmax 的求导

我们要求： $$ \frac{\partial y_i}{\partial x_j} $$ 情况 1: $i = j$ ，此时

  • $\frac{\partial e^{x_i}}{\partial x_j} = e^{x_i}$
  • $\frac{\partial S}{\partial x_j} = e^{x_i}$

$$ \begin{align*} \frac{\partial y_i}{\partial x_i} &= \frac{\partial}{\partial x_i} \left( \frac{e^{x_i}}{S} \right) \ &= \frac{\frac{\partial e^{x_i}}{\partial x_j} S-e^{x_i}\frac{\partial S}{\partial x_j}}{S^2} \ &= \frac{e^{x_i} S - e^{x_i} e^{x_i}}{S^2} \ &= \frac{e^{x_i}}{S} \left( 1 - \frac{e^{x_i}}{S} \right) \ &= \sigma_i (1 - \sigma_i) \end{align*} $$

情况 2: $i \neq j$

  • $\frac{\partial e^{x_i}}{\partial x_j} = 0$
  • $\frac{\partial S}{\partial x_j} = e^{x_j}$

$$

\begin{align*}

\frac{\partial y_i}{\partial x_j} &= \frac{\partial}{\partial x_j} \left( \frac{e^{x_i}}{S} \right) \ &= \frac{\frac{\partial e^{x_i}}{\partial x_j} S-e^{x_i}\frac{\partial S}{\partial x_j}}{S^2} \ &= -\frac{e^{x_i} e^{x_j}}{S^2} \ &= -y_i y_j \end{align*} $$ 统一形式（雅可比矩阵）

$$

\boxed{\frac{\partial y_i}{\partial x_j} = y_i (\delta_{ij} - y_j)}

$$

其中 $\delta_{ij}$ 是 Kronecker delta。

对应的 Jacobian 矩阵：

$$

J = \text{diag}(y) - y y^T

$$

总结

设

$$ \begin{cases} y = softmax(z) \ L = f(y) \end{cases} $$

其中，$L$ 表示 Loss，$f(.)$ 表示 Loss 函数，$y = \left[ y_1\ \ y_2\ \ y_3 \right]$，$z = \left[ z_1\ \ z_2\ \ z_3 \right]$，若现在我们想求 $\frac{\partial L}{\partial z_j}$，要怎么算呢？

根据链式法则，我们有 $\frac{\partial L}{\partial z_j} = \frac{\partial L}{\partial y} \frac{\partial y}{\partial z_j}$，所以我们分别来看这两项。

其中 $$ \begin{cases} \frac{\partial y_i}{\partial z_j} = y_i(1 - y_i)，& \text{当 } i = j \ \frac{\partial y_i}{\partial z_j} = -y_i y_j，& \text{当 } i \neq j \end{cases} $$ 现在就可以来推导 $\frac{\partial L}{\partial z_j}$，我们有：

$$

\frac{\partial L}{\partial z_j} = \frac{\partial L}{\partial y} \frac{\partial y}{\partial z_j} = \sum_{i=1}^l \frac{\partial L}{\partial y_i} \frac{\partial y_i}{\partial z_j} = y_j\left( dy_j - \sum_{j=1}^l y_j dy_j \right)

$$

举个例子，若我们现在想求 $\frac{\partial L}{\partial z_1}$，我们将 $\frac{\partial L}{\partial y} = \left[ m_1\ \ m_2\ \ m_3 \right]$ 代入上面公式，极为 $dy=\left[ dy_1\ \ dy_2\ \ dy_3 \right]=\left[ m_1\ \ m_2\ \ m_3 \right]$ 则有：

$$

\frac{\partial L}{\partial z_1} = m_1\left( y_1 - y_1^2 \right) - m_2 y_1 y_2 - m_3 y_1 y_3

$$

现在，针对所有的 $z$，我们将 $\frac{\partial L}{\partial z}$ 写成矩阵表达式有：

$$

\begin{aligned}

\frac{\partial L}{\partial z} &= \frac{\partial L}{\partial y} \frac{\partial y}{\partial z} = dy\left( diag(y) - y^T y \right) \

&= \left[ m_1\ \ m_2\ \ m_3 \right] \left( \begin{bmatrix} y_1 & 0 & 0 \ 0 & y_2 & 0 \ 0 & 0 & y_3 \end{bmatrix} - \begin{bmatrix} y_1 \ y_2 \ y_3 \end{bmatrix} \left[ y_1\ \ y_2\ \ y_3 \right] \right) \

&= \left[ m_1\ \ m_2\ \ m_3 \right] \begin{bmatrix} y_1 - y_1^2 & -y_1 y_2 & -y_1 y_3 \ -y_2 y_1 & y_2 - y_2^2 & -y_2 y_3 \ -y_3 y_1 & -y_3 y_2 & y_3 - y_3^2 \end{bmatrix}

\end{aligned}

$$

至此，大家记住这两个重要的结论：

$$ \begin{aligned}

\frac{\partial L}{\partial z} &= \frac{\partial L}{\partial y} \frac{\partial y}{\partial z} = dy\left( diag(y) - y^T y \right) \ \frac{\partial L}{\partial z_j} &= y_j\left( dy_j - \sum_{j=1}^l y_j dy_j \right) \end{aligned} $$