【论文笔记】Mamba：挑战Transformer地位的新架构

Mamba

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

摘要

背景

基础模型，它现在驱动着深度学习中大多数令人振奋的应用，几乎都是基于Transformer架构及其核心注意力模块。

存在的问题

许多subquadratic-time架构如线性注意力、门控卷积和循环模型以及structured state space models (SSMs)已经被开发，以解决Transformer在长序列上的计算效率低下问题，但它们在语言等重要模态上的表现不如注意力。

本文的做法

本文发现这些模型的一个关键弱点是它们不能执行基于内容的推理，对此，本文做了如下工作：

• 将SSM参数作为输入的函数，以离散的模态解决了它们的弱点，允许模型根据当前token沿序列长度维度选择性地传播或遗忘信息。
• 尽管这种变化阻碍了高效卷积的使用，但在循环模式下设计了一种硬件感知的并行算法。将这些选择性SSM集成到一个简化的端到端神经网络架构中，它不需要注意力，甚至不需要MLP模块。

Mamba具有快速推理能力(比Transformers高5倍吞吐量)和序列长度的线性伸缩能力，在实际数据上的性能提高到百万序列长度。

实验结果

作为一个通用的序列模型主干，Mamba在语言、音频和基因组学等多个模态上取得了最先进的性能。在语言建模方面，Mamba - 3B模型在预训练和下游评估方面都优于相同大小的Transformers，并且匹配了两倍大小的Transformers。

文章内容

Transformer的缺点

尽管当前的Foundatio model主要基于是Transformer，但它存在一些不足：

• 无法对有限窗口之外的任何东西进行建模
• 窗口长度的二次复杂度

尽管有很多关于注意力变体的研究，但这些往往都会以牺牲精度为代价，且这些研究在跨域问题上并不有效。

Structured state space sequence models (SSMs)介绍

SSMs可以理解为循环神经网络( RNNs )和卷积神经网络( ***Ns )的组合，在序列长度上具有线性或近似线性的复杂度。它们在特定的数据模态上效果不错，但是它们对文本等离散、信息密集的数据建模效果较差。

本文的工作

提出了一类新的选择性状态空间模型，它改进了先前的工作，以实现Transformer的建模能力，同时在序列长度上是线性复杂度。具体来说：

• Selection Mechanism ：之前的模型根据输入数据进行高效选择的能力不强，对此，本文基于输入对SSM参数进行参数化，设计了一种简单的选择机制。这使得模型可以过滤掉不相关的信息，无限期地记住相关信息。
• Hardware-aware Algorithm：所有的先验SSMs模型必须是时间和输入不变的，才能具有计算效率，为了克服这个问题，文章通过扫描而不是卷积来循环计算模型，但不实现expanded state，以避免在GPU内存层次结构的不同级别之间进行IO访问。
• Architecture：将之前的SSM架构的设计与Transformers的MLP模块合并为一个模块来简化先验深度序列模型架构，从而得到一个包含选择性状态空间的简单同构架构设计( Mamba )。

模型的优点：

• High quality：选择性在诸如语言和基因组学等密集模态上带来了强大的性能。
• Fast training and inference：在训练过程中，计算和记忆按序列长度线性缩放，并且在推理过程中自回归地展开模型只需要每一步恒定的时间，因为它不需要先前元素的缓存。
• Long context：质量和效率一起提高了在序列长度为1M的真实数据上的性能。

模型介绍

State Space Models

离散化： Structured state space sequence models (S4)将一个一维函数或者序列x(t)->y(t)通过一种隐性的隐状态h(t)进行映射，它最开始是针对连续数据的。而文本、基因序列都是离散的，因此需要将数据进行离散化，具体如下：

离散过程为：

关于以上公式的解释可以看这篇文章：通俗易懂解释Mamba

计算： 模型有两种计算模式：

• 线性递归：用于高效的自回归推理(其中输入在某一时刻被看作一个时间步)
• 全局卷积：用于高效的可并行化训练(其中整个输入序列被提前看到)

为什么S4有这两种模式，可以看这篇文章：通俗易懂解释Mamba

线性时间不变性：$(A,B,C,\Delta )$和$(\overline{A},\overline{B})$在所有时间步都是固定的，它与递归和卷积有着深刻的联系。

本文认为LTI模型在建模某些类型的数据时具有根本的局限性，因此本文将会消除LTI约束，同时克服效率瓶颈。

结构和维度： 结构化SSM的命名原因是由于高效地计算它们还需要在A矩阵上施加结构，最常用的结构形式是对角线。

A的维度为$N\times N$，B的维度为$N\times 1$，C的维度为$N\times 1$，其中N为数的个数。输入的一个x的维度为$B\times L\times D$，其中D为通道数，L为序列长度，B为batch数，在序列长度上计算需要O ( BLDN )的时间和内存。

一般的State Space Models：状态空间模型简单地表示了任何具有潜在状态的循环过程的概念。

它在不同的学科中被用来指称许多不同的概念，包括马尔可夫决策过程，动态因果建模，隐马尔可夫模型和线性动力系统。

SSM体系结构： SSMs是独立的序列变换，可以结合到端到端的神经网络架构中。一些知名的SSM架构为：

