【AI学习】Mamba学习（十二）：深入理解S4模型

#1024程序员节｜征文#

HiPPO的学习暂告一段落，按照“HiPPO->S4->Mamba 演化历程”，接着学习S4。

S4对应的论文：《Efficiently Modeling Long Sequences with Structured State Spaces》
文章链接：https://ar5iv.labs.arxiv.org/html/2111.00396，https://arxiv.org/abs/2111.00396

上一篇文章初步了解了S4（Structured State Space sequence model）这个序列模型，接下来深入学习一下。

文章的第二章讲述了算法背景，接着学习。从第二章可以看出，发展到S4模型，作者才引入了SSM，而HiPPO那篇论文是没有提到SSM的。

2、S4背景：状态空间

图1中的四个属性描述了状态空间模型（SSM）：经典的连续时间表示、用HiPPO框架解决长距离依赖问题、离散时间递归表示，以及可并行化的卷积表示。特别是2.4节介绍了SSM卷积核K，这是我们在第3节中理论贡献的重点。

2.1 状态空间模型：连续时间潜在状态模型

状态空间模型由简单的方程（1）定义。它将一维输入信号u(t)映射到N维潜在状态x(t)，然后投影到一维输出信号y(t)。

SSM在许多科学学科中广泛使用，并且与潜在状态模型如隐马尔可夫模型（HMM）相关。我们的目标是简单地将SSM作为一个黑盒表示，用于深度序列模型，其中A、B、C、D是梯度下降学习的参数。在本文的其余部分，为了表述方便，我们将省略参数D（或者等价地，假设D = 0），因为项Du可以被视为一个跳跃连接，并且容易计算。

前面说过，Du项相当于残差连接，而整个模型中一般也有残差连接，可以合并到整个模型，而在SSM方程这里省略。

2.2 用HiPPO解决长距离依赖问题

先前的工作发现基本的SSM（1）在实践中实际上表现非常差。直观地，一个解释是线性一阶ODEs解决为指数函数，因此可能遭受梯度在序列长度上呈指数增长（即，梯度消失/爆炸问题[32]）。为了解决这个问题，线性状态空间层（LSSL）利用了连续时间记忆的HiPPO理论[16]。HiPPO指定了一类特定的矩阵A∈RN×N，当它们被纳入（1）时，允许状态x(t)记忆输入u(t)的历史。这个类别中最重要的矩阵由方程（2）定义，我们将称之为HiPPO矩阵。例如，LSSL发现简单地将SSM从随机矩阵A修改为方程（2）将其在顺序MNIST基准测试上的性能从60%提高到98%。

2.3 离散时间SSM：递归表示

为了在离散输入序列（u0, u1, …）上应用而不是连续函数u(t)，（1）必须通过一个步长∆来离散化，该步长表示输入的分辨率。从概念上讲，输入uk可以被视为对隐含的连续信号u(t)的采样，其中uk = u(k∆)。
为了离散化连续时间SSM，我们遵循先前的工作，使用双线性方法[43]，将状态矩阵A转换为近似A。离散SSM是：

方程（3）现在是序列到序列的映射uk → yk，而不是函数到函数的映射。此外，状态方程现在是一个递归的xk，允许离散SSM像RNN一样计算。具体来说，xk ∈ RN可以被视为具有转换矩阵的隐藏状态。

从符号上讲，本文中我们使用、、…来表示由（3）定义的离散SSM矩阵。注意，这些矩阵是A和步长∆的函数；当它清晰时，我们为了符号方便而省略了这种依赖。

2.4 训练SSM：卷积表示

递归SSM（3）不适用于在现代硬件上进行训练，因为它是顺序的。相反，线性时不变（LTI）SSM（如（1））和连续卷积之间有一个众所周知的联系。相应地，（3）实际上可以写成一个离散卷积。
为了简单起见，让初始状态为x−1 = 0。然后展开（3）得到

这可以向量化为一个卷积（4），并为卷积核（5）提供一个明确公式。

换句话说，方程（4）是一个单一的（非循环）卷积，并且可以使用FFTs高效计算，前提是K是已知的。然而，计算（5）中的K并非易事，这是我们在第3节中技术贡献的重点。我们称K为SSM卷积核或滤波器。

个人说明

