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

Abstract

고성능의 NLP model 을 만드는 것은 비용(시간, 하드웨어)이 많이 듬 (성능이 모델의 크기와 비례)

NLP 에서 Transfer learning 을 사용하여 큰 모델을 사용하는걸 보편적으로 볼 수는 있지만 

model 을 deploy 하기 위해서는 제한된 환경이 주어질 경우가 많음

이 논문에서 제한하고 있는 Distil BERT 는 기존 BERT 의 크기 40% 감소, 속도는 60% 빠르며

그럼에도 불구하고 97% 의 성능을 유지하고 있음

 

 

1. Introduction

최근 2년 동안 NLP 에서 Transfer Learning 의 접근이 많아짐

(large-model 을 pretrained 하고 특정 task 에 맞춰 fine tuning 하는 방식)

 

 

 

초기 자연어 처리의 문제

[1] Gradient vanishing : 데이터가 길어짐에 따라 기울기가 사라짐

[2] 다음 입력 데이터 처리를 위해 이전 데이터를 필요로 하여 병렬화가 어려움

computer vision 에서는 pre-trained data 와 새로운 input 데이터가 비슷한 형태를 가진다면 일반적으로 trasnfer learning 이 잘 수행되었다. (CNN 은 특징을 잘 추출하는 방법이 사전에 많이 연구됨)

[3] 자연어 처리에서는 Corpus 에서 word-embedding 을 학습하여 재사용하는 정도에서 그침

 

2018 년도에 transformer 를 사용한 GPT, BERT를 발표하며

거대한 모델의 pre-train - finetuning 과 다양한 문제들에서 뛰어난 성능을 보인다는 것을 보여줌

 

 

 

large model 의 문제

pretrained 된 large model 을 사용하여 성능은 확보할 수 있지만 real-time 에서 추론을 해야하는 경우, deploy 를 할 경우에 computational, memory 문제가 발생하게 됨

이를 해결 하기 위해 본 논문에서 제시하고 있는 방법은 pretrained model 에 knowledge distillation 의 적용

triple loss 를 사용하여 40% 더 작고 60% 더 빠르며 성능에는 큰 차이가 없는 Transformer 를 얻음

 

 

 

2. Knowledge distillation

2.1 Knowledge distillation

---

[참고] https://intellabs.github.io/distiller/knowledge_distillation.html#knowledge-distillation

큰 모델에서 Output(softmax) 는 해당 모델이 학습한 지식을 담고 있을 것이다. 라는 가설에서 시작함

모델의 크기는 작지만 성능을 최대한 유지하는 것이 distillation 의 목표 (증류 : 농도를 진하게 만드는 것)

teacher model 은 사전에 학습이 완료된 모델임 (teacher 가 존재하는 상태에서 student 를 학습)

SoftLabels 와 softPrediction 이 유사하게 되도록 학습하게 됨 (qi) 각각은 T 를 나눠주는 과정이 포함되는데 분포를 완만하게 해줌 (정보를 풍부하게 해줌)

주로 teacher 모델의 softmax 함수의 경우 높은 확률로 올바른 class 를 가지고 있을 것이고 다른 클래스는 0에 가까운 값을 가지는 확률 분포일 것 (이미 제공된 정답 레이블 이상의 정보를 제공하지 않음)

분포가 완만 -> information 이 풍부 -> large model 일 수록 풍부한 정보를 갖고 있을 것이라는 가정 (teacher(freezed) 에서 각 block 에 T 를 나누어줌)

즉, T 가 높을 수록 information 이 더 풍부해짐

비슷하다고 예측한 class 에 대해 더 많은 정보를 줄 것 ( dark knowledge )

dark knowledge 를 통해 student model 을 transferring 하는 것이 distillation 의 동작 방식임

student 의 softmax 를 구할 때에는 teacher 과 동일한 T 를 사용함 (이때 발생하는 손실 "distillation loss")

 

teacher 의 soft label 외에도 학생의 예측된 확률과 실제 label(hard label) 사이의 손실을 계산

학생 손실에 대한 class 확률을 계산할 때에 T = 1

 

따라서 전체 손실 함수는 아래와 같이 구성됨

왼쪽 : 실제 label(hard label) 과 예측된 student 사이의 손실 계산(T = 1)

