MIMOMamba: From Scalar Duality to Matrix-Valued Attention

A selective state space model updates a latent state h_t with h_t = A_th_t−1 + B_tx_t, y_t = C_th_t. In its most general form this is a MIMO system, with A_t a full D×D matrix that can rotate and mix all D feature dimensions at once. The ¹ framework behind Mamba-2 gives this recurrence a dual, parallel attention form, but only after collapsing A_t to a scalar-identity matrix a_tI — the constraint that lets the recurrence commute across time and be reordered into an attention-like product. The price is that each of the D channels then evolves independently, pushing all cross-channel mixing onto a separate feed-forward layer after the SSM block². Linear-attention variants⁴⁵⁶⁷⁸ take a related but different shortcut, leaning on the associative property of matrix multiplication rather than commutativity. MIMOMamba's question is whether the commutativity route can be walked without shrinking the state matrices down to scalars.

A direct extension of SSD to MIMO faces a dichotomy: drop commutativity and lose the duality, or keep the scalar-identity constraint and lose all cross-dimensional interaction. MIMOMamba only needs A_t, B_t, and C_t for a given head to commute with each other, and realizes this by drawing all three from a matrix polynomial of a single, shared, time-invariant base matrix A: each is a linear combination of I, A, A², …, A^p with input-dependent coefficients. Matrix powers commute automatically (AⁱA^j = A^jAⁱ), so this construction supplies commutativity for free while letting every state matrix act as a full D×D transformation.

This restriction costs nothing in practice. The Cayley–Hamilton theorem bounds the polynomial degree at D−1, keeping the parameterization compact, and the centralizer theorem shows the set of matrices commuting with A equals exactly the polynomials of A whenever A is non-derogatory — true almost surely under random initialization. MIMOMamba's polynomial family is therefore, with probability 1, the entire space of commutativity-compatible transformations rather than a proper subset of it.

Unrolling the MIMO recurrence expresses output y_i as a sum over past inputs x_j, each weighted by C_iA_iA_i−1…A_j+1B_j. Because these matrices all commute, this product reorders into A_i:j+1·(C_iB_j), separating a state-dynamics term that depends only on the distance i−j from a projection term depending only on the absolute positions. Applied to every entry, this turns the MIMO recurrence into a block-wise Hadamard product Y = (L ○ (QK^⊤)) ▹ X, where Q stacks the C_t matrices, K the B_t matrices, and L is a block-lower-triangular causal mask. Setting D = 1 recovers exactly the scalar SSD attention form in Mamba-2, so the SISO case is a special case rather than being replaced.

The key difference from ordinary attention is that every entry of this L×L grid is a D×D matrix block, not a scalar: each off-diagonal block C_iA_i:j+1B_j rotates and mixes a past input across feature dimensions rather than just decaying it by a scalar weight, and because the ordered matrix products A_i:j+1 enforce sequential order structurally, the mask needs no external positional encoding. The factorization also preserves the efficiency SSMs are built on: the global transformation stays block D-semiseparable, the same property Mamba-2 exploits¹¹, and the polynomial parameterization turns the D×D matrix-product chain needed for parallel training into scalar polynomial multiplications evaluable by FFT-based convolution, compressing what must be communicated across GPUs from O(D²) down to O(p) coefficients per head, in the spirit of FlashAttention¹² and Monarch matrices¹³.

A separate constraint is the state expansion ratio N/D: larger expansion is known to matter for information-dense sequences⁹, but MIMOMamba's per-head state matches D for parameter efficiency, which would otherwise cap the ratio at N = D. Following multi-head attention³¹⁰, MIMOMamba instead runs H parallel MIMO SSM heads, each with its own base matrix and polynomial coefficients, giving an effective joint dimension of H·D without inflating per-head computation. The H outputs are concatenated and combined by a linear merge layer^a, letting heads specialize in different frequency bands, timescales, or channel-interaction patterns.

The benchmark generates sequences from a linear dynamical system whose state transition matrix is block-diagonal 2×2 rotation blocks, each rotating a pair of feature dimensions at its own frequency and decaying for stability, a direct test of whether a model can represent rotation-like mixing inside its own state recurrence rather than relegating it to a downstream feed-forward layer. MIMOMamba is compared against Gated DeltaNet¹⁴, Mamba-2¹, Mamba-3¹⁵, a standard Transformer³, PD-SSM¹⁶, and an LSTM¹⁷, using RMSE (lower is better) and R² (higher is better).

MIMOMamba improves over Mamba-2 on every metric (RMSE mean 0.687 vs. 0.717; R² 0.480 vs. 0.459), isolating the benefit of the MIMO formulation over the SISO duality it generalizes, and it beats every other baseline on every metric while using roughly a third of the Transformer's parameters at matched settings^b.

Ablations vary the three knobs the polynomial parameterization introduces: polynomial degree p, number of heads H, and the rank of the low-rank base-matrix factorization.

Together these support degree 4, 4 heads, and rank 4 as the default: each knob shows a clear, non-monotonic optimum rather than accuracy simply tracking parameter count.

A matrix polynomial restores the commutativity that scalar Mamba gets from an identity constraint, so channels can mix inside the state instead of staying decoupled. Deriving which cross-channel mixing patterns that algebra can express, and which still need attention, remains open.