在做深度学习的时候，softmax是一个经常遇到的计算公式来，但是还没深究过这个究竟是怎么算的。计算的基础公式是知道的，但是没想到在底层实现的时候还是有不少玄机。

背景

我们定义有 N 个值的数值集合
{x_n}_{n=1}^N
，我们想要求
z = \log \sum_{n=1}^N \exp{x_n}
的值，应该如何计算？

看上去这个问题可能比较容易，但是实际上我们在使用这个数值集合来进行概率分布的时候，通常会涉及到一个问题。

在神经网络中，softmax函数得到一个概率分布，softmax的形式为：
\frac{e^{x_j}}{\sum_{i=1}^n e^{x_i}}
这里的 x_j 是其中的一个值。最经典的loss函数就是使用cross entropy，那么就涉及到需要对这个对数
\log \left( \frac{e^{x_j}}{\sum_{i=1}^n e^{x_i}} \right) = \log(e^{x_j}) – \log \left( \sum_{i=1}^n e^{x_i} \right)
= x_j – \log \left( \sum_{i=1}^n e^{x_i} \right)
这里的分母是所有的和，也就是我们上面所求的 z ，即LogSumExp（之后简称LSE）。

数据溢出问题

因为我们通过softmax想要得到概率分布，假设我们目前有两个例子：一个数据集为 [1000,1000,1000] ，另一个数据集为 [−1000,−1000,−1000] 。

这两组数据在理论上应该得到的概率分布都应该是 [1/3,1/3,1/3] ，但是实际上如果我们直接去计算的时候，会发现：

>>> import math
>>> math.e**1000
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
OverflowError: (34, 'Result too large')
>>> math.e**-1000
0.0

可以发现，一个计算不出来，另一个得0，很明显这些结果都是有问题的。

LSE 求解

那么对于lse的求解需要换一种方式来进行：
\log \text{SumExp}(x_1 \ldots x_n) = \log \left( \sum_{i=1}^{n} e^{x_i} \right)
这个函数可以被重新写成减去一个常数 c 的形式，以增强数值稳定性：
\begin{aligned}
\log \left( \sum_{i=1}^{n} e^{x_i}\right) &= \log \left( \sum_{i=1}^{n} e^{x_i – c} e^c \right)\\
&= \log \left( e^c \sum_{i=1}^{n} e^{x_i – c} \right)\\
&= \log \left( \sum_{i=1}^{n} e^{x_i – c} \right) + \log(e^c)\\
&= \log \left( \sum_{i=1}^{n} e^{x_i – c} \right) + c
\end{aligned}
接下来说明了这个技巧如何应用于计算交叉熵损失函数：
\begin{aligned}
\text{loss}(\text{Softmax}(x_j, x_1, \ldots, x_n)) &= x_j – \log(\text{SumExp}(x_1, \ldots, x_n))\\
&= x_j – \log \left( \sum_{i=1}^{n} e^{x_i} \right)\\
&= x_j – \log \left( \sum_{i=1}^{n} e^{x_i – c} \right) – c
\end{aligned}
那么对于 [1000, 1000, 1000] 的softmax函数计算的交叉熵损失：
\text{loss}(\text{Softmax}(1000, [1000, 1000, 1000])) = 1000 – \log(3) – 1000
= -\log(3)
通过引入常数 ( c )，可以避免这一问题，其中 ( c ) 通常取为 \max{x_i}。

LSE函数是一个数学上的技巧，它对一组数值的指数求和后取对数，通常用于计算概率归一化（例如在softmax函数中）时数值稳定性问题。LSE函数的定义是
\text{LSE}(x_1, \ldots, x_n) = \log(\exp{x_1} + \ldots + \exp{x_n})
在数学上，LSE可以被近似为
\text{LSE}(0, x_1, \ldots, x_n) = \text{LSE}(c, x_1 – c, \ldots, x_n – c)
，其中 ( c ) 是一个常数。这个性质基于函数 ( f(x) ) 在 ( x ) 点的泰勒展开，其中
f(x) \approx f(c) + f'(c)(x – c)
这一性质在计算softmax函数时尤其有用，因为它允许我们减去一个常数 c 来增加数值稳定性。在选择常数 c 时，一个常见的策略是取 ( x ) 的最大值，即 c = \max{x_i}，因为这样可以减少指数函数计算中的数值溢出或下溢问题。

对于LSE函数，以下不等式总是成立的：

\max{x_1, \ldots, x_n} \leq \text{LSE}(x_1, \ldots, x_n) \leq \max{x_1, \ldots, x_n} + \log(n)
所以它实际上是针对max函数的一种平滑操作，从字面上理解来说，LSE函数才是真正意义上的softmax函数。而我们在神经网络里所说的softmax函数其实是近似于argmax函数的，也就是我们平时所用的softmax应该叫做softargmax。