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

본 포스트를 읽기 전에 DistilBERT에 사용된 메인 테크닉인 Knowledge Distillation에 대해서 먼저 학습하시길 바랍니다. 다음의 포스트를 참고하시오.

 

The overview of this paper

 NLP에서 large-scale의 pre-trained model을 활용하여 transfer learning을 처리하는 일이 흔해지면서, 이 거대한 규모의 모델을 한정된 자원으로 어떻게 구동할지는 아직도 어려운 문제로 남아있다. 그래서 이 논문에서는 작은 규모의 general purpose language representation model임에도 불구하고, 다양한 분야의 task에 대해 좋은 성능을 보여주는 DistilBERT를 제안하였다. 이 DistilBERT는 BERT에 비해 40% 작은 크기에 60% 더 빠른 속도를 보여주면서, 기존 BERT의 97%에 달하는 성능을 보여주었다. larger model이 pre-training 도중에 학습하는 inductive bias를 leverage하기 위해서, triple loss가 합쳐진 LM을 소개하였다. 논문에서 소개한 이 작고 빠른 모델은 on-device 환경에서도 충분히 구동하다는 것 또한 증명되었다.

 

 

Table of Contents

1. Introduction

2. Knowledge distillation

3. DistilBERT: a distilled version of BERT

4. Ablation study

 

 

1. Introduction

 요즘에 개발되는 pre-trained language model을 살펴보면 대부분이 수많은 파라미터를 가지고 있다. 그리고, 요즘의 많은 연구들이 더 큰 규모의 LM은 더 좋은 성능 향상을 가져온다는 것을 증명하고 있다.

 

그림 1. 최근에 개발된 LM들의 파라미터 수

 

 이러한 큰 규모의 LM에 대한 관심은 다음의 걱정을 불러일으킬 수 있다.

 

1. 비용이 많이 발생하게 된다. $($computationally & environmentally$)$
2. LM의 발달은 LM이 on-device 환경에서 구동할 수 있는 잠재력을 보여주고 있지만, 너무 큰 규모로 인해 넓은 방면의 task에 적용을 할 수 없다는 크나큰 단점을 보유하고 있다.

 이 논문에서는, knowledge distillation을 통해 pre-train된 작은 크기의 LM이 larger model에 버금가는 성능을 보여줄 수 있다는 가능성을 보여줬다. 그리고 논문의 general-purpose pre-trained model이 larger model의 가용성을 보존하면서, 다양한 task에 대해서 좋은 성능을 보여주도록 fine-tune될 수 있다는 것을 보여줬다. 또한, 이 모델이 모바일 기기에서도 충분히 돌아갈 수 있을 만큼 작다는 것을 보여줬다.

 triple loss를 사용함으로써, 40% 작은 크기의 Transformer가 더 큰 Transformer과 비슷한 성능을 보여주었다. ablation study에 따르면 triple loss의 세 가지 요소는 좋은 성능을 내기 위해 매우 중요하다는 것을 알 수 있었다.

 

2. Knowledge distillation

 Knowledge distillation은 압축 기술로, compact model이 '학생'의 입장으로, '교사'인 larger model 또는 ensemble model로부터 이들의 방식을 따라하면서 학습하는 방법이다.

 지도학습에서는, 분류 모델이 최적의 라벨에 대한 측정 확률을 최대화시킴으로써 class를 예측하도록 학습된다. 따라서, 기존의 훈련 목적은 모델의 예측 분포와 훈련 라벨의 분포 간에 cross-entropy를 최소화시키는 것이다. 따라서 모델은 훈련셋에 대해서 매우 잘 수행하게 된다. 결국에 모델은 옳은 클래스에 대해서는 높은 확률로 예측하고, 그 반대의 클래스에 대해서는 거의 0에 가까운 확률로 예측을 한다. 그러나 이러한 "0에 가까운" 확률 중 일부는 다른 확률보다 크며 모델의 일반화 기능과 테스트 세트에서 얼마나 잘 수행될 것인지를 부분적으로 반영한다.

 

Training loss

 학생은 교사의 soft target probabilities에 대해 distillation probabilities로 학습된다: $L_{ce}=\sum_{i} t_{i}*log(s_i)$ 여기서 $t_i$는 교사에 의해 측정된 확률이다. 이 식은 전체 교사의 분포를 활용하여 풍부한 학습을 제공한다. Hinton을 따라서 이 논문에서는 softmax-temperature: $p_i=\frac {exp(z_i/T)}{\sum_{j} exp(z_j/T)}$을 사용하였다. 여기서 $T$는 출력 분포의 매끄러움을 조정하고, $z_i$는 클래스 $i$에 대한 모델의 스코어이다. 훈련 중에는 학생과 교사에게 똑같은 temperature $T$가 적용되지만, 추론을 할 때에 $T$는 기존의 softmax를 회복하기 위해서 1로 고정된다.

 

 마지막 단계는 distillation loss인 $L_{ce}$와 이 케이스에서는 masked language model의 loss인 $l_{mlm}$인 지도 학습의 오차의 선형 결합이다. 논문에서는 여기에 cosine embedding loss인 $L_{cos}$를 추가하는 것이 이득이라고 하는데, 이는 왜냐하면 교사와 학생의 hidden state vector의 방향을 조정해주는 경향이 있기 때문이다.

 

3. DistilBERT: a distilled version of BERT

Student architecture

 이전 절에서 했던 것처럼 학생의 입장인 DistilBERT는 BERT와 똑같은 구조를 가지고 있다. 하지만, Token-type embedding과 pooler가 두 개의 요소에 의해 layer의 수가 줄어듬에 따라서 제거되었다. Transformer architecture에서 사용되는 연산자들은 현대의 선형대수학에 알맞게 조정된 것이기 때문에, 본 연구에서는 마지막 차원의 텐서에서 변동을 주는 것은 계산 비용에 별로 영향을 미치지 않는다는 것을 알아내었다. 그보다는 layer의 수와 같은 요소에 변동을 주는 것이 영향을 더 끼치기 때문에, layer의 수를 줄이는 것에 집중하였다.

 

Student initialization

 앞에서 설명한 최적화 및 아키텍처 선택 외에도 교육 절차의 중요한 요소는 수렴할 sub-network에 대한 올바른 초기화를 찾는 것이다. 교사와 학생 네트워크 간의 일반적인 차원의 이점을 챙기기 위해, 논문에서는 두 레이어 중 한 레이어를 취하여 교사로부터 학생을 초기화하였다.

 

Distillation

 논문에서는 지금까지 나온 BERT중 가장 최고의 성능을 보여주는 모델을 사용하였다. 따라서 DistilBERT는 dynamic masking을 사용하고, next sentence prediction을 배제함으로써 gradient accumulation을 통해 매우 큰 크기의 배치로부터 distill된다.

 

 DistilBERT와 기존의 BERT의 성능 비교에 대한 표는 다음과 같다.

 

표 1. DistilBERT와 기존의 모델들의 성능 비교

 

4. Ablation study

 이 절에서는 triple loss의 세 개의 요소의 영향력과 distilled model의 student initialization에서의 성능에 대해서 조사하였다. 이에 대한 결과는 다음의 표에 나타나있다. 표를 살펴보면, Masked Language Model loss를 제거하는 것이 다른 두 loss를 제거하는 것보다 덜 영향을 끼친다고 결론을 내릴 수 있다.

 

표 2. triple loss에서 각 loss를 제외했을 때의 성능 변화 비교

 

 

 

 

참고문헌

https://arxiv.org/abs/1910.01108

DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter