DistilBERT(a distilled version of BERT) 논문 리뷰

DistilBERT 배경 설명

•  DistillBERT는 Huggiing Face에서 2019년에 낸 논문으로, knowledge distillation을 사용하여, BERT의 문제점으로 지적되던, 큰 parameter 문제를 극복하기 위한 연구이다.
• 실제로 BERT를 실생활에서 이용할 때, 속도와 메모리 때문에 거의 DistilBERT를 많이 이용하는 것 같다.  

 

Abstract

• NLP 분야에서 large-scale pre-trained model을 사용한 transfer learning이 대세가 되었지만, model의 크기가 커서 연산 cost는 아직 크다.
• 이 논문에서는 DistilBERT라는 더 작지만, 성능은 떨어지지 않는 새로운 모델을 제안한다. 
• DistilBERT은 pre-training 단계에서 knowledge distillation을 이용하여, BERT의 모델 크기를 40%로 줄이고, 60%나 빠르지만, 성능은 97% 정도로 거의 떨어지지 않는다.
• lLarge language model이 학습한 내용을 배우기 위해, language modeling에 triplet loss를 사용하여 학습한다.

 

Introduction

• NLP 분야에서 transfer learning이 주류로 자리 잡은 이후, 성능 향상을 위해 모델의 크기를 점점 키우고 있다. 
• 하지만, 이런 흐름은 1) 학습을 위한 자원 사용이 계속 증가한다는 점 2) 연산과 메모리 소모가 더 커져서, real time 처리와 같은 응용을 어렵게 한다는 점이 있다. 
• 이 논문에서는 knowledge distillation을 사용하여, 기존 languager model보다 훨씬 더 작아, 가볍고, 속도가 빠르면서 성능은 좋은 모델을 제안한다.
•  triplet loss를 사용하여 학습된, 기존보다 40% 가량 작은 Transformer가 기존 방식으로 학습된 Transformer와 거의 비슷한 성능을 내면서, 60%나 빠르다는 것을 보인다.

 

Knowledge Distillation

• knowledge distillation은 compact model을 위한 압축방식에서 비롯된 것으로, larger model(일반적으로 teacher라고 부름)의 performance를 따를 수 있는 compact model(student라고 함)을 찾는 과정이다.
• classification 학습은 일반적으로 one-hot vector를 이용한 Cross-entropy를 이용한다. 따라서, 각 class에 분포될 확률을 가지고 있다. classification 학습에서 knowledge distillation은 기존 classification을 위한 학습뿐 아니라, large model의 확률 분포를 따르도록 compact model을 학습한다(distillation loss).
•  이 논문에서도 pre-training 과정에서 mask를 추정할때, distiilation loss의 개념을 활용한다. 추가적으로 compact model과 large model의 hidden states vector가 align 되도록 cosine embedding loss를 추가했다.

 

DistillBERT

[Student(compact model) architecture]

• 전체적인 모델은 BERT와 동일하지만, layer의 개수가 반으로 줄었고, token-type embeddings와 pooler를 삭제했다. 
• transformer의 hidden state dimension은 연산이 최적화된 라이브러리를 사용하기 때문에, 줄여도 큰 영향을 미치지 않아, layer를 줄이는데 집중했다.

[Student Initialization]

• Teacher와 Student가 구조가 비슷한 덕(dimension이 같아서)에 Teacher의 2개 layer마다 1개 layer의 parameter를 Student로 가져와서 Initialization 하였다. 

[Distillation]

• RoBERTa의 방법을 따라서 학습했다. 

2023.06.03 - [NLP 논문] - RoBERTa (A Robustly Optimized BERT Pretraining Approach) 논문 리뷰

RoBERTa (A Robustly Optimized BERT Pretraining Approach) 논문 리뷰

[Data & Computational Power]

• Original BERT처럼 English Wikipedia와 Toronto Book Corpus를 사용했다.
• 8대의 16GB V100로 90시간 학습했다. 

 

Experiments

• GLUE에서 성능 비교를 했을때, ELMo보다는 높고, Original BERT보다 40%의 paramter만 사용함에도 불구하고, BERT 성능에 97% 정도로 매우 비슷한 성능을 보인다.

• Downstream tasks들에서도 DistillBERT는 BERT와 과의 비슷한 성능을 보인다.

• Parameter와 속도를 비교해보았을때, original BERT에 비해 매우 적은 parameter를 사용하고, 빠르다는 것을 알 수 있다.

 

Reference 

DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter, Victor Sanh and Lysandre Debut and Julien Chaumond and Thomas Wolf, arXiv, 2019

 

총평

• 아이디어가 매우 간단하고, 명확해서 좋다. 
• BERT를 실용 단계에서 잘 활용할 수 있는 이런 연구들이 많이 나왔으면 좋겠다.