如何实现log_softmax()，以更好的速度和数值稳定性计算其值（和梯度）？

Question

MXNet和PyTorch都提供了计算log(softmax())的特殊实现，速度更快，数值更稳定。但是，我在这两个软件包中都找不到这个函数log_softmax()的实际Python实现。

有谁能解释一下这个函数是如何实现的，或者最好能给我指出相关的源代码？

Answer 1

• 数值误差。

>>> x = np.array([1, -10, 1000])
>>> np.exp(x) / np.exp(x).sum()
RuntimeWarning: overflow encountered in exp
RuntimeWarning: invalid value encountered in true_divide
Out[4]: array([ 0.,  0., nan])

有两种方法可以避免计算softmax时的数值误差。

• Exp归一化

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

>>> exp_normalize(x)
array([0., 0., 1.])

• 对数和 Exp

def log_softmax(x):
    c = x.max()
    logsumexp = np.log(np.exp(x - c).sum())
    return x - c - logsumexp

请注意，上式中b、c的合理选择是max(x)。有了这个选择，因exp而溢出是不可能的。移位后指数化的最大数字是0。

Answer 2

你可以找到其中一个CPU的实现 此处 和一个矢量化版本 此处 (这是日志版本，从 vec_host_softmax_lastdim).

你可以找到一个CUDA的实现 此处，然后调用 softmax_warp_forward.

它们都很相似，只是语法不同。如你所见，通常有一个标志定义是否使用log.来计算softmax，即用LogSoftMax代替SoftMax。