Diagonal Equivalence in State Space Models: Characterization, Retrieval Rank, and Selective Extensions

📄 arXiv: 2611.03101v1 📥 PDF

作者: Dennis Rochau

分类: cs.LG, math.OC, stat.ML

发布日期: 2026-11-02

备注: 81 pages, 6 figures, 33 tables. Code: https://github.com/d-rochau/OrbitSSM

💡 一句话要点

提出对角等价的状态空间模型以提升多输入多输出性能

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 状态空间模型 对角等价 多输入多输出 Mamba架构 机器学习优化 复杂系统建模 线性注意力 参数优化

📋 核心要点

  1. 现有的状态空间模型在处理多输入多输出时存在性能瓶颈，特别是在对角等价的表征上。
  2. 本文提出了一种新的分类方法，通过交换的转移族来表征对角等价，并结合多种Mamba架构进行优化。
  3. 实验结果表明，新的设计在多个大型工作负载上显著提高了训练时间，并在多令牌回忆任务中表现优异。

📝 摘要（中文）

本文对具有线性转移的实数和复数状态空间模型（SSMs）进行了分类，区分了固定基对角等价与输入输出等价，并在坐标变换下跟踪系数条件窗口。通过交换的转移族，我们对对角等价进行了表征，识别了该类的状态介导核与未归一化的标量模式线性注意力相结合。通过Mamba 3的多输入多输出（MIMO）构建，提升了局部接口的值，并在小模型多令牌回忆实验中超越了Mamba 1-3。实验结果显示，在大型H100 bfloat16工作负载上，融合的GPU内核在训练时间上显著优于官方Mamba 1/3。

🔬 方法详解

问题定义：本文旨在解决现有状态空间模型在多输入多输出（MIMO）配置下的性能不足，特别是在对角等价的表征和优化上存在的挑战。现有方法在处理复杂转移时效率较低，无法充分利用状态介导核。

核心思路：论文提出通过引入交换的转移族来表征对角等价，并结合Mamba架构的多种变体，优化状态空间模型的性能。通过时间共享的局部矩阵映射，提升了模型在多输入多输出场景下的表现。

技术框架：整体架构包括对角等价的分类、状态介导核的识别以及Mamba架构的多种配置。主要模块包括Mamba 1、Mamba 2和Mamba 3，分别针对不同的输入输出配置进行优化。

关键创新：最重要的创新在于通过对角等价的表征和状态介导核的结合，提升了模型在多输入多输出场景下的性能。这一方法与现有的输入输出等价方法有本质区别，能够更有效地处理复杂的状态转移。

关键设计：在设计中，使用了未归一化的标量模式线性注意力和乘法掩码，优化了参数设置和损失函数，确保在不同的配置下都能实现最佳性能。

📊 实验亮点

实验结果显示，在五个大型H100 bfloat16工作负载上，融合的GPU内核在训练时间上相较于官方Mamba 1/3有显著提升，尤其是在多令牌回忆任务中，新的设计在最难的层级上超越了Mamba 1-3，显示出更高的路由排名在门控或相位多样化配置中的优势。

🎯 应用场景

该研究的潜在应用领域包括复杂系统建模、控制理论以及机器学习中的多模态数据处理。通过提升状态空间模型的性能，能够在机器人控制、自动驾驶和智能系统等领域实现更高效的决策和控制，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.