这篇论文介绍了一种名为NSA（Natively Sparse Attention）的新型稀疏注意力机制，旨在提高长上下文建模的效率，同时保持模型的性能。
1. 研究背景：问题：长上下文建模对于下一代语言模型至关重要，但标准的注意力机制由于其高计算成本而带来了显著的挑战。难点：现有的稀疏注意力方法在实际部署中往往未能实现预期的加速效果，主要因为它们在推理阶段有效，但在训练阶段缺乏支持。相关工作：现有的稀疏注意力方法包括KV缓存驱逐方法、基于块的选择方法和采样、聚类或哈希选择方法，但这些方法在实际应用中存在局限性。
2. 研究方法：动态层次稀疏策略：NSA结合了粗粒度的令牌压缩和细粒度的令牌选择，以保持全局上下文意识和局部精度。具体来说，NSA通过将键和值组织成时间块，并通过三个注意力路径处理：压缩的粗粒度令牌、选择性保留的细粒度令牌和滑动窗口的局部上下文信息。算法设计：通过算术强度平衡的算法设计和现代硬件优化的实现，NSA在推理阶段实现了显著加速。NSA引入了两种关键创新：算术强度平衡的算法设计和端到端训练能力。硬件优化：NSA采用了专门的内核设计，以最大化其在实际应用中的效率。具体实现包括组中心数据加载、共享KV获取和外循环在网格上的调度。
3. 实验设计：预训练设置：NSA和全注意力模型在270B个8k长度文本上进行预训练，然后在32k长度文本上进行持续训练和监督微调。实验模型包括27B总参数，30层，每层2560的隐藏维度。基准测试：在一般基准测试中，NSA在9个指标中有7个超过全注意力模型，特别是在推理相关的基准测试中表现出显著优势（如DROP：+0.042，GSM8K：+0.034）。长上下文评估：在64k上下文的needle-in-a-haystack测试中，NSA实现了完美的检索准确性。在LongBench评估中，NSA的平均得分为0.469，超过全注意力模型（+0.032）和Exact-Top（+0.046）。
4. 结果与分析：训练速度：在8-GPU A100系统上，NSA在64k上下文长度下实现了9.0倍的加速（前向传播）和6.0倍的加速（反向传播）。解码速度：NSA在64k上下文长度下实现了11.6倍的加速，主要由于内存访问效率的提高。
5. 总体结论：NSA通过层次化的令牌压缩和块选择的稀疏注意力机制，在保持全注意力模型性能的同时，实现了显著的加速。NSA在一般基准测试、长上下文评估和推理能力方面均表现出色，验证了其在实际应用中的有效性和效率。