本地稀疏注意力:硬件对齐和本地可训练稀疏注意力Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse
Attention
长上下文建模对于下一代语言模型至关重要,然而标准注意力机制的高计算成本带来了显著的计算挑战。稀疏注意力为提高效率同时保持模型能力提供了一个有前途的方向。我们提出了NSA,一种本地可训练的稀疏注意力机制,它将算法创新与硬件对齐的优化相结合,实现了高效的长上下文建模。NSA采用动态分层稀疏策略,将粗粒度的标记压缩与细粒度的标记选择相结合,以保留全局上下文意识和局部精度。我们的方法通过两个关键创新推进了稀疏注意力设计:(1) 我们通过算术强度平衡的算法设计实现了显著的加速,同时针对现代硬件进行了实现优化。(2) 我们实现了端到端训练,减少了预训练计算而不损害模型性能。如图1所示,实验证明,使用NSA预训练的模型在一般基准、长上下文任务和基于指令的推理方面保持或超越了全注意力模型。与此同时,NSA在64k长度序列的解码、前向传播和反向传播中实现了显著的加速,验证了其在整个模型生命周期中的高效性。