Monarch: Expressive Structured Matrices for Efficient and Accurate Training

Abstract

Large neural networks excel in many domains, but they are expensive to train and fine-tune. A popular approach to reduce their compute/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 denseto-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 2× 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 2× 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.7× with comparable accuracy.

Spurious Valleys, NP-Hardness, and Tractability of Sparse Matrix Factorization with Fixed Support

“…These supports are interesting because they are those taken at the first two steps of the hierarchical algorithm in [25,44] for approximating a matrix by a product of N butterfly factors [25]. The first pair of support constraints (I 1 , J 1 ) is also equivalent to the recently proposed Monarch parameterization [9]. Both pairs (I 1 , J 1 ) and (I 2 , J 2 ) are proved to satisfy Theorem 3.…”

“…While revising this manuscript we heard about the work of Dao et al[9] introducing the "Monarch" class of structured matrices, essentially corresponding to the first stage of the recursion from[25,44].…”