状态空间模型

State Space Models

在机器学习和统计建模中，SSM 是一类非常常用的动态模型，用于表示系统随时间演化（或随其他序列索引变化）的过程。它将系统的“内在状态”与观测到的“输出”分离，并通过状态转移与观测方程来刻画整体系统的动态行为。

物理

以弹簧-质量-阻尼系统 (SMD)为例，在给定一个力作为系统输入量 $u(t)$ ，质量块 $M$ 随时间变化的位移 $x$，速度为 $v$，其中 $x$ 作为系统的输出量 $y(t)$

系统的状态向量可以表示为：

$$\mathbf{X}=\left[\begin{array}{c}{x}_1\\ x_2\end{array}\right]=\left[\begin{array}{c}x\\ v\end{array}\right]=\left[\begin{array}{c}y(t)\\ \frac{dy(t)}{dt}\end{array}\right]$$

根据牛顿第二定律：$F=Ma$，弹簧提供的力为 $-kx$，阻尼器提供的力为 $-bv$，外部力为 $u(t)$，有：

$$Ma=-kx-bv+u(t)$$$$M\frac{d^{2}x}{d{t}^2}+b\frac{dx}{dt}+kx=u(t)$$

可以表示为：

$$M{\dot{x}}_2+b{x}_2+k{x}_1=u(t)$$

求解该微分方程，可得：

$${\dot{x}}_1=x_2$$$${\dot{x}}_2=\frac{u(t)-b{x}_2-k{x}_1}{M}$$

即：

$$\dot{\mathbf{X}}=\left[\begin{array}{c}{\dot{x}}_1\\ {\dot{x}}_2\end{array}\right]=\left[\begin{array}{c}{x}_2\\ \frac{u(t)-b{x}_2-k{x}_1}{M}\end{array}\right]$$

即：

$$\dot{\mathbf{X}}=\left[\begin{array}{cc}0& 1\\ -\frac{k}{M}& -\frac{b}{M}\end{array}\right]\left[\begin{array}{c}{x}_1\\ x_2\end{array}\right]+\left[\begin{array}{c}0\\ \frac{1}{M}\end{array}\right]u(t)$$

输出 $y(t)$ 也可以写成：

$$y=\left[\begin{array}{cc}1& 0\end{array}\right]\left[\begin{array}{c}{x}_1\\ x_2\end{array}\right]+\left[\begin{array}{c}0\end{array}\right]u(t)$$

卡尔曼滤波

基本动态系统模型

卡尔曼滤波模型假设 $k+1$ 时刻的真实状态是从 $k$ 时刻的状态演化而来，符合状态方程：

$$x_k=F_k{x}_{k-1}+B_k{u}_K+w_k$$

其中：

• Fk 是作用在 x 上的状态变换模型
• B_k 是作用在控制器向量 $u_k$ 上的输入-控制模型
• w_k 是过程噪声，并假定其符合均值为 0，协方差矩阵为 $Q_k$ 的多元正态分布

算法逻辑

卡尔曼滤波器的状态由以下变量表示：

• ${\hat{X}}_{k|k}=\mathbb{E}(X_k|Y_1,Y_2,\cdots ,Y_k)$ 表示在时刻 $k$ 的状态的估计
• ${\hat{X}}_{k|k-1}=\mathbb{E}(X_{k-1}|Y_1,Y_2,\cdots ,Y_{k-1})$ 表示已知过去 $k-1$ 个时刻的状态，对 $k$ 时刻状态的预测。
• ${\hat{P}}_{k|k}$ 为后验估计误差协方差矩阵，度量估计值的精确程度。

1. 预测 ${\hat{X}}_{k-1|k-1}\Rightarrow {\hat{X}}_{k|k-1}$

预测步骤中，根据上一时刻的状态和控制量，预测当前时刻的状态。这个预测值是一个估计值，因为它还没有考虑当前时刻的观测值。预测值的误差协方差矩阵是通过上一时刻的误差协方差矩阵和系统噪声协方差矩阵计算得到的。

$${\hat{x}}_{k|k-1}=F_k{\hat{x}}_{k-1}+B_k{u}_k$$$$P_{k|k-1}=F_k{P}_{k-1|k-1}{F}_k^T+Q_k$$

2. 更新 ${\hat{X}}_{n|n-1}\Rightarrow {\hat{X}}_{n|n}$

更新步骤中，根据当前时刻的观测值和预测值，计算出当前时刻的状态估计值。这个估计值是一个更加准确的估计值，因为它已经考虑了当前时刻的观测值。状态估计值的误差协方差矩阵是通过预测步骤中计算得到的误差协方差矩阵、观测噪声协方差矩阵和卡尔曼增益计算得到的。

