Hugging Face 的 LOGO

状态空间模型（SSM）简介

社区文章 发布于 2024年7月19日
赞
160
Loïck BOURDOIS's avatar
Loïck BOURDOIS
lbourdois
法文版本可在我的博客上找到.

更新日志
2023-12-14：文章发布。
2024-04-08：修正了打字错误（我的英语不是很好😅）。
2024-06-11：添加了SSM系列博客文章第二篇的链接。
2024-07-18：修正了LaTex排版错误。
2024-09-23：重写了引言部分，并增加了关于深度学习中SSM起源的章节。

前言

我衷心感谢 Boris ALBAR、Pierre BEDU 和 Nicolas PREVOT 同意组建一个SSM主题的工作组，并在此过程中陪伴我探索这种模型。特别感谢 Boris ALBAR 抽出时间审阅这篇博客文章。


简介

状态空间模型（States Spaces Models）传统上用于控制理论中，通过状态变量对动态系统进行建模。

Aaron R. VOELKER 和 Chris ELIASMITH 探讨了大脑如何有效表示时间信息的问题。他们在2018年发表的论文“Improving Spiking Dynamical Networks: Accurate Delays, Higher-Order Synapses, and Time Cells”中发现，SSM是描述大脑中“时间细胞”（特别是海马体和皮层）的优秀模型。

从神经科学出发，他们将研究应用于深度学习领域，并因此（据我们所知）率先在深度学习中使用SSM。有关这项工作的更多细节，请参阅本博客文章末尾的“SSM历史”部分。

在本文中，我们将定义深度学习SSM的基础知识。为此，我们将基于 Albert GU 等人在2021年提出的S4模型，该模型发表在“高效建模长序列的结构化状态空间”中。这并不是一个在实践中原样使用的模型（现在有其他性能更好或更易于实现的SSM）。我们在这里使用它主要是出于教育目的。S4发表一周前，同一批作者发表的LSSL也是该主题的重要信息来源。我们将在未来的博客文章中探讨S4的各种发展。在此之前，让我们先深入了解SSM的基础知识。


深度学习中SSM的定义

让我们用下图来定义SSM：

image/png
图1：连续时不变SSM视图（来源：https://en.wikipedia.org/wiki/State-space_representation）

可以看出，SSM基于三个与时间tt相关的变量：

  • x(t)∈Cnx(t) \in \mathbb {C}^{n} 表示 nn 个状态变量，
  • u(t)∈Cmu(t) \in \mathbb {C}^{m} 表示 mm 个状态输入，
  • y(t)∈Cpy(t) \in \mathbb {C}^{p} 表示 pp 个输出，

我们还可以看到它由四个可学习矩阵组成：A,B,C\mathbf A, \mathbf B, \mathbf C 和 D\mathbf D。

  • A∈Cn×n\mathbf A \in \mathbb {C}^{n \times n} 是状态矩阵（控制潜态 xx），
  • B∈Cn×m\mathbf B \in \mathbb {C}^{n \times m} 是控制矩阵，
  • C∈Cp×n\mathbf C \in \mathbb {C}^{p \times n} 是输出矩阵，
  • D∈Cp×m\mathbf D \in \mathbb {C}^{p \times m} 是命令矩阵，

上图可简化为以下方程组：

x′(t)=Ax(t)+Bu(t)y(t)=Cx(t)+Du(t) \begin{aligned} x'(t) &= \mathbf{A}x(t) + \mathbf{B}u(t) \\ y(t) &= \mathbf{C}x(t) + \mathbf{D}u(t) \end{aligned}

注意：这里我们使用符号 x′x' 来表示 xx 的导数。文献中也可能使用符号 x˙ẋ。

同样，由于变量依赖于时间是隐含的，为了简化，前面的方程通常写成以下形式：

