Mamba系列日积月累（一）：状态空间模型SSM的离散化过程推导

本文首发于： Mamba系列日积月累（一）：状态空间模型SSM的离散化过程推导

最近Mamba系列（Mamba、VMamba、Vision Mamba）比较火，在同样具备高效长距离建模能力的情况下，Transformer具有平方级计算复杂度，而Mamba架构则是线性级计算复杂度，并且推理速度更快。

秉承着公众号科研的思路扩展视野的思路，笔者觉得需要学习一下相关内容，于是挑选了目前较新的Vision Mamba论文，准备开始学习。由于缺乏之前的基础知识储备，Preliminaries里面的状态空间模型及其离散化过程直接给我干蒙，想着不能出师未捷身先死，于是决定搜索相关资料，把这个过程弄明白，不过由于本人水平有限，如果内容存在错误，希望大家能给出指导进行纠正。

1. 背景基础知识

1.1 什么是状态空间模型（State Space Model，SSM）？

状态空间模型（State Space Model，简称SSM）是一种数学模型，用于描述和分析动态系统的行为。这种模型在多个领域都有应用，包括控制理论、信号处理、经济学和机器学习等。在深度学习领域，状态空间模型被用来处理序列数据，如时间序列分析、自然语言处理（NLP）和视频理解等。通过将序列数据映射到状态空间，可以更好地捕捉数据中的长期依赖关系。

状态空间模型的核心思想是将系统的当前状态（state）$x(t)\in {\mathbb{R}}^n$与输入（input）$u(t)\in {\mathbb{R}}^p$和输出（output）$y(t)\in {\mathbb{R}}^q$之间的关系用一组方程来表示：
$$\begin{array}{rl}& \dot{x}(t)=A(t)x(t)+B(t)u(t)\\ & y(t)=C(t)x(t)+D(t)u(t)\end{array}\text{\hspace{1em}(1)}$$

1. 状态方程（State Equation）：描述系统状态随时间的演变。状态方程通常包含当前状态和输入，以及可能的系统参数。数学上，状态方程可以表示为： $\dot{x}(t)=A(t)x(t)+B(t)u(t)$， 其中，$x(t)$是在时间步 $t$ 的系统状态，$\dot{x}(t)$是状态向量$x(t)$关于时间 $t$的导数，$u(t)$ 是在时间步 $t$的输入，$A(t)$是状态转移矩阵，$\dim [A(\cdot )]=n\times n$，$B$ 是输入矩阵，$\dim [B(\cdot )]=n\times p$。
2. 观测方程（Observation Equation）：描述系统输出与状态之间的关系。观测方程允许我们从系统的输出中观察到系统的状态。数学上，观测方程可以表示为： $y(t)=C(t)x(t)+D(t)u(t)$ 其中，$y(t)$ 是在时间步 $t$ 的系统输出，$C(t)$是观测矩阵，$\dim [C(\cdot )]=q\times n$，$D(t)$ 是前馈矩阵，$\dim [D(\cdot )]=q\times p$。

当式(1)中的所有矩阵均随着时间$t$而变化时，此时所表示的线性时变系统，而当所有矩阵都不随时间$t$​变化时，此时表示的是线性非时变系统，在Mamba系列中，实际上是线性非时变系统：
$$\begin{array}{rl}& \dot{x}(t)=Ax(t)+Bu(t)\\ & y(t)=Cx(t)+Du(t)\end{array}\text{\hspace{1em}(2)}$$

1.2 什么是离散化（Discretization）？

离散化（Discretization）是将连续的数学对象或过程转换为离散形式的过程。在不同的领域中，离散化有着不同的应用和含义，但核心思想是一致的：将连续的变量或函数映射到有限的、离散的集合中。这个过程在数学、工程、计算机科学和许多其他领域中都非常常见。

1.3 为什么需要离散化？

SSM作为一个连续时间系统，其难以直接集成到现代深度学习算法中：

