摘要 (Abstract)
长文本建模是下一代语言模型的关键能力，但标准注意力机制的高昂计算成本构成了巨大的挑战。稀疏注意力为提升效率同时保持模型性能提供了一个充满希望的方向。在这篇解读中，我将介绍我们提出的NSA（Natively Trainable Sparse Attention），一种将算法创新与硬件协同优化相结合的原生可训练稀疏注意力机制，以实现高效的长文本建模。NSA采用动态分层稀疏策略，结合了粗粒度的令牌压缩与细粒度的令牌选择，旨在同时保留全局上下文感知和局部细节精度。我们的方法通过两大核心创新推动了稀疏注意力的设计：首先，我们通过算术强度均衡的算法设计，并针对现代硬件进行实现优化，实现了显著的速度提升。其次，我们实现了端到端的训练能力，在不牺牲模型性能的前提下，大幅降低了预训练的计算量。实验证明，使用NSA预训练的模型在通用基准、长文本任务和基于指令的推理任务上，性能持平甚至超越了全注意力（Full Attention）模型。同时，在处理64k长度的序列时，NSA在解码、前向传播和反向传播等各个阶段都比全注意力机制取得了巨大的速度提升，验证了其在模型整个生命周期中的高效性。这项工作证明了，一个精心设计的、与硬件协同的稀疏注意力架构，不仅可以在效率上，更可以在模型能力上，与密集注意力分庭抗礼，甚至更胜一筹。
引言：长文本时代的“注意力”危机与我们的破局之路
大家好，我是袁经洋。今天，我想和大家聊聊我们在大型语言模型（LLM）领域做的一些探索，特别是关于如何让模型更“聪明”地处理超长文本。想象一下，如果一个AI能一口气读完一整本《三体》或者一个庞大的代码库，并在此基础上进行深入的推理和创作，那将是多么激动人心的场景。这正是我们努力的方向——长文本建模。
然而，理想很丰满，现实却很“骨感”。阻碍我们前进的最大“拦路虎”，就是注意力机制（Attention Mechanism）的计算量问题。它是现代LLM的基石，赋予了模型理解上下文的能力。但它的工作方式有点“暴力”——对于每一个新生成的词，它都要回头去“关注”前面出现过的所有词，计算一个相关性得分。当文本长度从几千增长到几万、几十万时，这种计算量会呈平方级暴增（\(O(N^2)\)），就像一场计算资源的“海啸”，让训练和推理变得异常缓慢和昂贵。
类比：一场盛大却低效的圆桌会议
想象一下，你正在主持一个有数万名与会者的圆桌会议。每当一位新成员（Query）要发言时，为了确保他了解所有背景信息，你要求他与在座的每一位旧成员（Key）都单独私聊一遍，评估每个人的重要性，然后根据这个重要性来决定听取谁的意见（Value）。这个过程无疑是极其严谨的，但效率也低得令人发指。当会议规模扩大，整个流程就会陷入瘫痪。这就是全注意力（Full Attention）机制面临的困境。
社区当然也意识到了这个问题，并提出了“稀疏注意力”这条路。核心思想很简单：我们真的需要让每个新成员都和所有人私聊吗？或许，大部分人提供的信息是冗余的，只有少数几个人（或几个小组）的意见是至关重要的。稀疏注意力的目标，就是精准地找出这些“关键人物”，只与他们进行深入交流，从而大大减少计算量。这就像在大会议中，我们只挑选几个关键代表发言，或者让每个小组派一个代表总结观点，效率自然就高了。
尽管想法很美好，但现有的稀疏注意力方法在实践中却遇到了两大瓶颈：
- 理论快，实践慢：很多算法虽然在理论上减少了计算步骤，但因为其操作方式（比如随机、零散地读取内存）与现代GPU的“脾气”不合，导致实际运行起来并没有快多少。GPU喜欢整齐、连续的数据块，就像一个喜欢批量处理订单的仓库管理员，你让他东拿一个西取一个，他效率自然就低了。
- 训练与推理的“割裂”：大多数方法都是在已经用全注意力训练好的模型上，到推理阶段才“强行”让它变得稀疏。这就像一个习惯了全面思考的人，突然被要求只看重点，他可能会感到不适，甚至做出错误的判断。模型也是如此，这种“后处理”式稀疏化，往往会导致性能下降。更重要的是，训练过程本身依然是昂贵的，我们希望能从一开始就用更高效的方式来训练模型。
正是为了攻克这两个难题，我们设计了NSA（Natively Trainable Sparse Attention），即“原生可训练稀疏注意力”。“原生”意味着它从训练之初就是稀疏的，模型天生就学会在稀疏的世界里思考；“可训练”意味着稀疏的模式不是固定的，而是模型自己学习得来的；而这一切，都建立在与硬件（特别是GPU）“心有灵犀”的协同设计之上。我们的目标是打造一个既快又强，贯穿训练和推理全流程的高效注意力机制。接下来，我将带大家深入了解NSA的内部构造和设计哲学。
NSA的核心设计：三管齐下的分层注意力策略
NSA的核心思想，可以用一句话概括：用一种更智能、更具层次感的方式来替代原来“一视同仁”的全注意力。我们不再直接使用原始的、冗长的Key和Value序列，而是为每个Query动态构建一个浓缩、高效的“信息摘要”集合 \(\tilde{K}_t, \tilde{V}_t\)。
传统注意力：\(o_t = \text{Attn}(q_t, K_{:t}, V_{:t})\)
NSA 的核心思想：\(o_t^* = \text{Attn}(q_t, \tilde{K}_t, \tilde{V}_t)\)，其中 \(\tilde{K}_t, \tilde{V}_t\) 是从 \(K_{:t}, V_{:t}\) 中智能提取的精华。
这个“信息摘要”是如何构建的呢？我们设计了三个并行的“注意力分支”，像三个各有所长的专家，协同为Query服务。这三个分支分别是：压缩注意力（Compression）、选择注意力（Selection） 和 滑动窗口注意力（Sliding Window）。
静态示意图：NSA 整体架构
下图描绘了NSA处理信息的流程。当一个查询（Query, qt）到来时，过去的键值对（Keys/Values, k:t, v:t）被分块，然后兵分三路进行处理，最终通过一个门控机制（Gated Output）汇总结果。
1. 压缩 (Compression)：提纲挈领，把握全局
首先是压缩分支。它的任务是从宏观上理解整个文本的脉络。我们把长长的历史信息（Key-Value序列）切成一个个连续的块（block），然后用一个可学习的神经网络（MLP），像一个高效的秘书一样，把每个块的内容压缩成一个“摘要” Key-Value 对。这样一来，原本成千上万的词元就被浓缩成了几百个“章节概要”。
压缩操作：\(\tilde{K}_{t}^{\text{cmp}} = \text{CompressFunc}(\{k_{i \cdot d+1 : i \cdot d + l}\})\)
其中 \(\text{CompressFunc}\) 是一个可学习的MLP，它将一个长度为 \(l\) 的块压缩成一个代表性的向量。
当Query到来时，它只需要和这些“章节概要”进行注意力计算，就能快速了解全文的宏观结构和主题分布，而无需陷入细节的泥潭。这极大地降低了计算量，同时保留了对全局上下文的感知能力。
动画1：令牌压缩 (Token Compression)
这个动画演示了压缩过程。一长串原始令牌（灰色方块）被分成小组，每个小组通过一个“压缩器”被凝聚成一个单一的、信息更密集的压缩令牌（紫色方块）。
状态: 待开始 | 已压缩块数: 0
2. 选择 (Selection)：精读重点，洞察细节
只看“章节概要”显然是不够的，很多关键细节隐藏在原文中。因此，我们需要第二个专家——选择分支。它的任务是找出那些包含最重要信息的“原始文本块”，并进行精读。
那么，如何判断哪些块是重要的呢？这里我们用了一个非常巧妙的“借力”技巧：我们直接利用上一步“压缩注意力”计算出的注意力分数！那些与当前Query高度相关的“章节概要”，它们所对应的原始文本块，很可能就包含了Query最需要的信息。这就像你读一本书的目录，发现某一章的标题特别吸引你，你就会直接翻到那一章去仔细阅读。
我们根据这些分数对所有原始文本块进行排序，选出得分最高的Top-N个块。然后，Query只与这几个被“高亮”出来的块里的所有原始词元进行标准的注意力计算。这样，我们就在不引入大量额外计算的前提下，实现了对关键细节的精准定位和处理。
类比：图书馆的智能检索
想象一下，你在一个巨大的图书馆里找资料（Query）。全注意力就像让你把图书馆里每一本书都从头到尾读一遍。而NSA的策略是：
- 压缩：先快速浏览所有书的摘要卡片（压缩Key），对整个图书馆的藏书有个大概了解。
- 选择：根据你的需求，发现有几张摘要卡片（高分压缩Key）特别相关。于是，你只把这几张卡片对应的原书（被选择的块）从书架上取下来，仔细阅读。
这样既不会错过全局信息，又能高效地找到并深入研究关键内容。
动画2：基于重要性分数的块选择 (Blockwise Selection)
此动画展示了选择机制。首先，查询（橙色圆圈）与压缩令牌（紫色）计算注意力，得到重要性分数。然后，系统根据这些分数选择最重要的原始数据块（绿色高亮），并只在这些块上执行精细的注意力计算。
状态: 待开始 | 已选择块数: 0 / 4
3. 滑动窗口 (Sliding Window)：聚焦当下，捕捉局部
最后，我们还需要一位专家来处理“眼前事”。在语言中，最近的上下文通常具有最高的优先级（比如，一个代词指代的对象往往就在前一句话）。为了确保模型能精确捕捉这种紧密的局部依赖关系，我们设立了第三个分支：滑动窗口。这个分支非常直接，它只关注Query之前的、固定长度（比如512个词元）的上下文。这就像一个人的短期记忆，永远保持着对最近发生事情的清晰印象。
将这三个分支独立出来，可以避免它们之间的“学习干扰”。比如，如果没有专门的滑动窗口，模型可能会偷懒，只学会处理局部信息，而忽略了对全局信息的学习。通过这种“分工合作”的架构，我们鼓励模型在压缩和选择分支中专注于学习更长程的、更复杂的依赖关系。
动画3：滑动窗口注意力 (Sliding Window)
这个动画模拟了滑动窗口。无论查询移动到哪里，注意力（黄色区域）始终只覆盖其紧邻的前一段固定长度的序列，确保对局部上下文的精确捕捉。
查询位置: 0 | 窗口覆盖范围: [0, 0]
最终，三个分支的输出会通过一个可学习的“门控机制”进行加权融合，得到最终的注意力结果。这个门控机制像一个总指挥，根据当前Query的特点，动态地决定更相信哪位“专家”的意见。
最终输出：\(o_{t}^{*} = \sum_{c \in \{\text{cmp, slc, win}\}} g_{t}^{c} \cdot \text{Attn}(q_{t}, \tilde{K}_{t}^{c}, \tilde{V}_{t}^{c})\)
其中 \(g_t^c\) 是一个由MLP计算出的门控分数，决定了每个分支的权重。
硬件协同：让算法与GPU“心有灵犀”
一个算法在理论上再高效，如果不能在实际硬件上跑得快，那也是纸上谈兵。我们设计的核心原则之一就是“硬件协同”（Hardware-aligned）。这意味着NSA的每一步操作都充分考虑了现代GPU（如图形处理单元）的特性。
GPU最喜欢处理的是连续、整齐、大块的数据。我们采用的“块选择”（Blockwise Selection）策略就是为此而生。相比于那些需要从内存中随机、零散地挑选单个词元的方法，我们的方法总是以“块”为单位进行读取和计算。这使得数据传输更流畅，并且能最大化利用GPU中为矩阵运算专门优化的Tensor Cores，从而实现接近理论极限的加速效果。
类比：仓库管理员的智慧
想象一个仓库管理员需要从货架上取100件商品。一个低效的管理员可能会拿着清单，在仓库里来来回回跑100趟，每次只取一件。而一个高效的管理员（就像我们的NSA核函数）会先规划好路线，把清单上位于同一货架区域的商品一次性取完。我们的“块选择”就是这个道理，它把需要访问的数据（Key-Value块）在内存中组织好，然后一次性、成块地加载到GPU的高速缓存（SRAM）中进行计算，大大减少了来回访问主内存（HBM）的次数，从而实现了巨大的速度提升。
我们为NSA专门设计和优化了Triton核函数。Triton是一种能让我们用类似Python的语言编写高性能GPU代码的工具。通过精心设计的循环调度和内存访问模式，我们确保了在训练（前向和反向传播）和推理的各个阶段，NSA都能获得实打实的速度提升。实验结果也证明了这一点：在处理64k长度的序列时，我们的训练速度比FlashAttention-2快了数倍，解码速度更是提升了超过11倍！
动画4：硬件友好的块数据加载 (Hardware-Aligned Loading)
此动画对比了两种数据加载方式。左边是低效的“随机访问”，需要多次、零散地从主内存（HBM）读取数据到高速缓存（SRAM）。右边是NSA采用的“块访问”，一次性加载连续的数据块，显著提高了效率和算术强度。
随机访问加载次数: 0 | 块访问加载次数: 0
实验效果：不仅更快，而且更强
我们进行了一系列详尽的实验来验证NSA的性能。我们基于一个270亿参数的先进模型架构（结合了GQA和MoE），分别训练了使用全注意力的基线模型和使用我们NSA的模型。结果令人振奋。
- 通用能力不降反升：在知识、推理、代码等一系列通用能力评测中，NSA模型在绝大多数任务上都超越了全注意力基线。这说明，从头开始进行稀疏训练，不仅没有损害模型能力，反而可能通过“强迫”模型关注最重要的信息，过滤掉噪声，从而提升了其学习效率和最终性能。
- 长文本理解能力卓越：在著名的“大海捞针”测试中，NSA可以在64k的超长文本中精准地找到并回答隐藏在任意位置的信息，实现了100%的准确率。在更复杂的长文本问答和代码理解基准LongBench上，NSA的平均分也显著高于包括全注意力在内的所有对比方法。
- 复杂推理能力更胜一筹：我们还对模型进行了数学推理能力的“极限挑战”。通过使用大量数学解题过程数据进行微调，我们发现NSA模型在解决复杂的数学竞赛题（AIME）时，表现明显优于全注意力模型。这表明我们的稀疏结构能更有效地捕捉和利用对复杂逻辑推理至关重要的长距离依赖关系。
静态示意图：性能与速度对比
以下图表直观展示了NSA与全注意力（Full Attention）在性能和速度上的对比。左图显示，在多个基准测试上，NSA的得分（红色）普遍高于或持平于全注意力（橙色）。右图显示，在处理64k长序列时，NSA在解码、前向和后向传播阶段的速度提升倍率。
结论与展望
通过NSA，我们证明了一条重要的路径：精心设计的、与硬件协同的、原生可训练的稀疏注意力，不仅可以解决长文本带来的计算效率瓶颈，甚至能够在模型能力上实现超越。它不再是全注意力的“廉价替代品”，而是一种更智能、更高效的全新范式。
我们相信，这只是一个开始。未来的语言模型将需要处理越来越长、越来越复杂的输入，从多模态数据流到持续的交互式对话。像NSA这样的高效架构将是支撑这些未来应用不可或缺的基础设施。我们期待看到更多沿着这条道路的探索，共同推动大型语言模型走向一个更广阔、更高效的新纪元。
动画5：未来的可能性 (Flowing into the Future)
这个最终的动画象征着NSA开启的无限可能性。无数的信息粒子（代表数据和思想）在由高效算法构成的无形力场中有序而优雅地流动，汇聚成智慧的洪流。这代表了我们对未来AI发展的愿景：更高效、更强大、更智能。
附录：技术细节
A. 算术强度与硬件优化
算术强度（Arithmetic Intensity）是衡量一个计算任务是“计算密集型”还是“内存密集型”的关键指标，其定义为计算操作次数与内存访问量的比值。现代GPU拥有极高的浮点运算能力（FLOPS），但其内存带宽（Memory Bandwidth）相对有限。这两者的比值决定了一个“临界算术强度”。
\(\text{Arithmetic Intensity} = \frac{\text{Total Floating Point Operations}}{\text{Total Bytes Accessed from Memory}}\)
当一个任务的算术强度高于此临界值时，我们称之为“计算绑定”（Compute-bound），其瓶颈在于GPU的计算速度。反之，则为“内存绑定”（Memory-bound），瓶颈在于从内存搬运数据的速度。
- 训练/预填充阶段：这个阶段涉及大量的矩阵乘法，算术强度高，是典型的计算绑定任务。优化目标是减少总的计算量。
- 自回归解码阶段：每一步只生成一个token，但需要加载整个KV缓存。计算量小而内存访问量大，算术强度极低，是典型的内存绑定任务。优化目标是减少内存访问量。
NSA的设计充分考虑了这两个阶段的不同特点。我们的块选择策略，通过共享KV块和组中心（Group-Centric）的数据加载，在训练时最大化了算术强度，减少了冗余计算。在解码时，由于我们只需要加载被选择的少数块、压缩块和滑动窗口，极大地减少了内存访问量，从而实现了显著的解码加速。
B. GQA/MQA 架构下的稀疏性挑战
为了提升解码效率，现代LLM广泛采用分组查询注意力（GQA）或多查询注意力（MQA）。在这类架构中，多组查询头（Query heads）共享同一份键值对（KV cache）。这就给稀疏注意力带来了新的挑战。
如果每个查询头都独立地选择自己感兴趣的KV块，那么在解码时，为了服务一个GQA组内的所有查询头，就需要加载这些头所选择的所有块的“并集”。这可能导致实际加载的KV缓存量远大于单个头所需的量，从而削弱了稀疏性带来的内存访问优势。
我们的解决方案是，在GQA组内强制实现“选择一致性”。我们聚合组内所有头的注意力分数来共同决定选择哪些块。这样，所有共享KV的查询头都会访问完全相同的稀疏块集合，避免了“并集”问题，确保了在GQA/MQA架构下依然能获得最大的解码速度提升。
组内共享的重要性分数：\(P^{\text{slc}}_{\text{group}} = \sum_{h=1}^{H} P^{\text{slc},(h)}\)
这确保了组内所有头（\(h=1, \dots, H\)）选择相同的块。