Monarch: Expressive Structured Matrices for Efficient and Accurate Training

By Tri Dao, Beidi Chen, Nimit Sharad Sohoni, Arjun D Desai, Michael Poli, Jessica Grogan, Alexander Liu, Atri Rudra, Christopher R'e. 2022

133 citations recorded by the source corpus.

Large neural networks excel in many domains, but they are expensive to train and fine-tune. A popular approach to reduce their compute or memory requirements is to replace dense weight matrices with structured ones (e.g., sparse, low-rank, Fourier transform). These methods have not seen widespread adoption (1) in end-to-end training due to unfavorable efficiency--quality tradeoffs, and (2) in dense-to-sparse fine-tuning due to lack of tractable algorithms to approximate a given dense weight matrix. To address these issues, we propose a class of matrices (Monarch) that is hardware-efficient (they are parameterized as products of two block-diagonal matrices for better hardware utilization) and expressive (they can represent many commonly used transforms). Surprisingly, the problem of approximating a dense weight matrix with a Monarch matrix, though nonconvex, has an analytical optimal solution. These properties of Monarch matrices unlock new ways to train and fine-tune sparse and dense models. We empirically validate that Monarch can achieve favorable accuracy-efficiency tradeoffs in several end-to-end sparse training applications: speeding up ViT and GPT-2 training on ImageNet classification and Wikitext-103 language modeling by 2x with comparable model quality, and reducing the error on PDE solving and MRI reconstruction tasks by 40%. In sparse-to-dense training, with a simple technique called"reverse sparsification,"Monarch matrices serve as a useful intermediate representation to speed up GPT-2 pretraining on OpenWebText by 2x without quality drop. The same technique brings 23% faster BERT pretraining than even the very optimized implementation from Nvidia that set the MLPerf 1.1 record. In dense-to-sparse fine-tuning, as a proof-of-concept, our Monarch approximation algorithm speeds up BERT fine-tuning on GLUE by 1.7x with comparable accuracy.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Model Architectures

Citations in the corpus

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, background): “[17] Tri Dao, Beidi Chen, Nimit Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré.”
Hungry Hungry Hippos: Towards Language Modeling with State Space Models (2022, methodology): “By the Cooley-Tukey decomposition of the DFT [3, 11] (also known as the four-step FFT algorithm), we can write FN = P(IN2 ⊗FN1)PD(IN1 ⊗FN2)P, where P denotes a fixed permutation that reshapes the input as a N1 ×N2 array and then transpose it, ⊗ denotes Kroneker product, D is a N ×N diagonal matrix (”