浅谈状态空间模型SSM基础理论

浅谈状态空间模型SSM基础理论

状态空间模型（State Space Model，简称SSM）是一种描述系统随时间变化的动态行为的模型，被广泛应用于控制理论、机器人学以及经济学等领域。通过一系列隐藏变量（即“状态”），SSM能够高效地捕捉时序数据的依赖关系，从而实现对系统行为的建模。

基本理论

状态空间模型由两种方程构成：状态方程和观测方程。这两种方程共同构建了时刻 $t$ 的输入 $x(t) \in \mathbb{R}$ 和输出 $y(t) \in \mathbb{R}$ 之间的关系。其中，隐藏状态 $h(t) \in \mathbb{R}^N$ 的维度为 $N$。

状态方程描述了系统状态随时间的变化：
$$
\begin{aligned}
h'(t) = \mathbf{A}h(t) + \mathbf{B}x(t)
\end{aligned}\tag{1}
$$

观测方程描述了系统输出与当前状态的关系：
$$
\begin{aligned}
y(t) = \mathbf{C}h(t) + \mathbf{D}x(t)
\end{aligned}\tag{2}
$$

其中：

• $h'(t)$ 是当前状态 $h(t)$ 的导数
• $\mathbf{A} \in \mathbb{R}^{N \times N}$ 是状态转换矩阵，用于描述状态随时间的变化情况
• $\mathbf{B} \in \mathbb{R}^{N \times 1}$ 是输入矩阵，控制输入对状态的影响
• $\mathbf{C} \in \mathbb{R}^{1 \times N}$ 是输出矩阵，表示基于当前状态产生的输出
• $\mathbf{D}$ 表示命令系数，决定输入对输出产生的直接影响

通常情况下，大部分状态空间模型的 $\mathbf{D}x(t) = 0$，在深度学习中可识别为 skip 连接。

离散化

为了满足各种真实世界场景中机器学习的需求，SSMs 模型需要经历离散化过程，将连续参数转换为离散参数。离散化方法通常将连续时间分为 $K$ 个离散化间隔。其中，最有代表性的是 Zero-Order Hold（ZOH），其假设在间隔 $[t_{k-1}, t_k]$ 内为常量。在 ZOH 离散化后，状态空间模型变为：

$$
\begin{aligned}
h_k = \bar{\mathbf{A}}h_{k-1} + \bar{B}x_k
\end{aligned}\tag{3}
$$

$$
\begin{aligned}
y_k = \mathbf{C}h_k
\end{aligned}\tag{4}
$$

其中：

• $\bar{\mathbf{A}} = \exp(\Delta \mathbf{A})$
• $\bar{\mathbf{B}} = (\Delta \mathbf{A})^{-1}(\exp(\Delta A) - I) \cdot \Delta \mathbf{B}$
• $k$ 为离散化时间步

由此可见，离散化 SSMs 与 RNN 具有相似的结构。因此，与自回归解码的 Transformer 相比，SSMs 的推理效率更高。

卷积运算

离散化状态空间模型属于一个线性系统，具有相关特性，并且与卷积运算无缝集成。具体来说，每个时间步可以独立计算输出：

$$
\begin{aligned}
y_0 = \mathbf{C}\bar{\mathbf{A}}^0\bar{\mathbf{B}}x_0
\end{aligned}\tag{5}
$$

$$
\begin{aligned}
y_1 = \mathbf{C}\bar{\mathbf{A}}^1\bar{\mathbf{B}}x_0 + \mathbf{C}\bar{\mathbf{A}}^0\bar{\mathbf{B}}x_1
\end{aligned}\tag{6}
$$

$$
\begin{aligned}
y_2 = \mathbf{C}\bar{\mathbf{A}}^2\bar{\mathbf{B}}x_0 + \mathbf{C}\bar{\mathbf{A}}^1\bar{\mathbf{B}}x_1 + \mathbf{C}\bar{\mathbf{A}}^0\bar{\mathbf{B}}x_2
\end{aligned}\tag{7}
$$

$$
\begin{aligned}
\cdots
\end{aligned}\tag{8}
$$

$$
\begin{aligned}
y_2 = \mathbf{C}\bar{\mathbf{A}}^k\bar{\mathbf{B}}x_0 + \mathbf{C}\bar{\mathbf{A}}^{k-1}\bar{\mathbf{B}}x_1 + \ldots + \mathbf{C}\bar{\mathbf{A}}^1\bar{\mathbf{B}}x_{k-1} + \mathbf{C}\bar{\mathbf{A}}^0\bar{\mathbf{B}}x_k
\end{aligned}\tag{9}
$$

通过创建卷积核 $\bar{\mathbf{K}} = (\mathbf{C}\bar{\mathbf{B}}, \ldots, \mathbf{C}\bar{\mathbf{A}}^{k}\bar{\mathbf{B}}, \ldots)$，循环计算可以转化为卷积形式：

$$
\begin{aligned}
y = \mathbf{x} * \mathbf{\bar{K}}
\end{aligned}\tag{10}
$$

其中：

• $\mathbf{x} = [x_0, x_1, \ldots]$ 和 $\mathbf{y} = [y_0, y_1, \ldots] \in \mathbb{R}^{L}$ 表示输入和输出序列
• $L$ 为序列长度

卷积运算使 SSMs 模型可以充分利用现代矩阵运算硬件，从而在训练过程中实现并行计算，这对于使用非线性激活函数的 RNN 是不可能的。在这种情况下，给定 $D$ 维输入 $x(k)$，SSMs 模型独立计算每一个维度，最终产生 $D$ 维度输出 $y(t)$。在这种情况下，输入矩阵、输出矩阵、命令矩阵以及状态转换矩阵均保持不变。

RNN、Transformer 以及 SSM 之间的关系

RNN、Transformer 以及 SSM 之间的关系如图 1 所示：

图 1 模型架构的描述

RNN 在非线性循环框架下计算，且每个计算只取决于之前的隐藏状态和当前输入。虽然这种方式导致 RNN 可快速推理，但阻碍了利用 GPU 并行计算的能力，即训练缓慢。与之相比，Transformer 的推理过程很缓慢。这是因为语言模型的自回归设计导致每个输出的计算均需要重复计算注意力分数。与 RNN 和 Transformer 模型相比，离散化 SSMs 支持循环运算和卷积运算，即不仅可高效推理也可高效计算。然而，经典的 SSMs 模型属于时间不变的，即 $A, B, C$ 以及 $\Delta$ 与模型输入无关。这限制了上下文意识的建模，从而导致 SSMs 在特定任务中表现次优，例如选择性复制。

确切地说，状态空间模型具有广泛的意义，表示具有任何隐式状态的循环过程的概念。在许多学科中，用于指代其他概念，包含马尔可夫决策过程、卡尔曼滤波、动力学因果建模、隐马尔可夫模型、线性动力学系统、以及深度学习中循环神经网络。