Diagonal Equivalence in State Space Models: New Insights into Retrieval Rank and Selective Scaling

Diagonal Equivalence in State Space Models: New Insights into Retrieval Rank and Selective Scaling

The field of State Space Models (SSMs) has recently seen a theoretical work that systematically examines the diagonal equivalence classification of linear transition SSMs, and on this basis proposes hybrid architectures and GPU kernel optimizations. The research attempts to answer a fundamental question: under coordinate transformations, which structural properties are truly invariant?

A Rigorous Characterization of Diagonal Equivalence

The paper divides linear transition SSMs over the real and complex domains into two classes of equivalence relations: diagonal equivalence under a fixed basis and input-output equivalence, and traces the window of coefficient conditions under coordinate transformations. Diagonal equivalence is characterized by a family of commuting transition matrices—matrices that are diagonalizable over the chosen domain. The value of this framework lies in elevating "whether models are equivalent" from parameter-level comparison to the structural level, providing a mathematical basis for correspondences between different SSM variants.

Connection to Linear Attention

Under time-shared token-local matrix mappings, the study identifies that the state-mediated kernel of such models corresponds to unnormalized scalar-mode linear attention with multiplicative masks. This provides a more precise language for bridging SSMs and Transformer-family attention mechanisms.

Linear Attention is an efficient alternative to standard softmax attention: it replaces the softmax kernel with a decomposable similarity function such as φ(q)·φ(k), thereby exploiting the associativity of matrix multiplication to rewrite the originally O(L²) sequence modeling into an incrementally recurrent linear-complexity form. The so-called "unnormalized scalar mode" means that attention weights are not normalized through softmax, and each channel or head maintains only a scalar state. "With multiplicative masks" corresponds to causality or block-local locality constraints, implemented by multiplying a mask matrix into the kernel function. This type of attention is mathematically equivalent to a linear recurrent state space model: the state is updated token by token, and the output is jointly determined by the state and the current input. Therefore, the state-mediated kernel of an SSM can be viewed as a parameterized special case of linear attention; the two share the same computational skeleton, differing only in state update rules and decay methods.

From MIMO to HybridMamba

The multi-input multi-output (MIMO) construction of Mamba 3 raises one value per local interface; SSD-MIMO routes independent values for cases where the local retrieval rank is greater than one. Retrieval rank becomes a key dimension distinguishing different designs.

HybridMamba combines the per-coordinate decay of Mamba 1, the head structure in the style of Mamba 2, and the complex rotations and dual-endpoint writes of Mamba 3, mixing real-valued state blocks with rotated state blocks having different decay laws. The researchers emphasize that the single-input single-output (SISO) configurations of Mamba 1-3 can be recovered as exact special cases—an important yardstick for measuring the generality of the new architecture.

Retrieval rank refers to the number of independent value channels that a model can write to and read from in parallel at a local interface, which can be understood as the "width" of state updates or the number of parallel memory slots. Under the SISO (single-input single-output) configuration, each interface carries only one value, and state transitions degenerate into per-coordinate scalar decay; while the MIMO (multi-input multi-output) configuration allows multiple values to be written at once, theoretically improving the model's ability to capture associations among multiple tokens. The key design point of SSD-MIMO is that when the local retrieval rank is greater than one, mutually independent values are routed for each rank, avoiding information aliasing between channels. Retrieval rank therefore becomes a key structural dimension distinguishing variants such as Mamba 1, Mamba 2, Mamba 3, and HybridMamba—it affects both expressive power and, directly, state size and memory footprint.

Kernel Performance and Memory Trade-offs

On five large-scale H100 bfloat16 workloads, the authors' fused GPU kernel, under the matched SISO recovery configuration, improves the geometric mean of training time relative to the official Mamba 1/3; Mamba 2 performs nearly on par within the SSM core scope. The memory trade-off depends on the choice of baseline and the measurement convention: the savings used for backpropagation and peak allocation are two different perspectives and cannot be generalized.

The geometric mean is commonly used to aggregate relative speedups across multiple benchmark tests; compared with the arithmetic mean, it better suppresses the distortion of overall conclusions by individual outlier tasks, making it suitable for measuring the comprehensive performance of a kernel on heterogeneous workloads. A fused GPU kernel refers to merging multiple operators such as state updates, gating, and output projection into a single kernel launch to reduce memory round-trips and launch overhead, which is a common means of accelerating SSM training. bfloat16 is a 16-bit floating-point format with a dynamic range similar to float32 but lower precision; it is widely used in large-model training to balance throughput and numerical stability. The so-called "backpropagation savings" and "peak allocation" measure, respectively, the reduction in memory usage throughout training and the instantaneous maximum memory demand; the two may yield opposite conclusions depending on the baseline model and statistical convention, so memory trade-offs cannot be discussed apart from a specific baseline.

Recall Experiments and Language Model Performance

In small-model multi-token recall experiments, both designs outperform Mamba 1-3 at the hardest level; higher routing rank performs best when gating or phase diversity configurations are present. However, on a single-seed language model of about 500 million parameters and 2B tokens, although both outperform Mamba 1/2 on test loss, the gains brought by routing do not translate into lower language model loss. The best configuration approaches Mamba 3 SISO in 10B-token training, but still lags behind the language model losses of Mamba 3 SISO and Mamba 3 MIMO.

This suggests a noteworthy gap: improvements in retrieval capability and language modeling loss do not transfer linearly. Whether the advantage on recall tasks can be realized in larger-scale pretraining remains an open question.

Editorial Commentary

The contribution of this research lies in connecting the abstract concept of diagonal equivalence with concrete retrieval rank design and kernel implementation into a single line, and using Mamba 1-3 as recoverable special cases to anchor the position of the new architecture. However, the lag in language model loss indicates that the boundary of gains from the routing mechanism is not yet clear. For readers interested in the evolution of SSM architectures and efficient training kernels, this work provides a reference with both theoretical depth and engineering detail.

One-line Summary

The paper unifies linear SSMs with a diagonal equivalence framework, proposes a hybrid architecture and fused kernel, but retrieval gains do not translate into language loss advantages.

Paragraph Summary

This research theoretically characterizes the diagonal equivalence classification of linear transition State Space Models (SSMs), uses families of commuting transition matrices to distinguish diagonal equivalence under a fixed basis from input-output equivalence, and reveals the correspondence between such models and unnormalized linear attention with multiplicative masks. At the architecture level, the paper analyzes the MIMO construction of Mamba 3 and the retrieval rank design of SSD-MIMO, and proposes HybridMamba that fuses the characteristics of Mamba 1-3, capable of recovering the SISO configurations of all three as exact special cases. Engineering-wise, the fused GPU kernel achieves geometric-mean training-time improvements relative to Mamba 1/3 on five H100 bfloat16 workloads, while memory trade-offs vary by baseline. Experiments show that increased routing rank is effective on small-model recall tasks, but does not translate into lower language model loss on a language model of about 500 million parameters.