Monarch Mixer：一种性能比Transformer更强的网络架构-CSDN博客

六年前谷歌团队在arXiv上发表了革命性的论文《Attention is all you need》。作为一种优势的机器学习网络架构Transformer技术迅速席卷全球。Transformer一直是现代基础模型背后的主力架构并且在不同的应用程序中取得了令人印象深刻的成功包括像BERT、ChatGPT和Flan-T5这样的预训练语言模型到像SAM和stable diffusion这样的图像模型。

尽管如此Transformer架构中的自注意力机制和MLP在处理长度很长的序列或者维数很大的模型的时候速度和效率会打折扣。这主要是因为tranformer构架的时空复杂性随序列的长度和训练模型的维数按平方的依赖关系生长即所谓“二次元”quadratic。

最近斯坦福大学和纽约州立大学布法罗分校的一个研究团队在arXiv上发表题为《Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture》论文。该论文提出了一种新的transformer的替代技术Monarch MixerM2。该方法去掉了 Transformer 中高成本的自注意力和 MLP代之以富有表现力的 Monarch 矩阵使之在语言和图像实验中以更低的成本取得了更优的表现。其复杂度随序列长度和模型维度的增长是低于二次元即所谓次二次元的sub-quadratic。

斯坦福大学和布法罗大学在axXiv上发表论文截图

该论文已入选 NeurIPS 2023 并获得 Oral Presentation 资格。

算法与原理

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

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

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

Monarch矩阵

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

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

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

图1Monarch 矩阵是一种简单、富有表现力且硬件效率高的次二次结构矩阵。Monarch Mixer (M2)使用Monarch=矩阵来混合输入首先沿着序列维度然后沿着模型维度。

图 2Monarch 乘法可以解释为多项式求值和插值。

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

表1RTX 4090 上各种混频器层的 FLOP 成本和利用率输入维数64K

实验测试结果

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

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

1非因果语言建模

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

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

图 3M2-BERT 使用 Monarch 矩阵在序列混合器中创建双向门控长卷积并使用 Monarch 矩阵替换维度混合器中的线性层。

在序列混合器中注意力被带残差卷积的双向门控卷积替代见图 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 的模型的性能表现。

表 3M2-BERT-base 与 BERT-base相比的平均 GLUE 分数以及参数和 GLUE 分数的变化

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

表 4M2-BERT-large 与 BERT-large 相比的平均 GLUE 得分以及变化参数和 GLUE 分数

从这些表中结果可以看到在 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 数这能反映推理时间。

表 5M2-BERT-base (80M) 吞吐量与 BERT-base 比较结果以token/毫秒为单位

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

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

表 6在不同输入序列长度下批量大小为 1 的 CPU 推理延迟以毫秒为单位。 在运行 Intel Cascade Lake 处理器的 GCP n2-standard-48 系列的 48 vCPU、96 GB RAM 实例上对 10 多个示例进行了平均测量

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

(2) ViT风格的图像分类

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

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

表 7ImageNet-1k 上的准确性。 ResNet-152 提供供参考。

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

3GPT风格因果语言建模

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

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

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

表 8针对不同数量的标记进行训练时PILE 上的困惑度。

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

下一步研究计划

在最近的博客中hazyresearch.stanford.edu/blog/2023-07-25-m2-bert作者列出了他们的下一步研究计划

1. 我们今天发布了 BERT 代码以及 80M 和 110M 模型的检查点代码使用序列长度 128 的标准配方进行了预训练 - 请继续关注更长的序列 查看我们的代码和检查点80M、110M。
2. 在接下来的几周内请留意进一步的发布因为我们将训练长序列 BERT 并开始追溯 Transformers 的历史 - 在 ImageNet、因果语言建模、T5 风格模型以及对长序列功能的探索。
3. 作为此版本的一部分您将找到一些用于 M2 层前向传递的优化 CUDA 代码我们将其用于基准测试——我们将在未来几周内继续优化并发布更新。 当我们探索计算权衡空间时期待有关这些的另一系列博客和材料
4. 当然更完整的论文将在 arXiv 中推出

小结

该论文的研究工作为机器学习领域带来了新的思路挑战了传统Transformer模型的优越性。他们的研究不仅探索了Monarch Mixer的理论基础还进行了一系列实验来验证其性能。这篇文章的发表为机器学习社区提供了一个全新的研究方向也让人们重新思考了在自然语言处理和计算机视觉任务中的模型选择。

总的来说Monarch MixerM2是一种具有次二次复杂度的新型模型架构能够在不使用传统Transformer中的注意力和MLP的情况下在自然语言处理和计算机视觉任务中表现出色。它的硬件效率和参数效率使其成为一个有望取代传统Transformer的新选择为深度学习研究领域带来了新的思考。

参考文献

Daniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, Christopher Ré. “Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture”Oct 18, 2023, https://arxiv.org/abs/2310.12109