Mamba: Linear-Time Sequence Modeling with Selective State Spaces(论文笔记)

What can I say?

2024年我还能说什么？

Mamba out!

曼巴出来了！

原文链接：

[2312.00752] Mamba: Linear-Time Sequence Modeling with Selective State Spaces (arxiv.org)

原文笔记：

What：

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

SSM——>S4——>S6——>Mamba

SSM:（不能处理离散序列，不能有选择地处理信息A,B,C,D的值都是固定不变的）

一般SSMs包括以下组成

• 映射输入序列x(t)，比如在迷宫中向左和向下移动
• 到潜在状态表示h(t)，比如距离出口距离和 x/y 坐标
• 并导出预测输出序列y(t)，比如再次向左移动以更快到达出口

然而，它不使用离散序列(如向左移动一次)，而是将连续序列作为输入并预测输出序列

这里的h'(t)不是h(t)导数，是新的状态。SSM 假设系统(例如在 3D 空间中移动的物体)可以通过两个方程从其在时间 时的状态进行预测「 当然，其实上面第一个方程表示成这样可能更好：，不然容易引发歧义 」

RNN本身就可以理解乘一个SSM

参数A春初这之前所有历史信息的浓缩精华（可以通过一系列系数组成的矩阵表示之），以基于A更新下一个时刻的空间状态hidden state

总之，通过求解这些方程，可以根据观察到的数据：输入序列和先前状态，去预测系统的未来状态

建立两个方程的统一视角：

S4：（离散化SSM、循环/卷积表示，基于HiPPO处理长序列，不能有选择的处理信息，A,B,C,D的值都是固定不变的）详见参考1

S6：SSM+Selection（可以处理离散数据，有选择地处理信息，训练效率降低）

在Mamaba中，作者让矩阵、矩阵、成为输入的函数(即可学习或可训练的)，让模型能够根据输入内容自适应地调整其行为

Mamba: （可以处理离散数据+有选择地处理信息+硬件感知算法+更简单的SSM架构）

Why：

1、Transformer在处理长序列时存在着计算效率低下的问题，无法对有限窗口之外的任何信息进行建模，以及相对于窗口长度的二次缩放。

2、用于解决问题一的方法如线性注意力，门控卷积和循环模型以及结构化状态空间模型在表现上（在语言等重要模态上）不如注意力，他们无法执行基于内容的推理。(效果差的原因：不能像attention机制那样全局感知；缺乏对离散数据的的建模能力（SSM））

3、传统的SSM模型在建模离散和信息密集数据（例如文本）方面效果较差。S4能离散建模泵选择处理，S6能离散建模和选择处理又遇到了新的问题：无法将高效的卷积应用于训练以通过并行的方式提高训练效率，之后作者也提出了优化方法(idea3))

4、鉴于1，2，3，我们要寻求一种在计算效率和模型效果之间的平衡，找到同时在这两个方面表现都很优秀的模型

总之，序列模型的效率与效果的权衡点在于它们对状态的压缩程度：

• 高效的模型必须有一个小的状态(比如RNN或S4)
• 而有效的模型必须有一个包含来自上下文的所有必要信息的状态(比如transformer)

而mamba为了兼顾效率和效果，选择性的关注必须关注的、过滤掉可以忽略的

Challenge：

1、基于内容的长时记忆以及推理(应该包含上下文所有的必要信息)，能够以依赖于输入的方式有效地选择数据（即专注于或忽略特定输入）

2、较短的处理时间（训练时间以及测试时间）

Idea：Mamba: （可以处理离散数据+有选择地处理信息+硬件感知算法+更简单的SSM架构）

1、通过离散化操作解决了传统SSM的弱点（准确率上）

2、为了解决Challenge1，作者提出：让SSM参数成为输入的函数，允许模型根据当前token有选择地沿序列长度维度传播或忘记信息。（准确率上）

3、idea1和2 阻碍了将高效卷积应用于训练以通过并行的方式提高训练效率，为此作者在循环模式下设计了一个硬件感知的并行算法。（性能上）

4、作者将所涉及的ssm集成到简化的端到端神经网络架构中，而无需注意力模块甚至MLP块

（模型原创性，泛化性）

