引言

本着“凡我不能创造的，我就不能理解”的思想，本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架，该框架类似PyTorch能实现自动求导。

要深入理解深度学习，从零开始创建的经验非常重要，从自己可以理解的角度出发，尽量不使用外部完备的框架前提下，实现我们想要的模型。本系列文章的宗旨就是通过这样的过程，让大家切实掌握深度学习底层实现，而不是仅做一个调包侠。
本系列文章首发于微信公众号：JavaNLP

在上篇文章中我们了解了逻辑回归中如何确保数值稳定，同样，Softmax回归也存在这个问题。

数值稳定的Softmax

我们先来回顾下Softmax函数的公式：
Softmax(xi)=exp⁡(xi)∑j=1nexp⁡(xj)(1)\text{Softmax}(x_i) = \frac{\exp(x_i)}{\sum_{j=1}^n \exp(x_j)} \tag{1} Softmax(xi)=∑j=1nexp(xj)exp(xi)(1)
但是Softmax存在上溢和下溢的问题。如果xix_ixi太大，对应的指数函数值也非常大，此时很容易就溢出，得到nan结果；如果xix_ixi太小，或者说负的太多，就会导致出现下溢而变成0，如果分母变成0，就会出现除0的结果。

上图是指数函数的图像，上升的非常迅速，哪怕x=1000x=1000x=1000，e1000e^{1000}e1000结果也是巨大的，同时不管xxx取什么，ex>0e^x > 0ex>0。

我们可以利用Softmax函数的特性来进行改进，该特性就是Softmax(xi−b)=Softmax(xi)\text{Softmax}(x_i - b) = \text{Softmax}(x_i)Softmax(xi−b)=Softmax(xi)，看一下它的证明：
Softmax(xi)=exp⁡(xi)∑j=1nexp⁡(xj)=exp⁡(xi−b)⋅exp⁡(b)∑j=1n(exp⁡(xj−b)⋅exp⁡(b))=exp⁡(xi−b)⋅exp⁡(b)exp⁡(b)⋅∑j=1nexp⁡(xj−b)=exp⁡(xi−b)∑j=1nexp⁡(xj−b)=Softmax(xi−b)(2)\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} Softmax(xi)=∑j=1nexp(xj)exp(xi)=∑j=1n(exp(xj−b)⋅exp(b))exp(xi−b)⋅exp(b)=exp(b)⋅∑j=1nexp(xj−b)exp(xi−b)⋅exp(b)=∑j=1nexp(xj−b)exp(xi−b)=Softmax(xi−b)(2)

即从Softmax的参数中减去一个常数bbb，不会改变函数值。

此时，我们令b=max⁡i=1nxib=\max_{i=1}^n x_ib=maxi=1nxi，即取xxx所有元素中的最大值。这样哪怕是xxx中最大的元素，会变成了exp⁡(0)=1\exp(0)=1exp(0)=1，这样避免了上溢。同时分母中也会有一个exp⁡(0)=1\exp(0)=1exp(0)=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])

这里我们有三个负的比较多的数，其中相对负的较少的是-800，其经过Softmax输出为1是比较合理的。

数值稳定的交叉熵损失

回顾一下交叉熵损失的公式：

其中zk=wk⋅x+bkz_k = w_k \cdot x + b_kzk=wk⋅x+bk。

(7)(7)(7)到(8)(8)(8)是因为∑k=1Kyk=1\sum_{k=1}^K y_k = 1∑k=1Kyk=1。最终得到的结果中没有了Softmax，如果它是数值稳定的，那么也就是说，以后我们计算交叉熵损失时，只要传入logits，并通过代码实现公式(8)(8)(8)就好了。

首先，公式(8)(8)(8)中的∑k=1Kykzk\sum_{k=1}^K y_k z_k∑k=1Kykzk是数值稳定的，因为它只是一个线性方程，那么另一项log⁡∑j=1Kexp⁡(zj)\log \sum_{j=1}^K \exp(z_j)log∑j=1Kexp(zj)是否为数值稳定的呢？

实际上log⁡∑j=1Kexp⁡(zj)\log \sum_{j=1}^K \exp(z_j)log∑j=1Kexp(zj)不仅是数值稳定的，而且还很常见，甚至它还有一个名字，叫作LogSumExp(LSE)。

