1 数学定义

softmax函数的定义为： \[ \sigma(\vec z)_i = \frac{e^{z_i}}{\sum_j e^{z_j}}, \qquad(1)\] 其中，\(\vec z\)为输入向量。

下面的代码展示了softmax函数的一种简单的实现方式。

import torch 

def my_softmax(z, dim):
    # z.shape == bn, n
    z = torch.exp(z)
    z = z / torch.sum(z, dim=dim)
    return z 

tensor_test = torch.rand(10, 100)
diff = (my_softmax(tensor_test, 0) - torch.softmax(tensor_test, 0)).abs().mean().item()
assert diff < 1e-6, diff
print('Difference: ', diff)

Difference:  6.180256750809576e-09

2 Softmax的上溢和下溢问题

softmax在计算机的实际运算过程中，容易遇到上溢和下溢问题。 设\(\vec z\)是softmax函数的输入。假如\(\vec z\)中的数值都极小（趋于负无穷大），这时公式 1的分母接近0，容易导致产生的数超出浮点型的上限，这被称为下溢；反之，若\(\vec z\)中存在特别大的数，由于函数\(e^{x}\)增长极快，输出的数也很可能超出浮点型的上限，这被称为上溢。

下面的代码展示了softmax的计算中出现上溢出和下溢出的情形。

out = my_softmax(tensor_test + 1e5, 0)
if torch.any(torch.isnan(out)):
    print('检测到上溢！')

out = my_softmax(tensor_test - 1e5, 0)
if torch.any(torch.isnan(out)):
    print('检测到下溢！')

out = torch.softmax(tensor_test + 1e5, 0)
out2 = torch.softmax(tensor_test - 1e5, 0)
if not (
    torch.any(torch.isnan(out)) or 
    torch.any(torch.isnan(out2))
):
    print('torch.softmax函数没有出现上溢下溢问题。')

检测到上溢！
检测到下溢！
torch.softmax函数没有出现上溢下溢问题。

尽管我们实现的简单的softmax函数会发生上溢和下溢，但torch的softmax函数没有出现问题。torch是如何做到的呢？

def my_softmax_2(z, dim):
    # z.shape == bn, n

    # 减去最大值
    z -= z.max(dim=dim)[0]

    z = torch.exp(z)
    z = z / torch.sum(z, dim=dim)
    return z 

out = my_softmax_2(tensor_test + 1e5, 0)
if torch.any(torch.isnan(out)):
    print('检测到上溢！')

out = my_softmax_2(tensor_test - 1e5, 0)
if torch.any(torch.isnan(out)):
    print('检测到下溢！')

diff = (my_softmax_2(tensor_test, 0) - torch.softmax(tensor_test, 0)).abs().mean().item()
assert diff < 1e-6, diff
print('Difference: ', diff)

Difference:  9.164214387347158e-10

3 Softmax函数求导

softmax是一个将向量映射为向量的函数。求其输出对输入的导数得到一个雅克比矩阵。为了方便，记\(s_i:=\sigma(\vec z)_i\). 计算该雅克比矩阵要分两种情况讨论。一是\(\frac{\partial s_i}{\partial z_i}\)（雅克比矩阵的对角线），二是\(\frac{\partial s_i}{\partial z_k}, i\neq k\)的情况。

第一种情况的导数计算如下： \[ \begin{aligned} s_i &= \frac{e^{z_i}}{\sum_j e^{z_j}} = \frac{1}{\sum_{j, j\neq i} e^{z_j - z_i} + 1}\\ \frac{\partial s_i}{\partial z_i} &= \frac{e^{z_i}\sum_j e^{z_j}}{(\sum_j e^{z_j})^2} - \frac{e^{z_i} e^{z_i}}{(\sum_j e^{z_j})^2} \\ &= s_i - s_i^2 \\ &= s_i(1 - s_i), \end{aligned} \] 第二种情况的导数\(\frac{\partial s_i}{z_k} (i \neq k)\)计算如下： \[ \begin{aligned} \frac{\partial s_i}{\partial z_k} &= -\frac{e^{z_i} e^{z_k}}{(\sum_j e^{z_j})^2} \\ &= -s_i s_k \end{aligned} \] 因此对softmax函数求导，其雅可比矩阵为： \[ \begin{pmatrix} s_1(1-s_1) & -s_1 s_2 & -s_1 s_3 & \cdots & -s_1 s_n \\ -s_2 s_1 & s_2(1-s_2) & -s_2 s_3 & \cdots & -s_2 s_n \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ -s_n s_1 & -s_n s_2 & -s_n s_3 & \cdots & s_n (1 - s_n) \\ \end{pmatrix} \]

观察该矩阵，我们可以发现在两种情况下，矩阵的将会接近0矩阵。一是 当softmax函数的输出为one hot时（即\(\exists s_i = 1\)且\(\forall j \neq i, s_j=0\)），得到的雅克比矩阵为\(0\)矩阵；另一种情况是softmax的输出维度较多，且个维度的值都较为均等的情况，此时\(s_1\approx s_2 \approx s_3 \cdots \approx s_n \approx 0\)。训练中，这两种情况将导致梯度消失，可能影响训练，因此是我们要提前了解的。

4 结语

记得在读硕士的最后一学期时，有一回面试一家公司。这家公司是国内小有名气的中厰。面试官问了softmax函数的问题，我便以上面所述的内容对之。结束时，面试官笑著说，记得几年前，我问同样的问题，几乎没几个人答得上来。但是，今年同样的问题，居然人人都能回答上来呢。

听到面试官的话，我也不禁感慨。看来这几年时间过去，深度学习/人工智能的赛道确实是越来越卷了。作为一名准备面试的学生，这个问题在我看来确实是常识。再过几年，它会不会成为一般本科生，甚至中小学生的常识呢。笑。