序言

在mamba之前，HiPPO、S4、DSS、S4D都是Albert Gu的工作，可以说是一己之力推动了SSM模型的发展。在S4到mamba之间，其实还经历了很多简化。

本篇选取中间的一篇工作S4D，来一窥目前的mamba仓库代码的雏形。（主要是比mamba清楚太多了，mamba文章和代码仓库完全看不懂，只能说被ICLR24拒稿是有原因的）

P.S. DSS和S4D都是同一年的NIPS，爽发啊

Motivation

S4的突出贡献就在于给出了HiPPO矩阵及相关卷积核高阶幂的求解办法 DPLR，但是从推导就可以看出，仍然是过于复杂了。

前置工作 DSS 发现，存在使用对角线状态矩阵简化原来的HiPPO矩阵的可能性，不过还是引入了一些有难度的操作（需要复值softmax，并且没有解释为什么对角线矩阵可以work）

本篇文章就进一步梳理了SSM使用对角线状态方程的最佳表达形式。如果A只有对角线的话那A的高阶幂就太好算了。

零阶保持离散化(ZOH)

我们首先来介绍一下零阶保持离散化方法

状态方程

$\frac{dh(t)}{dt} = Ah(t) + Bx(t)$

求解

使用积分因子法，首先将方程整理为标准形式：

$\frac{dh(t)}{dt} - Ah(t) = Bx(t)$

接下来计算积分因子 $\mu(t)$：

$\mu(t) = exp(\int -Adt) = e^{-At}$

将积分因子乘以方程两边：

$e^{-At}\frac{dh(t)}{dt}-Ae^{-At}h(t) = Be^{-At}x(t)$

左边变为：

$\frac{d}{dt}(e^{-At}h(t)) = Be^{-At}x(t)$

对两边积分：

$\int_{t_0}^t \frac{d}{d\tau}(e^{-A\tau}h(\tau))d\tau=\int_{t_0}^{t}Be^{-A\tau}x(\tau)d\tau$

积分后得到：

$e^{-At}h(t) - e^{-At_0}h(t_0) = \int_{t_0}^{t}Be^{-A\tau}x(\tau)d\tau$

解出 $h(t)$：

$h(t) = e^{A(t-t_0)}h(t_0) + \int_{t_0}^tBe^{A(t-\tau)}x(\tau)d\tau$

离散化

我们考虑 $ t = t_0 + \Delta$，则 $x(\tau)$ 在这段时间内保持恒定，不妨设为 $x_k$, 则

$\begin{aligned} h(t+\Delta) &= e^{A\Delta}h(t) + \int_t^{t+\Delta}Be^{A(t-\tau)}x_kd\tau \\ &= e^{A\Delta}h(t) + (\int_t^{t+\Delta}e^{A(t-\tau)}d\tau)Bx_k \\ &= e^{A\Delta}h(t) + (\int_0^{\Delta}e^{A\tau}d\tau)Bx_k \end{aligned}$

其中，$\int_0^{\Delta}e^{A\tau}d\tau$ 的结果可以如下计算得到：

$\frac{d}{dt}e^{At} = Ae^{At}$

两边积分得：

$\int Ae^{At}dt = e^{At} + C$

左乘 $A^{-1}$ 得：

$\int e^{At}dt = A^{-1}e^{At} + C$

对于定积分 $[0，\Delta]$，结果是：

$\int_0^{\Delta}e^{A\tau}d\tau = A^{-1}(e^{A\Delta}-I)$

综上，代入积分结果，原来的离散化方程为：

$h(t+\Delta) = e^{A\Delta}h(t) + A^{-1}(e^{A\Delta}-I)Bx_k$

所以

$\begin{aligned} \bar{A} &= e^{A\Delta} \\ \bar{B} &= A^{-1}(e^{A\Delta}-I)B \end{aligned}$

而 SSM 中的 $C$ 和 $D$ 是不需要离散化的，因为它本来就是在离散空间上更新的。其实也不一定需要是 $x’ = Ch + Dx$ 吧，也可以是 $y = Ch + Dx$，SSM最重要的只是维护一个状态罢了——这个状态是对记忆空间正交基的投影。

之前S4中用双线性离散化，就是因为不是对角线形式的 $\bar{A} = e^{A\Delta}$ 不好算。

矩阵指数

矩阵指数的定义参考标量指数函数的泰勒展开：

$e^{At} = I + At + \frac{(At)^2}{2!}+...$

系统的稳定性主要由矩阵指数 $e^{At}$ 决定，分析 $e^{At}$ 的敛散性比较复杂且经典，这里直接给出结论 —— 由矩阵 $A$ 的特征值的实部的正负决定

当矩阵 $A$ 的所有特征值的实部都小于0时，矩阵指数函数 $e^{At}$ 在 $t \rightarrow \infty$ 时收敛到零矩阵。

证明

意会一下，以后补上完整证明

任何矩阵 $A$ 都可以分解为 Jordan 标准型 $A = PJP^{-1}$，其中 $J$ 是 Jordan 矩阵。因此，$e^{At}=Pe^{Jt}P^{-1}$

每个 Jordan 块对应一个特征值 $\lambda$，其矩阵指数包含形如 $t^ke^{\lambda t}$的项（其中 $k$ 是多项式项的阶数）。当 $\lambda$ 的实部小于0时，指数衰减 $e^{\lambda t}$ 会压制多项式增长 $t^k$，使得每一项在 $t \rightarrow \infty$ 时趋于0

因此整个Jordan 矩阵的指数 $e^{Jt}$ 趋向于零矩阵，进而 $e^{At} = Pe^{Jt}P^{-1}$ 也趋向于零矩阵。