原文：
towardsdatascience.com/structured-state-space-models-visually-explained-86cfe2757386
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/a3551757a56352eabdce39ea6b99502a.png
图片由 Sascha Kirch 提供。
这是我的新多部分系列的第二部分 🐍 迈向图像、视频和时间序列的 Mamba 状态空间模型。
状态空间模型，几十年来在许多工程学科中广为人知，现在正在深度学习中崭露头角。在我们迈向 Mamba 选择性状态空间模型及其在研究中的最新成就的过程中，理解状态空间模型至关重要。而且，正如工程中经常发生的那样，正是细节使得理论概念在实践中的应用成为可能。除了状态空间模型之外，我们还必须讨论如何将它们应用于序列数据，如何处理长距离依赖，以及如何通过利用某些矩阵结构来有效地训练它们。
结构化状态空间模型为 Mamba 构建了理论基础。然而，它们与系统理论和高级代数之间的联系可能是采用这一新框架的障碍之一。因此，让我们将其分解，确保我们理解关键概念，并通过可视化来揭示这一新旧理论的一些启示。
即使你最终没有使用状态空间模型，了解一些技巧，比如为什么我们需要加快矩阵乘法，以及我们如何利用矩阵的某些结构来实现这一点，也肯定会提升你作为工程师或开发者的技能。
图像、视频和时间序列的 Mamba 状态空间模型
概述
第一部分结束后的内容
状态空间模型概述 2.1 状态空间模型的连续时间表示 2.2 状态空间模型的离散化 2.3 SSM 的递归和卷积表示
启用长距离关系 3.1 状态大小很重要 3.2 HiPPO 框架
对状态矩阵施加结构以提高效率 4.1 什么是结构化矩阵以及为什么它们很重要？ 4.2 整合一切 – S4：结构化状态空间序列模型
总结 5.1 结构化状态空间模型的局限性 5.2 继续阅读第三部分
进一步阅读与资源
1. 第一部分结束后的内容
迈向图像、视频和时间序列的 Mamba 状态空间模型
在第一部分，我们回顾了循环神经网络（RNNs）和 Transformer 的优缺点，以说明为什么我们需要一个新的模型架构。
我们说过 RNNs 在推理时间快但在训练时慢，而 Transformer 在训练时快但在推理时慢。我们想要的是一个在训练和推理时都快的模型，同时与 Transformer 的性能具有竞争力。
2. 状态空间模型概述
Mamba 通过学习其各种矩阵，在深度学习中构建了使用状态空间模型的想法。因此，在探讨“结构化”部分之前，让我们简要介绍一下状态空间模型。状态空间模型可以用连续时间表示来处理连续信号，或者可以离散化以处理离散数据序列。
这是我们主要感兴趣离散化状态空间模型，因为它像循环神经网络（RNNs）和 Transformer 一样，一个离散 SSM 处理数据序列，例如文本标记或模拟时间信号的样本。
2.1 状态空间模型的连续时间表示
连续时间状态空间模型描述了输入信号通过具有状态的系统传播与产生的输出信号之间的关系。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/637b0df4c587809caed2c024f6409c51.png
图 1：状态空间模型的高级功能。图片由 Sascha Kirch 提供。
这是一个从系统理论中借用的想法。输出取决于输入和系统的当前状态，而当前状态取决于先前状态和输入。
这种关系可以通过两个简单的方程有效地表达，其中A、B、C和D是矩阵（稍后我们将看到，这些矩阵将被学习），x(t)是输入信号，y(t)是输出信号，而h(t)和h’(t)分别是当前状态和更新后的状态。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/69ee5e889495623f4c4a775ee1db27fe.png
方程式 1：连续时间 SSM 的状态和输出方程。
矩阵D将输入x(t)转换并映射到输出y(t)，这通常包含在经典 SSM 的输出方程中。然而，当将 SSM 应用于深度学习时，我们移除了这种转换，并将其建模为一个简单的跳跃连接，从而简化了 SSM。
我们可以将这些方程表示在一个类似于这样的框图中：
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/d7be8818c5ab9d1cabb2f4c49d8c4c1e.png
图 2：连续时间状态空间模型的框图。图片由 Sascha Kirch 提供。
这是 SSM 的连续时间表示。然而，有两个困难：
找到模型状态h(t)的解析解具有挑战性，
由于我们在计算机上工作，我们通常处理离散信号，而不是连续信号。
2.2 状态空间模型的离散化
因此，我们实际上是在将一个函数x(t)映射到另一个函数y(t)，而不是将一个函数映射到另一个函数。这意味着我们需要一个可以处理离散信号的离散化状态空间模型。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/9d3e2bccb55815656e84af05a9b9a71e.png
图 3：离散状态空间模型的高级功能。图片由 Sascha Kirch 提供。
要做到这一点，我们需要对矩阵A和B进行离散化。这导致了我们状态空间模型的离散化表示，现在包括一个离散化的状态h[k]。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/7e4d8ba1d7f13ea4950e09dfd9f1e3ce.png
方程 2：离散化 SSM 的状态和输出方程。
为了离散化我们的状态空间模型，我们使用了一种在系统理论中经常使用的技巧：我们使用一个零阶保持（ZOH）模块将我们的离散信号x[k]转换为连续时间信号x(t)，然后将其输入到连续的 SSM 中。SSM 将生成一个连续输出y(t)，然后我们对其进行采样以获得离散输出序列y[k]：
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/6ee465a4db89f6566cd55487791ce9c3.png
图 4：离散化状态空间模型的框图。图片由 Sascha Kirch 提供。
为了离散化状态空间模型，我们必须对矩阵A和B进行离散化。最初，大多数基于深度学习的状态空间模型，如LSSL 和 S4，使用了双线性离散化，而后来的一些，如DSS 和 Mamba，则使用了零阶保持（ZOH）离散化。目前，我们关心的是我们需要进行离散化，并且可能存在多种方式进行离散化。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/1b3f17bec6b7a0735b8e73ab3b9206ac.png
方程 3：零阶保持（ZOH）离散化与双线性离散化的比较。
A和B是 SSM 的连续时间矩阵，I是单位矩阵，Δ 是表示输入分辨率的步长，在深度学习设置中是学习的。
注意，离散化使我们能够使用 SSM 的时间连续公式并将其应用于计算机中使用的离散信号。这意味着我们变得对分辨率不变，因为我们学习了底层的时间信号。查看 S4ND 论文以获取更多关于该内容的阅读材料。
Sascha Kirch 的论文解读
2.3 SSM 的循环和卷积表示
你有没有注意到状态方程中的循环性，因为h[k]依赖于h[k-1]？事实上，从这种意义上讲，状态空间模型非常类似于 RNN：
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/3e31a46e1712a05f9b7899d399c3cecb.png
图 5：离散状态空间模型的循环表示。图片由Sascha Kirch提供。
矩阵A、B和C是时间不变的。这意味着每个输入和每个状态都使用相同的参数进行转换，无论时间步长如何。线性时不变系统，或简称为 LTI 系统，具有一个非常整洁的特性：它们可以表示为卷积！
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/796037dd74ccf7776a8d82b190282883.png
方程 4：离散状态空间模型的卷积核。
核由矩阵C和离散矩阵A和B构建。请注意，核的长度L与输入序列相同。核与输入x卷积以获得输出y。为了在输出端保持与输入端相同的序列长度，通常需要对输入进行填充。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/9f2125122ea66dda6ea1f21584dffd73.png
图 6：离散状态空间模型的卷积表示示例。图片由Sascha Kirch提供。
这真是太棒了，因为我们知道卷积是可并行的，而且 GPU 和其他硬件加速器都经过优化，可以快速计算它们。
为了总结本节，利用线性时不变系统的特性，我们最终得到了三种不同的表示 SSM 的方法：
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/b8e552f8bacb5943542562fde1f398d4.png
图 7：状态空间模型的不同表示。图片由Sascha Kirch提供。
SSM 的循环和卷积表示是相同的，因此我们可以决定使用其中之一。
在训练期间，我们会选择卷积表示，因为我们可以通过训练数据提前访问整个序列x[k]和标签y[k]。
在推理期间，我们会选择循环表示，因为我们通常只能访问过去的样本（例如自回归文本生成）。注意，由于我们的下一个输出只依赖于前一个状态和当前输入，推理的缩放与序列长度成线性关系。
现在我们有一个具有高效训练和高效推理的模型。太棒了！😎
是 LSSL（线性状态空间层）论文 展示了我们可以训练一个 SSM，并强调了不同表示之间的联系。
每当 Sascha Kirch 发布内容时，都收到一封电子邮件 🚀_ 每当 Sascha Kirch 发布内容时，都收到一封电子邮件 🚀 想了解更多关于深度学习或只是保持最新状态…medium.com
3. 启用长距离关系
矩阵A负责从上一个状态创建下一个状态。系统的状态是考虑过去状态和新输入的多次更新后的结果。可以说，状态包含了整个输入历史的信息。不难想象，对于非常长的序列，可能很难准确地记住所有输入。那么，我们如何处理长距离依赖？
3.1 状态大小很重要
让我们先谈谈状态，它的压缩以及推理的效率。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/df87c9ded69107403e70630f9cb81647.png
图 8：RNN、SSM、Transformer 和期望模型的效率与性能对比。图片由 Sascha Kirch 提供。
RNN 和 Transformer 都可以从理论上处理无限长的数据序列（有时人们说：这些模型是无界的）。RNN通过将整个历史压缩成一个低维状态来实现这一点。这对于推理来说非常高效，因为现在下一个状态只依赖于上一个状态和当前输入。RNN 的性能取决于模型学习进行这种压缩的好坏。此外，对于长序列，尤其是由于梯度消失，即使对于使用门控机制（如 GRUs 和 LSTMs）的更高级 RNN 来说，这也是一项不容易完成的任务。Transformer的另一方面则完全不压缩历史。它通过其注意力机制考虑整个历史（如果不是像自回归生成中那样掩码，则包括未来）。这允许 Transformer 并行训练，但同时也随着序列长度的增加，推理时的内存和计算需求激增。
❓ 所以，再次问：我们如何处理具有高度压缩状态的长距离依赖，同时在训练过程中避免困难，并实现高性能？
一种解决方案：我们使用 HiPPO（高阶多项式投影算子）框架初始化我们的矩阵。
3.2 HiPPO 框架
主要思想是将状态方程中的矩阵A和B参数化，以便我们可以在任何给定点重建历史（即从时间t=0到t=now的输入信号）。具体来说，如果状态表示为某些正交多项式的某些系数（例如，勒让德多项式），我们只需评估多项式即可重建直到当前时间的输入信号。
例如：我们有一个输入函数x(t)。在任意给定的时间步ti，我们想要找到一个重建g(ti)，它近似于t<ti之前的所有t的x(t)。近似g(ti)完全由可以插入勒让德多项式以重建g(ti)的系数c(ti)描述。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/1f55c8ff10a1d43d02b9afca96445caf.png
图 9：HiPPO 框架中输入函数 x(t) 的近似 g(t)。图片由 Sascha Kirch 提供。
这一点很重要：我们不仅可以访问输入的压缩表示，我们还可以在任何给定时间重建输入。尽管如此，它仍然是一个具有线性复杂度所有优点的循环模型。
这些系数随时间演变，因为对于每个ti，我们都有不同的系数c(ti)。那么，我们如何得到这些系数呢？实际上，c(t)可以定义为常微分方程（ODE）：
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/0e4df073d832ed4b5803855d36302fc9.png
方程 5：HiPPO 框架系数的常微分方程（ODE）。
通过一些数学技巧（例如，假设x和g的正交性）以及选择不同时间步长的加权µ(t)，这个问题就被完全定义了。我们可以求解常微分方程（ODE）并得到满足我们加权选择的A(t)和B(t)，然后在我们状态空间模型（SSM）中使用它们。
已提供了不同的 HiPPO 矩阵作为不同过去样本加权和不同多项式基的解决方案。
LegT：具有滑动窗口的勒让德多项式
LegS：具有恒定加权的勒让德多项式。
LagT：具有指数衰减的拉格朗日多项式
LegS 是后续工作中使用最突出的一个。以下是矩阵A的解决方案：
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/1354cb00cf6b97a2a60bdcc2541087ec.png
方程 6：N=4 的 HiPPO 矩阵定义和示例。
一个简单的 Python 实现可能看起来像这样：
defget_hippo_A(N:int)->np.ndarray:A=np.zeros((N,N))forninrange(N):forkinrange(N):ifn>k:A[n,k]=(2*n+1)**0.5*(2*k+1)**0.5elifn==k:A[n,k]=(n+1)returnA之前提到的 LSSL 显示，如果用 HiPPO 矩阵初始化，可以训练一个 SSM，并且它在所展示的实验中甚至超过了 Transformer 的性能。然而，其昂贵的计算和内存需求使得 LSSL 作为通用序列模型不可行。直到 S4（结构化状态空间序列）模型的引入，这个问题才得到解决，我们将在下一章中讨论。
4. 对状态矩阵施加结构以提高效率
现在是时候讨论结构化状态空间模型中的结构是什么以及为什么它很重要了。我们首先尝试获得一些关于结构是什么以及为什么它很重要的非常基本的直觉，然后我们将更详细地研究 S4 模型。
4.1 结构化矩阵是什么？为什么它们很重要？
让我们考虑一个标准的NxN元素方阵。我们将关注两个方面：
我们需要多少内存来表示这个矩阵，
我们能有多快地乘以两个矩阵，这里的“快”可以由几个因素决定，比如能够并行化计算、利用冗余以及最初处理的值更少。
我们需要多少内存？让我们从存储矩阵的内存需求开始。我们可以有不同的矩阵类型。有时它们已经以某种形式存在，有时我们可以利用线性代数的规则将它们转换成某种形式。让我们看看一些常见的矩阵类型：
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/e7985566a1d7138200fcb0f8cb882d96.png
图 10：几种常见矩阵类型需要存储的值数量。图片由 Sascha Kirch 提供。
我们看到，在最坏的情况下，对于一个密集矩阵，我们需要在内存中存储N²个值。不仅如此，我们还必须从 CPU 加载N²个值到 GPU，然后在 GPU 上应用某些矩阵运算。
注意，我们在 CPU 端（即在 RAM 或磁盘上）需要存储的值数量与我们在 GPU 上用于计算的内存占用之间存在差异。为了使这一点更清晰，让我们看看低秩分解。一些矩阵可以被分解为 2 个NxR的低秩矩阵，其中R是这些矩阵的秩。我们现在只需要为每个矩阵存储NR个值，因此总共需要2NR个值。这也意味着我们只需要从 CPU 加载2NR个值而不是N²个值到 GPU。但是，如果我们想用原始矩阵进行计算，我们首先需要在 GPU 内部重建原始矩阵（有些人称之为实体化它们），这意味着我们仍然需要在 GPU 缓存中存储N²个值。
另一方面，像对角矩阵或稀疏矩阵这样的其他表示可以减少两个方面：磁盘上存储的值以及需要在 GPU 上执行矩阵运算所需的值的数量。
需要多少计算量？现在让我们考虑一个标准的矩阵乘法C=AB，其中A和B要么是密集矩阵，要么是对角矩阵，而C是矩阵乘法的结果：
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/6641a0d6c9b6097989854956ae8a78cd.png
图 11：不同矩阵类型矩阵乘法的操作数。图片由 Sascha Kirch 提供。
对于A和B都是密集矩阵的情况，使用我们所有人都知道的矩阵乘法标准算法，我们需要N³次乘法和(N-1)N²次加法。
在A是密集矩阵而B是对角矩阵的情况下，这减少到N²次乘法和0次加法。
在A和B都是对角矩阵的情况下，我们只需要N次乘法，即对角线元素的逐元素乘法。虽然这在硬件资源方面是可取的，但请注意，对角矩阵可能不如密集矩阵表达性强，因为它有更少的值。像所有事情一样，这是一个权衡。
计算的并行化如果你仔细观察密集矩阵乘法，你可能会注意到计算输出C的单行涉及到B的所有值，但只有A的单行。天真地，你可以在 GPU 的不同核心上并行计算C的每一行来加速计算。这意味着每个核心只需要N²次乘法和(N-1)N次加法，而不是单核心情况下的N³次乘法和(N-1)N²次加法，因为我们现在有一个向量-矩阵乘法，而不是矩阵-矩阵乘法。
简而言之为了总结，以NxN的方阵为例，通过尝试利用矩阵的某些特性，我们希望将它们带入一种结构化的形式，这样我们就不需要存储和加载超过N²个值，并且对于矩阵乘法，我们需要的乘法次数不超过N³，加法次数不超过(N-1)N²。
当然，还有许多其他我们可以利用的矩阵结构，它们各自都有优缺点。在整个关于 Mamba 状态空间模型的系列中，我们将遇到许多这样的结构，它们都服务于同一个目的：在保持最佳性能的同时，减少资源需求。
4.2 将所有内容整合 – S4：结构化状态空间序列模型
S4 通过对涉及的矩阵施加一定的结构，旨在使 SSMs 更高效。通过这样做，它可以应用已研究的代数规则来降低涉及计算的复杂性并节省内存，特别是对于卷积视图的核K的重复计算。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/2c807f740d7cd24a0e45bfea09281978.png
图 12：S4 结构化状态序列模型。图片由 Sascha Kirch 提供。
与 LSSL 相比，S4 已经显示出 30 倍的速度提升和 400 倍的内存使用减少，同时性能也得到了提升。
S4 论文涉及许多高级线性代数，在这篇博客文章中我将主要跳过。如果你想要深入了解数学和代码实现，我推荐你查看 The Annotated S4。
回想一下，为了计算 SSM 的卷积表示中的卷积核K，需要对离散化方阵A进行重复的矩阵乘法。记住，核的长度L与输入序列x相同，并且随着核中每个元素的加入，矩阵A的幂次会增加。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/796037dd74ccf7776a8d82b190282883.png
方程 6：离散化状态空间模型的卷积核。
为了明确起见，长度为L=101的序列将会有矩阵的100次幂！以及99，98，……，以此类推。当然，我们可以重用以前的结果，但这仍然需要大量的计算。
S4 的主要目标是通过在矩阵A上施加某种结构来减少计算核K所需的资源数量。
理想情况下，矩阵A将是一个对角矩阵，因为正如我们在图 11 中之前看到的，两个对角矩阵相乘会产生另一个对角矩阵，只需要N次乘法。
由于 HiPPO 矩阵不能以数值稳定的方式对角化（如 S4 论文中所述），接下来最好的办法是将它分解为正加低秩（NPLR）形式，然后可以通过共轭进一步简化为对角加低秩（DPLR）形式。请注意，我们希望保留 HiPPO 矩阵，不将其转换为另一个可对角化的矩阵，因为我们需要它来处理长距离依赖关系。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/59956aef50be463c0b68fc80c0273e6f.png
图 13：比较矩阵 A 的不同表示及其计算卷积核 K 的复杂度。图片由 Sascha Kirch 提供。
当矩阵A处于 DPLR 形式时，我们可以应用许多复杂的技巧。总之，核是在频域中计算的，而 DPLR 形式允许将矩阵重新排列成柯西矩阵，对于这种矩阵，存在许多硬件优化的算法。
此外，使用 Woodburry 标识符可以将A的多个幂次的计算转换为单个逆矩阵。最后，应用快速傅里叶变换（FFT）以获得最终的卷积滤波器K。
因此，之前是A, B, C的函数的SSM(A,B,C)，现在变成了SSM(**Λ-*PP,B,C)，其中Λ,P, B和C是 S4 的可学习参数。
如果你对原始实现感兴趣，这里提供了dplr()函数和nplr()函数，它们位于S4 仓库中。
5. 总结
5.1 摘要
状态空间模型（SSMs）是序列模型。
SSMs 可以被描述为循环模型或卷积模型
由于它们只依赖于当前输入和前一个状态，它们在循环表示中快速推断，并且因为这种表示线性扩展，所以它们是快速的。
由于它们的卷积表示可以并行化，因此它们训练得很快。
由于它们是 LTI 系统，任何时间步的每个输入都被同等对待，这意味着通过相同的矩阵进行转换。
结构化状态空间序列模型，即 S4，在其矩阵上使用一定的结构，使 SSM 更快，允许它们在计算和内存需求上不那么严格。
感谢！
🤝我想要感谢你，我故事书的读者，感谢你们的持续支持，为我点赞，并关注我，以便不错过我的任何最新文章。正是你们激励我继续写作，深入研究复杂话题，并通过写作成为更好的工程师/研究人员。你们太棒了！
5.2 结构化状态空间模型（Structured SSMs）的局限性
然而，有一个大问题：由于矩阵A、B和C是不变的，意味着对所有输入都是相同的，因此 SSM 无法对不同时间步的不同输入进行推理。
5.3 继续阅读第三部分
Mamba：选择性状态空间模型
在第三部分中，我们将发现 Mamba 如何通过向模型架构添加选择性来解决结构化状态空间模型在输入序列中无法区分样本的问题。❓但我们得到了什么，代价是什么？
6. 进一步阅读与资源
Mamba 状态空间模型用于图像、视频和时间序列
Sascha Kirch 的论文解读
博客文章
Mamba 和状态空间模型的视觉指南
S4 的注释
HiPPO：具有最优多项式投影的循环记忆
论文：
S4：使用结构化状态空间高效建模长序列 by A. Gu et.al.，2021 年 10 月 31 日
S5：简化状态空间层用于序列建模 by J. Smith et. al.，2022 年 8 月 9 日
HiPPO: 具有最优多项式投影的循环记忆 by A. Gu et.al., 17 Aug. 2020
LSSL: 结合循环、卷积和连续时间模型与线性状态空间层 by A. Gu et. al., 26 Oct. 2021
DSS: 对角状态空间与结构化状态空间一样有效 by A. Gupta et. al., 27 Mar. 2022
S4D: 对对角状态空间模型的参数化和初始化 by A. Gu et. al., 23 Jun. 2022
S4ND: 使用状态空间将图像和视频建模为多维信号 by E. Nguyen et. al., 12 Oct. 2020
代码
- S4 代码仓库
