DeepSeek再发重磅成果！全面降低预训练成本_deepseek减少了预训练-优快云博客

本文链接：https://blog.youkuaiyun.com/Python_cocola/article/details/145747406

2月18日，DeepSeek 重磅推出NSA：一种硬件对齐且原生可训练的稀疏注意力机制，用于超快速长文本训练和推理！****

NSA 的核心组件：

· 动态分层稀疏策略

· 粗粒度 token 压缩

· 细粒度 token 选择

凭借针对现代硬件的优化设计，NSA 加速了推理速度，同时降低了预训练成本——且不影响性能。它在通用基准测试、长文本任务和基于指令的推理方面与全注意力模型相匹配或优于后者。

原生稀疏注意力：硬件对齐且原生可训练的稀疏注意力

Jingyang Yuan*¹²,Huazuo Gao¹, Damai Dai¹, Junyu Luo², Liang Zhao¹, Zhengyan Zhang¹, Zhenda Xie¹, Y. X. Wei¹, Lean Wang¹, Zhiping Xiao³, Yuqing Wang¹, Chong Ruan¹, Ming Zhang2, Wenfeng Liang¹, Wangding Zeng¹

¹DeepSeek-AI

²Key Laboratory for Multimedia Information Processing, Peking University, PKU-Anker LLM Lab

³University of Washington {yuanjy, mzhang_cs} @pku.edu.cn**,**

{zengwangding, wenfeng.liang} @deepseek.com

摘要

长文本建模对于下一代语言模型至关重要，然而，标准注意力机制的高计算成本带来了巨大的算力挑战。稀疏注意力为在提高效率的同时保持模型能力提供了一个有前景的方向。我们提出了 NSA，一种原生可训练的稀疏注意力机制，它融合了算法创新和硬件对齐的优化，以实现高效的长文本建模。

NSA 采用了一种动态的分层稀疏策略，结合了粗粒度的 token 压缩和细粒度的 token 选择，以同时保留全局上下文感知和局部精度。我们的方法通过两项关键创新推进了稀疏注意力设计：（1）我们通过算术强度平衡的算法设计以及针对现代硬件的实现优化，实现了显著的加速。（2）我们实现了端到端训练，在不牺牲模型性能的前提下，减少了预训练的计算量。如图 1 所示，实验表明，使用 NSA 预训练的模型在通用基准、长文本任务和基于指令的推理方面，性能保持或超过了全注意力模型。

与此同时，NSA 在解码、前向传播和反向传播过程中，针对 64k 长度的序列，相比全注意力实现了显著的加速，验证了其在模型生命周期内的效率。

*在DeepSeek-AI实习期间的贡献。

1. 引言

研究界日益认识到长文本建模是下一代大型语言模型的关键能力，这受到了各种实际应用的驱动，包括深入推理 (DeepSeek-AI, 2025; Zelikman et al., 2022)、仓库级代码生成 (Zhang et al., 2023a; Zhang et al.) 和多轮自主代理系统 (Park et al., 2023)。最近的突破，包括 OpenAI 的 o-series 模型、DeepSeek-R1 (DeepSeek-AI, 2025) 和 Gemini 1.5 Pro (Google et al., 2024)，使得模型能够处理整个代码库、冗长的文档，维持数千 token 上下连贯的多轮对话，并执行跨越长距离依赖的复杂推理。然而，随着序列长度的增加，原始注意力 (Vanilla Attention) (Vaswani et al., 2017) 机制的高复杂度 (Zaheer et al., 2020) 成为一个关键的延迟瓶颈。理论估计表明，注意力机制的计算复杂度为在解码 64k 长度的上下文时，基于 Softmax 架构的计算占据了总延迟的 70-80%，这突显了对更高效注意力机制的迫切需求。

图 1 | 性能与效率对比：全注意力模型 vs. 我们的 NSA (原生稀疏注意力)。左图：尽管具有稀疏性，但在通用基准测试、长程上下文任务以及推理评估中，NSA 的平均表现超越了全注意力基线模型。右图：在处理 64k 长度的序列时，NSA 在所有阶段（解码、前向传播和反向传播）均实现了相对于全注意力模型显著的算力加速。

一种实现高效长程上下文建模的自然方法是利用 Softmax 注意力固有的稀疏性 (Ge et al., 2023; Jiang et al., 2023)。通过选择性地计算关键的查询-键对，可以在保持性能的同时显著降低算力开销。近期的进展通过多种策略证明了这种潜力：KV 缓存淘汰方法 (Li et al., 2024; Zhang et al., 2023b; Zhou et al., 2024)，分块 KV 缓存选择方法 (Tang et al., 2024; Xiao et al., 2024)，以及基于采样、聚类或哈希的选择方法 (Chen et al., 2024; Desai et al., 2024; Liu et al., 2024)。尽管这些策略前景广阔，但现有的稀疏注意力方法在实际部署中往往表现不足。许多方法未能实现与其理论收益相当的加速；此外，大多数方法主要关注推理阶段，缺乏有效的训练时支持来充分挖掘注意力的稀疏模式。

