碾压Transformer！SSM参数少85倍，长序列处理破16K！

在当今信息爆炸的时代，如何高效地处理和理解海量序列数据，一直是人工智能领域面临的核心挑战。无论是语音识别、自然语言处理，还是时间序列预测、视频分析，我们都迫切需要一种能够兼顾长序列建模能力、计算效率和模型泛化性的新方法。正是在这样的背景下，状态空间模型（State Space Models, SSM）开始在深度学习领域崭露头角，成为研究的热点。

状态空间模型最初起源于控制理论，它通过引入一套“状态变量”来描述一个动态系统的内部运行机制。这些状态变量虽然通常是不可直接观测的，但它们却能以一种简洁而强大的方式捕捉系统的历史信息和当前状态。想象一下，一个复杂的机器人在环境中移动，它的位置、速度、加速度等就是一组状态变量，而这些变量的变化则驱动着机器人的行为。类似地，在人类大脑中，科学家们也观察到了所谓的“时间细胞”现象，这些细胞能够有效编码和表征时间信息，使得大脑能够处理和记忆序列事件。

2018年，美国科学家Aaron R. Voelker和Chris Eliasmith在深入研究大脑如何有效表征时间信息时，发现状态空间模型能出色地描述大脑中“时间细胞”的运作机制，尤其是在海马体和皮层区域。他们从神经科学的研究出发，将这一发现创造性地应用于深度学习领域，可以说是最早将SSM引入深度学习的先驱。他们的研究不仅为我们理解大脑工作方式提供了新的视角，也为人工智能模型处理时序数据开辟了新的道路。

新媒网跨境了解到，自此之后，状态空间模型在深度学习领域引起了广泛关注，并涌现出许多令人瞩目的进展。本文将带大家一起探秘深度学习SSM的基础知识，特别是以2021年由美国学者Albert Gu等团队提出的S4模型为核心，深入浅出地讲解SSM的核心概念。S4模型虽然并非目前实践中直接使用的最新架构（更新、性能更优的模型层出不穷），但它作为开创性的工作，对理解SSM至关重要，具有极高的教育价值。

让我们从S4模型出发，共同揭开SSM的神秘面纱，探索它如何在复杂的序列世界中展现出强大的建模能力。

深度学习中的状态空间模型定义

要理解深度学习中的状态空间模型，我们可以借助一个直观的图示来描绘其核心构成。
image/png
图1：连续、时不变状态空间模型视图

从上图可以看出，一个状态空间模型主要由三个随时间t变化的变量组成：

x(t) ∈ C^n：代表n个状态变量，它们是模型的内部记忆或隐藏状态。
u(t) ∈ C^m：代表m个输入变量，是外部对系统施加的控制或信息。
y(t) ∈ C^p：代表p个输出变量，是模型根据当前状态和输入生成的结果。

同时，模型还包含四个可学习的矩阵：A、B、C和D，它们共同决定了系统的动态行为。

A ∈ C^(n × n)：称为状态矩阵，它控制着潜在状态x如何随时间演化，是模型“记忆”能力的关键。
B ∈ C^(n × m)：称为控制矩阵，它决定了输入u如何影响状态变量x。
C ∈ C^(p × n)：称为输出矩阵，它将内部状态x映射到可观测的输出y。
D ∈ C^(p × m)：称为前馈矩阵或命令矩阵，它允许输入u直接影响输出y，形成一个“跳跃连接”。

上述图示可以用以下一组微分方程来精确描述：

x'(t) = A x(t) + B u(t)
y(t) = C x(t) + D u(t)

这里，x'(t)表示状态x对时间t的导数。在文献中，有时也会看到ẋ来表示导数。为了简化书写，由于变量对时间的依赖性是隐含的，所以通常会写成更简洁的形式：

x' = A x + B u
y = C x + D u

在深度学习的SSM中，Du这一项通常被视为一个易于计算的跳跃连接。在很多实际应用中，为了进一步简化模型或侧重于状态演化，它有时被设为0。因此，方程组可以进一步简化为：

x' = A x + B u
y = C x

这个连续的系统方程，在实际的计算机处理之前，必须首先进行离散化。

离散化：从连续到数字的桥梁

