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

本文是LLM系列文章，针对《Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention》的翻译。

摘要

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

1 引言

2 对稀疏注意力方法的再思考

3 方法