Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture
Dec 15, 2023
Speakers
Simran Arora
Speaker · 0 followers
Jessica Grogan
Speaker · 0 followers
Isys Johnson
Speaker · 0 followers
Evan Sabri Eyuboglu
Speaker · 0 followers
Armin W. Thomas
Speaker · 0 followers
Benjamin Spector
Speaker · 0 followers
Michael Poli
Speaker · 0 followers
Atri Rudra
Speaker · 0 followers
Christopher Re
Speaker · 0 followers
About
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.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 dimens…
Organizer
NeurIPS 2023
Account · 645 followers
Like the format? Trust SlidesLive to capture your next event!
Professional recording and live streaming, delivered globally.
Sharing
Recommended Videos
Presentations on similar topic, category or speaker
Knowledge Distillation Performs Partial Variance Reduction
Watch later
Total of 0 viewers voted for saving the presentation to eternal vault which is 0.0%
HYTREL: Hypergraph-enhanced Tabular Data Representation Learning
Watch later
Pei Chen, …
Total of 0 viewers voted for saving the presentation to eternal vault which is 0.0%
Cross-modal Prompts: Adapting Large Pre-trained Models for Audio-Visual Downstream Tasks
Watch later
Haoyi Duan, …
Total of 0 viewers voted for saving the presentation to eternal vault which is 0.0%
Semantic Image Synthesis with Unconditional Generator
Watch later
JungWoo Chae, …
Total of 0 viewers voted for saving the presentation to eternal vault which is 0.0%
Federated Learning with Client Subsampling, Data Heterogeneity, and Unbounded Smoothness: A New Algorithm and Lower Bounds
Watch later
Total of 0 viewers voted for saving the presentation to eternal vault which is 0.0%
From Discrete Tokens to High-Fidelity Audio Using Multi-Band Diffusion
Watch later
Total of 0 viewers voted for saving the presentation to eternal vault which is 0.0%