为了解决这些局限性，有效稀疏注意力的部署必须应对两个关键挑战：（1）硬件对齐的推理加速：将理论上的计算量减少转化为实际的速度提升，需要在预填充和解码阶段进行硬件友好的算法设计，以缓解内存访问和硬件调度瓶颈；（2）训练感知的算法设计：实现与可训练算子的端到端计算，以降低训练成本，同时保持模型性能。这些要求对于实际应用中实现快速长程上下文推理或训练至关重要。当同时考虑这两个方面时，现有方法仍然存在明显的差距。

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

NSA 引入了两个对应于上述关键要求的核心创新：(1) 硬件对齐的系统：优化分块稀疏注意力，以充分利用 Tensor Core 并高效访问内存，确保均衡的算术密集度。(2) 训练感知的设计：通过高效的算法和反向算子，实现稳定的端到端训练。这种优化使 NSA 能够同时支持高效部署和端到端训练。

图 2 | NSA 架构概述。左图：框架通过三个并行注意力分支处理输入序列：对于给定的查询，先前的键和值被处理成用于粗粒度模式的压缩注意力、用于重要 token 块的选择性注意力和用于局部上下文的滑动注意力。右图：每个分支产生的不同注意力模式的可视化。绿色区域表示需要计算注意力分数的区域，而白色区域表示可以跳过的区域。

我们通过在真实世界语言语料库上进行全面的实验来评估 NSA。在具有 2600 亿 token 的 270 亿参数 Transformer 主干网络上进行预训练后，我们评估了 NSA 在通用语言评估、长程上下文评估和思维链推理评估中的性能。我们进一步比较了在 A100 GPU 上使用优化的 Triton (Tillet et al., 2019) 实现的内核速度。实验结果表明，NSA 实现了与全注意力基线模型相当或更优的性能，同时超越了现有的稀疏注意力方法。此外，与全注意力模型相比，NSA 在解码、前向和反向传播阶段均实现了显著的算力加速，并且加速比率随着序列长度的增加而提高。这些结果验证了我们的分层稀疏注意力设计有效地平衡了模型能力和算力效率。

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

现代稀疏注意力方法在降低 Transformer 模型理论计算复杂度方面取得了显著进展。然而，大多数方法主要在推理 (inference) 阶段应用稀疏性，同时保留预训练的全注意力骨干网络，这可能会引入架构偏差 (architectural bias)，限制了它们充分利用稀疏注意力优势的能力。在介绍我们的原生稀疏架构之前，我们通过两个关键视角系统地分析了这些局限性。

2.1. 高效推理的幻觉

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

· 阶段受限的稀疏性 (Phase-Restricted Sparsity)。诸如 H2O (Zhang et al., 2023b) 等方法在自回归解码 (autoregressive decoding) 期间应用稀疏性，但需要在预填充 (prefilling) 阶段进行计算密集型预处理（例如，注意力图计算、索引构建）。相比之下，像 MInference (Jiang et al., 2024) 这样的方法仅专注于预填充阶段的稀疏性。这些方法未能实现跨所有推理阶段的加速，因为至少有一个阶段的计算成本与全注意力相当。这种阶段专精化降低了这些方法在预填充主导型工作负载（如书籍摘要和代码补全）或解码主导型工作负载（如长链式思考 (chain-of-thought) (Wei et al., 2022) 推理）中的加速能力。

· 与先进注意力架构的不兼容性 (Incompatibility with Advanced Attention Architecture)。一些稀疏注意力方法无法适应现代解码高效架构，如多查询注意力 (Multiple-Query Attention, MQA) (Shazeer, 2019) 和分组查询注意力 (Grouped-Query Attention, GQA) (Ainslie et al., 2023)，这些架构通过在多个查询头之间共享键值 (KV) 对，显著降低了解码期间的内存访问瓶颈。例如，在 Quest (Tang et al., 2024) 等方法中，每个注意力头独立地选择其 KV 缓存子集。虽然这在多头注意力 (Multi-Head Attention, MHA) 模型中展示了一致的计算稀疏性和内存访问稀疏性，但在基于 GQA 等架构的模型中，情况却有所不同，在 GQA 中，KV 缓存的内存访问量对应于同一 GQA 组内所有查询头的选择的并集。这种架构特性意味着，虽然这些方法可以减少计算操作，但所需的 KV 缓存内存访问仍然相对较高。这种局限性迫使人们做出关键选择：虽然一些稀疏注意力方法减少了计算，但其分散的内存访问模式与来自先进架构的高效内存访问设计相冲突。

