Auto Byte

专注未来出行及智能汽车科技

微信扫一扫获取更多资讯

Science AI

关注人工智能与其他前沿技术、基础学科的交叉研究与融合发展

微信扫一扫获取更多资讯

再战Transformer!原作者带队的Mamba 2来了,新架构训练效率大幅提升

自 2017 年被提出以来,Transformer 已经成为 AI 大模型的主流架构,一直稳居语言建模方面 C 位。

但随着模型规模的扩展和需要处理的序列不断变长,Transformer 的局限性也逐渐凸显。一个很明显的缺陷是:Transformer 模型中自注意力机制的计算量会随着上下文长度的增加呈平方级增长。

几个月前,Mamba 的出现打破了这一局面,它可以随上下文长度的增加实现线性扩展。随着 Mamba 的发布,这些状态空间模型 (SSM) 在中小型规模上已经实现了与 Transformers 匹敌,甚至超越 Transformers。

Mamba 的作者只有两位,一位是卡内基梅隆大学机器学习系助理教授 Albert Gu,另一位是 Together.AI 首席科学家、普林斯顿大学计算机科学助理教授 Tri Dao。

Mamba 面世之后的这段时间里,社区反应热烈。可惜的是,Mamba 的论文却惨遭 ICLR 拒稿,让一众研究者颇感意外。

仅仅六个月后,原作者带队,更强大的 Mamba 2 正式发布了。

图片

  • 论文地址:https://arxiv.org/pdf/2405.21060

  • GitHub 地址:https://github.com/state-spaces/mamba

  • 论文标题:Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

总体而言,本文提出了 SSD(state space duality)框架,基于此,研究者设计了一个新的体系架构 Mamba-2,其核心层是对 Mamba 的选择性 SSM 的改进,速度提高了 2-8 倍,同时在语言建模方面继续与 Transformers 竞争。

Tri Dao 表示,他们构建了一个丰富的 SSD 理论框架,许多线性注意力变体和 SSM 是等效的,由此产生的模型 Mamba-2 比 Mamba-1 更好、更快。

图片

Mamba-2 的新算法使其能够利用更大的状态维度 (16 → 256),同时训练速度更快。在需要更大状态容量的任务上,例如 MQAR 任务,它比 Mamba-1 有了显著的改进。

图片

此外研究者还发现,最近新出的混合模型(Jamba、Zamba)增加了一些注意力层来提高模型质量。基于这些发现,研究者将 4-6 个注意力层与 Mamba-2 层混合,其表现优于 Transformer++ 和纯 Mamba-2,因而得出注意力和 SSM 是互补的。

图片

这项研究的贡献概括为:

本文展示了状态空间模型与一类称为半可分矩阵的结构化矩阵族之间的等价性。这一联系是 Mamba-2 框架的核心,揭示了状态空间模型的新属性和算法。

本文显著改进了线性注意力理论,首先通过张量收缩的语言对其循环形式提供了一个明确的证明,然后将其推广到一种新的结构化掩码注意力(SMA)家族。 

本文将 SSM(状态空间模型)和 SMA(结构化掩码注意力)联系起来,显示它们有一个很大的交集,彼此是对偶的,同时具有 SSM 式的线性形式和类似注意力的二次方形式。本文还证明了任何具有快速循环形式的核注意方法都是 SSM。

除了内在的理论价值外,研究者所提出的框架为理解和改进序列模型开辟了广阔的方向。

在算法层面。所提框架为计算 SSM 提供了新的高效且易于实现的算法。本文提出了一种基于半可分离矩阵块分解的 SSD 算法,该算法利用了 SSM 线性递推和二次对偶形式,在所有主要效率轴上获得了最优的权衡。基于 SSD 的实现比 Mamba 的优化选择性扫描实现快 2 到 8 倍,同时允许使用更大的循环状态大小(是 Mamba 的 8 倍甚至更高,且几乎不影响速度)。SSD 与优化过的 softmax 注意力实现(FlashAttention-2)具有高度竞争力,在序列长度 2k 时性能相当,在序列长度 16K 时速度快 6 倍。 

