【论文翻译】MODELING SEQUENCES WITH STRUCTURED STATE SPACES 第一章

这篇文章是Mamba作者博士论文 MODELING SEQUENCES WITH STRUCTURED STATE SPACES
的第一章介绍部分翻译。

1.1 深度序列模型

定义序列模型为参数化的序列到序列的转换，并作为深度学习模型的组件。状态空间序列模型是在一维序列上定义的简单线性映射。一个深度序列模型由核心的序列变换例如卷积，注意力或S4和额外的位置部分例如归一化层，线性层和残差连接组成。

$y=f_{\theta }(x)$ 输入$x$ 输出$y$ 是长度为$L$ 的在${\mathbb{R}}^D$上的特征向量，$\theta$ 是通过梯度下降学习到的参数。

Challenge 1：能力通用

RNN：需要快速更新隐藏状态例如在线处理和强化学习

CNN：建模均匀采样的感知信号，例如音频，图像和视频

Transformers：领域内建模密集，复杂交互例如语言

Challenge 2：计算效率

RNN很难在得益于并行性的现代硬件加速器上训练。CNN和Transformer因为不是stateful所以很难进行自回归推理。

Challenge 3：长程依赖

有限的上下文窗口，或者由于优化问题，例如RNN的梯度消失问题。正交和Lipschitz RNN，扩张卷积，和各种注意力变体。‘

1.2 状态空间模型

定义SSM为一个简单的序列模型，通过一个隐状态$x(t)\in {\mathbb{R}}^N$ 映射一个一维函数或者序列$u(t)$。

equation 1.1

$x^{{}^{\prime }}(t)=\text{A}x(t)+\text{B}u(t)$

$y(t)=\text{C}x(t)+\text{D}u(t)$

状态空间模型更广泛地指潜在变量如何在状态空间中进行演化建模的东西，传统上定义这些动态的概率模型。例如随机转移矩阵 。这些更广泛的SSM存在不同的风格，它们可以改变$x$的状态空间（例如连续，离散和混合空间），观测空间$y$，转换的动态，附加的噪声过程或者系统的线性。ssm通常指隐马尔可夫模型变体和线性动态系统的变体，例如分层狄利克雷过程和切换线性动态系统。

但是，1.1所示的状态空间模型在状态空间和动态上都是连续的，而且是完全线性和确定的，还没有被用作深度序列模型。这篇论文探讨了状态空间序列模型的许多优点，和如何使用它们来解决常见的序列建模问题。

1.2.1 一个通用的序列模型

SSM是一个简单基础的模型，有丰富的性质。它们和NDE，RNN和CNN等模型密切相关，并且实际可以用几种不同形式编写，因此可以获得通常需要专门的特殊模型的各种能力。

• SSM是连续的。SSM本身是一个微分方程。因此可以执行连续时间模型的独特应用，例如模拟连续过程，处理丢失的数据和适应不同的采样率。
• SSM是循环的。SSM可以使用标准技术离散化为一个线性循环，在推理期间模拟为一个每时间步有常数内存和计算量状态循环模型。
• SSM是卷积的。SSM是线性时不变系统，可以显式表示为连续卷积。除此之外，离散时间版本可以在训练期间使用离散卷积并行化，实现高效的训练。

因此，SSM作为一种通用的序列模型，在并行和串行settings和多种领域（例如音频，视觉，时间序列）中都是有效的。第二章介绍了SSM的背景和状态空间模型的性质。

然而，SSM的通用性也存在着trade-off，原始的SSM依然存在两个挑战，可能比其他模型更严重，阻碍了它们作为深度序列模型的使用。

(i)一般的SSM相比于同等大小的RNN和CNN慢很多。

(ii)它们很难记住长依赖，例如继承RNN的梯度消失问题。

我们用SSM的新算法和理论解决这些挑战。

1.2.2通过结构化SSM(S4)高效计算

不幸的是，状态表示引入的令人望而却步的计算和内存需求使得一般的SSM在实践中无法被用作深度序列模型。

对于SSM的状态维度$N$ 和序列长度$L$，单单计算完整的潜在空间状态需要$O(N^{2}L)$ 操作和$O(NL)$空间，相比于计算总体输出的$\Omega (L+N)$ 下界。因此，对于合理尺寸的模型，例如$N\approx 100$，一个SSM相比于CNN或RNN使用的内存多几个数量级，使得它在作为一般序列建模方案在计算上不实际。

克服这个计算瓶颈需要在状态矩阵$\text{A}$ 上施加结构使得它成为高效算法。我们介绍了具有各种形式的结构化矩阵$\text{A}$ 的结构化状态空间序列模型（S4）或简称为结构化状态空间的族，以及允许在其任何表示例如循环和卷积中进行高效优化。

Chapter3介绍了这些高效的S4模型的特点。第一个结构使用状态矩阵的对角参数化，非常简单和通用，足以表示几乎所有的SSM，然后我们通过一个低秩矫正项推广这一点，对于得到之后引入的一类特殊SSM使必要的。通过结合许多技术思想例如生成函数，线性代数变换和结构化矩阵乘法的结果，我们实现了具有$O(N+L)$时间和$O(N+L)$空间复杂度的算法，对于序列模型来说是紧密的。

1.2.3 通过HIPPO解决长距离依赖