这些局限性之所以产生，是因为许多现有的稀疏注意力方法侧重于 KV 缓存缩减或理论计算量缩减，但难以在先进框架或后端中实现显著的延迟降低。这促使我们开发能够结合先进架构和硬件高效实现 (hardware-efficient implementation) 的算法，以充分利用稀疏性来提高模型效率。

2.2. 可训练稀疏性的神话

我们对原生可训练稀疏注意力的追求，源于对仅推理方法的分析而得到的两个关键见解：

（1）性能退化 (Performance Degradation)：事后应用稀疏性迫使模型偏离其预训练的优化轨迹。正如 Chen et al. (2024) 所证明的那样，前 20% 的注意力只能覆盖总注意力分数的 70%，使得预训练模型中的检索头 (retrieval heads) 等结构在推理过程中容易受到剪枝的影响。

（2）训练效率需求 (Training Efficiency Demands)：高效处理长序列训练对于现代 LLM 开发至关重要。这包括在更长的文档上进行预训练以增强模型容量，以及后续的适应阶段，如长文本微调 (long-context fine-tuning) 和强化学习。然而，现有的稀疏注意力方法主要针对推理，在很大程度上忽略了训练中的计算挑战。这种局限性阻碍了通过高效训练开发更强大的长文本模型。此外，将现有稀疏注意力方法应用于训练的努力也暴露了挑战：

· 不可训练的组件 (Non-Trainable Components)。诸如 ClusterKV (Liu et al., 2024) (包括 k-means 聚类) 和 MagicPIG (Chen et al., 2024) (包括基于 SimHash 的选择) 等方法中的离散操作，在计算图中创建了不连续性。这些不可训练的组件阻止了梯度流经 token 选择过程，限制了模型学习最佳稀疏模式的能力。

· 低效的反向传播 (Inefficient Back-propagation)。一些理论上可训练的稀疏注意力方法在实际训练中效率低下。诸如 HashAttention (Desai et al., 2024) 等方法中使用的 token 粒度选择策略，导致在注意力计算期间需要从 KV 缓存中加载大量单独的 token。这种非连续的内存访问 (non-contiguous memory access) 阻碍了对诸如 FlashAttention 等快速注意力技术的有效适配，而 FlashAttention 依赖于连续的内存访问 (contiguous memory access) 和分块计算 (blockwise computation) 来实现高吞吐量。因此，实现被迫退回到较低的硬件利用率，显著降低了训练效率。

2.3. 原生稀疏性势在必行

推理效率和训练可行性方面的这些局限性，促使我们对稀疏注意力机制进行根本性的重新设计。我们提出了 NSA，一个原生稀疏注意力框架，旨在解决计算效率和训练需求。在接下来的章节中，我们将详细介绍 NSA 的算法设计和算子实现 (operator implementation)。

3. 方法论

我们的技术方法涵盖算法设计和内核优化 (kernel optimization)。在以下小节中，我们首先介绍我们方法论的背景。然后，我们介绍 NSA 的总体框架，然后是其关键算法组件。最后，我们详细介绍我们硬件优化的内核设计，以最大限度地提高实际效率。

3.1. 背景

注意力机制 (Attention Mechanism) 广泛应用于语言建模中，其中每个查询 token(qt)计算与所有先前键(k:t)的相关性得分，以生成值的加权和(v:t)。形式上，对于长度为 t 的输入序列，注意力操作定义为：

这里，αt,i 表示 qt 和 ki 之间的注意力权重，dk 是键的特征维度。随着序列长度的增加，注意力计算在整体计算成本中变得越来越占主导地位，给长文本处理带来了重大挑战。

算术强度 (Arithmetic Intensity) 是计算操作与内存访问的比率。它从本质上塑造了硬件上的算法优化。每个 GPU 都有一个关键的算术强度，它由其峰值计算能力和内存带宽决定，计算公式为这两者硬件限制的比率。对于计算任务，高于此关键阈值的算术强度将变为计算受限 (compute-bound)（受 GPU FLOPS 限制），而低于此阈值则变为内存受限 (memory-bound)（受内存带宽限制）。

具体来说，对于因果自注意力机制 (causal self-attention mechanism)，在训练和预填充阶段，批量矩阵乘法 (batched matrix multiplications) 和注意力计算表现出高算术强度，使得这些阶段在现代加速器上成为计算受限。相比之下，自回归解码 (auto-regressive decoding) 变得受内存带宽限制，因为它每次前向传播生成一个 token，同时需要加载整个键值缓存，导致算术强度较低。这导致了不同的优化目标 - 减少训练和预填充期间的计算成本，同时减少解码期间的内存访问。

