面向图像、视频和时间序列的 Mamba 状态空间模型（第一部分）

添加图片注释，不超过 140 字（可选）

欢迎来到雲闪世界。这是我的新系列🐍面向图像、视频和时间序列的 Mamba 状态空间模型的第一部分。 是Mamba 就是你所需要的一切？当然，人们长期以来一直认为，A. Vaswani 等人在 2017 年的Attention 中提出的 Transformer 架构就是你所需要的一切。毫无疑问，Transformer 一次又一次地彻底改变了深度学习领域。它的通用架构可以轻松适应各种数据模式，例如文本、图像、视频和时间序列，而且似乎你投入到 Transformer 的计算资源和数据越多，它的性能就越好。 然而，Transformer 的注意力机制有一个重大缺点：它的复杂度为O(N²)，这意味着它与序列长度成二次方关系。这意味着输入序列越大，所需的计算资源就越多，因此通常无法处理较大的序列。 ❓ 问题是：我们能做得更好吗？有没有办法在保持高性能的同时降低O(N²)复杂度？这个新架构会是什么样子？ 🏊‍♂️加入我关于 Mamba 状态空间模型的多部分深度探讨系列，我们将发现为什么 Mamba 是一种适用于许多应用程序和数据模式的有趣的新通用方法。

大纲 本系列内容是什么？

1. 为什么我们需要一个新模型？

2. 结构化状态空间模型

3. 曼巴来了

4. Vision Mamba — Mamba 在图像上的应用

5. VideoMamba — 理解视频

6. MambaMixer — 增加频道选择性

7. Mamba-2 — 缩小注意力机制与 SSM 之间的差距

8. 结论

9. 进一步阅读和资源

1. 这个系列是关于什么的？ 显然，这一切都与 Mamba 有关，据说该模型将取代强大的 Transformer！但它是如何工作的？论文中没有提到的细节是什么？你如何处理各种不同类型的数据模式？对于这些问题和更多问题，你将在本系列文章中找到答案。 具体来说，在第 1 部分中，我们将通过了解 RNN 和 Transformer 的优缺点来启动本系列，以激发对新架构的需求。此外，我将为本系列的所有后续部分提供入门知识，介绍许多基于 Mamba 成功的最新研究论文。 在第 2 部分中，我们将研究为 Mamba 提供基础的结构化状态空间模型 (SSM)。我们会发现，虽然它们解决了O(N²)复杂度问题，但还出现了其他障碍，例如如何处理长距离依赖关系以及如何提高其性能。 我们在第 3 部分继续介绍 Mamba，看看它如何形成一种新的通用架构，这种架构训练速度快、推理速度快，并且具有与强大的 Transformer 相似甚至在某些应用中更出色的性能。 进一步来说，我们将深入研究最近的研究论文，如Vision Mamba、VideoMamba、MambaMixer和Mamba-2，了解 Mamba 如何扩展到图像、视频和时间信号等各种数据模式，以及如何不断改进它。

2.为什么我们需要一个新模型？ 为了理解对新模型架构的需求，让我们快速看一下迄今为止最突出的两个序列模型：Transformer 和循环神经网络 (RNN)：

图 1：Transformers、RNN 的属性以及我们想要找到的新模型架构

最终，我们的目标是建立一个训练速度快、推理速度快的模型，同时占用较少的内存，并能访问该序列中其他样本的所有相关数据（不多也不少）。 让我们在非常高的层面上比较一下 RNN 和 Transformer：

图 2：RNN 工作原理与 Transformer 工作原理的高级抽象比较。

2.1 进一步了解 RNN 在 RNN 中，输入x通过隐藏状态h传播以获得输出y。请注意，输出y[t]仅取决于隐藏状态h[t]和当前输入x[t]，进一步注意当前隐藏状态h[t]取决于先前的隐藏状态h[t-1]和当前输入x[t]。这具有以下含义：

1. 过去的输入样本只能通过隐藏状态来访问，该隐藏状态压缩了过去数据点的信息，并且必须使用新输入为每个新样本进行更新。

2. 我们只能获取过去的样本。

3. 无论输入的长度如何，每个输出的计算次数保持不变，并且总计算次数与输入序列的长度成线性比例。

由于3 ， RNN 推理速度很快，但由于1和2 ，训练期间无法并行化。与 Transformer 相比，状态的高压缩可能会限制 RNN 的性能。由于压缩，有价值的信息可能会丢失，尤其是在考虑长序列时，而且状态需要使用基于梯度的算法进行训练，这可能会导致梯度消失等困难。 2.2 Transformer 怎么样？ 在 Transformer 中，输入序列中的每个样本都会直接关注自身和序列中的每个其他样本。在N 个样本的序列中，每个单独的样本都会计算其对其他样本的关注程度的分数，这就是O(N²)复杂度的原因。这意味着（对于 vanilla 注意力机制）：

