2022 年状态空间模型 (SSM) 历史

社区文章 发布于 2024 年 4 月 11 日
法文版本可在我的博客上找到。

引言

上一篇文章中,我们使用一个连续时间系统定义了“状态空间模型”(SSM)。我们将其离散化,以展示其递归和卷积视图。这里的兴趣在于能够以卷积方式训练模型,然后对非常长的序列进行递归推理。

image
图 1:《结构化状态空间:结合连续时间、循环和卷积模型》博客文章中的图片,作者 Albert GU 等人(2022)

这一愿景由 Albert GU 在他 2021 年发表的论文LSSLS4中提出。S4 相当于 Transformer 的“注意力即你所需”。
在本文中,我们将回顾 2022 年发表的 SSM 文献。2023 年出现的文献将在下一篇文章中列出。目的是展示这些模型在几个月内的不同演变,同时保持简洁(即,我不会详细介绍所列论文的所有细节)。在 2022 年期间,各项进展主要集中于应用不同的离散化算法,同时用更简单的矩阵替换HiPPO矩阵。


理论模型

在本节中,我们将回顾 S4 架构改进背后的理论工作。然后,在另一个部分,我们将探讨其在不同任务(音频、视觉等)中的具体应用。

S4 V2

2022 年 3 月 4 日,S4 的作者更新了他们的论文,其中包含了一节关于 HiPPO 矩阵重要性的内容(参见论文最新版本第 4.4 节)。
简而言之,它包含了在 CIFAR-10 序列数据集上进行消融实验后观察到的结果。作者尝试使用各种参数化,例如随机稠密矩阵和随机对角矩阵,而不是使用带有 HiPPO 矩阵的 SSM。

image
图 2:《S4 论文》图 3,在 CIFAR-10 验证集上的准确率

因此,HiPPO 的使用很重要。获得的性能是否归因于其特定的内在品质,或者任何低秩正规矩阵 (NPLR) 是否就足够了?

image
图 3:《S4 论文》图 4,在 CIFAR-10 验证集上使用不同初始化和参数化的准确率

用 HiPPO 初始化 NPLR 矩阵可以显著提高性能。因此,根据这些实验,HiPPO 矩阵对于获得高性能模型至关重要。

S4 的作者进一步发展了他们的工作,并于 2022 年 6 月 24 日在论文《如何训练你的 HiPPO》中进行了介绍。这是一篇超过 39 页的极其详细的论文。

在这篇论文中,作者将 MSS 更直观地解释为卷积模型,其中卷积核是特定基函数的线性组合,从而引出了几种泛化和新方法。
例如,他们证明 S4 的A矩阵生成指数缩放的勒让德多项式 (LegS)。这使得系统能够通过非常平滑的核更好地建模长期依赖关系。
作者还推导出了一个新的 SSM,它生成截断傅里叶函数 (FouT) 的近似值。这种方法推广了短时傅里叶变换和局部卷积(即标准 ConvNet)。这个 SSM 还可以编码最先进的函数来解决经典的记忆任务。
请注意,本文主要介绍的是 HiPPO-FouT,而 HiPPO-LegS 和 HiPPO-LegT(截断勒让德多项式)已在两年前的原始 HiPPO 论文中介绍过。

image
图 4:不同的 HiPPO 变体

颜色代表每种方法的前 4 个基函数Kn(tK_n(t)(卷积核)(我们邀请读者查看论文表 1,了解每种方法中Kn(tK_n(t) 等于什么)。

此外,作者还研究了时间步长,它独立于离散化概念,可以简单地解释为控制 SSM 核的依赖长度或宽度。作者还详细介绍了如何为给定任务选择一个好的值。

所进行的工作将 S4 在 TAY、DEHGHANI 等人(2020 年)的基准LRA上的结果提高了 5.5 个点以上。

image
图 5:S4 V2 基准测试结果

本文所产生的模型在文献中通常被称为“S4 V2”或“S4 更新版”,与“原始 S4”或“S4 V1”相对。

DSS:对角线状态空间

2022 年 3 月 27 日,Ankit GUPTA 在其论文《对角线状态空间与结构化状态空间同样有效》中介绍了“对角线状态空间”(DSS)。
看来,在这篇论文之后,他与 Albert GU 开始合作更新此论文(GU 随后作为合著者出现在文章的 v2 和 v3 版本中)并在 S4D 的框架内(参见下一节)。
最主要的一点是,这种方法比 S4 简单得多。事实上,DSS 基于对角线状态矩阵(即没有 S4 的低秩校正,也就是没有 HiPPO 矩阵),如果初始化得当,其性能优于原始 S4。此后,用对角线矩阵代替 HiPPO 矩阵来表示A\mathbf{A}已成为标准做法。

然而,让我们探讨一下本文中包含的一些复杂性/局限性。列出它们将有助于我们理解以下方法(旨在进一步简化)的贡献。


1. 离散化
DSS 使用与 S4 相同的微分方程系统

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

然而,它使用不同的离散化来获得卷积和递归视图,即零阶保持 (ZOH) 离散化,而不是双线性离散化,后者假设采样信号在每个采样点之间是恒定的。
下表比较了两种离散化在递归视图中A\mathbf{A}B\mathbf{B}C\mathbf{C}的值,以及卷积视图中卷积核的表达式

离散化 双线性 ZOH
递归 Aˉ=(IΔ2A)1(I+Δ2A\mathbf{\bar{A}} = (\mathbf {I} - \frac{\Delta}{2} \mathbf{A})^{-1}(\mathbf {I} + \frac{\Delta}{2} \mathbf{A})
Bˉ=(IΔ2A)1ΔB\mathbf {\bar{B}} = (\mathbf{I} - \frac{\Delta}{2} \mathbf{A})^{-1} \Delta \mathbf{B}
Cˉ=C\mathbf{\bar{C}} = \mathbf{C}
Aˉ=eAΔ\mathbf{\bar{A}} = e^{\mathbf{A}\Delta}
Bˉ=(AˉI)A1B\mathbf{\bar{B}} = (\mathbf{\bar{A}} - I)\mathbf{A}^{-1}\mathbf{B}
Cˉ=C\mathbf{\bar{C}} = \mathbf{C}
卷积 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}}) Kˉ=( CeAkΔ(eAΔI)A1B )0k<L\mathbf{\bar{K}} = (\ \mathbf{C} e^{\mathbf{A}\cdot k\Delta} (e^{\mathbf{A}\Delta} - I)\mathbf{A}^{-1}\mathbf{B}\ )_{0 \leq k < L}

对于 ZOH,经过计算,我们得到yk=j=0kCˉAˉjBˉukj=j=0kKˉjukjy_k = \sum_{j=0}^k \bar{C}\bar{A}^j\bar{B}\cdot u_{k-j} = \sum_{j=0}^k \bar{K}_j\cdot u_{k-j}