• 计算效率：现代深度学习框架和硬件通常是基于离散时间操作而设计的，对SSM进行离散化后，才能将其转化为可以在这些框架和硬件上高效运行的模型。
• 训练算法：大多数深度学习训练算法，如梯度下降和反向传播，都是为离散时间模型设计的。离散化使得这些算法可以直接应用于状态空间模型，简化了训练过程。
• 实际应用：在许多实际应用中，数据是离散的，如文本数据（单词序列）、时间序列数据（股票价格、传感器读数）等。离散时间模型更自然地与这些数据格式相匹配。
• 模型复杂度：离散化过程可以通过选择合适的时间步长 $T$ 来控制模型的复杂度。较小的时间步长可以提供更精细的控制，但计算成本更高；较大的时间步长可以减少计算量，但可能牺牲一些精度。

2. SSM离散化过程推导

这里再贴上状态方程公式
$$\dot{x}(t)=Ax(t)+Bu(t)\text{\hspace{1em}(3)}$$
为了进行离散化，我们首先要对状态方程(3)进行积分。

2.1 为什么在离散化过程中要先进行积分？

在离散化连续状态方程的过程中，积分是一个关键步骤，因为它涉及到状态变量随时间的累积效应，我们需要考虑在每个离散时间步长内状态变量是如何累积变化的。

在离散时间系统中，我们不能直接处理导数，因为离散时间点上没有导数的概念。相反，我们需要考虑在每个时间步长内状态变量的累积变化。这可以通过对连续时间积分进行离散化来实现，即将连续时间的积分转换为离散时间的求和。

在实际的数值模拟中，我们通常使用数值积分方法（如梯形法则、矩形法则、辛普森法则等）来近似连续时间积分。这些方法允许我们在离散时间点上近似连续时间的累积效应，从而得到离散时间状态方程。这个转换过程涉及到将连续时间的导数项替换为离散时间的差分项，这通常涉及到指数函数和采样间隔 $T$​ 的计算。

2.2 为什么不直接对$\dot{x}(t)$进行积分？

在式(3)中，假设我们直接对$\dot{x}(t)$进行积分的话，结果如下：
$$x(t)=x(0)+{\int }_0^t(Ax(\tau )+Bu(\tau ))d\tau \text{\hspace{1em}(4)}$$
此时，积分项中会包含$x(\tau )$项本身，由于我们是离散系统，我们是无法获取在一个连续的时刻（$0\to t$）内所有的$x(\tau )$值的，因此无法完成该积分结果的计算。

对于离散系统来说，我们希望将公式(4)这个积分表达式转变为以下形式：
$$x(k+1)=x(k)+\sum _{i=0}^k(Ax(i)+Bu(i))\Delta t\text{\hspace{1em}(5)}$$
这个形式要求我们对公式(3)进行一些改造，目标是消除$\dot{x}(t)$表达式中的$x(t)$本身。

2.2 状态方程的改造以及$\alpha (t)$的设计

为了消除$\dot{x}(t)$表达式中的$x(t)$本身，我们通常会构造一个新的函数$\alpha (t)x(t)$，通过对这个新函数进行求导，来简化相应的导数项。

我们对$\alpha (t)x(t)$​进行求导

$$\frac{d}{dt}[\alpha (t)x(t)]=\alpha (t)\dot{x}(t)+x(t)\frac{d\alpha (t)}{dt}\text{\hspace{1em}(6)}$$
我们将公式(3)代入到公式(6)中，替换$\dot{x}(t)$：

$$\frac{d}{dt}[\alpha (t)x(t)]=\alpha (t)(Ax(t)+Bu(t))+x(t)\frac{d\alpha (t)}{dt}\text{\hspace{1em}(7)}$$
我们进一步对公式(7)进行改写，合并$x(t)$的相关系数：

