Mamba LLM 架构简介：机器学习的新范式

探索 Mamba LLM 的强大功能，Mamba LLM 是来自一流大学的变革性架构，重新定义了 AI 中的序列处理。语言模型是一种经过训练的机器学习模型，用于在自然语言上执行概率分布。它们的架构主要由多层神经网络组成，例如递归层、前馈层、嵌入层和注意力层。这些层组合在一起以处理给定的输入文本并生成输出预测。

2023 年底，卡内基梅隆大学和普林斯顿大学的研究人员发表了一篇研究论文，揭示了一种名为 Mamba 的大型语言模型（LLM）的新架构。Mamba 是一种与序列建模有关的新状态空间模型架构。它是为了解决 transformer 模型的一些限制而开发的，尤其是在处理长序列时，并且已经显示出有希望的性能。

Mamba

Mamba 是一种新的 LLM 架构，它集成了结构化状态空间序列（S4）模型来管理冗长的数据序列。S4 结合了递归、卷积和连续时间模型的最佳功能，可以有效且高效地模拟长期依赖关系。这使它能够处理不规则采样的数据，具有无限的上下文，并在整个训练和测试过程中保持计算效率。

Mamba 扩展了 S4 范式，带来了几项值得注意的改进，尤其是在处理时变作方面。它的架构围绕着一种特殊的选择机制展开，该机制根据输入修改结构化状态空间模型（SSM）参数。

因此，Mamba 可以通过只关注序列中的关键信息来成功过滤掉不太重要的数据。根据 Wikipedia 的说法，“该模型从时不变框架过渡到时变框架，这会影响系统的计算和效率。

主要特点和创新

Mamba 偏离了传统的 attention 和 MLP 块，使其与众不同。这种简化使模型更轻、更快，并且与序列的长度呈线性缩放，这是其前辈都无法实现的壮举。

Mamba 的关键组件包括：

选择性状态空间（SSM）：根据当前输入有选择地处理信息的递归模型是 Mamba SSM 的基础。这使他们能够过滤掉无关数据并专注于相关信息，这可能会带来更高效的处理。
简化的架构： Mamba 用一个单一的、有凝聚力的 SSM 块取代了 Transformers 错综复杂的注意力和 MLP 块。这旨在加速推理并降低计算复杂性。
硬件感知并行性：Mamba 的性能可能会更好，因为它使用循环模式和专为硬件效率而创建的并行算法。
另一个关键组成部分是线性时间不变性（LTI）;LTI 是 S4 模型的核心功能之一。这一特性表明，模型的参数在所有时间步中保持不变，从而保持模型动力学的一致性。使用 LTI 构建序列模型更容易、更有效，LTI 是递归和卷积的基础。

Mamba LLM 架构详细信息

Mamba 的架构进一步强调了机器学习进步的重要性。它通过引入选定的状态空间模型（SSM）层来修改模型处理序列的方式。这使 Mamba 能够做两件极其重要的事情：

关注相关信息 – Mamba 可以通过为每个输入分配不同的权重来为任务确定更多预测数据的优先级。

动态适应输入 – 由于模型能够适应输入，Mamba 可以轻松处理各种序列建模工作。

因此，Mamba 可以以前所未有的效率处理序列，这使其成为涉及冗长数据序列的任务的完美选择。

Mamba 的设计理念基于对现代硬件功能的认识。它旨在充分利用 GPU 计算能力，保证：

**优化的内存使用：**通过设计 Mamba 的状态扩展以适应 GPU 的高带宽内存（HBM），可以缩短数据传输时间并加快处理速度。

**最大化并行处理：**Mamba 通过协调其计算与 GPU 计算的并行性质，达到了为序列模型建立新基准的性能水平。

Mamba VS Transformer

GPT-4 等 Transformer 的引入进入了自然语言处理（NLP）领域，并为多项自然语言任务建立了基准。长期以来，较长的序列一直是变压器的眼中钉，因为它们会严重阻碍其效率。

这个缺陷正是 Mamba 擅长的地方。也就是说，mamba 可以比变压器更快地处理冗长的序列，并且由于其独特的架构而更加简单。

Transformer 架构

Transformer 非常擅长处理数据序列，例如语言模型的文本。它们同时处理完整的序列，这与早期按顺序处理数据的模型不同。这种固有的功能使他们能够捕获数据中错综复杂的关系。它们使用注意力机制，使模型能够在生成预测时专注于各种序列段。使用三组权重来计算此关注度：从输入数据获取的值、键和查询。

