Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

长上下文建模对于下一代语言模型至关重要，但标准注意力机制的高计算成本带来了显著的计算挑战。稀疏注意力为在保持模型能力的同时提高效率提供了一个有前景的方向。我们提出了NSA（原生可训练稀疏注意力机制），它结合了算法创新与硬件对齐优化，以实现高效的长上下文建模。NSA采用动态分层稀疏策略，将粗粒度的标记压缩与细粒度的标记选择相结合，以保持全局上下文感知和局部精度。我们的方法通过两项关键创新推进了稀疏注意力设计：（1）我们通过算术强度平衡的算法设计实现了显著的加速，并针对现代硬件进行了实现优化。（2）我们支持端到端训练，减少了预训练计算量而不牺牲模型性能。如图1所示，实验表明使用NSA预训练的模型在通用基准测试、长上下文任务和基于指令的推理中保持或超过了全注意力模型的表现。同时，NSA在64k长度序列的解码、前向传播和反向传播过程中都比全注意力机制实现了显著的加速，验证了其在整个模型生命周期中的高效性。

该论文旨在解决长上下文建模中标准注意力机制带来的高计算成本问题，这是当前大型语言模型面临的一个重要挑战。虽然稀疏注意力机制提供了一个有希望的方向，但如何在保持模型性能的同时显著提高效率是一个新的研究课题。

论文提出了NSA（Natively trainable Sparse Attention），一种结合了算法创新和硬件优化的稀疏注意力机制。关键思路包括：1) 动态分层稀疏策略，通过粗粒度压缩和细粒度选择来平衡全局上下文感知和局部精度；2) 算术强度均衡的算法设计和硬件优化，以实现显著的速度提升；3) 支持端到端训练，减少预训练计算量而不牺牲模型性能。这与现有方法相比，不仅提高了效率，还确保了模型的有效性和灵活性。

实验结果显示，使用NSA预训练的模型在多个基准测试、长上下文任务和基于指令的推理方面表现优异或超过全注意力模型。此外，NSA在处理64k长度序列时，在解码、前向传播和反向传播过程中均实现了显著加速。该研究还提供了详细的实验设计，并验证了其在整个模型生命周期中的高效性。值得注意的是，论文提到了开源代码的可能性，鼓励社区进一步探索和改进。

近年来，关于稀疏注意力机制的研究逐渐增多。例如，《Reformer: The Efficient Transformer》引入了局部敏感哈希（LSH）注意力来降低计算复杂度；《Performer: Generalized Attention with RFF Kernels》则利用随机傅里叶特征核函数进行快速近似计算。其他相关工作还包括《Longformer: The Long-Document Transformer》和《Big Bird: Transformers for Longer Sequences》等，这些研究都在尝试解决长序列处理中的效率问题。