3.2. 总体框架

为了利用具有自然稀疏模式的注意力的潜力，我们建议将原始的键值对 k:t, v:t 在公式 (1) 中替换为更紧凑且信息密集的重映射 (remapped) 键值对 K~t, V~t，它们是针对每个查询 qt 生成的。具体来说，我们正式定义优化后的注意力输出如下：

在这里插入图片描述

其中，K~t, V~t 是基于当前查询 qt 和上下文记忆 k:t, V:t 动态构建的。我们可以设计各种映射策略来获得不同类别的 K~tc, V~tc，并将它们组合如下：

如图 2 所示，NSA 具有三种映射策略 c = {cmp, slc, win}，分别代表压缩 (compression)、选择 (selection) 和滑动窗口 (sliding window) 用于键和值。gtc ∈ ℝ 是对应策略 c 的门控分数 (gate score)，通过 MLP 和 sigmoid 激活函数从输入特征中导出。设 Nt 表示重映射键/值的总数：

我们通过确保 Nt ≪ t 来保持高稀疏性比率。

3.3. 算法设计

在本小节中，我们介绍我们的重映射策略 fk 和 fv 的设计：token 压缩、token 选择和滑动窗口。

3.3.1. Token 压缩

通过将连续的键或值块聚合为块级表示，我们获得了捕获整个块信息的压缩键和值。形式上，压缩键表示定义为：

在这里插入图片描述

其中 l 是块长度，d 是相邻块之间的滑动步长，φ 是一个可学习的 MLP，带有块内位置编码，用于将一个块中的键映射到单个压缩键。K~tcmp ∈ ℝdk × ⌊dt/l⌋ 是由压缩键组成的张量。通常，我们采用 d < l 以减轻信息碎片化。压缩值表示 V~tcmp 具有类似的公式。压缩表示捕获了更粗粒度、更高层次的语义信息，并减轻了注意力的计算负担。

3.3.2. Token 选择

仅使用压缩的键值对可能会丢失重要的细粒度信息，这促使我们有选择地保留单个的键值对。下面我们描述我们高效的 token 选择机制，它可以以低计算开销识别并保留最相关的 token。

分块选择 (Blockwise Selection)。我们的选择策略在空间连续的块中处理键值序列，这主要是基于两个关键因素：硬件效率考虑和注意力分数的固有分布模式。分块选择对于在现代 GPU 上实现高效计算至关重要。这是因为，与基于随机索引的读取相比，现代 GPU 架构对于连续的块访问表现出明显更高的吞吐量。此外，分块计算能够实现张量核心 (Tensor Cores) 的最佳利用。这种架构特性已将分块内存访问和计算确立为高性能注意力实现的基本原则，例如 FlashAttention 的基于块的设计就体现了这一点。分块选择也符合注意力分数的固有分布模式。先前的工作 (Jiang et al., 2024) 表明，注意力分数通常表现出空间连续性，这表明相邻的键往往具有相似的重要性级别。

我们在第 6.2 节中的可视化也显示了这种空间连续模式。为了实现分块选择，我们首先将键值序列划分为选择块。为了识别用于注意力计算的最重要块，我们需要为每个块分配重要性分数。下面我们介绍我们计算这些块级重要性分数的方法。

重要性分数计算 (Importance Score Computation)。计算块重要性分数可能会引入显著的开销。幸运的是，压缩 token 的注意力计算会产生中间注意力分数，我们可以利用这些分数来推导选择块的重要性分数，公式如下：

其中上标 (h) 表示头索引，H 是每个组中查询头的数量。这种聚合确保了同一组内各头之间块选择的一致性。

Top-n 块选择 (Top-n Block Selection)。在获得选择块重要性分数后，我们保留块重要性分数排名前 n 的稀疏块内的 token，公式如下：

3.3.3. 滑动窗口 (Sliding Window)

在注意力机制中，局部模式通常适应得更快，并且可能在学习过程中占据主导地位，从而可能阻止模型有效地从压缩和选择的 token 中学习。为了解决这个问题，我们引入了一个专用的滑动窗口分支，显式地处理局部上下文，允许其他分支（压缩和选择）专注于学习各自的特征，而不会被局部模式 “短路 (shortcutted)”。具体来说，我们在一个大小为 w 的窗口中维护最近的 token K~twin = kt-w:t, V~twin = vt-w:t，并将不同信息源（压缩 token、选择 token 和滑动窗口）的注意力计算隔离到不同的分支中。然后，这些分支的输出通过学习到的门控机制进行聚合。为了进一步防止跨注意力分支的 “短路学习”，同时保持边际计算开销，我们为三个分支提供独立的键和值后，我们按照公式 (5) 计算最终的注意力输出。结合上述的压缩、选择和滑动窗口机制，就构成了 NSA 的完整算法框架。

