Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture Article Swipe

MAKE IMPACT

Daniel Y. Fu , Simran Arora , Jessica Grogan , Isys Johnson , Sabri Eyuboglu , Armin W. Thomas , Benjamin Spector , Michael Poli , Atri Rudra , Christopher Ré · YOU? · · 2023 · Open Access · · DOI: https://doi.org/10.48550/arxiv.2310.12109

Machine learning models are increasingly being scaled in both sequence length and model dimension to reach longer contexts and better performance. However, existing architectures such as Transformers scale quadratically along both these axes. We ask: are there performant architectures that can scale sub-quadratically along sequence length and model dimension? We introduce Monarch Mixer (M2), a new architecture that uses the same sub-quadratic primitive along both sequence length and model dimension: Monarch matrices, a simple class of expressive structured matrices that captures many linear transforms, achieves high hardware efficiency on GPUs, and scales sub-quadratically. As a proof of concept, we explore the performance of M2 in three domains: non-causal BERT-style language modeling, ViT-style image classification, and causal GPT-style language modeling. For non-causal BERT-style modeling, M2 matches BERT-base and BERT-large in downstream GLUE quality with up to 27% fewer parameters, and achieves up to 9.1$\times$ higher throughput at sequence length 4K. On ImageNet, M2 outperforms ViT-b by 1% in accuracy, with only half the parameters. Causal GPT-style models introduce a technical challenge: enforcing causality via masking introduces a quadratic bottleneck. To alleviate this bottleneck, we develop a novel theoretical view of Monarch matrices based on multivariate polynomial evaluation and interpolation, which lets us parameterize M2 to be causal while remaining sub-quadratic. Using this parameterization, M2 matches GPT-style Transformers at 360M parameters in pretraining perplexity on The PILE--showing for the first time that it may be possible to match Transformer quality without attention or MLPs.

Type

preprint

Language

en

Landing Page

http://arxiv.org/abs/2310.12109

PDF

https://arxiv.org/pdf/2310.12109

OA Status

green

Cited By

14

Related Works

10

OpenAlex ID

https://openalex.org/W4387804367

Raw OpenAlex JSON

OpenAlex ID

https://openalex.org/W4387804367

Canonical identifier for this work in OpenAlex

DOI

https://doi.org/10.48550/arxiv.2310.12109

Digital Object Identifier

Title

Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture

Work title

Type

preprint

OpenAlex work type

Language

en

Primary language

Publication year

2023

Year of publication

Publication date

2023-10-18

Full publication date if available

Authors

Daniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, Christopher Ré

List of authors in order

Landing page

https://arxiv.org/abs/2310.12109

Publisher landing page

PDF URL

https://arxiv.org/pdf/2310.12109

Direct link to full text PDF

Open access

Yes

Whether a free full text is available

OA status

green

Open access status per OpenAlex

OA URL

https://arxiv.org/pdf/2310.12109

Direct OA link when available

Concepts

Quadratic growth, Computer science, Bottleneck, Algorithm, Quadratic equation, Quadratic programming, Theoretical computer science, Mathematics, Mathematical optimization, Embedded system, Geometry

Top concepts (fields/topics) attached by OpenAlex

Cited by

14

Total citation count in OpenAlex

Citations by year (recent)

2025: 9, 2024: 4, 2023: 1

Per-year citation counts (last 5 years)

Related works (count)

10

Other works algorithmically related by OpenAlex

Full payload