Mamba-360：综述作为长序列建模 Transformer 的替代模型 SSM

24年4月来自微软的论文“Mamba-360: Survey of State Space Models as Transformer Alternative for Long Sequence Modelling: Methods, Applications, and Challenges”。

序列建模是各领域的关键技术，包括自然语言处理 (NLP)、语音识别、时间序列预测、音乐生成和生物信息学。RNN 和 LSTM 历来主导机器翻译、命名实体识别 (NER) 等序列建模任务。然而，鉴于 Transformer 卓越的性能， 其导致这一范式的转变。然而，Transformer 遭受 𝑂(𝑁^2) 注意计算复杂性和处理归纳偏差的挑战。已经提出了几种变体来解决这些问题，它们使用谱网络（spectral network）或卷积，并且在一系列任务上表现良好。然而，它们在处理长序列方面仍然有困难。在此背景下，状态空间模型 (SSM) 已成为序列建模范式中有前途的替代方案，尤其是随着 S4 及其变体（如 S4nd、Hippo、Hyena、DSS、GSS、LRU、Liquid-S4和Mamba 等）的出现。

根据三个范式对基础 SSM 进行了分类，即门控架构、结构化架构和循环架构。本综述重点介绍 SSM 在视觉、视频、音频、语音、语言（尤其是长序列建模）、医学（包括基因组学）、化学（如药物设计）、推荐系统和时间序列分析（包括表格数据）等领域的多种应用。此外，整合 SSM 在 Long Range Arena (LRA)、WikiText、Glue、Pile、ImageNet、Kinetics-400、sstv2 等基准数据集以及 Breakfast、COIN、LVU 等视频数据集，以及各种时间序列数据集上的性能。

Mamba-360 工作的项目页面：https://github.com/badripatro/mamba360

本文根据结构性、循环性和门控性对状态空间模型 (SSM) 进行分类，如图所示。

状态空间模型 (SSM) 已成为 Transformer 的有力替代品，尤其适用于处理长序列。SSM 可以概念化为具有固定长度的 RNN，其长度不会随输入长度而增长。与 Transformer 相比，这在推理速度和计算/内存复杂度方面带来了显著的效率优势。然而，尽管 SSM 具有效率优势，但在某些数据模式中，尤其是在视觉任务中，SSM 往往无法弥补与最先进的 Transformer 的性能差距。SSM 的一个显著缺点是它们牺牲了某些序列处理任务所必需的核心能力，例如复制长输入序列 [63]、上下文学习和归纳头 [109]。

如图说明序列数据建模范式的演变过程，从循环神经网络 (RNN) 、卷积神经网络 (CNN) 到 Transformer 模型、状态空间模型 (SSM)，突出了捕捉时间动态、空间层次和复杂系统交互方面的进步。

弹簧-质量-阻尼器（spring-mass-damper）系统是用来说明动力学和控制理论原理的经典示例。其状态空间模型基本数学公式如下。考虑一个弹簧-质量-阻尼器系统，该系统由质量 𝑚 组成，质量 𝑚 通过弹簧常数 𝑘 和阻尼系数 𝑐 的阻尼器连接到墙壁。目标是使用状态变量描述系统的行为。质量的位移表示为 𝑥，其速度表示为 𝑥̇ ，施加在质量上的外力表示为 𝐹 。

状态空间模型一个观测方程 y = Cx + Du 和一个动态方程 𝑥̇ = Ax + B u。其使用一组一阶微分方程来表示系统的动态。它是描述线性时不变 (LTI) 系统的强大框架。其基本组成部分是：

• 状态向量 (x)：状态向量包含描述系统内部状态的状态变量。将其表示为 𝐱。
• 输入向量 (u)：输入向量表示系统的控制或外部输入。将其表示为 𝐮。
• 输出向量 (y)：输出向量包含可测量的关注量。将其表示为 𝐲。
• 系统动力学：状态动力学由一阶微分方程描述。在时不变 (LTI) 情况下，矩阵 𝐀、𝐁、𝐂 和 𝐃 随时间保持不变。

状态向量 𝐱 包含有关质量位置和速度的信息。输入向量 𝐮 可以表示施加到质量上的外力。输出向量 𝐲 通常是位移 𝑥1。在这个系统中，状态向量 𝐱 表示质量的位置和速度。状态动力学由从牛顿第二定律推导出来的运动方程控制。输入 𝑢 表示施加到质量上的任何外力，在本例中假定为零。

稳定性分析涉及检查矩阵 𝐀 的特征值。可以通过调整控制输入 𝐮 来实现控制设计，以实现所需的行为（例如，抑制振荡）。

本质上，线性动力学模型假设下一时间步的状态是前一时间步的状态以及可能的其他外生输入的线性组合。此外，它还假定输出是状态和输入向量的线性函数。

