【论文翻译】SSM状态空间模型的循环和卷积模式

这篇文章是Mamba作者博士论文 MODELING SEQUENCES WITH STRUCTURED STATE SPACES
的第二章的部分翻译，这部分关于SSM的卷积和循环模式表达。

2.3.2 递归表达（高效推理）

离散化后，离散SSM可以被定义为：
$$\begin{gathered} x_k=\overline{\text{A}}{x}_{k-1}+\overline{\text{B}}{u}_k \\ {y}_k=\text{C}{x}_k \end{gathered}$$
$\overline{\text{A}}\in {\mathcal{R}}^{N\times N}$ $\overline{\text{B}}\in {\mathcal{R}}^{N\times M}$ $u_k\in {\mathcal{R}}^N$ $x_k\in {\mathcal{R}}^M$

等式现在是sequence-to-sequence映射$(u_k)\mapsto (y_k)$ 而且状态等式是$x_k$的递归，可以像其他例如卷积RNN的循环网络一样计算，用RNN的语言说，$x_k\in {\mathcal{R}}^N$可以被看作有转移矩阵$\overline{\text{A}}$的隐状态。

循环模式解释

循环状态$x_k\in {\mathcal{R}}^N$携带了时间$k$之前的所有输入信息，换句话说，当展开上式的递归，只有循环状态需要保留。因此SSM和常规的循环模型有高效的基于状态的推理：不像CNN和Transformer，它们可以被推测一个（潜在无穷）输入序列每时间步仅需要常数计算和空间。

对于在线和自回归情形下，特别有用。值得注意的例子是自回归生成模型例如现代大语言模型，生成需要一次展开模型一步。

2.3.3 卷积表达（高效训练）

循环SSM在现代硬件上训练是不实际的，因为其序列性。而例如线性时不变(LTI)SSM和连续卷积有众所周知的联系。对应的，上面的等式也可以被写成离散卷积。

定义初始状态$x_{-1}=0$。因为离散卷积是线性的，它可以被分析展开，推出SSM的一个卷积形式。
$$\begin{array}{rlrlr}{x}_0& =\overline{B}{u}_0& x_1=\overline{AB}{u}_0+\overline{B}{u}_1& x_2={\overline{A}}^2\overline{B}{u}_0+\overline{AB}{u}_1+\overline{B}{u}_2& \dots \\ y_0& =C\overline{B}{u}_0& y_1=C\overline{AB}{u}_0+C\overline{B}{u}_1& y_2=C{\overline{A}}^2\overline{B}{u}_0+C\overline{AB}{u}_1+C\overline{B}{u}_2& \dots \end{array}$$
换句话说，根据输入$u$可以推出输出$y$的一个简单的闭式解：
$$\begin{array}{rlr}{y}_k& =C{\overline{A}}^k\overline{B}{u}_0+C{\overline{A}}^{k-1}\overline{B}{u}_1+\cdots +C\overline{A}\overline{B}{u}_{k-1}+C\overline{B}{u}_k.& (2.6)\end{array}$$
可以用一个显式kernel被向量化成一个简单卷积：
$$\begin{array}{rl}y& =u\ast \overline{K}\mathrm{where}\overline{K}=(C\overline{B},C\overline{A}\overline{B},\dots ,C{\overline{A}}^k\overline{B},\dots ).(2.7)\end{array}$$
上式是一个简单（non-circular）卷积，$\overline{\text{K}}$有多种称呼，包括SSM（卷积）滤波器或核，或者简称为state space kernel（SSK）

计算可以在整个序列长度上并行化，而不是从每时刻从$u_k$计算$y_k$,因为$y_k$的闭式解存在，从完整的输入序列$u$可以并行计算出完整的输出序列$y$

卷积模式解释

SSSM可以非常相似地解释为CNN中的线性卷积层。主要的区别是状态空间核实际上是无限长的。

Remark 2.2 注意到 SSM核$\overline{\text{K}}$实际上是无限长的，但是受限于输入的长度$L$。它可以进一步截断为更短的长度以表现地更像一个传统CNN（例如强调局部性，或者为了计算高效）。我们通常重载表示，使用$\overline{\text{K}}$直接指代输入长度$L$的截断版本。
$$\overline{K}:={\mathcal{H}}_L(\overline{A},\overline{B},C):={\left(C{\overline{A}}^k\overline{B}\right)}_{k\in [L]}=(C\overline{B},C\overline{A}\overline{B},\dots ,C{\overline{A}}^{L-1}\overline{B}).(2.8)$$
这导致了相比于传统CNN的其他不同。首先SSSM的卷积形式可以选择通过3个FFT(假设核$\overline{\text{K}}$是已知的)高效计算。在数学上等价于一个传统的"密集卷积"但是取决于核大小有不同的计算复杂复杂度。

Proposition 2.8 （FFT-卷积）令$u,k$分别为长度为$L$的$1-D$序列.圆周卷积$u\ast k$等价于${\mathcal{F}}^{-1}(F(u)\circ \mathcal{F}(k))$，$\mathcal{F}$是傅里叶变换。

Proposition 2.9 （因果卷积） 长度$L$的$u$和长度为$K$的核$k$的因果（非循环）卷积可以padding到长度$L+K$来引用循环卷积。

Corollary 2.10 （卷积复杂度）对于一个序列长度为$L$核长度为$K$

