深入理解softmax

​

前言

本文代码基于Pytorch实现。

一、softmax的定义及代码实现

1.1 定义

$$softmax(x_i)=\frac{exp(x_i)}{{\sum }_j^{n}exp(x_j)}$$

1.2 代码实现

def softmax(X):
    '''
    实现softmax
    输入X的形状为[样本个数，输出向量维度]
    '''
    return torch.exp(X) / torch.sum(torch.exp(X), dim=1).reshape(-1, 1)

>>> X = torch.randn(5, 5)
>>> y = softmax(X)
>>> torch.sum(y, dim=1)
tensor([1.0000, 1.0000, 1.0000, 1.0000])

二、softmax的作用

softmax可以对线性层的输出做规范化校准：保证输出为非负且总和为1。
因为如果直接将未规范化的输出看作概率，会存在2点问题：

1. 线性层的输出并没有限制每个神经元输出数字的总和为1；
2. 根据输入的不同，线性层的输出可能为负值。
   这2点违反了概率的基本公理。

三、softmax的上溢出(overflow)与下溢出(underflow)

3.1 上溢出

当$x_i$的取值过大时，指数运算的取值过大，若超出精度表示范围，则上溢出。

>>> torch.exp(torch.tensor([1000]))
tensor([inf])

3.2 下溢出

当向量$\mathbf{x}$的每个元素$x_i$的取值均为绝对值很大的负数时，则$exp(x_i)$的数值很小超出了精度范围向下取0，分母${\sum }_{j}exp(j)$的取值为0。

>>> X = torch.ones(1, 3) * (-1000)
>>> softmax(X)
tensor([[nan, nan, nan]])

3.3 避免溢出

参考¹中的技巧：

1. 找到向量$\mathbf{x}$中的最大值:
   $$c=max(\mathbf{x})$$
2. $softmax$的分子、分母同时除以$c$
   $$softmax(x_i-c)=\frac{exp(x_i-c)}{{\sum }_j^{n}exp(x_j-c)}=\frac{exp(x_i)exp(-c)}{{\sum }_j^{n}exp(x_i)exp(-c)}=softmax(x_i)$$
   经过上述变换，分子的最大取值变为了$exp(0)=1$，避免了上溢出；
   分母中至少会$+1$，避免了分母为0造成下溢出。
   $$\begin{gathered} \sum _j^{n}exp(x_j-c)=exp(x_i-c)+exp(x_2-c)+...+exp(x_{max}-c) \\ =exp(x_1-c)+exp(x_2-c)+...+1 \end{gathered}$$

def softmax_trick(X):
    c, _ = torch.max(X, dim=1, keepdim=True)
    return torch.exp(X - c) / torch.sum(torch.exp(X - c), dim=1).reshape(-1, 1)
>>> X = torch.tensor([[-1000, 1000, -1000]])
>>> softmax_trick(X)
tensor([0., 1., 0.])
>>> softmax(X)
tensor([[0., nan, 0.]])

pytorch的实现中已经做过了防止溢出的处理，所以，其运行结果与softmax_trick一致。

import pytorch.nn.functional as F
>>> X = torch.tensor([[-1000., 1000., -1000.]])
>>> F.softmax(X, dim=1)
tensor([[0., 1., 0.]])

3.4 Log-Sum_Exp Trick²（取log操作）

1. 避免下溢出
对数运算可以将相乘变为相加，即：$log(x_1{x}_2)=log(x_1)+log(x_2)$。 当两个很小的数$x_1、x_2$相乘时，其乘积会变得更小，超出精度则下溢出；而对数操作将乘积变为相加，降低了下溢出的风险。
2. 避免上溢出
$log-softmax$的定义：
$$\begin{array}{rl}log-softmax& =log[softmax(x_i)]\\ & =log(\frac{exp(x_i)}{{\sum }_j^{n}exp(x_j)})\\ & =x_i-log[\sum _j^{n}exp(x_j)]\end{array}$$
令$y=log{\sum }_j^{n}exp(x_j)$，当$x_j$的取值过大时，$y$存在上溢出的风险，因此，采用与3.3中同样的Trick:
$$\begin{array}{rl}y& =log\sum _j^{n}exp(x_j)\\ & =log\sum _j^{n}exp(x_j-c)exp(c)\\ & =c+log\sum _j^{n}exp(x_j-c)\end{array}$$
当$c=max(\mathbf{x})$时，可避免上溢出。
此时，$log-softmax$的计算公式变为：（其实等价于直接对3.3节的Trick取对数）
$$log-softmax=(x_i-c)-log\sum _j^{n}exp(x_j-c)$$
代码实现：

def log_softmax(X):
	c, _ = torch.max(X, dim=1, keepdim=True)
	return X - c - torch.log(torch.sum(torch.exp(X-c), dim=1, keepdim=True))
>>> X = torch.tensor([[-1000., 1000., -1000.]])
>>> torch.exp(log_softmax(X))
tensor([[0., 1., 0.]])
# pytorch API实现
>>> torch.exp(F.log_softmax(X, dim=1))
tensor([[0., 1., 0.]])

3.5 log-softmax与softmax的区别³

结合3.3节的Trick及我自己的理解：

1. 在pytorch的实现中，softmax的运算结果等价于对log_softmax的结果作指数运算

>>> X = torch.tensor([[-1000., 1000., -1000.]])
>>> torch.exp(F.log_softmax(X, dim=1)) == F.softmax(X)
tensor([[True, True, True]])

2. 使用$log$运算之后求导更方便，可以加快反向传播的速度⁴
   $$\begin{array}{rl}\frac{\partial }{\partial x_i}logsoftmax& =\frac{\partial }{\partial x_i}[x_i-log\sum _j^{n}exp(x_j)]\\ & =1-softmax(x_i)\end{array}$$

​

---

1. BTTB你不知道的softmax ↩︎

2. The Log-Sum_Exp Trick ↩︎

3. log-softmax与softmax的区别 ↩︎

4. 动手学深度学习：softmax回归 ↩︎