오른쪽 : teacher 의 softmax 와 studnet 의 softmax 사이의 손실 계산 (T = t) 

실험적으로 a 의 가중을 더 작게 하였을 때 성능이 더 잘 나옴

(위의 blog 기반으로 정리한 내용)

---

 

본 논문에서 student 는 teacher 의 soft target probability 를 배움

학생은 선생의 soft target probability 를 배움

기존의 CrossEntropy 를 이용한 학습은 모델의 예측 확률 분포를 정답 one-hot 분포에 맞춰 학습하게 됨

앞서 말했든 피팅이 잘 된 모델이라면 특정 class 의 확률은 높고 나머지는 거의 zero 에 가까운 확률 분포가 출력됨

이 때에 모델의 generalization 능력에 기여하는 부분이 near-zero 부분임

 

teacher model 이 출력하는 확률 분포 자체를 training 해 student model 은 자신보다 복잡한 모델들만이 배울 수 있는 signal 또한 배울 수 있음

 

 

2.2 Training loss

Training objective = Distillation Loss(CE) + Masked Language Modeling Loss + Cosine Embedding Loss

Distillation Loss 와 Cosine Embedding Loss 를 학습하는 동시에 studnet 의 MLM 을 통한 LM 도 학습하게 됨

최종 훈련의 목표는 Distillation Loss 와 supervised training 의 training loss(MLM) 의 조합 (선형)

마지막의 Cosine Embedding Loss 는 student / teacher 의 hidden vector 가 바라보는 방향을 일치 시켜주는데 도움을 주는 loss 임

 

 

 

Cosine Similarity

[참고] https://wikidocs.net/24603

Cosine Similarity 는 두 벡터간의 Cos 각도를 이용하여 구할 수 있는 두 벡터의 유사도를 의미함

두 벡터의 방향이 완전히 동일한 경우에는 1, 90도의 각을 이루면 0, 180도로 반대 방향일 경우 -1의 값을 가짐

값이 1에 가까울수록 유사도가 높다고 판단함

cosine simillarity 공식

문서의 유사도를 구하는 공식 이 값이 1에 가깝도록 update

 

 

3. DistilBERT : a distilled version of BERT

3.1 Student architecture

DistilBERT의 student단은 BERT 와 동일한 아키텍처를 가짐

token형 embedding 과 pooler 는 제거하고 Layer 수는 2배 감소함

Transformer 아키텍처에서 마지막 차원에 대한 변동이 계산 효율에 미치는 영향이 작다고 판단 (레이어의 다른 요인의 변형보다) 따라서 레이어 수를 줄이는데 중점을 둠

3.2 Student initialization

최적화 및 아키텍처 선택 외에도 올바른 initialization 을 찾는것이 중요함

teacher 과 student network 간의 공통 차원을 활용하여 두 레이어 중 한 레이어를 가져와 teacher 로 부터 student 를 초기화 

3.3 Distillation

Liu et al. 에서 제안한 BERT 모델 훈련을 위한 모범 사례를 적용

DistillBERT 는 큰 batch 를 사용하여 기울기를 축적, dynamic masking 을 적용하며 NSP 를 제거함

3.4 Data and compute power

원래 BERT 모델과 동일한 corpus 에서 DistilBERT 훈련함

BERT 와 동일하게 8개의 16GB V100 GPU 에서 90시간 동안 training 함

RoBERTa 모델은 V100 에서 하루정도 train함

4. Experiments

General Language Understanding

일반 언어 이해 및 평가 benchmark 에서 DistillBERT 의 기능을 평가함

9개의 데이터 세트 사용 (fine-tuning 시에 ensembling 이나 multi-tasking 을 사용하지 않음)

DistilBERT 는 ELMo(BiLSTM) 과 비교하였을 때 9개의 task 에서 성능이 우세하였음

BERT 와 비교하였을 때에는 97% 정도의 성능을 보존하고 40% 정도 더 작은 parameter 을 사용함

 

4.1 Downstream task benchmark

Downstream tasks

분류 작업 (감정 분류 및 Q/A task) Table 2 에서 확인해본 결과 BERT 에 비해 0.6% 정도의 정확도 감소 40% 정도의 모델 크기 감소

