[논문 리뷰] MixKD: Towards Efficient Distillation of Large-scale Language Models

Kevin J Liang, Weituo Hao외 5명|arXiv (Cornell University)|2020. 11. 01.

Topic Modeling참고 문헌 52인용 수 30

한 줄 요약

MixKD는 믹스업으로 데이터를 보강하고 보간된 샘플에서 교사에게 질의함으로써 대규모 언어 모델의 지식 증류를 개선하고, GLUE 작업에서 컴팩트한 학생 모델의 일반화와 성능을 향상시킵니다.

ABSTRACT

Large-scale language models have recently demonstrated impressive empirical performance. Nevertheless, the improved results are attained at the price of bigger models, more power consumption, and slower inference, which hinder their applicability to low-resource (both memory and computation) platforms. Knowledge distillation (KD) has been demonstrated as an effective framework for compressing such big models. However, large-scale neural network systems are prone to memorize training instances, and thus tend to make inconsistent predictions when the data distribution is altered slightly. Moreover, the student model has few opportunities to request useful information from the teacher model when there is limited task-specific data available. To address these issues, we propose MixKD, a data-agnostic distillation framework that leverages mixup, a simple yet efficient data augmentation approach, to endow the resulting model with stronger generalization ability. Concretely, in addition to the original training examples, the student model is encouraged to mimic the teacher's behavior on the linear interpolation of example pairs as well. We prove from a theoretical perspective that under reasonable conditions MixKD gives rise to a smaller gap between the generalization error and the empirical error. To verify its effectiveness, we conduct experiments on the GLUE benchmark, where MixKD consistently leads to significant gains over the standard KD training, and outperforms several competitive baselines. Experiments under a limited-data setting and ablation studies further demonstrate the advantages of the proposed approach.

연구 동기 및 목표

• 정확도를 희생하지 않으면서 모델 크기와 추론 비용을 줄여 대규모 언어 모델의 효율적 배치를 촉진한다.
• 작업 데이터가 제한될 때 증강을 통해 학습 데이터를 풍부히 하여 과적합과 기억화를 다룬다.
• 교사–학생 학습을 위해 mixup을 활용해 증강 샘플을 생성하는 데이터에 얽매이지 않는 증류 프레임워크(MixKD)를 제안한다.
• 증강된 KD를 통한 일반화 향상을 이론적으로 정당화한다.
• 저데이터 구간에서 특히 GLUE 작업에서 실증적 이득을 보여주고, 기본 KD 및 관련 방법과 비교한다.

제안 방법

• 큰 교사(BERT)로부터 더 작은 학생(BERT-3 또는 BERT-6)으로 지식 증류를 수행한다.
• 단어 임베딩에 mixup을 적용하여 학생과 교사 지도에 대해 증강 입력을 생성한다 (x' = λx_i + (1-λ)x_j; y' = λy_i + (1-λ)y_j).
• 원래 데이터에 대한 표준 교차 엔트로피, mixup 데이터에 대한 학생의 교차 엔트로피, 그리고 mixup 데이터에서 교사와 학생 간의 증류 손실을 결합한 공동 목표로 학습한다 (L = L_MLE + α_SM L_SM + α_TMKD L_TMKD).
• mixup 샘플에서 교사를 질의하고 해당 샘플에 대한 학생 예측에 대해 증류 손실을 최소화한다 (L_TMKD).
• 증강 데이터가 경험적 오차와 일반화 오차 사이의 차이를 줄이고 일반화를 향상시키는 조건을 보여주는 이론적 결과를 제공한다.

실험 결과

연구 질문

• RQ1 KD(MixKD)의 mixup 기반 증강이 표준 KD에 비해 작은 학생 모델의 일반화를 향상시키고 일반화 갭을 줄이는가?
• RQ2 다양한 데이터 가용성에서 MixKD가 GLUE 작업에서 어떻게 수행되며 DistilBERT, PKD 및 다른 기준선과 어떻게 비교되는가?
• RQ3 KD에서 mixup 증강을 백트랜스레이션(backtranslation) 및 다른 데이터 증강 기법과 결합하는 영향은 무엇인가?
• RQ4 제한된 데이터 구간에서 특히 하이퍼파라미터와 mixup 비율에 대해 MixKD의 민감도는 어느 정도인가?

주요 결과

• MixKD는 GLUE 작업에서 표준 KD 및 관련 기준선보다 일관되게 개선되며, 제한된 데이터 환경에서 특히 큰 이점을 보인다.
• SM+TMKD+BT를 포함한 6레이어 MixKD 학생은 교사 성능의 큰 부분을 포착하면서도 추론 속도가 훨씬 빠르고 매개변수 수가 적은 경향을 보인다.
• SST-2에서 SMTKD와 백트랜스레이션을 사용한 MixKD는 처음부터 학습하거나 표준 KD와 비교했을 때 교사와의 격차를 상당 부분 줄인다.
• 이론적 결과는 mixup를 통한 데이터 증강이 여러 데이터 생성 시나리오에서 일반화 갭을 줄일 수 있음을 보여주어 실증적 이득을 뒷받침한다.
• 임베딩 시각화는 MixKD가 학생이 보간된 샘플을 데이터 매니폴드와 정렬되도록 안내하고 잠재 공간 구성을 개선함을 시사한다.