原生稀疏注意力：Deepseek又一重磅开源！实现高效的大模型超长文本建模！

最新推荐文章于 2025-07-27 20:28:43 发布

Sherlock Ma

最新推荐文章于 2025-07-27 20:28:43 发布

阅读量1.3k

点赞数 12

CC 4.0 BY-SA版权

分类专栏： AIGC 大语言模型混合专家模型文章标签：人工智能 nlp AIGC gpt agi chatgpt pytorch

本文链接：https://blog.youkuaiyun.com/sherlockMa/article/details/145762389

1.摘要

长上下文建模是下一代语言模型的关键，然而标准注意机制的高计算代价给计算带来了巨大挑战.稀疏注意为在保持模型性能的同时提高效率提供了一个有前景的方向。因此作者提出了NSA，一种可本地训练的稀疏注意机制，它将算法创新与硬件优化相结合，以实现高效的长上下文建模。

NSA采用动态分层稀疏策略，将粗粒度令牌压缩与细粒度令牌选择相结合，以保持全局上下文感知和局部精度。

如图1所示，实验表明，经过NSA预训练的模型在一般基准测试、长上下文任务和基于指令的推理中保持或超过了全注意模型。同时，NSA在解码、前向传播和后向传播过程中的速度比Full Attention有了显著的提高，验证了其在整个模型生命周期中的有效性。

2.简介

研究界越来越多地认识到，长上下文建模是下一代大型语言模型的关键能力，它的应用范围从深度推理到多种现实应用的驱动、存储库级代码生成和多轮自主代理系统。

最近的突破使模型能够处理整个代码库、冗长的文档，在数千个令牌上保持连贯的多回合对话，并在长范围的依赖关系上执行复杂的推理。然而，高度复杂性的注意机制随着序列长度的增加而成为关键的延迟瓶颈。理论估计表明，在解码64 k长度的上下文时，softmax架构的注意力计算占总延迟的70-80%，强调了对更高效的注意力机制的迫切需要。

一种有效的长上下文建模方法是利用softmax注意力的固有稀疏性，其中选择性地计算关键Query-Key对可以在保持性能的同时显著地减少计算开销。尽管这些策略很有前途，但现有的稀疏注意方法在实际部署中往往存在不足。许多方法都没有达到理论上的加速比;而且，大多数方法主要集中在推理阶段，缺乏有效的训练时间支持来充分利用注意力的稀疏性模式。

为了解决这些限制，有效的稀疏注意的部署必须解决两个关键挑战：

硬件对齐的推理加速：将理论计算减少转换为实际速度提高需要在预填充和解码阶段期间的硬件友好算法设计，以减轻存储器访问和硬件调度瓶颈;
训练感知算法设计：使用可训练的运算符实现端到端计算，在保持模型性能的同时降低训练成本。这些需求对于真实世界的应用程序实现快速的长上下文推理或训练是至关重要的。

为了实现更有效和高效的稀疏注意力，作者提出了NSA，一个原生可训练的稀疏注意力架构，集成了分层令牌建模。如图2所示，NSA通过将键和值组织到时间块中，并通过三种注意路径（压缩的粗粒度标记、选择性保留的细粒度标记和本地上下文信息的滑动窗口）处理它们，减少了每次查询的计算量。然后，我们实现了专门的内核，以最大限度地提高其实际效率。

作者通过对真实世界语言语料库的综合实验来评估NSA。在具有260 B令牌的27 B参数Transformer骨干上进行预训练，作者评估了NSA在一般语言，长上下文评估和思维链推理评估中的表现。实验结果表明，NSA实现了相当或更好的性能完全注意基线，同时优于现有的稀疏注意方法。此外，与Full Attention相比，NSA在解码、前向和后向阶段提供了显著的加速比，对于较长的序列，加速比会增加。这些结果验证了分层稀疏注意设计有效地平衡了模型能力和计算效率。

3.重新思考稀疏注意力方法

现代稀疏注意力方法在降低Transformer模型的理论计算复杂度方面取得了重大进展。然而，大多数方法在推理过程中主要应用稀疏性，同时保留预训练的Full Attention骨干，这可能会引入架构偏差，限制其充分利用稀疏注意力优势的能力。在介绍原生稀疏架构之前，我们通过两个关键镜头系统地分析了这些限制。

有效推理的错觉

尽管在注意力计算中实现了稀疏性，但许多方法未能实现推理延迟的相应减少，主要是由于两个挑战：

相位限制稀疏度（Phase-Restricted Sparsity）。例如H2O的方法在自回归解码期间应用稀疏性，而在预填充期间需要计算密集的预处理（例如，注意力图计算、索引构建）。相比之下，像MInference仅关注预填充稀疏性。这些方法无法在所有推理阶段实现加速，因为至少有一个阶段的计算成本与全注意相当。
与高级注意力架构不兼容。一些稀疏注意力方法无法适应现代解码高效架构，如多查询注意力（MQA）和分组查询注意力（GQA），这通过跨多个查询头共享KV来显著减少解码期间的存储器访问瓶颈。尽管这些方法可以减少计算操作，但所需的KV高速缓存存储器访问仍然相对较高。虽然一些稀疏注意方法减少了计算量，但它们分散的内存访问模式与高级架构中高效的内存访问设计相冲突。

可训练稀疏性

在本2节中，作者探讨了现有稀疏注意力方法在训练阶段的局限性，指出这些方法大多专注于推理阶段的稀疏性优化，而在训练阶段的支持不足，导致无法充分发挥稀疏注意力的潜力。作者指出，将稀疏性应用于预训练好的全注意力模型会导致模型性能下降，因为这种后处理方式会偏离模型原有的优化轨迹。此外，训练阶段的计算效率对于现代大型语言模型的开发至关重要，但现有方法主要关注推理阶段，未能有效解决训练阶段的计算挑战。

作者还指出，一些稀疏注意力方法包含不可训练的组件，例如基于聚类或哈希的离散操作，这些操作会阻断梯度流，限制模型学习最优稀疏模式的能力。此外，即使某些稀疏注意力方法理论上可训练，但在实际训练中存在效率问题。例如，基于哈希的稀疏注意力方法在训练时需要加载大量分散的KV缓存，这与现代高效注意力技术（如FlashAttention）所依赖的连续内存访问和块计算不兼容，导致硬件利用率低下，训练效率显著降低。

基于这些分析，作者强调了开发一种原生可训练稀疏注意力机制的必要性，这种机制需要同时满足高效推理和训练的要求。这为后续介绍NSA（Native Sparse Attention）的设计和实现奠定了基础，NSA通过算法创新和硬件对齐优化，实现了高效的训练和推理，同时保持了模型性能。