$$K_k=P_{k\mid k-1}{H}_k^T{(H_k{P}_{k\mid k-1}{H}_k^T+R_k)}^{-1}$$$${\hat{x}}_{k\mid k}={\hat{x}}_{k\mid k-1}+K_k(z_k-H_k{\hat{x}}_{k\mid k-1})$$$$P_{k\mid k}=(I-K_k{H}_k)P_{k\mid k-1}$$

以上五个公式为卡尔曼滤波的核心公式。更新步骤更加简洁便于理解的形式为，首先计算以下三个量：

$$\{\begin{array}{ll}{\hat{y}}_k=z_k-H_k{\hat{x}}_{k\mid k-1}& \text{(测量残差)}\\ \\ S_k=H_k{P}_{k\mid k}{H}_k^T+R_k& \text{(测量残差协方差) }\\ \\ K_k=P_{k\mid k-1}{H}_k^T{S}_k^{-1}& \text{(最优卡尔曼增益) }\end{array}$$

然后用它们来更新滤波器变量：

$$\{\begin{array}{ll}{\hat{x}}_{k\mid k}={\hat{x}}_{k\mid k-1}+K_k{\hat{y}}_k& \text{(更新的状态估计) }\\ \\ P_{k\mid k}=(I-K_k{H}_k)P_{k\mid k-1}& \text{(更新的协方差估计) }\end{array}$$

---

两个方程：状态方程与输出方程

状态方程

$$h(t)=\mathbf{A}h(t-1)+\mathbf{B}x(t)$$

输出方程

$$y(t)=\mathbf{C}h(t)+\mathbf{D}x(t)$$

其中：

• $\mathbf{A}\in {\mathbb{R}}^{d\times d}$：状态转移矩阵
• $\mathbf{B}\in {\mathbb{R}}^{d\times m}$：输入映射矩阵
• $\mathbf{C}\in {\mathbb{R}}^{p\times d}$：观测矩阵
• $\mathbf{D}\in {\mathbb{R}}^{d\times m}$：前馈矩阵

离散化

在实际应用中往往是离散的数据，因此要对模型进行离散化表示。

欧拉法

双线性（bilinear method）

采用双线性变换将连续 SSM 转换为离散形式：

$$x_k=\bar{\mathbf{A}}{x}_{k-1}+\bar{\mathbf{B}}{u}_k$$$$y_k=\mathbf{C}{x}_k+\mathbf{D}{u}_k$$$$\bar{\mathbf{A}}=(\mathbf{I}-\mathrm{\Delta }/2\cdot \mathbf{A}{)}^{-1}(\mathbf{I}+\mathrm{\Delta }/2\cdot \mathbf{A})$$$$\bar{\mathbf{B}}=(\mathbf{I}-\mathrm{\Delta }/2\cdot \mathbf{A}{)}^{-1}\mathrm{\Delta }\mathbf{B}$$

零阶保持法（Zero-Order Hold Technoloty）

SSM 的 RNN 表示

SSM 的卷积表示

$$y_k=x_k\ast kernel$$

假设初始状态 $x_{-1}=0$，有：

$$\begin{array}{rl}{y}_0& =\mathbf{C}\bar{\mathbf{B}}{u}_0\\ y_1& =\mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}}{u}_0+\mathbf{C}\bar{\mathbf{B}}{u}_1\\ y_2& =\mathbf{C}{\bar{\mathbf{A}}}^2\bar{\mathbf{B}}{u}_0+\mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}}{u}_1+\mathbf{C}\bar{\mathbf{B}}{u}_2\\ \cdots \\ y_k& =\mathbf{C}{\bar{\mathbf{A}}}^k\bar{\mathbf{B}}{u}_0+\mathbf{C}{\bar{\mathbf{A}}}^{k-1}\bar{\mathbf{B}}{u}_1+\cdots +\mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}}{u}_{k-1}+\mathbf{C}\bar{\mathbf{B}}{u}_k\end{array}$$

将其写成卷积表达：

$$y_k=\left[\begin{array}{llll}\mathbf{C}{\bar{\mathbf{A}}}^k\bar{\mathbf{B}}& \mathbf{C}{\bar{\mathbf{A}}}^{k-1}\bar{\mathbf{B}}& \dots & \mathbf{C}\bar{\mathbf{B}}\end{array}\right]\left[\begin{array}{c}{u}_0\\ u_1\\ \dots \\ u_k\end{array}\right]=\bar{\mathbf{K}}u$$

