본 논문은 Hugging Face에서 2019년에 작성한 논문이다.
논문 링크: https://arxiv.org/abs/1910.01108
1. Introduction
2017년부터 (논문이 작성된) 2019년까지 pre-trained language model에 대한 연구가 활발히 진행되어 왔다.
여러 논문들이 모델의 크기와 성능은 비례한다는 연구 결과를 제시해 온 바 있다.
이런 large 모델들은 다음과 같은 문제점을 가진다.
- 모델 계산 자원을 기하급수적으로 확장할 때 드는 환경 비용
- 제한된 연산 능력과 메모리 크기를 가지는 edge device에서의 사용 불가
본 논문에서는 훈련 비용이 작을 뿐만 아니라, 적은 연산 자원으로도 빠른 추론을 수행하는 모델 DistilBERT를 소개한다.
Knowledge Distillation에 대한 지식이 있는 분들이라면, 어떻게 모델이 학습되었는지 벌써 감이 오실 거라 생각한다.
이 용어가 생소하신 분들이라도 아래 설명을 보면 쉽게 이해할 수 있을 것이다!
DistilBERT 모델은 base 모델 BERT에 비해 크기가 약 40% 작으며, 추론 속도는 60% 더 빠르다고 한다.
저자들은 실험을 통해 DistilBERT가 여러 downstream task들에 적용 가능하다는 것을 보였고, edge device에서도 활용될 수 있다고 언급한다.
2. Knowledge Distillation
Knowledge distillation은 모델을 압축하는 하나의 기법이다.
크기가 작은 student 모델이 상대적으로 크기가 크고 좋은 성능을 보이는 teacher 모델을 모방하도록 학습하는 방법이다.
loss를 설명하기 전에 classificatrion에서 loss가 어떻게 계산되는지 짚고 넘어가 보겠습니다.
분류 모델의 결과가 아래와 같다고 가정해 봅시다.
| 레이블 | Dog | Cat | Pig |
| 예측값 (확률) | 0.1 | 0.7 | 0.2 |
| 실제값 (원-핫 인코딩) | 0 | 1 | 0 |
이 경우엔 제일 높은 예측값을 갖는 Cat을 최종 결괏값으로 내보내겠지요.
분류 문제의 경우 cross-entropy 수식을 loss function으로 사용해서 모델을 학습시키곤 합니다.
쉽게 말해 정답 레이블(1)인 경우엔 확률값이 높아지도록, 그렇지 않은 경우(0)엔 값이 작아지도록 말이죠.
이러한 방법론을 염두에 두고, DistilBERT에서 사용한 3가지 loss function에 대해 살펴봅시다.
Loss 함수에 관한 부분은 링크 사이트를 참고하였습니다! https://alexnim.com/coding-projects-knowledge-distillation.html
(1) Distillation (Cross Entropy) Loss
어떤 자연어 텍스트가 들어간 후 BERT 모델은 마스킹된 영역에 들어갈 단어에 대한 확률을 계산합니다.
그림에 보시면 shop이 92%로 가장 높은 확률을 가지는 것을 알 수 있지요.
하지만 market 또한 빈칸에 들어가도 문장이 완성이 되겠지요?
한쪽으로 예측이 편향되는 모델을 만드는 것의 장점이 있겠지만, 자칫하면 오답인 결과를 생성할 가능성도 있겠지요?
이걸 방지하고자 나온 방법론이 바로 soft-max temperature입니다!
일반 soft-max 수식에 output 분포의 완만함을 결정해 주는 T라는 상수를 추가한 형태이지요.
상수 T값에 대한 확률 분포의 차이를 관찰해 봅시다.
확률값이 매우 높았던 shop은 조금 낮아졌지만, market에 대한 확률은 높아진 것을 확인할 수 있지요.
점점 부드럽게 예측한다-라는 느낌을 받을 수 있겠습니다.
저자들은 T를 3으로 설정해서 teacher 모델과 student 모델의 확률값을 계산했다고 합니다.
ti를 teacher 확률값, si를 student 확률값이라고 했을 때, 이를 cross-entropy 수식을 사용해 첫 번째 loss 값을 계산합니다.
이런 원리로 student와 teacher 모델의 결과 차이가 작아지는 방향으로 학습이 진행됩니다.
(2) Supervised Training (Masked Language Modelin) Loss
식 자체는 cross-entropy를 사용하기 때문에 (1)과 유사합니다.
다만, student 모델의 결괏값에 대해선 상수 T를 1로 설정하여 일반적인 softmax 확률값을 구하고,
이를 일반적인 class label (0, 1, 0, 0, 0..)과 비교하며 loss를 계산합니다.
(3) Cosine Embedding Loss
수식은 1 - cos(T(x), S(x)) 일 것인데........ 자세한 설명이 없다.
내가 예측한 바에 따르면, student와 teacher 모델 각각 최종 word가 나올 것이다.
두 word를 vector로 변환한 후, cosine similarity를 구해서 이를 최소화하는 방향으로 학습시킨 것이 아닐까 싶다.
3. DistilBERT: a distiled version of BERT
DistilBERT의 구조, 초기화 방법, 증류 방법, 훈련 상세 정보에 대해 알아보도록 합시다.
Student Architecture
모델 자체의 일반적인 구조는 BERT와 같고, 몇몇 차이점은 아래와 같다.
- token-type embedding 제거: 각각의 단어들이 첫 번째 문장에 해당하는지, 두 번째 문장에 해당하는지 표시하는 token_type_ids를 사용하지 않았다는 뜻이다.
- pooler 제거: BERT의 마지막 Transformer 레이어의 출력을 입력으로 받아서 이를 하나의 벡터로 변환하는 과정을 제거한 것이다. 보통 pooler는 해당 레이어의 출력을 가져와서 선형 변환을 적용하여 [CLS] 토큰의 임베딩을 만듭니다.
- 레이어 1/2: BERT-base의 경우 12개의 레이어를 가진다. DistilBERT는 그의 절반인 6개의 레이어로만 이루어져 있다.
Student Initialziation
모델의 가중치를 초기화하는 것 또한 학습에서 중요한 부분이다.
Student 모델은 Teacher 모델의 각 층 중 하나씩을 제거하여 생성했다고 한다.
Distillation
BERT 모델을 distillation 하는 과정에서, 한 배치당 약 4천 개의 데이터를 할당하여 gradient accumulation의 장점을 학습에 활용했다고 한다.
Data and Compute Power
BERT와 동일하게 English Wikipedia와 Toronto Book Corpus 데이터셋을 사용해 훈련을 진행했다.
16GB V100 GPU 8개로 학습했을 때 약 90시간(약 4일)이 걸렸다고 한다.
32GB V100 GPU 1024개로 하루동안 학습한 RoBERTa 모델에 비하면 훈련 비용 또한 상대적으로 작다는 것을 알 수 있다.
4. Experiments
자 이제 성능이 어떤지 확인할 차례입니다!
General Language Understanding
GLUE (General Language Understanding Evaluation) 벤치마크에서 ELMo와 BERT-base 모델과 비교한 결과입니다.
9개 데이터셋에 대한 평균인 Score 열의 값을 봐도 BERT-base에 비해 크게 뒤처지지 않는 성능을 가지는 것을 볼 수 있습니다.
저자들은 40%나 더 적은 파라미터 개수를 가지고도 BERT의 97% 정도의 성능을 낼 수 있다는 것을 강조합니다.
Downstream Tasks
이 실험에선 감정 분류용 IMDb, QA 태스크에 사용되는 SQuAD 데이터셋을 사용했습니다.
- IMDb: BERT에 비해 accuracy가 0.6%밖에 떨어지지 않았습니다!
- SQuAD: 정답이 정확하게 일치하는지 나타내는 지표 EM과 F1-score를 평가지표로 사용했습니다. 이 지표들에서도 크게 뒤처지지 않는 성능을 보였으며, pre-training 단계에서만 distillation을 수행한 DistilBERT보다 fine-tuning 단계에서 한 번 더 distillation을 수행한 DistilBERT(D) 모델의 성능이 약간 더 우수한 것을 확인할 수 있습니다.
Size and Inference Speed
STS-B 데이터셋에서 동일한 CPU(Intel Xeon E5-2690 v3 Haswell @2.9 GHz)로 성능을 비교했을 때, 파라미터 수는 40%, 추론 속도는 60% 로 줄어든 것을 확인할 수 있습니다.
저자들은 또한 당시 최신 핸드폰 아이폰 7+ 에서 QA 태스크를 수행해 본 결과, DisilBERT가 BERT보다 71% 더 빨랐고, 모델이 전체 사용한 메모리 용량은 207MB뿐이었다고 합니다. (+ quantization으로 더 감소시킬 수 있다고 언급)
Ablation Study
3가지 loss 중에서 첫 번째 소개해드렸던 L(ce)가 제거되었을 때 성능이 가장 큰 폭으로 떨어졌네요.
파라미터를 teacher 모델의 것이 아닌 랜덤으로 초기화하였을 때 또한 score 값이 줄어든 것을 확인할 수 있습니다!
5. Conclusion
본 논문에서는 distillation 기법을 사용해 BERT 모델을 압축시킨 모델 DistillBERT를 제안하였습니다.
기존 BERT 모델에서 일부 요소를 제거한 DisillBERT는 Triple loss를 사용해서 모델의 가중치를 업데이트하는 등 다양한 방식으로 모델의 크기를 축소시켰다.
그 결과, BERT의 97%의 성능은 보존하면서 모델의 파라미터 수는 40% 줄이고, 추론 시간은 60% 단축한 모델을 만들 수 있었다.
이 논문은 '트랜스포머를 활용한 자연어 처리' 책에서 2장 텍스트 분류 섹션을 공부하며 알게 된 모델이다.
아무래도 단순히 코드만 실행하고 '이런 모델이 있구나~' 정도로 넘어가는 것은 공부의 의미가 퇴색될 것 같아서 논문을 읽게 되었다.
다음 글은 DistilBERT를 이용한 문장 감정 분류 코드에 대해 작성할 예정이다!
저의 글이 언제나 여러분께 도움이 되길 바라겠습니다 ~ ❣️
궁금한 점이 있다면 언제든 댓글 남겨주세요. 😍