【长文阅读】MAMBA作者博士论文＜MODELING SEQUENCES WITH STRUCTURED STATE SPACES＞-Chapter3 Computing SSMs

Chapter 3 Computing Structured SSMs

Gu A. Modeling Sequences with Structured State Spaces[D]. Stanford University, 2023.

本文是MAMBA作者的博士毕业论文，为了理清楚MAMBA专门花时间拜读这篇长达330页的博士论文，由于知识水平有限，只能尽自己所能概述记录，并适当补充一些相关数学背景，欢迎探讨与批评指正。内容多，分章节更新以免凌乱。

第3章讨论了结构化状态空间模型（SSM），特别是S4模型的计算。它概述了SSM的计算挑战，并提出了结构化SSM的需求。这一章介绍了两类结构化SSM：S4-Diag（对角线结构状态空间模型）和S4-DPLR（对角线加低秩参数化）。章节详细讨论了这些模型、它们的算法和计算复杂性，并对这些模型进行了深入比较。此外，还简要调查了其他类型的结构化SSM，涉及SSM的各个方面，包括参数化细节和计算效率。

3.1 Motivation: Computational Difficulty of SSMs

这一节讨论状态空间模型（SSM）的计算难度。首先，定义了最优效率的SSM，并介绍了两种不同的计算模式：循环模式和卷积模式。接着讨论了这两种模式的优化计算问题，强调了即使对于结构化的矩阵，实现高效的SSM计算也是具有挑战性的。为了解决这些挑战，文档提出了一些解决方案和思路，比如使用对角矩阵和利用特定的算法来减少计算复杂度。这些讨论为引入结构化SSM（特别是S4模型）提供了动机和理论基础。

定义了两种模式：

定义3.1（循环模式计算）：展开 SSM 递归的一个步骤需要与离散状态矩阵 $A$ 进行矩阵向量乘法 (MVM)。如果可以在 $O(N)$ 时间内计算该 MVM，则 SSM 具有最佳递归模式而不是 $O(N^2)$。

定义3.2（卷积模式计算）: SSM卷积模式需要具体化状态空间内核$\overline{\mathbf{K}}$,如果该 SSK 可以在 $O(L+N)$ 时间内计算，则 SSM 具有最佳卷积模式。

本节内容就是简单说明，为后续引出结构化方式做铺垫，简要概述一下，不展开。

3.1.1 General Recurrent Computation：这部分讨论了在循环模式下计算状态空间模型（SSM）的一般方法。它探讨了循环网络的效率和它们处理长序列数据时的挑战，特别是在资源受限的情况下。

3.1.2 General Convolutional Computation：这一小节集中讨论了卷积模式下的SSM计算。卷积模式在处理时间序列数据时更加高效，但它们也面临复杂度和优化的问题。

3.1.3 Structured SSMs：此部分介绍了结构化状态空间模型（如S4）作为一种解决上述计算问题的方法。它探索了通过特定结构化设计来优化SSM的可能性，强调了结构化方法在计算效率和模型性能上的优势

3.2 Diagonally Structured State Space Models

3.2节详细探讨了对角线结构化状态空间模型（Diagonally Structured State Space Models），简称为S4-Diag。这部分首先介绍了使用对角矩阵来简化SSM计算的理论基础。接着，它提出了S4-Diag模型，并讨论了该模型的两个主要问题：递归模式下的计算和卷积核的计算。对于递归模式，由于使用对角矩阵，使得计算变得更加简单。对于卷积核，使用范德蒙德矩阵的乘法进行计算，这也简化了计算过程。这一节还讨论了S4-Diag模型的时间和空间复杂度，并提供了实现示例。通过这些讨论，展示了对角线结构化SSM在简化计算和提高效率方面的优势。

引理 3.3：共轭是 SSM 上的等价关系

$(\mathbf{A},\mathbf{B},\mathbf{C})\sim \left({\mathbf{V}}^{-1}\mathbf{A}\mathbf{V},{\mathbf{V}}^{-1}\mathbf{B},\mathbf{C}\mathbf{V}\right)$

对上式补充说明

证明可以从状态空间模型(SSM)的两种表达形式出发。首先，有两个SSM，其状态分别用 $x$ 和 $\tilde{x}$ 表示：

1.对于第一个SSM：

$\begin{array}{rl}{x}^{\prime }& =\mathbf{A}x+\mathbf{B}u\\ y& =\mathbf{C}x\end{array}$

2.对于第二个SSM（经过变换）：

$\begin{array}{rl}{\tilde{x}}^{\prime }& ={\mathbf{V}}^{-1}\mathbf{A}\mathbf{V}\tilde{x}+{\mathbf{V}}^{-1}\mathbf{B}u\\ y& =\mathbf{C}\mathbf{V}\tilde{x}\end{array}$

