• 深入理解Mamba是如何解决Transformer在处理长序列时遇到的计算挑战
  • 发布于 1个月前
  • 246 热度
    0 评论
  • 远行de风
  • 0 粉丝 23 篇博客
  •   
在这篇关于 Mamba 的文章中,我们来探索这个创新的状态空间模型(state-space model,SSM)如何在序列建模领域带来革命性的变革。Mamba 是由 Albert Gu 和 Tri Dao 开发的,因其在语言处理、基因组学、音频分析等领域的复杂序列时表现出色而受到关注。Mamba 采用线性时间序列建模和选择性状态空间,因此在这些不同的应用领域都展现出了卓越的性能。

论文链接:https://arxiv.org/abs/2312.00752

本文,我们将深入了解 Mamba 是如何解决传统 Transformer 在处理长序列时遇到的计算挑战的。Mamba 采用在状态空间模型中的选择性方法,这不仅加快了推理速度,还实现了与序列长度的线性缩放,显著提高了处理能力。

Mamba 的特别之处在于它的快速处理能力、选择性的 SSM 层,以及受 FlashAttention 启发的硬件友好设计。这些特点使得 Mamba 在许多现有模型中脱颖而出,包括那些基于 Transformer 架构的模型,这使得 Mamba 成为机器学习领域的一个重要进步。

一.Transformer vs Mamba
Transformer,如 GPT-4,已在自然语言处理领域设定了高标准。但是,它们在处理较长的序列时效率会下降。这就是 Mamba 的优势,它不仅能更高效地处理长序列,而且其独特的架构还简化了整个处理过程。Transformer 擅长处理数据序列,如文本,用于语言模型。与之前按顺序处理数据的模型不同,Transformer 能够同时处理整个序列,捕获数据内部的复杂关系。

它们利用注意力机制,使模型在做出预测时能够关注序列的不同部分。这种注意力是通过计算输入数据中的三组权重:查询、键和值来实现的。序列中的每个元素都会与其他元素进行比较,从而得到一个权重,这个权重表示在预测序列中下一个元素时每个元素应得到的 “注意力”。

Transformer 由两个主要部分组成:编码器用于处理输入数据,解码器用于生成输出。编码器包含多个层,每层包含两个子层:一个多头自我注意力机制和一个简单的、逐位置的全连接前馈网络。每个子层都使用规范化和残差连接来帮助训练深度网络。

解码器也包含类似于编码器的层,但增加了一个对编码器输出执行多头注意力的子层。解码器的顺序性质保证了对一个位置的预测只能基于之前的位置,保持了自回归的特性。

相比之下,Mamba 采取了不同的方法。尽管 Transformer 通过使用更复杂的注意力机制来解决长序列的问题,但 Mamba 采用选择性状态空间,提供了一个更加高效的解决方案。

下面是一个关于 Transformer 工作原理的高层次概述:
输入处理:Transformer 首先将输入数据转换为模型能够理解的格式,通常是通过位置信息丰富的嵌入表示。
注意力机制:核心在于注意力机制,它计算一个得分来决定在理解当前元素时,应该将多少关注放在输入序列的其他部分。
编码器 - 解码器架构:Transformer 由一个编码器组成,用来处理输入,以及一个解码器,用来生成输出。每个部分都包含多层,以提炼模型对输入的理解。
多头注意力:编码器和解码器中的多头注意力允许模型同时关注序列的不同部分,这提高了它从多样化上下文中学习的能力。
逐位置前馈网络:经过注意力处理后,一个简单的神经网络独立且统一地处理每个位置的输出,通过残差连接与输入结合,并进行层规范化。
输出生成:接着解码器预测一个输出序列,这个过程受到编码器上下文的影响以及它到目前为止生成的内容。

Transformer 凭借其并行处理序列的能力和强大的注意力机制,在翻译和文本生成等任务中表现出色。

Mamba 模型则采用不同的方式运作,它使用选择性状态空间来处理序列。这种方法解决了 Transformer 在处理长序列时的计算效率问题。Mamba 的设计实现了更快的推理速度,并且随着序列长度的线性缩放,为序列建模树立了一个新的范式,特别适用于处理日益增长的序列。

二.Mamba 独特之处

Mamba 最独特的地方在于它脱离了传统的注意力和 MLP(多层感知器)模块。这种简化使得模型更轻便、更快速,并且能随着序列长度线性地缩放,这是之前模型无法做到的。


Mamba 的关键特性包括:
选择性 SSM:通过选择性 SSM,Mamba 能够过滤掉不相关信息,专注于重要数据,从而更高效地处理序列。
硬件意识算法:Mamba 采用了针对现代硬件(尤其是 GPU)优化的并行算法,这提高了计算速度,相比传统模型减少了内存需求。
简化架构:Mamba 通过整合选择性 SSM 并去除注意力和 MLP 模块,提供了一个更简单、更统一的结构,从而带来更好的可扩展性和性能。

在语言、音频和基因组学等多个领域,Mamba 展示了优越的性能,无论是在预训练还是特定领域任务中都有出色表现。例如,在语言建模方面,Mamba 的性能可媲美或超越更大的 Transformer 模型。

Mamba 的代码和预训练模型可以在 GitHub 上公开获取,供社区使用。
GitHub:https://github.com/state-spaces/mamba

