【长文阅读】MAMBA作者博士论文＜MODELING SEQUENCES WITH STRUCTURED STATE SPACES＞-Chapter6 Combining Orthogonal and

Chapter 6 Combining Orthogonal and Structured State Space Models

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

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

在论文第 2 章和第 3 章中，定义了结构化 SSM 模型（S4），它结合了连续、循环和卷积特性，同时计算效率非常高。然而，没有定义这些模型的具体实例来解决远程依赖等建模挑战（第 1 章）。另一方面，在第 4 章和第 5 章中，通过 HIPPO 框架定义了具有特殊数学属性的 SSM。特别是，它们解决了在线记忆问题，作为一种对远程依赖关系进行建模的方法（第 4 章），并且可以被视为将输入分解到平滑基函数正交系统上的正交 SSM（第 5 章）。然而，没有讨论如何计算这些模型。本章结合了这两个概念：展示了最重要的正交 SSM，对应于命名的 HIPPO 方法（第 5.1.1 节），都可以编写为结构化 SSM。因此，S4 模型类同时具有最佳效率（定义 3.1 和 3.2），并且可以使用 HIPPO 矩阵进行实例化，以获得额外的理论属性和功能，例如远程记忆。

6.1 节重申了使用 HIPPO 解决 SSM 计算效率问题的主要挑战，并提出了部分解决方案。

6.2节包含了作者最重要的结果：命名的HIPPO矩阵都可以完全写成第3章中的DPLR矩阵。这允许使用HIPPO方法实例化S4模型，解决了它们的缺点。

6.3 节进一步展示了最重要的 HIPPO 矩阵可以通过对角矩阵来近似的理论结果，从而允许 S4D 也可以使用 HIPPO 进行实例化，还提供了其他对角线变体。

6.1 Overview: Motivation and Partial Progress

1. 结构化状态空间模型（SSMs）的有效性：结构化状态空间模型（SSMs）并非总是有效的。由于状态空间模型（SS(S)Ms）的通用性，它们可能会继承递归和卷积在处理长期依赖关系时的问题（参见第1.2.3节）。例如，作为一种递归，通过$\bar{A}$的重复乘法可能会遭受众所周知的梯度消失问题。
2. SSMs的原则性效率问题：即使是原则上设计良好的SSMs也不一定高效。HIPPO框架描述了如何在连续时间内针对某一度量$\omega$记忆函数，并提供了几种具体的命名家族实例（如LegS、LegT、FouT）。

6.1.1 A Resolution to Problem 1

这一节主要讨论了HIPPO算子的计算效率问题

1. HIPPO算子的计算效率：尽管第5章的一般结果暗示对于任何度量$\omega$（确切地说，是具有对应正交多项式族的任何度量）都存在hippo(ω)，从而产生一个线性时不变的状态空间模型（LTI SSM），但这些模型的效率尚不明确。论文中解决了这个问题，专注于递归模式。
2. 定理6.1（非正式）：记忆操作符hippo(ω)总是具有形式$x^{\prime }(t)=\mathbf{A}x(t)+\mathbf{B}u(t)$，其中$A$是一个低递归宽度（low recurrence-width，LRW）状态矩阵。尽管LRW矩阵超出了这一章的讨论范围，但它们是一种结构化矩阵。

这些信息对于理解HIPPO算子在不同模式下的计算效率具有重要意义。需要注意的是，为了全面理解这些概念，可能需要结合整个文件的上下文。

6.1.2 An Attempt for Problem 2

1. 推论 6.3: 对于与经典正交多项式（OPs）相对应的$\omega$，hippo(ω)是3-准可分离的。准可分离矩阵是一类具有额外算法属性的结构化矩阵，它们已知具有高效的（线性时间）矩阵向量乘法（MVM）。
2. 定理 6.4: 对于任何$k-$准可分离矩阵$A$（其中$k$为常数）和任意的$B、C$，Krylov函数${\mathcal{K}}_L(A,B,C)$可以在准线性时间和空间$O˜(N+L)$以及对数深度内计算，意味着它是可并行化的。此算法在精确算术模型中操作，不考虑位复杂度或数值稳定性。
3. 推论 6.5: 定理6.4的算法是在精确算术上的，并且在浮点数上数值不稳定，因为它需要在状态维度N中指数级大的中间值。
4. 备注 6.1: 定理6.4被包含在内是出于几个目的，但在实践中并未使用，并且理解本章其余部分的主要S4-HIPPO联系不需要理解此结果。定理6.4的一个含义是，给定一个特定矩阵（例如HIPPO），可能有许多可能的结构可以满足，这些结构可用于潜在的高效利用。

