一文弄懂LogSumExp技巧

引言

今天来学习下LogSumExp(LSE)​​¹​​​技巧，主要解决计算Softmax或CrossEntropy​​²​​时出现的上溢(overflow)或下溢(underflow)问题。

我们知道编程语言中的数值都有一个表示范围的，如果数值过大，超过最大的范围，就是上溢；如果过小，超过最小的范围，就是下溢。

什么是LSE

LSE被定义为参数指数之和的对数：
$\text{LSE}(x_1,\cdots,x_n) = \log \sum_{i=1}^n \exp(x_i) =\log \left(\exp(x_1) + \cdots + \exp(x_n) \right)$
输入可以看成是一个n维的向量，输出是一个标量。

为什么需要LSE

在机器学习中，计算概率输出基本都需要经过Softmax函数，它的公式应该很熟悉了吧
$\text{Softmax}(x_i) = \frac{\exp(x_i)}{\sum_{j=1}^n \exp(x_j)} \tag{1}$
但是Softmax存在上溢和下溢大问题。如果$x_i$太大，对应的指数函数也非常大，此时很容易就溢出，得到​​​nan​​​结果；如果$x_i$太小，或者说负的太多，就会导致出现下溢而变成0，如果分母变成0，就会出现除0的结果。

此时我们经常看到一个常见的做法是(其实用到的是指数归一化技巧, exp-normalize​​³​​​)，先计算$x$中的最大值$b = \max_{i=1}^n x_i$，然后根据
$\begin{aligned} \text{Softmax}(x_i) &= \frac{\exp(x_i)}{\sum_{j=1}^n \exp(x_j)} \\ &= \frac{\exp(x_i - b) \cdot \exp(b)}{\sum_{j=1}^n \left (\exp(x_j - b) \cdot \exp(b) \right)} \\ &= \frac{\exp(x_i - b) \cdot \exp(b)}{ \exp(b) \cdot \sum_{j=1}^n \exp(x_j - b) } \\ &= \frac{\exp(x_i - b)}{\sum_{j=1}^n \exp(x_j - b)} \\ &= \text{Softmax}(x_i - b) \end{aligned} \tag{2}$
这种转换是等价的，经过这一变换，就避免了上溢，最大值变成了$\exp(0)=1$；同时分母中也会有一个1，就避免了下溢。

我们通过实例来理解一下。

def bad_softmax(x):
  y = np.exp(x)
  return y / y.sum()
 
x = np.array([1, -10, 1000])
print(bad_softmax(x))

... RuntimeWarning: overflow encountered in exp
... RuntimeWarning: invalid value encountered in true_divide
array([ 0.,  0., nan])

接下来进行上面的优化，并进行测试：

def softmax(x):
  b = x.max()
  y = np.exp(x - b)
  return y / y.sum()
 
print(softmax(x))

array([0., 0., 1.])

我们再看下是否会出现下溢：

x = np.array([-800, -1000, -1000])
print(bad_softmax(x))
# array([nan, nan, nan])
print(softmax(x))
# array([1.00000000e+00, 3.72007598e-44, 3.72007598e-44])

嗯，看来解决了这个两个问题。

等等，不是说LSE吗，怎么整了个什么归一化技巧。

好吧，回到LSE。

我们对Softmax取对数，得到：
$\begin{aligned} \log \left( \text{Softmax}(x_i) \right) &= \log \frac{\exp(x_i)}{\sum_{j=1}^n \exp(x_j)} \\ &= x_i - \log \sum_{j=1}^n \exp(x_j) \\ \end{aligned} \tag{3}$
因为上面最后一项也有上溢的问题，所以应用同样的技巧，得
$\log \sum_{j=1}^n \exp(x_j) = \log \sum_{j=1}^n \exp(x_j - b) \exp(b) = b + \log \sum_{j=1}^n \exp(x_j - b) \tag{4}$
$b$同样是取$x$中的最大值。

这样，我们就得到了LSE的最终表示：
$\text{LSE}(x) = b + \log \sum_{j=1}^n \exp(x_j - b) \tag{5}$
此时，Softmax也可以这样表示：
$\text{Softmax}(x_i) = \exp \left( x_i - b - \log \sum_{j=1}^n \exp(x_j - b) \right) \tag{6}$
对LogSumExp求导就得到了exp-normalize(Softmax)的形式，
$\frac{\partial \left (b + \log \sum_{j=1}^n \exp(x_j - b) \right )}{\partial x_j} = \frac{\exp(x_i - b)}{\sum_{j=1}^n \exp(x_j - b)} \tag{7}$

那我们是使用exp-normalize还是使用LogSumExp呢？

如果你需要保留Log空间，那么就计算$\log(\text{Softmax})$，此时使用LogSumExp技巧；如果你只需要计算Softmax，那么就使用exp-normalize技巧。

怎么实现LSE

实现LSE就很简单了，我们通过代码实现一下。

def logsumexp(x):
  b = x.max()
  return b + np.log(np.sum(np.exp(x - b)))
 
def softmax_lse(x):
  return np.exp(x - logsumexp(x))

上面是基于LSE实现了Softmax，下面测试一下：

> x1 = np.array([1, -10, 1000])
> x2 = np.array([-900, -1000, -1000])
> softmax_lse(x1)
array([0., 0., 1.])
> softmax(x1)
array([0., 0., 1.])
> softmax_lse(x2)
array([1.00000000e+00, 3.72007598e-44, 3.72007598e-44])
> softmax(x2)
> array([1.00000000e+00, 3.72007598e-44, 3.72007598e-44])

最后我们看一下数值稳定版的Sigmoid函数

数值稳定的Sigmoid函数

我们知道Sigmoid函数公式为：
$\sigma(x) = \frac{1}{1 + \exp(-x)} \tag{8}$
对应的图像如下：

其中包含一个$\exp(-x)$，我们看一下$e^x$的图像：

从上图可以看出，如果$x$很大，$e^x$会非常大，而很小就没事，变成无限接近$0$。

当Sigmoid函数中的$x$负的特别多，那么$\exp(-x)$就会变成$\infty$，就出现了上溢；

那么如何解决这个问题呢？$\sigma(x)$可以表示成两种形式：
$\sigma(x) = \frac{1}{1 + \exp(-x)} = \frac{\exp(x)}{1 + \exp(x)} \tag{9}$
当$x \geq 0$时，我们根据$e^{x}$的图像，我们取$\frac{1}{1 + \exp(-x)}$的形式；

当$x < 0$时，我们取$\frac{\exp(x)}{1 + \exp(x)}$

# 原来的做法
def sigmoid_naive(x):
  return 1 / (1 + math.exp(-x))
  
# 优化后的做法
def sigmoid(x):
  if x < 0:
    return math.exp(x) / (1 + math.exp(x))
  else:
    return 1 / (1 + math.exp(-x))

然后用不同的数值进行测试：

> sigmoid_naive(2000)
1.0
> sigmoid(2000)
1.0
> sigmoid_naive(-2000)
OverflowError: math range error
> sigmoid(-2000)
0.0

References

---

1. ​​The Log-Sum-Exp Trick ​​​ ​​↩︎​​
2. ​​一文弄懂交叉熵损失 ​​​ ​​↩︎​​
3. ​​Exp-normalize trick​​​ ​​↩︎​​