[논문 리뷰] DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter

Victor Sanh, Lysandre Debut외 2명|arXiv (Cornell University)|2019. 10. 02.

Topic Modeling참고 문헌 18인용 수 4,569

한 줄 요약

이 논문은 사전 훈련 중 지식 정련을 통해 BERT의 언어 이해 능력의 97%를 유지하면서도 크기가 작고, 빠르고, 더 효율적인 BERT의 버전인 DistilBERT를 소개한다. 마스크된 언어 모델링, 정련, 코사인 거리 손실을 조합한 삼중 손실를 사용함으로써 DistilBERT는 모델 크기를 40% 줄이고 추론 시간을 60% 감소시켰으며, GLUE 및 후속 작업에서 강력한 성능을 유지한다.

ABSTRACT

As Transfer Learning from large-scale pre-trained models becomes more prevalent in Natural Language Processing (NLP), operating these large models in on-the-edge and/or under constrained computational training or inference budgets remains challenging. In this work, we propose a method to pre-train a smaller general-purpose language representation model, called DistilBERT, which can then be fine-tuned with good performances on a wide range of tasks like its larger counterparts. While most prior work investigated the use of distillation for building task-specific models, we leverage knowledge distillation during the pre-training phase and show that it is possible to reduce the size of a BERT model by 40%, while retaining 97% of its language understanding capabilities and being 60% faster. To leverage the inductive biases learned by larger models during pre-training, we introduce a triple loss combining language modeling, distillation and cosine-distance losses. Our smaller, faster and lighter model is cheaper to pre-train and we demonstrate its capabilities for on-device computations in a proof-of-concept experiment and a comparative on-device study.

연구 동기 및 목표

• 일반 목적의 언어 모델로서 크기가 작고, 빠르며, 더 효율적이면서도 NLP 작업에서 높은 성능을 유지할 수 있도록 개발하는 것.
• BERT와 같은 대규모 사전 훈련 모델을 훈련하고 배포하는 데 드는 계산 비용과 환경 비용을 해결하는 것.
• 사전 훈련 단계에서 지식 정련을 탐구하는 것. 특히, 미세 조정이나 작업별 적응을 위한 것만이 아니라.
• 이동 및 엣지 디바이스에서 실행할 수 있도록 충분히 작아서 효율적인 디바이스 내 추론을 가능하게 하는 모델을 만들기 위한 것.
• 사전 훈련 중 정련이 강력한 일반화 및 전이 능력을 지닌 컴act 모델을 생성할 수 있음을 보여주는 것.

제안 방법

• DistilBERT는 레이어 수를 반으로 줄이고 토큰 유형 임베딩 및 풀러를 제거함으로써 BERT의 40% 더 적은 파라미터를 가진 정련된 버전이다.
• 학생 모델는 교사 BERT-base 모델의 매번 다른 레이어를 가져와 초기화되어 지식을 유지하고 수렴을 향상시킨다.
• 삼중 손실 함수는 세 가지 구성 요소를 조합한다: 마스크된 언어 모델링(MLM) 손실, 교사의 소프트 레이블에 대한 정련 손실, 그리고 은닉 상태 방향을 일치시키기 위한 코사인 임베딩 손실.
• 정련 손실은 교사의 소프트 레이블 분포에서 지식을 학생 모델로 전달하기 위해 온도 조절된 소프트맥스를 사용한다.
• 훈련은 BERT와 동일한 위키백과 및 BookCorpus 데이터셋을 사용하며, 대용량 배치 경량화 및 동적 마스킹을 적용한다.
• 모델은 8×16GB V100 GPU를 사용해 약 90시간 동안 훈련되었으며, 전체 BERT 훈련에 비해 계산 예산을 크게 줄였다.

실험 결과

연구 질문

• RQ1 사전 훈련 중 지식 정련이 BERT의 성능 대부분을 유지하면서도 더 작은 언어 모델을 생성할 수 있는가?
• RQ2 정련된 모델에서 높은 성능을 달성하기 위해 손실 함수의 어떤 구성 요소가 필수적인가?
• RQ3 사용자 작업에 특화된 정련과 비교해 사전 훈련 중 정련을 수행했을 때 후속 작업 성능는 어떻게 되는가?
• RQ4 정련된 모델는 이동 및 엣지 디바이스에서 효율적으로 실행할 수 있을 정도로 작고 빠른가?
• RQ5 학생 모델를 교사의 가중치에서 초기화하는 것이 수렴과 성능 향상에 기여하는가?

주요 결과

• DistilBERT는 GLUE 벤치마크의 매크로 스코어에서 BERT-base의 97% 성능을 달성했으며, 파라미터 수는 40%에 불과하다.
• STS-B 작업에서 DistilBERT는 BERT-base 대비 60% 더 빠른 추론 속도를 보였으며, CPU에서의 추론 시간은 410ms로 BERT의 668ms보다 짧았다.
• IMDb 감성 분류 작업에서 DistilBERT는 92.82%의 정확도를 기록했으며, BERT-base의 93.46%에 비해 0.64% 뒤져 있었다.
• SQuAD 1.1에서 DistilBERT는 EM/F1 점수로 77.7/85.8을 기록했으며, BERT-base의 81.2/88.5에 비해 3.5점 이내였다.
• 제거 실험 결과, 삼중 손실의 어떤 구성 요소(MLM, 정련, 코사인 손실)를 제거해도 성능 저하가 발생했으며, 특히 코사인 손실이 기여도가 높았다.
• 모바일 앱 개념 증명에서 DistilBERT는 iPhone 7 Plus에서 BERT-base 대비 71% 더 빠르게 실행되었으며, 모델 크기는 207MB였다.