这些部分提供了关于HIPPO结构化状态空间模型在理论和实践中的应用及其潜在限制的重要洞察

6.2 DPLR Structure of HIPPO Matrices

这一节主要讨论了HIPPO状态空间模型（SSMs）中的DPLR（Diagonal Plus Low-Rank）结构。

1. DPLR结构的目的：此部分的目的是为了解决问题2，并计算HIPPO SSMs的状态空间核心。为此，需要将HIPPO SSMs简化为一个已知高效的结构化状态空间模型（SSM）。
2. 理想的结构化方案：理想情况下，如果HIPPO SSMs能被转换为对角线结构化（如第3.2节所述），这将是最理想的。根据提议3.4，可以期望应用状态空间转换，即将HIPPO SSMs转换为对角线结构化形式。

6.2.1 HIPPO Cannot be Diagonalized

1. HIPPO矩阵的定义和重要性：HIPPO指定了一类特定的矩阵$\mathbf{A}\in {\mathbb{R}}^{N\times N}$，当这些矩阵纳入(2.1a)时，它允许状态$x(t)$记忆输入$u(t)$的历史。这个类别中最重要的矩阵是HIPPO-LegS，它在指数度量上进行正交化。

   

2. HIPPO矩阵不能被对角化的问题：遗憾的是，对HIPPO矩阵进行简单的对角化是不可行的，因为这会导致在状态大小$N$中的项指数级增长，使得状态空间转换在数值上不可行（例如，定理3.3中的CV将无法计算）。定理6.6说明，HIPPO矩阵A在方程(6.1)中通过矩阵${\mathbf{V}}_{ij}=\left(\begin{array}{c}i+j\\ i-j\end{array}\right)$对角化，特别地，${\mathbf{V}}_{3i,i}=\left(\begin{array}{c}4i\\ 2i\end{array}\right)\approx 2^{4i}$。

这一节强调了HIPPO矩阵在对角化过程中面临的数值稳定性问题，说明了即使理论上可能存在对角化的方法，但在实际应用中却不可行。

6.2.2 Normal Plus Low-Rank (NPLR) Forms of HIPPO

这一节主要讨论了HIPPO矩阵的NPLR（正常加低秩）形式

1. NPLR形式的动机：由于HIPPO矩阵不是正常矩阵，它不能被完美条件化（即酉）矩阵对角化。这一发现促使了对NPLR矩阵类别的定义。这类矩阵可以被分解为正常矩阵和低秩矩阵的和。
2. 定理 6.7：NPLR表示的HIPPO矩阵：所有在第4章和第5章中定义的HIPPO矩阵都有一个NPLR表示，这与DPLR（对角加低秩）表示在酉等价。具体地，矩阵A可以表示为 $\mathbf{A}=\mathbf{V}\mathbf{\Lambda }{\mathbf{V}}^{\ast }-\mathbf{P}{\mathbf{Q}}^{\top }=\mathbf{V}\left(\mathbf{\Lambda }-\left({\mathbf{V}}^{\ast }\mathbf{P}\right){\left({\mathbf{V}}^{\ast }\mathbf{Q}\right)}^{\ast }\right){\mathbf{V}}^{\ast }$，其中$\mathbf{V}\in {\mathbb{C}}^{N\times N}$是酉的，$\Lambda$是对角线的，$\mathbf{P},\mathbf{Q}\in {\mathbb{R}}^{N\times r}$是低秩因子分解。这些矩阵HIPPO-LegS、LegT、FouT都满足$r=1$或r = 2。特别是方程(6.1)是NPLR的，其中$r=1$。

这一节强调了HIPPO矩阵在结构化状态空间模型中的NPLR形式，这是一种将矩阵分解为正常和低秩部分的方法，有助于理解HIPPO矩阵的结构特性和数学性质。

