ICML2026
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.