Regularity and Stability Properties of Selective SSMs with Discontinuous Gating
Nikola Zubić, Davide Scaramuzza
cs.LG, math.DS, math.OC, stat.ML
2025-05-17
选择性 SSM(如 Mamba)的门控在线调制让稳定性一直没被理解透。这篇用控制论给出稳定性证明,并推出可微训练正则,7 个时序数据集上把 Mamba 核不稳定降约 92%、精度几乎无损。
选择性状态空间模型(selective SSM),比如 Mamba,已经是长序列建模的主力。但它的稳定性一直没被理解透:它的状态空间系数会被一个随 token 变的门控信号在线调制,使得这个递推既不是线性时不变系统,也不是经典的非线性系统,既有的分析(表达能力、Lyapunov 指数那类)给不出一个统一的能量视角。这篇要给这类系统补上稳定性的理论,并把它落到一个能用的训练正则上。
作者用控制论里的无源性(passivity)、耗散性(dissipativity)和输入-状态稳定性(ISS)来分析连续时间的选择性 SSM。关键的一步是把两类输入分开:选择信号 x(·) 通过 A、B、C 矩阵调度系统动态,驱动输入 u(·) 进入供给率(supply rate)。这种分开让分析能抓住「门控在线变」这件事。
四个主要结果。一,严格耗散下齐次轨迹指数遗忘(定理 4.1)。二,冻结选择信号(x≡0)后,子系统是一个真正无源的线性时变系统,其最小可用存储必然是二次的,且即使在门控不连续时也满足 AUCloc 正则性(定理 4.2),这一条是为了处理分段常数门控的不连续性。三,在通用二次存储假设下,给出参数化 LMI(线性矩阵不等式)、通用核约束,以及「不可逆遗忘」性质,即一旦某个状态方向能量为零,结构上就不可能靠任何选择再让它变得可观测。四,给出在所有容许选择调度下一致全局 ISS 的充分条件(定理 7.1)。
落到实践:第 6 节推导出一个一步离散耗散不等式和采样块 LMI,专门针对 Mamba 的 selective-scan 递推。这个矩阵的最大正特征值对模型参数是可微的,于是可以直接当成训练时的正则项。
实验在 7 个标准时序数据集(ETTh1/h2、ETTm1/h2、weather、electricity、traffic)上做,每个 4 个预测视界,共 28 组。
加了这个 LMI 正则后,Mamba 核的 LMI 违反平均降约 92%,而且 28 组全部下降。代价很小:干净 MSE 成本不到未正则基线的 0.018%。保守性消融里,把对角矩阵 P=diag(q) 设成可学习,在 28 组上全部 Pareto 改进,准确率没有可测损失。
一个「中性结果」要说一下:在带路由的消融(8.5 节)里,这个惩罚会稳定地把 Mamba 分支的权重往下压,但训练期 Mamba 权重和局部 LMI 违反的相关性在总体上基本为零、跨数据集符号还不一致。也就是说,正则确实压住了 Mamba 核的不稳定,但并没有表现出「自动把流量从不稳的 Mamba 路由开」的效果。
对用 Mamba 类架构训长序列模型的人,这篇给了一个有理论背书的训练正则:几乎不损精度(0.018%)就能把 Mamba 核的数值不稳定大幅压下去(92%),而且正则是可微的、即插即用。更关键的是它把「Mamba 稳不稳」从一个经验问题变成了一个可计算、可监测的量(采样 LMI 违反),训练时能直接盯着这个指标。
作者把自己的范围划得很清楚,这是诚实的部分也是局限所在。第一,这套判据只针对 Mamba 的 selective-scan 核,不是带归一化、残差、路由的整网 passivity 证书。第二,用的是采样、局部线性化的判据,不是精确的离散 passivity。第三,只在一个架构族(SST/Mamba)上验证,能不能推广到别的选择性 SSM 变体没建立。再加上那个「中性结果」:正则压住了核的不稳定,但没能在路由层面自动把流量绕开不稳的分支,说明这套理论和「整网训练动态」之间还有没接上的地方。