Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access
This addresses the limitation of RNNs in long-context modeling for applications requiring efficient memory access, though it is incremental as it builds on existing RNN and attention mechanisms.
The paper tackled the problem of enabling Recurrent Neural Networks (RNNs) to randomly access historical context without sacrificing efficiency, by proposing Hierarchical Sparse Attention (HSA) and RAMba, which achieved perfect accuracy in passkey retrieval across 64 million contexts with pre-training on only 4K-length contexts and improvements on downstream tasks with nearly constant memory footprint.
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.