[논문 리뷰] Prune Once for All: Sparse Pre-Trained Language Models

Ofir Zafrir, Ariel Larey외 3명|arXiv (Cornell University)|2021. 11. 10.

Topic Modeling인용 수 26

한 줄 요약

이 논문은 Prune Once for All (Prune OFA)을 소개합니다. 이는 가중치 가지치기와 지식 증류를 결합하여 다양한 아키텍처에 무관한 희소 사전 학습 Transformer 언어 모델을 학습시키는 방법으로, 최소한의 정확도 손실로 높은 희소성(예: 85–90%)을 달성하고, 필요시 양자화를 옵션으로 제공합니다.

ABSTRACT

Transformer-based language models are applied to a wide range of applications in natural language processing. However, they are inefficient and difficult to deploy. In recent years, many compression algorithms have been proposed to increase the implementation efficiency of large Transformer-based models on target hardware. In this work we present a new method for training sparse pre-trained Transformer language models by integrating weight pruning and model distillation. These sparse pre-trained models can be used to transfer learning for a wide range of tasks while maintaining their sparsity pattern. We demonstrate our method with three known architectures to create sparse pre-trained BERT-Base, BERT-Large and DistilBERT. We show how the compressed sparse pre-trained models we trained transfer their knowledge to five different downstream natural language tasks with minimal accuracy loss. Moreover, we show how to further compress the sparse models' weights to 8bit precision using quantization-aware training. For example, with our sparse pre-trained BERT-Large fine-tuned on SQuADv1.1 and quantized to 8bit we achieve a compression ratio of $40$X for the encoder with less than $1\%$ accuracy loss. To the best of our knowledge, our results show the best compression-to-accuracy ratio for BERT-Base, BERT-Large, and DistilBERT.

연구 동기 및 목표

• 모델 규모 증가와 환경 비용으로 인해 효율적이고 배포 가능한 Transformer LMs의 필요성이 커지는 점을 동기화합니다.
• 아키텍처에 무관한 방법(Prune OFA)을 제안하여 전이 학습 능력을 유지하는 희소한 사전 학습 LM을 학습합니다.
• 희소한 사전 학습 모델이 다수의 Downstream 태스크에서 최소한의 정확도 손실로 미세조정될 수 있음을 보여줍니다.
• 후속 양자화가 정확도에 modest한 영향을 주며 모델 크기를 추가로 감소시키고 재현 가능한 도구와 모델을 공개합니다.

제안 방법

• 단일 사전 학습/지식 전달 과정에서 비구조적 가중치 가지치기를 사용하여 희소한 사전 학습 LM을 얻습니다.
• 가지치기 과정에서 Gradual Magnitude Pruning (GMP)과 Learning Rate Rewinding (LRR), 그리고 지식 증류(KD)를 도입합니다.
• 다운스트림 미세조정 동안 희소성 패턴을 유지하기 위해 패턴-락(pattern-lock) 메커니즘을 적용합니다.
• 영문 위키피디아에서 사전 학습을 수행한 뒤, SQuADv1.1, GLUE 태스크로 전이하여 KD를 통해 성능을 유지합니다.
• 선택적으로 Quantization-Aware Training (QAT)을 적용하여 8비트 양자화 희소 모델을 얻습니다.
• 재현성을 위한 스크립트와 희소 사전 학습 모델이 포함된 오픈 소스 압축 라이브러리를 제공합니다.

실험 결과

연구 질문

• RQ1 사전 학습 중 가지치기가 다운스트림 태스크로 전달될 때 최소한의 정확도 손실로 희소한 사전 학습 언어 모델을 만들 수 있는가?
• RQ2 가지치기 중 GMP, LRR, KD를 결합하면 태스크별 가지치기보다 전달 성능이 향상되는가?
• RQ3 패턴-락을 통해 희소성 패턴을 보존하는 것이 미세조정 시 정확도 유지에 도움이 되는가?
• RQ4 다운스트림 양자화(8비트 QAT)가 희소한 사전 학습 모델의 압축-정확도 트레이드오프에 어떤 영향을 미치는가?

주요 결과

모델	희소성	KD를 이용한 전이	SQuAD	MNLI (m/mm)	SST-2	QNLI	QQP
Prune OFA (BERT-Base)	85%	Yes	78.59	86.63	81.67	82.53	91.34	89.95
Prune OFA (BERT-Base)	85%	No	78.00	86.16	82.45	83.05	88.82	87.79
Prune OFA (BERT-Base)	85%	Yes	81.10	88.42	82.71	83.67	91.46	90.34

• Prune OFA는 높은 희소도(85–90%)를 달성하면서도 SQuADv1.1 및 GLUE 태스크에서 밀집 기반 기준선 및 기존 가지치기 방법에 비해 전달 성능이 경쟁력 있게 유지됩니다.
• 전이 시 KD를 사용하면 결과가 향상되며, KD와 LRR 및 패턴-락의 결합은 추가 이득을 제공하고 대부분의 태스크에서 정확도 감소를 최소화합니다.
• 희소 모델에 대한 양자화 인식 학습은 정확도를 약간 감소시키지만(전체 정밀도 희소 모델 대비 평균 약 0.67% 상대적 손실), 크기 축소를 크게 달성하여 압축-정확도 비율을 개선합니다.
• 90% 희소성에서 BERT-Large의 경우 대부분의 태스크에서 약 1% 내의 정확도 손실 범위이고, 매개변수 효율성 측면에서 Dense BERT-Base를 능가하는 경우도 있습니다(비제로 매개변수 수 기준).
• 저자들은 재현 가능한 연구를 촉진하기 위해 압축 라이브러리와 희소 사전 학습 모델을 공개합니다.