DeepSeek新作-Native Sparse Attention

NSA概述

长文本建模的重要性与挑战

• 长文本建模的重要性：长文本建模对于下一代语言模型至关重要。这意味着模型需要能够处理和理解长篇幅的文本内容，例如长篇文章、书籍、复杂的对话等，这对于语言模型的推理、生成和理解能力提出了更高的要求。

• 标准注意力机制的挑战：传统的注意力机制（如Transformer中的全注意力机制）在处理长文本时面临巨大的计算成本。这是因为全注意力机制需要计算每个词与其他所有词之间的关系，计算量与文本长度的平方成正比。例如，对于长度为64,000（64k）的文本序列，计算量会非常庞大，导致训练和推理速度缓慢。

稀疏注意力的解决方案

• 稀疏注意力的概念：稀疏注意力是一种改进的注意力机制，它通过减少注意力计算的范围或频率来降低计算成本，同时尽量保持模型的性能。例如，它可能只计算部分词之间的关系，而不是所有词之间的关系。

• NSA（Native Sparse Attention）的提出：文章提出了一种新的稀疏注意力机制——NSA。它通过算法创新和硬件优化来实现高效的长文本建模。NSA的核心思想是结合动态分层稀疏策略，既保留全局上下文信息，又保持局部细节的精确性。

NSA的核心技术特点

• 动态分层稀疏策略：

  • 粗粒度标记压缩：通过减少需要处理的标记（token）数量来降低计算量。例如，将长文本划分为较大的块（block），并对每个块进行压缩处理。

  • 细粒度标记选择：在压缩的基础上，选择关键的标记进行更精细的处理，以保留重要的局部信息。

  • 全局与局部的平衡：这种策略既保留了全局上下文感知能力（通过粗粒度压缩），又保持了局部细节的精确性（通过细粒度选择）。

• 关键创新点：

  • 算术强度平衡的算法设计：通过优化算法，使其更适合现代硬件（如GPU、TPU）的计算特性，从而实现显著的加速。

  • 端到端训练：传统的稀疏注意力机制可能需要复杂的预处理或额外的预训练步骤，但NSA可以直接进行端到端训练，减少了预训练的计算量，同时不牺牲模型性能。

实验结果与效率验证

• 性能表现：实验表明，使用NSA预训练的模型在多个任务中表现优异，包括通用基准测试、长文本任务和基于指令的推理任务。这意味着NSA在保持高效计算的同时，没有损失模型的性能。

• 效率提升：在处理64k长度的序列时，NSA在解码、前向传播和反向传播三个阶段都显著优于全注意力机制。这验证了NSA在整个模型生命周期（从训练到推理）中的高效性。

总结

• NSA是一种新型的稀疏注意力机制，旨在解决长文本建模中的高计算成本问题。

• 它通过动态分层稀疏策略和硬件优化，在保持模型性能的同时，显著提高了计算效率。

• 实验表明，NSA在长文本任务中表现出色，同时在训练和推理阶段都实现了显著的加速。

这种技术的提出为下一代语言模型的长文本处理提供了一个新的方向。

1 Introduction

1. 长文本建模的重要性

• 背景：长文本建模是下一代大型语言模型的关键能力，因为它能够支持多种现实世界的应用场景，例如：

  • 深入推理：例如DeepSeek-AI（2025）和Zelikman等人的研究（2022）。

  • 仓库级代码生成：例如Zhang等人的研究（2023a）。

  • 多轮自主代理系统：例如Park等人的研究（2023）。

• 最近的突破：一些最新的模型（如OpenAI的o系列模型、DeepSeek-R1、Gemini 1.5 Pro）已经能够处理整个代码库、长篇文档、维持数千标记的连贯对话，以及执行复杂的长距离依赖推理。

2. 传统注意力机制的瓶颈

• 问题：传统的自注意力机制（如Transformer中的注意力机制）在处理长文本时面临高复杂性问题，这成为了延迟的瓶颈。

• 具体表现：当解码长度为64k的上下文时，基于softmax的注意力计算占据了总延迟的70%到80%。这表明需要更高效的注意力机制。

3. 稀疏注意力的潜力与挑战

• 稀疏注意力的潜力：稀疏注意力通过选择性地计算关键的查询-键对，可以在保留性能的同时显著减少计算开销。最近的研究通过多种策略展示了这种潜力，例如：

  • 键值缓存驱逐方法。

  • 分块键值缓存选择方法。

  • 基于采样、聚类或哈希的选择方法。

• 现有方法的不足：

  • 许多方法未能实现理论上的加速效果。

  • 大多数方法主要关注推理阶段，缺乏训练阶段的支持，无法充分利用注意力的稀疏模式。

