状态空间模型的对角等价：检索秩与选择性扩展新解

状态空间模型（SSM）领域近期出现一项理论工作，系统梳理了线性转移 SSM 的对角等价分类，并在此基础上提出混合架构与 GPU 内核优化。该研究试图回答一个基础问题：在坐标变换下，哪些结构性质是真正不变的？

对角等价的严格刻画

论文将实数和复数域上的线性转移 SSM 分为两类等价关系：固定基下的对角等价与输入-输出等价，并追踪坐标变换下系数条件的窗口。对角等价通过一族可交换的转移矩阵来刻画——这些矩阵在选定域上可对角化。这一框架的价值在于，它把"模型是否等价"从参数层面的比较提升到结构层面，为不同 SSM 变体之间的对应关系提供了数学依据。

与线性注意力的联系

在时间共享的 token 局部矩阵映射下，该研究识别出此类模型的状态中介核与带乘法掩码的未归一化标量模式线性注意力相对应。这为 SSM 与 Transformer 系注意力机制之间的桥接提供了更精确的语言。

线性注意力（Linear Attention）是标准 softmax 注意力的高效替代形式：它把 softmax 核替换为 φ(q)·φ(k) 这类可分解的相似度函数，从而利用矩阵乘法的结合律，将原本 O(L²) 的序列建模改写为可增量递推的线性复杂度形式。所谓「未归一化标量模式」，指注意力权重不经过 softmax 归一化，且每个通道或头只维护一个标量状态。「带乘法掩码」则对应因果性或分块局部性约束，通过在核函数中乘入掩码矩阵实现。这类注意力在数学上等价于一个线性递推状态空间模型：状态随 token 逐位更新，输出由状态与当前输入共同决定。因此 SSM 的状态中介核可以视为线性注意力的一种参数化特例，二者共享同一计算骨架，只是状态更新规则与衰减方式不同。

从 MIMO 到 HybridMamba

Mamba 3 的多输入多输出（MIMO）构造在每个局部接口上提升一个值；而 SSD-MIMO 为局部检索秩大于一的情况路由独立的值。检索秩成为区分不同设计的关键维度。

HybridMamba 则把 Mamba 1 的逐坐标衰减、Mamba 2 风格的头结构，以及 Mamba 3 的复数旋转与双端点写入组合在一起，混合了实数状态块与不同衰减律的旋转状态块。研究者强调，Mamba 1-3 的单输入单输出（SISO）配置可作为其精确特例被恢复——这是衡量新架构通用性的重要标尺。

检索秩（retrieval rank）指模型在局部接口上能够并行写入和读取的独立值通道数，可理解为状态更新的「宽度」或「并行记忆槽位」数量。SISO（单输入单输出）配置下每个接口只承载一个值，状态转移退化为逐坐标的标量衰减；而 MIMO（多输入多输出）配置允许一次写入多个值，理论上能提升模型对多 token 关联的捕获能力。SSD-MIMO 的设计要点在于：当局部检索秩大于一时，为每个秩路由相互独立的值，避免通道间信息混叠。检索秩因此成为区分 Mamba 1、Mamba 2、Mamba 3 及 HybridMamba 等变体的关键结构维度——它既影响表达能力，也直接影响状态大小与内存占用。

内核性能与内存权衡

在五个 H100 bfloat16 大规模工作负载上，作者实现的融合 GPU 内核在匹配的 SISO 恢复配置下，相对官方 Mamba 1/3 提升了训练时间的几何平均值；Mamba 2 在 SSM 核心范围内的表现接近持平。内存方面的权衡取决于基线选择与度量口径：用于反向传播的节省量与峰值分配是两个不同视角，不能一概而论。

几何平均值（geometric mean）常用于聚合多个基准测试的相对加速比，相比算术平均更能抑制个别离群任务对总体结论的拉偏，适合衡量内核在异构工作负载上的综合表现。融合 GPU 内核（fused kernel）指把状态更新、门控、输出投影等多个算子合并进单一 kernel 启动，以减少显存往返与启动开销，这是 SSM 训练加速的常见手段。bfloat16 是 16 位浮点格式，动态范围与 float32 相近但精度较低，在大模型训练中被广泛用于平衡吞吐与数值稳定性。所谓「反向传播的节省量」与「峰值分配」，分别衡量训练全程显存占用的降低幅度与瞬时最大显存需求，二者随基线模型和统计口径不同可能给出相反的结论，因此内存权衡不能脱离具体基线来谈。

召回实验与语言模型表现

在小模型多 token 召回实验中，两种设计在最难层级上均优于 Mamba 1-3；更高的路由秩在有门控或相位多样性配置时表现最佳。然而在约 5 亿参数、2B token 的单种子语言模型上，两者虽在测试损失上优于 Mamba 1/2，但路由带来的增益并未转化为更低的语言模型损失。最佳配置在 10B token 训练中接近 Mamba 3 SISO，但仍落后于 Mamba 3 SISO 与 Mamba 3 MIMO 的语言模型损失。

这提示了一个值得关注的落差：检索能力的提升与语言建模损失之间并非线性传导。召回任务上的优势能否在更大规模预训练中兑现，仍是开放问题。

编辑点评

该研究的贡献在于把对角等价这一抽象概念与具体的检索秩设计、内核实现串成一条线，并用 Mamba 1-3 作为可恢复特例来锚定新架构的位置。不过，语言模型损失上的落后说明路由机制的收益边界尚不清晰。对于关注 SSM 架构演进与高效训练内核的读者，这篇工作提供了兼具理论深度与工程细节的参考。