• 这里，${\tilde{x}}^{\prime }$ 表示下一状态。
• $V^{-1}AV\tilde{x}$ 是状态变换的主要部分。首先，$\tilde{x}$ 乘以 $V$，将状态空间转换到新的基。然后，应用原始的状态转换矩阵 $A$，最后通过 $V^{-1}$ 把状态空间转换回原来的基。
• $V^{-1}Bu$ 表示对控制信号 $u$ 的影响，也经过了类似的基变换。
• $CV\tilde{x}$表示输出是如何从变换后的状态 $\tilde{x}$产生的。首先，状态 $\tilde{x}$乘以 $V$，进行基的变换，然后乘以输出矩阵$C$ 产生最终的输出。

通过 $V$ 将第二个 SSM 乘以 $V$ 后，这两个SSM变得一样，即 $x=V\tilde{x}$。因此，这两个SSM计算的是相同的操作符 $\mathrm{u}\to \mathrm{y}$，但在状态 $x$ 中通过 $V$ 进行了基的变换。换句话说，它们表达了相同的从 u 到 y 的映射。这也被称为 SSM 文献中的状态空间转换。

这些公式通过 $V$和 $V^{-1}$ 实现了状态空间的基变换。这意味着虽然状态空间的表示形式改变了，但系统的基本动态（即从输入$u$ 到输出 $y$ 的映射）保持不变。这种转换有助于简化问题，或者在特定情况下使得问题更易于分析。在实际应用中，选取合适的变换矩阵 $V$ 可以使状态空间更加规范化，例如使得矩阵 $A$更接近于对角矩阵形式，从而简化计算和分析。

命题3.4 ：在所有 $N\times N$ 维的复数矩阵集合 $C^{N\times N}$ 中，可对角化矩阵的集合$D$ 是稠密的，并且具有完全测度（即其补集的测度为0）。

换句话说，这个命题意味着（几乎）所有的状态空间模型（SSM）都等价于某个对角SSM。这表明，通过适当的变换，大多数SSM都可以被简化为对角形式，这在数学和工程应用中是一个重要的结论，因为对角矩阵通常更易于处理和分析。特别是，计算 $\overline{\mathbf{K}}$成为一种经过充分研究的结构化矩阵乘法，具有高效的时间和空间复杂度。

3.2.1 S4D: Diagonal SSM Algorithms

本小节介绍了对角SSM（状态空间模型）算法，即S4D（Structured State Space Model for Diagonal matrices）。以下是该部分的关键内容：

1. 对角SSM的符号表示：
   • 当$A$是对角矩阵时，$A_n$ 表示其对角线上的元素。
   • 在单输入单输出（SISO）情况下，$B\in R^{N\times 1}$ 和 $C\in R^{1\times N}$，因此可以直接使用 $B_n$ 和 $C_n$来表示它们的元素。
2. S4D递归算法：
   • 在对角SSM上计算任何离散化都很简单，因为对角矩阵上的解析函数可以简化为对其对角线元素逐个执行。
   • 通过对角矩阵进行矩阵-向量乘法（MVM）也很简单，因为它可以简化为逐元素乘法。因此，对角SSM轻松满足了定义3.1。
3. S4D卷积核：范德蒙德矩阵乘法：
   • 当$A$是对角矩阵时，计算卷积核变得非常简单： $K_{\ell }={\sum }_{n=0}^{N-1}{C}_n{A}_n^{\ell }{B}_n$
   • 这里 $K$ 可以表示为 $\left(B^{\top }\circ C\right)\cdot V_L(A)$，其中 $ \circ$ 是哈达马德积，$\cdot$ 是矩阵乘法，$V$ 是范德蒙德矩阵，$V_L(A{)}_{n\ell }=A_n^{\ell }$。

总结来说，S4D是一种处理对角SSM的算法，其核心优势在于通过利用对角矩阵的简单性质，简化了计算过程。对角矩阵的解析函数和乘法运算都可以高效地逐元素进行，而范德蒙德矩阵乘法则进一步简化了卷积核的计算。

这里补充说明一些基础数学知识，方便小伙伴看懂：

补充说明1：对角矩阵上的解析函数

指的是对于一个对角矩阵，对其每个对角线上的元素单独应用某个解析函数，并保持矩阵的其他位置为零。这种操作利用了对角矩阵的特性，即除了对角线上的元素外，其他位置的元素都是零。

以 $A$ 为一个对角矩阵，其对角线上的元素为 $A_{11},A_{22},\dots ,A_{nn}$，那么对 $A$ 应用解析函数ff