独立于计算问题，一个基本的SSM依然在经验上表现不佳，并且难以对长距离依赖关系进行建模。直觉上，一个解释使线性一阶的ODE求解为指数函数，因此在序列长度上受梯度指数缩放的影响。也可以从线性循环角度解释，涉及到反复乘一个递归矩阵，这是众所周知的RNN的梯度消失和爆炸问题。

在Chapter4，我们从SSM退一步，转而研究如何使用第一原理的循环模型对LRDs进行建模。我们开发了一个称为HIPPO的数学框架，形式化和解决了一个我们称之为在线函数近似（记忆）的问题。这就产生了一些旨在通过保持连续函数的历史压缩来增量地记忆连续函数的方法。这些最终的方法被证明是SSM的特定形式，尽管它们的动机完全独立。

Chapter5细化了这个框架，并将其更严格地联系到SSM抽象。引入了正交SSM地概念，并将该概念推广到了HIPPO并引生出更多地实例和理论结果，例如如何以一种principled的方式初始化SSM的参数。

HIPPO概述

考虑一个输入函数$u(t)$ ，一个固定的概率测度$w(t)$ 和一个N个正交基函数如多项式函数的序列。在每个时刻$t$，$u$在时刻$t$之前的历史可以被投影到这个基上，产生一个系数向量$x(t)$表示$u$的历史相对于所提供度量$w$的最优逼近。函数$u(t)\in \mathbb{R}$ 到系数$x(t)$的映射成为关于度量$w$的高阶多项式投影算子。在很多情况下，它的形式是$x^{{}^{\prime }}=\text{A}x(t)+\text{B}u(t)$以及$(\text{A},\text{B})$的闭式解

结合HIPPO和S4

HIPPO提供了一个数学工具来创建有重要属性的SSM，而S4是关于计算表达的。Chapter 6 在形式上将它们结合在一起而且表明它们可以结合起来两全其美。我们表明用HIPPO生成的用于处理远程依赖关系的特殊矩阵实际上可以用Chapter 3的特定结构化形式编写。这提供了包含HIPPO的S4具体实例，它产生了一个具有丰富功能的通用序列模型，高效并且擅长远距离推理。

1.2.4 应用，消融和拓展

通用序列建模能力

Chapter 7 提供了S4方法在各种领域和任务的综合经验验证，当合并到一个通用简单深度神经网络中，其在众多基准中取得SOTA。

特别的亮点和功能包括：

• 通用序列建模

  在没有结构改变的情况下，S4在语音分类上超过了audio CNN，在时间序列预测问题上超过了专门的informer模型，搭配一个2D ResNet在序列CIFAR上取得超过90%的准确率

• 长程依赖

  在高效序列模型的LRA基准测试中，S4和所有基线一样快，同时比所有Transformer变体的平均准确率高出25%以上。S4是第一个解决困难LRA Path-X任务(length-16384)的模型，与所有先前工作的50%随机猜测相比，实现了96%的准确度

• 采样分辨率变化

  像特殊的NDE方法一样，S4可以适应时间序列采样频率的变化而无需重新训练。

• 快速自回归生成的大规模生成模型

  在CIFAR-10密度估计上，S4可以和最好的自回归模型竞争(2.85 bits per dim)在WikiText-103语言建模上，S4实际上缩小了和Transformers的差距(0.5困惑度以内)，达成了非注意力模型的SOTA。像RNN一样，S4可以使用它的潜在状态来生成像素和token，比CIFAR-10/WikiText-103上的标准自回归模型快60倍。

  理论消融

  我们对S4的处理讨论了训练SSM的许多技术细节，例如如何仔细地初始化每个参数和如何同HIPPO框架协作。这些都是全面的经验分析和消融，验证了我们的SSM理论的各个方面。

  例如，我们验证了HIPPO大大提高了SSM的建模能力，在标准序列模型基准测试中，与原始的SSM实例相比，性能提升了15%。在算法上，我们的S4算法比原始SSM提高了几个数量级（例如速度快30倍， 内存使用减少400倍）。

  应用：音频波形生成

  作为一个具有多种属性的序列模型原型，S4可以被整合到不同的神经网络体系结构中，并以多种方式使用。Chapter 8 展示了S4在原始音频波形生成中的应用，一个体照型的问题是音频波形的高采样率，它引入了围绕S4构建的SASHIMI多尺度结构

  ，在无条件音频和语音生成中，在多种生成模型范式如自回归和diffusion中取得了SOTA。该应用凸显了S4的灵活能力，包括高效训练，快速自回归生成和对建模连续信号建模的强归纳偏置。

  扩展：计算机视觉中的多维信号

  虽然我们主要关注一维序列，但一些数据形式天然是高维的，例如图像(2d)和视频(3d)。序列模型的灵活性也可以应用于这些设置。Chapter 9 介绍了S4ND是S4从一维到多维(N - D)信号的扩展。S4ND继承了S4的特性，例如直接对底层连续信号建模，并且具有更好地处理输入分辨率变化等相关优点，并且是第一个在大型视觉任务(如ImageNet)上具有竞争性的连续模型。

1.3 论文注意事项

简要指出一些可供选择地阅读路径

• Part I (Chapters 2 and 3) 是S4模型家族的独立描述，对于实践者来说，包含实现和使用这些模型所必需的所有材料。在Part III的应用前，Part II并不严格需要。
• Part II 的第一章 (Chapter 4)完全独立于 Part I, 是从另外的角度去驱动SSM.可以以任意的顺序去阅读，历史上是这一工作历程的先驱。