4. 稀疏注意力的核心挑战

• 硬件对齐的推理加速：将理论计算减少转化为实际速度提升，需要设计硬件友好的算法，以缓解内存访问和硬件调度瓶颈。

• 训练感知的算法设计：需要实现端到端计算，通过可训练的算子减少训练成本，同时保持模型性能。

5. NSA（Native Sparse Attention）的提出

设计思路：NSA是一种可原生训练的稀疏注意力架构，通过层次化标记建模来减少每个查询的计算量。

• 层次化处理：将键和值组织成时间块，并通过三条注意力路径处理：

  • 压缩的粗粒度标记：减少计算量。

  • 选择性保留的细粒度标记：保留关键信息。

  • 滑动窗口：用于局部上下文信息。

• 专用内核实现：通过优化内核实现最大化实际效率。

• 核心创新：

  • 硬件对齐系统：优化分块稀疏注意力，以利用Tensor Core和内存访问，确保算术强度平衡。

  • 训练感知设计：通过高效的算法和反向算子实现稳定的端到端训练。

这张图展示了NSA（Native Sparse Attention）架构的概览。旨在提高处理输入序列的效率。图中分为左右两部分：

左侧：

• 框架结构：NSA的框架通过三个并行的注意力分支处理输入序列。

  • 压缩注意力（Compressed Attention）：处理粗粒度模式，适用于给定查询的前导键和值。

  • 选择性注意力（Selected Attention）：关注重要的标记块。

  • 滑动注意力（Sliding Attention）：用于处理局部上下文。

这些分支通过不同的注意力模式来处理输入数据，从而提高效率。

右侧：

• 注意力模式的可视化：

  • 压缩注意力掩码（Compressed Attention Mask）：显示哪些区域需要计算注意力分数（绿色区域），哪些区域可以跳过（白色区域）。

  • 选择性注意力掩码（Selected Attention Mask）：同样显示需要计算和可以跳过的区域。

  • 滑动注意力掩码（Sliding Attention Mask）：也展示了注意力计算的区域分布。

总的来说，这张图说明了NSA架构如何通过三种不同的注意力机制来提高处理效率，并通过可视化展示了这些机制在计算注意力分数时的区域分布。绿色区域表示需要计算注意力分数的区域，而白色区域表示可以跳过的区域。

6. NSA的实验验证

• 实验设置：在27B参数的Transformer骨干网络上进行预训练，使用260B标记。

• 评估内容：在通用语言评估、长文本评估和推理链评估方面验证NSA的性能。

• 实验结果：

  • NSA在性能上与全注意力基线相当或更优。

  • 比现有的稀疏注意力方法表现更好。

  • 在解码、前向和反向阶段实现了显著加速，且加速比随着序列长度的增加而增加。

2 Rethinking Sparse Attention Methods

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

2.1 The Illusion of Efficient Inference

稀疏注意力方法旨在通过减少注意力计算中的冗余操作来降低Transformer模型的计算复杂度和推理延迟。然而，尽管这些方法在理论计算复杂度上取得了进展，但在实际推理延迟的减少上却面临挑战。

两大挑战

（1）阶段限制的稀疏性（Phase-Restricted Sparsity）

• 问题描述：许多稀疏注意力方法只在推理的某个阶段（如预填充或解码）应用稀疏性，而其他阶段仍然保持较高的计算复杂度。

  • 预填充阶段：例如，H2O方法在自回归解码阶段应用稀疏性，但在预填充阶段需要进行计算密集型的预处理（如注意力图计算和索引构建）。

  • 解码阶段：MInference方法则只关注预填充阶段的稀疏性，而解码阶段的计算复杂度仍然较高。

• 影响：由于这些方法未能在所有推理阶段实现稀疏性，因此在某些工作负载下（如需要大量预填充的任务，如书籍摘要；或需要大量解码的任务，如长链式推理）无法实现显著的加速。

（2）与先进注意力架构的不兼容性（Incompatibility with Advanced Attention Architectures）

• 问题描述：一些稀疏注意力方法无法适应现代高效的解码架构，如多查询注意力（MQA）和分组查询注意力（GQA）。这些架构通过在多个查询头之间共享键值（KV）对，显著减少了内存访问瓶颈。

  • MQA和GQA的优势：在这些架构中，KV对被共享，从而减少了每个查询头的内存访问量。

  • 稀疏注意力方法的局限性：例如，Quest方法在多头注意力（MHA）模型中表现出稀疏性，但在GQA架构中，KV缓存的内