Monarch Mixer：一种基于GEMM的简单次二次架构

机器学习模型正日益在序列长度和模型维度上同时扩展，以达到更长的上下文和更好的性能。然而，Transformer 等现有架构在这两个轴上均呈二次方扩展。我们提问：是否存在能在序列长度和模型维度上均以次二次方式扩展且性能良好的架构？我们引入 Monarch Mixer (M2)，一种在序列长度和模型维度上使用相同次二次原语的架构：Monarch 矩阵，这是一类简单的富有表现力的结构化矩阵，可涵盖许多线性变换，在 GPU 上实现高硬件效率，并以次二次方式扩展。作为概念验证，我们在三个领域探索 M2 的性能：非因果 BERT 风格语言建模、ViT 风格图像分类以及因果 GPT 风格语言建模。对于非因果 BERT 风格建模，M2 在下游 GLUE 质量上匹配 BERT-base 和 BERT-large，且参数最多减少 27%，并在序列长度 4K 时实现高达 9.1$\times$ 的吞吐量的提升。在 ImageNet 上，M2 以仅一半的参数比 ViT-b 准确率高出 1%。因果 GPT 风格模型带来了一个技术挑战：通过掩码强制因果性引入了二次瓶颈。为缓解该瓶颈，我们基于多元多项式求值与插值开发了 Monarch 矩阵的新理论视角，使得我们能将 M2 参数化为因果的同时保持次二次。使用该参数化，M2 在 The PILE 上的预训练困惑度以 360M 参数匹配 GPT 风格 Transformer——首次表明无需注意力或 MLP 也可能匹配 Transformer 的质量。