MIMOMamba: From Scalar Duality to Matrix-Valued Attention
Yanbo Li, Richard Cornelius Suwandi, Feng Yin, Yiyong SUN, Wei Huang, Wenqiang Pu
Abstract
The state space duality (SSD) framework, central to modern state-space models (SSMs) such as Mamba, has established an efficient attention-like mechanism by leveraging the commutative property of linear recurrences. However, existing formulations are limited to single-input single-output (SISO) systems that enforce commutativity with a restrictive scalar-identity constraint, which prevents cross-dimensional interactions within the state dynamics. In this work, we generalize SSD to the multi-input multi-output (MIMO) setting by introducing a matrix polynomial parameterization. This approach not only provides a principled way to ensure commutativity for generalized duality but also induces a shared algebraic structure across state transitions, thereby significantly reducing parameter redundancy. Building on this foundation, we present MIMOMamba, a multi-head SSM architecture that captures rich cross-dimensional dynamics while retaining linear-time training. Empirical evaluations on a sequence modeling benchmark demonstrate that MIMOMamba achieves performance comparable to or surpassing that of standard Transformers while utilizing substantially fewer parameters.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on15
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra et al.NeurIPS 2022 · 5,493 citations
- Efficiently Modeling Long Sequences with Structured State SpacesAlbert Gu, Karan Goel, Christopher RéICLR 2022 · 3,482 citations
- Transformers are RNNs: Fast Autoregressive Transformers with Linear AttentionAngelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, François FleuretICML 2020 · 2,665 citations
- FlashAttention-2: Faster Attention with Better Parallelism and Work PartitioningTri DaoICLR 2024 · 2,600 citations
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space DualityTri Dao, Albert GuICML 2024 · 1,407 citations
Related papers
- On Structured State-Space DualityJerry Yao-Chieh Hu, Xiwen Zhang, Ali ElSheikh, Weimin Wu et al.ICML 2026
- Achilles' Heel of Mamba: Essential difficulties of the Mamba architecture demonstrated by synthetic dataTianyi Chen, Pengxiao Lin, Zhiwei Wang, Zhi-Qin John XuNeurIPS 2025 · 4 citations
- SSDi8: Accurate and Efficient 8-bit Quantization for State Space DualityHyunwoo Kim, Byoungchan Ko, Minseok Kang, Minwoo Kim et al.ICLR 2026 · 3 citations
- Efficient Unstructured Pruning of Mamba State-Space Models for Resource-Constrained EnvironmentsIbne Farabi Shihab, Sanjeda Akter, Anuj SharmaEMNLP 2025 · 1 citation
- VSSD: Vision Mamba With Non-Causal State Space DualityYuheng Shi, Mingjia Li, Minjing Dong, Chang XuICCV 2025 · 20 citations
