面向高效长期内存访问的硬件对齐层次稀疏注意力机制

循环神经网络 (Recurrent Neural Networks, RNNs) 与 Transformer 相比的一个关键优势是其线性计算和空间复杂度，使其在处理长序列时具有更快的训练和推理速度。然而，RNNs 本质上无法随机访问历史上下文，仅仅集成注意力机制可能会削弱其效率优势。为克服这一限制，我们提出层次稀疏注意力 (Hierarchical Sparse Attention, HSA)，一种新颖的注意力机制，为 RNNs 引入长距离随机访问灵活性，同时保持其效率和长度泛化优势。HSA 将输入划分为若干块，选择 top-k 个块并层次化地聚合信息。核心创新在于基于每个块内部的细粒度 token 级信息学习 token 到块的相关性。这一方法在 both in-domain 和 out-of-domain context 长度上显著提升块选择精度。为提高效率，我们进一步引入硬件对齐的 kernel 设计。将 HSA 与 Mamba 结合，我们提出 RAMba，该模型在仅进行 4K 长度预训练的情况下，即可在 6400 万长度的上下文中实现 passkey 检索的 perfect accuracy，并在各种下游任务上取得显著提升，内存占用几乎保持恒定。这些结果表明 RAMba 在长上下文建模方面具有巨大的潜力。