State space models improved with new routing and memory designs
Diagonal Equivalence in State Space Models: Characterization, Retrieval Rank, and Selective Extensions
Machine Learning
Summary
This work studies a special type of mathematical model used in AI called state-space models, which help computers understand sequences like language. The authors identify new ways to organize and connect these models, improving how they handle multiple inputs and recall past information. They test their ideas on large GPUs and find that these new designs can train faster and remember more compared to earlier versions. However, some improvements in handling information do not always lead to better overall language understanding.
What this means in practice
- •For machine learning engineers: Develop faster training pipelines for large sequence models by leveraging improved GPU kernel implementations of state-space designs.
- •For natural language processing teams: Use improved state-space models with enhanced multi-input handling to boost small-scale sequence recall in language understanding tasks.
Authors
Dennis Rochau
Abstract
We classify real and complex state-space models (SSMs) with linear transitions, distinguishing fixed-basis diagonal equivalence from input-output equivalence and tracking coefficient-conditioning windows under coordinate changes. We characterize diagonal equivalence through commuting transition families whose members are diagonalizable over the chosen field. With time-shared token-local matrix maps, we identify this class's state-mediated kernels with unnormalized scalar-mode linear attention with multiplicative masks. Mamba 3's multiple-input, multiple-output (MIMO) construction lifts one value per local interface; SSD-MIMO routes independent values for local retrieval rank above one. HybridMamba combines Mamba 1's per-coordinate decay, Mamba 2-style heads and Mamba 3's complex rotations and two-endpoint writes, mixing real and rotating state blocks with different decay laws. They recover Mamba 1-3 single-input, single-output (SISO) configurations as exact special cases. On five large H100 bfloat16 workloads, our fused GPU kernels improve geometric-mean training time over official Mamba 1/3 in matched SISO recoveries, with Mamba 2 near parity at SSM-core scope. Memory tradeoffs depend on the baseline and metric: saved-for-backward versus peak allocation. In small-model multi-token recall experiments, both designs outperform Mamba 1-3 on the hardest tier; higher routing ranks perform best in gated or phase-diverse configurations. In single-seed language models with about 500M parameters, both improve on Mamba 1/2 in test loss at 2B tokens, but the routing gains do not translate into lower language-model loss. Our best configurations approach Mamba 3 SISO through 10B tokens while remaining behind it and Mamba 3 MIMO in language-model loss.