Mamba with Hierarchical Memory: Solving Representation Bottleneck in Long Sequence Modeling

2026-08-03

AI Summary

This work addresses the limitation of existing recurrent linear attention models, such as Mamba, which struggle to effectively capture long-sequence semantics due to their fixed-capacity recurrent state. Inspired by human hierarchical memory mechanisms, the authors introduce a lightweight working memory module into the Mamba architecture that extracts and compresses paragraph-level semantics from the backbone hidden states into retrievable long-term memory representations. This enables parameterized, task-adaptive hierarchical semantic processing. The proposed approach endows Mamba with cross-task generalization capabilities for the first time, overcoming its long-context modeling bottleneck. Experimental results demonstrate that, with only a 2% increase in parameters and minimal training overhead, the method improves retrieval success rates by 34.3–37.1% on Passkey Retrieval and LongBench-E tasks and boosts inference accuracy by 1.6–14.2%.

Abstract

Recurrent linear attention models (RLAs) such as Mamba offer efficient linear-time sequence modeling as an alternative to Transformers, yet their fixed-capacity recurrent states limit long-sequence modeling. Drawing inspiration from hierarchical human memory, we propose Hierarchical Memory Mamba (HMM) to address this limitation. Building upon a pre-trained Mamba backbone, HMM integrates a lightweight working memory that extracts slow paragraph-level semantics (PLS) from the fast sensory memory embedded in the backbone's hidden states. The PLS is subsequently compressed into persistent long-term memory for task-relevant retrieval. The hierarchical processing of semantic information overcomes the representation bottleneck of RLAs and endows HMM cross-task generalization through parametric learning, which is not observed in other long-context enhanced Mamba variants. Evaluations on Passkey Retrieval and LongBench-E tasks demonstrate that HMM improves retrieval success by 34.3--37.1% and reasoning accuracy by 1.6--14.2% over strong Mamba-based models, while adding only 2% extra parameters and with minimal training overhead.

Problem

• long sequence modeling
• representation bottleneck
• recurrent linear attention
• hierarchical memory
• Mamba

Innovation

• Hierarchical Memory
• Mamba
• Long Sequence Modeling
• Recurrent Linear Attention
• Cross-task Generalization