1 数学定义

softmax函数的定义为： $$\sigma (\overrightarrow{z}{)}_i=\frac{e^{z_i}}{\sum _j{e}^{z_j}},(1)$$ 其中，$\overrightarrow{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在计算机的实际运算过程中，容易遇到上溢和下溢问题。 设$\overrightarrow{z}$是softmax函数的输入。假如$\overrightarrow{z}$中的数值都极小（趋于负无穷大），这时公式 1的分母接近0，容易导致产生的数超出浮点型的上限，这被称为下溢；反之，若$\overrightarrow{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 (\overrightarrow{z}{)}_i$. 计算该雅克比矩阵要分两种情况讨论。一是$\frac{\partial s_i}{\partial z_i}$（雅克比矩阵的对角线），二是$\frac{\partial s_i}{\partial z_k},i\ne k$的情况。

第一种情况的导数计算如下： $$\begin{array}{rl}{s}_i& =\frac{e^{z_i}}{\sum _j{e}^{z_j}}=\frac{1}{\sum _{j,j\ne 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{array}$$ 第二种情况的导数$\frac{\partial s_i}{z_k}(i\ne k)$计算如下： $$\begin{array}{rl}\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{array}$$ 因此对softmax函数求导，其雅可比矩阵为： $$\left(\begin{array}{ccccc}{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\\ ⋮& ⋮& ⋮& \ddots & ⋮\\ -s_n{s}_1& -s_n{s}_2& -s_n{s}_3& \cdots & s_n(1-s_n)\end{array}\right)$$

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

4 结语

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

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