NL-055, DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter (2019-Arxiv)

*Distilling
1.Pre-trained BERT 모델로 student 모델 초기화
2.Task data로 Fine-tuned BERT을 student에 distilling을 한다.
3.이 때는 teacher의 predicted prob을 이용한다.
4.Fine-tuned할 때는 distilling을 이용안하는 듯하다.

■ Comment
  • Hugging Face라는 회사에서 낸 논문이데 HuggingFace는 다양한 BERT 버전을 사용하기 쉽게 코드화해서 배포한 것으로 유명하다.
  • 세계적으로 다양한 언어로 구현된 BERT들이 있고 모델들 한국어 BERT로는 SKT의 BERT 가 들어가 있다.
  • BERT뿐만 아니라 기존의 존재하는 코드를 파이토치로 쉽게 쓸 수 있게 해주고 다른 논문을 리프로듀싱하는 것도 있는 것 같다.
  • 이 논문은 직접 학습시켜서 낸 논문이고 당연히 라이브러리화 되어 있다.
  • 하지만 논문이 어디에 억셉된 것은 아니고 실제로 논문 내용을 봐도 상당히 간결하고 자세한 내용들은 없고 방법론은 기존의 방법론과 크게 차이 있어 보이지는 않는다.

0. Abstract

  • 큰 스케일의 pre-trained model이 NLP에서 연구되고 있는데 이러한 큰 모델들을 계산 학습 혹은 인퍼런스 버짓이 정해진 제약이 있는 edge단에서 작동하는 것은 챌린지하다.
  • 이 논문에서는 DistilBERT라고 하는 작고 general language representation 모델을 제안한다.
    • 이는 많은 부분에서 finetune시 좋은 성능을 보여준다.
  • 다른 이전의 연구들은 task-specific 모델에 맞는 distillation을 수행하지만 여기서는 Pre-training을 할 때 Knowledge distillation을 활용한다.
  • 그 결과 BERT 사이즈는 40% 줄이고 lanugage understanding 능력은 99% 유지하고 60% 빠른 모델을 보여준다.
  • Pre-training 시 inductive bias을 활용하기 위해 LM을 결합한 triple loss, distillation, cosine-distance loss 3가지를 소개한다.
  • 우리의 작고 빠르고 가벼운 모델은 학습하기 싸고 on-device에서 계산 기능을 보여준다.

1. Introduction

  • 최근 NLP에서 큰 pre-trained LM이 연구가 활발히 되고 있다.
  • 최근 트렌드인 큰 모델들은 몇 가지 컨선이 있다.
    • 비용이 지수적으로 커지는 것
    • 실제 on-device에서 작동하면 novel하고 흥미롭지만, 증가하는 계산량과 메모리 요구가 적용하는데 방해가 된다.
  • 이 논문에서는 KD을 통하여 훨씬 작은 LM을 만들어서 downstream 작업에서 비슷한 성능을 낸다.
  • 이 모델은 훨씬 가볍고 인퍼런스 시간이 훨씬 빠르다.
  • 이 논문의 목적은 pre-trained 모델이 여러 다운스트림 작업에서 fine-tuned 했을 때 성능이 좋은 것이다.
  • 또한 edge나 모바일 디바이스에서 수행되기에 충분히 작음을 보여준다.
  • Triple loss을 통하여 40% 작은 Transforemr와 60% 빠른 인퍼런스 시간을 보여주며 성능은 비슷하다.

2. Knowledge distillation

  • Knowledge distillation
    • student모델이 큰 teacher 모델의 행동을 따라하도록 배우게 하는 압축 기술이다.
    • supervised learning에서 일반적으로 classification 모델은 gold labels의 확률을 최대화 하도록 학습된다.
    • 이는 cross-entropy을 최소화하도록 학습 하는 것이고 training labels의 분포와 model predicted 분포를 가까이 하는 의미이다.
    • 그 결과 correct labels에서는 가장 높은 확률을 나머지에서는 0에 가까운 확률을 내뱉게 된다.
  • Training loss
    • student은 distilaation loss을 다음과 같이 가진다.
      • ti는 teacher가 예측한 확률 값, si는 student가 예측한 확률 값이다.
      • 즉 real label이 아닌 teahcer의 확률값을 이용한 cross entropy loss라 보면 된다.
    • 이 학습 시그널은 full teacher 분포를 레버리징하는 결과를 보여준다.
    • 또한 학습할 때는 Temperature을 이용한 cross-entropy을 사용하고 inference할 때는 T=1을 이용한다.
  • 전체 loss는 와 함께 masked language modeling 와 cosine embedding loss 을 결합하여 학습한다.
  • 은 BERT에서 사용한 것이고 는 student와 teacher의 hidden states vectors가 같은 방향을 보도록 align 시켜주는 경향이 있다고 한다.

