人工智能
Mamba：重新定义序列建模和超越Transformer架构
在这篇关于 Mamba 的文章中，我们将探讨这个创新的状态空间模型（SSM）如何革命性地改变序列建模。由Albert Gu和Tri Dao开发的Mamba，以其在语言处理、基因组学和音频分析等领域处理复杂序列的效率而著称。其线性时间序列建模和选择性状态空间确保了在这些多样化的模态中获得卓越的性能。
我们将深入探讨Mamba如何克服传统Transformer面临的计算挑战，特别是在处理长序列时。其在状态空间模型中的选择性方法使得推理速度更快，且随着序列长度的增加而线性扩展，从而显著提高了吞吐量。
Mamba的独特之处在于其快速处理能力、选择性SSM层和受FlashAttention启发的硬件友好设计。这些特点使Mamba能够超越许多现有的模型，包括基于Transformer的模型，使其成为机器学习领域的一项值得注意的进步。
Transformer vs Mamba
Transformer，如GPT-4，在自然语言处理中设定了基准。然而，当处理更长的序列时，其效率会下降。这就是Mamba领先的地方，它能够更高效地处理长序列，并具有简化整个过程的独特架构。
Transformer擅长处理数据序列，例如语言模型中的文本。与之前的模型不同，Transformer可以同时处理整个序列，能够捕捉数据中的复杂关系。
它们使用注意力机制，使模型能够在预测时关注序列的不同部分。
这种注意力是使用三个权重集计算的：查询、键和值，这些权重是从输入数据中派生的。序列中的每个元素都与其他元素进行比较，提供一个权重，表示每个元素在预测序列中的下一个元素时应该获得的“注意力”。
Transformer维护两个主要块：编码器，处理输入数据，和解码器，生成输出。编码器由多个层组成，每个层包含两个子层：多头自注意力机制和一个简单的位置感知全连接前馈网络。归一化和残差连接在每个子层中使用，以帮助训练深度网络。
解码器也具有类似于编码器的层，但添加了一个执行多头注意力机制的第三个子层，用于编码器的输出。解码器的顺序性确保预测仅考虑之前的位置，保持自回归属性。
相比之下，Mamba模型采取了不同的方法。虽然Transformer通过使用更复杂的注意力机制来处理长序列的问题，但Mamba使用选择性状态空间，提供了一种更高效的序列建模方法。
Mamba的独特之处在于其快速处理能力、选择性SSM层和受FlashAttention启发的硬件友好设计。这些特点使Mamba能够超越许多现有的模型，包括基于Transformer的模型，使其成为机器学习领域的一项值得注意的进步。
Transformer vs Mamba
Transformer，如GPT-4，在自然语言处理中设定了基准。然而，当处理更长的序列时，其效率会下降。这就是Mamba领先的地方，它能够更高效地处理长序列，并具有简化整个过程的独特架构。
以下是Transformer工作原理的高级概述：
- 输入处理:Transformer首先将输入数据编码为模型可以理解的格式，通常使用包含每个元素在序列中的位置的嵌入。
- 注意力机制:在其核心，注意力机制计算一个分数，表示在理解当前元素时应该关注输入序列的其他部分的程度。
- 编码器-解码器架构:Transformer模型由编码器处理输入和解码器生成输出组成。每个都由多个层组成，这些层完善了模型对输入的理解。
- 多头注意力:在编码器和解码器中，多头注意力使模型能够同时从不同表示空间关注序列的不同部分，提高了其从多样化上下文中学习的能力。
- 位置感知前馈网络:在注意力之后，一个简单的神经网络独立地处理每个位置的输出。它通过残差连接与输入相结合，然后进行层归一化。
- 输出生成:解码器然后预测输出序列，受到编码器上下文和到目前为止生成的内容的影响。
Transformer能够并行处理序列及其强大的注意力机制使其在翻译和文本生成等任务中非常强大。
相比之下，Mamba模型通过使用选择性状态空间来处理序列。这种方法解决了Transformer在处理长序列时的计算效率问题。Mamba的设计使得推理速度更快，且随着序列长度的增加而线性扩展，建立了一个新的序列建模范式，这可能在序列变得越来越长时更加高效。
Mamba
Mamba真正独特之处在于其从传统注意力和MLP块中脱颖而出。这种简化导致了一个更轻、更快的模型，它随着序列长度的增加而线性扩展，这是其前辈无法匹敌的壮举。
Mamba的关键特点包括：
- 选择性SSM:这些使Mamba能够过滤不相关的信息并关注相关数据，增强其处理序列的能力。这种选择性对于高效的内容基于推理至关重要。
- 硬件感知算法:Mamba使用一个针对现代硬件（尤其是GPU）优化的并行算法。这一设计使得计算速度更快，相比传统模型，内存需求更低。
- 简化架构:通过集成选择性SSM并消除注意力和MLP块，Mamba提供了一个更简单、更同质的结构。这导致了更好的可扩展性和性能。
Mamba在包括语言、音频和基因组学在内的各个领域表现出色，在预训练和特定领域任务中都表现出色。例如，在语言建模中，Mamba的性能与更大的Transformer模型相匹配或超过了它们。
Mamba的代码和预训练模型在GitHub上公开可用。
标准复制任务对于线性模型来说很简单。选择性复制和归纳头需要动态、内容感知的内存用于LLM。
结构化状态空间（S4）模型近期作为一种有前途的序列模型类别出现，涵盖了RNN、CNN和经典状态空间模型的特征。S4模型从连续系统中汲取灵感，特别是一种将一维函数或序列映射到隐式潜在状态的系统。在深度学习的背景下，它们代表了一项重大创新，提供了一种新的方法来设计高效且高度适应性的序列模型。
S4模型的动态
SSM（S4）这是基本的结构化状态空间模型。它接受一个序列，并使用学习的参数、、和延迟参数生成一个输出。转换涉及将参数离散化（将连续函数转换为离散函数），并应用SSM操作，这是一个时间不变的操作——它不会随时间变化。
离散化的重要性
离散化是一个关键过程，它通过固定公式将连续参数转换为离散参数，使S4模型能够与连续时间系统保持联系。这赋予了模型额外的属性，例如分辨率不变性，并确保适当的归一化，从而提高了模型的稳定性和性能。离散化还与RNN中的门控机制相似，这对于管理网络中的信息流至关重要。
线性时间不变性（LTI）
S4模型的一个核心特征是其线性时间不变性。这意味着模型的动态在时间上保持一致，所有时间步的参数都保持不变。LTI是递归和卷积的基石，提供了一个简化但强大的框架，用于构建序列模型。
克服基本限制
S4框架传统上受到其LTI性质的限制，这限制了其在需要自适应动态的数据建模方面的能力。最近的研究论文提出了一种方法，通过引入时间变化的参数，克服了这些限制，从而消除了LTI的约束。这使得S4模型能够处理更广泛的序列和任务，显著扩大了其适用性。
术语“状态空间模型”广泛涵盖了任何涉及潜在状态的递归过程，并被用来描述各个学科中的各种概念。在深度学习的背景下，S4模型或结构化SSM指的是一类经过优化的模型，用于高效计算，同时保留了建模复杂序列的能力。
S4模型可以集成到端到端的神经网络架构中，作为独立的序列转换。它们可以被视为类似于CNN中的卷积层，为各种神经网络架构中的序列建模提供了基础。
SSM vs SSM + 选择
序列建模中选择性的动机
结构化SSM
该论文认为，序列建模的一个基本方面是将上下文压缩为可管理的状态。能够选择性地关注或过滤输入的模型提供了一种更有效的维护压缩状态的方法，从而导致更高效、更强大的序列模型。这种选择性对于模型在序列维度上自适应地控制信息流动至关重要，这是处理语言建模和其他任务的复杂任务所必需的能力。
选择性SSM增强了传统SSM，使其参数能够依赖于输入，从而引入了一种以前不可用于时间不变模型的自适应性。这导致了时间变化的SSM，它们不能使用卷积进行高效计算，而是依赖于线性递归机制，这是与传统模型的一个重大偏离。
性能亮点
Mamba在每个评估结果中都是最好的
在性能方面，Mamba在推理速度和准确率方面都表现出色。其设计使得上下文的利用更好，这在DNA和音频建模中得到了体现，超越了之前的模型在需要长距离依赖的复杂任务中的表现。其多功能性也体现在多任务零样本评估中，设定了此类模型在效率和可扩展性方面的新标准。
开始使用Mamba
对于那些有兴趣利用Mamba的人，技术要求包括Linux操作系统、NVIDIA GPU、PyTorch 1.12+和CUDA 11.6+。安装涉及使用pip命令安装Mamba存储库中的必要包。如果与PyTorch版本出现兼容性问题，使用–no-build-isolation标志与pip可以提供帮助。这些模型是在大量数据集（如Pile和SlimPajama数据集）上训练的，旨在满足各种计算需求和性能基准。
Mamba提供了多个级别的接口，从选择性SSM层到Mamba块和完整的语言模型结构。Mamba块是架构的主要模块，使用了因果Conv1d层，可以轻松集成到神经网络设计中。Python中的示例用法演示了实例化Mamba模型并处理数据的简单性和灵活性。
预训练的Mamba模型在Hugging Face上可用，参数数量从130M到2.8B不等，训练数据包括Pile和SlimPajama数据集。这些模型旨在满足多样化的计算和性能需求，遵循GPT-3的尺寸标准。用户可以期待这些模型具有高吞吐量和准确率，使Mamba成为各种应用（包括但不限于语言建模）的有力竞争者。
Mamba的影响
Mamba代表了序列建模领域的一项重大进步，提供了一个比传统Transformer架构更强大的替代方案，用于处理信息密集的数据。其设计符合现代硬件的需求，优化了内存使用和并行处理能力。Mamba代码库和预训练模型的开源可用性使其成为AI和深度学习领域研究人员和开发人员的强大工具。