离散化是状态空间模型从理论走向实践最关键的一步，甚至可以说是SSM效率的核心所在。通过这一步骤，SSM得以从连续的时间域转换到离散的时间域，从而展现出它的两种强大视角：递归视图和卷积视图。如果说本文有什么核心内容值得牢记，那便是离散化在SSM中的独特作用。

图2：图片来自Albert Gu等学者于2022年发表的博客文章《结构化状态空间：结合连续时间、循环和卷积模型》

在后续的文章中，我们还会探讨多种不同的离散化方法。事实上，不同的SSM架构之间，其核心区别之一就体现在离散化策略上。为了更好地理解SSM的递归和卷积视图，我们将在本文中采用S4模型所提出的离散化方法进行阐述。

递归视图：记忆的链条

要将连续系统离散化，我们可以采用梯形法则。这个方法的原理是将函数f在时间段[t_n, t_n+1]下的曲线区域近似为一个梯形，并计算其面积T = (t_n+1 - t_n) (f(t_n) + f(t_n+1)) / 2。

将SSM方程组的第一行x'_n = A x_n + B u_n视为函数f，并设时间步长Δ = t_n+1 - t_n，我们可以得到：

x_n+1 - x_n = (Δ/2) (A x_n + B u_n + A x_n+1 + B u_n+1)

通过一系列代数推导，我们可以得到离散化的状态更新方程。这里我们假设在很小的时间步长Δ内，控制向量u_n+1 ≈ u_n（即输入在一个短时窗内近似恒定）。

x_n+1 - (Δ/2)A x_n+1 = x_n + (Δ/2)A x_n + Δ B u_n+1
⟺ (I - (Δ/2) A) x_n+1 = (I + (Δ/2) A) x_n + Δ B u_n+1
⟺ x_n+1 = (I - (Δ/2) A)^(-1) (I + (Δ/2) A) x_n + (I - (Δ/2) A)^(-1) Δ B u_n+1

我们已经得到了离散化的SSM方程！为了使其表达更清晰，我们引入新的符号：

Ā = (I - (Δ/2) A)^(-1)(I + (Δ/2) A)
B̄ = (I - (Δ/2) A)^(-1) Δ B
C̄ = C

于是，我们得到了简洁的离散化递归方程：

x_k = Ā x_k-1 + B̄ u_k
y_k = C̄ x_k

这里带有上划线的矩阵符号（Ā, B̄, C̄）是S4模型中引入的约定，用来特指离散化后的矩阵，这在深度学习应用SSM领域已成为一种惯例。这个递归视图清晰地展示了当前状态x_k如何依赖于前一时刻的状态x_k-1和当前输入u_k，体现了模型对历史信息的记忆和传递。

卷积视图：高效并行处理

上述的递归关系，通过迭代展开，也可以被巧妙地表达为一种卷积形式。让我们从系统的第一个方程开始迭代：

x_k = Ā x_k-1 + B̄ u_k
y_k = C̄ x_k

第0步： x_0 = B̄ u_0 (假设初始状态为零)
第1步： x_1 = Ā x_0 + B̄ u_1 = Ā B̄ u_0 + B̄ u_1
第2步： x_2 = Ā x_1 + B̄ u_2 = Ā (Ā B̄ u_0 + B̄ u_1) + B̄ u_2 = Ā^2 B̄ u_0 + Ā B̄ u_1 + B̄ u_2

我们可以观察到，x_k可以表示为关于u_0, u_1, ..., u_k的函数。现在，我们将其代入系统的第二个方程，计算输出y_k：

第0步： y_0 = C̄ x_0 = C̄ B̄ u_0
第1步： y_1 = C̄ x_1 = C̄ (Ā B̄ u_0 + B̄ u_1) = C̄ Ā B̄ u_0 + C̄ B̄ u_1
第2步： y_2 = C̄ x_2 = C̄ (Ā^2 B̄ u_0 + Ā B̄ u_1 + B̄ u_2) = C̄ Ā^2 B̄ u_0 + C̄ Ā B̄ u_1 + C̄ B̄ u_2

至此，一个明显的卷积核结构浮现出来：K̄_k = (C̄ B̄, C̄ Ā B̄, ..., C̄ Ā^k B̄)，它作用于输入序列u_k，形成了K̄ * u的卷积形式。与矩阵类似，这里的K̄上的上划线同样表示它是经过离散化得到的卷积核。这个卷积核在文献中通常被称为SSM卷积核，其长度与整个输入序列等同。这种卷积核的计算可以通过快速傅里叶变换（FFT）高效实现，这将在后续文章中进一步探讨。