张量具体怎么对齐的等我研究明白再更新

Model：

原文翻译：

（红字：出现的问题，黄字：想法重点，绿字：效果，篮字：看参考）

Abstract

现在为深度学习中大多数令人兴奋的应用程序提供动力的基础模型，几乎普遍基于Transformer架构及其核心的注意力模块。许多subquadratic-time(次二次时间)架构，如线性注意力、门控卷积和循环模型，以及结构化状态空间模型(ssm)已经被开发出来，以解决Transformer在长序列上的计算效率低下问题，但它们在语言等重要模态上的表现不如注意力。我们发现这些模型的一个关键弱点是它们无法执行基于内容的推理，并进行了一些改进。首先，简单地让SSM参数成为输入的函数，通过离散模态解决了它们的弱点，允许模型根据当前令牌有选择地沿序列长度维度传播或忘记信息。其次，尽管这种变化阻碍了高效卷积的使用，但我们在循环模式下设计了一个硬件感知的并行算法。我们将这些选择性ssm集成到简化的端到端神经网络架构中，而无需注意力模块甚至MLP块。Mamba具有快速推理(比Transformers高5倍的吞吐量)和序列长度的线性缩放，并且其性能在真实数据高达百万长度的序列上有所提高。作为一个通用的序列模型主干，Mamba在语言、音频和基因组学等多种模式上实现了最先进的性能。在语言建模方面，我们的Mamba-3B模型在预训练和下游评估方面都优于同等大小的变形金刚，并且与变形金刚的两倍大小相匹配。

1 Introduction

基础模型 (FM) 或在大量数据上预训练的大型模型，然后适应下游任务，已成为现代机器学习的有效范式。这些FM的主干通常是序列模型，对来自语言、图像、语音、音频、时间序列和基因组学等多个领域的任意输入序列进行操作(Brown等人，2020;Dosovitskiy等人，2020;Ismail Fawaz等人，2019;Oord等人，2016;Poli等人2023;Sutskever, Vinyals和Quoc V Le 2014)。虽然这个概念与模型架构的特定选择无关，但现代 FM 主要基于一种类型的序列模型：Transformer (Vaswani et al. 2017) 及其核心注意力层 (Bahdanau, Cho, and Bengio 2015) 自注意力的功效归因于它能够在上下文窗口中密集地路由信息，使其能够对复杂的数据进行建模。然而，此属性带来了根本的缺点：无法对有限窗口之外的任何信息进行建模，以及相对于窗口长度的二次缩放。大量研究似乎在更有效的注意力变体上来克服这些缺点（Tay、Dehghani、Bahri 等 2022），但通常以牺牲使其有效的重要属性为代价。然而，这些变体都没有被证明在跨领域的规模上在经验上是有效的。

