DeepSeek V4架构猜想：从NSA到DSA的演进与基于层间相似度的Indexer优化思路

最近在思考从 NSA（Nested Sparse Attention）到 DSA（Dynamic Sparse Attention）的技术发展脉络。这两种都是大模型推理中用于优化注意力计算（特别是长上下文场景）的重要方法。

NSA的局限

NSA 本质上是 Block Selection（块选择）与 SWA（Sliding Window Attention，滑动窗口注意力）的混合体。当 NSA 与 MLA（Multi-Query Latent Attention）结合时，会遇到一个核心问题：MLA 本身的训练与推理不一致。在训练或 Prefill（预填充）阶段，MLA 使用标准的 MHA（Multi-Head Attention，多头注意力），而在 Decode（解码）阶段则切换为 MQA（Multi-Query Attention，多查询注意力）。如果在 Prefill 阶段采用每个注意力头独立选择（per head selection）的策略，那么在 Decode 阶段，被选中的 Key-Value 缓存就需要被放大最多 128 倍，这带来了额外的开销和复杂性。

DSA的潜力与挑战

DSA 采用了不同的思路。它将 Query、Key、Value 映射到一个低维空间进行检索（这个检索模块称为 Indexer）。虽然检索本身是平方级复杂度，但由于映射后的维度很低，整体计算量得到了压缩。与此同时，核心的 Softmax 和 FA（全注意力计算）被限定在每次只处理 2048 个 token，因此 FA 部分可以保持恒定的计算耗时。

DSA 的理念很美好，但它自身也面临挑战：Indexer 模块本身是平方复杂度。随着序列长度的持续增加，Indexer 的计算耗时最终会超过 FA 部分的耗时。而且，Indexer 涉及排序、TopK 选取以及离散的内存访问，这些操作的效率在硬件上难以优化。因此，下一步的优化重点很可能就落在 Indexer 的效率提升上。

一个优化思路：利用层间相似度

之前基于 DeepSeek-V3.2 的研究发现了一个有趣的现象：Transformer 相邻层的 hidden states 所选择的重点 token 相似度很高。那么，这个特性该如何利用呢？

统计意义上的预测行为在在线推理系统中，通常只能用于提升吞吐量，很难直接降低延迟。因为现象本身缺乏规范性——你无法直接、确定性地指出哪些 token 一定会被重复选中。

那么，我们是否有可能基于前一层（第 l 层）的 Indexer 排序结果，通过适当扩充召回 token 的数量，来为下一层（第 l+1 层）提供候选呢？这是一种在 Precision（精确率）和 Recall（召回率）之间优先保障 Recall 的策略。

假设对于某一层 l，其对应的 hidden state 为 x_l，经过 Indexer 得到的排序结果为：K = [k0， k1, ..., ks]，其中 s 是序列长度（当 s ≤ 2048 时默认使用全量 FA 计算）。

具体的算法流程可以设计如下：

第一步：对于第 l 层，选择排序结果 K 中的前 2048 个（top{k=2048}(K)）用于计算 FA。

第二步：对于第 l+1 层，计算 score = cos(x_l， x_(l+1))，即两层 hidden state 的余弦相似度。

1. 当 score ≥ threshold（某个阈值）时，直接复用第 l 层的排序结果，但在此基础上多补偿 128 个 token，即使用 top{k=2048+128}(K)。这里的 128 是一个可调的超参数，可以基于实际效果进行设定。（有研究，例如小米的 SWA 论文，指出 128 这个补偿值效果更好？）
2. 当 score < threshold 时，则重新为第 l+1 层计算 Indexer，并以此结果为新的起点，重复第一步和第二步的流程。

上面的算法还可以进一步强化：强制连续的 N 个 layer 共享同一个 Indexer，并走上述的补偿逻辑。在模型训练时，可以额外通过损失函数来约束这 N 个 layer 的 hidden state 保持较高的余弦相似度。这是一种更强的假设和优化目标。

此思路可能带来的好处

1. 性能有保障：最差情况会退化为朴素的 DSA（每层独立计算 Indexer）；最好情况是所有层共享一个 Indexer，性能介于二者之间。也可以通过训练损失函数的约束，达到一个理想的均衡点。
2. 符合模型特性：利用 Transformer 层间的内在相似度来实现动态的 Indexer 计算，这符合 mHC（Model Hamiltonian Continuum）等研究路线。尤其是当模型层数不断加深后，相邻层的相似度通常会很高。
3. 有利于硬件优化：由于存在跨层的 Indexer 复用（尽管有补偿），预测性变得更强。系统可以预先根据 Indexer 结果对 Key-Value 缓存进行重排，从而实现更高效的连续内存访问，这对 GPU 等硬件至关重要。
4. 训练开销可控：在模型训练阶段，不同的 Pipeline Parallelism（流水线并行）分组之间，只需要额外传递一个 hidden state 和 Indexer 的 topk 结果，通信开销增加不大。

对于 AIGC 和 大模型 领域的开发者而言，深入理解并参与这类底层 Transformer 注意力机制的优化，是提升模型实际推理效率的关键。欢迎在云栈社区交流更多关于模型加速与架构演进的想法。

来源：https://zhuanlan.zhihu.com/p/1996718177191757819 （已获作者授权发布）