x′=Ax+Buy=Cx+Dumtable> \begin{aligned} x' &= \mathbf{A}x + \mathbf{B}u \\ y &= \mathbf{C}x + \mathbf{D}u \end{aligned}

这个系统可以进一步简化，因为在深度学习SSM中，Du=0\mathbf{D}u = 0 被看作是一个易于计算的跳跃连接（skip connection）。

x′=Ax+Buy=Cx \begin{aligned} x' &= \mathbf{A}x + \mathbf{B}u \\ y &= \mathbf{C}x \end{aligned}

此系统是连续的。因此，在将其提供给计算机之前，必须先对其进行离散化。


离散化

离散化是SSM中最重要的一点，如果不是最重要的话。该架构的所有效率都体现在这一步，因为它使我们能够从SSM的连续视图过渡到其另外两种视图：递归视图和卷积视图。
如果这篇文章有什么值得记住的，那就是这一点。

image
图2：图片来自 Albert GU 等人（2022）的博客文章Structured State Spaces: Combining Continuous-Time, Recurrent, and Convolutional Models

我们将在稍后的文章中看到有几种可能的离散化方法。这是各种现有SSM架构之间的主要区别之一。
对于本文，让我们应用S4中提出的离散化来阐释SSM的两个额外视图。


SSM的递归视图

为了对连续情况进行离散化，我们使用梯形方法，其原理是将函数 ff 在区间 [tn,tn+1][t_n , t_{n+1}] 上代表曲线下的区域近似为梯形并计算其面积 TT：T=(tn+1−tn)f(tn)+f(tn+1)2T=(t_{n+1} - t_n){\frac {f(t_n)+f(t_{n+1})}{2}}。

我们有：xn+1−xn=12Δ(f(tn)+f(tn+1))x_{n+1} - x_n = \frac{1}{2}\Delta(f(t_n) + f(t_{n+1})) 其中 Δ=tn+1−tn\Delta = t_{n+1} - t_n。
如果 xn′=Axn+Bunx'_n = \mathbf{A}x_n + \mathbf{B} u_n（SSM方程的第一行），对应于 ff，那么

xn+1=xn+Δ2(Axn+Bun+Axn+1+Bun+1)⟺xn+1−Δ2Axn+1=xn+Δ2Axn+Δ2B(un+1+un)(∗)⟺(I−Δ2A)xn+1=(I+Δ2A)xn+ΔBun+1⟺xn+1=(I−Δ2A)−1(I+Δ2A)xn+(I−Δ2A)−1ΔBun+1 \begin{aligned} x_{n+1} & = x_n + \frac{\Delta}{2} (\mathbf{A}x_n + \mathbf{B} u_n + \mathbf{A}x_{n+1} + \mathbf{B} u_{n+1}) \\ \Longleftrightarrow x_{n+1} - \frac{\Delta}{2}\mathbf{A}x_{n+1} & = x_n + \frac{\Delta}{2}\mathbf{A}x_{n} + \frac{\Delta}{2}\mathbf{B}(u_{n+1} + u_n) \\ (*) \Longleftrightarrow (\mathbf{I} - \frac{\Delta}{2} \mathbf{A}) x_{n+1} & = (\mathbf{I} + \frac{\Delta}{2} \mathbf{A}) x_{n} + \Delta \mathbf{B} u_{n+1}\\ \Longleftrightarrow x_{n+1} & = (\mathbf{I} - \frac{\Delta}{2} \mathbf{A})^{-1} (\mathbf{I} + \frac{\Delta}{2} \mathbf{A}) x_n + (\mathbf{I} - \frac{\Delta}{2} \mathbf{A})^{-1} \Delta \mathbf{B} u_{n+1} \end{aligned}

（*）un+1unu_{n+1} \overset{\Delta}{\simeq} u_n（控制向量在小的 Δ\Delta 范围内被认为是常数）。

我们刚刚得到了离散化的SSM！
为了使其完全显式化，我们设：