假设x=[x1,x2,⋯,xn]x=[x_1,x_2,\cdots,x_n]x=[x1,x2,⋯,xn]，那么
LSE(x)=log⁡∑i=1nexp⁡(xi)(9)\text{LSE}(x) = \log \sum_{i=1}^n \exp(x_i) \tag 9 LSE(x)=logi=1∑nexp(xi)(9)
我们可以对LSE(x)\text{LSE}(x)LSE(x)使用对Softmax中同样的技巧，即令xxx中所有元素减去同样一个常数bbb。
LSE(x−b)=log⁡∑i=1nexp⁡(xi−b)=log⁡∑i=1nexp⁡(xi)⋅exp⁡(−b)=log⁡exp⁡(−b)∑i=1nexp⁡(xi)=log⁡∑i=1nexp⁡(xi)+log⁡exp⁡(−b)=log⁡∑i=1nexp⁡(xi)−b(10)\begin{aligned} \text{LSE}(x - b) &= \log \sum_{i=1}^n \exp(x_i - b) \\ &= \log \sum_{i=1}^n \exp(x_i) \cdot \exp(-b) \\ &= \log \exp(-b) \sum_{i=1}^n \exp(x_i) \\ &= \log \sum_{i=1}^n \exp(x_i) + \log \exp(-b)\\ &= \log \sum_{i=1}^n \exp(x_i) - b \end{aligned} \tag{10} LSE(x−b)=logi=1∑nexp(xi−b)=logi=1∑nexp(xi)⋅exp(−b)=logexp(−b)i=1∑nexp(xi)=logi=1∑nexp(xi)+logexp(−b)=logi=1∑nexp(xi)−b(10)
这里也令b=max⁡i=1nxib=\max_{i=1}^n x_ib=maxi=1nxi，那么xxx中的最大元素也变成了exp⁡(0)=1\exp(0)=1exp(0)=1，所以说上式是数值稳定的。

根据公式(10)(10)(10)，我们改变公式(9)(9)(9)得到LSE\text{LSE}LSE的最终形式：
LSE(x)=b+LSE(x−b)=b+log⁡∑i=1nexp⁡(xi−b)(11)\text{LSE}(x) = b + \text{LSE}(x-b) = b + \log \sum_{i=1}^n \exp(x_i - b) \tag{11} LSE(x)=b+LSE(x−b)=b+logi=1∑nexp(xi−b)(11)
那么现在我们就基于此((7)(7)(7)和(11)(11)(11))来优化交叉熵损失函数：

def logsumexp(x, axis=-1):b = x.max(axis=axis, keepdims=True)return b + (x - b).exp().sum(axis=axis, keepdims=True).log()def cross_entropy(input: Tensor, target: Tensor, reduction: str = "mean") -> Tensor:''':param input: logits:param target: 真实标签one-hot向量:param reduction::return:'''N = len(target)axis = -1errors = target.sum(axis=axis, keepdims=True) * logsumexp(input, axis=axis) - (target * input).sum(axis=axis, keepdims=True)if reduction == "mean":loss = errors.sum() / Nelif reduction == "sum":loss = errors.sum()else:loss = errorsreturn loss

这里我们要注意维度，这种实现是可以支持批数据的。作为logits的input的维度是N×KN \times KN×K，NNN为批大小，KKK为类别个数。使用keepdims=True避免广播的消耗，广播虽好用，但会引入额外的运算。

负对数似然损失与交叉熵损失的关系

再来看下交叉熵损失的公式：
LCE(y^,y)=−∑k=1Kyklog⁡y^k(12)L_{CE} (\hat y,y) = -\sum_{k=1}^K y_k \log \hat y_k \tag{12} LCE(y^,y)=−k=1∑Kyklogy^k(12)
那负对数似然损失(negative log likelihood loss,NLLLoss)呢？

我们知道在多分类中，会经过Softmax得到概率，有
y^=exp⁡(zi)∑j=1Kexp⁡(zj)1≤j≤K(13)\hat y =\frac{\exp(z_i)}{\sum_{j=1}^K \exp(z_j)} \quad 1 \leq j \leq K \tag{13} y^=∑j=1Kexp(zj)exp(zi)1≤j≤K(13)
这里假设有KKK个类别，即y^\hat yy^向量是一个长度为KKK的向量，向量中每个元素代表属于一个类别的概率。

