Auto Byte

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

微信扫一扫获取更多资讯

Science AI

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

微信扫一扫获取更多资讯

比Transformer更好,无Attention、MLPs的BERT、GPT反而更强了

本文探索了 Monarch Mixer (M2) ,这是一种在序列长度和模型维度上都是次二次的新架构,并且在现代加速器上具有很高的硬件效率。

从 BERT、GPT 和 Flan-T5 等语言模型到 SAM 和 Stable Diffusion 等图像模型,Transformer 正以锐不可当之势席卷这个世界,但人们也不禁会问:Transformer 是唯一选择吗?

斯坦福大学和纽约州立大学布法罗分校的一个研究团队不仅为这一问题给出了否定答案,而且还提出了一种新的替代技术:Monarch Mixer。近日,该团队在 arXiv 公布了相关论文和一些检查点模型及训练代码。顺带一提,该论文已入选 NeurIPS 2023 并获得 Oral Presentation 资格。

图片

论文地址:https://arxiv.org/abs/2310.12109

代码地址:https://github.com/HazyResearch/m2

该方法去掉了 Transformer 中高成本的注意力和 MLP,代之以富有表现力的 Monarch 矩阵,使之在语言和图像实验中以更低的成本取得了更优的表现。

这并不是斯坦福大学第一次提出 Transformer 的替代技术。今年六月该校的另一个团队还曾提出过一种名为 Backpack 的技术,参阅机器之心文章《斯坦福训练 Transformer 替代模型:1.7 亿参数,能除偏、可控可解释性强》。当然,这些技术要取得真正的成功,还需要研究社区的进一步检验并在应用开发者手中变成切实好用的产品。

下面我们看看这篇论文中对 Monarch Mixer 的介绍以及一些实验结果。

论文介绍

自然语言处理计算机视觉领域,机器学习模型已能处理更长的序列和更高维度的表征,从而支持更长的上下文和更高的质量。然而,现有架构的时间和空间复杂性在序列长度和 / 或模型维度上呈二次增长模式,这会限制上下文长度并提升扩展成本。举个例子,Transformer 中的注意力和 MLP 会随序列长度和模型维度呈二次扩展模式。

针对这一问题,斯坦福大学和纽约州立大学布法罗分校的这个研究团队声称找到了一种高性能的架构,其复杂度随序列长度和模型维度的增长是次二次的(sub-quadratic)。

他们的研究灵感来自 MLP-mixer 和 ConvMixer;这两项研究观察到:许多机器学习模型的运作方式都是沿序列和模型维度轴对信息进行混合,并且它们往往对两个轴使用了单个算子。

寻找表现力强、次二次且硬件效率高的混合算子的难度很大。举个例子,MLP-mixer 中的 MLP 和 ConvMixer 中的卷积都颇具表现力,但它们都会随输入维度二次扩展。近期有一些研究提出了一些次二次的序列混合方法,这些方法使用了较长的卷积或状态空间模型,而且它们都会用到 FFT,但这些模型的 FLOP 利用率很低并且在模型维度方面依然是二次扩展。与此同时,不损质量的稀疏密集 MLP 层方面也有一些颇具潜力的进展,但由于硬件利用率较低,某些模型实际上可能还比密集模型更慢。

基于这些灵感,这个研究团队提出了 Monarch Mixer (M2),其使用到了一类富有表现力的次二次结构化矩阵:Monarch 矩阵。

Monarch 矩阵是一类泛化了快速傅立叶变换(FFT)的结构化矩阵,并且研究表明其涵盖了范围广泛的线性变换,包括哈达玛变换、托普利兹矩阵、AFDF 矩阵和卷积。它们可通过分块对角矩阵的积进行参数化,这些参数被称为 Monarch 因子,与排列交织。

它们的计算是次二次扩展的:如果将因子的数量设为 p,则当输入长度为 N 时,计算复杂度为 图片,从而让计算复杂度可以位于 p = log N 时的 O (N log N) 与 p = 2 时的 图片之间。