序列中的每个元素都相对于其他每个元素进行加权，以指示它应该有多少权重（或 “注意力”）来预测序列中的下一个元素。Transformer 由两个主要块组成：创建输出的 decoder 和处理输入数据的 encoder。编码器由几层组成 - 每层都有两个子层：一个基本的、按位置的、完全连接的前馈网络和一个多头自注意力机制。为了帮助训练深度网络，每个子层都使用残差连接和归一化。

与编码器一样，解码器由两层和两个子层组成，但它也增加了第三个子层，用于处理编码器输出上的多头注意力。由于解码器的 sequential 性质，解码器的 autoregressive 属性被保留下来，这将对位置的预测限制为仅考虑较早的位置。

因此，Transformers 试图通过利用更复杂的注意力过程来解决冗长序列的问题，但 Mamba 采取了不同的方法。

Mamba 架构

Mamba 利用选择性状态空间。这种方法解决了Transformers 在长序列下计算效率低下的问题。Mamba 的架构使更快的推理和线性序列长度扩展成为可能，为序列建模创造了一种新的范式，随着序列越来越长，这种范式可能会被证明会更加有效。由于我们在上面深入探讨了 Mamba 的架构，因此我们不会在这里深入讨论。

特征	变压器	曼巴
建筑	基于注意力	基于 SSM
复杂性	高	降低
推理速度	O（n）	O（1）
训练速度	O（注2）	O（n）

值得注意的是，尽管 SSM 与 Transformer 相比具有许多优势，但后者可以处理比 SSM 存储在内存中更长的序列，需要更少的数据来学习类似的任务，并且在需要从输入上下文检索或复制的任务中优于 SSM，即使参数更少。

开始使用 Mamba

如果您有兴趣使用 Mamba 或在项目中利用它，则必须具备以下条件：

Linux的
英伟达图形处理器
PyTorch 1.12+ 版本
CUDA 11.6+ 的

要从 Mamba 存储库安装所需的软件包，请使用一些简单的 pip 说明：

[Option] ：在 Mamba 块内部使用的简单因果 Conv1d 层的有效实现。pip install causal-conv1d>=1.2.0
pip install mamba-ssm：核心 Mamba 包。

它也可以通过 from this repository从源代码构建。pip install .

如果 PyTorch 版本导致兼容性问题，可以使用与 switch 一起使用来帮助。这些模型是在 Pile 和 SlimPajama 数据集等大型数据集上训练的，旨在满足各种计算要求和性能基准。pip``--no-build-isolation

Mamba 模型具有多个接口级别，但主模块是包装选择性 SSM 的 Mamba 架构块。

# Source: Mamba Repository
import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape

Mamba 的应用

Mamba LLM 的推出是 LLM 架构领域的一个重大潜在转变。Mamba 更快、更高效且可扩展，可以毫不费力地以高性能标准处理长序列，这解释了为什么它将在塑造复杂 AI 系统的未来方面发挥关键作用。

也就是说，下一波 AI 创新可能由 Mamba 的有效性和性能带来，它为创建越来越复杂的模型和应用程序铺平了道路。它的潜在影响力是巨大的，包括音频和语音处理应用程序、长篇文本分析、内容创建、实时语言翻译等。

这可能会带来革命性的行业包括：

医疗： Mamba 可以通过快速分析遗传数据来加快开发个性化健康药物的过程。

**金融：**可以部署 Mamba 来分析长期市场趋势，从而获得更准确的库存预测。

顾客服务： Mamba 能够为监控长篇讨论的聊天机器人提供支持，从而改善客户沟通

结论

Mamba 不仅为当前的序列模型提供了增量改进;它重新定义了什么是可能的。随着它的推出，人工智能的历史将翻开新的篇章，计算效率低下和序列长度限制最终变得过时。在过去的几年里，我们看到了 AI 从 RNN 到Transformers，再到现在的 Mamba，每一步都离实现能够与人类相媲美的 AI 进行深度思考和信息处理更近一步。Mamba 体现了革命性的精神，通过其选定的状态空间方法和线性时间缩放推动 AI 领域向前发展。Mamba 标志着人工智能前景广阔的开始。这是一种为未来设计的范式，并将以其无限的潜力对 AI 产生重大影响。