Monarch:用于高效精确训练的表达型结构化矩阵
机器学习
2022-04-04 v1
摘要
大型神经网络在许多领域表现出色,但训练与微调代价高昂。降低其计算或内存需求的流行方法是用结构化矩阵(如稀疏、低秩、傅里叶变换)替代稠密权重矩阵。这些方法尚未被广泛采用:(1)在端到端训练中由于不利的效率—质量权衡;(2)在稠密到稀疏微调中由于缺乏可处理的算法来近似给定稠密权重矩阵。为解决这些问题,我们提出一类矩阵(Monarch),其硬件高效(被参数化为两个块对角矩阵的乘积以获得更好硬件利用率)且具表达力(可表示许多常用变换)。令人惊讶的是,用 Monarch 矩阵近似稠密权重矩阵的问题虽非凸,却有解析最优解。Monarch 矩阵的这些性质开启了训练和微调稀疏与稠密模型的新途径。我们经验性验证 Monarch 在若干端到端稀疏训练应用中可实现有利的准确性—效率权衡:在 ImageNet 分类与 Wikitext-103 语言建模上以可比模型质量将 ViT 和 GPT-2 训练加速 2 倍,并将 PDE 求解与 MRI 重建任务误差降低 40%。在稀疏到稠密训练中,通过称为“反向稀疏化”的简单技术,Monarch 矩阵作为有用中间表示在 OpenWebText 上将 GPT-2 预训练加速 2 倍而无质量下降。同样技术在 BERT 预训练上比即便设定 MLPerf 1.1 纪录的 Nvidia 高度优化实现还快 23%。在稠密到稀疏微调中,作为概念验证,我们的 Monarch 近似算法在 GLUE 上将 BERT 微调加速 1.7 倍且精度相当。
引用
@article{arxiv.2204.00595,
title = {Monarch: Expressive Structured Matrices for Efficient and Accurate Training},
author = {Tri Dao and Beidi Chen and Nimit Sohoni and Arjun Desai and Michael Poli and Jessica Grogan and Alexander Liu and Aniruddh Rao and Atri Rudra and Christopher Ré},
journal= {arXiv preprint arXiv:2204.00595},
year = {2022}
}