• Linear attention：一个包含递归的自注意力的近似，可以看作是一个退化的线性SSM。
• H3：可以被看作是由两个门控连接夹在一个SSM中的体系结构。还在主SSM层之前插入了一个标准的局部卷积，是一个shift-SSM。
• Hyena：使用与H3相同的架构，但将S4层替换为MLP参数化的全局卷积
• Ret***：在架构中添加一个额外的门，并使用更简单的SSM，允许一个可选择的可并行化计算路径，使用多头注意力( MHA )的变体代替卷积。
• RWKV：基于另一种线性注意力近似( attention-free Transformer )设计的用于语言建模的RNN。其主要的" WKV "机制涉及LTI重现，可视为两个SSM的比率。

Selective State Space Models

本文的灵感来源

本文认为序列建模的一个基本问题是将上下文压缩到一个较小的状态。Transformer没有压缩上下文，导致推理过程缓慢；递归模型具有有限的状态，实现了常数时间推理和线性时间训练，但是它们的有效性受限于这种状态对上下文的压缩程度。

本文重点关注合成任务的两个运行例子：

• Selective Copying task：通过改变token的位置来记忆，从而实现修改Copying task。它要求内容感知推理能够记忆相关的token并过滤掉不相关的token
• Induction Heads task：假设可以解释LLMs的大部分上下文学习能力。它需要上下文感知的推理来知道何时在适当的上下文中产生正确的输出。

这些任务揭示了LTI模型的失败模式，具体如下：

• recurrent view：constant dynamics不能让它们从上下文中选择正确的信息，也不能以输入依赖的方式影响沿序列传递的隐藏状态。
• convolutional view：由于缺乏内容意识，在选择性复制任务中存在困难。更具体地说，输入到输出之间的间距是变化的，不能用静态卷积核来建模。

序列模型的效率和有效性的tradeoff是由它们压缩状态的程度来表征的：

• 高效的模型必须有一个小的状态
• 有效的模型必须有一个包含来自上下文的所有必要信息的状态

Improving SSMs with Selection

$s_B$，$s_C$，$s_{\Delta }$都是线性层，以x为输入，如下图。

S6的参数有一个长度维度L，这意味着模型已经从时不变（输出跟输入的时间无关）变为时变。这失去了与卷积的等价性，影响了它的效率。
$s_{\Delta }$和${\tau }_{\Delta }$的选择与RNN门控机制有关.

Efficient Implementation of Selective SSMs

隐藏状态维度较大的模型应该更高效，但速度较慢，因此本文希望在不付出速度和内存成本的情况下，最大化隐藏状态维度。

为了高效地运行，本文将SSM参数$(A,B,C,\Delta )$直接从慢速HBM加载到快速SRAM中，在SRAM中进行离散化和递归，然后将最终输出的尺寸( B，L，D)写回HBM。
本文还避免保存反向传播所必需的中间状态。本文应用了经典的重计算技术来减少内存需求：当输入从HBM加载到SRAM时，中间状态不存储，而是在后向通道中重新计算。

A Simplified SSM Architecture

选择性SSM是独立的序列变换，可以灵活地集成到神经网络中。

激活函数使用Silu或者Swish激活函数，同时使用了Layer Norm。

Selection Mechanisms的解释

可变间距： 选择性允许过滤掉感兴趣的输入之间可能出现的不相关的噪声标记。例如’ um '等语言填充词的存在。

过滤上下文： 许多序列模型并没有随着时间的延长而得到改善，按道理应该上下文越多性能越好。一种解释是，许多序列模型在必要时不能有效地忽略不相关的上下文；一个直观的例子是全局卷积(和一般的LTI模型)。选择性模型可以在任意时刻简单地重置它们的状态以去除无关的历史，因此它们的性能原则上随上下文长度单调地提高。

边界重置：在多个独立序列拼接在一起的场景中，Transformer可以通过实例化特定的注意力掩码来保持它们之间的分离，而LTI模型则会在序列之间引入信息。选择性SSM可以在边界处重置状态。

$\Delta$的解释：一般而言，$\Delta$控制着关注或忽略多少当前输入xt的平衡。它可以泛化到RNN的门控机制，一个大的$\Delta$重置状态h并聚焦在当前输入x上，而一个小的$\Delta$保持状态并忽略当前输入。

A的解释：A参数也可能是选择性的，但它最终只通过其和$\Delta$的交互作用对模型产生影响。

B和C的解释：在SSM中，修改B和C使其具有选择性，可以更细粒度地控制输入xt进入状态ht或状态进入输出yt。这可以解释为允许模型分别基于内容(输入)和上下文(隐藏状态)来调节递归动态。

结论

Mamba的作用：

• 可以对token或者特征进行有效选择，过滤无用的信息，保留感兴趣的信息
• Mamba在模型的效率和性能之间得到了权衡，因为它能够很好地过滤无用的信息，最大化隐藏状态维度
• Mamba使用扫描代替卷积，加快了训练速度

论文信息：

• 论文里面的$\Delta$相当于RNN中的门控进制
• Mamba是输入依赖的，里面的A，B，C都是通过输入数据得到的，因此实现了选择性。