架构设计。采用 SSM 等新架构的一个主要障碍是针对 Transformers 量身定制的生态系统,例如用于大规模训练的硬件高效优化和并行技术。本文框架允许使用已建立的惯例和技术来构建 SSM 的架构设计选择词汇表,并进一步改进它们。

本文还对 Mamba 块做了一些修改,这些修改允许实现张量并行,其主要思想包括引入分组值注意力 (GVA,grouped-value attention) 头结构。

将修改后的并行 Mamba 块与作为内部 SSM 层的 SSD 结合使用,形成了 Mamba-2 架构。研究者在与 Mamba 相同的设置中研究了 Mamba-2 的 Chinchilla 扩展法则,发现它在困惑度和实际运行时间方面均优于 Mamba 和 Transformer++。研究者还在 Pile 数据集上训练了一系列 Mamba-2 模型,结果显示 Mamba-2 在标准下游评估中匹配或超过 Mamba 和开源的 Transformers。例如,在 Pile 上训练了 3000 亿 token 的 2.7B 参数的 Mamba-2 在性能上超过了在同一数据集上训练的 2.8B 参数的 Mamba 和 Pythia 以及 6.9B 参数的 Pythia。 

系统优化:SSD 框架连接 SSM 和 transformer,允许利用为 transformer 开发的丰富的系统优化工作。

图片

SSD

Mamba-2 的核心贡献是新的 SSD(state space dual)层。SSD 层可以被定义为选择性 SSM 的特例。与 Mamba 相比,Mamba-2 的改动会略微降低表达能力,但却显著提高了训练效率,特别是允许在现代加速器上使用矩阵乘法单元。

图片

图片

SSD 层的对偶注意力:

图片

除了最新的 SSD 层,研究者也对 Mamba 的神经网络架构做了一些小的改变,Mamba-2 架构如下所示。

图片

Mamba-2 在网络架构上的主要变化是从顺序生成变为并行生成 SSM 参数,并且 Mamba-2 更适合张量并行等扩展方法。

通过提供状态空间模型的显式矩阵变换形式,研究团队揭示了理解和使用它们的新方法。从计算的角度来看,任何计算状态空间模型前向传播的方法都可以看作是半可分离矩阵上的矩阵乘法算法。半可分离矩阵视角为 SSD 提供了一个视角,其中双重模式分别指的是线性时间半可分离矩阵乘法算法和二次时间朴素矩阵乘法。

图片

研究团队定义了结构化状态空间模型和结构化注意力,讨论了它们的属性,并表明它们都有二次算法和线性算法。

图片

图片

自最初的 Mamba 论文研究了合成任务 —— 如:合成复制和归纳 Head 以来,许多后续工作开始研究更难的关联回忆任务。由 Zoology 和 Based 系列工作引入的 MQAR(multi-query associative recall)任务已成为事实上的标准。

图片

通过运行一个比文献中通常报告的版本要难得多的任务,该团队发现 Mamba-2 明显优于 Mamba-1,而改善性能的一个原因是状态大小(比 Mamba-1 大约 16 倍)。

在这篇文章中,作者深入探讨了模型背后的理论。

从两个完全不同的角度推导出 SSD 的「对偶性」:

  • 一个从 SSM 的角度出发;

  • 另一个从注意力机制的角度出发。

SSD 框架提供了状态空间模型、注意力机制和结构化矩阵之间丰富的联系。

虽然 SSD 模型可以被视为框架内每个分支的具体实例,但 SSD 框架本身更加通用,为未来的工作开辟了许多方向。

图片

SSD 框架(红色,蓝色):状态空间模型(即半可分矩阵)和结构化掩码注意力机制包含了大量高效的序列模型。它们的交集是 SSD 模型(紫色)。

SSD 算法

