专栏名称: Dots机构投资者社区

机构投资者组成的社区，深度点评财经事件

DeepSeek革命性NSA注意力机制问世！梁文锋上阵，长文本推理能力飙升

Dots机构投资者社区 · 公众号 · · 2025-02-19 08:15

正文

本文转自微信公众号“新智元”，作者：新智元。

编辑：编辑部

【导读】 DeepSeek联手两大机构祭出神作，再次惊艳全世界。创始人亲自率队，提出了革命性注意力机制NSA，在通用、长文本、思维链推理基准测试中，刷新SOTA碾压全注意力，很有可能，NSA是对Transformer注意力机制的重大优化。

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块，以及滑动注意力以用于局部 ‍ 上下文。

上图右：可视化每个分支产生的不同注意力模式。绿色区域表示需要计算注意力分数的区域，而白色区域表示可以跳过的区域。

方法概述

具体而言，在注意力机制中，每个查询token q_t与所有前面的键k_:t计算相关性得分，以生成值v_:t的加权和。形式上，对于长度为t的输入序列，注意力操作定义为：

其中Attn表示注意力函数：

这里，𝛼_𝑡,𝑖表示q_𝑡和k_𝑖之间的注意力权重，𝑑_𝑘是键的特征维度。随着序列长度的增加，注意力计算在整体计算成本中的占比越来越大，给长序列上下文处理带来了显著的挑战。

算术强度是计算操作与内存访问的比率，它本质上决定了硬件上的算法优化。

对于计算任务，当算术强度超过这个临界阈值时，任务变为计算受限（受限于GPU的FLOPS），而低于该阈值时则变为内存受限（受限于内存带宽）。

具体到因果自注意力机制，在训练和预填充阶段，批量矩阵乘法和注意力计算表现出高算术强度，使得这些阶段在现代加速器上成为计算受限阶段。

相比之下，自回归解码阶段则为内存带宽受限，因为每次前向传递只生成一个token，同时需要加载整个键值缓存，导致算术强度较低。

就这导致了不同的优化目标——在训练和预填充阶段减少计算成本，而在解码阶段减少内存访问。

NSA的整体框架

为了利用注意力机制在自然稀疏模式下的潜力，研究人员提出用更紧凑、信息更密集的表示键值对

，来替换方程1中的原始键值对k_:𝑡，v_:𝑡，并根据每个查询q_𝑡进行优化。

具体而言，将优化后的注意力输出正式定义如下：

其中，

是基于当前查询q_𝑡和上下文记忆k_:𝑡，v_:𝑡动态构建的。

可以设计不同的映射策略，来生成不同类别的

，并将它们组合如下：

如图2所示，NSA采用了三种映射策略 C = {cmp,slc,win}，分别表示对键值对的压缩（compression）、选择（selection）和滑动窗口（sliding window）。

是与策略c对应的门控得分（gate score），通过MLP和sigmoid激活函数从输入特征中学习得到。

令𝑁_𝑡表示重映射后的键值对总数：

通过确保𝑁_𝑡 ≪ 𝑡，就能维持高稀疏率，从而优化计算效率。

接下来，就是重新映射策略的设计，包括token压缩、token选取和滑动窗口的具体算法。

关键算法组件

核心算法是重映射策略fK和fV的设计：token压缩、token选择和滑动窗口。

token压缩：将连续的键或值块聚合为块级表示，得到压缩后的键和值。压缩表示能够捕获较粗粒度的高级语义信息，并减少注意力机制的计算负担。
token选择：用于识别并保留最相关的token，同时减少计算开销。其中，按块选择在空间连续的块中处理键和值序列：首先将键值序列划分为选择块；为每个块分配重要性分数，识别出最重要的块以用于注意力计算。而Top-𝑛块选择在获得选择块重要性得分后，按块重要性得分排序，选择top-𝑛稀疏区块中的token，被选中的键和值会参与注意力计算。
滑动窗口：专门的滑动窗口分支，明确处理局部上下文，使得其他分支（压缩和选择）可以专注于学习各自的特征，而不会被局部模式所干扰。

这三种策略都是为了提高计算和内存效率。token压缩减少了数据量，token选择减少了计算量，滑动窗口限制了需要关注的序列长度。

硬件优化的内核设计

新研究使用Triton实现了与硬件对齐的稀疏注意力内核。

遵循当前最先进的LLM，专注于具有共享KV缓存的架构，为稀疏选择注意力引入了专门的内核设计。

为了解决内存访问效率低下问题，关键优化在于不同的查询分组策略：对于查询序列上的每个位置，将GQA组内的所有查询头（它们共享相同的稀疏KV块）加载到SRAM中。

NSA的内核设计采用了以下策略：

查询加载（Grid Loop）：内核按GQA（Grid Query Attention）组加载查询。每个GQA组包含共享相同稀疏键值块的查询头。这种方法减少了冗余的键值传输，提高了内存访问效率。
键值获取（Inner Loop）：在每个查询加载后，内核获取相应的稀疏键值块。这些键值块存储在高带宽内存（HBM）中。内核将这些键值块传输到片上静态随机存取存储器（SRAM）中，以进行注意力计算。
注意力计算：内核在SRAM中执行注意力计算。绿色块表示存储在SRAM中的数据，蓝色块表示存储在HBM中的数据。这种内存层次结构优化了数据传输和计算效率。

DeepSeek革命性NSA注意力机制问世！梁文锋上阵，长文本推理能力飙升

正文

【导读】 DeepSeek联手两大机构祭出神作，再次惊艳全世界。创始人亲自率队，提出了革命性注意力机制NSA，在通用、长文本、思维链推理基准测试中，刷新SOTA碾压全注意力，很有可能，NSA是对Transformer注意力机制的重大优化。

关键算法组件

硬件优化的内核设计

请到「今天看啥」查看全文