DeepSeek推出原生稀疏注意力机制NSA,重塑长文本建模
发布时间:2025-02-19AI寒武纪2月18日消息,DeepSeek研究团队推出原生稀疏注意力(Native Sparse Attention,NSA)机制,旨在实现高效的长上下文建模。
长上下文建模对于下一代语言模型至关重要,但传统注意力机制的计算复杂度随着序列长度的增加而呈平方级增长,成为制约模型发展的关键瓶颈。同时,现有稀疏注意力方法存在推理效率的“假象”,仅在自回归解码时稀疏,预填充阶段仍需大量计算,且与现代高效解码架构不兼容;在训练阶段存在性能退化、训练效率低、有不可训练组件、反向传播效率低下等问题。
NSA 机制融合算法创新与硬件优化,采用动态分层稀疏策略,结合粗粒度的token压缩与细粒度的token选择,保证了模型对全局上下文的感知和局部信息的精准性。通过算术强度平衡的算法设计,针对现代硬件进行实现优化,显著提升了计算速度。同时支持端到端训练,不仅在推理阶段高效,还能减少预训练的计算量,且不牺牲模型性能。
NSA 架构采用分层 Token 建模,通过压缩注意力、选择注意力和滑动窗口注意力三个分支并行处理,实现训练推理一体化。同时设计了硬件优化的 Kernel以实现效率最大化。
实验结果显示,NSA在性能超越传统Full Attention模型的同时大幅提升长文本处理效率,在处理64长度的序列时,速度最高提升11.6倍,验证了其在模型生命周期各个阶段的效率优势。
NSA 的硬件友好设计和训推一体化特性,有望加速下一代 LLM 在长文本处理领域的应用落地。
论文链接:https://arxiv.org/abs/2502.11089
参考链接:https://mp.weixin.qq.com/s/jpS3pr3fb1hslIUf7yAHeQ
来源:AI寒武纪