Aˉ=(I−Δ2A)−1(I+Δ2A)Bˉ=(I−Δ2A)−1ΔBCˉ=C \begin{aligned} \mathbf{\bar{A}} &= (\mathbf {I} - \frac{\Delta}{2} \mathbf{A})^{-1}(\mathbf {I} + \frac{\Delta}{2} \mathbf{A}) \\ \mathbf {\bar{B}} &= (\mathbf{I} - \frac{\Delta}{2} \mathbf {A})^{-1} \Delta \mathbf{B} \\ \mathbf {\bar{C}} &= \mathbf{C}\\ \end{aligned}

我们便得到：

xk=Aˉxk−1+Bˉukyk=Cˉxk \begin{aligned} x_k &= \mathbf{\bar{A}}x_{k-1} + \mathbf{\bar{B}}u_k \\ y_k &= \mathbf{\bar{C}}x_k \end{aligned}

带横线的矩阵表示法在S4中被引入，用于表示离散情况下的矩阵，此后在应用于深度学习的SSM领域中已成为一种约定。


SSM 的卷积视图

这个递推关系可以写成卷积形式。为此，只需迭代系统的方程

xk=Aˉxk−1+Bˉukyk=Cˉxk \begin{aligned} x_k &= \mathbf{\bar{A}}x_{k-1} + \mathbf{\bar{B}}u_k \\ y_k &= \mathbf{\bar{C}}x_k \end{aligned}

我们从系统的第一行开始
步骤 0: x0=Bˉu0x_0 = \mathbf{\bar{B}} u_0
步骤 1: x1=Aˉx0+Bˉu1=AˉBˉu0+Bˉu1x_1 = \mathbf{\bar{A}}x_{0} + \mathbf{\bar{B}}u_1 = \mathbf{\bar{A}} \mathbf{\bar{B}} u_0 + \mathbf{\bar{B}}u_1
步骤 2: x2=Aˉx1+Bˉu2=Aˉ(AˉBˉu0+Bˉu1)+Bˉu2=Aˉ2Bˉu0+AˉBˉu1+Bˉu2x_2 = \mathbf{\bar{A}}x_{1} + \mathbf{\bar{B}}u_2 = \mathbf{\bar{A}} (\mathbf{\bar{A}} \mathbf{\bar{B}} u_0 + \mathbf{\bar{B}}u_1) + \mathbf{\bar{B}}u_2 = \mathbf{\bar{A}}^{2} \mathbf{\bar{B}} u_0 + \mathbf{\bar{A}} \mathbf{\bar{B}} u_1 + \mathbf{\bar{B}}u_2
我们有 xkx_k，可以写成由 (u0,u1,...uk)(u_0, u_1, ... u_k) 参数化的函数 ff。

我们来看系统的第二行，现在可以把刚才计算出的 xkx_k 值代入
步骤 0: y0=Cˉx0=CˉBˉu0y_0 = \mathbf{\bar{C}} x_0 = \mathbf{\bar{C}} \mathbf{\bar{B}} u_0
步骤 1: y1=Cˉx1=Cˉ(AˉBˉu0+Bˉu1)=CˉAˉBˉu0+CˉBˉu1y_1 = \mathbf{\bar{C}} x_1 = \mathbf{\bar{C}} ( \mathbf{\bar{A}} \mathbf{\bar{B}} u_0 + \mathbf{\bar{B}}u_1) = \mathbf{\bar{C}} \mathbf{\bar{A}} \mathbf{\bar{B}} u_0 + \mathbf{\bar{C}} \mathbf{\bar{B}}u_1
步骤 2: y2=Cˉx2=Cˉ(Aˉ2Bˉu0+AˉBˉu1+Bˉu2)=CˉAˉ2Bˉu0+CˉAˉBˉu1+CˉBˉu2y_2 = \mathbf{\bar{C}} x_2 = \mathbf{\bar{C}}(\mathbf{\bar{A}}^{2} \mathbf{\bar{B}} u_0 + \mathbf{\bar{A}} \mathbf{\bar{B}} u_1 + \mathbf{\bar{B}}u_2 ) = \mathbf{\bar{C}}\mathbf{\bar{A}}^{2} \mathbf{\bar{B}} u_0 + \mathbf{\bar{C}}\mathbf{\bar{A}} \mathbf{\bar{B}} u_1 + \mathbf{\bar{C}}\mathbf{\bar{B}}u_2
我们可以观察到卷积核 Kˉk=(CˉBˉ,CˉAˉBˉ,...,CˉAˉkBˉ)\mathbf{\bar{K}} _k = (\mathbf{\bar{C}} \mathbf{\bar{B}}, \mathbf{\bar{C}} \mathbf{\bar{A}} \mathbf{\bar{B}}, ..., \mathbf{\bar{C}} \mathbf{\bar{A}}^{k} \mathbf{\bar{B}}) 适用于 uku_k，因此为 K∗uK \ast u。