M2 使用了 Monarch 矩阵来沿序列和模型维度轴混合信息。这种方法不仅易于实现,而且硬件效率也很高:使用支持 GEMM(广义矩阵乘法算法)的现代硬件就能高效地计算分块对角 Monarch 因子。

图片

该研究团队实现了一个 M2 层来进行概念验证 —— 完全使用 PyTorch 编写,代码行数不到 40(包括 import 软件包),而且其只需依赖矩阵乘法、转置、reshape 和逐元素乘积(见图 1 中部的伪代码);结果,对于大小为 64k 的输入,这些代码在一台 A100 GPU 上实现了 25.6% 的 FLOP 利用率。在 RTX 4090 等更新的架构上,对于同样大小的输入,一个简单的 CUDA 实现就能实现 41.4% 的 FLOP 利用率。

图片

有关 Monarch Mixer 的更多数学描述和理论分析请参看原论文。

实验

该研究团队在 Transformer 已占主导地位的三个任务上对 Monarch Mixer 和 Transformer 进行了比较:BERT 风格的非因果掩码语言建模任务、ViT 风格的图像分类任务、GPT 风格的因果语言建模任务。

在每个任务上,实验结果表明新提出的方法在不使用注意力和 MLP 的前提下均能达到与 Transformer 相媲美的水平。他们还在 BERT 设置中评估了新方法相较于强大 Transformer 基准模型的加速情况。

非因果语言建模

对于非因果语言建模任务,该团队构建了一种基于 M2 的架构:M2-BERT。M2-BERT 可以直接替代 BERT 风格的语言模型,而 BERT 是 Transformer 架构的一大主力应用。对于 M2-BERT 的训练,使用了在 C4 上的掩码语言建模,token 化器则是 bert-base-uncased。

M2-BERT 基于 Transformer 骨干,但其中的注意力层和 MLP 被 M2 层替换,如图 3 所示。

图片

在序列混合器中,注意力被带残差卷积的双向门控卷积替代(见图 3 左侧)。为了恢复卷积,该团队将 Monarch 矩阵设置为 DFT 和逆 DFT 矩阵。他们还在投射步骤之后添加了逐深度的卷积。

在维度混合器中,MLP 中两个密集矩阵被替换成了学习得到的分块对角矩阵(1 阶 Monarch 矩阵,b = 4)。

研究者预训练了 4 个 M2-BERT 模型:其中两个是大小分别为 80M 和 110M 的 M2-BERT-base 模型,另外两个是大小分别为 260M 和 341M 的 M2-BERT-large 模型。它们分别相当于 BERT-base 和 BERT-large。

表 3 给出了相当于 BERT-base 的模型的性能表现,表 4 给出了相当于 BERT-large 的模型的性能表现。

图片

从表中可以看到,在 GLUE 基准上,M2-BERT-base 的表现可以媲美 BERT-base,同时参数还少了 27%;而当两者参数数量相当时,M2-BERT-base 胜过 BERT-base 1.3 分。类似地,参数少 24% 的 M2-BERT-large 与 BERT-large 表现相当,而参数数量一样时,M2-BERT-large 有 0.7 分的优势。

表 5 给出了相当于 BERT-base 的模型的前向吞吐量情况。其中报告的是在 A100-40GB GPU 上每毫秒处理的 token 数,这能反映推理时间。

图片

可以看到,M2-BERT-base 的吞吐量甚至超过了经过高度优化的 BERT 模型;相较于在 4k 序列长度上的标准 HuggingFace 实现,M2-BERT-base 的吞吐量可达其 9.1 倍!

表 6 则报告了 M2-BERT-base (80M) 和 BERT-base 的 CPU 推理时间 —— 结果是直接运行这两个模型的 PyTorch 实现得到的。

图片