然后,通过快速傅里叶变换 (FFT) 在O(L log(L)O(L~log(L)) 中计算uuKˉ\bar{K}yy,其中LL是序列长度,同时计算两个L1L-1次多项式的乘积。



2. DSSsoftmax 和 DSSexp

简短版

GUPTA 提出了一个方案,以获得与 S4 具有相同表达能力的 DSS,从而产生了两种不同的 DSS:DSSexp 和 DSSsoftmax。有关它们的信息总结在下表中

方法 DSSexp DSSsoftmax
卷积视图 K=KˉΔ,L(Λ,I1iN,w~)K = \bar{K}_{\Delta, L}(\Lambda,\mathbb{I}_{1 \leq i \leq N},\widetilde{w})
=w~Λ1(eΛΔI)elementwise-exp(P)= \widetilde{w} \cdot \Lambda^{-1} (e^{\Lambda\Delta} - I) \cdot \text{elementwise-exp}(P)
K=KˉΔ,L(Λ, ((eLλiΔ1)1)1iN,w)K = \bar{K}_{\Delta, L}(\Lambda,\ ((e^{L\lambda_i\Delta} - 1)^{-1})_{1\leq i \leq N}, w)
=wΛ1row-softmax(P)= w \cdot \Lambda^{-1} \cdot \text{row-softmax}(P)
递归视图 Aˉ=diag(eλ1Δ,,eλNΔ)\bar{A} = \mathrm{diag}(e^{\lambda_1\Delta}, \ldots, e^{\lambda_N\Delta})
Bˉ=(λi1(eλiΔ1))1iN\bar{B} = \left(\lambda_i^{-1} (e^{\lambda_i\Delta} - 1) \right)_{1\leq i \leq N}
Aˉ=diag(eλ1Δ,,eλNΔ)\bar{A} = \mathrm{diag}(e^{\lambda_1\Delta}, \ldots, e^{\lambda_N\Delta})
Bˉ=(eλiΔ1λi(eλiΔL1))1iN\bar{B} = \left( {e^{\lambda_i\Delta} - 1 \over \lambda_i (e^{\lambda_i\Delta L} - 1)} \right)_{1\leq i \leq N}
解释 作为 LSTM 的遗忘门 如果(λ)<<0\Re(\lambda)<<0:保留局部信息,
如果(λ)>>0\Re(\lambda)>>0:可以捕获非常长距离的信息

因此,我们在这里处理的是 ℂ,而不是 ℝ。


长版

GUPTA 提出了以下方案,以获得与 S4 具有相同表达能力的 DSS

令长度为 LL 的给定状态空间 (A,B,C(A, B, C) 和采样时间 Δ>0\Delta > 0 的核为 KR1×LK \in \mathbb{R}^{1\times L},其中 ACN×NA \in \mathbb{C}^{N \times N}C\mathbb{C} 上是对角化的,其特征值为 λ1,,λN\lambda_1,\ldots,\lambda_N,且对于所有 iiλi0\lambda_i \neq 0eLλiΔ1e^{L\lambda_i\Delta} \neq 1。令 PCN×LPi,k=λikΔP \in \mathbb{C}^{N \times L} P_{i,k} = \lambda_i k\DeltaΛ\Lambda 为对角矩阵,其对角线元素为 λ1,,λN\lambda_1,\ldots,\lambda_N。那么存在 w~,wC1×N\widetilde{w}, w \in \mathbb{C}^{1\times N},使得

  • (a): K = KˉΔ,L(Λ, (1)1iN, w~) = w~Λ1(eLambdaΔI)elementwise-exp(P)K\ =\ \bar{K}_{\Delta, L}(\Lambda,\ (1)_{1 \leq i \leq N},\ \widetilde{w})\ =\ \widetilde{w} \cdot \Lambda^{-1} (e^{Lambda\Delta} - I) \cdot \text{elementwise-exp}(P)
  • (b): K  = KˉΔ,L(Λ, ((eLambdaiΔ1)1)1iN, w) =  wΛ1row-softmax(P)K\ \ =\ \bar{K}_{\Delta, L}(\Lambda,\ ((e^{Lambda_i\Delta} - 1)^{-1})_{1\leq i \leq N},\ w)\ =\ \ w \cdot \Lambda^{-1} \cdot \text{row-softmax}(P)

(a) 表明我们可以通过 Λ,w~CN\Lambda, \widetilde{w} \in \mathbb{C}^N 对状态空间进行参数化,并按所示计算核。不幸的是,在实践中,ΛΛ 元素的实部在学习过程中可能会变为正值,导致长输入不稳定。为了解决这个问题,作者提出了两种方法:DSSexp 和 DSSsoftmax。

2.1 卷积视角
在 DSSexp 中,ΛΛ 的实部必须为负值。然后我们有 Λ=elementwise-exp(Λre)+iΛim\Lambda = - \text{elementwise-exp}(\Lambda_\mathrm{re}) + i\cdot \Lambda_\mathrm{im}Δ=exp(Δlog)R>0\Delta = \mathrm{exp}(\Delta_{\log}) \in \mathbb{R}_{> 0}KK 则根据提案的 (a) 部分中给出的公式进行计算。
在 DSSsoftmax 中,ΛΛ 的每一行都按这些元素的总和进行归一化。我们有 Λ=Λre+iΛim\Lambda = \Lambda_\mathrm{re} + i\cdot \Lambda_\mathrm{im}Δ=exp(Δlog)R>0\Delta = \mathrm{exp}(\Delta_{\log}) \in \mathbb{R}_{> 0}KK 则根据提案的 (b) 部分中所示公式进行计算。
请注意,C\mathbb{C} 上的 softmax 不一定在 softmax (0,iπ)(0, i \pi) 期间定义,作者使用了修正版的 softmax 来防止此问题(参见论文附录 A.2)。

2.2 循环视角
在 DSSexp 中,使用上表中的循环公式,我们得到 Aˉ=diag(elambda1Δ,,elambdaNΔ)\bar{A} = \mathrm{diag}(e^{lambda_1\Delta}, \ldots, e^{lambda_N\Delta})Bˉ=(λi1(elambdaiΔ1))1iN\bar{B} = \left(\lambda_i^{-1} (e^{lambda_i\Delta} - 1) \right)_{1\leq i \leq N},其中两个等式中,λi\lambda_i 是 Lambda 的第 i 个特征值。
由于 Aˉ\bar{A} 是对角矩阵,我们可以独立计算 xkx_k,如下所示:xi,k=eλiΔxi,k1+λi1(eλiΔ1)ukx_{i,k} = e^{\lambda_i\Delta} x_{i,k-1} + \lambda_i^{-1} (e^{\lambda_i\Delta} - 1)u_k
然后可以推断,如果 λiΔ0\lambda_i|\Delta \approx 0,我们有 xi,kxi,k1x_{i,k} \approx x_{i,k-1},这允许我们在许多时间步长上复制历史信息。另一方面,如果 Re(λi)Δ0\mathrm{Re}(\lambda_i)\Delta \ll 0,那么 i,kλi1uk_{i,k} \approx -\lambda_i^{-1}u_k,并且会忘记先前时间步的信息,类似于 LSTM 中的“遗忘”门。

在 DSSsoftmax 中,使用上表中的循环公式,我们得到:Aˉ=diag(eambda1Δ,,eambdaNΔ)\bar{A} = \mathrm{diag}(e^{ambda_1\Delta}, \ldots, e^{ambda_N\Delta})barB=(eambdaiΔ1λi(eambdaiΔL1))1iNbar{B} = \left( {e^{ambda_i\Delta} - 1 \over \lambda_i (e^{ambda_i\Delta L} - 1)} \right)_{1\leq i \leq N}
因此 xi,k=eλiΔxi,k1+uk(eλiΔ1)λi(eλiΔL1)x_{i,k} = e^{\lambda_i\Delta} x_{i,k-1} + {u_k(e^{\lambda_i\Delta} - 1) \over \lambda_i (e^{\lambda_i\Delta L} - 1)}
请注意 eλiΔe^{\lambda_i\Delta} 可能会不稳定。我们必须根据 Re(λ)\mathrm{Re}(\lambda) 的符号,通过引入中间状态 x~k\widetilde{x}_{k} 来计算两种不同情况。
• 如果 Re(λ)0\mathrm{Re}(\lambda) \leq 0x~k=eλΔx~k1+uk   ,   xk=x~k(eλΔ1)λ(eλΔL1)\widetilde{x}_{k} = e^{\lambda\Delta}\cdot \widetilde{x}_{k-1} + u_k \ \ \ ,\ \ \ x_k = \widetilde{x}_k \cdot {(e^{\lambda\Delta} - 1) \over \lambda (e^{\lambda\Delta L} - 1) }

特别地,如果 Re(λ)0\mathrm{Re}(\lambda) \ll 0,则 x~kuk\widetilde{x}_k \approx u_kxkuk/λx_k \approx u_k / \lambda,这意味着模型侧重于局部信息(忽略之前的步骤)。

• 如果 Re(λ)>0\mathrm{Re}(\lambda) > 0x~k=x~k1+ekλΔuk    ,    xk=x~kelambdaΔ(k(L1))λeλΔ1eλΔL1\widetilde{x}_{k} = \widetilde{x}_{k-1} + e^{-k\lambda\Delta} \cdot u_k \ \ \ \ ,\ \ \ \ x_k = \widetilde{x}_k \cdot {e^{lambda\Delta (k-(L-1))} \over \lambda}\cdot {e^{-\lambda\Delta}-1 \over e^{-\lambda\Delta L} - 1 }

类似地,如果 Re(λ)0\mathrm{Re}(\lambda) \gg 0,则 x~0u0\widetilde{x}_0 \approx u_0x~kx~k1u0\widetilde{x}_k \approx \widetilde{x}_{k-1} \approx u_0xk<L10x_{k < L-1} \approx 0xL1u0/λx_{L-1} \approx u_0 / \lambda,模型可以捕获非常远距离的信息。实际上,S4D 的作者指出,如果 LL 非常大(当 tt \rightarrow \inftyK(t)=Cexp(A).B)K(t) = C \exp(A^\intercal).B) 中),Re(λ)0\mathrm{Re}(\lambda) \gg 0 不起作用。