三种视图的优势与局限

状态空间模型的三种不同视图——连续、递归和卷积，各自拥有独特的优势和局限性。理解这些特点对于在不同应用场景中灵活选择和优化模型至关重要。

图3：图片来自Albert Gu等学者在S4发布前一周发表的论文《通过线性状态空间层结合循环、卷积和连续时间模型》

连续视图

优点：
1. 自动处理连续数据： 能够天然地处理音频信号、时间序列等连续型数据，即使采样不规则或存在时间偏移，也能轻松应对。这对于许多现实世界的传感器数据和物理过程建模具有巨大的实用价值。
2. 数学分析可行性高： 允许进行严谨的数学分析，例如计算精确的轨迹或构建复杂的记忆系统（如HiPPO），从而为模型的理论研究和性能优化提供坚实基础。
缺点：
1. 训练和推理速度慢： 在连续时间域直接进行计算通常效率低下，无论是模型训练还是推理，都可能消耗大量计算资源和时间。

递归视图

优点：
1. 天然的序列数据归纳偏置： 对于序列数据具有固有的亲和性，能够原则上处理无限长的上下文信息。这使得它在处理需要长距离依赖的语言、语音等任务上表现出色。
2. 高效的推理： 在推理阶段，只需更新当前状态，时间复杂度恒定，非常适合在线或实时应用。
缺点：
1. 训练速度慢： 由于其内在的顺序性，递归计算难以并行化，导致模型训练速度较慢，尤其是在处理长序列时。
2. 梯度消失或爆炸： 在训练过长序列时，容易出现梯度消失或梯度爆炸问题，使得模型难以学习长距离依赖关系。

卷积视图

优点：
1. 局部、可解释的特征： 能够提取数据中的局部特征，并且这些特征通常具有较好的可解释性。这对于理解模型决策过程非常有帮助。
2. 高效的并行训练： 卷积操作可以高度并行化，因此在训练阶段效率极高，能够充分利用现代硬件的计算能力。
缺点：
1. 在线或自回归场景速度慢： 在在线推理或自回归生成任务中，每当有新的数据点输入，都需要重新计算整个输入序列的卷积，导致效率低下。
2. 固定上下文大小： 传统的卷积核大小固定，这意味着其能够捕捉的上下文信息长度是有限的，难以处理极长的序列依赖。

因此，根据具体的应用阶段（训练或推理）以及数据的特性，我们可以灵活地在这些视图之间切换，以充分发挥模型的最优性能。例如，在追求训练速度时，我们更倾向于采用可并行化的卷积视图；而在需要高效在线推理时，递归视图则是不二之选；对于那些采样不规则的连续数据，连续视图的优势便体现得淋漓尽致。这种多视图的灵活性正是SSM模型设计巧妙之处。新媒网跨境认为，这种多维度、高效率的建模方式，预示着未来人工智能在处理复杂动态系统方面将拥有更广阔的应用前景。

矩阵学习：SSM的核心驱动力

在前面构建的卷积核中，C̄（一个行向量）和B̄（一个列向量）都是可学习的参数。关于Ā，我们看到在卷积核中，它以时间步长k的幂次形式出现。直接计算Ā的高次幂可能非常耗时，因此，我们希望Ā能够具备某种特性，使其计算变得高效。最理想的情况是Ā为对角矩阵：

A = [
 λ_1  0  ⋯  0
  0  λ_2  ⋯  0
  ⋮   ⋮  ⋱  ⋮
  0   0  ⋯  λ_n
]
⇒ A^k = [
 λ_1^k  0  ⋯  0
  0  λ_2^k  ⋯  0
  ⋮   ⋮  ⋱  ⋮
  0   0  ⋯  λ_n^k
]

根据线性代数的谱定理，这正是正规矩阵的特点。除了之前提到的离散化选择之外，如何定义和初始化Ā矩阵，也是区分各种SSM架构的关键点之一，这一点我们会在后续的文章中进行详细探讨。

