人工智能

Mamba：重新定义序列建模和 Outforming Transformers 架构

发布时间 2023 年 12 月 18 日

阿尤什·米塔尔米塔尔

在这篇关于曼巴我们将探讨这种创新的状态空间模型 (SSM) 如何彻底改变序列建模。Mamba 由 Albert Gu 和 Tri Dao 开发，因其在语言处理、基因组学和音频分析等领域处理复杂序列的高效性而著称。其具有选择性状态空间的线性时间序列建模确保了其在这些不同模态下均具有卓越的性能。

我们将深入探讨 Mamba 如何克服传统 Transformer 所面临的计算挑战，尤其是在处理长序列时。其在状态空间模型中的选择性方法可以实现更快的推理速度，并随着序列长度线性缩放，从而显著提高吞吐量。

Mamba 的独特之处在于其快速的处理能力、选择性 SSM 层以及受 FlashAttention 启发的硬件友好型设计。这些特性使 Mamba 的性能超越了许多现有模型，包括基于 Transformer 方法的模型，使其成为机器学习领域一项令人瞩目的进步。

变形金刚 vs 曼巴

像 GPT-4 这样的 Transformer 已经在自然语言处理领域树立了标杆。然而，它们的效率会随着序列的延长而下降。而 Mamba 在这方面则取得了突破，它能够更高效地处理长序列，其独特的架构也简化了整个流程。

转换器擅长处理数据序列，例如语言模型的文本。与以前按顺序处理数据的模型不同，Transformers 同时处理整个序列，使它们能够捕获数据中的复杂关系。

他们使用注意力机制，使模型在进行预测时能够关注序列的不同部分。

这种注意力机制使用三组权重来计算：查询、键和值，这些权重源自输入数据。序列中的每个元素都会与其他所有元素进行比较，从而提供一个权重，该权重表示每个元素在预测序列中的下一个元素时应获得的重要性或“注意力”。

Transformer 维护两个主要模块：处理输入数据的编码器和生成输出的解码器。该编码器由多个层组成，每个层包含两个子层：多头自注意力机制和简单的位置式全连接前馈网络。每个子层都使用归一化和残差连接来帮助训练深度网络。

解码器也具有与编码器类似的包含两个子层的层，但增加了第三个子层，用于对编码器的输出执行多头注意力机制。解码器的顺序性确保对某个位置的预测只能考虑先前的位置，从而保留了自回归特性。

与变形金刚相比，曼巴模型采取了不同的方法。 Transformers 通过使用更复杂的注意力机制来处理长序列问题，而 Mamba 使用选择性状态空间，提供了更多计算能力

以下是变压器功能概述：

输入处理：Transformers 首先将输入数据编码为模型可以理解的格式，通常使用嵌入来合并序列中每个元素的位置。
注意机制：注意力机制的核心是计算一个分数，该分数表示在理解当前元素时应将多少注意力放在输入序列的其他部分上。
编码器-解码器架构：Transformer 模型由一个用于处理输入的编码器和一个用于生成输出的解码器组成。每个模型都包含多个层，用于细化模型对输入的理解。
多头注意力：在编码器和解码器中，多头注意力允许模型同时关注来自不同表示空间的序列的不同部分，从而提高其从不同上下文中学习的能力。
位置式前馈网络：经过关注后，一个简单的神经网络会单独且相同地处理每个位置的输出。这通过残差连接与输入相结合，然后进行层归一化。
输出生成：然后，解码器会预测输出序列，该序列受到编码器的上下文及其迄今为止生成的内容的影响。

Transformer 并行处理序列的能力及其强大的注意力机制使其对于翻译和文本生成等任务非常强大。

相比之下，Mamba 模型的运作方式有所不同，它使用选择性状态空间来处理序列。这种方法解决了 Transformer 在处理长序列时计算效率低下的问题。Mamba 的设计能够实现更快的推理，并随序列长度线性扩展，为序列建模树立了新的范式，使其能够更加高效，尤其是在序列越来越长的情况下。

曼巴

Mamba 真正的独特之处在于它脱离了传统的注意力和 MLP 块。这种简化带来了更轻、更快的模型，可以随序列长度线性扩展——这是其前辈无法比拟的壮举。

曼巴的主要特点包括：

选择性SSM：这些允许 Mamba 过滤不相关的信息并专注于相关数据，从而增强其对序列的处理。这种选择性对于有效的基于内容的推理至关重要。
硬件感知算法：Mamba 采用针对现代硬件（尤其是 GPU）优化的并行算法。与传统模型相比，这种设计可以实现更快的计算速度并降低内存需求。
简化架构：通过集成选择性 SSM 并消除注意力和 MLP 块，Mamba 提供了更简单、更均匀的结构。这会带来更好的可扩展性和性能。

Mamba 在语言、音频和基因组学等多个领域表现出了卓越的性能，在预训练和特定领域任务中表现出色。例如，在语言建模中，Mamba 的性能可与大型 Transformer 模型相媲美或超过。

Mamba 的代码和预训练模型已公开供社区使用，网址为 GitHub上.

