新智元
整合编辑：太平洋科技
发布于：2025-02-19 15:47
|
2月18日，DeepSeek官方再发布新论文！ 在这篇论文中，来自DeepSeek、北大和华盛顿大学的研究人员提出了一种全新的注意力机制NSA。
具体来说，它是一种面向硬件且支持原生训练的稀疏注意力机制（Sparse Attention），专为超快长上下文训练与推理设计。其核心组成包括—— 动态分层稀疏策略 粗粒度的token压缩 细粒度的token选择 通过针对现代硬件的优化设计，NSA（原生可训练稀疏注意力机制）在保证性能的同时显著提升了推理速度，并有效降低了预训练成本。 在通用基准测试、长文本处理以及基于指令的推理任务中，它的表现均能达到甚至超越传统全注意力（Full Attention）模型的水平。
值得一提的是，DeepSeek创始人梁文锋这次不仅亲自上阵，并且还亲自提交了论文。
论文地址：https://arxiv.org/abs/2502.11089 目前，这成为了该论文的最大亮点之一。 网友们也不忘艾特整天放小道消息的奥特曼：看，DeepSeek又有新论文了。
部署高效稀疏注意力机制的关键挑战 随着序列长度的增加，注意力计算在整体计算成本中变得越来越占主导地位，这对长上下文处理提出了重大挑战。 尤其在深入推理、代码库级别的代码生成以及多轮自主智能体系统等多种现实应用的推动下，这一趋势尤为明显。 近期OpenAI的o系列模型、DeepSeek-R1和Gemini 1.5 Pro，在数千个token的范围内保持连贯的多轮对话，同时执行复杂的长距离依赖推理了。 然而，随着序列长度的增加，传统注意力机制的高计算复杂度，成为关键的时延瓶颈。 理论估算表明，在处理64k长度上下文时，使用softmax结构的注意力计算占总时延迟的70–80%，更凸显了对更高效注意力机制的迫切需求。 部署高效的稀疏注意力机制必须应对两个关键挑战：硬件对齐的推理加速和训练感知的算法设计。 为此，研究团队提出了NSA，一种原生可训练的稀疏注意力架构，并集成了分层token建模。同时，还设计了专用计算内核，以最大化实际计算效率。 如图2所示，NSA通过将键值对组织为时间块（temporal blocks），并通过三条注意力路径进行计算，从而降低每个查询的计算量。 NSA引入了两个核心创新点： 硬件对齐系统：优化块级稀疏注意力，使其充分利用Tensor Core并优化内存访问，从而实现平衡的算术强度。 训练感知设计：通过高效算法和反向传播运算符实现稳定的端到端训练，使NSA能够同时支持高效推理与完整训练流程。
上图左：该框架通过三个并行的注意力分支处理输入序列。对于给定的查询，前面的键和值被处理成压缩注意力以用于粗粒度模式、选择注意力以用于重要token块，以及滑动注意力以用于局部 上下文。 上图右：可视化每个分支产生的不同注意力模式。绿色区域表示需要计算注意力分数的区域，而白色区域表示可以跳过的区域。 革命性NSA，碾压全注意力 接下来，研究人员从三个方面评估了NSA，将其与全注意力基线，以及SOTA稀疏注意力方法进行比较： 通用基准性能 长文本基准性能 思维链推理性能 如下图4所示，NSA和全注意力基线的预训练损失曲线，呈现出稳定和平滑的下降，NSA始终优于全注意力模型。
通用基准评估研究人员在一套全面的基准测试中评估了预训练的NSA和全注意力基线，覆盖了知识、推理和编码能力。 结果如表1所示，尽管NSA具有稀疏性，但它在总体性能上表现优越，在9个指标中有7个超过了包括全注意力在内的所有基线。 这表明，尽管NSA在较短序列上可能未能充分利用其效率优势，但它仍然表现出强大的性能。 值得注意的是，NSA在推理相关的基准测试中显示出显著的提升（DROP: +0.042，GSM8K: +0.034）。 这进一步证明了，预训练有助于模型发展专门的注意力机制。 NSA稀疏注意力预训练机制迫使模型集中关注最重要的信息，可能通过过滤掉来自无关注意力路径的噪声来增强性能。 不同评估中的一致表现，也验证了NSA作为通用架构的稳健性。
长文本评估如图5所示，NSA在64k上下文中的「大海捞针」测试中，在所有位置上都实现了完美的检索准确率。 这一性能得益于层次化的稀疏注意力设计，该设计结合了压缩token以实现高效的全局上下文扫描，以及用于精确局部信息检索的选择token。 粗粒度的压缩在低计算成本下识别相关的上下文块，而对选择token的token级注意力则确保了关键细粒度信息的保留。 这种设计，使NSA能够同时保持全局感知能力和局部精确性。
此外，研究人员还在LongBench基准上，对NSA进行了评估。 如表2所示，NSA达下了最高的平均分0.469，优于所有基准（比全注意力高0.032，比Exact-Top高0.046）。
思维链推理评估考虑到RL在小规模模型上的效果有限，研究人员采用了从DeepSeek-R1进行知识蒸馏的方法，使用100亿个32k长度的数学推理轨迹进行监督微调（SFT）。 由此，产生了两个模型：Full Attention-R（全注意力基准）和NSA-R（稀疏变体）。 在AIME 24基准测试上，研究人员评估了这两个模型。
结果如下表3所示，在8k上下文设置下，NSA-R比Full Attention-R获得了更高的准确率（高出0.075），这种优势在16k上下文中依然保持（高出0.054）。 这些结果验证了原生稀疏注意力的两个关键优势： 预训练的稀疏注意力模式，能够有效捕获复杂数学推导中至关重要的长程逻辑依赖关系； 架构采用硬件对齐设计，在增加推理深度的同时保持足够的上下文密度，避免了灾难性遗忘。 在不同上下文长度下的持续优异表现，证实了当稀疏注意力原生集成到训练流程中时，它在高级推理任务中的可行性。
64k上下文，前向传播9倍速飙升 研究团队在配备8个A100 GPU的系统上评估了NSA相对于全注意力机制（Full Attention）的计算效率。 训练速度为了评估训练速度，研究团队将基于Triton的NSA实现与基于Triton的FlashAttention-2进行了比较，以确保在同一后端进行公平的速度比较。 实验结果表明，随着上下文长度的增加，NSA实现了逐渐提高的加速，在64k上下文长度下实现了高达9.0倍的前向加速和6.0倍的反向加速。值得注意的是，随着序列长度的增加，速度优势变得更加明显。 这种加速源于NSA的硬件对齐算法设计，旨在最大限度地提高稀疏注意架构的效率：分块内存访问模式通过合并加载最大限度地利用了Tensor Core。内核中精细的循环调度消除了冗余的KV传输。 下图6比较了Triton实现的NSA和全注意力与Triton实现的FlashAttention-2，以确保在同一后端进行公平的速度比较。
解码速度注意力机制的解码速度主要受限于内存访问瓶颈，这与KV缓存的加载量密切相关。 随着解码长度的增加，本文中的方法显著降低了延迟，在64k上下文长度下实现了高达11.6倍的加速。 如下表4所示NSA在长序列解码时具有显著的效率优势。通过减少内存访问量，NSA能够显著提高解码速度，尤其是在处理较长的上下文时。 这种内存访问效率的优势也会随着序列的增长而放大。
讨论 研究团队反思了NSA的开发过程，并讨论了在探索不同稀疏注意力策略过程中获得的关键见解。 虽然本文中的方法展示了有希望的结果，但了解替代策略遇到的挑战并分析注意力分布模式，为未来的研究方向提供了有价值的背景。 研究团队首先考察了促使他们做出设计选择的替代token选择策略所面临的挑战，然后通过可视化来深入了解注意力分布模式。 替代token选择策略研究人员考察了像ClusterKV这样的基于聚类的策略。这些方法将来自同一集群的Key和Value存储在连续的内存区域中。虽然理论上对于训练和推理是可行的，但它们面临着三个显著的挑战。 首先是动态聚类机制引入了不可忽略的计算开销。其次，算子优化困难，集群间的不平衡加剧了这种情况，特别是在混合专家（MoE）系统中，倾斜的专家并行（EP）组执行时间导致了持续的负载不平衡。 最后是实现约束，需要强制性的周期性重新聚类和chunk-sequential训练协议。这些因素结合起来造成了巨大的瓶颈，严重限制了它们在实际部署中的有效性 研究人员还考虑了与NSA不同的分块key、value选择策略，例如Quest和InfLLM。 然而，现有方法面临两个关键问题。首先由于选择操作是不可微的，因此基于神经网络的重要性得分计算依赖于辅助损失，这增加了算子开销，并且通常会降低模型性能。其次启发式无参数的重要性得分计算策略的召回率较低，导致性能欠佳。 研究人员在一个具有类似架构的3B参数模型上评估了这两种方法，并将它们的损失曲线与NSA和全注意力进行了比较。 对于基于辅助损失的选择方法，为每个块引入了额外的查询和代表性key，以估计块重要性得分。这些得分由原始查询和每个块内的key之间的平均注意力得分监督。 对于启发式无参数选择方法，遵循Quest的策略，实现了使用查询和key块的坐标式min-max乘积的直接选择，而没有引入额外的参数。 研究人员还探索了一种冷启动训练方法，其中在过渡到启发式分块选择之前，对初始的1000步应用全注意力。这两种方法的损失都较差。 下图7展示了在3B参数模型上，使用全注意力和不同token选择策略时的训练损失曲线比较。
可视化研究人员对预训练的27B全注意力模型的注意力图进行了可视化（如下图8所示）。 可视化结果揭示了一个有趣的模式：注意力分数倾向于表现出分块聚类的特性，即相邻的key通常显示出相似的注意力分数。
这一观察结果启发了NSA的设计，表明基于空间连续性选择key块可能是一种很有前景的方法。分块聚类现象表明，序列中相邻的tokens可能与查询tokens共享某些语义关系，尽管这些关系的确切性质需要进一步研究。 并促使研究人员探索一种在连续token块上操作的稀疏注意力机制，而不是在单个token上操作，旨在提高计算效率并保留高注意力模式。 正如之前讨论的，在设计NSA架构时，研究人员面临着基于Key-Clustering的策略和其他分块选择策略的挑战，包括计算开销、算子优化困难和召回率低等问题。 结论 本文中，研究团队提出了NSA，一种面向硬件对齐的稀疏注意力架构，用于高效的长上下文建模。 通过将分层token压缩与块级token选择集成到一个可训练的架构中，架构在保持全注意力性能的同时，实现了加速的训练和推理。 该方法通过以下方式推进了当前最佳技术水平：在通用基准测试中达到了全注意力的性能，长上下文评估中的建模能力更胜一筹，推理能力得到增强，同时计算延迟显著降低，实现了可观的加速。 参考资料：JHNYZ https://arxiv.org/abs/2502.11089 本文来源：新智元 |
原创栏目
IT百科
网友评论
聚超值•精选
新智元
硬件编年史