经验表明，如果SSM使用随机初始化的A矩阵，其性能往往不尽如人意。然而，如果基于HiPPO（High-Order Polynomial Projection Operator，高阶多项式投影算子）矩阵进行初始化，模型性能将获得显著提升。例如，在MNIST序列基准测试中，性能可以从60%大幅提升到98%。HiPPO矩阵由S4模型的作者们在2020年的一篇论文中首次提出，并在LSSL论文（2021年）以及S4的附录中被详细阐述。其具体公式如下：

A = [
  1
 -1  2
  1 -3  3
 -1  3 -5  4
  1 -3  5 -7  5
 -1  3 -5  7 -9  6
  1 -3  5 -7  9 -11  7
 -1  3 -5  7 -9 11 -13  8
  ⋮                 ⋱
]
⇒ A_nk = {
  (-1)^(n-k) (2k+1)  if n > k
  k+1               if n = k
  0                 if n < k
}

（这里展示的是HiPPO-LegT版本，关于其不同形式的更多细节，将在后续文章中探讨。）

值得注意的是，HiPPO矩阵本身并非对角矩阵，但它可以分解为一 Normal 矩阵加上一个低秩矩阵（在论文中简称为NPLR，即Normal Plus Low Rank）。S4的作者们在论文中证明，这种类型的矩阵（特别是其幂次）可以通过三种技术高效计算：截断生成级数、柯西核和伍德伯里恒等式（详情可参考论文中的算法1）。具体证明NPLR矩阵可以像对角矩阵一样高效计算的细节，可以在LSSL论文的附录（B和C部分）中找到。

此后，S4的作者们在2022年的论文《如何训练你的HiPPO》中，又对HiPPO矩阵的初始化方式进行了修改。基于这篇论文的模型，通常被称为“S4 V2”或“S4更新版”，以区别于“原始S4”或“S4 V1”。在接下来的文章中，我们将看到其他研究人员（尤其是Ankit Gupta）提出了使用对角矩阵代替NPRL矩阵的方法，这种方法如今因其实现更简单而受到青睐。

实验成果：SSM的强大表现

为了更直观地感受状态空间模型（SSM）的巨大潜力，让我们通过S4模型在多个任务和基准测试上取得的实验结果进行分析。

首先，我们来看一个音频任务，以WARDEN于2018年提出的语音指令（Speech Commands）基准测试为例。

图4：图片来自Albert Gu等学者于2022年发表的论文《对角状态空间模型的参数化与初始化》（S4D），该论文发布于S4之后，但以更结构化的形式再现了S4在该基准测试上的结果。

从上表我们可以观察到几个关键点：

在参数数量大致相同的情况下，S4模型的性能远超其他卷积神经网络（ConvNet）类模型，至少提升了13%以上。
要达到与S4相当的性能水平，卷积神经网络需要多达85倍的参数量，这凸显了S4在参数效率上的巨大优势。
一个在16K Hz数据上训练的卷积神经网络，在应用于8K Hz重采样数据时，性能会大幅下降。然而，S4模型在此重采样任务中仍能保持95%的性能，这得益于其连续视图的特性，只需在测试阶段将Δ值减半即可。

接下来，我们转向时间序列任务（该任务是在S4的修订版中引入的）。

图5：图片来自S4论文附录

S4论文的作者们采用了ZHOU等人在2020年提出的Informer模型的方法论，并展示了S4模型在50种配置中的40种上超越了该Transformer模型。表中展示的是单变量框架下的结果，但在多变量框架下也观察到了相似的卓越表现（详见附录表14）。

再来看看视觉任务，我们关注KRIZHESKY在2009年提出的sCIFAR-10基准测试。

图6：图片来自S4论文附录

S4模型以仅仅10万个参数就在sCIFAR-10上取得了当时的最先进水平（SoTA）。尽管论文没有详细说明其他方法的参数数量，但这一成就无疑彰显了SSM在视觉任务上的强大潜力。

最后，我们考察文本任务，以TAY等学者在2020年提出的长距离竞技场（Long Range Arena, LRA）基准测试为例。

图7：图片来自S4论文附录

LRA基准测试包含6项任务，其中Path-X任务的序列长度高达16K tokens。S4是第一个成功解决这项任务的模型，充分展示了其处理超长序列任务的卓越性能。直到两年多以后，AMOS等学者才在2023年的论文中指出，由VASWANI等人在2017年引入的Transformer模型（未与SSM混合）也能解决这项任务。然而，与SSM不同的是，Transformer模型无法通过65K tokens的PathX-256任务，这再次突出了SSM在超长序列处理方面的独特优势。

