Random Long-Context Access for Mamba via Hardware-aligned Hierarchical Sparse Attention

• URL: http://arxiv.org/abs/2504.16795v1
• Date: Wed, 23 Apr 2025 15:15:06 GMT
• Title: Random Long-Context Access for Mamba via Hardware-aligned Hierarchical Sparse Attention
• Authors: Xiang Hu, Jiaqi Leng, Jun Zhao, Kewei Tu, Wei Wu,
• Abstract summary: We propose textbfHierarchical textbfSparse textbfAttention (HSA), a novel attention mechanism that enhances RNNs with long-range random access flexibility.<n>HSA divides inputs into chunks, selecting the top-$k$ chunks and hierarchically aggregates information.<n>By combining HSA with Mamba, we introduce RAMba, which achieves perfect accuracy in passkey retrieval across 64 million contexts.
• Score: 43.3704626107852
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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 \textbf{H}ierarchical \textbf{S}parse \textbf{A}ttention (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, selecting 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.

Related papers

• Modality Agnostic Efficient Long Range Encoder [14.705955027331674]
  We address the challenge of long-context processing on a single device using generic implementations.<n>To overcome these limitations, we propose MAELRE, a unified and efficient transformer architecture.<n>We demonstrate that MAELRE achieves superior accuracy while reducing computational cost compared to existing long-context models.
  arXiv  Detail & Related papers  (2025-07-25T16:19:47Z)
• RAT: Bridging RNN Efficiency and Attention Accuracy in Language Modeling [17.437929000395112]
  We introduce an intermediate design called rat between recurrence and attention mechanisms.<n>It partitions the input into chunks, applies a simple linear recurrence within each chunk to capture local dependencies, and then performs softmax attention across chunks to model long-range interactions.<n>With a chunk size of 16, the rat layer achieves a (7times) improvement in training speed with 100K token sequences and (9times) in generation at 4K sequence length.
  arXiv  Detail & Related papers  (2025-07-06T15:08:49Z)
• Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers [58.98923344096319]
  REFORM is a novel inference framework that efficiently handles long contexts through a two-phase approach.<n>It achieves over 50% and 27% performance gains on RULER and BABILong respectively at 1M context length.<n>It also outperforms baselines on Infinite-Bench and MM-NIAH, demonstrating flexibility across diverse tasks and domains.
  arXiv  Detail & Related papers  (2025-06-01T23:49:14Z)
• FLAMES: A Hybrid Spiking-State Space Model for Adaptive Memory Retention in Event-Based Learning [16.60622265961373]
  FLAMES is a hybrid framework integrating structured state-space dynamics with event-driven computation.<n>By bridging neuromorphic computing and structured sequence modeling, FLAMES enables scalable long-range reasoning in event-driven systems.
  arXiv  Detail & Related papers  (2025-04-02T00:08:19Z)
• Optimal Gradient Checkpointing for Sparse and Recurrent Architectures using Off-Chip Memory [0.8321953606016751]
  We introduce memory-efficient gradient checkpointing strategies tailored for the general class of sparse RNNs and Spiking Neural Networks.<n>We find that Double Checkpointing emerges as the most effective method, optimizing the use of local memory resources while minimizing recomputation overhead.
  arXiv  Detail & Related papers  (2024-12-16T14:23:31Z)
• SparseAccelerate: Efficient Long-Context Inference for Mid-Range GPUs [0.0]
  We introduce SparseAccelerate, a dynamic sparse attention method that adapts its sparsity patterns based on input characteristics.<n> Experimental results show that SparseAccelerate achieves up to a 1.04x reduction in Time-To-First-Token (TTTF) latency at 32K tokens.
  arXiv  Detail & Related papers  (2024-12-09T04:27:03Z)
• Squeezed Attention: Accelerating Long Context Length LLM Inference [64.11145320159126]
  We propose Squeezed Attention as a mechanism to accelerate LLM applications where a large portion of the input prompt is fixed. We use K-means clustering offline to group the keys for the fixed context based on semantic similarity and represent each cluster with a single centroid value. We then compute exact attention using only these important keys from the fixed context, thereby reducing bandwidth and computational costs.
  arXiv  Detail & Related papers  (2024-11-14T18:54:19Z)
• Taipan: Efficient and Expressive State Space Language Models with Selective Attention [100.16383527459429]
  Long-context language modeling is a significant challenge in Natural Language Processing (NLP) Recent State Space Models (SSMs) such as Mamba offer alternatives with constant memory usage, but they underperform in tasks requiring extensive in-context retrieval. We introduce Taipan, a novel hybrid architecture that combines Mamba-2 with Selective Attention Layers (SALs) Our experiments demonstrate Taipan's superior performance across various scales and tasks, offering a promising solution for efficient long-context language modeling.
  arXiv  Detail & Related papers  (2024-10-24T09:25:37Z)
• Efficient Length-Generalizable Attention via Causal Retrieval for Long-Context Language Modeling [42.67141329779589]
  Grouped Cross Attention can generalize to 1000 times the pre-training context length.<n>Experiments show that GCA-based models achieve near-perfect accuracy in passkey retrieval for 16M context lengths.
  arXiv  Detail & Related papers  (2024-10-02T15:18:34Z)
• Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers [58.5711048151424]
  We introduce SPARSEK Attention, a novel sparse attention mechanism designed to overcome computational and memory obstacles. Our approach integrates a scoring network and a differentiable top-k mask operator, SPARSEK, to select a constant number of KV pairs for each query. Experimental results reveal that SPARSEK Attention outperforms previous sparse attention methods.
  arXiv  Detail & Related papers  (2024-06-24T15:55:59Z)
• Training-Free Exponential Context Extension via Cascading KV Cache [49.608367376911694]
  We introduce a novel mechanism that leverages cascading sub-cache buffers to selectively retain the most relevant tokens.<n>Our method reduces prefill stage latency by a factor of 6.8 when compared to flash attention on 1M tokens.
  arXiv  Detail & Related papers  (2024-06-24T03:59:17Z)
• Hierarchical Context Merging: Better Long Context Understanding for Pre-trained LLMs [61.40047491337793]
  We present Hierarchical cOntext MERging (HOMER), a new training-free scheme designed to overcome the limitations of large language models. HomeR uses a divide-and-conquer algorithm, dividing long inputs into manageable chunks. A token reduction technique precedes each merging, ensuring memory usage efficiency.
  arXiv  Detail & Related papers  (2024-04-16T06:34:08Z)
• Heterogenous Memory Augmented Neural Networks [84.29338268789684]
  We introduce a novel heterogeneous memory augmentation approach for neural networks. By introducing learnable memory tokens with attention mechanism, we can effectively boost performance without huge computational overhead. We show our approach on various image and graph-based tasks under both in-distribution (ID) and out-of-distribution (OOD) conditions.
  arXiv  Detail & Related papers  (2023-10-17T01:05:28Z)
• Landmark Attention: Random-Access Infinite Context Length for Transformers [45.69864961773124]
  We present a novel approach that allows access to the complete context while retaining random-access flexibility. Our method uses a landmark token to represent each block of the input and trains the attention to use it for selecting relevant blocks. We demonstrate that our method can obtain comparable performance with Transformer-XL while significantly reducing the number of retrieved tokens in each step.
  arXiv  Detail & Related papers  (2023-05-25T17:53:42Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.