[논문 리뷰] Monarch: Expressive Structured Matrices for Efficient and Accurate Training

Tri Dao, Beidi Chen외 8명|arXiv (Cornell University)|2022. 04. 01.

Advanced Neural Network Applications인용 수 8

한 줄 요약

Monarch는 두 개의 블록 대각 행렬의 곱으로 파arameterized된 구조적 행렬의 클래스를 제안하며, 딥러닝을 위한 하드웨어 효율적이고 표현력 있는 표현을 가능하게 한다. 이는 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.

연구 동기 및 목표

• 엔드 투 엔드 딥러닝 학습에서 기존의 구조적 행렬의 비효율성과 제한된 표현력 문제를 해결한다.
• 밀도-희소 미세조정 과정에서 밀도 가중치 행렬을 구조적 형태로 변환하기 위한 타당한 근사 알고리즘이 부족한 문제를 해결한다.
• 최적화된 GPU 커널과 호환되는 파arameterization을 통해 희소 모델의 하드웨어 효율적 학습을 가능하게 한다.
• 모델 품질을 유지하면서 학습 속도를 향상시키는 희소-밀도 미세조정을 위한 새로운 중간 표현을 개발한다.
• Monarch 행렬이 PDE 해소 및 MRI 재구성과 같은 과학적 기계학습 작업에서의 유용성을 입증한다.

제안 방법

• GPU에서 효율적인 배치 행렬 곱셈을 가능하게 하기 위해, 두 개의 블록 대각 행렬의 곱(적절한 순열을 포함)으로 구성된 Monarch 행렬을 제안한다.
• 최적화된 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배 빠르게 하며 정확도 저하 없이 성능을 유지한다.
• 동일한 기법은 Nvidia의 최적화된 MLPerf 1.1 구현보다 BERT 사전학습 속도를 23% 빠르게 한다.
• Monarch 기반의 밀도-희소 미세조정은 GLUE에서 BERT 미세조정 속도를 1.7배 빠르게 하며 정확도는 유사하게 유지한다.
• Monarch 근사의 분석적 해법은 밀도 모델에서 구조적 모델로 지식을 효율적이고 정확하게 전달할 수 있도록 한다.