3.4. 内核设计

为了在训练和预填充阶段达到 FlashAttention 级别的加速，我们基于 Triton 实现了硬件对齐的稀疏注意力内核。考虑到多头注意力（MHA）在解码时内存密集且效率低下，我们遵循当前最先进的大型语言模型（LLM）的发展趋势，专注于具有共享 KV 缓存的架构，如分组查询注意力（GQA）和多查询注意力（MQA）。虽然压缩和滑动窗口注意力计算与现有的 FlashAttention-2 内核兼容，但我们为稀疏选择注意力引入了专门的内核设计。

如果我们遵循 FlashAttention 将时间上连续的查询块加载到 SRAM 中的策略，那将导致低效的内存访问，因为一个块内的查询可能需要不相交的 KV 块。为了解决这个问题，我们的关键优化在于不同的查询分组策略：对于查询序列上的每个位置，我们将 GQA 组内的所有查询头（它们共享相同的稀疏 KV 块）加载到 SRAM 中。图 3 说明了我们的前向传播实现。所提出的内核架构具有以下关键特征：1.组中心数据加载2.共享 KV 获取3.网格上的外循环。

这种设计通过以下方式实现了接近最佳的算术强度：（1）通过组级别共享消除冗余的 KV 传输，以及（2）平衡 GPU 流式多处理器之间的计算工作负载。

图3 | NSA的内核设计。该内核按GQA组加载查询（网格循环），获取对应的稀疏KV块（内部循环），并在SRAM上执行注意力计算。绿色块表示SRAM上的数据，而蓝色表示HBM上的数据。

4. 实验

我们通过三个方面评估了原生稀疏注意力（NSA）：（1）通用基准测试性能，（2）长文本基准测试性能，以及（3）链式思考推理性能，并与全注意力基线和最先进的稀疏注意力方法进行了比较。我们将稀疏计算范式的效率分析推迟到第 5 节，在那里我们将详细讨论训练和推理速度。

4.1. 预训练设置

遵循当前最先进 LLM 的常用实践，我们的实验采用结合了分组查询注意力（GQA）和混合专家（MoE）的骨干网络，总参数量为 270 亿，其中 30 亿为活跃参数。该模型由 30 层组成，隐藏层维度为 2560。对于 GQA，我们将组数设置为 4，总共有 64 个注意力头。对于每个头，查询、键和值的隐藏维度配置为dq= dk= 192，dv= 128。对于 MoE，我们利用 DeepSeekMoE (Dai et al., 2024; DeepSeek-AI, 2024) 结构，包含 72 个路由专家和 2 个共享专家，并将 top-k 专家设置为 6。为了确保训练稳定性，第一层中的 MoE 由 SwiGLU 形式的 MLP 替代。

图4 | 在27B参数模型上，Full Attention与我们的NSA在预训练过程中的损失对比。两种模型都表现出稳定的收敛，且NSA的损失值较低。

表1 | Full Attention基准模型与NSA在通用基准测试中的预训练性能对比，包括知识类（MMLU、MMLU-PRO、CMMLU）、推理类（BBH、GSM8K、MATH、DROP）和编程类（MBPP、HumanEval）任务。尽管稀疏性较高，NSA在大多数基准测试中仍表现出优越的平均性能。

所提架构在计算成本与模型性能之间实现了有效的权衡。对于NSA，我们设置了压缩块大小 𝑙 = 32，滑动步长 d= 16，选择的块大小l′ = 64，选择的块数量 n= 16（包括固定激活的1个初始块和2个局部块），以及滑动窗口大小w = 512。Full Attention和稀疏注意力模型均在270B tokens的8k长度文本上进行了预训练，随后在32k长度文本上使用YaRN（Peng et al., 2024）进行继续训练和监督微调，以实现长序列适应。两种模型均训练至完全收敛，以确保公平比较。如图4所示，NSA和Full Attention基准模型的预训练损失曲线均表现出稳定和平滑的下降，且NSA始终优于Full Attention模型。

4.2. 基线方法