3. 初始化
ww 的实部和虚部从 N(0,1)\mathcal{N}(0,1) 初始化,Δlog\Delta_{\log} 的元素从 exp(U (log(0.001),log(0.1)))\exp(\mathcal{U}~(log(0.001), log(0.1))) 初始化,Δ\Delta 通过 HiPPO 矩阵的特征值初始化。作者们想知道是否可以为 Δ\Delta 找到更简单的初始化方法。然而,他们指出随机化会导致更差的结果。

关于结果,DSS 已在 LRA 和 WARDEN (2018) 的 Speech Commands 上进行了测试。

image
图 6:LRA 上的 DSS 结果

对于此基准,DSS(softmax 或 exp 版本)实现了比原始 S4 更好的平均结果。DSSsoftmax 似乎比 DSSexp 表现稍好。这篇论文的另一个有趣之处在于,它是第一篇重现 S4 结果的论文,从而证实 SSM 通过了此基准测试。

image
图 7:Speech Commands 上的 DSS 结果

在 Speech Commands 上,S4 保持了其对 DSS 的优势。

深入探讨
官方实现可在 GitHub 上获取。
这篇论文是 NeurIPS 2022 上的 spotlight talk 的主题。



S4D:对角线 S4

2022 年 6 月 23 日,GU、GUPTA 等人在他们的论文 On the Parameterization and Initialization of Diagonal State Space Models 中介绍了 S4D。
DSS 状态矩阵 A\mathbf{A} 的初始化依赖于 S4 HiPPO 矩阵的特定近似。虽然 S4 矩阵具有处理长程依赖的数学解释,但对角线近似的有效性在理论上仍未得到解释。
通过 S4D,作者们引入了一种对角线 SSM,它结合了 S4 计算和参数化的优点以及 DSS 初始化的优点。结果是一种非常简单、理论上合理且经验上有效的方法。

论文表 1 对这三种方法进行了比较

image
图 8:S4、DSS 和 S4D 的比较

S4D 可以使用 S4 的双线性离散化或 DSS 的 ZOH 离散化。

在 S4D 中,方程 y=uKy = u \ast \mathbf{\overline{K}} 的离散化卷积核计算如下:
K=n=0N1CnAnBn    K=(BC)VL(A)\mathbf{\overline{K}}_\ell = \sum_{n = 0}^{N-1} \mathbf{C}_n \mathbf{\overline{A}}_n^\ell \mathbf{\overline{B}}_n \implies \mathbf{\overline{K}} = (\mathbf{\overline{B}}^\top \circ \mathbf{C}) \cdot \mathcal{V}_L(\mathbf{\overline{A}})
其中:
\circ 表示 Hadamard 矩阵积
\cdot 是一个经典的矩阵积,
VL\mathcal{V}_L范德蒙矩阵,即:V=[1α1α12α1n11α2α22α2n11α3α32α3n11αmαm2αmn1]\mathcal{V} = \begin{bmatrix} 1&\alpha _{1}&{\alpha _{1}}^{2}&\dots &{\alpha _{1}}^{n-1}\\ 1&\alpha _{2}&{\alpha _{2}}^{2}&\dots &{\alpha _{2}}^{n-1}\\ 1&\alpha _{3}&{\alpha _{3}}^{2}&\dots &{\alpha _{3}}^{n-1}\\ \vdots &\vdots &\vdots & \vdots \\ 1&\alpha _{m}&{\alpha _{m}}^{2}&\dots &{\alpha _{m}}^{n-1} \end{bmatrix}

换句话说,对于任意iijj,第ii行第jj列的系数是Vi,j=αij1\displaystyle V_{i,j}={\alpha _{i}}^{j-1}

最后,在 S4D 中,

K=[B0C0BN1CN1][1A0A02A0L11A1A12A1L11AN1AN12AN1L1]where VL(A)n,=An \mathbf{\overline{K}} = \begin{bmatrix} \mathbf{\overline{B}}_0 \mathbf{C}_0 & \dots & \mathbf{\overline{B}}_{N-1} \mathbf{C}_{N-1} \end{bmatrix} \begin{bmatrix} 1 & \mathbf{\overline{A}}_0 & \mathbf{\overline{A}}_0^2 & \dots & \mathbf{\overline{A}}_0^{L-1} \\ 1 & \mathbf{\overline{A}}_1 & \mathbf{\overline{A}}_1^2 & \dots & \mathbf{\overline{A}}_1^{L-1} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & \mathbf{\overline{A}}_{N-1} & \mathbf{\overline{A}}_{N-1}^2 & \dots & \mathbf{\overline{A}}_{N-1}^{L-1} \\ \end{bmatrix} \qquad \text{where } \mathcal{V}_L(\mathbf{\overline{A}})_{n, \ell} = \mathbf{\overline{A}}_n^\ell

所有计算都在O(N+L)O(N+L)操作和空间中完成。

各种矩阵的参数化如下:

  • A=exp((A))+i(A)\mathbf{A} = -\exp(\Re(\mathbf{A})) + i \cdot \Im(\mathbf{A}).
    作者指出,指数函数可以用任何正函数代替。
  • B=1\mathbf{B} = 1然后进行训练。
  • C\mathbf{C} 随机初始化,标准差为1,然后进行训练。

请注意,S4 考虑实数,而 S4D 通过将状态大小参数化为 N/2N/2 并隐式地向参数添加共轭对来考虑复数。这样我们就有相当于 NN 个实数参数,确保输出是实数。

关于初始化,作者介绍了两种方法:

  • S4D-Inv 是 S4-LegS 的近似值:An=12+iNπ(N2n+11)\quad \mathbf{A}_n = -\frac{1}{2} + i \frac{N}{\pi} \left( \frac{N}{2n+1}-1 \right)
  • S4D-Lin 是 S4-FouT 的近似值:An=121+iπn\quad \mathbf{A}_n = -\frac{1}{2}\mathbf{1} + i \pi n

我们邀请读者参阅论文的第 4 部分,了解这些方程的更多细节。
从可解释性的角度来看,An\mathbf{A}_n的实部控制权重衰减率。An\mathbf{A}_n的虚部控制基函数Kn(t)=etABK_n(t) = e^{t\mathbf{A}}\mathbf{B}的振荡频率。