1. 不同样本之间根本没有压缩，状态的大小为NxN。

2. 每个样本可以同时处理之前和之后的样本。

3. 如果我们进行自回归采样（例如下一个词预测），我们需要再次重新计算所有注意力图以考虑新样本。

尤其是2是我们能够在训练期间很好地并行化 Transformer 的原因，而1则让它的预测性能如此出色。您可以使用这种让一切都关注一切的蛮力方法解决几乎所有问题，让 Transformer 弄清楚它应该关注什么来提高其性能。由于没有压缩，不难想象我们也会浪费大量资源来计算不相关的事情。 ❓ 问题是我们如何才能两全其美？快速训练、快速采样和内存效率。 🤔Mamba 有可能成为实现这一目标的候选人吗？ 3. 结构化状态空间模型 本系列第二部分的入门知识 让我们从基础开始。Mamba 建立在使用状态空间模型进行深度学习的理念之上。状态空间模型是系统理论中众所周知的模型，它描述了通过状态系统传播的输入信号与产生的输出信号之间的关系。

图 3：连续时间状态空间模型的框图。

由于我们在计算机上处​​理离散数据，因此我们需要离散化状态空间模型。有了离散状态空间模型，我们可以发现两种不同的表示：循环表示和卷积表示。这两种表示是相等的，因为标准 SSM 是线性和时不变 (LTI) 系统。

图 4：状态空间模型的不同表示

我们可以利用这种二元性，使用循环表示进行推理，使用卷积表示进行训练。这意味着推理时间与序列长度成线性比例，同时我们可以在现代硬件加速器上有效地训练模型。 👉您将在第二部分中找到以下内容： 我们将深入研究状态空间模型和状态方程的公式。我们将继续进行离散化，并揭示 SSM 的不同表示。然后，我们换个角度，发现如何使用 HiPPO 框架解决长距离依赖关系，最后，我们到达 S4 模型，即结构化状态空间序列模型，这是迈向 Mamba 的里程碑。我们还将讨论为什么我们甚至关心所涉及矩阵的结构。

4. 曼巴来了 本系列第 3 部分的入门知识 借助 S4（结构化状态空间序列模型），我们最终得到了一个可训练的状态空间模型，其推理时间与输入序列的长度成线性比例，并且可以利用卷积表示进行有效训练。不仅如此，与强大的 Transformer 相比，它还表现出了卓越的性能（至少在所展示的实验中）。 S4 是一个线性系统，我们利用它在表示之间进行切换，这也是它最大的弱点。由于矩阵是时不变的，因此输入的每个样本都被平等对待。这是一个很大的限制，因为我们可能希望以不同的方式处理重要的输入，或者可能忽略不相关的输入。

图 5：选择性状态空间模型中涉及的矩阵的详细视图。

Mamba 正是这样做的：它使所涉及的矩阵随时间变化，以引入我们想要的选择性。因此，我们现在有一个随时间变化的系统，我们无法再使用卷积表示，因此我们失去了快速高效训练的可能性。 Mamba 通过进一步引入选择性并行扫描来解决这个问题，这是一种结合核融合、并行关联扫描和梯度重新计算的技术，可以加速模型的训练。 它进一步提出了一个通用架构，即 Mamba Block，它就像 Transformer 块一样可以堆叠在一起构建深度网络。

图 6：Mamba-Block 的详细框图。

👉您将在第 3 部分中找到以下内容： 我们将讨论 Mamba 的选择机制，并了解其如何影响所涉及矩阵的维度。我们将进一步深入研究选择性并行扫描，并揭示内核融合、并行关联扫描和梯度重新计算背后的原理。此外，我们将了解 Mamba Block，我将向您展示模型的实际训练权重隐藏在哪里以及它们具有哪些维度（老实说，我花了很长时间才明白这一点）

5. Vision Mamba — Mamba 应用于图像 本系列第 4 部分的入门知识 有了新的序列建模通用架构 Mamba Block，历史开始重演。就像 Transformer 通过引入Vision Transformer (ViT)应用于图像一样，Mamba 应用于图像，从而产生了 Vision Mamba (Vim)。

图 7：比较 ViT-Vision Transformer 和 ViM-Vision Mamba 模型。

