[论文解读] Monarch: Expressive Structured Matrices for Efficient and Accurate Training

Tri Dao, Beidi Chen+8 more|arXiv (Cornell University)|Apr 1, 2022

Advanced Neural Network Applications被引用 8

一句话总结

Monarch 提出了一类结构化矩阵，其参数化形式为两个块对角矩阵的乘积，可在深度学习中实现硬件高效且表达能力强的表示。该方法在 ViT 和 GPT-2 中实现了最高达 2 倍的训练速度提升，在 PDE/MRI 重建任务中将误差降低 40%，并通过反向稀疏化和基于 Monarch 的密集到稀疏微调技术，使 BERT 预训练速度提升 23% 而准确率无下降。

ABSTRACT

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.

研究动机与目标

• 解决现有结构化矩阵在端到端深度学习训练中效率低下且表达能力有限的问题。
• 克服在密集到稀疏微调过程中将密集权重矩阵转换为结构化形式时缺乏可 tractable 近似算法的问题。
• 通过与优化 GPU 内核兼容的参数化方法，实现稀疏模型的硬件高效训练。
• 开发一种新型稀疏到密集微调的中间表示，保持模型质量的同时加速训练。
• 展示 Monarch 矩阵在科学机器学习任务（如 PDE 求解和 MRI 重建）中的实用性。

提出的方法

• 提出 Monarch 矩阵作为两个块对角矩阵（允许排列）的乘积，实现在 GPU 上的高效批量矩阵乘法。
• 利用优化的 BMM 例程，在训练中实现最高达 2 倍于密集矩阵乘法的速度提升。
• 证明 Monarch 矩阵能够表示傅里叶变换、卷积和余弦变换等关键变换，确保高表达能力。
• 尽管问题具有非凸性，仍推导出近似密集权重矩阵为 Monarch 矩阵的解析最优解。
• 引入“反向稀疏化”——一种先使用 Monarch 矩阵训练，再微调为密集权重的技术，加速密集模型的预训练。
• 在密集到稀疏微调中应用基于 Monarch 的投影，利用解析近似高效地从密集模型迁移知识。

实验结果

研究问题

• RQ1 是否存在一类结构化矩阵，既能实现硬件高效，又具备足够的表达能力以表示深度学习和科学计算中常用的变换？
• RQ2 Monarch 近似问题的解析可解性是否能实现实用且高效的密集到稀疏微调？
• RQ3 Monarch 矩阵是否能在不牺牲模型质量的前提下，加速稀疏模型的端到端训练？
• RQ4 基于 Monarch 的中间表示是否能提升大规模语言模型中稀疏到密集微调的效率？
• RQ5 Monarch 矩阵是否能提升数据受限的科学任务（如加速 MRI 和 PDE 求解）中的重建精度？

主要发现

• Monarch 矩阵在 ImageNet 和 Wikitext-103 上对 ViT 和 GPT-2 的训练速度最高提升 2 倍，且模型准确率相当。
• 基于 Monarch 的 MRI 重建（mSENSE）相比基线方法将重建误差降低 40%，即使仅使用一次训练扫描。
• 使用 Monarch 矩阵的反向稀疏化技术使 OpenWebText 上 GPT-2 预训练速度提升 2 倍，且无准确率下降。
• 相同技术使 BERT 预训练速度比 Nvidia 的优化版 MLPerf 1.1 实现快 23%。
• 基于 Monarch 的密集到稀疏微调使 GLUE 上 BERT 微调速度提升 1.7 倍，且准确率相当。
• Monarch 近似的解析解实现了从密集模型到结构化模型的知识高效且精确的迁移。