最后,作者提出了一些结果:

  1. 用 Softmax 代替 Vandermonde 计算模型,效果没有太大差异。
  2. 训练 B 总是能得到更好的结果。
  3. 两种可能的离散化之间没有显著差异。
  4. 限制 A 的实部会带来更好的结果(尽管不显著)。
  5. 所有测试的初始化修改都降低了结果。即,对虚部应用系数或使用随机虚部/使用随机实部/使用虚部和随机实部。

由于这种方法比其他方法更容易实现(Vandermade 只需要两行代码),S4D 已取代 S4 成为常用方法(实际上,例如在 Mamba 论文的表 6 中,作者使用术语 S4 来指代 S4D)。


深入探讨
官方实现可在 GitHub 上获取。
2022 年 12 月 1 日,GUPTA 等人发表了 DSS 的续篇《Simplifying and Understanding State Space Models with Diagonal Linear RNNs》,其中介绍了 DLR。他们摆脱了离散化步骤,提出了一种基于对角线性 RNN(DLR)的模型,该模型可以在大约 100 万个位置上运行(与经典 RNN 不同)。该模型的代码可在 GitHub 上获取。



GSS: 门控状态空间(Gated State Space)

在 S4D 发布的五天后,即 2022 年 6 月 27 日,MEHTA、GUPTA 等人在他们的论文《Long Range Language Modeling via Gated State Spaces》中介绍了 GSS。
在这项工作中,他们专注于对英语书籍、Github 源代码和 ArXiv 数学文章中的自回归序列进行建模(而 GSS 之前的研究特别关注序列分类任务)。他们表明,他们的称为门控状态空间(GSS)的层训练速度明显快于 DSS(快 2 到 3 倍)。他们还表明,利用自注意力来建模局部依赖关系可以进一步提高 GSS 的性能。

image
图 9:DSS 与 GSS 的比较。模型在 4K 长度的序列上进行训练,然后在长达 65K 令牌的序列上进行评估。

基于 SSM(S4/DSS)在 TPU 上训练速度慢于预期的观察结果,作者修改了架构,以减少特定操作的维度,这些操作被证明是瓶颈。这些修改的灵感来自于关于门控单元效率的经验观察(DAUPHIN 等人(2016)的《Language Modeling with Gated Convolutional Networks》,SHAZEER(2020)的《GLU Variants Improve Transformer》等)。更具体地说,作者借鉴了 HUA 等人(2022)的论文《Transformer Quality in Linear Time》。后者表明,他们的 FLASH 模型,用门控单元替换 Transformer 中的前馈层,可以在质量损失最小的情况下使用更低的单元峰值注意力。他们将此组件称为门控注意力单元(GAU)。

image
图 10:门控注意力单元。
这与论文中的图不完全相同:我进行了水平平移,使输入位于底部而不是顶部,以便于与下面可见的 Mega 图进行并行比较。

因此,GSS 的作者将门控单元的使用扩展到 SSM,并观察到在执行 FFT 操作时维度有所降低。

image
图 11:GAU 在 SSM 中的应用。

值得注意的是,与 HUA 等人不同,作者没有观察到使用 RELU² 或 Swish 激活函数而不是 GELU 有多少优势,因此保留了 GELU。
此外,DSS 使用固定的时间步长 为 1(作者观察到这减少了创建核所需的计算时间并简化了它们的计算)。
一个特别有趣的点是,与 S4 和 DSS 中观察到的不同,模型在语言建模任务上的性能对初始化不那么敏感,从而可以通过随机初始化状态空间变量成功训练模型。这是一个非常重要的结果,因为它表明 HiPPO 矩阵(S4)和 HiPPO 初始化(DSS)都不是必需的。

至于 GSS-Transformer 混合模型,它只是稀疏地将传统 Transformer 块与 GSS 层交织在一起。混合模型实现了比纯 SSM 模型更低的困惑度。

图 12:GSS-Transformer 混合模型的性能。

深入探讨
官方实现可在 GitHub 上获取。



Mega

2022 年 9 月 21 日,MA、ZHOU 等人发表了《Mega: Moving Average Equipped Gated Attention》。
Mega 是一种单头注意力机制的 Transformer,使用 GAU 门控系统,并配备了阻尼指数移动平均(EMA)以融入位置归纳偏置。

image
图 13:Mega 概述,根据我对论文的理解设计的图。
作者将 GAU 的 RELU² 替换为更稳定的 Laplace 函数(参见论文图 4)。

作者还提出了一个变体,Mega-chunk,它有效地将整个序列分成几个固定长度的块。在这里,他们采用了 FLASH 模型中已经存在并解释的原理(参见该模型论文的图 4)。这提供了线性的时间复杂度和空间复杂度,同时质量损失最小。

image
图 14:Mega 分块

这提供了线性复杂度,只需将注意力局部应用于每个固定长度的块。
更准确地说,我们将查询、键和值的序列分成长度为 cc 的块。例如,Q=Q1,...Qk\mathbf{Q} = {\mathbf{Q}_1, ... \mathbf{Q}_k},其中 k=nck = -\frac{n}{c} 是块的数量。注意力操作单独应用于每个块,给出关于 nn 的线性复杂度 O(kc2)=O(nc)\mathcal{O}(kc^2) = \mathcal{O}(nc)
然而,这种方法存在一个关键限制,即损失其他块的上下文信息。但 EMA 子层通过捕获每个令牌附近的局部上下文信息来缓解这个问题,其结果用作注意力子层的输入。这样,块级注意力利用的有效上下文可以扩展到块边界之外。

Mega 极具竞争力,因为它成为 LRA 上最好的模型。

image
图 15:Mega 在 LRA 基准测试中的结果

那么,Transformer 在一篇关于 SSM 的博客文章中做什么呢?让我们看看阻尼 EMA 以了解 Mega 和 S4D 之间的联系。

“经典” EMA 提醒
“经典”EMA 的方程为 𝐲t=𝜶𝐱t+(1𝜶)𝐲t1𝐲_t = 𝜶 ⊙ 𝐱_t + (1-𝜶) ⊙ 𝐲_{t-1},其中 𝜶𝜶[0,1]d[0,1]^d 中,表示 EMA 系数,代表权重衰减程度,⊙ 为 Hadamard 矩阵积。
𝜶 越高,旧观察值的折现速度越快。
因此,我们在此施加了一个归纳偏差:两个令牌之间的依赖权重随时间呈指数衰减,且衰减因子 𝜶 与输入无关。此特性有利于局部依赖,并限制了长期依赖。
EMA 计算可以表示为 n 个单独的卷积,可以通过 FFT 高效计算。

Mega 中使用的 EMA
Mega 使用多维“阻尼”EMA。也就是说,在“阻尼”EMA 方程中,𝐲t=𝜶𝐱t+(1𝜶𝜹)𝐲t1𝐲_t = 𝜶 ⊙ 𝐱_t + (1-𝜶 ⊙ 𝜹) ⊙ 𝐲_{t-1} 中引入了参数 𝜹𝜹,其在 [0,1]d[0, 1]^d 中表示阻尼因子,xx 通过 Rd×hR^{d \times h} 的扩展矩阵 𝜷𝜷 扩展到 hh 维。
该方程变为 𝐲t,j=𝜼j𝐡t(j)𝐲_{t, j} = 𝜼_j^{\intercal} 𝐡_t^{(j)},其中 𝐡t(j)=𝜶j𝐮t(j)+(1𝜶j𝜹j)𝐡t1(j)𝐡_t^{(j)} = 𝜶_j ⊙ 𝐮_t^{(j)}+ (1-𝜶_j ⊙ 𝜹_j) ⊙ 𝐡_{t-1}^{(j)},而 𝜼Rd×h𝜼 \in R^{d \times h} 是投影矩阵,它将 hh 维隐藏状态返回到 1 维输出 𝐲t,jR𝐲_{t, j} \in \mathbb{R}

