Softmax の定義

逆温度 $\lambda>0$ の $\text{softmax}:\mathbb R^d\to\mathbb R^d$ は次のように定義されます。

$$\text{softmax}(\mathbf{x}) = \frac{1}{\sum\limits_{i=1}^d\exp(\lambda x_i)} \begin{bmatrix} \exp(\lambda x_1) \\ \dots \\ \exp(\lambda x_d) \end{bmatrix}$$

Softmax は log-sum-exp の勾配

log-sum-exp 関数 $\text{lse}:\mathbb R^d\to\mathbb R$ は次のように定義されます。

$$\text{lse}(\mathbf{x}) = \frac{1}{\lambda} \log\sum\limits_{i=1}^d\exp(\lambda x_i)$$

$\text{softmax}(\mathbf{x})$ は、log-sum-exp 関数の勾配です。 つまり、次が成り立ちます。

Proposition 1.
$\nabla\text{lse}(\mathbf{x}) = \text{softmax}(\mathbf{x})$

証明

$$\begin{align*} \nabla\text{lse}(\mathbf{x}) &= \frac{1}{\lambda} \nabla\log\sum\limits_{i=1}^d\exp(\lambda x_i) \\ &= \frac{1}{\lambda} \frac{1}{\sum\limits_{i=1}^d\exp(\lambda x_i)} \nabla\sum\limits_{i=1}^d\exp(\lambda x_i) \\ &= \frac{1}{\lambda} \frac{\lambda}{\sum\limits_{i=1}^d\exp(\lambda x_i)} \begin{bmatrix} \exp(\lambda x_1) \\ \dots \\ \exp(\lambda x_d) \end{bmatrix} \\ &= \frac{1}{\sum\limits_{i=1}^d\exp(\lambda x_i)} \begin{bmatrix} \exp(\lambda x_1) \\ \dots \\ \exp(\lambda x_d) \end{bmatrix} \\ &= \text{softmax}(\mathbf{x}) \end{align*}$$

Softmax の微分

$\text{softmax}(\mathbf{x})$ の微分は、次のように計算できます。

Proposition 2.
$J_{\text{softmax}}(\mathbf{x}) = \lambda(\text{diag}(\text{softmax}(\mathbf{x})) - \text{softmax}(\mathbf{x})\text{softmax}(\mathbf{x})^T)$

証明

$\displaystyle \frac{\partial (\text{softmax}(\mathbf{x}))_i}{\partial x_j}$ を計算します。

$i = j$ の場合

$$\begin{align*} \frac{\partial (\text{softmax}(\mathbf{x}))_i}{\partial x_i} &= \frac{\partial}{\partial x_i}\left(\frac{\exp(\lambda x_i)}{\sum\limits_{k=1}^d\exp(\lambda x_k)}\right) \\ &= \frac{\lambda \exp(\lambda x_i)\sum\limits_{k=1}^d\exp(\lambda x_k) - \lambda\exp(\lambda x_i)^2}{\left(\sum\limits_{k=1}^d\exp(\lambda x_k)\right)^2} \\ &= \lambda\text{softmax}(\mathbf x)_i - \lambda\text{softmax}(\mathbf x)_i^2 \\ \end{align*}$$

$i \neq j$ の場合

$$\begin{align*} \frac{\partial (\text{softmax}(\mathbf{x}))_i}{\partial x_j} &= \frac{\partial}{\partial x_j}\left(\frac{\exp(\lambda x_i)}{\sum\limits_{k=1}^d\exp(\lambda x_k)}\right) \\ &= -\frac{\lambda \exp(\lambda x_i)\exp(\lambda x_j)}{\left(\sum\limits_{k=1}^d\exp(\lambda x_k)\right)^2} \\ &= -\lambda \text{softmax}(\mathbf x)_i\text{softmax}(\mathbf x)_j \\ \end{align*}$$

よって、微分は次のように計算できます。

$$J_{\text{softmax}}(\mathbf{x}) = \lambda(\text{diag}(\text{softmax}(\mathbf{x})) - \text{softmax}(\mathbf{x})\text{softmax}(\mathbf{x})^T)$$

References

• The Lipschitz Constant of Self-Attention

---