1. DSA原理

 DSA(Deepseek Sparse Attention) 是一种稀疏Attention技术，主要用来节省Attention的计算量，用TopK个元素的Attention来代替原生的Attention, 原理有点像MLP部分的MoE。

DSA关键机制是新增索引器（Lighting  Indexer) 模块，如图中绿色部分。用来为每个Token选出TopK个最想着的Token（例如128K个上下文中选出2048个）。后续的Attention计算只需要在这TopK上进行，计算复杂度从降低到。

2.DSA训练过程

在正常的数据训练过程中，新增密集预热阶段（Dense Warm-up Stage)，用于Lighting Indexer的训练，核心思想是先模仿再适应。

第一阶段：密集预热（Dense Warm-up Stage）

 这个阶段的目标是让全新的闪电索引器学会"如何挑选重要Token"。具体做法是让它在密集注意力（即传统的全量注意力）的环境下学习。

• 训练对象：只训练闪电索引器（Lightning Indexer），冻结模型中其他所有参数。这意味着索引器是作为一个独立、可插拔的组件被先训练出来的。

• 学习目标：索引器会接收当前查询（Query），并为上下文中所有的历史Token快速打分，选出一部分。它的学习目标是让自己的"挑选结果"和原本密集注意力机制所关注的重点尽可能一致。技术报告提到，这个阶段使用了约21亿（2.1B）个Token的数据进行训练。

• 实现方式：索引器在训练时被设计得极其轻量，通常使用FP8低精度和少量的注意力头（Attention Heads）来运行，这样即使它为全部Token打分，总的计算开销也能被控制住。

第二阶段：稀疏训练 (Sparse Training Stage)

在索引器"毕业"后，就进入第二阶段，让整个模型适应新的稀疏模式。这个阶段就是常规训练阶段。

• 训练对象：优化所有模型参数，让整个模型（包括索引器和主模型）都学会在稀疏注意力的模式下工作。

• 关键机制：现在，索引器会直接为每个Token选出最相关的 Top-k 个历史Token计算注意力。

• 分离优化：为了训练的稳定性，索引器的输入会从主计算图中分离出来。这意味着索引器和主模型会各自通过不同的损失函数（Loss）进行优化：索引器优化其挑选的"得分"，主模型则优化最终的语言建模损失。

• 训练方法：同时训练闪电索引器（Lightning Indexer） 和主模型（Main Model） 的全部参数。索引器与主模型的优化是分离的。具体做法是将索引器的输入从主计算图中剥离（detach），这让二者可以专注于各自的目标：

• 索引器（Indexer）：只通过索引器损失（Indexer Loss） 进行优化。这个损失函数（通常为KL散度）会持续监督它，确保在稀疏模式下选出的Top-k个Token仍然是最重要的那些。

• 主模型（Main Model）：仅通过标准的语言建模损失（Language Modeling Loss） 进行优化，学习如何高效地利用索引器选出的那部分关键信息来处理长文本。这部分和常规训练方法差不多。

第一阶段和第二阶段训练索引器方法对比

第二阶段索引器损失 LILI 的计算公式可以表示为：

公式中各部分含义如下：

• t：代表当前的查询（Query）Token位置。

• St：代表索引器为当前查询选出的Top-k个最重要Token的索引集合。

• I_t,St​：代表索引器对属于集合 St 内Token计算出的原始分数（Logits）。

• Softmax(I_t,St)：将索引器在选中的 k 个Token上的分数转换为一个概率分布。

• p_t,St​​：代表主模型（Main Model）在同样的 St 这 k 个Token上的注意力分布。这个分布通常通过将主模型的注意力分数在注意力头上求和，再进行L1归一化得到。它充当了索引器学习的“目标”或“标签”。

• D_KL：代表KL散度（Kullback-Leibler Divergence），用于衡量索引器分布与主模型目标分布之间的差异。通过最小化这个损失，索引器被训练去“模仿”主模型在它自己选出的Token上的关注模式。