除了与全注意力进行比较外，我们还评估了几种最先进的推理阶段稀疏注意力方法：H2O (Zhang et al., 2023b)、infLLM (Xiao et al., 2024)、Quest (Tang et al., 2024) 和 Exact-Top，它首先计算完整的注意力分数，然后选择每个查询对应的 top-n 分数键，并基于这些位置计算注意力。这些方法涵盖了不同的稀疏注意力范式，包括 KV 缓存驱逐、查询感知选择和精确 top-n 稀疏选择。对于通用评估，由于大多数样本的长度都在稀疏注意力基线的局部上下文窗口内，这些方法实际上等同于全注意力。因此，在这种设置下，我们仅展示 NSA 和全注意力基线之间的比较结果。在长文本评估中，我们对所有基线方法进行比较，并将所有稀疏注意力方法的稀疏度设置为相同，以确保公平比较。对于需要长文本监督微调的链式思考推理评估，由于稀疏注意力基线不支持训练，我们将比较限制为全注意力。

表 2 | 我们的 NSA 和基线模型在 LongBench 基准测试上的性能比较，包括单文档问答（QA）、多文档问答（QA）、合成任务和代码任务子类。NSA 的性能优于包括全注意力机制在内的大多数基线模型。

4.3. 性能比较

· 通用评估。我们在涵盖知识、推理和编码能力的综合基准测试套件上评估了预训练的 NSA 和全注意力基线，这些基准测试包括 MMLU (Hendrycks et al., 2020)、MMLU-PRO (Wang et al., 2024)、CMMLU (Li et al., 2023)、BBH (Suzgun et al., 2022)、GSM8K (Cobbe et al., 2021)、MATH (Hendrycks et al., 2020)、DROP (Dua et al., 2019)、MBPP (Austin et al., 2021) 和 HumanEval (Chen et al., 2021)。结果如表 1 所示。尽管具有稀疏性，NSA 仍实现了出色的整体表现，在 9 项指标中的 7 项上优于包括全注意力在内的所有基线。这表明，虽然 NSA 可能未能在较短的序列上充分利用其效率优势，但它表现出了强大的性能。值得注意的是，NSA 在与推理相关的基准测试中表现出显著的提升（DROP：+0.042，GSM8K：+0.034），这表明我们的预训练有助于模型发展专门的注意力机制。这种稀疏注意力预训练机制迫使模型专注于最重要的信息，可能通过过滤掉来自不相关注意力路径的噪声来提高性能。跨多个评估的一致性能也验证了 NSA 作为通用架构的鲁棒性。

· 长文本评估。如图 5 所示，NSA 在 64k 上下文的“大海捞针”测试中，在所有位置都实现了完美的检索准确率 (Kamradt, 2023)。这种性能源于我们的分层稀疏注意力设计，它结合了用于高效全局上下文扫描的压缩 token 和用于精确局部信息检索的选择 token。粗粒度的压缩以低计算成本识别相关的上下文块，而对选定 token 的 token 级注意力则确保了关键细粒度信息的保留。这种设计使 NSA 能够同时保持全局感知和局部精度。我们还在 LongBench (Bai et al., 2023) 上针对最先进的稀疏注意力方法和全注意力基线评估了 NSA。为了确保一致的稀疏性，我们将所有稀疏注意力基线中每个查询激活的 token 数量设置为 2560 个 token，这对应于 NSA 在处理 32k 序列长度时激活的平均 token 数量。

遵循 Stream-LLM (Xiao et al., 2023)，此 token 预算包括前导的 128 个 token 和 512 个局部 token。由于某些子集在所有模型上的得分都很低，可能无法提供有意义的比较，因此我们从 LongBench 中排除了某些子集。如表 2 所示，NSA 实现了 0.469 的最高平均分，优于所有基线（比全注意力高 +0.032，比 Exact-Top 高 +0.046）。

这种改进源于两项关键创新：（1）我们的原生稀疏注意力设计，实现了预训练期间稀疏模式的端到端优化，促进了稀疏注意力模块与其他模型组件之间的同步适应；以及（2）分层稀疏注意力机制实现了局部和全局信息处理之间的平衡。

图5 | 在64k上下文长度下，针在干草堆中的检索准确率。NSA通过其分层稀疏注意力设计实现了完美的准确率。

值得注意的是，NSA 在需要对长文本进行复杂推理的任务上表现出卓越的性能，在多跳 QA 任务（HPQ 和 2Wiki）上比全注意力提高了 +0.087 和 +0.051，在代码理解方面超过了基线的性能（LCC：+0.069），并在段落检索方面优于其他方法（PassR-en：+0.075）。这些结果验证了 NSA 处理各种长文本挑战的能力，其原生预训练的稀疏注意力在学习任务最优模式方面提供了额外的优势。

· 链式思考推理评估。为了评估 NSA 与高级下游训练范式的兼容性，我们研究了它通过后训练获得链式思考数学推理能力的能力。鉴于强化学习在较小规模模型上的效果有限，我们采用了来自 DeepSeek-R1 的知识蒸馏，使用 10B token 的 32k 长度数学推理轨迹进行监督微调（SFT）。这产生了两个可比较的模型：Full Attention-R（全注意力基线）和 NSA-R（我们的稀疏变体）。