6.2.3 DPLR Form

这一节主要讨论了HIPPO矩阵的对角加低秩（DPLR）形式。

1. DPLR形式的定义：NPLR（正常加低秩）矩阵可以转换为DPLR（对角加低秩）形式。在方程(6.2)中的表达式 $\mathbf{\Lambda }-\left({\mathbf{V}}^{\ast }\mathbf{P}\right){\left({\mathbf{V}}^{\ast }\mathbf{Q}\right)}^{\ast }$ 被称为HIPPO矩阵的DPLR形式。这种形式现在是在复数域$\mathbb{C}$上，而不是实数域R上。
2. DPLR形式的复杂性：定理3.5和3.6描述了当$A$处于DPLR形式时，状态空间模型（SSMs）的复杂性。这一部分强调了对于状态空间模型，将其表示为DPLR形式所涉及的数学和计算上的挑战。

此节强调了HIPPO矩阵的对角加低秩形式，这是一种特殊的数学表示方法，有助于理解HIPPO矩阵在状态空间模型中的应用和计算特性。

6.2.4 Hurwitz DPLR Form

这一节主要讨论了HIPPO矩阵的Hurwitz对角加低秩（DPLR）形式

1. Hurwitz DPLR形式的定义：在第3.3.3节中，讨论了DPLR形式的一个小改动，这个改动确保矩阵$A$是负半定的，因此稳定。这表明即使经过这种修改，HIPPO矩阵仍然满足这种更受限的结构。
2. 推论 6.8：将定理 6.7 改进为 Hurwitz 形式：第5章中所有三个HIPPO矩阵都与形式为 $\mathbf{A}=\mathbf{\Lambda }-\mathbf{P}{\mathbf{P}}^{\ast }$ 的矩阵酉等价，其中$\mathbf{\Lambda }$是对角线的，$\mathbf{P}\in {\mathbb{R}}^{N\times r}$，$r=1$或r = 2。此外，$\mathbf{\Lambda }$的所有项的实部为0（对于LegT和FouT）或−1/2（对于LegS）。

这一节强调了HIPPO矩阵的Hurwitz DPLR形式，这是一种特殊的数学表示方法，有助于确保矩阵的稳定性，同时保留了HIPPO矩阵的主要特性。

6.3 Diagonal Approximations of HIPPO

该节指出，HIPPO方法可以被写成一种结构化的DPLR形式，并且当通过S4结构整合到状态空间模型（SSM）中时，可以非常高效地计算。这表明HIPPO方法的结构化表示不仅可以保持其特性，还可以在计算上更加高效。

6.3.1 Forms of HIPPO-LegS

这一节主要讨论了HIPPO-LegS矩阵的不同形式。

1. HIPPO-LegS矩阵的对角线近似：这一节提出了对HIPPO-LegS矩阵的对角线近似方法。这种近似是通过简单地移除DPLR（对角加低秩）表示中的低秩部分来定义的。这种近似保留了HIPPO的理论特性，并且能够在实践中保持完整的S4-DPLR + HIPPO模型的性能。

2. HIPPO-LegS的具体表示：

   

这一节提供了对HIPPO-LegS矩阵如何通过对角线近似来简化其表示和计算的深入洞察，这对于理解HIPPO方法在实际应用中的潜在优势和改进空间非常重要。

6.3.2 S4D Instantiations

讨论了HIPPO-LegS的不同形式，包括S4D-LegS、S4D-Inv、S4D-Lin。这些形式是对HIPPO-LegS矩阵的不同近似和实现，旨在简化计算并提高效率。

1. S4D-LegS：这是对HIPPO-LegS矩阵进行对角线近似的结果。它保持了原始S4模型的数学解释，并能在无限状态维度的极限下恢复相同的SSM基础。
2. S4D-Inv：这是进一步简化S4D-LegS的方法，通过分析A(D)的结构，特别是它的实部和虚部。
3. S4D-Lin：这是对S4-LegS的另一种更简单的近似，采用不同的缩放法则来处理虚部。

这些变体各有其特点和优势，旨在在保持理论性能的同时简化模型的实现和计算。每种方法都有其独特的数学解释和实际应用场景