Introduction
Model Compression 은 수많은 파라미터로 이루어진 네트워크 내에 내재되어 있을 불필요한 파라미터들을 제거하여 효율적으로 좋은 성능을 내기 위한 기술입니다. 일반적으로 Language Model(LM)에서는 baseline으로 pre-training 이후 fine-tuning을 하는 방식을 사용하고 있습니다.
하지만 본 논문에서는 이러한 2-staged approach는 compact한 모델이 end task에 집중할 때 더 잘 동작한다는 가정을 기반으로 하기 때문에 제한된 일반화만이 가능하다는 점을 지적하였습니다. 이러한 문제를 해결하고 더 좋은 성능을 얻기 위해 추가적으로 model compression 방법 중 하나인 distillation을 적용하는 방식을 제안하였습니다.
pre-training과 distillation을 함께 사용하는 방식은 이전에도 제안이 되었지만 일반화가 어려웠고, 본 논문에서 제안한 방식은 이러한 점을 해결했다는 점에서 차이가 있습니다.
본 논문에서 제안한 Pre-trained Distillation 알고리즘은 위와 같습니다. 저자들은 기존의 2-stage 사이에 잘 학습된 모델의 출력을 라벨로 활용하여 출력의 차이를 줄이는 방식으로 학습하는 distillation을 추가하였습니다. 다양한 조건에서의 실험을 통해 제안한 알고리즘이 기존의 baseline 보다 대부분의 경우에서 좋은 성능을 보여주는 것을 확인할 수 있었습니다.
Contribution
- Pre-trained Distillation 을 통해 성능 향상을 이끌어냄
- 기존의 baseline과 본 논문에서 제안한 방식이 다양한 조건에서 어떠한 성능을 보여주는지 실험을 통해 확인
Background
Knowledge Distillation
Knowledge Distillation은 잘 학습된 모델의 지식을 다른 모델로 전달하기 위한 기술입니다. 복잡한 모델의 모든 파라미터가 유효한 것은 아니라는 가정에서 시작하여 복잡한 모델의 출력과 상대적으로 단순한 구조를 가지는 모델의 출력이 같아지도록 학습시키는 방식입니다. 일반적으로 지식을 전달하는 모델을 Teacher, 전달받는 모델을 Student 라고 합니다.
Patient Knowledge Distillation
단순히 최종 출력의 softmax 사이의 loss를 구하는 것 외에 중간 결과의 loss 까지 줄이는 방식으로 학습시키는 방식입니다. Student 모델은 Teacher 모델의 일부 레이어를 복사해오는 것을 통해 초기화되고, 복사한 레이어의 출력에 대한 loss를 줄이는 방식으로 학습합니다.
DistillBert
Teacher 모델의 일부 레이어를 복사해온 다음 distillation, language modeling, cosine distance 의 loss를 더한 값을 학습에 사용하는 방식을 사용합니다. 레이어를 복사해온 다음에는 LM Knowledge Distillation으로 pre-training을 이어나간 다음 fine tuning을 수행합니다.
Proposed Method
Pretrained Distillation
1. Masked Language Model(MLM) 을 통해 pre-training을 수행합니다. MLM은 input에서 무작위로 token을 mask 시킨 다음 mask 된 token을 예측하는 모델입니다.
2. Teacher 모델의 출력을 soft label로 사용하여 distillation을 수행합니다. 출력의 값만 활용하는 일반적인 Knowledge Distillation 을 사용합니다.
3. 마지막으로 라벨이 있는 데이터를 사용하여 fine tuning을 수행합니다.
Comparison to concurrent works
Difference as distillations
최근의 distillation 방식은 Teacher 모델의 레이어를 복사하는 방식을 사용하기 때문에 구조적으로 자유롭지 못합니다. 반면에 본 논문에서 제시한 방법은 레이어의 weight를 복사하는 대신 pre-training을 수행하기 때문에 모델이 더 자유로운 구조를 가질 수 있습니다.
Performance comparison in GLUE leaderboard
pre-training + distillation 방식을 통해 학습했을 때 가장 좋은 성능을 보여주는 것을 확인할 수 있습니다.
- GLUE Dataset(optional)
- MNLI(Multi-Genre Natural Language Inference): entailment classification task
- QQP(Quora Question Pairs): Quora에 올라온 질문 페어가 의미적으로 동일한지 확인
- QNLI(Question Natural Language Inference): paragraph의 answer 포함여부 확인
- SST-2(Stanford Sentiment Treebank): 단문장 이진분류문제
- MRPC(Microsoft Research Paraphrase Corpus): 문장쌍의 유사성 확인
- RTE(Recognizing Textual Entailment): MNLI와 유사하나 데이터가 적음
Analysis
다양한 조건의 Pre-training Distillation(PD) 을 적용했을 때의 모델의 복잡도에 따른 정확도를 보여주는 그래프입니다.
- word embedding에 대해서만 적용했을 때 단순히 distillation에만 적용한 것보다는 좋지만 레이어까지 적용한 것에 비해서는 좋은 결과를 얻지는 못했습니다.
- 다른 주목할 점은 PD를 적용한 모델을 truncate 했을 때 얕은 모델일수록 성능이 떨어지는 것을 확인할 수 있습니다.
Student 모델의 24가지의 depth와 width 조합에 따른 성능을 보여주는 그래프입니다.
- 모든 구조에서 PD가 월등한 성능을 보여주는 것을 확인할 수 있습니다.
- pre-training의 효과는 parameter의 수가 많을수록 크게 성능이 좋아는 것을 확인할 수 있습니다.
다양한 데이터셋에 대해 서로 다른 방식을 적용했을 때의 성능을 보여주는 그래프입니다.
- 다른 데이터셋에 대해서도 본 논문에서 제안한 방식의 성능이 가장 좋은 것을 확인할 수 있습니다.
- Transfer 데이터셋의 크기와 PD의 성능이 비례하는 것을 확인할 수 있습니다. 주목할 점은 Transfer 데이터셋의 크기가 커졌을 때 Distillation이 PF 보다 좋은 성능을 보여주는 것을 확인할 수 있습니다.
(MNLI, RTE: < 1.3m, SST-2: 1.7m, Amazon Book Reviews: 8m)
Transfer 데이터의 크기에 따른 성능을 보여주는 그래프입니다.
- 복잡한 모델에 distillation 을 적용하는 것보다 단순한 모델에 본 논문에서 제안한 방식을 적용했을 때 더 좋은 성능을 보여주는 것을 확인할 수 있습니다.
- 데이터셋의 사이즈가 증가할수록 Teacher 모델의 성능에 가까워지는 것을 확인할 수 있습니다.
서로 다른 도메인의 training set과 transfer set을 사용하여 학습했을 때 distillation이 어떻게 영향을 받는지 확인하기 위해 Spearman Rank Correlation Coefficient (두 데이터셋 사이의 의존성을 수치화한 것) 에 따른 성능을 보여주는 그래프입니다. transfer set의 크기를 고정하고 distillation에 사용되는 데이터의 종류를 변경하면서 실험했다고 합니다.
- 어느 정도의 의존성을 가진다면 다른 방식에 비해 월등한 성능을 보여주는 것을 확인할 수 있습니다.
pre-training에 사용하는데이터와 transfer 데이터를 동일하게 했을 때의 성능을 보여주는 그래프입니다.
- pre-training 과 distillation이 독립적으로 성능 향상에 도움이 되는 것을 확인할 수 있습니다.
Conclusion
일반적인 학습 방법에 distillation 방법을 추가한 단순한 방법이지만 다양한 조건에서의 실험을 통해 각각의 요소들이 어떤 영향을 미치는지 확인할 수 있었다는 점에 의미가 있는 논문인 것 같습니다.
새로운 모델을 학습시키거나 가벼운 모델이 필요할 때 최대한 복잡한 모델로 학습시키고 본 논문에서 제안한 방식을 참고하여 최적화를 시키면 도움이 될 것 같습니다.
Share article