DeepSeek近日发布论文Native Sparse Attention：硬件对齐且可原生训练的稀疏注意力机制

最新推荐文章于 2025-03-10 00:15:00 发布
最新推荐文章于 2025-03-10 00:15:00 发布
阅读量1.3k 收藏 29
点赞数 49
分类专栏： 学术研究综述 DeepSeek 文章标签： 人工智能 nlp 深度学习 自注意力机制 长序列建模
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.youkuaiyun.com/MnivL/article/details/145834326
版权

长序列建模一直是自然语言处理领域的难题，而 NSA 的出现，为解决这一难题提供了新的思路。

论文地址：https://arxiv.org/pdf/2502.11089v1

核心贡献

首次实现硬件对齐的稀疏注意力机制，显著提升长上下文建模效率。

技术突破点

  1. 创新方法：用动态分层稀疏策略替代传统全注意力机制
  2. 关键优势：在64k长度序列上，解码、前向传播和反向传播速度提升显著
  3. 基础原理：类似"给AI添加决策流程图"，通过压缩和选择关键信息减少计算量

现实影响

  • 短期应用：加速长文本处理，如代码生成和文档分析
  • 长期潜力：推动通用AI发展，实现更高效的长上下文推理
  • 风险警示：可能在某些任务上引入信息丢失，需进一步优化

延伸思考

  • 该研究与马斯克的脑机接口是否有结合可能？
  • 如果技术开源，开发者最应该关注哪个模块？

资源索引

原文精简翻译

1. 引言

随着下一代大语言模型对长上下文建模的需求日益增加，标准注意力机制的高计算成本成为了一个显著的挑战。稀疏注意力机制为提高效率提供了有前景的方向，同时保持模型的能力。我们提出了NSA（Natively trainable Sparse Attention），一种硬件对齐且可原生训练的稀疏注意力机制，通过算法创新和硬件优化实现高效的长上下文建模。

​

图1：全注意力模型与NSA的性能和效率对比

​

图1展示了NSA在64k长度序列处理中的显著计算加速，同时在通用基准、长上下文任务和推理评估中超越全注意力基线。

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

现代稀疏注意力方法在减少Transformer模型的理论计算复杂度方面取得了显著进展。然而，大多数方法主要在推理阶段应用稀疏性，而保留了预训练的全注意力骨干，这可能引入架构偏差，限制其充分利用稀疏注意力优势的能力。

2.1 高效推理的错觉

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

  • 阶段限制的稀疏性：例如H2O在自回归解码期间应用稀疏性，但在预填充阶段需要计算密集的预处理。
  • 与高级注意力架构的不兼容性：一些稀疏注意力方法无法适应现代解码高效架构，如多查询注意力（MQA）和分组查询注意力（GQA）。

2.2 可训练稀疏性的神话

我们分析推理方法的两个关键见解促使我们追求原生可训练的稀疏注意力：

  • 性能下降：事后应用稀疏性迫使模型偏离其预训练的优化轨迹。
  • 训练效率需求：高效处理长序列训练对于现代LLM开发至关重要。

3. 方法论

我们的技术方法涵盖算法设计和内核优化。我们首先介绍方法背景，然后介绍NSA的整体框架，接着是其关键算法组件，最后是我们硬件优化的内核设计。

3.1 背景

注意力机制广泛用于语言建模，其中每个查询标记计算与所有先前键的相关性分数，以生成值的加权和。

3.2 整体框架

为了利用自然稀疏模式的潜力，我们提出用更紧凑和信息密集的表示键值对替换原始键值对。

​

图2：NSA架构概述

​

图2展示了NSA的架构框架，通过三个并行注意力分支处理输入序列：压缩注意力、选择注意力和滑动窗口注意力。

3.3 算法设计

我们介绍了重新映射策略和的设计：标记压缩、标记选择和滑动窗口。

3.3.1 标记压缩