我们如何表示负对数似然损失呢？

假设某个数据对应的真实类别为ccc，那么该数据的似然就是y^c\hat y_cy^c。

而真实类别一般通过ont-hot向量表示(假设在真实类别yyy中只有第ccc个元素为111，其他元素都为000)。那么该数据的似然可以表示为：
∏j=1Ky^jyj\prod_{j=1}^K \hat y_j ^{y_j} j=1∏Ky^jyj

虽然一个连乘的形式，假设真实类别为ccc，只有yc=1y_c=1yc=1，其他都是000，最终的似然依然是y^c\hat y_cy^c。

那么负对数似然就是先取对数，再加上负号，即：
−∑j=1Kyjlog⁡y^j(14)- \sum_{j=1}^K y_j \log \hat y_j \tag{14} −j=1∑Kyjlogy^j(14)
这里yyy是一个one-hot向量，可以看成是真实标签的概率分布，而y^\hat yy^就是模型预测的概率。

可以看到，多分类下的负对数似然公式就是交叉熵的公式，这又给了我们计算交叉熵损失的另一种方式——通过LogSoftmax。

LogSoftmax

把公式(13)(13)(13)代入(14)(14)(14)得：
−∑j=1Kyjlog⁡y^j=−∑j=1Kyjlog⁡exp⁡(zj)∑k=1Kexp⁡(zk)(15)- \sum_{j=1}^K y_j \log \hat y_j = - \sum_{j=1}^K y_j \log \frac{\exp(z_j)}{\sum_{k=1}^K \exp(z_k)} \tag{15} −j=1∑Kyjlogy^j=−j=1∑Kyjlog∑k=1Kexp(zk)exp(zj)(15)
其中log⁡exp⁡(zj)∑k=1Kexp⁡(zk)\log \frac{\exp(z_j)}{\sum_{k=1}^K \exp(z_k)}log∑k=1Kexp(zk)exp(zj)也有一个名字，称为LogSoftmax，我们上面知道Softmax有数值稳定的形式，所以LogSoftmax也是数值稳定的。
log⁡exp⁡(zj)∑k=1Kexp⁡(zk)=zj−log∑k=1Kexp⁡(zk)(16)\log \frac{\exp(z_j)}{\sum_{k=1}^K \exp(z_k)} = z_j- log\sum_{k=1}^K \exp(z_k) \tag{16} log∑k=1Kexp(zk)exp(zj)=zj−logk=1∑Kexp(zk)(16)
稍进行变换，上式又出现了LogSumExp，所以我们还可以基于LogSumExp来实现。

所以(15)(15)(15)可以写成：
−∑j=1Kyjuj(17)- \sum_{j=1}^K y_j u_j \tag{17} −j=1∑Kyjuj(17)
其中uj=log⁡y^j=logsoftmax(zj)u_j = \log \hat y_j = \text{logsoftmax}(z_j)uj=logy^j=logsoftmax(zj)

公式(17)(17)(17)只是一个线性方程，由于uju_juj是数值稳定的，因此(17)(17)(17)也是数值稳定的。

这样，我们又有另一种数值稳定版的交叉熵损失，其实PyTorch中的CrossEntropyLoss就是基于LogSoftmax和负对数似然损失实现的。

Note that this case is equivalent to the combination of LogSoftmax and NLLLoss.

(CrossEntropyLoss)在这种情况下(应该是当类别为one-hot向量时，不过传入的是类别的索引，而不是ont-hot向量)等同于组合LogSoftmax和NLLLoss。

同时在NLLLoss的文档中也说了：

Obtaining log-probabilities in a neural network is easily achieved by adding a LogSoftmax layer in the last layer of your network. You may use CrossEntropyLoss instead, if you prefer not to add an extra layer.

在一个神经网络中获取对数概率是很容易的，只要增加一个LogSoftmax层到网络的最后即可。如果你不想增加额外的层，你也可以直接CrossEntropyLoss(此时传入logits)。

下面我们就来实现LogSoftmax和NLLLoss：

def logsoftmax(x, axis=-1):return x - logsumexp(x, axis)

首先实现了LogSoftmax，这里的x也是logits。然后我们实现负对数似然损失：

