DeepSeek 新注意力架构NSA概要

研究背景：

实现高效长上下文建模的自然方法是利用 softmax 注意力的固有稀疏性，通过选择性计算关键 query-key 对，可以显著减少计算开销，同时保持性能。最近这一路线的进展包括多种策略：KV 缓存淘汰方法、块状 KV 缓存选择方法以及基于采样、聚类或哈希的选择方法。尽管这些策略前景广阔，现有的稀疏注意力方法在实际部署中往往表现不佳。许多方法未能实现与其理论增益相媲美的加速；此外，大多数方法主要关注推理阶段，缺乏有效的训练时支持以充分利用注意力的稀疏模式。

挑战：

部署有效的稀疏注意力必须应对两个关键挑战：
1、硬件对齐的推理加速：将理论计算减少转化为实际速度提升，需要在预填充和解码阶段设计硬件友好的算法，以缓解内存访问和硬件调度瓶颈；
2、训练感知的算法设计：通过可训练的操作符实现端到端计算，以降低训练成本，同时保持模型性能。

关键技术

为了实现更有效和高效的稀疏注意力，DeepSeek 研究人员提出了一种原生可训练的稀疏注意力架构 NSA，它集成了分层 token 建模。
NSA引入了两个核心创新以对应于上述关键需求：
1、硬件对齐的系统：优化块状稀疏注意力以利用 Tensor Core 和内存访问，确保算术强度平衡；
2、训练感知的设计：通过高效算法和反向操作符实现稳定的端到端训练。这一优化使 NSA 能够支持高效部署和端到端训练。

如上图所示，NSA 通过将键和值组织成时间块（temporal blocks）并通过三条注意力路径处理它们来减少每查询计算量：压缩的粗粒度 token、选择性保留的细粒度 token 以及用于局部上下文信息的滑动窗口。随后，作者实现了专门的核以最大化其实际效率。

实验测试结果

1. 一般基准性能：尽管 NSA 可能无法充分利用其在较短序列上的效率优势，但它依然表现出了强劲的性能。值得注意的是，NSA 在推理相关基准测试中表现出了显著的提升（DROP：+0.042，GSM8K：+0.034），这表明 DeepSeek 的预训练有助于模型发展出专门的注意力机制。通过过滤掉不相关的注意力路径中的噪音，这种稀疏注意力预训练机制可迫使模型专注于最重要的信息，有可能提高性能。
2. 长上下文基准性能： NSA 具有处理各种长上下文任务中的能力，其原生预训练的稀疏注意力在学习任务最优模式方面提供了额外的助益。
3. 思维链推理性能评估：在 8k 上下文设置下，NSA-R 的准确性显著高于 Full Attention-R（+0.075），这一优势在 16k 上下文设置下仍然保持（+0.054）。

总结：

原生稀疏注意力的两个关键优势：
（1）预训练的稀疏注意力模式能够高效捕捉对复杂数学推导至关重要的长距离逻辑依赖关系；
（2）该架构的硬件对齐设计保持了足够的上下文密度，以支持不断增长的推理深度，而不会出现灾难性遗忘。在不同上下文长度下的一致优势证实了稀疏注意力在原生集成到训练流程中时，对于高级推理任务的可行性。

参考链接：
https://arxiv.org/abs/2502.11089
https://www.jiqizhixin.com/articles/2025-02-18-11