【AI学习】Mamba学习（十四）：S5模型的矩阵A对角化

上一篇简单了解了S5模型，S5在两个主要方面简化了S4层。首先，S5使用一个多输入多输出（MIMO）SSM代替S4中的许多独立SISO SSM。其次，S5使用高效且广泛实现的并行扫描。这消除了S4使用的卷积和频域方法的需要。

而S5模型之所以能够实现并行扫描，是因为之前有研究发现，HiPPO矩阵的对角近似可以实现相当的表现。于是，HiPPO的A矩阵可以以对角化矩阵的方式进行处理。

文章提到"与DSS层和S4D层的最新工作一致，我们发现HiPPO矩阵的对角近似可以实现相当的表现"。

简单看看DSS和S4D这两个工作

DSS：Diagonal State Spaces

文章：《Diagonal State Spaces are as Effective as Structured State Spaces》

摘要：建模序列数据中的长期依赖关系是实现许多模态（如文本、视觉、音频和视频）中人类水平性能的基本步骤。虽然基于注意力的模型是建模短距离交互的流行且有效的选择，但它们在需要长期推理的任务上的表现一直不尽如人意。在一个令人兴奋的结果中，Gu等人[GGR22]提出了结构化状态空间（S4）架构，在多个领域的长期任务上取得了比最先进的模型更大的提升。S4的核心主张是通过对角线加低秩结构参数化状态矩阵，允许高效计算。在这项工作中，我们展示了即使没有低秩修正，也可以匹配S4的性能，因此假设状态矩阵是对角的。我们的对角状态空间（DSS）模型在长期范围竞技场任务上与S4的性能相匹配，在语音命令数据集上的语音分类上也是如此，同时在概念上更简单，更直接地实现。

引言

