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 scalesub-quadraticallyalong sequence length and model dimension?
We introduceMonarch 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 ofM2in three domains: non-causal BERT-style language modeling, ViT-style image classification, and causal GPT-style language modeling.
For non-causal BERT-style modeling,M2matches BERT-base and BERT-large in downstream GLUE quality with up to 27% fewer parameters, and achieves up to 9.1higher throughput at sequence length 4K.
On ImageNet,M2outperforms 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 parameterizeM2to be causal while remaining sub-quadratic. Using this parameterization,M2matches 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.