节前发版：Deepseek v3.2 exp

加班快乐...

论文原文
推理代码

架构

与Deepseek-V3.1相比，新一般的架构更改仅仅在后续训练中引入了新的稀疏注意力机制DSA。

DSA：deepseek稀疏注意力

主要包括两个部分：一个ligtning indexer（索引器）和一个细粒度的token选择机制。

Lightning indexer

Step 1: 计算索引分数。

计算了 当前询问 Q token \(h_t\in \mathbb{R}^d\) 与一个 前序token \(h_s\in\mathbb{R}^d\) 的索引分数，决定了Qtoken将会选择哪一个token。

\[I_{t,s}=\sum\limits_{j=1}^{H^I}w_{t,j}^I\cdot\text{ReLU}\left(q_{t,j}^I\cdot k_s^I\right) \]

其中我们有：

• \(H^I\) 索引头的数目。
• \(q_{t,j}^I\in\mathbb{R}^{d^I}\) 和 \(w_{t,j}^I\in\mathbb{R}\) 从Q token \(h_t\) 中导出。
• \(k_s^I\in R^{d^I}\) 从前序的 \(h_s\) 中导出。

作者选择了ReLU来提升吞吐率。即使lightning indexer仅有很少数量的头并且可以在FP8上部署，其计算效率也是非常显著的。

Step 2: 选择前k个索引分数最高的 \(c_s\), 计算注意力输出。

给定了索引分数 \(I_{t,s}\)，我们的细粒度token索引机制将会仅仅取出那些具有前k个索引分数的token。随后，注意力输出 \(o_t\) 将会在当前Q token \(h_t\) 和稀疏化选出的 \(c_s\) 中进行。

\(c_s\) 其实是MLA中低秩投影计算出来的向量，用于减少KVCache的存储开销，提高推理效率。

\[u_t = \text{Attn}(h_t,\{c_s\,|\, I_{t,s}\in\text{Top-k}(I_{t,:})\}) \]

下面是新旧结构的对比。上图为新的结构。下图为曾经的旧结构。

在MLA下实例化DSA

为了考虑从v3.1继续训练，需要基于MLA上实例化DSA。在kernel层面，每一个KV项都需要在多个查询之间共享，提升计算效率。因此，我们在MLA的MQA模式上部署了DSA。这样每一个潜在层(latent vector)将会在每个头之间共享（多个头共用一个潜在向量 \(c_i\)， 也就是多个头——多个Query， 共用一个KV）。

训练

从v3.1-Terminus 后继续训练，上下文长度扩展到128K。

Step 1: 稠密 warm-up 阶段

用于初始化lightning indexer。继续保持稠密注意力机制，其余参数全部冻结，仅剩下lightning indexer进行训练。

为了保持indexer输出与原先的主要注意力分布对齐，对于第t个查询token，我们首先将多个头的主要注意力分数进行相加，然后在序列维度上进行 L1-正则化，生成目标分布 \(p_{t,:}\in\mathbb{R}^t\). 基于 \(p_{t,:}\), 我们设置一个 KL-散度 loss作为我们训练indexer的优化目标。

\[\mathcal{L}^I = \sum\limits_{t}\mathbb{D}_{KL}\left(p_{t,:}||\text{Softmax}(I_{t,:})\right) \]

作者声称采用了 \(10^{-3}\) 的学习率训练了1000步。每一步具有128K长度的16个序列，总共2.1B个token。

Step 2: 稀疏训练阶段

在进行稠密训练之后，进入到了细粒度的token选择，并以此来优化整体模型的参数，来获得DSA的稀疏模式。在这一阶段，我们不在选择所有的token，而是通过上文的方式选择通过indexer判断出来的，索引分数最大的K个token：
\(\mathcal{S}_t=\{s\,|\,I_{t,s}\in\text{Top-k}(I_{t,:})\}\)

\[\mathcal{L}^I = \sum\limits_{t}\mathbb{D}_{KL}\left(p_{t,S_t}||\text{Softmax}(I_{t,S_t})\right) \]

需要值得注意的是我们将indexer的输入从计算图中分离，也就是分开indexer和DSA的其他部份，分别进行优化。

• indexer仍然仅仅根据 \(\mathcal{L}^I\) 进行优化。
• 其他部分通过模型其他部分的loss进行优化。

稀疏训练采用学习率 \(7.3\times 10^{-6}\)，每个query选择2048个KV token。训练15000步，具有480个长度为128K的token，总共是943.7B token数量。

Step 3: 后训练

后训练与先前deepseek-v3的后训练类似，主要有两步：

1. 专家知识蒸馏。
2. 混合RL训练。

专家知识蒸馏

• 对于每个任务我们都训练了一个专门的针对这个领域知识的模型，这些模型都是从相同的预训练v3.2基座模型的ckpt而来。
• 针对写作任务和通用问答任务，我们划分了5个领域：数学，竞赛类编程，通用因果逻辑，多智能体编码，多智能体搜索。
• 对于每个专家，我们都通过大规模强化学习方式进行训练。
• 并且，我们部署了不同的模型来生成针对思维链(CoT)的训练数据，以及直接回答(非思维链模式)的训练数据。
• 当专家模型完成后，他们将被用于为最后的ckpt生成领域专用的知识。最终ckpt在各个领域与专家模型的差距将通过后续的强化学习来进行弥补。

混合强化学习

• 与v3.1相同，仍然采用的是GRPO强化学习方式。
• 与前面分不同阶段强化学习不同的是，作者将多个阶段的RL学习(因果，智能体，人类对齐训练)混合到了一起。
• 优势是可以讲多个领域的表现有效进行平衡并且设法克服在多阶段训练中造成的灾难性遗忘问题。
• 对于因果和智能体任务，我们部署了基于规则的结果奖励，长度惩罚以及语言一致性奖励。
• 对于生成式任务，我们部署了一个生成式奖励模型，将按照自己的规则进行评估。
• reward进行了两方面的权衡：(1) 长度vs准确度。(2)一致性vs准确度。

评估结果

• 推理开销从原先的 \(O(L^2)\) (原先需要计算所有的 token，长度为 L) 变成 \(O(Lk)\) (Q token长度不变，但是KV低秩投影token通过lightning indexer选择K个)。对于lightning indexer，其计算复杂度仍然为 \(O(L^2)\)，但是因为其具有的头数量比原先的MLA头数量少，因此常数因子的减少也显著提升了其计算效率。