• 一个传统的密集卷积有复杂度$O(LK)$
• 一个$FFT$卷积有复杂度$O((L+K)log(L+K))$ 有更多的常数

第二，卷积核是隐式的，因为它是无限长的，他是从一系列更小参数中产生的，在这里通过式2.7$(\Delta ,\text{A},\text{B},\text{C})\to \overline{\text{K}}$许多其他方法创建隐式核，但是通常缺乏SSSM的其他（连续性和循环性）性质。

最后我们注意到当计算长卷积$y=u\ast \overline{\text{K}}$通过标准FFT卷积技术是快速的，但生成2.8中的隐式核$\overline{\text{K}}$是不容易的，是Chapter 3关注的重点。

Remark 2.3 我们称为状态空间核的东西也被称为"Markov parameters"在SSM文献中。我们的术语强调SSM和SSSM之间的区别，后者更接近CNN，参数被称为卷积滤波器或者核。

2.3.4 SSM表达的总结

总之，SS(S)M层作为序列模型的前向传播由一个离散化步骤构成，接着不同的等价方法计算离散SSM。我们称之为第一离散化的计算图之后是2.5为循环模式，相似地，我们称第一离散化的计算图和计算式2.7为卷积模式。

Remark 2.4 我们有时把这些简称为“RNN”模式和"CNN"模式。注意一层SSSM层并不是真正的神经网络-仅仅是一个线性序列变换-因此它们本身并不是一个RNN或者CNN。

Remark 2.5 这部分展示了从循环角度的离散化， 之后从展开递归推导离散卷积形式。离散卷积形式也可以通过直接离散化连续卷积形式得到。

$y(t)=(K\ast u)(t)={\int }_0^{\infty }\text{C}{e}^{s\text{A}}\text{B}u(t-s)ds$

这涉及保持模型卷积结构的方式近似积分（这种情况通常被称为正交）

Remark 2.6 如2.2.1指出的，我们在这个部分只考虑T-SSM，对于卷积和SSM之间的等效是必要的。我们注意到对于通常的时变SSM,推广这两种模式是可能的。

2.3.5 关于SSM维度

在定义2.1中，我们定义序列模型为映射${\mathcal{R}}^{L\times D}\to {\mathcal{R}}^{L\times D}$。从另一个角度，SSM和因此SSSM被定义为一个${\mathcal{R}}^{L\times M}\to {\mathcal{R}}^{L\times M}$。我们把$D$看作输入的维度（独立于模型定义）$M$作为一个定义模型的超参数。我们约定$M$必须被整除$D$，对于$D>M$的输入，通过切分模型维度$D$到$M$，对于每个这些输入${\mathcal{R}}^{L\times M}$最后把$D/M$个SSM的输出拼接在一块。

注意这实际上在模拟多头注意力是如何实现的，因此这也可以被解释为一个多头SSM，其中$H=D/N$是头的数目。

Definition 2.11 不同模型维度的SSM变体

• 单输入单输出（SISO）SSM $M=1$

• 多输入多输出（MIMO）SSM $M>1$

• 我们通常使用MIMO来 指代最大MIMO $M=D$

• 多头SSM强调$M<D$的情况，我们定义$H=M/D$

  注意一个SISO SSM是一个有最大head数$H=D$的多头SSM

Definition 2.12 （SSM 维度）重要维度的术语：

• $D$是模型的尺寸或者维度
• $N$是状态尺寸或维度，或者模糊地称为“SSM 状态尺寸”
• $NH=ND/M$是总的状态尺寸

例如，一个（最大）MIMO SSM有总状态尺寸$N$，而一个SISO SSM有总状态尺寸$ND$

当我们把SSM和其它模型做比较时，对RNN和CNN风格进行对比是有用的。

Remark 2.7 （与RNN维度的关系）传统地，RNN相对应于$MIMO(M=DorH=1)$的情况。总的状态维度是$N$，对于模型维度$D$来说是一个独立的超参数。通过投影矩阵$\overline{\text{B}}$和$\text{C}$输入在这些维度间变换。

许多最近的RNN变种转变成$SISO$类型，其$D$个通道被一个$1-D$递归独立地处理。值得注意的例子包括LSTM的cell state,QRNN,SRU,indRNN和RWKV。对于维度，这些模型对应一个$M=1,H=D$和$N=1$的SSM，总的状态大小是$D$

Remark 2.8（和CNN维度的关系）对于一个$M=D$的$MIMOSSM$ ，$\text{C}{\overline{\text{A}}}^i\overline{\text{B}}\in {\mathcal{R}}^{M\times M}$和$\overline{\text{K}}$是一个$(L,M,M)$的张量。在标准CNN形式下解释被$\overline{\text{K}}$卷积，准备些维度可以表示$(sequence,out\_channels,in\_channels)$

一个有$D|M$的多头SSM在CNN文献中被称作分组卷积。

当$M=1$或者输入维度是独立通道没有交互，在CNN文献中称为深度（或深度可分离）卷积。

在本文剩余部分中，我们只关注$M=1$和$H=D$的$SISO$情况

主要动机是一维输入下，在线记忆的理论框架。然而，我们注意到文献中许多流行在经验上有效的模型也采取了独立输入维度的版本，例如RWKV RNN和ConvNeXt CNN。我们还注意到，S4的大多数后序模型都保留了SISO形式，尽管出现了一些切换到MIMO公式的变体，如S5。