$\bar{\mathbf{K}}$ 被称为 SSM 卷积核

---

S4（结构化序列状态空间模型）

Structured State Space for Sequences

Gu A, Goel K, Ré C. Efficiently modeling long sequences with structured state spaces[J]. arXiv preprint arXiv:2111.00396, 2021.

https://srush.github.io/annotated-s4/

ICLR2022

核心改进：

• 使用低秩修正来优化矩阵 $\mathbf{A}$ 的性质，使其更容易对角化。
• 使用 HiPPO 矩阵解决长距离依赖问题
• 将连续 SSM 转换为离散形式

HiPPO 矩阵：

$${\mathbf{A}}_{nk}=-\{\begin{array}{ll}(2n+1{)}^{1/2}(2k+1{)}^{1/2}& \text{ if }n>k\\ n+1& \text{ if }n=k\\ 0& \text{ if }n<k\end{array}$$

S4 的参数化

为了避免 HiPPO 矩阵运算量过大，作者提出将 HiPPO 矩阵转换为：正规矩阵+低秩矩阵

NPLR(Normal Plus Low-Rank)

$$\mathbf{A}=\mathbf{V}\mathbf{\Lambda }{\mathbf{V}}^{\ast }-\mathbf{P}{\mathbf{Q}}^{\mathrm{\top }}=\mathbf{V}(\mathbf{\Lambda }-\left({\mathbf{V}}^{\ast }\mathbf{P}\right){\left({\mathbf{V}}^{\ast }\mathbf{Q}\right)}^{\ast }){\mathbf{V}}^{\ast }$$

矩阵 $\mathbf{A}$ 转换为 NPLR 后：

• RNN 形式

  • S4 1 个 step 的计算复杂度依旧为 $O(N)$

• 卷积形式

  • $\bar{\mathbf{K}}$ 计算复杂度从 $O(N^{2}L)$ 降低为 $O(N+L)$

---

S4D（对角状态空间模型）

Gu A, Goel K, Gupta A, et al. On the parameterization and initialization of diagonal state space models[J]. Advances in Neural Information Processing Systems, 2022, 35: 35971-35983.

S4D 模型是一种对角 SSM，它结合了 S4 的计算和参数化优势以及 DSS 的初始化方法。

S4D 模型的核心在于其对角状态矩阵的参数化和初始化，这使得模型在保持简单性的同时，也具备了强大的性能。

S4D 模型的卷积核 $\bar{\mathbf{K}}$ 可以通过 Vandermonde 矩阵乘法来高效计算：

$$\begin{array}{rl}\bar{\mathbf{K}}& =\left[\begin{array}{llll}\mathbf{C}\bar{\mathbf{B}}& \mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}}& \dots & \mathbf{C}{\bar{\mathbf{A}}}^{L-1}\bar{\mathbf{B}}\end{array}\right]\\ & =\left[\begin{array}{lll}{\bar{\mathbf{B}}}_0{\mathbf{C}}_0& \dots & {\bar{\mathbf{B}}}_{N-1}{\mathbf{C}}_{N-1}\end{array}\right]\left[\begin{array}{ccccc}1& {\bar{\mathbf{A}}}_0& {\bar{\mathbf{A}}}_0^2& \dots & {\bar{\mathbf{A}}}_0^{L-1}\\ 1& {\bar{\mathbf{A}}}_1& {\bar{\mathbf{A}}}_1^2& \dots & {\bar{\mathbf{A}}}_1^{L-1}\\ ⋮& ⋮& ⋮& \ddots & ⋮\\ 1& {\bar{\mathbf{A}}}_{N-1}& {\bar{\mathbf{A}}}_{N-1}^2& \dots & {\bar{\mathbf{A}}}_{N-1}^{L-1}\end{array}\right]\\ & =({\bar{\mathbf{B}}}^{\mathrm{\top }}\circ \mathbf{C})\cdot {\mathcal{V}}_L(\bar{\mathbf{A}})\end{array}$$

---

S5

Smith J T H, Warrington A, Linderman S W. Simplified state space layers for sequence modeling[J]. arXiv preprint arXiv:2208.04933, 2022.

ICLR2023

S4 采用的离散 SSM 形式为：

$$x_k=\bar{\mathbf{A}}{x}_{k-1}+\bar{\mathbf{B}}{u}_k$$

该形式为单输入单输出，对于输入向量 $u=[u_1,u_2,\cdots ,u_k]$，S4 直接对每个分量独立地应用上述公式，然后将结果拼接起来。

而 S5 采用了多输入多输出（MIMO）形式。

---

S6（Mamba）

见Mamba