【LLM】Mamba和State Space Models详解

Mamba和State Space Models详解1. Transformer的问题1.1 Transformers的核心组件1.2 一份带有训练的Blessing……1.3 还有带推理的Curse1.4 RNN是解决方案吗2. 状态空间模型SSM2.1 什么是状态空间2.2 什么是状态空间模型2.3 从连续信号到离散信号2.4 循环表示2.5 卷积表示2.7 三种表示方式的对比2.8 矩阵A的重要性3. Mamba——一种选择性SSM3.1 它试图解决什么问题3.2 有选择地保留信息3.3 扫描操作3.4 硬件感知算法3.5 Mamba Block4. 结论5. 资源Transformer架构一直是大型语言模型LLM取得成功的重要组成部分。从开源模型如Mistral到闭源模型如ChatGPT如今几乎所有正在使用的大型语言模型都采用了这一架构。为进一步提升大语言模型的性能人们开发了新的架构这些架构甚至可能超越Transformer架构。其中一种方法便是Mamba一种状态空间模型。Mamba 是在论文《Mamba具有选择性状态空间的线性时间序列建模》中提出的。您可以在其仓库中找到官方实现和模型检查点。在本文中我将从语言建模的角度介绍状态空间模型这一领域并逐一探讨相关概念以帮助大家建立对该领域的直观理解。随后我们将深入探讨 Mamba 如何可能对 Transformer 架构构成挑战。作为一份可视化指南将助您更好地理解 Mamba 和状态空间模型1. Transformer的问题为了说明Mamba为何是一种如此有趣的架构我们先来简要回顾一下Transformer并探讨其一个缺点。Transformer 将任何文本输入视为由token组成的sequence。Transformer的一个主要优势在于无论它接收到何种输入都能回溯序列中任意先前的标记以获取其表示。1.1 Transformers的核心组件请记住Transformer 由两种结构组成一组用于表示文本的编码器块以及一组用于生成文本的解码器块。这两种结构结合起来可用于多项任务包括翻译。我们可以采用这种结构仅使用解码器来构建生成模型。这种基于Transformer的模型——生成式预训练TransformerGPT利用解码器模块来完成对输入文本的生成。让我们来看看它是如何运作的1.2 一份带有训练的Blessing……一个解码器块由两个主要组成部分构成掩码自注意力机制后接前馈神经网络。自注意力机制是这些模型能够取得如此出色效果的一个重要原因。它使得模型能够以快速的训练速度对整个序列进行无压缩的全局视角分析。那么它是如何工作的呢它会创建一个矩阵将每个token与之前出现过的每一个token进行比较。矩阵中的权重由token对之间的相关性决定。在训练过程中这个矩阵是一次性生成的。在计算“name”与“is”之间的注意力之前无需先计算“My”与“name”之间的注意力。这使得并行化成为可能从而极大地加快了训练速度1.3 还有带推理的Curse然而存在一个缺陷在生成下一个token时即使我们已经生成了一些tokens仍需重新计算entire sequence的注意力。为长度为L的序列生成tokens大约需要L²次计算如果序列长度增加这可能会带来高昂的成本。这种需要重新计算整个序列的做法是Transformer架构的一个主要瓶颈。让我们来看看一种“经典”技术——循环神经网络——是如何解决这一推理速度慢的问题的。1.4 RNN是解决方案吗循环神经网络RNN是一种基于序列的网络。在序列的每个时间步它都会接收两个输入时间步t的输入以及前一时间步t-1的隐藏状态从而生成下一个隐藏状态并预测输出。RNN具有一种循环机制能够将信息从上一个时间步传递到下一个时间步。我们可以对这一过程进行“展开”可视化以使其更加清晰明了。在生成输出时RNN只需考虑前一个隐藏状态和当前输入。这避免了像Transformer那样重新计算所有先前的隐藏状态。换句话说RNN的推理速度非常快因为它与序列长度呈线性 scaling理论上它的上下文长度甚至可以是无限的。为了说明这一点让我们将RNN应用于我们之前使用过的输入文本。每个隐藏状态都是所有先前隐藏状态的聚合通常是一种压缩后的视图。然而这里存在一个问题……请注意当生成名字“Maarten”时最后一个隐藏状态已不再包含关于单词“Hello”的信息。由于RNN每次只考虑一个先前状态因此随着时间推移它们往往会逐渐遗忘信息。RNN这种序列化的特性也带来了另一个问题训练无法并行进行因为必须按顺序逐个步骤地完成。与Transformer相比RNN的问题恰恰相反它的推理速度极快但却无法并行化。我们能否找到一种架构既能像Transformer那样实现训练的并行化又能在推理时保持与序列长度呈线性 scaling 的性能当然可以这正是Mamba所提出的方案。不过在深入探讨其架构之前让我们先来了解一下状态空间模型的世界吧。2. 状态空间模型SSM状态空间模型SSM与Transformer和RNN一样能够处理信息序列这些序列既可以是文本也可以是信号。在本节中我们将介绍状态空间模型的基本概念以及它们如何应用于文本数据。2.1 什么是状态空间状态空间包含完整描述一个系统所需的最少变量。它是一种通过定义系统的可能状态来对问题进行数学建模的方法。让我们来简化一下这个概念。想象一下我们正在迷宫中导航。“状态空间”就是一张描绘所有可能位置状态的地图。地图上的每个点都代表迷宫中的一个独特位置并附带具体信息比如你距离出口有多远。“状态空间表示”是对这张地图的一种简化描述。它展示了你当前所处的位置当前状态你下一步可能前往的地方未来可能的状态以及哪些变化会带你进入下一个状态向右或向左移动。尽管状态空间模型使用方程和矩阵来追踪这种行为但这其实只是一种用来追踪你当前所处位置、可前往的地点以及如何抵达这些地点的方法。在我们的示例中描述状态的变量——即X和Y坐标以及到出口的距离——可以被表示为“state vectors”。听起来很熟悉吧这是因为语言模型中的嵌入或向量也常被用来描述输入序列的“state”。例如你当前位置的向量状态向量可能看起来有点像这样在神经网络中系统的“state”通常指其隐藏状态而在大型语言模型的背景下这正是生成新token时最重要的方面之一。2.2 什么是状态空间模型状态空间模型是用来描述这些状态表示并根据某些输入预测其下一状态的模型。传统上在时刻tSSMs:将输入序列x ( t ) x(t)x(t)——例如在迷宫中向左下方移动映射到潜在状态表示h ( t ) h(t)h(t)——例如到出口的距离以及x/y坐标并由此推导出预测输出序列y ( t ) y(t)y(t)——例如再次向左移动以更早到达出口。然而它并非使用离散序列例如向左移动一次而是以连续序列为输入并预测输出序列。SSM假设诸如物体在三维空间中运动之类的动态系统可以通过两个方程从其在时刻t tt的状态进行预测。h ′ ( t A h ( t ) B x ( t ) h(tAh(t)Bx(t)h′(tAh(t)Bx(t)y ( t ) C h ( t ) D x ( t ) y(t)Ch(t)Dx(t)y(t)Ch(t)Dx(t)通过求解这些方程我们假设能够揭示出统计规律从而根据观测数据输入序列和先前状态预测系统的状态。它的目标是找到这种状态表示h ( t ) h(t)h(t)以便我们能够从输入序列过渡到输出序列。这两个方程是状态空间模型的核心。在本指南中将多次引用这两个方程。为了使它们更易于理解我们对它们进行了颜色编码以便您能快速查阅。状态方程描述了状态如何变化通过矩阵A以及输入如何影响状态通过矩阵B。正如我们之前所见h ( t ) h(t)h(t)表示在任意时刻 t 的隐状态表示x ( t ) x(t)x(t)则表示某种输入。输出方程描述了状态如何通过矩阵 C 转换为输出以及输入如何通过矩阵 D 影响输出。注意矩阵A、B、C和D通常也被称为参数因为它们是可学习的。将这两个方程可视化我们得到以下架构让我们逐步了解这一通用技术以理解这些矩阵如何影响学习过程。假设我们有一个输入信号x ( t ) x(t)x(t)该信号首先会与矩阵B相乘而矩阵B描述了输入如何影响系统。更新后的状态类似于神经网络的隐藏状态是一个潜在空间其中包含了环境的核心“知识”。我们将该状态与矩阵A相乘矩阵A描述了所有内部状态之间的连接方式从而表征了系统的内在动力学。正如你可能注意到的矩阵A是在创建状态表示之前应用的并在状态表示更新之后进行更新。然后我们使用矩阵C来描述状态如何转化为输出。最后我们可以利用矩阵D实现从输入到输出的直接信号传递。这种连接方式通常也被称为**“跳跃连接”**。由于矩阵D类似于跳跃连接因此SSM通常被视为不带跳跃连接的以下形式回到简化的视角我们现在可以将矩阵A、B和C作为SSM的核心来重点研究。可以表示为可以像之前那样更新原始方程并添加一些鲜艳的色彩以表明每个矩阵的用途。这两条方程共同旨在根据观测数据预测系统的状态。由于输入预计为连续信号状态空间模型的主要表示形式是连续时间表示。2.3 从连续信号到离散信号如果你处理的是连续信号那么寻找状态表示h(t)在理论上会面临很大挑战。此外由于我们通常处理的是离散输入比如文本序列因此希望对模型进行离散化处理。为此我们采用了零阶保持技术Zero-order hold technique。其工作原理如下首先每当我们接收到一个离散信号时便将其值保持下来直至接收到新的离散信号。这一过程生成了一种连续信号可供SSM使用我们保持该值的时间长短由一个名为步长∆step size的新可学习参数来表示。它代表了输入的分辨率。现在我们有了输入的连续信号就可以生成一个连续的输出并仅按输入的时间步长对数值进行采样。这些采样值就是我们的离散化输出从数学上讲我们可以按如下方式应用零阶保持器它们共同使我们能够从连续的SSM过渡到离散的SSM这种离散SSM采用一种新的表述形式不再是函数到函数的映射x(t) → y(t)而变成了序列到序列的映射xₖ → yₖ。在这里矩阵A和B现在代表了模型的离散化参数。我们用k代替t来表示离散化的时间步长以便在提及连续与离散状态空间模型时更加清晰。注意在训练期间仍保存的是矩阵A的连续形式而非离散化后的版本。在训练过程中连续表示会被离散化。2.4 循环表示我们所采用的离散化状态空间模型使我们能够以特定的时间步长来表述问题而非连续信号。正如我们之前在循环神经网络中所见这种递归方法在此处非常有用。如果我们考虑离散时间步长而非连续信号便可以重新以时间步长来表述该问题在每个时间步我们计算当前输入Bxₖ如何影响先前状态Ahₖ₋₁然后计算预测输出Chₖ。这种表示方式或许已经显得有些熟悉了我们可以像之前处理RNN那样以同样的方式来处理它。我们可以将其展开或展开成如下请注意我们如何能够利用这种离散化版本并基于RNN的底层方法论来实现。2.5 卷积表示我们还可以用卷积来表示SSM。还记得在经典的图像识别任务中我们如何应用滤波器内核以提取聚合特征吗由于我们处理的是文本而非图像因此我们需要采用一维视角运用来自不同领域的技术能打造出一条有趣的流水线。用于表示此“滤波器filter”的内核源自SSM公式让我们来探讨一下这个内核在实际中的工作原理。与卷积类似我们也可以利用SSM内核遍历每组标记并计算输出这也说明了填充可能对输出产生的影响。我调整了填充的顺序以改善可视化效果但通常会在句子末尾应用填充。下一步内核会向右移动一次以执行计算的下一步在最后一步我们可以看到内核的全部效果将SSM表示为卷积的一大优势在于它能够像卷积神经网络CNN一样并行训练。然而由于核大小固定其推理速度不如循环神经网络RNN那样快速且不受限制。2.7 三种表示方式的对比这三种表示方法——连续、循环和卷积——各自都具有一系列不同的优缺点有趣的是我们现在实现了基于循环SSM的高效推理以及基于卷积SSM的可并行训练。借助这些表示方法我们可采用一种巧妙的技巧根据任务的不同选择相应的表示。在训练过程中我们使用可并行化的卷积表示而在推理阶段则采用高效的循环表示。该模型被称为线性状态空间层Linear State-Space Layer, LSSL。这些表示法具有一项重要特性即线性时不变性Linear Time Invariance, LTI。所谓线性时不变性是指SSM的参数A、B和C在所有时间步长上均保持固定。这意味着对于SSM生成的每一个标记矩阵A、B和C都是相同的。换句话说无论你向SSM输入何种序列A、B和C的值始终保持不变。我们所采用的是一种静态表示它并不具备内容感知能力。在探讨Mamba如何解决这一问题之前让我们先来了解一下拼图的最后一块——矩阵A。2.8 矩阵A的重要性可以说SSM公式中最重要的方面之一就是矩阵A。正如我们之前在循环表示中所看到的它会捕捉前一状态的信息以构建新状态。本质上矩阵A生成隐藏状态因此创建矩阵A的差异可能就在于是仅记住前几个标记还是捕捉迄今为止我们所见过的每一个标记。尤其是在循环表示的背景下因为这种表示方式只回看之前的状态。那么我们该如何以一种能够保留大容量记忆上下文大小的方式创建矩阵A呢我们使用“饥饿的河马Hungry Hungry Hippo”或者称作HiPPO3用于高阶多项式投影算子。HiPPO试图将其迄今为止所见的所有输入信号压缩成一个系数向量。它使用矩阵A构建一种状态表示能够很好地捕捉近期的标记并对较早的标记进行衰减。其公式可表示如下假设我们有一个方阵A这将给我们使用HiPPO构建矩阵A被证明比将其初始化为随机矩阵要好得多。因此与旧信号初始标记相比它能更准确地重建较新的信号近期标记。HiPPO矩阵的核心思想在于它能够生成一种记忆其历史的隐藏状态。从数学上讲它通过追踪Legendre多项式的系数来实现这一功能从而能够近似表示此前的所有历史。随后HiPPO被应用于我们之前所见到的循环和卷积表示以处理长距离依赖关系。其成果便是结构化状态空间序列模型Structured State Space for Sequences, S4一类能够高效处理长序列的SSM。它由三个部分组成状态空间模型用于处理长距离依赖关系的HiPPO用于构建循环和卷积表示的离散化方法这类SSM具有多种优势具体取决于您选择的表示方式循环型 vs. 卷积型。它还能通过基于HiPPO矩阵高效地处理长文本序列并有效存储记忆。注意如果想深入了解如何计算HiPPO矩阵并自行构建S4模型的技术细节我强烈建议您仔细阅读《带注释的S4》。3. Mamba——一种选择性SSM至此我们终于涵盖了理解Mamba独特之处所需的所有基础知识。状态空间模型可用于对文本序列建模但仍然存在一些我们希望避免的缺点。在本节中我们将详细介绍Mamba的两大主要贡献一种选择性扫描算法使模型能够过滤不相关信息一种面向硬件的算法通过并行扫描、内核融合和重新计算实现中间结果的高效存储。这两者共同构成了选择性SSM或S6模型这些模型可像自注意力机制一样用于构建Mamba模块。在深入探讨这两项主要贡献之前让我们先了解一下它们为何必不可少。3.1 它试图解决什么问题状态空间模型甚至包括S4结构化状态空间模型在语言建模与生成中至关重要的某些任务上表现欠佳尤其是对特定输入进行聚焦或忽略的能力。我们可以通过两个合成任务来说明这一点即选择性复制selective copying和归纳头任务(induction heads)。在选择性复制任务中状态空间模型的目标是按顺序复制输入中的部分内容并将其输出然而递归/卷积SSM 在这项任务中表现不佳因为它具有线性时不变特性。正如我们之前所见SSM 生成的每个标记所对应的矩阵 A、B 和 C 均相同。因此SSM 无法进行内容感知推理因为它会因矩阵 A、B 和 C 固定不变而对每个标记一视同仁。这构成了一个问题因为我们希望 SSM 能够针对输入提示进行推理。SSM表现不佳的第二个任务是归纳头其目标是重现输入中所发现的模式在上述示例中我们实际上是在进行一次性提示试图“教会”模型在每个“Q:”之后都给出“A:”的回复。然而由于SSM是时间不变的它无法从历史中选择哪些先前的标记进行调用。让我们以矩阵B为例来说明这一点。无论输入x为何矩阵B始终保持完全不变因此它与x无关同样无论输入如何变化A和C也始终保持不变。这体现了我们迄今所见的SSM的静态特性。相比之下这些任务对Transformer来说相对容易因为它们会根据输入序列动态地调整注意力。它们能够有选择地“关注”或“留意”序列的不同部分。SSM在这些任务上表现不佳凸显了时不变SSM所存在的根本问题矩阵A、B和C的静态特性导致其在内容感知方面存在缺陷。3.2 有选择地保留信息SSM的递归表示会创建一个状态较小且非常高效的模型因为它能够压缩整个历史信息。然而与通过注意力矩阵完全不压缩历史信息的Transformer模型相比它的表现要逊色得多。Mamba的目标是兼得两者之长既拥有小型状态又具备与Transformer模型相当的强大能力。正如上文所暗示的它是通过将数据有选择地压缩到状态中来实现这一功能的。当你输入一个句子时其中往往包含一些诸如停用词之类的词语这些词语本身并没有太多实际意义。为了有选择地压缩信息我们需要使参数依赖于输入。为此我们首先来探讨一下在训练过程中SSM中输入与输出的维度在结构化状态空间模型S4中矩阵A、B和C与输入无关因为它们的维度N和D是静态的不会发生变化。相反Mamba 通过引入输入的序列长度和批次大小使矩阵 B 和 C 乃至步长 ∆ 都依赖于输入这意味着对于每个输入标记我们现在都有不同的B和C矩阵从而解决了内容感知问题注意矩阵A保持不变因为我们希望状态本身保持静态但其受外界影响的方式通过B和C则应是动态的。它们共同选择性地决定在隐藏状态中保留什么、忽略什么因为此时它们已开始依赖输入。步长∆越小就越倾向于忽略特定的词语而更多地利用先前的上下文步长∆越大则越注重输入的词语而非上下文。3.3 扫描操作由于这些矩阵现在是动态的因此无法再采用卷积表示法进行计算因为卷积表示法假定核是固定的。我们只能使用递归表示法从而失去卷积带来的并行化优势。为了实现并行化让我们来探讨一下如何通过递归方式计算输出每个状态都是前一状态乘以A与当前输入乘以B之和。这种操作称为扫描运算可轻松通过for循环来实现。相比之下由于每个状态的计算都依赖于前一状态因此并行化似乎是不可能的。然而Mamba借助并行扫描算法使这一目标成为可能。它利用结合律假定操作的顺序并不重要。因此我们可以分段计算序列并逐步将各部分合并起来动态矩阵B和C与并行扫描算法共同构成了选择性扫描算法以体现递归表示所具有的动态性和快速性。3.4 硬件感知算法近期GPU的一个缺点是其小型但高效能的SRAM与大型但效率稍低的DRAM之间的传输I/O速度有限。频繁地在SRAM和DRAM之间复制数据已成为一个瓶颈。Mamba与Flash Attention一样试图减少我们从DRAM到SRAM以及反向访问的次数。它通过**内核融合kernel fusion**实现这一目标使模型能够避免写入中间结果并持续进行计算直至任务完成。我们可以通过可视化Mamba的基础架构查看DRAM和SRAM分配的具体实例在此以下内容被融合为一个内核步长为∆的离散化步骤选择性扫描算法与C的乘法运算硬件感知算法的最后一部分是重新计算recomputation。中间状态不会被保存但对反向传播计算梯度却是必需的。因此作者在反向传播过程中重新计算这些中间状态。尽管这看似效率不高但实际上其成本远低于从速度较慢的DRAM中读取所有这些中间状态。我们现已涵盖了其架构的所有组成部分相关示意图见其文章中的以下图片图选择性SSM。来源Gu, Albert, and Tri Dao. “Mamba: Linear-time sequence modeling with selective state spaces.” arXiv preprint arXiv:2312.00752 (2023).这种架构通常被称为选择性SSM或S6模型因为它本质上是采用选择性扫描算法计算得出的S4模型。3.5 Mamba Block我们迄今所探讨的这种选择性SSM可以像我们在解码器块中表示自注意力机制那样以一个模块的形式实现。与解码器类似我们可以堆叠多个Mamba模块并将它们的输出用作下一个Mamba模块的输入它首先通过线性投影来扩展输入嵌入input embeddings。随后在应用选择性SSM之前会先进行卷积操作以防止对各个标记进行独立计算。The Mamba Block选择性SSM具有以下特性通过离散化创建的递归SSM对矩阵A进行HiPPO初始化以捕捉长程依赖关系选择性扫描算法用于有选择地压缩信息面向硬件的算法以加速计算在查看代码实现时我们可以进一步扩展这一架构并探讨一个端到端示例的具体模样请注意一些变化例如引入了归一化层以及用于选择输出标记的Softmax函数。将所有这些整合起来我们便实现了快速的推理与训练甚至还能支持无限长度的上下文。借助这种架构作者发现其性能不仅与同等规模的Transformer模型相当有时甚至超越了后者4. 结论至此我们关于状态空间模型以及采用选择性状态空间模型的惊艳Mamba架构的旅程便告一段落。希望本文能帮助您更好地理解状态空间模型——尤其是Mamba——的潜力。谁又知道未来它是否会取代Transformer呢但眼下看到如此迥异的架构获得应有的关注实在令人振奋5. 资源希望这能为你提供一个易于理解的Mamba与状态空间模型入门介绍。如果你想深入学习我推荐以下资源The Annotated S4 is a JAX implementation and guide through the S4 model and is highly advised!A great YouTube video introducing Mamba by building it up through foundational papers.The Mamba repository with checkpoints on Hugging Face.An amazing series of blog posts (1, 2, 3) that introduces the S4 model.The Mamba No. 5 (A Little Bit Of…) blog post is a great next step to dive into more technical details about Mamba but still from an amazingly intuitive perspective.And of course, the Mamba paper! It was even used for DNA modeling and speech generation.1 Gu, Albert, and Tri Dao. “Mamba: Linear-time sequence modeling with selective state spaces.” arXiv preprint arXiv:2312.00752 (2023).2 Gu, Albert, et al. “Combining recurrent, convolutional, and continuous-time models with linear state space layers.” Advances in neural information processing systems 34 (2021): 572-585.3 Gu, Albert, et al. “Hippo: Recurrent memory with optimal polynomial projections.” Advances in neural information processing systems 33 (2020): 1474-1487.4 Voelker, Aaron, Ivana Kajić, and Chris Eliasmith. “Legendre memory units: Continuous-time representation in recurrent neural networks.” Advances in neural information processing systems 32 (2019).5 Gu, Albert, Karan Goel, and Christopher Ré. “Efficiently modeling long sequences with structured state spaces.” arXiv preprint arXiv:2111.00396 (2021).