通过将键或值的顺序块聚合为块级表示，我们获得压缩的键和值，捕捉整个块的信息。

3.3.2 标记选择

仅使用压缩键和值可能会丢失重要的细粒度信息，促使我们选择性地保留单个键和值。

3.3.3 滑动窗口

在注意力机制中，局部模式通常适应得更快，并可能主导学习过程，防止模型从压缩和选择标记中有效学习。为了解决这个问题，我们引入了一个专门的滑动窗口分支，明确处理局部上下文。

3.4 内核设计

为了在训练和预填充期间实现FlashAttention级别的加速，我们在Triton上实现了硬件对齐的稀疏注意力内核。

​

图3：NSA的内核设计

​

图3展示了NSA的内核设计，通过GQA组加载查询（网格循环），获取相应的稀疏KV块（内部循环），并在SRAM上执行注意力计算。

4. 实验

我们通过三个视角评估NSA：通用基准性能、长上下文基准性能和链式思维推理性能，与全注意力基线和最先进的稀疏注意力方法进行比较。

4.1 预训练设置

我们的实验采用结合分组查询注意力（GQA）和专家混合（MoE）的骨干，总参数为27B，活动参数为3B。

​

图4：全注意力与NSA在27B参数模型上的预训练损失对比

​

图4展示了NSA和全注意力基线的预训练损失对比，NSA始终表现更好。

4.2 基线方法

除了与全注意力进行比较外，我们还评估了几种最先进的推理阶段稀疏注意力方法：H2O、infLLM、Quest和Exact-Top。

4.3 性能比较

模型SQAMQASyntheticCodeAvg.
H2O0.4280.4290.3080.1120.101
InfLLM0.4740.5170.3560.3060.250
Quest0.4950.5610.3650.2950.245
Exact-Top0.5020.6050.3970.3210.288
Full Attn0.5120.6230.4090.3500.305
NSA0.5030.6240.4320.4370.356

表2展示了NSA在LongBench上的性能比较，包括单文档QA、多文档QA、合成和代码任务类别。NSA在大多数基准上表现优于全注意力。

5. 效率分析

我们在8-GPU A100系统上评估NSA与全注意力的计算效率。

5.1 训练速度

我们比较了基于Triton的NSA注意力和全注意力与Triton-based FlashAttention-2的实现。

​

图6：NSA内核与FlashAttention-2内核的对比

​

图6展示了NSA内核在64k上下文长度上实现了9.0倍的前向和6.0倍的反向加速。

5.2 解码速度

注意力的解码速度主要由内存访问瓶颈决定，这与KV缓存加载量密切相关。

上下文长度8192163843276865536
全注意力8192163843276865536
NSA2048256035845632
预期加速4×6.4×9.1×11.6×

表4展示了解码期间每次注意力操作的内存访问量（以等效标记数表示）。由于解码的低算术强度和内存限制，预期加速与内存访问量近似线性。

6. 讨论

在本节中，我们反思NSA的开发过程，并讨论从不同稀疏注意力策略探索中获得的关键见解。

6.1 替代标记选择策略的挑战

在设计NSA之前，我们探索了将现有稀疏注意力方法适应训练阶段。然而，这些尝试遇到了各种挑战，促使我们设计不同的稀疏注意力架构。

​

图7：不同标记选择策略的训练损失对比

​

图7展示了在3B参数模型上，NSA与全注意力和不同标记选择策略的训练损失对比。NSA表现更好。

6.2 可视化

为了探索Transformer注意力分布的潜在模式并为我们的设计寻求灵感，我们可视化了预训练的27B全注意力模型的注意力图。

​

图8：全注意力Transformer的注意力图可视化

​

图8展示了注意力图的可视化，浅色区域表示较高的注意力值。如图所示，注意力分数呈现出块状聚类分布。

7. 相关工作