前面介绍部分就提到“SSM由于具体的理论原因，它尚未适用于深度学习。深度 SSM 实际上即使在简单的任务上也很困难，但当配备最近为解决连续时间记忆（就是HiPPO）问题而导出的特殊状态矩阵 A 时，可以表现得非常好”。第二章部分进一步阐释，“基本的 SSM在实践中表现非常差。直观地说，一种解释是线性一阶常微分方程解为指数函数，因此可能会受到序列长度中梯度指数缩放的影响（即消失/爆炸梯度问题）。”其实，熟悉RNN的训练就会知道，RNN的训练就面临梯度爆炸或消失的问题。

S4 这里选取的矩阵 A 为 HiPPO-LegS的矩阵形式， 但是LegS 所满足的 ODE 是矩阵A、B下面要除以t的。就是下面的形式：

不知道为啥HiPPO-LegS的矩阵形式直接用到SSM方程效果就可以很好。
这里面有点奇怪，回头需要再看看相关论文：《Legendre memory units: Continuous-time representation in recurrent neural networks. In Advances in Neural Information Processing Systems, pages 15544–15553, 2019.》
苏神在文章《重温状态空间模型SSM：HiPPO的高效计算（S4）》也进行了分析，这种修改之后，导致时间度量无法对整个历史平均的取权重，而是历史信息呈现指数衰减的效应，就是更看重时间较近的信息。

离散时间 SSM这里的离散化还是使用了双线性的近似方式，没有采用解析解。要在后面Mamba模型才使用离散化的解析解。

在前面介绍部分，就提到，深层 SSM 原则上可以解决 LRD 问题，但是由于状态表示引起的计算和内存要求过高，LSSL 在实践中不可行。以及，“尽管提出了 LSSL 的 理论上有效的算法，但这些算法在数值上是不稳定的。特别是，特殊 A 矩阵在线性代数意义上是高度非正规的，这阻碍了传统算法技术的应用。因此，尽管 LSSL 表明 SSM 具有很强的性能，但它们目前作为通用序列建模解决方案在计算上是不切实际的。”这就引出了下一章，具体的S4方法。

注：正规矩阵的一个重要性质是它可以经过一个‌酉变换变为对角矩阵。厄米矩阵（Hermitian matrices）和酉矩阵（unitary matrices）都是正规矩阵的典型例子。

3、Method: Structured State Spaces (S4)

具体方法，咱们直接跳到最终方法部分：

这个方法，就是把HiPPO 矩阵A分解为正规矩阵和低秩矩阵的和。这个分解的目的，就是为了方便的实现矩阵A的对角化。
为什么要对角化？为了计算简便。前面的公式5的卷积形式，矩阵A的L次幂，如果是对角矩阵，计算就可以大大简化。

文章3.1节，给出了这种对角化动机的目的。

但是文章接着说，不幸的是，由于数值问题，对角化的简单应用在实践中不可行。尽管有其他方法，数值上也不稳定。

那怎么办，3.2节给出方法：将HiPPO 矩阵A分解为正态矩阵和低秩矩阵的和。这样处理后获得了一个反对称矩阵，“重点来了，反对称矩阵不单单一定可以对角化，它一定可以被正交矩阵（复数域叫做酉矩阵）对角化！酉矩阵一般数值稳定性都非常好”。

这里的生成函数如何理解？其实也简单，熟悉卷积运算的就知道，卷积运算计算量大，可以先做FFT，在频域变成乘法，然后IFFT。这是利用FFT的简化卷积运算经常使用的方法。只不过，这里傅立叶变换所需要的实际是“截断生成函数”，将无限长度截断为L。

最后，再总结一下S4的架构细节。

具体第三章涉及的公式推导，可以参见苏神在文章《重温状态空间模型SSM：HiPPO的高效计算（S4）》中的详细推导。

最后

再重温一下上篇《Mamba学习（十一）：S4》提到的问题和S4的方法：

解决的主要问题

• 序列建模中长距离依赖（LRDs）的有效处理。
• 现有模型在处理非常长序列时的局限性。
• 先前基于SSM的方法在计算和内存需求上的高成本问题。

方法

• 提出了S4模型，它基于SSM的新参数化，通过将状态矩阵A分解为低秩和正规项的和，使得A可以被稳定地对角化。
• 利用Woodbury identity和Cauchy核的计算，将SSM的计算复杂度从O(N^2L)降低到O(N+L)，其中N是状态维度，L是序列长度。
• 在多个任务和数据集上验证了S4模型的性能，包括顺序CIFAR-10、WikiText-103语言建模、图像分类和时间序列预测等。