Self Supervised Learning를 여행하는 히치하이커를 위한 안내서
저는 요즘 모두의 연구소에서 Self Supervised Learning(SSL)을 공부하고 SSL에 대해 논문을 쓰는 것을 목표로 하는 SSL LAB(쓸랩)의 연구원으로도 활동하고 있습니다. 벌써 2개의 논문을 끝냈습니다. 공부하고 있는 논문들이 SSL을 공부하는 좋은 로드맵이 될 것 같아 살짝 정리를 해보는 것이 좋을 것이라 생각했습니다. 논문 설명과 논문에서 공헌한 부분의 코드를 설명하는 시리즈가 될 것 같습니다.
미리 알아두면 좋은 지식
- Distilling the Knowledge in a Neural Network
: 지난 번 SSL 히치하이거 시리즈에서 소개한 논문입니다. Knowledge Distillation을 소개하는 논문이고 DistilBERT 역시 큰 뼈대는 같은 방식을 사용하고 있기 때문에 한 번 읽어보시는 것이 좋습니다.
- BERT
: 2018년 구글에서 소개한 NLP 분야의 pre-trained 모델입니다. Transformer의 인코더를 쌓아올린 아키텍처를 가지고 있습니다. BERT-base는 110M개의 파라미터, BERT-Large는 무려 340M개의 파라미터를 가지고 있습니다. 해당 논문은 BERT의 모델을 변형한 DistilBERT를 Student로 사용하여 학습하는 방법을 이야기하기 때문에 BERT의 아키텍처(Embedding, Transformer, pooler 등)에 대해 아는 것이 좋습니다.
DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter
[그림 2]에서 볼 수 있듯 거대한 사이즈의 pre-trained 모델들이 공개되고 있고, 그것을 이용해 transfer learning을 하는 것은 NLP 분야에서 일반적인 과정이 되었습니다.
이러한 모델에는 몇 가지 문제가 있습니다.
- 계산(computational) requirements가 많이 필요하다는 것
- On-device로 모델을 이용하게 된다면 새롭고 흥미로운 언어 처리 응용 프로그램을 가능하게 할 가능성이 있지만 계산 및 메모리 요구 사항이 모델 채택을 방해한다는 것
그래서 DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter 논문에서는 Knowledge Distillation 방법을 적용해 기존의 BERT의 97% language understanding capacity를 유지하면서 모델 사이즈의 40%를 줄이고 60%는 더 빠른 모델을 만드는 법을 제안했습니다.
Architecture
우선 [그림 3]을 보며 지난번에 이야기했던 Distilling the Knowledge in a Neural network의 아키텍처와 학습과정을 먼저 떠올려봅시다. 우선 pre-trained 모델인 Teacher와 Knowledge Distillation으로 학습할 Student가 있을 것입니다. Teacher를 freeze 시킨 상태로 Distillation Loss, Student Loss로 학습을 시킵니다. 그리고 Softmax를 계산할 때 사용하는 하이퍼파라미터 temperature T도 잊지말아야하죠.
DistilBERT도 이러한 과정을 이용해 학습되었습니다. Teacher, Student, Loss function이 어떻게 되어있는지 비교해봅시다. Teacher는 기존의 BERT를 이용했습니다. 그리고 Student는 어떨까요? 분명 Teacher 모델인 BERT보다 더 작은 사이즈인 모델일 것입니다.
Student architecture 논문에서는 이 Student 모델인 DistilBERT를 기존의 BERT 모델에서 token type embedding과 pooler 레이어를 제거했습니다. BERT는 Transformer의 인코더를 쌓아올린 구조를 가지고 있다고 했습니다. Transformer를 주로 이루고 있는 linear layer, layer normalisation 등의 레이어는 modern linear algebra frameworks 들로 이미 최적화되어 있기 때문에 layer 갯수를 줄이는 것이 dimension 을 computation efficiency에 큰 영향을 준다는 이유로 두 레이어가 제거되었습니다.
Student Initialization 논문의 저자들은 student architecture 선택외에도 학습이 잘되기 위해서 weight initialization이 잘 되어있어야 한다고 합니다. 따라서 Teacher 모델과 Student 모델의 dimension이 같다는 장점을 이용해 Teacher의 weight를 Student에 사용했습니다.
Training loss DistilBERT는 총 3개의 loss를 linear한 결합으로 사용하고 있습니다. $L_{CE}$, $L_{mlm}$, $L_{cos}$ 가 존재합니다. $L_{CE}$는 teacher의 softmax-temperature 값인 $t_i$, student의 softmax-temperature 값인 $s_i$으로 최종 $L_{CE} = \sum_{i}t_i ∗ log(s_{i})$를 계산합니다. Softmax-temperature는 Distilling the Knowledge in a Neural network 방식을 그대로 사용했습니다. $L_{mlm}$는 Student는 BERT 기반 모델이기 때문에 BERT의 masked language modeling loss를 의미합니다. 기존의 Knowledge Distillation 모델에서는 Distillation을 해주는 Loss, Student를 학습하는 Loss로 총 2가지 Loss만을 사용했는데 DistilBERT에서는 $L_{cos}$를 추가하여 총 3개의 loss를 가지고 있습니다. 나머지 한 개의 Loss의 역할은 Teacher와 Student 모델의 마지막 hidden state의 vector 방향을 align하게 해주는 효과가 있어서 추가했다고 합니다.
- Teacher : 기존 BERT
- Student : DistilBERT (기존 BERT에서 Token-type embedding 과 Pooler 제거)
- Loss : $L_{CE}$, $L_{mlm}$, $L_{cos}$
DistilBERT 결과 [그림 4] DistilBERT의 결과입니다. Table 1을 보면 기존 BERT랑 비교해서 모델의 97%의 퍼포먼스를 유지한다는 것을 알 수 있고 Table 2를 통해 Downstream Task의 퍼포먼스도 크게 차이나지 않는다는 것을 알 수 있습니다. 마지막으로 Table 3을 통해서 ELMo와 BERT-base보다 훨씬 더 적은 파라미터를 가지고 inference time이 더 빠르다는 것을 알 수 있습니다.
DistilBERT Code
DistilBERT Loss, BERT와의 다른 점을 코드를 보며 이해해봅시다. 이곳의 코드를 참고했습니다. 지난번에 Knowledge Distillation을 하기위해 Teacher, Student, Distiller(이곳에 두 모델을 보내서 학습)을 만들었습니다. 아래의 코드는 train.py의 일부입니다.
[코드 line 1 -15]
Student 모델을 만들고 Student Initialization을 하는 것을 볼 수 있습니다. 그리고 마지막 줄 15에서 device에 모델을 로드합니다.
[코드 line 17 -21]
Pre-trained된 Teacher 모델을 만들고 device에 모델을 로드합니다.
[코드 line 23 - 27]
freeze_pos_embeddings는 Student 모델이 ['roberta', 'gpt2'] 일 때, freeze_token_type_embeddings는 Student 모델이 ['roberta'] 일 때만 해당되므로 넘어가겠습니다.
[코드 line 29 - 34]
Teacher 모델과 Student 모델의 configuration이 동일한지 체크합니다.
[코드 line 36 - 43]
Distillation 과정을 가지고 있는 Distiller 클래스에 Teacher 모델, Student 모델, 데이터 등을 넣고 학습을 시작합니다.
이제 Distiller 클래스에서 어떻게 DistilBERT가 학습되는지 확인해봅시다. 코드가 꽤 길지만 겁먹지 말고 필요한 부분을 찾아봅시다.
[코드 line 53 - 54]
Teacher, Student 모델 넘겨받기
[코드 line 81 - 94]
Masked language modeling loss인 $L_{mlm}$을 위한 과정
[코드 line 111 - 116]
Loss를 계산해줄 객체 생성, 아래에서 논문에서 나온 3개의 Loss를 어떻게 코딩했는지 살펴봅시다.
[코드 line 324 - 365]
[331 - 332] Student 모델은 train 상태로 Teacher 모델은 eval 상태로(freeze)
[334 - 365] Epoch 수만큼 학습시킵니다.
[코드 line 366 - 459]
loss_ce, loss_mlm, loss_cos를 제외한 나머지 loss와 alpha는 GPT2, roberta 모델일 때 필요하기 때문에 넘어가겠습니다.
[406 - 412] $L_{CE}$ 계산, $t_{i}$와 $s_{i}$를 먼저 구한 뒤 softmax with temperature를 사용하여 계산합니다.
[415 - 417] $L_{mlm}$ 계산
[429 - 443] Teacher와 Student 모델의 마지막 hidden state의 vector 방향을 align하는 목적으로 사용되기 때문에 hidden_states의 마지막 벡터 s_hidden_states, t_hidden_states로 Loss를 계산합니다.
모델을 직접 학습해보지는 않았지만 코드를 이용해서 논문에서의 내용이 어떻게 구현되어있는지 살펴보았습니다. DistilBERT 논문을 읽고 코드를 보며 느낀 점은 학습 방식이 Distilling the Knowledge in a Neural Network과 크게 달라지지 않았지만 논문의 목표는 Knowledge Distilation을 사용하여 BERT의 사이즈를 줄이고 비슷한 퍼포먼스를 유지하는 것이기 때문에 BERT에 대해 잘 모른다면 DistilBERT의 아키텍처와 코드를 이해하기 어렵다는 것입니다. 또한 Knowledge Distilation의 아이디어를 간단하게 어떤 모델이든 적용해볼 수 있겠다는 생각이 들었습니다.
모두의 연구소에서 진행하는
"함께 콘텐츠를 제작하는 콘텐츠 크리에이터 모임"
COCRE(코크리) 2기 회원으로 제작한 글입니다.
코크리란? 🐘
Reference
[그림 1] https://upload.wikimedia.org/wikipedia/commons/a/a0/Milky_Way_libya.jpg
[그림 2] DistilBERT 논문
[그림 3] https://intellabs.github.io/distiller/knowledge_distillation.html
[그림 4] DistilBERT 논문
'DL|ML' 카테고리의 다른 글
| SpaceNet Challenge 1 (0) | 2022.07.21 |
|---|---|
| Self Supervised Learning를 여행하는 히치하이커를 위한 안내서 (3) - Representation Learning withContrastive Predictive Coding (0) | 2022.06.29 |
| Self Supervised Learning를 여행하는 히치하이커를 위한 안내서 (1) - Distilling the Knowledge in a Neural network (0) | 2022.06.10 |
| Foursquare - Location Matching 컴피티션 캐글 노트북 번역 (0) | 2022.05.07 |
| model.fit()에서 벗어나기! (2) (0) | 2022.04.20 |