最近，结构化状态空间模型 (SSM) (Gu, Goel and Ré 2022; Gu, Johnson, Goel, et al. 2021) 已成为序列建模的一种有前途的架构。这些模型可以解释为循环神经网络 (RNN) 和卷积神经网络 (CNN) 的组合，灵感来自经典状态空间模型 (Kalman 1960)。此类模型可以非常有效地计算为递归或卷积，序列长度具有线性或接近线性缩放。此外，它们具有在某些数据模式下对远程依赖关系建模的原则机制(Gu, Dao, et al. 2020)，并主导了远程竞技场等基准(Tay, Dehghani, Abnar等人，2021年)。许多风格的SSMs(Gu，Goel和Ré 2022；Gu，Gupta等人。2022；Gupta，Gu和Berant 2022；Y。李等人。 2023；马等人。 2023；奥维托等人。 2023；史密斯、华灵顿和林德曼 2023）在涉及音频和视觉等连续信号数据的领域取得了成功（Goel 等人 2022；Nguyen、Goel 等人，2022；Saon、Gupta 和 Cui 2023）。然而，它们在建模离散和信息密集数据（例如文本）方面效果较差。

我们提出了一类新的选择性状态空间模型，它在几个轴上改进了先前的工作，以实现 Transformer 的建模能力，同时在序列长度上线性缩放。

选择机制。首先，我们确定了先前模型的一个关键限制：能够以依赖于输入的方式有效地选择数据（即专注于或忽略特定输入）。基于基于重要合成任务(选择性复制和诱导头等)的直觉，我们设计了一种简单的选择机制:根据输入参数化SSM参数。这允许模型过滤掉不相关的信息并无限地记住相关信息。（长时记忆并且能正确选择）

硬件感知算法。这种简单的改变对模型的计算提出了技术挑战;事实上，为了提高计算效率，所有先前的ssm模型都必须是时间和输入不变的。我们使用硬件感知算法克服了这一点，该算法使用扫描而不是卷积来循环计算模型，但不实现扩展状态，以避免在GPU内存层次结构的不同级别之间进行IO访问。由此产生的实现在理论上(序列长度线性缩放，与所有基于卷积的ssm的伪线性相比)和现代硬件上(在A100 gpu上快3倍)都比以前的方法快。

体系结构。我们通过将先前的SSM架构(Dao, Fu, Saab, et al. 2023)的设计与transformer的MLP块合并为一个块来简化先前的深度序列模型架构，从而得到一个包含选择性状态空间的简单且同质的架构设计(Mamba)。

选择性ssm，以及扩展的Mamba架构，是具有关键属性的完全循环的模型，这些属性使它们适合作为在序列上操作的一般基础模型的特征提取网络。这些属性有(i)高质量:选择性使语言和基因组学等密集模式的表现强劲。(ii)快速训练和推理:在训练期间，计算和内存在序列长度上呈线性扩展，并且在推理期间自回归展开模型只需要每步恒定的时间，因为它不需要先前元素的缓存(因为只需要先前状态)。(iii)长上下文:质量和效率共同提高了实际数据的性能，最长可达1M序列长度。

在以下几种类型的模式和设置下，我们通过经验验证了曼巴在预训练质量和特定领域任务性能方面作为一般序列FM骨干的潜力:

• 合成。在复制和诱导头等重要的合成任务中，作为大型语言模型的关键，Mamba不仅很容易解决它们，而且可以无限长地推断解决方案(>1M令牌)。
• 音频和基因组学。Mamba 在预训练质量和下游指标（例如，在具有挑战性的语音生成数据集上将 FID 减少一半以上）建模音频波形和 DNA 序列的先前最先进模型（例如，在具有挑战性的语音生成数据集上减少 FID 一半以上）。在这两种情况下，它的性能都随着更长的上下文而提高，最高可达百万长序列。
• 语言建模。Mamba 是第一个真正实现 Transformer 质量性能的线性时间序列模型，无论是在预训练困惑度还是下游评估中。通过高达 1B 参数的缩放定律，我们表明 Mamba 超过了大量基线的性能，包括基于 LlaMa 的非常强大的现代 Transformer 训练配方（Touvron 等人 2023）。与大小相似的 Transformer 相比，我们的 Mamba 语言模型具有 5 倍的生成吞吐量，而 Mamba-3B 的质量与其大小的两倍匹配（例如，在常识推理方面，与Pythia-3B相比平均高出4分，甚至超过了Pythia-7B）。

模型代码和预训练的检查点是开源的在https://github.com/state-spaces/mamba

图1:(概述)结构化的SSMs独立的将输入x的每个channel通过高维的潜在状态h映射到输出y，先前的 SSM 通过巧妙的替代计算路径避免通过需要时间不变性来具体化这种大有效状态（DN，乘以批量大小 B 和序列长度 L）：(Δ, A, B, C) 参数在时间上是恒定的。我们的选择机制添加了依赖于输入的动态，这也需要仔细的硬件感知算法，以在更高效的 GPU 内存层次结构级别上仅具体化扩展状态。

2 状态空间模型

结构化状态空间模型 (S4) 是最近一类深度学习序列模型，与 RNN 和 CNN 以及经典状态空间模型广泛相关。它们受到一个特定的连续系统 (1) 的启发，该系统通过隐式潜在状态 ℎ(t) ∈ ℝ^N将一个一维函数或序列x(t) ∈ ℝ 映射到y(t)∈ℝ。

具体来说，S4 模型由四个参数（Δ、A、B、C）定义，这些参数分两个阶段定义序列到序列的转换。

离散化。（这段看下原文吧）第一阶段通过固定公式A拔=