Transformer架构[VSP`17]在机器学习领域的许多领域都取得了成功。通过去噪目标预训练的Transformer，在自然语言处理中已成为标准，展现出令人印象深刻的语言和世界知识[BMR`20, BMH`21]。这种配方也导致了视觉[RPG`21, RKH`21]和语音[DXX18, BZMA20]领域的显著发展。

Transformer的上下文化组件是多头注意力层，对于长度为L的输入，其复杂度为Ω(L^2)。当模型需要捕获长输入中各个部分的长距离交互时，这变得令人望而却步。为了缓解这个问题，提出了几种计算和内存需求较低的Transformer变体[QML20, KKL20, RSVG21, BPC20, GB20, KVPF20, WLK20, CLD21, ZS21, GDG21]（参见[TDBM20]的调查）。尽管如此，所有这些模型在正式评估和量化模型执行长期推理能力的基准测试中（如长期范围竞技场[_TDA21]和SCROLLS[SSI22]）都报告了不佳的性能。

在一个最近的突破性结果中，Gu等人[GGR22]提出了S4，这是一个序列到序列模型，它使用线性状态空间而不是注意力进行上下文化。它在需要长期推理的领域（如文本、图像和音频）的任务上表现出色。例如，在长期范围竞技场上，它比表现最好的Transformer变体提高了19个准确点。其卓越的能力不仅限于文本和图像，还延续到时间序列预测、语音识别和音频生成[GGDR22]。

尽管S4取得了成就，但其设计复杂，以HiPPO理论为中心，这是长期建模的数学框架[VKE19, GDE20, GJG21]。[GGR22]表明，与使用特定的HiPPO矩阵初始化状态空间参数相比，具有各种替代初始化的状态空间模型表现不佳。为了利用这个矩阵，他们使用对角加低秩（DLPR）结构参数化学习的状态空间，并因此需要采用几个简化步骤和线性代数技术来有效地计算状态空间输出，这使得S4难以理解、实现和分析。

在这项工作中，我们展示了可以匹配S4的性能，同时使用一个更简单的、完全对角的参数化状态空间。我们确认随机对角状态空间效果较差，但我们观察到确实存在有效的对角状态矩阵：简单地移除DPLR HiPPO矩阵的低秩部分仍然可以保持其性能。利用这个想法，我们提出的对角状态空间（DSS）模型强制状态矩阵为对角线，使其在制定、实现和分析上显著更简单，同时被证明与一般状态空间一样具有表达能力。与S4相比，DSS不假设任何专业背景，除了基本的线性代数之外，并且可以用几行代码实现。我们的实现适合在一页纸上，提供在§A.5（图6）。

我们在长期范围竞技场（LRA）上评估了DSS的性能，这是一系列序列级分类任务，输入长度多样（1K-16K），需要在广泛的模态上进行相似性、结构和视觉空间推理，如文本、自然/合成图像和数学表达式。尽管其简单性，DSS在6个LRA任务上的平均准确率达到81.88，与S4的最新性能（80.21）相当。此外，DSS保持了比最佳Transformer变体高出20个百分点的舒适领先优势（81.88 vs 61.41）。

在音频领域，我们在原始语音分类上评估了DSS的性能。在语音命令数据集[War18]上，该数据集由长度为16K的原始音频样本组成，我们再次发现DSS的性能与S4相当（98.2 vs 98.1）。

总结来说，我们的结果表明，DSS是建模文本、图像和音频等模态中长期交互的一种简单有效的方法。我们相信，DSS的有效性、效率和透明度可以显著促进状态空间模型在其基于注意力的同行中的采用。

方法

Diagonal State Spaces

文章首先提出了一个关键命题，即在一定技术假设下，对角状态空间与一般状态空间具有相同的表达能力。这个命题基于给定的状态空间核K，其中A是对角可解的，具有非零特征值λ1, ..., λN，并且对于所有的i，e^(LλiΔ) ≠ 1。基于这个命题，文章展示了如何将一般状态空间参数化为对角状态空间，并且证明了对于任何给定的状态空间，都存在一个对角状态空间能够计算出相同的核。

Diagonal State Space (DSS) Layer

接下来，文章描述了DSS层的结构。DSS层接收长度为L的H维向量序列作为输入，并产生同样维度的输出。该层的参数包括Λre, Λim, ∆log, 和W。对于每个坐标h，使用上述方法之一计算状态空间核Kh。然后，使用离散卷积公式计算输出yh。

在文章的第3.1节中，介绍了两种方法来模拟对角状态空间，即DSSEXP和DSSSOFTMAX。这两种方法都是为了解决在训练过程中可能出现的数值不稳定性问题，同时保持模型的表达能力。下面分别解释这两种方法：

DSSEXP（Diagonal State Space Exponential）

DSSEXP方法基于命题1(a)来模拟状态空间，但它限制了Λ矩阵元素的实部必须为负。这样可以避免在计算过程中出现数值爆炸的问题，因为当实部为正时，指数函数的值可能会变得非常大。DSSEXP的参数包括Λ的实部和虚部（Λre, Λim），权重向量（rw），以及对数时间间隔（∆log）。Λ是通过以下方式计算的：

然后，状态空间核（SSM kernel）K通过命题1(a)来计算。

DSSSOFTMAX（Diagonal State Space Softmax）

DSSSOFTMAX方法则是通过另一种方式来限制elementwise-exp§的元素，即通过对每一行进行softmax归一化。这种方法不需要限制Λ的实部，因此理论上可以表达更广泛的动态范围。DSSSOFTMAX的参数包括Λ的实部和虚部（Λre, Λim），权重向量（w），以及对数时间间隔（∆log）。Λ和∆的计算方式与DSSEXP相同。然后，状态空间核K通过命题1(b)来计算。

在DSSSOFTMAX中，为了避免在复数域上的softmax操作导致的数值不稳定问题，文章提出了一种修正的softmax函数（softmaxϵ），该函数通过添加一个小的正数ϵ来避免除以零的情况，并确保softmax操作在整个复数域上都是有定义且可微的。

Initialization of DSS Layer

文章强调了DSS层参数初始化的重要性，并提出了一种特定的初始化方法，称为Skew-Hippo初始化。这种初始化方法涉及对Λ的实部和虚部以及∆log的特定设置，以确保模型能够获得最佳性能。

Λ在CN中被初始化为HiPPO矩阵的正常部分的特征值[GGR22]。具体来说，Λre, Λim被初始化，以便得到的Λ是以下2N × 2N矩阵的N个特征值的向量，这些特征值具有正的虚部：

States of DSS via the Recurrent View

最后，文章讨论了如何在DSS模型中通过递归视角计算状态。这包括在自回归解码等情况下，如何有效地计算状态空间模型的状态。文章提供了DSSEXP和DSSSOFTMAX两种方法的详细计算步骤，以及如何根据状态矩阵的特征值来决定信息的传递和遗忘。

S4D

文章：On the Parameterization and Initialization of Diagonal State Space Models

文章链接：https://arxiv.org/pdf/2206.11893

摘要：状态空间模型（SSM）最近被证明作为一种深度学习层非常有效，是序列模型如RNN、CNN或Transformer的一个有希望的替代方案。展示这一潜力的首个版本是S4模型，它通过使用一种称为HiPPO矩阵的预定状态矩阵，在涉及长期依赖的任务上表现出色。虽然这在数学上有可解释的机制来模拟长期依赖，但它引入了一个自定义的表示和算法，可能难以实现。另一方面，一个名为DSS的S4的变体表明，将状态矩阵限制为完全对角线仍然可以在使用基于近似S4矩阵的特定初始化时保持原始模型的性能。这项工作旨在系统地理解如何参数化和初始化这些对角线状态空间模型。虽然从经典结果来看，几乎所有的SSM都有等价的对角线形式，但我们展示初始化对性能至关重要。我们通过数学上解释为什么DSS有效，展示了S4矩阵的对角线限制在无限状态维度的极限下惊人地恢复了与S4相同的内核。我们还系统地描述了在参数化和计算对角线SSM时的各种设计选择，并进行了受控的实证研究，对这些选择进行了消融效应分析。我们最终的模型S4D是一个简单的S4的对角线版本，其内核计算只需要2行代码，并且在几乎所有设置中与S4的性能相当，在图像、音频和医疗时间序列领域取得了最先进的结果，在长距离竞技场基准测试中平均达到了85%。

重点方法

文章的重点方法是对角状态空间模型（Diagonal State Space Models，简称DSS），特别是提出了一个新的模型S4D。以下是文章中提到的关于S4D的关键点：

简化的模型结构：S4D是一个对角线版本的S4模型，它通过限制状态矩阵为对角线形式来简化模型结构，这样做可以显著降低模型的复杂性，使其更易于实现和理解。
性能保持：尽管S4D简化了模型结构，但它仍然能够保持与原始S4模型相似的性能。
数学分析：文章提供了对DSS初始化的新的数学分析，解释了为什么对角线近似能够在状态维度趋向于无穷大时产生与S4相同的动态效果。
参数化和计算：文章系统地描述了在参数化和计算对角线SSM时的不同设计选择，并提出了S4D方法，该方法基于Vandermonde矩阵乘法，比DSS更简单、更高效。
实证研究：通过受控的实证研究，文章对这些设计选择进行了消融效应分析，验证了S4D方法的有效性。
广泛的应用：S4D在多个领域（图像、音频和医疗时间序列）和长距离竞技场基准测试中表现出色，证明了其广泛的适用性和强大的性能。
代码实现：S4D的内核计算只需要两行代码，这表明了其实现的简洁性。

总的来说，S4D方法的核心在于通过简化状态空间模型的状态矩阵为对角线形式，同时保持模型的性能，从而提供了一个既简单又有效的深度学习序列模型。

初始化对角状态矩阵

关键问题仍然是：哪些对角状态矩阵A实际上是有效的？我们首先评论对角SSM的局限性，然后提供三个在实证中表现良好的S4D实例。

对角SSM的表达性和局限性 ：我们首先对[11]中提到的对角SSM的表达性提供一个简化的观点。首先，众所周知，几乎所有矩阵都可以在复平面上对角化。因此，使用复值矩阵对于使用对角SSM至关重要。