BERT의 파생모델 [DistilBERT]

DistilBERT

• 사전 학습된 Large BERT의 성능은 확실히 뛰어나지만 그에 걸맞는 무게를 자랑한다. 그렇기에 downstream task를 위해 fine-tuning을 마친 모델 또한 부담스럽고 무거운 모델이 되므로 주어진 컴퓨팅 리소스 내에서 모델을 실행하기가 어려울 수 있다. 가령 모바일 환경에서 모델을 실행해야할 경우 말이다. 이 문제를 해결하기 위해 지식 증류라는 개념이 도입되었고, 이를 기반으로 제작된 모델이 DistilBERT이다.
• DistilBERT는 BERT에서 레이어 수를 줄이면서 성능 저하를 최소화 하는 방향으로 설계되었다고 한다. 결과적으로 BERT 모델 사이즈의 40% 수준으로 감소할 수 있었고, 97% 수준의 성능을 유지하면서 60% 가량 빠른 속도를 뽑아낼 수 있다고 한다.

https://arxiv.org/pdf/1910.01108.pdf

지식 증류란?

• 증류란, 간단히 말해서 물과 같은 용액을 가열해서 내부의 불순물만 쏙 빼내는 것을 말한다. 만약 여러 성분이 담긴 용액이라면 각각의 성분을 분리하는 것을 말한다.
• 이 개념을 large BERT에 적용해보면 커다란 모델이 학습을 통해 얻은 지식의 액기스만 빼낸 결과로 학습한 모델이 DistilBERT라고 이해할 수 있다.

 

HOW?

• DistilBERT의 학습 Flow는 간단히 말해서 BERT 모델이 선생님(the teacher)이 되어 학생(the student) 모델인 DistilBERT에게 지금까지 배운 내용을 전수해 준다는 개념이 담겨있다.
• → 실제 논문에서 선생님과 학생으로 구분해서 설명하고 있다.

 

Distillation Loss

Soft target

• 선생님의 교육은 MLM task의 예측 결과를 학생과 공유하는 방식으로 진행된다. 여기서 선생님의 MLM 출력 결과는 기존 softmax 대신 softmax-temperature로 불리는 함수를 거쳐 상대적으로 스무딩한 결과 분포를 예측하도록 한다.
• $P_i = \frac{exp(\frac{z_i}{T})}{\sum_jexp(\frac{z_j}{T})}$
• 위 함수는 기존 소프트맥스 함수의 z값을 T(Temperature)로 나눈 결과에 해당한다. 이를 적용할 경우 가령 T = 1의 결과(소프트맥스)가 [0.997, 0.000, 0.002, 0.001, 0.000] 이었다면 T = 5의 결과는 [0.637, 0.021, 0.191, 0.128, 0.021]가 된다.

https://arxiv.org/pdf/1910.01108.pdf

• 논문에서는 near-zero에 해당하는 값들 즉 정답 클래스에 해당하는 분포값 외 나머지 타겟값에도 작지만 위아래가 존재하고 이러한 분포값이 일반적인, 보편적인 모델 능력을 반영하고 있다고 본다. 그래서 smoothness를 좌우하는 T값을 적용하여 MASK 토큰의 예측 분포인 모든 레이블(토큰)의 분포를 학생이 학습하도록 한다.
• 사전 학습을 완료한 BERT를 예측하는 과정에서 softmax-temperature 함수가 적용된 결과를 Soft target이라 부르며 학생의 학습 목표 label이 된다.

Soft prediction

• 학생 모델은 아무런 학습도 받지 않은 BERT architecture 기반 모델이다. 앞에서 언급한 바와 같이 BERT 축소판으로 준비된다.
• 선생님 모델에 들어간 입력 데이터를 그대로 학생 모델에 넣어 출력 역시 softmax-temperature 함수를 최종적으로 적용한다. 이 결과를 Soft prediction이라 부른다.
• 최종적으로 Soft target과 soft prediction 간 cross-entropy를 loss function으로 적용하여 학생 모델의 weight를 업데이트 한다.(학습한다.) 이를 distillation loss라고도 부른다.

MLM Loss

• 학생 모델은 distillation loss와 더불어 기존 MLM 학습 방식도 병행해서 진행한다. 즉 익히 알고 있는 target label(one-hot)과 softmax 간 croos-entropy를 말한다.

Cosine Embedding Loss

• 마지막으로 소프트 타깃과 소프트 예측결과 (두 벡터) 간 코사인 유사도 측정을 통해 두 벡터가 같은 방향을 향하도록 조정하는 방식으로 모델을 학습한다.