证明多维“阻尼”EMA 可以通过卷积计算,因此可以通过 FTT 计算(通过将 d=1d = 1 用于 𝜶𝜶𝜹𝜹

我们有 𝐲t=𝜼𝐡t𝐲_t = 𝜼^{\intercal} 𝐡_t,其中 𝐡t=𝜶j𝐮t+(1𝜶𝜹)𝐡t1𝐡_t = 𝜶_j ⊙ 𝐮_t + (1-𝜶 ⊙ 𝜹) ⊙ 𝐡_{t-1}。我们设 ϕ=1𝜶𝜹ϕ = 1-𝜶 ⊙ 𝜹
则:𝐡t=𝜶𝐮t+(1𝜶𝜹)𝐡t1=𝜶𝜷𝐱t+ϕ𝐡t1𝐡_t = 𝜶 ⊙ 𝐮_t + (1-𝜶 ⊙ 𝜹) ⊙ 𝐡_{t-1} = 𝜶 ⊙ 𝜷 𝐱_t + ϕ ⊙ 𝐡_{t-1}
    且 𝐲t=𝜼𝐡t=𝜼(𝜶𝜷𝐱t+ϕ𝐡t1)𝐲_t = 𝜼^{\intercal} 𝐡_t = 𝜼^{\intercal} (𝜶 ⊙ 𝜷 𝐱_t + ϕ ⊙ 𝐡_{t-1})
然后,展开上述两个方程,我们明确得到
步骤 0:𝐡1=𝜶𝜷𝐱1+ϕ𝐡0𝐡_1 = 𝜶 ⊙ 𝜷𝐱_1 + ϕ ⊙ 𝐡_0
步骤 1:𝐡2=𝜶𝜷𝐱2+ϕ𝐡1𝐡_2 = 𝜶 ⊙ 𝜷𝐱_2 + ϕ ⊙ 𝐡_1 =𝜶𝜷𝐱2+ϕ(ϕ𝐡0+𝜶𝜷𝐱1)=𝜶𝜷𝐱2+ϕ2𝐡0+ϕ𝜶𝜷𝐱1= 𝜶 ⊙ 𝜷 𝐱_2 + ϕ ⊙ (ϕ ⊙ 𝐡_0 + 𝜶 ⊙ 𝜷 𝐱_1) = 𝜶 ⊙ 𝜷 𝐱_2 + ϕ^2 ⊙ 𝐡_0 + ϕ ⊙ 𝜶 ⊙ 𝜷 𝐱_1
...

同理
步骤 0:𝐲1=𝜼𝜶𝜷𝐱1+ϕ𝐡0)=𝜼𝜶𝜷𝐱1+𝜼ϕ𝐡0𝐲_1 = 𝜼^{\intercal} 𝜶 ⊙ 𝜷 𝐱_1 + ϕ ⊙𝐡0) = 𝜼^{\intercal} 𝜶 ⊙ 𝜷 𝐱_1 + 𝜼^{\intercal} ϕ ⊙ 𝐡_0
步骤 1:𝐲2=𝜼𝜶𝜷𝐱2+𝜼ϕ𝐡1𝐲_2 = 𝜼^{\intercal} 𝜶 ⊙ 𝜷 𝐱_2 + 𝜼^{\intercal} ϕ ⊙ 𝐡_1 =𝜼𝜶𝜷𝐱2+𝜼ϕ(𝜶𝜷𝐱1+ϕ𝐡0)=𝜼𝜶𝜷𝐱2+𝜼ϕ2𝐡0= 𝜼^{\intercal} 𝜶 ⊙ 𝜷 𝐱_2 + 𝜼^{\intercal} ϕ ⊙ (𝜶 ⊙ 𝜷 𝐱_1 + ϕ ⊙ 𝐡_0) = 𝜼^{\intercal} 𝜶 ⊙ 𝜷 𝐱_2 + 𝜼^{\intercal} ϕ^2 ⊙ 𝐡_0
...
步骤 tt𝐲t=𝜼𝜶𝜷𝐱t+...+𝜼ϕt1𝜶𝜷𝐱t1+𝜼ϕt𝐡0𝐲_t = 𝜼^{\intercal} 𝜶 ⊙ 𝜷 𝐱_t + ... + 𝜼^{\intercal} ϕ_{t-1} ⊙ 𝜶 ⊙ 𝜷 𝐱_{t-1} + 𝜼^{\intercal} ϕ^t ⊙ 𝐡_0