对于线性模型来说，标准复制任务很简单。选择性复印和感应头需要用于法学硕士的动态、内容感知存储器。

结构化状态空间 (S4) 模型最近成为一类有前途的序列模型，包含 RNN、CNN 和经典状态空间模型的特征。 S4 模型的灵感来自连续系统，特别是通过隐式潜在状态映射一维函数或序列的系统。在深度学习的背景下，它们代表了一项重大创新，为设计高效且高度适应性的序列模型提供了一种新的方法。

S4 车型的动力学

SSM（S4） 这是基本的结构化状态空间模型。需要一个序列 x 并产生输出 y 使用学习到的参数 A, B, C，和一个延迟参数 Δ转换涉及离散化参数（将连续函数转换为离散函数）并应用 SSM 操作，该操作是时不变的，这意味着它不会随着时间步骤而改变。

离散化的意义

离散化是通过固定公式将连续参数转化为离散参数的关键过程，使S4模型能够与连续时间系统保持联系。这赋予模型额外的属性，例如分辨率不变性，并确保适当的归一化，从而增强模型的稳定性和性能。离散化还与 RNN 中的门控机制相似，这对于管理网络中的信息流至关重要。

线性时不变性 (LTI)

S4 模型的核心特征是其线性时不变性。此属性意味着模型的动态随着时间的推移保持一致，所有时间步长的参数都是固定的。 LTI 是递归和卷积的基石，为构建序列模型提供了一个简化但功能强大的框架。

克服基本限制

S4 框架传统上受到其 LTI 性质的限制，这对需要自适应动态的数据建模提出了挑战。最近的研究论文提出了一种通过引入时变参数来克服这些限制的方法，从而消除了 LTI 的约束。这使得 S4 模型能够处理更加多样化的序列和任务，从而显着扩展了其适用性。

“状态空间模型”这一术语广泛涵盖任何涉及潜在状态的循环过程，并已用于描述多个学科的各种概念。在深度学习的背景下，S4 模型（或称结构化 SSM）是指一类特定的模型，这些模型经过优化，可实现高效计算，同时保留了对复杂序列进行建模的能力。

S4 模型可以集成到端到端神经网络架构中，充当独立的序列转换。它们可以被视为类似于 CNN 中的卷积层，为各种神经网络架构中的序列建模提供骨干。

SSM 与 SSM + 选择

序列建模中选择性的动机

结构化SSM

该论文认为，序列建模的一个基本方面是将上下文压缩为可管理的状态。可以选择性地关注或过滤输入的模型提供了维持这种压缩状态的更有效方法，从而产生更高效和更强大的序列模型。这种选择性对于模型自适应地控制信息如何沿着序列维度流动至关重要，这是处理语言建模及其他领域的复杂任务的基本能力。

选择性 SSM 通过允许其参数依赖于输入来增强传统 SSM，这引入了以前时不变模型无法实现的一定程度的自适应性。这导致时变 SSM 无法再使用卷积来进行高效计算，而是依赖于线性递归机制，这与传统模型存在显着偏差。

SSM + 选择 (S6) 该变体包括选择机制，为参数添加输入依赖性 B 与 C，和一个延迟参数 Δ。这允许模型有选择地关注输入序列的某些部分 x。考虑到选择，对参数进行离散化，并使用扫描操作以随时间变化的方式应用 SSM 操作，该扫描操作顺序处理元素，随时间动态调整焦点。

曼巴的性能亮点

Mamba 在每一项评估结果上都是同类最佳

在性能方面，Mamba 在推理速度和准确性方面均表现出色。其设计能够更好地利用更长的上下文，这在 DNA 和音频建模中均有体现，在需要长距离依赖关系的复杂任务上，其性能优于先前的模型。其多功能性也在跨多个任务的零样本评估中得到凸显，为此类模型在效率和可扩展性方面树立了新的标杆。

曼巴入门

对于那些有兴趣利用 Mamba 的人来说，技术要求包括 Linux 操作系统、NVIDIA GPU、PyTorch 1.12+ 和 CUDA 11.6+。安装涉及简单的 pip 命令来安装必要的软件包曼巴存储库。如果 PyTorch 版本出现兼容性问题，使用 pip 的 –no-build-isolation 标志会有所帮助。这些模型在 Pile 和 SlimPajama 数据集等广泛数据集上进行训练，旨在满足各种计算需求和性能基准。

Mamba 提供不同级别的接口，从选择性 SSM 层到 Mamba 模块以及完整的语言模型结构。Mamba 模块是该架构的主要模块，它利用因果 Conv1d 层，可以轻松集成到神经网络设计中。提供的 Python 使用示例演示了如何实例化 Mamba 模型并通过该模型处理数据，凸显了该系统的简洁性和灵活性。

预训练的 Mamba 模型可在拥抱脸，参数大小从 130M 到 2.8B 不等，在广泛的 Pile 数据集和 SlimPajama 数据集上进行训练。这些模型旨在满足不同的计算和性能要求，并遵守 GPT-3 的尺寸标准。用户可以期望这些模型具有高吞吐量和准确性，使 Mamba 成为各种应用程序的竞争选择，包括但不限于语言建模。