让RNN具备长序列随机访问能力，同时保持高效性。
Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access
- 分块选择+层级聚合，动态聚焦关键上下文信息
- 在6400万上下文下实现100%准确的密钥检索
- 适合需要超长上下文建模的高效模型应用
与Transformer相比，循环神经网络（RNN）具有线性计算和空间复杂度，能更快地训练和推理长序列。然而，RNN无法随机访问历史上下文，简单引入注意力机制会损害其效率优势。为此，我们提出层次稀疏注意力（HSA），一种增强RNN长程随机访问能力的新机制，同时保留其高效性和长度泛化能力。HSA将输入分块，选择前k个相关块并进行层级聚合。核心创新在于基于每块内的细粒度标记级信息学习标记到块的相关性，从而提升跨域内外上下文长度的块选择精度。为提升效率，我们进一步设计了硬件对齐的核函数。结合Mamba，我们提出RAMba，在仅用4K长度上下文预训练的情况下，实现了6400万上下文中的完全准确密钥检索，并在多个下游任务中取得显著提升，内存占用几乎恒定。结果表明RAMba在长上下文建模中具有巨大潜力。
原文摘要 · Abstract (English)
A key advantage of Recurrent Neural Networks (RNNs) over Transformers is their linear computational and space complexity enables faster training and inference for long sequences. However, RNNs are fundamentally unable to randomly access historical context, and simply integrating attention mechanisms may undermine their efficiency advantages. To overcome this limitation, we propose Hierarchical Sparse Attention (HSA), a novel attention mechanism that enhances RNNs with long-range random access flexibility while preserving their merits in efficiency and length generalization. HSA divides inputs into chunks, selects the top-$k$ chunks and hierarchically aggregates information. The core innovation lies in learning token-to-chunk relevance based on fine-grained token-level information inside each chunk. This approach enhances the precision of chunk selection across both in-domain and out-of-domain context lengths. To make HSA efficient, we further introduce a hardware-aligned kernel design. By combining HSA with Mamba, we introduce RAMba, which achieves perfect accuracy in passkey retrieval across 64 million contexts despite pre-training on only 4K-length contexts, and significant improvements on various downstream tasks, with nearly constant memory footprint. These results show RAMba's huge potential in long-context modeling.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。