[論文レビュー] 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は、2つのブロック対角行列の積としてパrameterizedされた構造的行列のクラスを提案し、深層学習におけるハードウェア効率的で表現力豊かな表現を可能にする。Vision Transformer (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カーネルと互換性を持つパrameterizationを用いて、スパースモデルのハードウェア効率的学習を可能にすること。
• スパースから全結合への微調整のための新規中間表現を開発し、モデル品質を維持しながら学習を高速化すること。
• PDEの解法やMRI再構成といった科学的機械学習タスクにおけるMonarch行列の実用性を示すこと。

提案手法

• GPU上で効率的なバッチ行列乗算を可能にするために、2つのブロック対角行列（置換を含む）の積として定義されるMonarch行列を提案する。
• 最適化されたBMMルーチンを活用し、学習において密度行列乗算よりも最大2倍の高速化を達成する。
• Fourier変換、畳み込み、コサイン変換といった主要な変換を表現できることを示し、高い表現力を保証する。
• 問題の非凸性にもかかわらず、密度行列を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）は、1回のトレーニングスキャンでのみ、ベースライン手法と比較して再構成誤差を40%低減する。
• Monarch行列を用いた逆スパース化により、OpenWebTextでのGPT-2事前学習が2倍高速化され、精度に劣化は認められない。
• 同じ手法により、Nvidiaの最適化されたMLPerf 1.1実装と比較して、BERTの事前学習が23%高速化される。
• Monarchベースの全結合からスパースへの微調整により、GLUEでのBERT微調整が1.7倍高速化され、精度は同等を維持する。
• Monarch近似の解析的解法により、密度モデルから構造的モデルへの知識の効率的かつ正確な転送が可能になる。