与矩阵类似，我们对 Kˉ\mathbf{\bar{K}} 应用横线，以表明它是离散化后获得的卷积核。在文献中，它通常被称为 SSM 卷积核，其大小等同于整个输入序列。
此卷积核通过快速傅里叶变换 (FFT) 计算，将在未来的文章中进行解释。


三种视图的优缺点

image
图 3: 图片来自 Albert GU 等人撰写的论文 结合递归、卷积和连续时间模型与线性状态空间层，该论文在 S4 发布前一周发布

SSM 的不同视图各有优缺点——让我们仔细看看。

对于连续视图，优缺点如下
✓ 自动处理连续数据（例如音频信号、时间序列）。这在处理采样不规则或时间偏移的数据时具有巨大的实际优势。
✓ 数学上可行的分析，例如通过计算精确轨迹或构建记忆系统（HiPPO）。
✗ 训练和推理都极其缓慢。

对于递归视图，这些是递归神经网络众所周知的优缺点，即
✓ 对序列数据具有天然的归纳偏置，并且原则上上下文不受限制。
✓ 高效推理（恒定时间状态更新）。
✗ 学习缓慢（缺乏并行性）。
✗ 训练过长序列时梯度消失或爆炸。

对于卷积视图，我们这里讨论的是卷积神经网络（我们这里是在其一维版本的上下文中）众所周知的优缺点，即
✓ 局部、可解释的特征。
✓ 高效（可并行化）训练。
✗ 在在线或自回归上下文中速度缓慢（每个新数据点都必须重新计算整个输入）。
✗ 上下文大小固定。

因此，根据过程的阶段（训练或推理）或我们拥有的数据类型，可以从一种视图切换到另一种视图，以便回到一个有利的框架中，从而最大限度地利用模型。
我们更喜欢用于通过并行化进行快速训练的卷积训练视图，用于高效推理的递归视图，以及用于处理连续数据的连续视图。



学习矩阵

在上面开发的卷积核中，Cˉ\mathbf{\bar{C}} 和 Bˉ\mathbf{\bar{B}} 是可学习的标量。
关于 Aˉ\mathbf{\bar{A}}，我们看到在我们的卷积核中，它在时间 kk 处表示为 kk 的幂。这可能非常耗时，因此我们正在寻找一个固定的 Aˉ\mathbf{\bar{A}}。为此，最好的选择是让它成为对角矩阵

A=[λ10⋯00λ2⋯0⋮⋮⋱⋮00⋯λn]⇒Ak=[λ1k0⋯00λ2k⋯0⋮⋮⋱⋮00⋯λnk] \mathbf{A} = \begin{bmatrix} \lambda_{1} & 0 & \cdots & 0 \\ 0 & \lambda_{2} & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \lambda_{n} \end{bmatrix} \Rightarrow \mathbf{A^k} = \begin{bmatrix} \lambda_{1}^k & 0 & \cdots & 0 \\ 0 & \lambda_{2}^k & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \lambda_{n}^k \end{bmatrix}