我们回顾了通过稀疏注意力提高注意力计算效率的现有方法。这些方法可以根据其核心策略大致分为三类：固定稀疏模式、动态标记修剪和查询感知选择。

7.1 固定稀疏模式

滑动窗口是一种常用方法，允许查询仅在固定窗口内计算注意力。

7.2 动态标记修剪

H2O实现了在解码期间减少KV缓存内存使用的自适应方法。

7.3 查询感知选择

Quest采用块状选择策略，其中每个块的重要性通过查询与键块的坐标最小最大值乘积来估计。

8. 结论

我们提出了NSA，一种硬件对齐的稀疏注意力架构，用于高效的长上下文建模。通过将分层标记压缩与块状标记选择结合在可训练架构中，我们的架构在保持全注意力性能的同时实现了加速训练和推理。

附录A AIME结果示例

提示：

我们的结果：

基线结果：

专业术语词典

1. 稀疏注意力（Sparse Attention）

  • 解释：一种减少注意力机制计算量的方法，通过只计算部分关键的查询-键对，而不是所有可能的对。类似于“选择性阅读”，只关注最重要的部分。

2. 全注意力（Full Attention）

  • 解释：传统的注意力机制，计算每个查询与所有键的相关性。类似于“逐字阅读”，计算量大但信息全面。

3. 硬件对齐（Hardware-Aligned）

  • 解释：算法设计时考虑硬件的特性（如内存访问模式、计算单元利用率），以最大化硬件性能。类似于“为特定机器定制工具”。

4. 解码（Decoding）

  • 解释：在生成文本时，模型逐步预测下一个词的过程。类似于“逐词写作”。

5. 前向传播（Forward Propagation）

  • 解释：神经网络中，输入数据通过各层计算得到输出的过程。类似于“信息从输入到输出的流动”。

6. 反向传播（Backward Propagation）

  • 解释：神经网络训练时，根据损失函数计算梯度并更新模型参数的过程。类似于“根据错误调整学习方法”。

7. KV缓存（KV Cache）

  • 解释：在解码过程中，存储键（Key）和值（Value）的内存区域，用于加速注意力计算。类似于“临时存储重要信息的笔记本”。

8. 算术强度（Arithmetic Intensity）

  • 解释：计算操作数与内存访问数的比值，用于衡量算法在硬件上的效率。类似于“计算量与数据搬运量的平衡”。

9. 分组查询注意力（Grouped-Query Attention, GQA）

  • 解释：一种注意力机制变体，多个查询共享相同的键和值，以减少内存访问和计算量。类似于“多人共享同一份参考资料”。

10. 专家混合模型（Mixture of Experts, MoE）

  • 解释：一种模型架构，包含多个专家网络，每个输入由部分专家处理。类似于“分诊系统，不同问题由不同专家解决”。

11. 滑动窗口（Sliding Window）

  • 解释：一种局部注意力机制，只计算查询附近一定范围内的键值对。类似于“只关注当前段落的内容”。

12. 注意力分数（Attention Score）

  • 解释：衡量查询与键之间相关性的数值，用于加权值。类似于“关键词的重要性评分”。

13. 链式思维推理（Chain-of-Thought Reasoning）

  • 解释：一种推理方法，模型逐步生成中间步骤以解决问题。类似于“逐步推导数学题的解法”。

14. 预训练（Pretraining）

  • 解释：在大规模数据上训练模型，使其学习通用特征。类似于“学习基础知识”。

15. 微调（Fine-Tuning）

  • 解释：在特定任务数据上进一步训练预训练模型，以适应具体任务。类似于“针对特定领域进行专项训练”。

16. FlashAttention

  • 解释：一种高效的注意力计算实现，通过优化内存访问和计算顺序来加速注意力机制。类似于“优化后的快速阅读方法”。

17. Triton

  • 解释：一种用于编写高效GPU内核的编程语言和编译器，常用于深度学习模型的加速。类似于“为GPU定制的编程工具”。