3. DistilBERT: a distilled version of BERT

3.1 Student architecture

  • DistilBERT는 token-type embedding과 pooler은 삭제한다.
  • layers의 수는 2배로 줄인다.
  • Transformer 구조의 operation은 (linear layer와 layer normalisation) 고도화로 최적화 되어 있는 프레임워크다.
  • 그리고 논문 저자들은 last dimension은 (hidden size dimension) 계산 효율성에 작은 영향을 준다고 한다.
  • 따라서 여기서는 layers의 수를 줄이는데 집중한다고 한다.

3.2 Student initialization

  • 모델의 구조와 optimization과 더불어 초기화를 통하여 sub-networkf가 제대로 수렴하도록 하는 것도 중요하다.
  • teacher와 student가 동일한 차원을 가지는 점을 이용하여 teacher의 2개의 layer중 하나를 student layer로 초기화 한다.
    • 아마도 이전의 연구 논문에서 말하는 방법들은 쓴 듯 (TinyBERT 혹은 PKD-skip과 같은 것)

3.3 Distillation

  • 우리는 BERT 모델을 4k의 매우 큰 배치를 활용하여 gradient accumulation으로 distill 했다고 한다.
  • 이 때 NSP없이 dynamic masking을 했다고 한다.
    • dynamic masking은 어떤 입력에 대해 masking 위치를 정하는 것이 아니고 학습할 때 같은 입력이 들어왔다고 해도 masking을 다르게 해주는 느낌

3.4 Data and compute power

  • 기존의 BERT와 같은 corpus으로 학습을 하였고 영어 위키피다와 Toronto Book Corpus을 concat하여 사용하였다.
  • 8 16GB V100 GPU를 이용하여 학습하였고 90시간 정도 걸렸다고 한다.
  • 비교를를 위해 RoBERTa 모델을 1024 32GB V100에서 하루 동안 학습했음.

4. Experiments

  • General Language Understanding
    • GLUE으로 DistilBERT을 평가해보았다.
    • Table 1의 결과를 보면, DistlBERT는 ELMo 보다 성능이 좋고 BERT의 99% 정도 성능을 이끌어내며 파라미터는 기존보다 40% 더 적은 파라미터를 가진다

4.1 Downstream task benchmark

    • Downstream tasks
      • DistilBERT을 다양한 다운스트림 task에 적용하였고 table 2가 그것의 결과이다.
      • 여기서 DistillBERT (D)는 fine-tunning할 때도 teacher 모델을 이용했다는 것이다.
      • 즉 SQuAD에 fine-tuned된 BERT 모델을 활용하여 추가적인 loss term (KD)을 구성했다는 것이다.
      • 이 때는 2개의 step이 있으며 pre-trainin phase와 adaptation phase에서의 distillation이 있다.
    • Size and inference speed
    • Ond device computation
      • QA 모바일 어플리케이션에서 on-the-edge 어플리케이션에서 DistilBERT을 사용할 수 있는지 연구를 했다.
      • 아이폰 7 플러스에서 평균적인 인퍼런스 속도를 BERT-base와 비교를 했다.
      • Tokenization step을 제외하고 DistilBERT는 BERT보다 71% 빠르다.
      • 전체 모델의 크기는 207MB이다. (이는 quantization으로 더 줄일 수 있다.)

    4.2 Ablation Study

    5. Related work

    • Task-specific distillation
      • Distilling task-specific knowledge from bert into simple neural networks.
      • Making neural machine reading comprehension faster
      • Well-read students learn better: The impact of student initialization on knowledge distillation
    • Multi-distillation
      • Model compression with multi-task knowledge distillation for web-scale question answering system
      • Small and practical bert models for sequence labeling
    • Other compression techniques
      • Are sixteen heads really better than one?
      • Deep learning with limited numerical precision

    6. Conclusion and future work

    • DistilBERT을 설명하고 기존 pre-trained BERT 보다 40% 작고 60% 빠르고 99% 성능을 유지한다.
    • ablation study을 통하여 다양한 구성 요소에 대해 분석하고 distillation을 성공적으로 학습시켰다.
    • 게다가 DistilBERT는 edge 어플리케이션에서 적절한 옵션임을 증명했다.
    Reference

    댓글

    댓글 쓰기