根据线性代数的谱定理，这正是正规矩阵的类别。
除了上面提到的离散化选择之外，Aˉ\mathbf{\bar{A}} 的定义和初始化方式也是文献中开发的各种SSM架构的差异点之一，我们将在下一篇博文中详细阐述。事实上，经验表明，用随机A\mathbf{A}矩阵初始化的SSM会导致糟糕的结果，而基于**HiPPO**矩阵（高阶多项式投影算子）的初始化则能获得非常好的结果（在MNIST序列基准测试中从60%提高到98%）。

**HiPPO**矩阵由S4的作者在之前的一篇论文（2020年）中引入。它也包含在S4作者的LSSL论文（2021年）以及S4的附录中。其公式如下：

A=[1−121−33−13−541−35−75−13−57−961−35−79−117−13−57−911−138⋮⋱]⇒Ank={(−1)n−k(2k+1)n>kk+1n=k0n<k \mathbf{A} = \begin{bmatrix} 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 \\ \vdots & & & & & & & & & \ddots \\ \end{bmatrix} \\ \Rightarrow \mathbf{A}_{nk} = \begin{cases}% (-1)^{n-k} (2k+1) & n > k \\ k+1 & n=k \\ 0 & n<k \end{cases}

该矩阵不是正规矩阵，但它可以分解为正规矩阵和低秩矩阵之和（在论文中简称为NPLR，即Normal Plus Low Rank）。作者在他们的论文中证明，这种类型的矩阵可以通过三种技术（参见论文中的算法1）高效计算：截断生成级数、柯西核和伍德伯里恒等式。

关于NPLR矩阵可以被高效计算为对角矩阵的详细证明，可以在论文附录（参见B和C部分）中找到。
S4的作者随后在他们的论文《如何训练你的HiPPO》（2022年）中修改了**HiPPO**矩阵（关于如何初始化它）。该论文产生的模型在文献中通常被称为“S4 V2”或“S4更新版”，以区别于“原始S4”或“S4 V1”。
在下一篇文章中，我们将看到其他作者（特别是Ankit GUPTA）建议使用对角矩阵而不是NPRL矩阵，这种方法现在更受青睐，因为它实现起来更简单。



实验结果

让我们通过分析S4在各种任务和基准测试中的一些结果来结束这篇博文，以了解SSM的潜力。

我们首先从音频任务和WARDEN（2018）的基准测试Speech Commands开始。

image
图4：来自Albert GU等人（2022）的论文On the Parameterization and Initialization of Diagonal State Space Models的图片，该模型也称为S4D，在S4之后发布，但以更结构化的形式再现了S4在该基准上的结果（S4D的结果已从图片中删除，以免剧透下一篇文章；）

从这个表格中可以观察到几点。
首先，在参数数量大致相等的情况下，S4的表现（至少高出13%）明显优于其他模型，这里是ConvNet类型。
其次，为了达到相同的性能，ConvNet需要多85倍的参数。
第三，在16K Hz上训练的ConvNet，当应用于8K Hz数据时，结果非常差。相比之下，S4在重新采样时保留了95%的性能。这可以解释为SSM的连续视图，在测试阶段，只需将Δ\Delta值减半即可。


让我们继续一个时间序列任务（在S4的一个修订版中引入）。

image
图5：来自S4附录的图片

论文的作者采用了ZHOU等人（2020年）的Informer模型的，并表明他们的模型在50种配置中的40种上优于该transformer。表格中显示的结果是单变量框架下的，但在多变量框架下也观察到了相同的结果（附录中的表14）。


让我们继续一个视觉任务，以及KRIZHESKY（2009）的基准测试sCIFAR-10。

image
图6：来自S4附录的图片

S4在sCIFAR-10上以仅10万个参数实现了SOTA（作者没有具体说明其他方法的参数数量）。


最后，我们来看一个文本任务，以及TAY等人（2020年）的基准测试Long Range Arena (LRA)。

image
图7：来自S4附录的图片

LRA包含6项任务，其中包括长度为16K tokens的Path-X，S4是第一个成功解决该任务的模型，展示了其在超长序列任务上的性能。
两年多以后，AMOS等人才在他们的论文Never Train from Scratch: Fair Comparison of Long-Sequence Models Requires Data-Driven Priors（2023年）中表明，由Ashish VASWANI等人（2017年）引入的transformer（且未与SSM混合）也可以解决此任务。然而，与SSM不同的是，它们无法通过65K tokens的PathX-256。

注意，S4在文本方面有一个缺点：与Transformer（标准版，更优化的版本具有更低的困惑度）在MERITY等人（2016年）的WikiText-103上相比，S4的困惑度更高。

image
图8：来自S4附录的图片

这可能是由于文本的非连续性（它并非从语音或时间序列等底层物理过程中采样而来）。我们将在专门讨论2023年SSM发展的文章中看到，这一点已经成为大量研究的焦点，并且SSM现在已经成功弥补了这一差距。


结论

SSM是具有三种视图的模型。一种连续视图，以及在离散化后的循环视图和卷积视图。
这种架构的挑战在于，根据处理过程的阶段（训练或推理）和处理的数据类型，知道何时偏向于一种视图而不是另一种。
这种模型具有高度通用性，因为它可应用于文本、视觉、音频和时间序列任务（甚至图形）。
它的一个优点是能够处理非常长的序列，通常比其他模型（ConvNet或transformers）使用更少的参数，同时速度也非常快。
正如我们将在接下来的文章中看到的，现有各种SSM架构之间的主要区别在于基本SSM方程的离散化方式，或者A\mathbf A矩阵的定义。


深入了解

SSM 历史

LMU由VOELKER、KAJIĆ和ELIASMITH于2019年12月发表，比S4早两年，可被视为S4的鼻祖。在这篇论文中，作者通过提出HOCHREITER和SCHMIDHUBER的LSTM的替代方案来启动循环视图，LSTM存在当处理步骤数量过高（根据变体不同，限制在100到5000之间）时梯度消失的问题。在论文中，他们展示了他们的模型能够处理超过10万步（VOELKER甚至在他的论文第6.1节中达到了超过10亿步）。为此，他们基于常微分方程x'(t) = Ax(t) + Bu(t)（在论文中x表示为m），并通过欧拉法对其进行离散化。矩阵A和B通过Padé近似获得，这极大地启发了HiPPO框架。该动力系统的关键特性是x通过勒让德多项式（直到d-1阶）表示u的滑动窗口。我们邀请读者查阅论文第2节以获取详细信息。
如引言所述，这篇论文是将同作者于2018年发表的更侧重于神经科学的模型应用于深度学习的成果。
最后，我们来提一下CHILKURI和ELIASMITH于2021年2月发布的LMU后续工作。在这篇论文中，他们展示了如何高效地计算他们的模型。为此，他们通过非顺序地重写其常微分方程来并行化训练（特别是论文第3页），从而可以使用标准控制理论工具（参见论文方程22和ÅSTRÖM和MURRAY以获取详细信息），然后将其视为卷积。他们取得了比SANH等人（2019）的DistillBERT更好的结果，参数数量减少了一半，并对text8数据集进行了字符级建模。还要注意的是，作者通过ZOH（零阶保持）离散化了他们的SSM，我们将在下一篇博客文章中更详细地讨论这一点。


SSM 资源

要了解更多关于SSM的信息，请查看


S4 资源

关于S4，请查阅以下资源


HiPPO 资源

有关**HiPPO**矩阵的更多信息，请查阅以下资源

参考文献

社区

注册或登录以发表评论

赞
160