通常,矩阵乘法(matmul)的 FLOPs 速度要比非矩阵乘法 FLOPs 快得多(高达 16 倍):A100 GPU 具有 312 TFLOPS 的 BF16 矩阵乘法性能,但只有 19 TFLOPS 的 FP32 算术性能,而 H100 具有 989 TFLOPS 的 BF16 矩阵乘法性能,但只有 67 TFLOPS 的 FP32 算术性能。

Mamba-2 的主要目标之一是「利用张量核心加速 SSM」。 

在绑定参数并引入 Head 结构后,Mamba-1 中的 SSM 变成了 SSD,这是一种更具限制性的形式,具有类似注意力的公式。并且由于 SSD 连接 SSM 和结构化矩阵,计算 SSM 的高效算法直接对应于「token-mixing」或「sequence-mixing」矩阵 M 的不同分解。

图片

因此,可以通过寻找替代的矩阵乘法方式,例如通过各种方式对其进行分解,从而创建计算 SSM 的新算法。

通过精心选择块大小,对这个矩阵进行简单块分解,就可以集 SSD 线性递归和二次注意力对偶形式的两种优势于一身。

而这也就是 SSD 算法的起源,它有 4 个步骤,并且对于这个算法有两种完全不同的诠释。

SSD 算法:分块矩阵分解

首先将半可分 SSM 矩阵划分为大小为 Q×Q 的块,然后,利用半分矩阵的性质来分解每个低秩的非对角块: 

  1. (橙色)每个对角块是一个更小的半可分矩阵,可以以喜欢的方式计算这个乘法,特别是使用 SSD 的二次(类似注意力机制)形式。 

  2. (绿色)总共有 T/Q 个不同的绿色块,通过批处理矩阵乘法来计算。

  3.  (黄色)注意,黄色项本身是一个 1 - 半可分矩阵,这一步等价于对某些修改后的 A 因子的 SSM 扫描。

  4. (蓝色)与绿色类似,通过批处理矩阵乘法来计算。

SSD 算法:分块和状态传递

该算法的另一种诠释涉及「推理 SSM 如何在实际序列上进行操作」。

首先将输入序列分割成大小为 Q 的块,步骤可以分为:

  1. 分块内部输出:计算每个块的局部输出(假设初始状态(对于块)为 0,则每个块的输出是多少?)

  2. 块状态:计算每个块的最终状态(假设初始状态(对于块)为 0,则每个块的最终状态是多少?)

  3. 传递状态:计算所有块的最终状态的递归 - 使用任何所需的算法,例如并行或顺序扫描(考虑到所有先前输入,每个块的实际最终状态是多少?)

  4. 输出状态:对于每个块,根据其真实的初始状态(在步骤 3 中计算),仅从初始状态得出的输出计算贡献

可以看到,大部分算法(步骤 1、2 和 4)利用了矩阵乘法(因此利用了张量核心),而且可以并行计算。

只有步骤 3 需要扫描,但它只操作一个非常短的序列,通常只需要很少时间。

系统及扩展优化

张量并行

图片

使用张量并行对 Mamba-1 进行大规模训练的一项困难是,每层都需要 2 次 all-reduce,而在 Transformer 中,每个注意力或 MLP 层只需 1 次 all-reduce。这是因为 SSM 的一些参数是内部激活的函数,而不是层的输入函数。在 Mamba-2 中,由于采用了「并行投影」结构,所有 SSM 参数都是层输入的函数,因此可以轻松地将张量并行应用于输入投影:将输入投影和输出投影矩阵分割成 2、4、8 个碎片,具体取决于张量并行度。研究者使用 grouped norm,分组数除以张量并行度,这样每个 GPU 都能单独完成归一化。这些变化导致每层只需 1 次 all-reduce,而不是 2 次。

序列并行

图片