Vision Mamba 旨在使用 Mamba 状态空间模型定义新的通用视觉主干。 为了将图像表示为一个序列，因为 Mamba 毕竟仍然是一个序列到序列模型，它被分成多个块，展开，然后线性投影并与位置编码相结合，然后输入到 Mamba 编码器块堆栈中。Mamba 的线性缩放定律对于高分辨率图像尤其有趣。因为 ViT 中引入的图像修补是Transformer 中注意力模块复杂度O(N²)的直接结果。由于 Mamba 具有O(N)复杂度，因此可以使用一小部分计算资源来处理较小的块或较高的分辨率。事实上，在没有任何注意力机制的情况下，Vision Mamba 比DeiT （数据高效的 Vision Transformer）快 2.8 倍，并且在分辨率为 1024x1024 的图像上节省高达 86.8% 的 GPU 内存。 与大多数时间信号不同，图像不具备这种属性，因为大多数时间信号中的每个新样本都与其过去的样本具有因果关系。在图像中，当前像素之前和之后的各个像素位置之间存在长距离依赖关系。为了处理这种非因果表示，Vision Mamba 引入了一个双向编码器块，其中包括两个 SSM：一个使用正常序列推断，另一个使用反转序列推断。

图 8：ViM — Vision Mamba 编码器模型中使用的 Mamba-Block 的详细框图

👉您将在第 4 部分中找到以下内容： 我们将仔细研究如何将图像转换为序列。我们将讨论补丁化的主要动机，并检查 Mamba 如何处理较小的补丁。我们将进一步回顾位置嵌入和 CLS 标记的插入。最后，我们将讨论双向 MambaBlock，它的内涵是什么以及它是如何实现的。我们将在第 4 部分结束时解释实验，看看它与以前的方法相比如何。

6. VideoMamba — 理解视频 本系列第 5 部分入门 Mamba 必须快速跟上深度学习的最新发展。深度学习拥有比以往更多的计算资源、数据和基础模型，能够处理以前被认为过于复杂而无法处理的数据模式。因此，Mamba 需要多长时间才能处理视频只是时间问题（而这一次是在 ViM 在 arxiv 上发布后不到 2 个月）。 VideoMamba 是 Vision Mamba 的扩展，但应用于视频。它对视频进行分类，而不是对图像进行分类。

图片 9：VideoMamba 的架构。

视频为输入数据添加了额外的时间维度。以前的方法使用 3D 卷积或 Transformers 来处理视频，但这两种方法在计算和内存使用方面都非常昂贵。视频不仅会随着时间引入长距离依赖关系，而且视频也非常冗余，因为在帧之间，像素通常不会改变或只会发生轻微变化（想象一个包含蓝天的场景）。处理这些冗余的一个技巧是关键帧提取，以减少需要处理的帧数。同样，具有O(N)线性复杂度的 Mamba派上了用场。 VideoMamba 的架构与 Vision Mamba 非常相似。它对输入进行补丁处理，进行线性投影，并添加位置编码，然后将其输入到 Mamba Block 中。现在仔细看看位置嵌入。补丁的空间位置嵌入的索引会随着每一帧重置，并且每一帧都会获得一个额外的时间位置嵌入，该嵌入对于帧内的所有补丁都是恒定的。 请记住，对于 Vision Mamba，双向块被引入来处理给定帧内像素的非因果关系。你可以想象，当添加额外的时间维度时，这会变得更加复杂。他们在实验中表明，在处理时间维度之前，先扫描空间维度效果更好。 👉您将在第 5 部分中找到以下内容： 第 5 部分首先概述了模型架构以及它是如何从 Vision Mamba 演变而来的。然后，我们将仔细研究位置嵌入和时空输入的不同扫描方向。我们将讨论 VideoMamba 如何通过自我提炼机制避免过度拟合，并找出如何改进短期动作和长期依赖关系。

7. MambaMixer — 增加频道选择性 本系列第 6 部分的入门知识 我们在本文中尚未提及的一个小细节：Mamba 及其变体将 SSM 单独应用于每个通道。这意味着信息无法从一个通道流向另一个通道。或者换句话说：通道之间没有选择，这限制了 Mamba 在多维数据（例如具有 3 个颜色通道的图像）或多变量时间信号中查找关系的能力。 MambaMixer 通过一个巧妙的小技巧解决了这个问题：它在将特征图输入第二个 SSM 之前对其进行转置：

图 10：转置输入矩阵的效果。