18. 内存访问瓶颈（Memory Access Bottleneck）

  • 解释：计算任务中，内存访问速度限制了整体性能的现象。类似于“搬运数据的速度跟不上计算速度”。

19. 块状选择（Blockwise Selection）

  • 解释：将键和值序列分成连续块，并选择最重要的块进行计算。类似于“按章节选择重要内容”。

20. 注意力图（Attention Map）

  • 解释：可视化注意力分数的矩阵，显示查询与键之间的相关性。类似于“关键词与上下文的关联图”。

DeepSeek 25.2.16 最新论文NSA：原生稀疏注意力：硬件对齐且可原生训练的稀疏注意力解读
轻口味的专栏
02-20 381
NSA（Native Sparse Attention）作为一种高效且原生可训练的稀疏注意力机制，具有显著的计算效率提升和性能优化能力，因此在许多需要处理长文本或大规模序列数据的应用场景中具有巨大的应用潜力。通过这些设计，NSA能够在基于深度学习的链式思维数学推理任务中表现出色，尤其是在更长的推理链中，表现出显著的性能提升。通过这种层次化的稀疏注意力设计，NSA能够在保持全局上下文感知的同时，实现高效的局部信息检索，从而在64k上下文的“针尖在干草堆中”测试中达到完美的检索准确率。
AI学习指南DeepSeek篇(12)-论文导读 Native Sparse Attention
俞兆鹏的博客
02-25 61
论文标题发布时间：2025 年 2 月论文地址arXiv链接主要内容动态分层稀疏策略：结合粗粒度的 token 压缩和细粒度的 token 选择，保持全局上下文感知的同时确保局部精确性。硬件对齐优化：通过算术强度均衡的算法设计，实现显著的加速，并针对现代硬件进行了优化实现。端到端训练支持：支持端到端训练，降低预训练计算成本，而不会牺牲模型性能。实验结果：在多个基准测试中，NSA 的性能与全注意力（Full Attention）模型相当甚至更优，同时在解码、前向传播和反向传播阶段均实现了显著的加速。
稀疏Attention
rosefun96的博客
10-29 6815
1. 模型 Self Attention是O(n2)O(n^2)O(n2)的，那是因为它要对序列中的任意两个向量都要计算相关度，得到一个n2n^2n2大小的相关度矩阵： 左边显示了注意力矩阵，右变显示了关联性，这表明每个元素都跟序列内所有元素有关联。 所以，如果要节省显存，加快计算速度，那么一个基本的思路就是减少关联性的计算，也就是认为每个元素只跟序列内的一部分元素相关，这就是稀疏Attenti...
稀疏自注意力（Sparse Self-Attention）
m0_51200050的博客
06-12 4023
稀疏自注意力以及举例说明
稀疏注意力：让长文本建模更高效
gs80140的专栏
03-10 1275
稀疏注意力（Sparse Attention）通过限制注意力范围，减少计算开销，同时保持模型性能，为处理长序列提供了高效的解决方案。通过动态分层稀疏策略、连接模式设计、Hilbert 曲线置换和 8-bit 量化等技术，稀疏注意力显著降低了计算复杂度和内存占用，同时保留了关键信息，确保模型性能不下降。
Deepseek Natively Sparse Attention
asd8705的专栏
02-20 742
提出一种高效、可训练的稀疏注意力机制，以提高长文本处理的计算效率，同时保持模型性能。：可根据任务需求设定不同的窗口大小（例如代码生成任务可能需要更大的窗口）。NSA 通过以下三种不同的注意力计算路径，让 Transformer。是可学习的门控参数（Gating Mechanism）。- 这样，原始 Key-Value 变成了。- 这样可以降低计算量，同时保留重要信息。，为长文本建模提供了新的解决方案。的方式，将 Key-Value。，提高计算效率，同时保留重要信息。，大幅减少了 Key 的数量。
深度学习中的稀疏注意力
热门推荐
不要给自己设限，尝试更多可能（思所向 皆可往）
01-07 1万+
稀疏注意力的实现并不仅限于DeepSpeed。虽然DeepSpeed提供了一种高效的稀疏注意力实现，但其他深度学习框架和库也可能提供稀疏注意力的支持。例如，Hugging Face的Transformers库就包含了一些稀疏注意力的实现，如Longformer和BigBird。
《论文阅读》Sparse Sinkhorn Attention
未知丶的博客
03-03 2613
留个笔记自用 PCT: Point Cloud Transformer 做什么 点云的概念：点云是在同一空间参考系下表达目标空间分布和目标表面特性的海量点集合，在获取物体表面每个采样点的空间坐标后，得到的是点的集合，称之为“点云”（Point Cloud）。 点包含了丰富的信息，包括三维坐标X，Y，Z、颜色、分类值、强度值、时间等等，不一一列举。 一般的3D点云都是使用深度传感器扫描得到的，可以简单理解为相比2维点，点云是3D的采样 做了什么 设计了一种sparse attention的计算方法，引入了一
DeepSeek-NSA 论文阅读【Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention】
qq_41685627的博客
02-20 1227
Long-context modeling is crucial for next-generation language models, yet the high computational cost of standard attention mechanisms poses significant computational challenges. 对于下一代生成式语言模型而言，长文本建模至关重要，然而标准注意力机制的高计算成本带来了显著的计算挑战。
DeepSeek最新成果：原生稀疏注意力（Native Sparse Attention-Aligned and Natively Trainable Sparse Attention）.pdf
02-22
本文介绍了一种新型的可训练稀疏注意力机制，即原生稀疏注意力（Native Sparse Attention，简称NSA），它通过结合算法创新与硬件对齐优化来实现高效的长距离上下文建模。 NSA机制采用了一种层级的稀疏策略，将粗...
梁文锋亲自挂名DeepSeek发布新论文
sexy19910923的博客
02-18 1172
由 DeepSeek 联合创始人梁文锋亲自挂名的研究团队，在 arXiv 上发表了一篇题为“Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention”的论文，提出了一种全新的注意力机制架构 NSA（Native Sparse Attention，原生稀疏注意力）。
探索开放源代码：`sparse_attention` - 实现高效注意力机制的新范式
gitblog_00028的博客
04-25 702
探索开放源代码：sparse_attention - 实现高效注意力机制的新范式 去发现同类优质开源项目:https://gitcode.com/ 在深度学习领域，尤其是自然语言处理中，注意力机制（Attention Mechanism）已成为了不可或缺的一部分。它允许模型聚焦于输入序列中的关键部分，从而提高理解和预测的准确性。然而，随着序列长度的增长，全连接的注意力机制计算成本会急剧增加。为了解...
Sparse transformer - 之Sparse attention 稀疏注意力
weixin_56336619的博客
06-10 5419
Sparse Transformer: 主要目的是减少传统Transformer 的时间复杂度和空间复杂度。 通过top-k选择，将注意退化为稀疏注意。这样，保留最有助于引起注意的部分，并删除其他无关的信息。这种选择性方法在保存重要信息和消除噪声方面是有效的。注意力可以更多地集中在最有贡献的价值因素上。 论文《Generating Long Sequences with Sparse Transformers》 self-attention是O(n^2)，因为对序列中任意两个位置的均建立起关联； 节省显.
【DeepSeek】Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
weixin_44376877的博客
03-05 594
长上下文建模(Long-context modeling)对于下一代语言模型至关重要，但标准注意力机制(Attention Mechanism)的高计算成本带来了显著的挑战。稀疏注意力(Sparse Attention)为提升效率同时保持模型能力提供了方向。本研究提出了NSA（Natively trainable Sparse Attention），一种可原生训练的稀疏注意力机制，通过算法创新与硬件优化实现高效长上下文建模。NSA采用动态分层稀疏策略，结合粗粒度标记压缩(Token Compression)
《Sparse and Constrained Attention for Neural Machine Translation》
wwx123521的博客
03-21 601
背景：在NMT中，单词有时会从源中删除或在翻译中重复生成。 思路：普通的注意力计算中，在得到每个时间步隐藏层状态的分数后，对其softmax归一化处理得到权重值。sparsemax：在源语句单词上的稀疏，计算时只考虑一些单词，其余单词的概率为0。constrained softmax：在时间步上的稀疏，返回一个接近sofrmax(z)的分布，注意力概率受限于上界值，本文将这两个稀疏结合在一起，产...
【论文笔记】Understanding Long Programming Languages with Structure-Aware Sparse Attention
m0_61899108的博客
07-30 682
基于编程的预训练语言模型(PPLM)，如CodeBERT，在许多与下游代码相关的任务中取得了巨大的成功。但由于Transformer中自我注意的内存和计算复杂性随着序列长度的增加而呈二次曲线增长，因此PPLM通常将代码长度限制在512。然而，现实应用中的代码通常很长，例如代码搜索，现有的PPLM无法有效地处理这些代码。为解决这一问题，本文提出结构感知的稀疏注意力Transformer模型SASA，这是面向长代码序列的Transformer模型优化方法，致力于提升长代码场景下的效果和性能。...
论文阅读【SWINBERT: End-to-End Transformers with Sparse Attention for Video Captioning】
hei_hei_hei_的博客
05-17 2020
SWINBERT: End-to-End Transformers with Sparse Attention for Video Captioning 发表
大模型训练之加速篇 -attention优化【稀疏-＞线性化-＞分块计算MQA/GQA-＞计算(flash atten/dec)-＞缓存MLA】
weixin_36378508的博客
09-21 9717
加速
稀疏注意力机制：
最新发布
03-11
### 稀疏注意力机制的应用及实现 #### 应用场景 稀疏注意力机制是一种特殊的注意力形式，在处理大规模数据集时表现出色。这种机制允许模型仅关注输入序列中的一部分位置，而不是像传统全连接注意力建模那样考虑每一个可能的位置组合[^1]。 对于自然语言处理任务而言，采用稀疏注意力可以有效减少计算复杂度并加速训练过程。具体来说，在文本摘要生成、机器翻译等应用场景里，当面对较长文档或者句子时，利用局部敏感哈希算法或其他策略筛选出最有可能影响当前预测结果的关键词或短语作为重点考察对象，从而达到降维的目的[^4]。 #### 实现方式 为了实现稀瘦注意力功能，通常会引入特定的设计思路和技术手段： - **分块策略**：将整个输入划分为若干个小片段，并假设同一区域内元素间关联较强而跨区域联系较弱； - **基于距离衰减因子**：给定一对位置i,j,定义一个随|i-j|增大呈指数下降趋势的权重系数w(i,j)，以此体现远近关系； - **多尺度建模**：构建多层次的感受野结构，既能捕捉近距离依赖又能兼顾全局模式特征。 以下是Python代码示例展示了如何创建简单的稀疏自注意力层: ```python import torch.nn as nn from transformers import SparseAttentionConfig class SparseSelfAttention(nn.Module): def __init__(self, config:SparseAttentionConfig): super().__init__() self.sparse_attn = nn.MultiheadAttention(embed_dim=config.hidden_size, num_heads=config.num_attention_heads, batch_first=True) # 设置为稀疏配置 self.sparse_config = config def forward(self,x): attn_output,_=self.sparse_attn(x,x,x,need_weights=False) return attn_output ``` 此段代码使用了`transformers`库中的SparseAttentionConfig类来指定参数设置，实现了带有稀疏特性的自我注意模块[^5]。
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

歌刎

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付：¥1
获取中
扫码支付

您的余额不足，请更换扫码支付或充值

打赏作者

实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值