人工智能
Mamba：重新定义序列建模和超越Transformer架构
在这篇关于Mamba的文章中，我们将探讨这款创新性的状态空间模型（SSM）如何革命性地改变序列建模。由Albert Gu和Tri Dao开发，Mamba以其在语言处理、基因组学和音频分析等领域处理复杂序列的效率而著称。其线性时间序列建模和选择性状态空间确保了在这些多样化的模态中获得卓越的性能。
我们将深入探讨Mamba如何克服传统Transformer面临的计算挑战，特别是在处理长序列时。其选择性状态空间模型允许更快的推理和线性扩展，显著提高了吞吐量。
Mamba的独特之处在于其快速处理能力、选择性SSM层和受FlashAttention启发的硬件友好设计。这些特点使Mamba能够超越许多现有的模型，包括基于Transformer的模型，使其成为机器学习领域的一项值得注意的进步。
Transformer与Mamba
Transformer，如GPT-4，在自然语言处理中设定了基准。然而，随着序列长度的增加，其效率会下降。这就是Mamba领先的地方，它能够更高效地处理长序列，并且其独特的架构简化了整个过程。
Transformer擅长处理数据序列，例如语言模型中的文本。与之前的模型不同，Transformer可以同时处理整个序列，允许它们捕获数据中的复杂关系。
它们使用注意力机制，这使得模型可以在预测时关注序列的不同部分。
这种注意力是使用三个权重集计算的：查询、键和值，这些权重来自输入数据。序列中的每个元素都与其他元素进行比较，提供一个权重，表示每个元素在预测序列中的下一个元素时应该获得的重要性或“注意力”。
Transformer维护两个主要块：编码器，处理输入数据，和解码器，生成输出。编码器由多个层组成，每个层包含两个子层：多头自注意力机制和简单的位置感知全连接前馈网络。归一化和残差连接用于每个子层，以帮助训练深层网络。
解码器也具有类似的层结构，但添加了一个执行多头注意力机制的第三个子层，用于编码器的输出。解码器的顺序性确保预测只能考虑之前的位置，保持自回归性质。
相比之下，Mamba模型采取了不同的方法。虽然Transformer通过使用更复杂的注意力机制来处理长序列的问题，但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） 这是基本的结构化状态空间模型。它接受一个序列 x 并使用学习到的参数 A、B、C 和延迟参数 Δ 生成输出 y。转换涉及将连续参数离散化（将连续函数转换为离散函数）并应用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在每个评估结果中都是最好的
在性能方面，Mamba在推理速度和准确性方面都表现出色。其设计使得它能够更好地利用更长的上下文，这在DNA和音频建模中得到了体现，超越了之前的模型在需要长距离依赖的复杂任务中的表现。其多功能性也体现在零样本评估中，横跨多个任务，设定了新的标准，体现出其在效率和可扩展性方面的优势。
开始使用Mamba
对于那些有兴趣利用Mamba的人来说，技术要求包括Linux操作系统、NVIDIA GPU、PyTorch 1.12+和CUDA 11.6+。安装涉及使用pip命令安装Mamba仓库中的必要包。如果出现与PyTorch版本的兼容性问题，使用–no-build-isolation标志可以帮助解决。这些模型是在大型数据集（如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和深度学习领域研究人员和开发人员的一种可访问且强大的工具。
