MIMOMamba: From Scalar Duality to Matrix-Valued Attention

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 \textbf{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.

Modern AI sequence models face a dilemma: Transformers are powerful but become prohibitively expensive on long inputs, while recent fast alternatives such as Mamba gain speed by treating each feature dimension as an isolated one‑dimensional channel and outsourcing cross‑feature mixing to a separate layer. When features interact in ways that change over time — as in physical simulations, audio, or sensor signals — this separation discards useful structure that the model could otherwise exploit. We propose MIMOMamba, an architecture that lets the model directly capture how features interact over time without giving up speed. The technical obstacle is that richer interactions normally break the algebraic property (matrix commutativity) that makes Mamba's fast computation possible. Our key idea is simple: require every time‑varying state matrix to be a polynomial of a single shared base matrix. Because polynomials of the same matrix automatically commute, the fast "dual" computation is restored, and we prove the restriction loses no expressivity within the relevant solution space under standard random initialization. On an ocean‑physics forecasting task with strongly coupled depth layers, MIMOMamba beats a standard Transformer using only one‑third of its parameters, and surpasses recent state‑space models on a time‑series classification benchmark. The result points to efficient long‑sequence models that learn coupled dynamics structurally, rather than by brute force.