表3 | 在监督微调后的AIME指令评估。我们的NSA-R在8k和16k序列长度下均表现出优于Full Attention-R的性能。

图6 | 基于Triton的NSA核心与基于Triton的FlashAttention-2核心的比较。我们的实现显著减少了所有上下文长度下的延迟，随着输入长度的增加，改进效果愈加明显。

我们在极具挑战性的美国邀请数学考试 (AIME 24) 基准测试中评估了这两个模型。我们使用 0.7 的采样温度和 0.95 的 top-p 值为每个问题生成 16 个答案，并获得平均分。为了验证推理深度的影响，我们进行了两个生成上下文限制的实验：8k 和 16k 个 token，测量扩展的推理链是否提高了准确性。模型预测的示例比较在附录 A 中提供。

5. 效率分析

我们在 8-GPU A100 系统上评估了 NSA 相对于全注意力的计算效率。在效率分析中，我们还将模型配置为：GQA 组数 g=4，每组头数 h=16，查询/键维度 dk=192，值维度 dv =128。遵循第 4 节中的相同设置，我们将 NSA 压缩块大小设置为 l=32，滑动步长 d=16，选择块大小 l′=64，选择块计数 n=16，滑动窗口大小 w=512。

5.1. 训练速度

我们将基于 Triton 实现的 NSA 注意力和全注意力与基于 Triton 的 FlashAttention-2 进行了比较，以确保在同一后端进行公平的速度比较。如图 6 所示，随着上下文长度的增加，我们的 NSA 实现了逐渐更大的加速，在 64k 上下文长度下，前向传播加速高达 9.0 倍，反向传播加速高达 6.0 倍。值得注意的是，随着序列长度的增加，速度优势变得更加明显。这种加速源于我们对硬件友好的算法设计，以最大限度地提高稀疏注意力架构的效率：（1）分块内存访问模式通过合并加载最大限度地提高了 Tensor Core 的利用率，（2）内核中精细的循环调度消除了冗余的 KV 传输。

表4 | 解码过程中每次注意力操作的内存访问量（以等效token数量表示）。由于解码的算术强度较低且受内存限制，预期的加速与内存访问量呈线性关系。

5.2. 解码速度

注意力的解码速度主要由内存访问瓶颈决定，这与 KV 缓存加载量密切相关。在每个解码步骤中，我们的 NSA 只需要加载最多 ⌊ds/l⌋ 个压缩 token、nl′ 个选择 token 和 w 个邻居 token，其中 s 是缓存的序列长度。如表 4 所示，随着解码长度的增加，我们的方法显著减少了延迟，在 64k 上下文长度下实现了高达 11.6 倍的加速。这种内存访问效率的优势也随着序列的加长而放大。

6. 讨论

在本节中，我们将回顾 NSA 的开发过程，并讨论从我们对不同稀疏注意力策略的探索中获得的关键见解。虽然我们的方法展示了有希望的结果，但了解替代策略遇到的挑战和分析注意力模式为未来的研究方向提供了宝贵的背景。我们首先考察了促使我们做出设计选择的替代 token 选择策略的挑战，然后通过可视化深入了解注意力分布模式。

6.1. 替代 Token 选择策略的挑战

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

· 基于键聚类的策略。我们研究了基于聚类的策略，如 ClusterKV (Liu et al., 2024)。这些方法将来自同一集群的键和值存储在连续的内存区域中。虽然理论上对训练和推理都是可行的，但它们面临三个重大挑战：（1）动态聚类机制引入了不小的计算开销；（2）算子优化困难，因集群间不平衡而加剧，尤其是在混合专家 (MoE) 系统中，其中倾斜的专家并行 (EP) 组执行时间导致持续的负载不平衡；（3）实现约束，源于对强制性周期性重聚类和块顺序训练协议的需求。这些综合因素造成了巨大的瓶颈，极大地限制了它们在实际部署中的有效性。

图 7 | 比较使用全注意力机制和不同 Token 选择策略的 3B 参数模型的训练损失。我们的 NSA 实现了更好的性能。注意力分数块模式的位置可视化

图 8 | 全注意力变压器上的注意力映射可视化。浅色区域表示较高的注意力值。如图所示，注意力分数显示出分块聚类分布。