在对超长序列进行训练时,可能需要沿着序列长度进行分割,并将不同部分分配给不同的设备。序列并行主要有两种形式:对于残差和归一化操作,用 reduce-scatter、残差 + 归一化、然后 all-gather,取代张量并行中的 all-reduce。由于 Mamba-2 使用与 Transformer 相同的残差和归一化结构,因此这种形式的序列并行无需修改即可直接应用。对于注意力或 SSM 操作,又称上下文并行(CP)。对于注意力,可以使用环形注意力沿序列维度进行分割。对于 Mamba-2,SSD 框架再次提供了帮助:使用相同的蒯分解,可以让每个 GPU 计算其本地输出和最终状态,然后在更新每个 GPU 的最终输出之前,在 GPU 之间传递状态(使用发送 / 接收通信原语)。

实验结果

该研究在 MQAR 的一种具有挑战性的版本上,使用更难的任务、更长的序列和更小的模型进行了对比实验。基线包括标准的多头 softmax 注意力以及 Based 架构,实验结果如图 8 所示。

图片

下表显示了 Mamba-2 在一系列下游零样本评估任务上的性能:

图片

感兴趣的读者可以阅读论文原文,了解更多研究内容。

产业TransformerMamba 2
相关数据
机器学习技术

机器学习是人工智能的一个分支,是一门多领域交叉学科,涉及概率论、统计学、逼近论、凸分析、计算复杂性理论等多门学科。机器学习理论主要是设计和分析一些让计算机可以自动“学习”的算法。因为学习算法中涉及了大量的统计学理论,机器学习与推断统计学联系尤为密切,也被称为统计学习理论。算法设计方面,机器学习理论关注可以实现的,行之有效的学习算法。

自注意力技术

自注意力(Self-attention),有时也称为内部注意力,它是一种涉及单序列不同位置的注意力机制,并能计算序列的表征。自注意力在多种任务中都有非常成功的应用,例如阅读理解、摘要概括、文字蕴含和语句表征等。自注意力这种在序列内部执行 Attention 的方法可以视为搜索序列内部的隐藏关系,这种内部关系对于翻译以及序列任务的性能非常重要。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

SSD技术

一种计算机视觉模型。论文发表于 2015 年(Wei Liu et al.)

注意力机制技术

我们可以粗略地把神经注意机制类比成一个可以专注于输入内容的某一子集(或特征)的神经网络. 注意力机制最早是由 DeepMind 为图像分类提出的,这让「神经网络在执行预测任务时可以更多关注输入中的相关部分,更少关注不相关的部分」。当解码器生成一个用于构成目标句子的词时,源句子中仅有少部分是相关的;因此,可以应用一个基于内容的注意力机制来根据源句子动态地生成一个(加权的)语境向量(context vector), 然后网络会根据这个语境向量而不是某个固定长度的向量来预测词。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

神经网络技术

(人工)神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型,那时候研究者构想了「感知器(perceptron)」的想法。这一领域的研究者通常被称为「联结主义者(Connectionist)」,因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型,它们都是前馈神经网络:卷积神经网络(CNN)和循环神经网络(RNN),其中 RNN 又包含长短期记忆(LSTM)、门控循环单元(GRU)等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习,但也有一些为无监督学习设计的变体,比如自动编码器和生成对抗网络(GAN)。

矩阵分解技术

矩阵分解是一种将矩阵简化为其组成部分的方法。这种方法可以简化更复杂的矩阵运算,这些运算可以在分解的矩阵上执行,而不是在原始矩阵本身上执行。它的衍生Non-negative matrix factorization也被用于降维等操作上。

算术技术

算术(英语:arithmetic)是数学最古老且最简单的一个分支,几乎被每个人使用着,从日常生活上简单的算数到高深的科学及工商业计算都会用到。一般而言,算术这一词指的是记录数字某些运算基本性质的数学分支。

分块技术

将标注好词性的句子按句法结构把某些词聚合在一起形成比如主语、谓语、宾语等等。

推荐文章
暂无评论
暂无评论~