本文探索了 Monarch Mixer (M2) ，这是一种在序列长度和模型维度上都是次二次的新架构，并且在现代加速器上具有很高的硬件效率。
从 BERT、GPT 和 Flan-T5 等语言模型到 SAM 和 Stable Diffusion 等图像模型，Transformer 正以锐不可当之势席卷这个世界，但人们也不禁会问：Transformer 是唯一选择吗？
斯坦福大学和纽约州立大学布法罗分校的一个研究团队不仅为这一问题给出了否定答案，而且还提出了一种新的替代技术：Monarch Mixer。近日，该团队在 arXiv 公布了相关论文和一些检查点模型及训练代码。顺带一提，该论文已入选 NeurIPS 2023 并获得 Oral Presentation 资格。
代码地址：github.com/HazyResearc...
该方法去掉了 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 大不相同的模型也可能在因果语言建模取得出色表现。
了解更多内容，请参考原论文。