这样，SSM 首先应用于 token 轴，然后应用于 channel 轴。各个块分别称为 TokenMixer 和 ChannelMixer。TokenMixer 和 ChannelMixer 组合起来构成 MambaMixer 块。

图片 11：MambaMixer Block 的总体架构

理论上，每个块可以有多个 SMM，用于对输入块进行不同的扫描方向。多个 MambaMixerBlocks 堆叠在一起以构建 MambaMixer。为了进一步提高性能和训练稳定性，先前的 TokenMixer 和 ChannelMixer 块的输出通过学习加权平均值进行组合。 在论文中，他们使用 MambaMixer Block 构建了两个模型：Vision MambaMixer (ViM2) 和 Time Series MambaMixer (TSM2)。顾名思义，ViM2 应用于图像，而 TSM2 应用于多变量时间序列。

图 12：ViM2（Vision MambaMixer）和 TSM2（Time Series MambaMixer）的架构。

本文受到 MLP Mixer 论文的强烈启发，该论文旨在简化 Vision Transformer (ViT) 的架构。 👉您将在第 6 部分中找到以下内容： 我们从第 6 部分开始，探讨不同渠道之间不交换信息的问题。然后，我们仔细研究实现代币和渠道选择性的通用架构。之后，我们使用 MambaMixer 检查了两个提议的实现，从而产生了 ViM2 和 TSM2 架构。

8. Mamba-2——缩小注意力机制和 SSM 之间的差距 本系列第 7 部分的入门知识 到目前为止，Mamba 已成功应用于各种应用程序和数据模式。它已成为 Transformer 架构的强大竞争对手，在某些任务和基准测试中超越了 Transformer。不过，SSM 仍存在一些缺点：它们面对的是专门针对 Transformer 及其注意力机制量身定制的深度学习基础设施。此外，Mamba 的选择性并行扫描不利用现代 GPU 的矩阵乘法单元。 Mamba-2 旨在重新制定 SSM，以展示与注意力机制的对偶性，以利用这个优化的生态系统，并利用现代硬件加速器的快速矩阵乘法单元。 具体来说，如果用TxT半可分矩阵来表述 SSM ，那么就可以实现这一点。

图 13：SSM 作为半可分矩阵。

输出Y的计算变成了输入X与 SSM 的矩阵乘法，SSM 由矩阵A、B和C定义，它们形成一个半可分矩阵。有了这样的矩阵，就可以进行块分解，也就是说，可以将矩阵划分为QxQ维度的子矩阵。 对角线上的块是三角矩阵。这些块中没有一个共享相同的条目，这意味着它们都可以并行计算。非对角线上的块可以使用低秩分解进一步简化。请注意，相同的因子出现在多个块中，这意味着它们只需要计算一次就可以重复使用，从而进一步加快速度。 通过定义这种状态空间对偶（SSD），作者能够

• 展示 SSM 与注意力之间的联系，

• 使用半可分离矩阵变换重新表述 SSM，从而可以实现更快的算法，使用更少的内存，并利用 GPU 专用的矩阵乘法核心，

• 新 Mamba-2 Block 的设计还允许应用张量并行或序列并行等现代并行化技术，

• 将 Transformers 中突出的多头模式的各种模拟概念引入 SSM，例如多头 SSM (MHS)。

总而言之，SSD 框架最终具有比注意力机制和 SSM 更小的内存占用，并且由于它在其公式核心中使用矩阵乘法，因此速度比 SSM 快 2 到 8 倍。 👉您将在第 7 部分中找到以下内容： 在第 7 部分中，我们将重新讨论为什么半可分离矩阵可以提高速度和内存占用。我们还将看到，低秩分解可以解释为状态空间模型，以及我们如何利用现代硬件加速器。我们将进一步了解 Mamba-2 块，以及它如何实现突出的并行化技术。当然，您还将了解可以实现的多头模式。

9. 结论 Mamba 状态空间模型是 Transformers 和 RNN 的有前途的替代方案，可以构建可应用于任何数据模态的新一代通用架构。 在短时间内，类似 Mamba 的模型已经成功应用于深度学习的许多问题。 与系统理论的紧密联系、所涉及的高级代数以及现代深度学习生态系统针对 Transformer 量身定制的事实可能会成为研究人员采用这一新模型系列的障碍。但由于一些研究人员的辛勤工作和奉献精神，类似 Mamba 的模型在现代硬件加速器上变得越来越快、越来越高效，并且它与标准深度学习层的集成使它越来越容易被纳入您自己的项目。

感谢关注雲闪世界。（Aws解决方案架构师vs开发人员&GCP解决方案架构师vs开发人员）