不过，S4在文本处理方面也存在一个不足。在MERITY等人在2016年提出的WikiText-103数据集上，S4的困惑度（Perplexity）高于标准Transformer模型（经过优化的Transformer版本困惑度更低）。

图8：图片来自S4论文附录

这可能与文本数据本身的非连续性有关，它不像语音或时间序列那样是从潜在的物理过程采样而来。然而，新媒网跨境获悉，在2023年SSM领域涌现的大量研究工作已经成功弥补了这一差距，使得SSM在文本处理方面也取得了显著进步。

状态空间模型（SSM）的发展历程

状态空间模型在深度学习领域的应用并非一蹴而就，其发展也经历了一个从探索到完善的过程。在S4模型发布的两年前，也就是2019年12月，美国学者VOELKER、KAJIĆ和ELIASMITH团队推出了Legendre Memory Units（LMU）模型，这可以被视为S4模型的早期先驱。

LMU论文中，研究人员提出了LSTM的一种替代方案。由HOCHREITER和SCHMIDHUBER于1997年提出的LSTM，虽然解决了传统循环神经网络的梯度消失问题，但在处理过长序列时（通常限制在100到5000步之间），梯度消失问题依然存在。而LMU模型则展现出处理超过10万步的能力（VOELKER在其博士论文的6.1节中甚至提到可以处理超过10亿步的序列），这在当时是一个突破性的进展。

LMU模型的核心思想是基于常微分方程x'(t) = Ax(t) + Bu(t)进行离散化（论文中x表示为m），并通过欧拉方法实现。其中的矩阵A和B则是通过Padé逼近法获得的，这一方法也为后来的HiPPO框架提供了重要灵感。这个动态系统的关键特性在于，x能够通过勒让德多项式（最高达到d-1次）来表示输入u的滑动窗口。感兴趣的读者可以查阅论文的第2节了解更多细节。如前所述，LMU是Aaron R. Voelker和Chris Eliasmith在2018年发表的神经科学相关模型的深度学习应用。

LMU的后续工作在2021年2月由CHILKURI和ELIASMITH团队进一步推动。他们在这篇论文中展示了如何高效地计算LMU模型。为了实现这一点，他们通过非顺序重写常微分方程来并行化训练（特别是论文第3页），使得模型能够利用标准的控制理论工具（参见论文方程22，以及ÅSTRÖM和MURRAY的详细论述），并将其视为一种卷积形式。在Text8数据集的字符级别建模任务中，他们的模型在参数量仅为DistillBERT一半的情况下，取得了优于SANH等人于2019年提出的DistillBERT的性能。此外，值得注意的是，论文作者通过ZOH（Zero Order Hold，零阶保持）方法对SSM进行离散化，这种方法我们将在后续的文章中进行更详细的探讨。这些研究共同构成了SSM在深度学习领域蓬勃发展的重要基石。

总结与展望

状态空间模型（SSM）无疑是一种充满活力和前景的AI模型，它以三种独特视角——连续视图、递归视图和卷积视图——为我们处理序列数据提供了强大的工具。SSM的巧妙之处在于，它能够根据任务需求和数据特点，灵活地在这些视图之间切换，从而在训练和推理的不同阶段发挥出各自的最佳优势。

这种模型展现出惊人的多功能性，其应用领域涵盖了文本、视觉、音频、时间序列乃至图数据等多个方面。SSM的一大核心优势在于其处理超长序列的能力，通常在参数量更少的情况下，能比传统模型（如卷积神经网络或Transformer）实现更高的计算效率。这种“小而美”的特性，使其在资源受限的环境下更具吸引力。

正如我们所见，现有SSM架构之间的主要差异，往往体现在基础SSM方程的离散化方法，以及A矩阵的定义方式上。这些细微的调整和优化，共同推动着SSM技术的不断进步。

展望未来，随着对SSM理论和实践的深入研究，我们有理由相信，状态空间模型将在人工智能领域扮演越来越重要的角色。它有望在自动驾驶、智能医疗、智能制造等前沿领域，为构建更高效、更智能的系统提供关键技术支持。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/ssm-tops-transformer-85x-fewer-params-16k-seq.html