class NLLLoss(_Loss):def __init__(self, reduction: str = "mean") -> None:super().__init__(reduction)def forward(self, input: Tensor, target: Tensor) -> Tensor:''':param input: 概率的对数:param target: 类别one-hot向量:return:'''errors = - target * inputif self.reduction == "mean":loss = errors.sum() / len(input)elif self.reduction == "sum":loss = errors.sum()else:loss = errorsreturn loss

为了验证我们的实现是正确的，我们编写测试代码：

import numpy as np
import torchimport metagrad.functions as F
from metagrad.loss import NLLLoss
from metagrad.tensor import Tensordef test_simple_nll_loss():x = np.array([[0, 1, 2, 3], [4, 0, 2, 1]], np.float32)t = np.array([3, 0]).astype(np.int32)mx = Tensor(x, requires_grad=True)mt = Tensor(np.eye(x.shape[-1])[t])  # 需要转换成one-hot向量tx = torch.tensor(x, dtype=torch.float32, requires_grad=True)tt = torch.tensor(t, dtype=torch.int64)my_loss = NLLLoss()torch_loss = torch.nn.NLLLoss()# 先调用各自的log_softmax转换为对数概率ml = my_loss(F.log_softmax(mx), mt)tl = torch_loss(torch.log_softmax(tx, dim=-1, dtype=torch.float32), tt)assert ml.item() == tl.item()ml.backward()tl.backward()assert np.allclose(mx.grad.data, tx.grad.data)def test_nll_loss():N, CLS_NUM = 100, 10  # 样本数，类别数x = np.random.randn(N, CLS_NUM)t = np.random.randint(0, CLS_NUM, (N,))mx = Tensor(x, requires_grad=True)mt = Tensor(np.eye(x.shape[-1])[t])  # 需要转换成one-hot向量tx = torch.tensor(x, dtype=torch.float32, requires_grad=True)tt = torch.tensor(t, dtype=torch.int64)my_loss = NLLLoss()torch_loss = torch.nn.NLLLoss()# 先调用各自的log_softmax转换为对数概率ml = my_loss(F.log_softmax(mx), mt)tl = torch_loss(torch.log_softmax(tx, dim=-1, dtype=torch.float32), tt)assert ml.item() == tl.item()ml.backward()tl.backward()assert np.allclose(mx.grad.data, tx.grad.data)

============================= test session starts ==============================
collecting ... collected 2 itemstest_nll_loss.py::test_simple_nll_loss PASSED                            [ 50%]
test_nll_loss.py::test_nll_loss PASSED                                   [100%]============================== 2 passed in 0.78s ===============================

最后我们基于LogSoftmax实现Softmax回归。

基于LogSoftmax实现Softmax回归

import matplotlib.pyplot as plt
import numpy as np
from sklearn import datasets
from sklearn.model_selection import train_test_split
from tqdm import tqdmfrom metagrad.loss import NLLLoss
from metagrad.module import Module, Linear
from metagrad.optim import SGD
from metagrad.tensor import Tensor
import metagrad.functions as Fclass SoftmaxRegression(Module):def __init__(self, input_dim, output_dim):self.linear = Linear(input_dim, output_dim)def forward(self, x: Tensor) -> Tensor:# 输出概率的对数return F.log_softmax(self.linear(x))def generate_dataset(draw_picture=False):iris = datasets.load_iris()X = iris['data']y = iris['target']names = iris['target_names']  # 类名feature_names = iris['feature_names']  # 特征名if draw_picture:x_min, x_max = X[:, 0].min() - 0.5, X[:, 0].max() + 0.5y_min, y_max = X[:, 1].min() - 0.5, X[:, 1].max() + 0.5plt.figure(2, figsize=(8, 6))plt.clf()for target, target_name in enumerate(names):X_plot = X[y == target]plt.plot(X_plot[:, 0], X_plot[:, 1],linestyle='none',marker='o',label=target_name)plt.xlabel(feature_names[0])plt.ylabel(feature_names[1])plt.xlim(x_min, x_max)plt.ylim(y_min, y_max)plt.axis('equal')plt.legend()fig = plt.gcf()fig.savefig('iris.png', dpi=100)y = np.eye(3)[y]X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=2)return Tensor(X_train), Tensor(X_test), Tensor(y_train), Tensor(y_test)if __name__ == '__main__':X_train, X_test, y_train, y_test = generate_dataset(True)epochs = 2000model = SoftmaxRegression(4, 3)  # 4个特征 3个输出optimizer = SGD(model.parameters(), lr=1e-1)# 采用负对数似然损失loss = NLLLoss()losses = []for epoch in range(int(epochs)):outputs = model(X_train)l = loss(outputs, y_train)optimizer.zero_grad()l.backward()optimizer.step()if (epoch + 1) % 20 == 0:losses.append(l.item())print(f"Train -  Loss: {l.item()}")# 在测试集上测试outputs = model(X_test)correct = np.sum(outputs.numpy().argmax(-1) == y_test.numpy().argmax(-1))accuracy = 100 * correct / len(y_test)print(f"Test Accuracy:{accuracy}")

