共计 6493 个字符，预计需要花费 17 分钟才能阅读完成。
Transformer 体系结构曾经成为大型语言模型 (llm) 胜利的次要组成部分。为了进一步改良 llm，人们正在研发可能优于 Transformer 体系结构的新体系结构。其中一种办法是 Mamba（一种状态空间模型）。
Mamba: Linear-Time Sequence Modeling with Selective State Spaces 一文中提出了 Mamba，咱们在之前的文章中也有具体的介绍。
在本篇文章中，通过将绘制 RNN，transformer，和 Mamba 的架构图，并进行具体的比照，这样咱们能够更具体的理解它们之间的区别。
为了阐明为什么 Mamba 是这样一个乏味的架构，让咱们先介绍 Transformer。
Transformer
Transformer 将任何文本输出视为由令牌组成的序列。
transformer 的一个次要长处是，无论它接管到多长的输出，它都应用序列中的任何令牌信息（无论序列有多长）来对输出数据进行解决。
这就是咱们在论文中看到的注意力机制的作用，然而为了取得全局信息，注意力机制在长序列上十分消耗显存，这个咱们前面说。
Transformer 由两个构造组成，一组用于示意文本的编码器块和一组用于生成文本的解码器块。这些构造能够用于多种工作，包含翻译。
咱们能够采纳这种构造来创立仅应用解码器的生成模型。比方基于 Transformer 的 GPT，应用解码器块来实现一些输出文本。
单个解码器块由两个次要局部组成，一个是自注意力模块，另一个是前馈神经网络。
注意力创立一个矩阵，将每个令牌与之前的每个令牌进行比拟。矩阵中的权重由令牌对之间的相关性决定。
它反对并行化，所以能够极大地放慢训练速度!
然而当生成下一个令牌时，咱们须要从新计算整个序列的注意力，即便咱们曾经生成了一些新的令牌。
为长度为 L 的序列生成令牌大概须要 L²的计算量，如果序列长度减少，计算量可能会很大。并且在这里须要计算所有令牌的注意力，所以如果序列很长，那么内存占用也会很大。所以须要从新计算整个序列是 Transformer 体系结构的次要瓶颈。当然也有很多技巧来晋升注意力机制的效率，这里咱们临时不提，只看最经典的原始论文。
RNN
上面咱们介绍更早的序列模型 RNN。循环神经网络 (RNN) 是一种基于序列的网络。它在序列的每个工夫步长取两个输出，即工夫步长 t 的输出和前一个工夫步长 t - 1 的暗藏状态，以生成下一个暗藏状态并预测输入。
RNN 有一个循环机制，容许它们将信息从上一步传递到下一步。咱们能够“开展”这个可视化，使它更明确。
在生成输入时，RNN 只须要思考之前的暗藏状态和以后的输出。这样不会从新计算以前的暗藏状态，这正 Transformer 没有的。
这种流程能够让 RNN 进行疾速推理，因为的工夫与序列长度线性扩大! 并且能够有有限的上下文长度（实践上），因为每次推理他只取一个暗藏状态和以后输出，内存的占用是十分稳固的。
咱们将 RNN 利用于之前应用过的输出文本。
每个暗藏状态都是以前所有暗藏状态的聚合。然而这里就呈现了问题，在生成名称“Maarten”时，最初一个暗藏状态不再蕴含对于单词“Hello”的信息（或者说最早的信息会被坐进的信息笼罩）。这会导致随着工夫的推移，rnn 会遗记信息，因为它们只思考前一个状态。
并且 rnn 的这种程序性产生了另一个问题。训练不能并行进行，因为它须要按程序实现每一步。
与 Transformer 相比，rnn 的问题齐全相同! 它的推理速度十分快，但不能并行化导致训练很慢。
人们始终在寻找一种既能像 Transformer 那样并行化训练，可能记住先前的信息，并且在推理工夫还是随序列长度线性增长的模型，Mamba 就是这样宣传的。
在介绍 Mamba 之前，让咱们还须要介绍以下状态空间模型
The State Space Model (SSM)
状态空间模型(SSM)，像 Transformer 和 RNN 一样，能够解决序列信息，比方文本，也包含信号。
状态空间是蕴含可能齐全形容一个零碎的起码数量变量的概念。它是一种通过定义零碎可能的状态来数学示意问题的形式。
比如说咱们正在通过一个迷宫。” 状态空间 ” 就是所有可能地位（状态）的地图。每个点代表迷宫中的一个独特地位，具备特定的细节，比方你离进口有多远。
“ 状态空间示意 ” 是对这个地图的简化形容。它展现了你以后所处的地位（以后状态），以及下一步能够去哪里（可能的将来）。
尽管状态空间模型应用方程和矩阵来跟踪这种行为，形容状态的变量，在咱们的例子中是 X 和 Y 坐标以及到进口的间隔，能够示意为“状态向量”。
听起来相熟吗? 这不就是强化学习中的状态吗，我集体认为是能够这么了解的，那么怎么和序列无关呢？
因为语言模型中的嵌入或向量也常常用于形容输出序列的“状态”。例如，你以后地位的向量 (状态向量) 可能看起来像这样:
在神经网络中，“状态”通常是指其暗藏状态，在大型语言模型的背景下，这是生成新标记的一个最重要的方面之一。
状态空间模型（SSMs）是用于形容这些状态示意并依据某些输出进行下一个状态预测的模型。
在工夫 t，状态空间模型(SSMs)：
- 将输出序列 x(t)（例如，在迷宫中向左和向下挪动）映射到潜在状态示意 h(t)（例如，到进口的间隔和 x / y 坐标），
- 并推导出预测的输入序列 y(t)（例如，再次向左挪动以更快地达到进口）。
这里就与强化学习中应用离散序列（如仅向左挪动一次）不同，它将间断序列作为输出并预测输入序列。
ssm 假如动静零碎，例如在三维空间中挪动的物体，能够通过两个方程从工夫 t 的状态预测。
通过求解这些方程，假如能够揭示基于观测数据 (输出序列和先前状态) 预测零碎状态的统计原理。
它的指标是找到这个状态示意 h(t)这样咱们就能够从一个输出序列到一个输入序列。
这两个方程就是是状态空间模型的外围。状态方程形容了基于输出如何影响状态 (通过矩阵 B) 的状态变动(通过矩阵 A)。
h(t)示意任意时刻 t 的潜在状态示意，而 x(t)示意某个输出。
输入方程形容了状态如何转化为输入(通过矩阵 C)，以及输出如何影响输入(通过矩阵 D)。
矩阵 A、B、C 和 D 通常被称为参数，因为它们是可学习的。将这两个方程可视化，咱们能够失去如下架构:
上面咱们看看这些矩阵如何影响学习过程。
假如咱们有一个输出信号 x(t)这个信号首先乘以矩阵 B 它形容了输出如何影响零碎。
更新状态 (h) 是蕴含环境外围“常识”的潜在空间。咱们将状态与矩阵 A 相乘，矩阵 A 形容了所有外部状态是如何连贯的，因为它们代表了零碎的潜在示意。
这里能够看到，在创立状态示意之前利用矩阵 A，并在状态示意更新之后更新矩阵 A。
而后应用矩阵 C 来形容如何将状态转换为输入。
最初利用矩阵 D 提供从输出到输入的间接信号。这通常也被称为跳过（残差）连贯。
因为矩阵 D 相似于跳过连贯，所以 SSM 通常被视为为不进行跳过连贯的局部
回到咱们的简化视图，当初能够将重点放在矩阵 A、B 和 C 上，它们是 SSM 的外围。
更新原始方程并增加一些色彩来示意每个矩阵的目标
这两个方程依据观测数据预测零碎的状态。因为冀望输出是间断的，SSM 是间断工夫示意。
然而因为文字都是离散的输出，咱们还须要将模型离散化。这里就要应用 Zero-order hold 技术
每次咱们接管到一个离散信号，都会保障他的值不变，直到接管到一个新的离散信号再扭转。这个过程创立了一个 SSM 能够应用的间断信号:
咱们放弃该值的工夫由一个新的可学习参数示意，称为步长∆。这样就失去了一个间断的信号并且能够只依据输出的工夫步长对值进行采样。
这些采样值就是咱们的离散输入! 在数学上，咱们能够利用 Zero-order hold 如下:
因为咱们 SSM 解决的是离散信号，所以这里不是一个函数到函数，x(t)→y(t)，而是一个序列到序列，xₖ→yₖ，咱们用公式示意如下:
矩阵 A 和 B 当初示意模型的离散参数，用 k 代替 t 来示意离散的工夫步长。
离散化的 SSM 容许在特定的工夫步中解决信息。就像咱们之前在循环神经网络（RNNs）中看到的那样，循环办法在这里也十分有用，能够将问题从新表述为工夫步骤：
在每个工夫步长，咱们计算以后输出 (Bxₖ) 如何影响前一个状态(Ahₖ₁)，而后计算预测输入(Chₖ)。
这种示意看起来是不是有点相熟？其实他的解决办法和 RNN 一样
也能够这样开展:
这种技术与 RNN 相似，疾速推理和慢速训练。
另一种 ssm 的示意是卷积的示意。咱们利用过滤器 (核) 来取得聚合特色:
因为咱们解决的是文本而不是图像，所以我只有一维的视角:
咱们用来示意这个“过滤器”的核是由 SSM 公式推导进去的:
能够应用 SSM 核遍历每一组令牌并计算输入:
上图也阐明了 padding 可能对输入产生的影响，所以咱们个别都会在开端 padding 而不是在后面。第二步核被挪动一次来执行下一步计算:
在最初一步，咱们能够看到核的残缺成果:
卷积的一个次要益处是它能够并行训练。然而因为核大小是固定，它们的推理不如 rnn 疾速并且对序列长度有限度。
下面的三种 SMM 都有各自的优缺点
这里能够应用一个简略的技巧，即依据工作抉择示意。在训练过程中应用能够并行化的卷积示意，在推理过程中，咱们应用高效的循环示意:
听起来有点奇幻，然而有人就是实现进去了，这个模型叫做 Linear State-Space Layer (LSSL)
https://proceedings.neurips.cc/paper_files/paper/2021/hash/05…
它联合了线性动静零碎实践和神经网络的概念，能够无效地捕捉数据中的时序信息和动静特色。LSSL 基于线性动静零碎实践，这种零碎能够用状态空间模型表示。在这个模型中，零碎的行为由状态变量的演变和内部管制信号的影响决定。状态变量是零碎的外部示意，能够捕捉零碎的动静个性。
这些示意都有一个重要的个性，即线性时不变性(LTI)。LTI 示意 ssm 参数 A、B 和 C 对于所有工夫步长都是固定的。这意味着对于 SSM 生成的每个令牌，矩阵 A、B 和 C 都是雷同的。
也就是说无论给 SSM 的序列是什么，A、B 和 C 的值都放弃不变。这样就失去了一个不感知内容的动态示意。然而动态示意没有任何意义对吧，所以 Mamba 解决的就是这个问题，然而在介绍 Mamba 之前，咱们还有一个知识点须要强调，那就是矩阵 A
因为 SSM 公式中最重要的就是矩阵 a。正如咱们之前在循环示意中看到的那样，它捕捉了对于前一个状态的信息来构建新状态，如果矩阵 a 如果跟 RNN 一样会忘记掉十分靠前的信息那么 SMM 将没有任何的意义，对吧。
矩阵 A 产生暗藏状态:
如何保留大上下文大小的形式创立矩阵 A 呢?
HiPPO 的模型联合了递归记忆（Recurrent Memory）和最优多项式投影（Optimal Polynomial Projections）的概念，这种投影技术能够显著改善递归记忆的性能，特地是在解决长序列和长期依赖关系时。
https://proceedings.neurips.cc/paper/2020/hash/102f0bb6efb3a6…
应用矩阵 A 来构建一个状态示意，该状态示意能够很好地捕捉最近的令牌并衰减较旧的令牌。其公式可示意为:
具体的具体内容咱们就不介绍了，有趣味的查看原论文。
这样咱们就基本上解决了所有的问题：1、状态空间模型；2、解决近程依赖关系；3、离散化和并行计算
如果想深刻理解无关如何计算 HiPPO 矩阵和本人构建 S4 模型建议您浏览正文的 S4。
https://srush.github.io/annotated-s4/
Mamba
下面介绍完所有必要的基础知识，最初就是咱们的重点了
Mamba 有两个次要奉献:
1、选择性扫描算法，模型能够过滤无关和无关的信息
2、硬件感知算法，通过并行扫描、核交融和重计算无效地存储 (两头) 后果。
在探讨这两个次要奉献之前，咱们先看看一下为什么它们是必要的。
状态空间模型，S4(Structured State Space Model)，在语言建模和生成中的某些工作上体现不佳
比方在选择性复制工作中，SSM 的指标是按程序复制输出和输入的局部:
(循环 / 卷积)SSM 在这个工作中体现不佳，因为它是线性时不变的。对于 SSM 生成的每个令牌，矩阵 A、B 和 C 都是雷同的。
因为它将每个令牌平等地视为固定的 a、B 和 C 矩阵的后果，所以 SSM 不能执行内容感知推理
SSM 体现不佳的第二个工作是重现输出中发现的模式:
咱们的提醒在“教”模型在每个“Q:”之后提供“A:”响应。然而因为 ssm 是工夫不变的，它不能抉择从其历史中获取先前的令牌。
以矩阵 B 为例不论输出 x 是什么，矩阵 B 放弃完全相同，并且与 x 无关:
同理无论输出是什么，A 和 C 也不变，这就是咱们下面说的动态。
而 Transformers 能够依据输出序列动静地扭转注意力。能够选择性地“看”或“留神”序列的不同局部，再加上地位编码，这使得 Transformers 对于这种工作十分的简略。
ssm 在这些工作上的蹩脚性能阐明了定常 ssm 的潜在问题，矩阵 A、B 和 C 的动态个性导致了内容感知问题。
选择性地保留信息
SSM 的循环示意创立了一个十分无效的小状态，因为它压缩了整个历史信息，所以与不压缩历史 (注意力矩阵) 的 Transformer 模型相比，它的性能要弱得多。
Mamba 的指标是取得 Transformer 一样弱小的“小”状态
通过有选择地将数据压缩到状态，当输出一个句子时，通常会有一些信息，比方进展词，这些信息没有太多的意义。
咱们先看看 SSM 在训练期时的输出和输入维度:
在结构化状态空间模型 (S4) 中，矩阵 a、B 和 C 独立于输出，因为它们的维度 N 和 D 是动态的，不会扭转。
而 Mamba 通过联合输出的序列长度和批量大小，使矩阵 B 和 C，甚至步长∆依赖于输出:
这意味着对于每个输出标记，有不同的 B 和 C 矩阵，这解决了内容感知问题! 这里矩阵 A 放弃不变，因为心愿状态自身放弃动态，但影响它的形式 (通过 B 和 C) 是动静的。
也就是说它们一起选择性地抉择将什么保留在暗藏状态中，什么须要疏忽，这都是由输出确定的。
较小的步长∆导致疏忽特定的单词，而是更多地应用之前的上下文，而较大的步长∆则更多地关注输出单词而不是上下文:
扫描操作
这些矩阵当初是动静的了所以它们不能应用卷积示意来计算，只能应用循环进行解决，这就使得无奈进行并行化。
为了实现并行化，咱们先看看循环的输入:
每个状态都是前一个状态 (乘以 A) 加上以后输出 (乘以 B) 的和。这被称为扫描操作，能够很容易地通过 for 循环计算出来。然而并行化仿佛是不可能的，因为每个状态只有在咱们有前一个状态时能力计算出来。
然而 Mamba 应用并行扫描算，通过关联属性假设执行操作的程序无关紧要。这样就能够计算局部序列并迭代组合它们:
这样还有一个益处是因为程序不重要，也能够省略掉 Transformer 的地位编码。
硬件感知的算法
最近 gpu 的一个毛病是它们在小但高效的 SRAM 和大但效率稍低的 DRAM 之间的传输 (IO) 速度无限。在 SRAM 和 DRAM 之间频繁地复制信息成为瓶颈。
Mamba 的 DRAM 和 SRAM 调配的具体实例如下:
中间状态不被保留，但对于反向流传计算梯度是必要的。作者从新计算了反向传递过程中的中间状态。只管这看起来效率很低，但它比从绝对较慢的 DRAM 读取所有这些中间状态的老本要低得多。
这里咱们就不具体阐明了，因为这部分我也没太钻研过
Mamba 块
选择性 SSM 能够作为一个块，就像在 Transformer 中的的注意力模块一样。咱们能够重叠多个块，并应用它们的输入作为下一个曼巴块的输出:
最初一个端到端（输出到输入）的例子蕴含了归一化层和抉择输入标记 softmax。
这样就失去了疾速的推理和训练，而且是“有限”长度上下文的模型
总结
看完这篇文章，我心愿你能对 Mamba 和状态空间模型有肯定的理解，最初咱们以作者的发现为结尾：
作者发现模型与雷同尺寸的 Transformer 模型的性能相当，有时甚至超过了它们!
https://avoid.overfit.cn/post/94105fed36de4cd981da0b916c0ced47
作者：Maarten Grootendorst