当序列较短时,数据局部性的影响依然主导着 FLOP 的减少情况,而过滤器生成(BERT 中没有)等操作的成本更高。而当序列长度超过 1K 时,M2-BERT-base 的加速优势就渐渐起来了,当序列长度达 8K 时,速度优势可达 6.5 倍。

图像分类

在非因果建模方面,为了验证新方法在图像上也有在语言上一样的优势,该团队还评估了 M2 在图像分类任务上的表现。

表 7 给出了 Monarch Mixer、ViT-b、HyenaViT-b 和 ViT-b-Monarch(用 Monarch 矩阵替换了标准 ViT-b 中的 MLP 模块)在 ImageNet-1k 上的性能表现。

图片

Monarch Mixer 优势非常明显:只需一半的参数量,其表现就能胜过原始 ViT-b 模型。而更让人惊讶的是,参数更少的 Monarch Mixer 很能胜过 ResNet-152;要知道,ResNet-152 可是专门针对 ImageNet 任务设计的。

因果语言建模

GPT 风格的因果语言建模是 Transformer 的一大关键应用。该团队为因果语言建模构建了一个基于 M2 的架构:M2-GPT。

对于序列混合器,M2-GPT 组合使用了来自 Hyena 的卷积过滤器、当前最佳的无注意力语言模型以及来自 H3 的跨多头参数共享。他们使用因果参数化替换了这些架构中的 FFT,并完全移除了 MLP 层。所得到的架构完全没有注意力,也完全没有 MLP。

他们在因果语言建模的标准数据集 PILE 上对 M2-GPT 进行了预训练。结果见表 8。

图片

可以看到,尽管基于新架构的模型完全没有注意力和 MLP,但其在预训练的困惑度指标上依然胜过 Transformer 和 Hyena。这些结果表明,与 Transformer 大不相同的模型也可能在因果语言建模取得出色表现。

了解更多内容,请参考原论文。

工程Monarch Mixer (M2)
相关数据
机器学习技术

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

基准技术

一种简单的模型或启发法,用作比较模型效果时的参考点。基准有助于模型开发者针对特定问题量化最低预期效果。

参数技术

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

伪代码技术

伪代码,又称为虚拟代码,是高层次描述算法的一种方法。它不是一种现实存在的编程语言;它可能综合使用多种编程语言的语法、保留字,甚至会用到自然语言。 它以编程语言的书写形式指明算法的职能。相比于程序语言它更类似自然语言。它是半形式化、不标准的语言。

计算机视觉技术

计算机视觉(CV)是指机器感知环境的能力。这一技术类别中的经典任务有图像形成、图像处理、图像提取和图像的三维推理。目标识别和面部识别也是很重要的研究领域。

自然语言处理技术

自然语言处理(英语:natural language processing,缩写作 NLP)是人工智能和语言学领域的分支学科。此领域探讨如何处理及运用自然语言;自然语言认知则是指让电脑“懂”人类的语言。自然语言生成系统把计算机数据转化为自然语言。自然语言理解系统把自然语言转化为计算机程序更易于处理的形式。

图像分类技术

图像分类,根据各自在图像信息中所反映的不同特征,把不同类别的目标区分开来的图像处理方法。它利用计算机对图像进行定量分析,把图像或图像中的每个像元或区域划归为若干个类别中的某一种,以代替人的视觉判读。

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用,如语音识别,机器翻译,词性标注,句法分析和资讯检索。由于字词与句子都是任意组合的长度,因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题),也使得在语料库中估算字串的机率变得很困难,这也是要使用近似的平滑n元语法(N-gram)模型之原因。

机器之心机构

机器之心,成立于2014年,是国内最具影响力、最专业、唯一用于国际品牌的人工智能信息服务与产业服务平台。目前机器之心已经建立起涵盖媒体、数据、活动、研究及咨询、线下物理空间于一体的业务体系,为各类人工智能从业者提供综合信息服务和产业服务。

https://www.jiqizhixin.com/
分块技术

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

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