$$\frac{d}{dt}[\alpha (t)x(t)]=(A\alpha (t)+\frac{d\alpha (t)}{dt})x(t)+B\alpha (t)u(t)\text{\hspace{1em}(8)}$$
由于我们的目的是消除导数项中的$x(t)$，因此，我们令$x(t)$的系数项为0即可：
$$A\alpha (t)+\frac{d\alpha (t)}{dt}=0\text{\hspace{1em}(9)}$$
此时，我们可以得到$\alpha (t)$的表达式：
$$\alpha (t)=e^{-At}\text{\hspace{1em}(10)}$$
将$\alpha (t)$的表达式代入公式(8)可以得到：
$$\frac{d}{dt}[e^{-At}x(t)]=B{e}^{-At}u(t)\text{\hspace{1em}(11)}$$
这时我们已经完成了在导数项中消除$x(t)$的目标，对$e^{-At}x(t)$进行积分：
$$e^{-At}x(t)=x(0)+{\int }_0^t{e}^{-A\tau }Bu(\tau )d\tau \text{\hspace{1em}(12)}$$
对公式(12)进行整理：

$$x(t)=e^{At}x(0)+{\int }_0^t{e}^{A(t-\tau )}Bu(\tau )d\tau \text{\hspace{1em}(13)}$$

2.3 离散时间积分近似

2.3 状态方程的离散化

在离散系统中，我们需要将公式(13)转化为离散形式，大致步骤如下：

• 参数定义：采样时刻$t_k$和$t_{k+1}$，其中$k$是采样索引，$T$是采样间隔，即$T=t_{k+1}-t_k$​

• 积分区间离散化：在连续时间积分中，我们通常有一个积分区间，例如从$t$ 到 $t+△t$。在离散时间系统中，我们需要将这个区间划分为$k$ 个等长的子区间，每个子区间的长度为$T$​​。

  在某个子区间内，公式(13)的形式变为：
  $$x(t_{k+1})=e^{A(t_{k+1}-t_k)}x(t_k)+{\int }_{t_k}^{t_{k+1}}{e}^{A(t_{k+1}-\tau )}Bu(\tau )d\tau \text{\hspace{1em}(14)}$$

• 近似积分：对于每个子区间来说，考虑使用数值积分方法来近似积分，这里考虑对$u(t)$应用零阶保持法，即假设$u(t)$在采样时刻$t_k$和$t_{k+1}$之间是恒定的，此时，我们可以将$u(t)$当做常数项从积分项中取出：
  $${\int }_{t_k}^{t_{k+1}}{e}^{A(t-\tau )}Bu(\tau )d\tau ={\int }_{t_k}^{t_{k+1}}{e}^{A(t_{k+1}-\tau )}d\tau Bu(t_k)\text{\hspace{1em}(15)}$$

• 离散时间状态方程构建：将公式(15)的积分结果代入到公式(14)中，同时使用$T=t_{k+1}-t_k$​进行化简，我们可以得到：
  $$x(t_{k+1})=e^{AT}x(t_k)+{\int }_{t_k}^{t_{k+1}}{e}^{A(t_{k+1}-\tau )}d\tau Bu\left(t_k\right)\text{\hspace{1em}(16)}$$
  引入新变量$\lambda =t_{k+1}-\tau$，对原积分进行简化得到：
  $$x(t_{k+1})=e^{AT}x(t_k)+Bu\left(t_k\right){\int }_0^T{e}^{A\tau }d\tau \text{\hspace{1em}(17)}$$
  这里涉及到矩阵作为指数的积分，这个部分我是查阅一些资料得到的结果：
  $${\int }_0^T{e}^{A\tau }d\tau =A^{-1}(e^{AT}-I)\text{\hspace{1em}(18)}$$
  最终我们得到了离散时间状态方程：
  $$x(t_{k+1})=e^{AT}x(t_k)+(e^{AT}-I)A^{-1}Bu\left(t_k\right)\text{\hspace{1em}(19)}$$

3. SSM离散化结果

对比公式(19)和Vision Mamba论文中的离散化结果：

两者形式基本一致，至此，我们完成了SSM的离散化过程的完整推导。