LLM学习笔记-48-关于投机解码-DFlash & DSpark(之六)- 训练部分

本篇笔记记录DSPARK的训练相关内容。

DSPARK的训练目标由三部分组成：是三项加权和（骨干与两个头联合训练）：

下面是论文中的原文描述，其实已经比较清楚：

训练时，我们从每个目标序列随机采样多个锚点位置，构成 -token 块作为训练数据。目标模型在整个训练过程中冻结；草稿模型共享其嵌入层和语言模型头并保持它们冻结，只更新主干草稿器、顺序块和置信度头。

训练目标由三项组成：交叉熵损失 、分布匹配损失 、置信度损失 。三者都按位置权重  加权（Chen et al., 2026），该权重强调块内更靠前的位置——在基于前缀的验证下，它们对期望接受长度贡献更大。1）交叉熵损失  训练草稿器预测正确的下一个 token：

其中  是真实 token， 是草稿分布。

2）分布匹配损失  惩罚草稿与目标分布之间的总变差距离：

由于总变差距离是接受率的直接代理——逐步接受概率等于 （Leviathan et al., 2023）——最小化  直接最大化期望接受率。

3）置信度损失  是一个二元交叉熵，训练置信度头去预测式 (8) 中的软接受标签 ：

总体目标是三项的加权组合（默认权重 ，，）：

1.2 补充理解

我们这里补充一些理解即可：

🟠（一）具体计算过程的理解：我们计算loss的输入需要： ， 和 真实 token  ，置信度头输出 

其中， 和  的含义：它们就是草稿模型（Draft Model）和目标模型（Target Model）在一次前向传播（Forward Pass）后，经过 Softmax 层输出的概率分布向量。

一旦我们同时拿到了  和 （两个形状为  的概率向量），以及真实的 token 标签 ，就可以直接套用公式了：

(1) 交叉熵损失 这个最简单，只需要从草稿分布  中，把真实 token  对应的那个概率值挑出来，取负对数即可：

(2) 分布匹配损失 将两个概率向量  和  的对应元素相减，取绝对值，求和再乘 （或者直接求  范数）。这就是在计算这两个向量在几何空间上的距离：

(3) 置信度损失 首先，利用上面算出来的  和  计算出软标签 ：

然后，将  作为目标值（Target），去计算置信度头输出  的二元交叉熵：

而训练的过程就是：拿着这个 ，去和老师模型给出的  比较（算 TV 损失），同时拿着  去查真实答案（算 CE 损失），并利用  和  的差异算出一个软标签 ，来训练能输出标量  的置信度头。

🟠（二）对交叉熵损失 L_ce 计算的理解：

在信息论中，交叉熵衡量的是使用预测分布  来编码真实分布  所需要的平均比特数。对于离散分类问题（如词表大小为  的下一个 token 预测）：真实分布  通常是 One-hot 向量（只在真实 token  处为 1，其余为 0）。交叉熵公式为：

由于  是 One-hot 的，只有当  时 ，其余项全为 0。因此，公式退化为：

物理意义：交叉熵越小，说明模型给真实 token 分配的概率  越接近 1，模型预测得越准。

🟠（三）对位置权重w_k 的理解

所有损失在计算时都引入了位置权重 。由于投机采样采用基于前缀的验证（一旦某个位置被拒绝，后面的草稿全部作废），越靠前的位置对期望接受长度的贡献越大，因此  随着位置  的增加呈指数衰减，赋予前面的 token 更高的权重。