SQuAD 에서 DistilBERT 는 full BERT 에 비해 3.9 points 정도의 성능 감소

 

본 논문에서는 fine-tuning DistilBERT 에서 adaptation phase 단계에서 (적응 단계) distillation 의 또다른 단계를 추가할 수 있는지 연구함

SQuAD 에서 fine-tune 된 BERT 모델을 사용하여 SQuAD 를 fine-tune DistilBERT 는 loss 의 추가적인 term 에 대한 teacher 로 사용됨

모델의 size 는 79.8 F1

전체 모델 크기의 3점 이내에서 성능을 가질 수 있었음

size and inference speed

DistilBERT 의 속도 향상 / 크기의 trade off 를 조사함

CPU 환경에서 batch size 는 1 로 실험함 DistilBERT 는 BERT 에 비해 40% 정도 작은 parameter 를 가지며 60% 정도 더 빠르다는 것을 보여줌

On device computation

DistilBERT 는 Q/A 기반의 app 등을 배포하는데 사용될 수 있음

BERT-base 와 DistilBERT 를 smartphone 환경에서 test 함

DistilBERT 가 BERT 보다 71% 정도 빨랐으며 모델의 크기는 207MB 정도였음(Q/A 에 한정)

4.2 Ablation study

3개의 loss 와 구성요소의 영향에 대해 조사함

GLUE에 대한 score 를 살펴보면

Masked LM loss 를 제거하는 것은 거의 영향이 없음

두가지의 distillation loss 가 성능에 큰 영향을 미침

5. Related work

5.1 Task-specific distillation

이전에는 대부분 task 별 distillation 설정을 구축하는 것에 중점을 둠

fine-tune classification model BERT 를 LSTM 기반의 분류기에 transfer 함

예를 들어 Chatterjee 는 BERT 에서 이전에 초기화된 더 작은 Transformer 모델의 SQuAD 에서 fine-tune 된 BERT 모델을 distillation 함

현재 본 논문에서 말했듯 이러한 task 별 distillation 보다 범용적인 distillation 을 사용하는 것이 더 좋다는 것을 확인

task 별 distillation이 아닌 pre-train distillation

더 작은 학생을 훈련시키기 위한 pre-train 목표는 distillation 을 통해 fine-tune 됨

(teacher 의 knowledge 를 활용하여 추가적으로 pre-trained 하는 것이 유익함, distillation signal 과 함께)

5.2 Multi-distillation

teacher 의 ensemble 지식을 결합하여 distillation 을 정규화 하기 위한 다중 작업 학습

대규모의 Q/A model 집합에서 간단한 Q/A model 을 학습함

유사하게 다중 언어 모델에서 distillation 을 사용하였고 ablation study 에서 초기화 및 추가적인 loss 를 통해 teacher 의 knowledge 를 활용하면 상당히 많은 성능 향상을 얻을 수 있었음

5.3 Other compression techniques

큰 모델을 압축하기 위한 연구

test 를 할 때에 self-attention 에서 일부 head 를 제거하는 것 (가중치에서 불필요한것들을 쳐내는 방법)

성능을 크게 저하시키지 않으면서 일부 레이어를 줄일 수 있음

별도의 연구에서는 양자화를 활용하여 더 작은 모델을 도출

 

6. Conclusion and future work

40% 더 작고 60% 더 빠르며 언어 이해 능력은 97% 를 유지한 DistilBERT

knowled distillation 을 통해 LM 을 효과적으로 학습하며 다양한 구성 요소를 분석할 수 있음

DistillBERT 가 edge 를 위한 매력적인 옵션임을 입증함

 

느낀점

knoweldge distillation 을 이해하고 적용한 것이 주된 논문의 내용이라고 생각함

다른 논문들에 비해 자체적으로 알고리즘을 설계하여 성능 향상을 얻은 case 는 아니었음

그럼에도 불구하고 확실히 experiment 를 통해 성능을 입증하였고 논문을 통해 다른 사람들의 시간을 많이 아낄 수 있었으므로 충분히 의미가 있는 논문이었다는 생각이 들었음

이렇게 Computer Vision 와 같이 NLP 와의 지식을 알고있는 경우 이를 NLP task 에 접목시켜 새로운 논문을 쓸 수도 있겠구나 생각하였음