· 其他分块选择策略。我们还考虑了与 NSA 不同的分块键值选择策略，例如 Quest (Tang et al., 2024) 和 InfLLM (Xiao et al., 2024)。这些方法依赖于计算每个块的重要性分数，并根据它们与 qt 的相似性选择 top-n 个块。然而，现有方法面临两个关键问题：（1）由于选择操作是不可微的，因此基于神经网络的重要性分数计算依赖于辅助损失，这增加了算子开销，并且通常会降低模型性能；（2）启发式无参数重要性分数计算策略的召回率较低，导致性能欠佳。我们使用架构类似的 30 亿参数模型评估了这两种方法，并将它们的损失曲线与 NSA 和全注意力进行了比较。对于基于辅助损失的选择方法，我们为每个块引入了额外的查询和代表性键，以估计块重要性分数。这些分数由原始查询与每个块内键之间的平均注意力分数进行监督。对于启发式无参数选择方法，遵循 Quest 的策略，我们使用查询和键块的坐标逐个最小值-最大值的乘积来实现直接选择，而无需引入额外的参数。我们还探索了一种冷启动训练方法，其中在最初的 1000 步中应用全注意力，然后再过渡到启发式分块选择。如图 7 所示，这两种方法都表现出较差的损失。

6.2. 可视化 (Visualization)

为了探索 Transformer 注意力分布中的潜在模式，并为我们的设计寻求灵感，我们在图 8 中可视化了我们预训练的 27B 全注意力模型的注意力图。可视化揭示了有趣的模式，其中注意力分数倾向于表现出分块聚类特征，附近的键通常显示相似的注意力分数。这一观察结果启发了我们的 NSA 设计，表明基于空间连续性选择键块可能是一种有前景的方法。分块聚类现象表明，序列中相邻的 token 可能与查询 token 共享某些语义关系，尽管这些关系的具体性质还需要进一步研究。这一观察结果促使我们探索一种在连续 token 块而不是单个 token 上运行的稀疏注意力机制，旨在提高计算效率并保留高注意力模式。

7. 相关工作 (Related Works)

我们回顾了现有的通过稀疏注意力提高注意力计算效率的方法。这些方法可以根据其核心策略大致分为三类：（1）固定稀疏模式，（2）动态 token 剪枝，和（3）查询感知选择。我们介绍了每个类别中的几个代表性工作。

7.1. 固定稀疏模式 (Fixed Sparse Pattern)

滑动窗口 (Sliding Window) 是一种常用的方法，它允许查询仅在固定窗口内计算注意力。StreamingLLM (Xiao et al., 2023) 通过维护上下文中两个关键部分来解决处理长文本流的挑战：一个注意力汇 (attention sink)（早期 token）和一个局部上下文窗口。虽然这些方法有效地降低了内存和计算成本，但它们忽略上下文的刚性模式限制了它们在需要完全理解上下文的任务上的性能。

7.2. 动态 Token 剪枝 (Dynamic Token Pruning)

H2O (Zhang et al., 2023b) 实现了一种自适应方法，以减少解码期间的 KV 缓存内存使用量。该方法根据 token 最近的效用（根据注意力分数），动态地驱逐 (evicts) 被认为对未来预测不太重要的 token。SnapKV (Li et al., 2024) 也引入了一种 token 剪枝策略，通过有选择地仅保留最关键的特征来减少 KV 缓存，从而实现高效的内存使用。SnapKV 在预填充期间通过注意力权重分析和投票来识别重要特征，然后通过将选定的压缩特征与最近的上下文相结合来更新 KV 缓存，以保持 prompt 的一致性。

7.3. 查询感知选择 (Query-Aware Selection)

Quest (Tang et al., 2024) 采用了一种分块选择策略，其中每个块的重要性通过查询和键块的坐标逐个最小值-最大值的乘积来估计。结果分数有助于选择 top-n 个重要的键值块以进行注意力计算。InfLLM (Xiao et al., 2024) 将固定模式与检索相结合，通过维护注意力汇、局部上下文和可检索块。该方法从每个块中选择代表性键来估计块的重要性。HashAttention (Desai et al., 2024) 通过使用学习到的函数将查询和键映射到 Hamming 空间，将关键 token 识别表述为推荐问题。ClusterKV (Liu et al., 2024) 通过首先聚类键，然后根据查询-集群相似性选择最相关的集群进行注意力计算来实现稀疏性。

8. 结论 (Conclusion)

我们提出了 NSA，一种硬件对齐的稀疏注意力架构，用于高效的长文本建模。通过在可训练架构中集成分层 token 压缩和分块 token 选择，我们的架构实现了加速的训练和推理，同时保持了全注意力的性能。NSA 通过以下方面推进了最先进水平：在通用基准测试中，性能与全注意力基线相当；在长文本评估中，超越了建模能力；以及增强了推理能力，所有这些都伴随着可衡量的计算延迟降低和显著的加速。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述