因此 𝐲=K𝐱+𝜼ϕt𝐡0𝐲 = \mathbf{K} * 𝐱 + 𝜼^{\intercal} ϕ^t ⊙ 𝐡_0,其中 K=(𝜼(𝜶𝜷),𝜼(ϕ𝜶𝜷),...,𝜼(ϕt𝜶𝜷)Rn\mathbf{K} = (𝜼^{\intercal} (𝜶 ⊙ 𝜷), 𝜼^{\intercal} (ϕ ⊙ 𝜶 ⊙ 𝜷), ..., 𝜼^{\intercal}(ϕ^t ⊙ 𝜶 ⊙ 𝜷) \in \mathbb{R}^n
K\mathbf{K} 在 Mega 中通过范德蒙乘积计算,这让我想起了 S4D 中使用的方法。


深入探讨
官方实现可在 GitHub 上获取。
该模型也可在 Transformers 上获取。
有关 Mega 和 S4 之间联系的更多详细信息,读者可查阅 Albert GU 和 Mega 作者在 Open Review 上对 Mega 提交的评论中交换的消息。总而言之,通过将 SSM 离散化步骤与阻尼 EMA 相关联,可以将 Mega 视为混合 SSM/Attention,从而简化 S4 为实值而非复数。



Liquid-S4:液体结构状态空间模型

2022 年 9 月 26 日,HASANI, LECHNER 等人在线发布了 Liquid Structural State-Space Models,介绍了 Liquid-S4。在本文中,作者使用结构 SSM (S4) 公式来获得线性液体网络的实例,这些实例同时具有 S4 和 LTC(液体时间常数)的近似能力。
LTC 神经网络是连续时间因果神经网络,具有输入依赖的状态转换模块,使其能够在推理过程中学习适应输入。这可以看作是一种选择机制。
要了解更多关于液体网络的信息,您可以查阅同一作者之前的一篇论文:液体时间常数网络 (2021)。
在 Liquid-S4 的上下文中,您只需要知道 LTC 在每个时间步的状态由以下公式给出:

dx(t)dt=[A+Bf(x(t),u(t),t,θ)]Liquid time-constantx(t)+Bf(x(t),u(t),t,θ). \frac{d\textbf{x}(t)}{dt} = - \underbrace{\Big[\mathbf{A} + \mathbf{B} \odot f(\textbf{x}(t),\textbf{u}(t), t, \theta)\Big]}_\text{Liquid time-constant} \odot \textbf{x}(t) + \mathbf{B} \odot f(\textbf{x}(t), \textbf{u}(t), t, \theta).

其中

  • x(N×1)(t)\textbf{x}^{(N \times 1)}(t) 是大小为 NN 的隐状态向量,
  • u(m×1)(t)\textbf{u}^{(m \times 1)}(t) 是具有 mm 个特征的输入信号,
  • A(N×1)\mathbf{A}^{(N \times 1)} 是时间常数状态转换机制,
  • B(N×1)\mathbf{B}^{(N \times 1)} 是一个偏差向量
  • \odot 表示 Hadamard 积
  • f(.)f(.) 是一个受 θ\theta 参数化的有界非线性函数。

在实践中,使用液体网络的 SSM 通过以下微分方程系统进行公式化:

x=(A+Bu)x+Buy=Cx \begin{aligned} x' &= (\mathbf{A} + \mathbf{B} u) x + \mathbf{B}u\\ y &= \mathbf{C}x \end{aligned}

这个动态系统可以通过与 S4 相同的参数化有效求解,从而产生一个额外的卷积核,该卷积核考虑了移位信号的相似性。由此产生的模型是 Liquid-S4。我们用一点数学来解释这一点。

Liquid-S4 的循环视图是通过使用梯形法则(双线性形式)离散化系统获得的。结果是

xk=(A+B uk) xk1+B uk,      yk=C xk x_k = \big( \overline{\textbf{A}} + \overline{\textbf{B}}~u_k\big)~x_{k-1} + \overline{\textbf{B}}~u_k,~~~~~~y_k = \overline{\textbf{C}}~x_k

与 S4 一样,通过在时间上展开递归视图(假设 x1=0x_{-1} = 0)获得卷积视图

x0=Bu0y0=CBu0x1=ABu0+Bu1+ B2u0u1y1=CABu0+CBu1+CB2u0u1x2=A2Bu0+ABu1+Bu2+ AB2u0u1+ AB2u0u2+ B2u1u2+ B3u0u1u2y2=CA2Bu0+CABu1+CBu2+ CAB2u0u1+ CAB2u0u2+ CB2u1u2+ CB3u0u1u2,  yk=CAkBu0+CAk1Bu1+CABuk1+CBuk +p=2P(k+1p) of uiui+1  up CA(k+1pi)Bpuiui+1  up  for iZ and i0             y= Ku+Kliquiducorrelations \begin{align} x_0 &= \overline{\textbf{B}} u_0 \\ \nonumber y_0 &= \overline{\textbf{C}}\overline{\textbf{B}} u_0 \\\\ \nonumber x_1 &= \overline{\textbf{A}}\overline{\textbf{B}} u_0 + \overline{\textbf{B}} u_1 {\color{violet} +~ \overline{\textbf{B}}^2 u_0 u_1} \\ \nonumber y_1 &= \overline{\textbf{C}}\overline{\textbf{A}}\overline{\textbf{B}} u_0 + \overline{\textbf{C}}\overline{\textbf{B}} u_1 {\color{violet} + \overline{\textbf{C}}\overline{\textbf{B}}^2 u_0 u_1} \\\\ \nonumber x_2 &= \overline{\textbf{A}}^2\overline{\textbf{B}} u_0 + \overline{\textbf{A}}\overline{\textbf{B}} u_1 + \overline{\textbf{B}} u_2 {\color{violet} +~ \overline{\textbf{A}}\overline{\textbf{B}}^2 u_0 u_1 +~ \overline{\textbf{A}}\overline{\textbf{B}}^2 u_0 u_2 +~ \overline{\textbf{B}}^2 u_1 u_2 +~ \overline{\textbf{B}}^3 u_0 u_1 u_2 } \\ \nonumber y_2 &= \overline{\textbf{C}}\overline{\textbf{A}}^2\overline{\textbf{B}} u_0 + \overline{\textbf{C}}\overline{\textbf{A}}\overline{\textbf{B}} u_1 + \overline{\textbf{C}}\overline{\textbf{B}} u_2 {\color{violet} +~ \overline{\textbf{C}}\overline{\textbf{A}}\overline{\textbf{B}}^2 u_0 u_1 +~ \overline{\textbf{C}}\overline{\textbf{A}}\overline{\textbf{B}}^2 u_0 u_2 +~ \overline{\textbf{C}}\overline{\textbf{B}}^2 u_1 u_2 +~ \overline{\textbf{C}}\overline{\textbf{B}}^3 u_0 u_1 u_2 },~~ \dots \nonumber\\\\ y_k &= \overline{\textbf{C}}\overline{\textbf{A}}^k\overline{\textbf{B}} u_0 + \overline{\textbf{C}}\overline{\textbf{A}}^{k-1}\overline{\textbf{B}} u_1 + \dots \overline{\textbf{C}}\overline{\textbf{A}}\overline{\textbf{B}} u_{k-1} + \overline{\textbf{C}}\overline{\textbf{B}} u_k \nonumber~+ \\ & {\color{violet} \sum_{p=2}^\mathcal{P} \forall \binom{k+1}{p} \text{~of~} u_i u_{i+1}~\dots~u_p~ \overline{\textbf{C}}\overline{\textbf{A}}^{(k+1-p-i)}\overline{\textbf{B}}^p u_i u_{i+1}~\dots~u_p} \nonumber ~~\text{for} ~i \in \mathbb{Z} \text{~and~} i \geq 0\\\\ &~~~~~\rightarrow~~~~~~~~y =~ \overline{\textbf{K}} * u \nonumber +{\color{violet} \overline{\textbf{K}}_{\text{liquid}} * u_{\text{correlations}}} \end{align}


你可以在公式中看到两种颜色。它们对应于两种类型的权重配置

  1. 黑色部分,是独立时间输入的权重,即 S4 的卷积核。
  2. 紫色部分,是与输入信号所有阶自相关相关的权重。这是一个额外的输入相关核,作者称之为液体核。

最后,卷积核表达如下:

KliquidRL~:=KL(C,A,B):=(CA(L~ip)Bp)i[L~], p[P]=(CAL~2B2,,CBp) \overline{\textbf{K}}_{\text{liquid}} \in \mathbb{R}^{\tilde{L}} := \mathcal{K}_L(\overline{\textbf{C}},\overline{\textbf{A}},\overline{\textbf{B}}) := \big(\overline{\textbf{C}}\overline{\textbf{A}}^{(\tilde{L}-i-p)}\overline{\textbf{B}}^p\big)_{i \in [\tilde{L}],~ p \in [\mathcal{P}]} = \big( \overline{\textbf{C}}\overline{\textbf{A}}^{\tilde{L}-2}\overline{\textbf{B}}^2, \dots, \overline{\textbf{C}}\overline{\textbf{B}}^p \big)

作者随后指出,这可以通过类似于 S4 中应用的过程(HiPPO、Woodbury、逆傅里叶变换等)高效计算。有关更多细节,请参阅论文中的算法 1。


在 LRA 上测试,这种方法似乎是最好的。只有 Mega(早几天发布,因此未出现在论文中)表现更好

图 16:Liquid-S4 在 LRA 基准测试上的结果

HASANI、LECHNER 等人还将他们的模型应用于 PIMENTEL 等人的 Speech Commands、sCIFAR 和 BIDMC Vital Signs 数据集,并建立了新的最先进水平。

深入探讨
官方实现可在 GitHub 上获取。
Open Review 上的交流。
官方 LTC 实现可在 GitHub 上获取。



S5:序列建模的简化状态空间层

按时间顺序,SMITH、WARRINGTON 和 LINDERMAN 引入 S5 模型的论文 Simplified State Space Layers for Sequence Modeling 于 2022 年 8 月 9 日发布,早于 Mega 和 Liquid-S5。然而,我在这之后才开始讨论这篇论文,因为在 2022 年 10 月 6 日,S5 的作者更新了他们的论文,将他们的模型在 LRA 上的性能比 V1 提高了 5 个百分点以上。更重要的是,他们提出了一个涵盖 2022 年发布的所有 SSM 的比较。从 V2 版本来介绍 S5 似乎更合适。

在 S5 中,作者提出用多输入/多输出 (MIMO) SSM 代替 S4 的单输入/单输出 (SISO) SSM 库的公式,该 MIMO SSM 具有减小的潜在维度。

图 17:S4 与 S5 内部行为比较

MIMO 系统减小的潜在维度允许使用并行扫描算法,这简化了将 S5 层作为序列到序列转换所需的计算。因此,由此产生的模型失去了 SSM 的卷积视图,而只关注循环视图(通过 ZOH 离散化获得)。作者的方法是在时域而不是频域中操作。他们使用 HiPPO 矩阵的对角近似,从而为其 MIMO 系统实现高效的初始化和参数化。

图 19:S4 与 S5 行为的完整比较

由于未来其他 SSM(特别是 Mamba)也会使用并行扫描,让我们仔细看看它在 S5 中的工作原理,以便从本文开始就熟悉此算法。最简单的方法是使用论文附录 H 中给出的示例,其中作者将其应用于长度为 L=4L = 4 的序列。

要计算并行扫描,需要两件事

  • 扫描将操作的初始元素。
    我们将长度为 LL 的序列的初始元素定义为 c1:Lc_{1:L},使得每个元素 ckc_k 是元组 ck=(ck,a,ck,b):=(A,Buk)c_k = (c_{k,a}, c_{k,b}) := (\overline{\mathbf{A}},\enspace \overline{\mathbf{B}}u_k)。在 L=4L = 4 的情况下,我们因此有 (A,Bu1),(A,Bu2),(A,Bu3)(\overline{\mathbf{A}},\enspace \overline{\mathbf{B}}u_1), (\overline{\mathbf{A}}, \enspace \overline{\mathbf{B}}u_2),(\overline{\mathbf{A}},\enspace \overline{\mathbf{B}}u_3)(A,Bu4)(\overline{\mathbf{A}},\enspace \overline{\mathbf{B}}u_4)
  • 用于组合元素的二元结合运算符。数学上,二元结合运算符是 IJK=(IJ)K=I(JK)I \bullet J \bullet K = (I \bullet J)\bullet K = I \bullet (J \bullet K)

如果我们按顺序进行扫描,令 s0:=(I,0)s_0:=(\mathbf{I},0),我们必须执行 4 次计算才能得到 4 个输出 sis_i
s1=s0c1=(I,0)(A,Bu1)=(AI,A0+Bu1)=(A,Bu1)s_1 = s_0 \bullet c_1 = (\mathbf{I},\enspace 0) \bullet (\overline{\mathbf{A}},\enspace \overline{\mathbf{B}}u_1) = (\overline{\mathbf{A}}\mathbf{I},\enspace \overline{\mathbf{A}}0+\overline{\mathbf{B}}u_1) = (\overline{\mathbf{A}},\enspace \overline{\mathbf{B}}u_1)
s2=s1c2=(A,Bu1)(A,Bu2)=(A2,ABu1+Bu2)s_2 = s_1 \bullet c_2 =(\overline{\mathbf{A}},\enspace \overline{\mathbf{B}}u_1) \bullet (\overline{\mathbf{A}},\enspace \overline{\mathbf{B}}u_2) = (\overline{\mathbf{A}}^2,\enspace \overline{\mathbf{A}}\overline{\mathbf{B}}u_1 + \overline{\mathbf{B}}u_2 )
s3=s2c3=(A2,ABu1+Bu2)(A,Bu3)=(A3,A2Bu1+ABu2+Bu3)s_3 = s_2 \bullet c_3 = (\overline{\mathbf{A}}^2,\enspace \overline{\mathbf{A}}\overline{\mathbf{B}}u_1 + \overline{\mathbf{B}}u_2 ) \bullet (\overline{\mathbf{A}},\enspace \overline{\mathbf{B}}u_3) = (\overline{\mathbf{A}}^3,\enspace \overline{\mathbf{A}}^2\overline{\mathbf{B}}u_1 + \overline{\mathbf{A}}\overline{\mathbf{B}}u_2 + \overline{\mathbf{B}}u_3)
s4=s3c4=(A3,A2Bu1+ABu2+Bu3)(A,Bu4)=(A4,A3Bu1+A2Bu2+ABu3+Bu4).s_4 = s_3 \bullet c_4 = (\overline{\mathbf{A}}^3,\enspace \overline{\mathbf{A}}^2\overline{\mathbf{B}}u_1 + \overline{\mathbf{A}}\overline{\mathbf{B}}u_2 + \overline{\mathbf{B}}u_3) \bullet (\overline{\mathbf{A}},\enspace \overline{\mathbf{B}}u_4)\\ = (\overline{\mathbf{A}}^4, \enspace \overline{\mathbf{A}}^3\overline{\mathbf{B}}u_1 + \overline{\mathbf{A}}^2\overline{\mathbf{B}}u_2 + \overline{\mathbf{A}}\overline{\mathbf{B}}u_3 + \overline{\mathbf{B}}u_4).

为了获得 xix_i 状态,我们必须取每个 sis_i 元组的第二个元素。

按顺序处理并非最有效的方式,因为可以通过并行扫描(parallel scan)来并行化计算递推关系。以下是它在我们长度为 LL = 4 的序列中的工作原理示意图。

image
图 19:S5 中并行扫描的工作原理

同样,要获得 xix_i 状态,我们必须取每个 sis_i 元组的第二个元素。
你会注意到这里可以并行计算 s2s_2i4i_4,然后并行计算 s1s_1s3s_3s4s_4。这将顺序计算的数量从 4 减少到仅 2。并行扫描的复杂度是 O(log(L))O(log(L))

image
图 20:并行扫描的工作原理(一般情况)。基于 Scott Linderman 的动画。

在性能方面,S5 在 LRA 上排名第二。

image
图 21:S5 在 LRA 基准测试上的结果

除了 LRA,S5 的作者还在 语音命令(Speech Commands)摆锤回归数据集(pendulum regression dataset)以及 sMNISTpsMNISTsCIFAR 上比较了他们的模型。完整的结果可在论文的附录中找到,其中还包含一项消融研究(ablation study)。

深入探讨
官方实现可在 GitHub 上获取。
Open Review 上的讨论。

SGConv

2022 年 10 月 17 日,LI、CAI 等人在其论文 What Makes Convolutional Models Great on Long Sequence Modeling? 中指出,他们发现 S4 过于复杂,需要复杂的参数化和初始化方案(= HiPPO)。因此,对于先验知识有限的人来说,它不那么直观且难以使用。所以他们的目标是通过专注于其卷积视角来揭示 S4 的神秘面纱。他们识别出 S4 受益的两个关键原则,这些原则足以构成一个高性能的全局卷积模型:

  1. 卷积核的参数化必须高效,即参数数量必须与序列长度呈次线性增长。
  2. 核必须具有递减结构,其中与最近邻居的卷积权重大于与最远邻居的卷积权重。
image
图 22:遵守作者提出的两个关键原则意味着卷积核应与图中所示的相似。

基于这两个原则,他们提出了一种高效的卷积模型,名为结构化全局卷积(Structured Global Convolution,SGConv)。

图 23:SGConv 将卷积核构建为连续的、长度更长、范数更低的弦函数的连接。这种形式的优点是它可以在频域实现非常快的卷积。

SGConv 作者报告称,在多项任务(文本、音频、图像)上取得了比 S4 更好的结果。我们在此不赘述细节。只看 LRA 的结果:

image
图 24:SGConv 在 LRA 上的结果。

确实,从这张表格中我们可以看到 SGConv 的表现优于 S4 的两个版本。然而,令人好奇的是,作者没有将 Mega、Liquid-S4 或 S5 纳入比较,这些模型都通过使用由递减指数函数之和构成的卷积核取得了更好的结果。
更重要的是,尽管所有在 LRA 上评估的模型都将数据视为 1D 序列,但 SGConv 却隐式地为图像任务(包括 PathX)引入了 2D 归纳偏置,这值得商榷。

最终,SGConv 的表现似乎与最新的 SSM 变体相似,但失去了 S4 的循环视图。
然而,这篇论文似乎是第一篇完全专注于 SSM 卷积视图的论文。

深入探讨
官方实现可在 GitHub 上获取。
Open Review 上的讨论。



其他模型

2022 年发表了另外两篇“理论”论文:WANG 等人引入 BiGS 的 Pretraining Without Attention 和 FU、DAO 等人引入 H3 的 Hungry Hungry Hippos: Towards Language Modeling with State Space Models
由于这些模型发表较晚(分别为 2022 年 12 月 20 日和 28 日),并且它们的 V2 版本是在 2023 年发布,我将在 SSM 系列的下一篇文章中介绍这些模型。



SSM 应用

SaShiMi

在 2022 年 2 月 20 日发表的论文 It's Raw! Audio Generation with State-Space Models 中,GOEL、GU 等人将 S4 应用于因果音频生成。
与基于文本、频谱图等进行条件生成的方法不同,这是一种直接作用于输入信号的方法,可以与 OORD 等人(2016)的 WaveNet 进行比较。
SaShiMi 可以在单个 V100 GPU 上直接对超过 100K(8 秒音频)的序列进行训练,这与 WaveNet 等模型面临的上下文长度限制形成对比。它有效利用了这种长上下文来提高密度估计。
作者在各种基准测试中比较了他们的模型,包括钢琴音乐生成和语音(数字发音)。
生成的音频可在此处查看:https://hazyresearch.stanford.edu/sashimi-examples/

image
图 25:Sashimi 架构概述

深入探讨
官方实现可在 GitHub 上获取。



ViS4mer

2022 年 4 月 4 日,ISLAM 和 BERTASIUS 在他们的论文 Long Movie Clip Classification with State-Space Video Models 中引入了 ViS4mer。
这是一种用于(长)视频分类的 S4 和 Transformer 混合模型。更确切地说,该模型使用标准的 Transformer 编码器进行短程时空特征提取,以及多尺度时间 S4 解码器进行长程时间推理。结果表明,该模型比 Transformer 快 2.6 倍,内存效率高 8 倍。
据我所知,这是第一篇强调 SSM 和 Transformer 混合优势的论文。

图 26:ViS4mer 解码器概述

ViS4mer 在 WU 和 KRÄHENBÜHL (2021) 的基准测试 Long Video Understanding (LVU) 中的 9 项长视频分类任务中,有 6 项取得了最佳结果。LVU 包括对时长 1 到 3 分钟的视频进行分类。该模型还表现出良好的泛化能力,尽管数据量少了 275 倍,但在 Breakfast 和 COIN 程序活动数据集上取得了有竞争力的结果。

深入探讨
官方实现可在 GitHub 上获取。



CCNN

2022 年 6 月 7 日,ROMERO、KNIGGE 等人在他们的论文 Towards a General Purpose CNN for Long Range Dependencies in ND 中介绍了 CCNN。
他们的出发点是卷积神经网络功能强大,但需要针对每个任务进行专门调整:

  • 输入长度:32x32,1024x1024 → 如何建模长程依赖?
  • 输入分辨率:8kHz,16kHz → 长程依赖,分辨率无关性?
  • 输入维度:1D,2D,3D → 如何定义卷积核?
  • 任务:分类、分割等 → 如何定义高低采样策略?
    那么,是否有可能设计一个独特的架构,无论维度、分辨率和输入长度如何,都能独立解决任务,而无需修改架构?答案是肯定的,这要归功于使用连续卷积核的 CCNN。

ROMERO、KNIGGE 等人从 S4 中汲取灵感,创建了一种高效残差块的变体,他们称之为 S4 块。然而,与仅适用于 1D 信号的 S4 不同,CCNN 可以轻松建模 ND 信号。

图 27:CCNN 概述

深入探讨
官方实现可在 GitHub 上获取。
论文演示幻灯片可在此处获取:https://app.slidebean.com/sbp/gk5j826nq7/CCNN#1



SSSDS4\mathbf{SSSD^{S4}}

2022 年 8 月 19 日,LOPEZ ALCARAZ 和 STRODTHOFF 在其论文 Diffusion-based Time Series Imputation and Forecasting with Structured State Space Models 中提出了一种 S4 和扩散模型的混合模型,用于预测时间序列中的缺失数据。他们的模型名为 SSSDS4\mathbf{SSSD^{S4}}(或简称 SSSD)。

image
图 28:SSSD 概述

深入探讨
官方实现可在 GitHub 上获取。



S4ND

2022 年 10 月 12 日,NGUYEN、GOEL、GU 等人提出了 S4ND: Modeling Images and Videos as Multidimensional Signals Using State Spaces

该模型将 S4(1D)扩展到多维连续信号,如图像和视频(而 ConvNets 和 ViT 在离散像素上学习)。为此,他们将标准 S4 ODE 转换为多维 PDE:

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

变为

t(1)x(t(1),t(2))=(A(1)x(1)(t(1),t(2)),x(2)(t(1),t(2)))+B(1)u(t(1),t(2))t(2)x(t(1),t(2))=(x(1)(t(1),t(2)),A(2)x(2)(t(1),t(2)))+B(2)u(t(1),t(2))y(t(1),t(2))=C,x(t(1),t(2)) \begin{aligned} \frac{\partial}{\partial t^{(1)}} x(t^{(1)}, t^{(2)}) &= (\mathbf{A}^{(1)} x^{(1)}(t^{(1)}, t^{(2)}), x^{(2)}(t^{(1)}, t^{(2)})) + \mathbf{B}^{(1)} u(t^{(1)}, t^{(2)}) \\ \frac{\partial}{\partial t^{(2)}} x(t^{(1)}, t^{(2)}) &= (x^{(1)}(t^{(1)}, t^{(2)}), \mathbf{A}^{(2)} x^{(2)}(t^{(1)}, t^{(2)})) + \mathbf{B}^{(2)} u(t^{(1)}, t^{(2)}) \\ y(t^{(1)}, t^{(2)}) &= \langle \mathbf{C}, x(t^{(1)}, t^{(2)}) \rangle \end{aligned}

其中 A(τ)CN(τ)×N(τ)\mathbf{A}^{(\tau)} \in \mathbb{C}^{N^{(\tau)} \times N^{(\tau)}}, B(τ)CN(τ)×1\mathbf{B}^{(\tau)} \in \mathbb{C}^{N^{(\tau)} \times 1}, CCN(1)×N(2) \mathbf{C} \in \mathbb{C}^{N^{(1)} \times N^{(2)}},线性偏微分方程的初始条件为 x(0,0)=0x(0, 0) = 0

根据所测试的数据集,作者获得了与ViT或ConvNext相似或更好的结果。
S4ND 的主要优点是它可以通过不同的采样率在不同分辨率下运行。作者在两个实验中强调了这一特点:

  1. 在零样本情况下,S4ND 在 8×88\times8 图像上进行训练,并在 32×3232\times32 图像上进行测试时,其性能优于 Conv2D 40 多个百分点。
  2. 通过渐进式尺寸调整,S4ND 可以将训练速度提高 22%,而最终准确度与仅在高分辨率下训练相比下降约 1%。
image
图 30:S4ND 在 2D 图像中的示例

深入探讨
官方实现可在 GitHub 上获取。
一个完全无用但有趣的资料:作者将他们的 TeX 文件命名为 Darude S4NDstorm




结论

因此,我们回顾了 2022 年发布的各种 SSM 模型。在这一年中,主要工作集中于通过各种方法(对角化、门控、LTC 等)改进/简化 S4。2022 年也首次出现了 SSM 的应用。
通过 SGConv 和 S5,我们也可以看到一种现象的开始,正如我们将在下一篇文章中看到的,这种现象在 2023 年将变得更加明显。即,只关注 SSM 卷积视图(例如 Hyena 及其派生模型)或只关注 SSM 递归视图(例如 Mamba)的工作的出现。


参考文献

引用

@inproceedings{ssm_in_2022_blog_post,
作者 = {Loïck BOURDOIS},
标题 = {2022年状态空间模型(SSM)的演变},
年份 = {2023},
网址 = {https://lbourdois.github.io/blog/ssm/ssm_en_2022}
}



Visitors

社区

注册登录 发表评论