[논문 리뷰] LightPAFF: A Two-Stage Distillation Framework for Pre-training and Fine-tuning

Kaitao Song, Hao Sun외 5명|arXiv (Cornell University)|2020. 04. 27.

Topic Modeling참고 문헌 17인용 수 9

한 줄 요약

LightPAFF는 사전 훈련된 대규모 모델(BERT, GPT-2, MASS 등)의 지식을 사전 훈련 단계와 미세조정 단계 모두에서 더 작은 학생 모델로 전이하는 이단계 지식 정복 프레임워크입니다. 이로 인해 모델 크기를 최대 5배 작게 만들고 추론 속도를 5~7배 빠르게 하며, 원래의 대규모 모델과 유사한 성능을 달성하여 자원이 제한된 시스템에 효율적으로 구현할 수 있습니다.

ABSTRACT

While pre-training and fine-tuning, e.g., BERT~\citep{devlin2018bert}, GPT-2~\citep{radford2019language}, have achieved great success in language understanding and generation tasks, the pre-trained models are usually too big for online deployment in terms of both memory cost and inference speed, which hinders them from practical online usage. In this paper, we propose LightPAFF, a Lightweight Pre-training And Fine-tuning Framework that leverages two-stage knowledge distillation to transfer knowledge from a big teacher model to a lightweight student model in both pre-training and fine-tuning stages. In this way the lightweight model can achieve similar accuracy as the big teacher model, but with much fewer parameters and thus faster online inference speed. LightPAFF can support different pre-training methods (such as BERT, GPT-2 and MASS~\citep{song2019mass}) and be applied to many downstream tasks. Experiments on three language understanding tasks, three language modeling tasks and three sequence to sequence generation tasks demonstrate that while achieving similar accuracy with the big BERT, GPT-2 and MASS models, LightPAFF reduces the model size by nearly 5x and improves online inference speed by 5x-7x.

연구 동기 및 목표

• 고비용의 메모리 및 지연 시간으로 인해 온라인 또는 모바일 환경에서 큰 사전 훈련 모델(BERT, GPT-2 등)을 구현하는 데 도전하는 문제를 해결합니다.
• 기존 지식 정복 방법은 미세조정 단계에서만 모델을 압축하여 심각한 정확도 저하를 초래하는 한계를 극복합니다.
• 마스크된 언어 모델링(BERT), 인과적 언어 모델링(GPT-2), 마스크된 시퀀스-투-시퀀스 모델링(MASS) 등 다양한 사전 훈련 목표에 적용 가능한 통합 프레임워크를 개발합니다.
• 사전 훈련된 모델과 미세조정된 교사 모델의 지식을 활용하여 경량 모델이 일반적인 언어 이해 및 생성 능력을 유지할 수 있도록 합니다.
• 다양한 NLP 작업(분류, 언어 모델링, 시퀀스-투-시퀀스 생성 포함)에서 성능을 손상시키지 않은 채 효율적이고 저지연 추론을 가능하게 합니다.

제안 방법

• 사전 훈련 데이터를 사용하여 사전 훈련 단계에서 지식 정복을 적용하고, 작업별 데이터를 사용하여 미세조정 단계에서 다시 적용합니다.
• 교사 모델보다 파rameter 수가 적은 학생 모델을 사용하며, 두 단계 모두에서 교사 모델의 로짓과 은닉 표현을 모방하도록 훈련합니다.
• 훈련 중에 지식 정복 손실과 진짜 레이블에서 유도된 표준 교차 엔트로피 손실을 균형 잡기 위해 하이퍼파ram터 λ를 도입합니다.
• 매개변수 수준의 지식 정복을 적용하여 학생 모델과 교사 모델의 매개변수 간 L2 거리 최소화를 수행하며, 정확도 향상을 위해 추가로 가우시안 노이즈를 적용합니다.
• BERT, GPT-2, MASS 모두에 동일한 방식으로 정복 목표를 설정하여 다양한 사전 훈련 목표와 모델 아키텍처에 걸쳐 일관된 적용이 가능하도록 합니다.
• 이단계 파이프라인을 사용합니다: (1) 교사 모델 사전 훈련, (2) 하류 작업에 대해 교사 모델 미세조정, (3) 사전 훈련된 교사 모델을 기반으로 가벼운 사전 훈련된 학생 모델 생성, (4) 하류 작업에 대해 학생 모델을 정복을 통해 미세조정.

실험 결과

연구 질문

• RQ1 사전 훈련 및 미세조정 단계 모두에 지식 정복을 적용할 경우, 모델 크기와 추론 지연 시간을 크게 줄일 수 있을까? 성능은 유지되는가?
• RQ2 이단계 정복은 단일 단계 정복(오직 미세조정 단계에서만)과 비교해 경량 모델의 정확도와 강건성 측면에서 어떻게 다를까?
• RQ3 사전 훈련 작업의 어려움(예: 마스크된 vs. 인과적 언어 모델링)이 최적의 정복 하이퍼파ram터와 모델 성능에 영향을 미치는가?
• RQ4 제안된 LightPAFF 프레임워크는 BERT, GPT-2, MASS와 같이 다양한 사전 훈련 목표를 가진 다양한 사전 훈련 모델에 균일하게 적용 가능한가?
• RQ5 모두의 단계에서 정복을 적용할 경우, 오직 미세조정 단계에서만 정복을 적용하는 것보다 일반화 능력과 강건성이 얼마나 향상되는가?

주요 결과

• LightPAFF는 BERT, GPT-2, MASS의 모델 크기를 거의 5배 작게 줄였지만 원래의 대규모 모델과 유사한 성능을 유지합니다.
• 모든 평가된 모델에서 온라인 추론 속도가 5배에서 7배 빨라졌으며, 엣지 장치에 실용적으로 구현할 수 있도록 했습니다.
• 언어 이해(SST-2, QQP, 다의어 해소), 언어 모델링(WikiText-2, PTB, WikiText-103), 시퀀스-투-시퀀스 생성(WMT Zh-En, En-De, En-Fr)에 각각 3개씩 포함된 총 9개의 하류 작업에서 압축된 모델이 일관된 성능 향상을 보였습니다.
• 이단계 정복 접근법은 더 나은 일반화 능력을 보였으며, 매개변수 변형에 대한 강건성 향상과 사전 훈련 및 미세조정 단계에서의 안정적인 성능을 입증했습니다.
• 최적의 정복 하이퍼파ram터 λ는 모델에 따라 다릅니다: BERT는 λ = 1.0, MASS는 λ = 0.7, GPT-2는 λ = 0.4이며, 이는 사전 훈련 작업의 어려움과 정보 유지 능력의 차이를 반영합니다.
• 사전 훈련 작업의 어려움과 성능 간 상관관계가 있습니다: BERT(평균적으로 85%의 입력 토큰을 볼 수 있음)가 가장 쉽고, 다음으로 MASS(75%), GPT-2(50%)가 어려운 편이며, 이는 마스크 예측 작업에서 관찰된 정확도 저하와 일치합니다.