对于线性模型来说,标准复制任务很简单。选择性复印和感应头需要用于 LLM 的动态、内容感知存储器。
结构化状态空间(Structured State Space,S4)模型最近成为了序列模型领域的一种有前景的新类别,它融合了循环神经网络(RNN)、卷积神经网络(CNN)和传统状态空间模型的特性。S4 模型借鉴了连续系统的思想,特别是那些将一维函数或序列通过隐含状态进行映射的系统。在深度学习领域,S4 模型代表了一项重大创新,为设计高效且适应性强的序列模型提供了全新的方法。

S4 模型的动态
基本的结构化状态空间模型(S4)接受一个序列 x,并通过学习到的参数 A、B、C 以及一个延迟参数 Δ 来生成输出 y。这个转换过程包括对这些参数进行离散化(即将连续的函数转换成离散的函数)并应用 SSM 操作,这个操作是时间不变的,也就是说在不同的时间步长上是不会改变的。

离散化的重要性
离散化是一个关键过程,它通过一些固定的公式将连续的参数转换成离散的形式,使得 S4 模型能够保持与连续时间系统的联系。这赋予了模型一些额外的特性,如分辨率不变性,并确保了适当的规范化,从而增强了模型的稳定性和性能。离散化的过程也类似于 RNN 中用于管理信息流的门控机制。

线性时间不变性(LTI)
S4 模型的核心特性之一是它们的线性时间不变性。这意味着模型的动态在时间上保持一致,其参数在所有时间步长上都是固定的。LTI 是递归和卷积的基础,为构建序列模型提供了一个简单但强大的框架。

克服基本限制
传统上,S4 框架由于其线性时间不变的本质而受到限制,在需要适应性动态的数据建模方面面临挑战。最近的研究论文提出了一种方法,通过引入时变参数来克服这些限制,从而摆脱了 LTI 的束缚。这使得 S4 模型能够处理更多样化的序列和任务,显著扩大了其适用范围。

“状态空间模型” 这一术语广泛用于描述涉及潜在状态的任何递归过程,在多个学科中有不同的含义。在深度学习的背景下,S4 模型或结构化 SSM 指的是一类专门为高效计算而优化的模型,同时它们也能够处理复杂的序列。

S4 模型可以被集成到端到端的神经网络架构中,作为独立的序列转换组件。它们在某种程度上类似于 CNN 中的卷积层,为各种神经网络架构中的序列建模提供了基础支撑。


三.序列建模中选择性的重要性

这篇论文指出,序列建模的一个关键方面是将上下文信息压缩成可管理的状态。能够选择性关注或过滤输入的模型,更有效地维护这种压缩状态,从而实现更高效、强大的序列模型。这种选择性对于模型在处理复杂的语言建模等任务时,适应性地控制信息沿序列的流动至关重要。

Selective SSM 提升了传统 SSM 的能力,使其参数可以依赖于输入,引入了一种以前时间不变模型无法实现的适应性。这导致了时变 SSM 的出现,它们不再使用卷积进行高效计算,而是依赖线性递归机制,这与传统模型有显著不同。

SSM + Selection (S6) 这种变体包含一个选择机制,使参数 B 和 C 以及延迟参数 Δ 依赖于输入。这使得模型能够选择性地关注输入序列 x 的特定部分。参数在考虑选择的情况下被离散化,并使用扫描操作以时变方式应用 SSM,顺序处理元素,并随时间动态调整关注点。

四.Mamba 性能亮点

就性能而言,Mamba 在推理速度和准确性方面表现出色。它的设计使得模型更好地利用更长的上下文,这在 DNA 和音频建模中得到了证明,超越了以前的模型在处理需要长期依赖的复杂任务上的表现。它还在跨多个任务的零样本评估中表现出多功能性,为此类模型在效率和可扩展性方面设定了新标准。

五.开始使用 Mamba
对于那些有意使用 Mamba 的用户,需要的技术条件包括 Linux 操作系统、NVIDIA GPU、PyTorch 1.12+ 和 CUDA 11.6+。安装过程包括使用 pip 命令从 Mamba 仓库安装所需的包。如果出现与 PyTorch 版本的兼容性问题,可以使用 pip 的 –no-build-isolation 选项来解决。这些模型经过了大型数据集如 Pile 和 SlimPajama 的训练,旨在满足多样化的计算需求和性能标准。
GitHub:https://github.com/state-spaces/mamba

Mamba 提供了从选择性 SSM 层到 Mamba 块再到完整语言模型结构的不同级别的接口。Mamba 块作为该架构的主要模块,利用了因果 Conv1d 层,并可轻松集成到神经网络设计中。提供的 Python 示例展示了如何实例化一个 Mamba 模型并通过它处理数据,突出了系统的简单性和灵活性。

在 Hugging Face 上,有多种规模的预训练 Mamba 模型可用,参数范围从 130M 到 2.8B,这些模型在 Pile 和 SlimPajama 数据集上接受了训练。这些模型旨在满足多样化的计算和性能需求,遵循 GPT-3 的规模标准。用户可以期望从这些模型中获得高吞吐量和准确性,使 Mamba 成为多种应用的强有力选择,包括但不限于语言建模。

Hugging Face:https://huggingface.co/papers/2312.00752


六.Mamba 的影响
Mamba 代表了序列建模领域的一次飞跃,为处理信息密集型数据提供了一种强大的替代 Transformer 架构的方案。其设计符合现代硬件的需求,优化了内存使用和并行处理能力。Mamba 的代码库和预训练模型的开源可用性使其成为 AI 和深度学习领域研究人员和开发人员的一个易于接入且强大的工具。
原文链接:https://www.unite.ai/mamba-redefining-sequence-modeling-and-outforming-transformers-architecture/
用户评论