Train -  Loss: 0.8603581190109253
Train -  Loss: 0.7593004107475281
Train -  Loss: 0.6828514337539673
Train -  Loss: 0.6106860637664795
Train -  Loss: 0.5394479036331177
...
Train -  Loss: 0.10196995735168457
Train -  Loss: 0.10157666355371475
Train -  Loss: 0.10119038075208664
Train -  Loss: 0.1008109301328659
Train -  Loss: 0.10043811798095703
Train -  Loss: 0.10007175803184509
Train -  Loss: 0.09971167147159576
Train -  Loss: 0.0993577092885971
Train -  Loss: 0.0990096777677536
Test Accuracy:100.0

总结

本文我们了解了Softmax中的数值稳定问题，知道了PyTorch中CrossEntropyLoss的底层实现逻辑。但是可能对似然估计、交叉熵损失不太理解，下篇文章打算写一篇长文来探讨这些概念。

References

一文弄懂LogSumExp技巧

从零实现深度学习框架——Softmax回归中的数值稳定相关推荐

深度学习基础--SOFTMAX回归（单层神经网络）
深度学习基础–SOFTMAX回归(单层神经网络) 最近在阅读一本书籍–Dive-into-DL-Pytorch(动手学深度学习),链接:https://github.com/newmonkey/Div ...
python学习框架图-从零搭建深度学习框架（二）用Python实现计算图和自动微分
我们在上一篇文章<从零搭建深度学习框架(一)用NumPy实现GAN>中用Python+NumPy实现了一个简单的GAN模型,并大致设想了一下深度学习框架需要实现的主要功能.其中,不确定性最 ...
深度学习：Softmax回归
在前面,我们介绍了线性回归模型的原理及实现.线性回归适合于预测连续值,而对于分类问题的离散值则束手无策.因此引出了本文所要介绍的softmax回归模型,该模型是针对多分类问题所提出的.下面我们将从so ...
【动手学深度学习】Softmax 回归 + 损失函数 + 图片分类数据集
学习资料: 09 Softmax 回归 + 损失函数 + 图片分类数据集[动手学深度学习v2]_哔哩哔哩_bilibili torchvision.transforms.ToTensor详解 | 使用 ...
从零实现深度学习框架——深入浅出Word2vec(下)
引言本着"凡我不能创造的,我就不能理解"的思想,本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架,该框架类似PyTorch能实现自动求导. 要深入理解深度学 ...
从零实现深度学习框架——GloVe从理论到实战
引言本着"凡我不能创造的,我就不能理解"的思想,本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架,该框架类似PyTorch能实现自动求导.
从零实现深度学习框架——Seq2Seq从理论到实战【实战】
引言本着"凡我不能创造的,我就不能理解"的思想,本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架,该框架类似PyTorch能实现自动求导.
从零实现深度学习框架——RNN从理论到实战【理论】
引言本着"凡我不能创造的,我就不能理解"的思想,本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架,该框架类似PyTorch能实现自动求导.
从零实现深度学习框架——从共现矩阵到点互信息
引言本着"凡我不能创造的,我就不能理解"的思想,本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架,该框架类似PyTorch能实现自动求导.
从零实现深度学习框架——LSTM从理论到实战【理论】
引言本着"凡我不能创造的,我就不能理解"的思想,本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架,该框架类似PyTorch能实现自动求导.

从零实现深度学习框架——Softmax回归中的数值稳定

引言