在本节内容中,我们将实现损失函数和梯度截断

实现之前,我们先进行简单的介绍,损失函数自不必多说,它是为我们模型提供优化方向的重要指标,梯度截断是我们稳定训练的基础

cross entropy loss

首先就是为什么我们会选择交叉熵而不是其它损失函数,例如MSE或者MAE。我们可以从两个角度去得出这个结论。

KL散度

两个分布的KL散度定义为: $$ KL(P||Q)=H(P,Q)-H(P) $$ KL散度衡量了两个分布的距离或者是差异程度,当Q和P的分布完全相同时,有$H(P,Q)=H(P)$,此时KL散度为0。

我们假定我们有数据分布$D_{data}$,模型的预测分布为$D_{model}$,我们可将上式写为 $$ KL(D_{data}||D_{model})=H(D_{data},D_{model})-H(D_{data}) $$ 所以我们希望$KL(D_{data}||D_{model})$尽可能小,使得我们的模型能尽可能拟合数据分布(值得注意的是,这里的是数据分布而不是真实的语言分布,因为我们无法得知真实的语言分布,所以只能通过尽可能多的数据来接近)

其中我们的可优化项只有$H(D_{data},D_{model})$,因为$H(D_{data})$是一个定值,所以只能优化交叉熵

写成我们训练常用的形式为: $$ minH(D_{data},D_{model})=minE_{p \sim D_{data}}(-log(fn(y|x))) $$ 其中$fn(y|x)$表示的即是模型分布的取样结果,x可以理解为提示词或者为前面的token序列,y为预测得到的token。

真实训练中我们使用的算法: $$ loss(y|x)=\frac {1} {batchsize}\sum^{batchsize}\frac {1} {seqlen}\sum_{t=1}^{seqlen}-plog(fn(y_{t}|x_{:t-1})) $$ 其中由于我们的$p \sim D_{data}$,所以事实上这里的概率分布向量为一个独热向量,只有label上唯一的下一个token的p为1,所以上式可写为 $$ loss(y|x)=\frac {1} {batchsize}\sum^{batchsize}\frac {1} {seqlen}\sum_{t=1}^{seqlen}-log(fn(y_{t}|x_{:t-1})) $$ 模型的输出后往往需要使用softmax等归一化函数得到logits,所以我们往往会习惯性认为需要先对输出进行softmax,再对其取负对数得到损失,但这种方法往往会导致softmax上溢的问题,当某个数值非常大时,我们使用softmax取指数时数值容易爆炸导致数值上溢,这里有一个常用的trick,我们可以通过 $$ loss=-log(softmax(O))=-log\frac {e^{o_{label}}} {\sum_{j=1}^{vocabsize} e^{o_j}}=-O_{label}+log({\sum_{j=1}^{vocabsize} e^{o_j}}) $$

通过这个技巧,我们巧妙地避免了计算每个数值的指数,从而避免了softmax上溢的问题