为了对大型序列进行建模，用状态空间模型 (SSM) 而不是多头自注意，因为它很复杂。状态空间模型 [44, 41] 通常被称为线性时不变系统，它通过隐藏空间 h(𝑡) 将输入刺激 𝑥(𝑡) 映射到响应 𝑦(𝑡)。结构化状态空间序列模型 (S4) 是一类最新的深度学习序列模型，与 RNN、CNN 和经典状态空间模型广泛相关。从数学上讲，连续时间潜状态空间可以建模为线性常微分方程，使用演化参数 𝐴 和投影参数 𝐵 和 𝐶 如下：
𝑥̇ = Ax(t) + B u(t), y(t) = Cx(t) + Du(t)

SSM 的离散形式使用时间尺度参数 Δ 将连续参数 A、B 和 C 转换为离散参数 𝐴、𝐵 和 𝐶，使用固定公式 𝐴 = 𝑓𝐴(Δ, 𝐴), 𝐵 = 𝑓𝐵(Δ, 𝐴, 𝐵)。𝑓𝐴, 𝑓𝐵 是使用零阶保持 (ZOH) 进行此转换的离散化规则。
xk = Axk-1 + B uk, yk = Cxk

如图描绘了状态空间模型的概念，该模型通过一系列一阶微分方程描述系统动力学。

尽管取得了进展，基于注意力机制的 Transformer 仍然难以处理长序列，导致路径-X 任务等长距离基准测试中仍存在未解决的挑战。为了解决这些限制，状态空间模型 (SSM) 为该问题提供了一种有前途的替代方法，S4 等先驱模型是首批有效解决路径-X 问题的模型之一。SSM 可以有效地对长序列进行建模，同时捕获长期依赖关系。关键状态空间模型的分类，如图所示：
• 结构化 SSM：这些模型基于 S4 及其变体，包括 Hippo、H3、Hye-naHierarchy、Liquid-S4、S4nd、DSS 和 Global 卷积及其变体，包括 LongConv、FFTFlashConv 和 SG-Conv，以及某些基础模型，如 LD-Stack 及其衍生模型 S5。它们提供了一种处理长距离依赖关系的原则性方法。
• 循环 SSM：这些模型基于 RNN 及其变体，例如 RWKV、LRU 和 HGRN，为基于注意的序列建模方法提供了替代方案。
• 门控 SSM：GSS、Mega 和 TNN 属于这一类，它们利用门控技术来提高长序列的性能。
• 各种 SSM：MambaFormer、Mamba-Byte 和 Mamba-MoE 探索标准注意机制之外的各种技术，将不同类别的思想结合起来，实现高效的序列建模。


但需要注意的是，例如，Mamba 既源自 Hippo（一种结构化 SSM），也采用了门控技术。图中用箭头表示。同样，结构化类别中的基础模型之一 GSS 也源自 DSS，但也使用门控。同样，S5 源自基础模型之一 LDStack 和 S4。状态空间模型（SSM）为处理长序列提供了有希望的解决方案，它们的效率和有效性使它们在某些情况下成为基于注意的 Transformer 的宝贵替代方案。

1. 结构化状态空间模型 (SSM) 包含各种序列建模的创新方法，包括 S4、HiPPO、H3 和 Liquid-S4。这些模型利用多项式投影算子、多输入多输出系统和卷积核等复杂机制来有效捕获长距离依赖关系。它们在各种基准测试中都表现出了极具竞争力的性能，展示了它们在处理序列数据方面的有效性，并提高了计算效率。
2. 门控状态空间 (GSS)、Toeplitz 神经网络 (TNN) 和 Mamba 代表了门控 SSM 领域的创新方法。GSS 利用门控单元优化 FFT 操作，实现高效的序列处理和具有竞争力的性能。TNN 引入位置编码的 Toeplitz 矩阵进行token混合，显著降低了时空复杂度，同时保持了最先进的结果。Mamba 结合门控 MLP 和硬件-觉察算法解决了传统 SSM 中的计算效率低下问题，与传统 Transformer 相比，具有线性时间复杂度和更高的效率。
3. 循环 SSMS 包括线性循环单元 (LRU) 和分层门控循环神经网络 (HGRN)。LRU 及其扩展 Griffin 和 Hawk 突出了线性循环、MLP 块和注意机制在增强长序列建模方面的有效性。相反，HGRN 为线性 RNN 引入动态遗忘门，从而显著提高了效率，并在一系列基准测试中具有竞争力。
4. MoE 已成为提升大语言模型 (LLM) 性能的重要方法。人们做出了许多努力将 MoE 与状态空间模型相结合，并催生出 BlackMamba [4]、MoE-Mamba [124] 和 Jamba [90] 等创新。
5. 在上下文学习 (ICL) 的探索中，重点在于成功的任务表现与训练数据中存在的信息之间的关系。一个重要的问题是，是否可以利用从该类实例中派生的数据，训练模型有效地参与特定函数类（例如线性函数）的上下文学习。一些研究讨论了这个问题，特别是仔细研究了标准 Transformer 和 Mamba 等专用架构的功能。这些研究共同揭示了 Transformers、Mamba 和 MambaFormer 等专门的架构在有效的上下文学习方面的潜力，有助于更深入地了解它们在该领域的能力和局限性。

最后，提一下SSM在各个领域的应用，如下方面：音频语音、视频、医学、视觉、多模态、语言、时间序列、表格、图、推荐系统和强化学习等。