DistilBERT : BERT의 파생, BERT의 경량화

DistilBERT 배경 설명

•  DistillBERT는 Huggiing Face에서 2019년에 낸 논문으로, knowledge distillation을 사용하여, BERT의 문제점으로 지적되던, 큰 parameter 문제를 극복하기 위한 연구이다.
• 실제로 BERT를 실생활에서 이용할 때, 속도와 메모리 때문에 거의 DistilBERT를 많이 이용하는 것 같다.  
• 기존 BERT에 비해서 속도와 메모리 면에서 뛰어난 성능을 보이는 DistillBERT

2018년, 자연어처리 영역의 위대한 한 획을 그은 BERT를 대표로하여, 자연어처리 영역은 대규모의 사전학습 언어모델(large-scale pre-trained langauge models)을 활용한 전이학습(transfer learning) 방식이 주를 이루고 있다.

당연하게도 언어모델이 커질수록 많은 파라미터들을 사용하게 되며, 그에 따라 훌륭한 성능을 달성할 수 있게 된다.

 

 

우선, 주요 언어모델들의 파라미터와 이에 대한 사전학습 비용을 가볍게 살펴보자.

이제는 국민 언어모델이 된 BERT를 먼저 보자. BERT-base는 약 110,000,000(1억 1천만)개의 파라미터를, BERT-large는 약 340,000,000(3억 4천만)개의 파라미터를 가지고 있다. Google research에 따르면, BERT-large를 사전학습 하는 데 16개의 Cloud TPU로 4일이 꼬박 걸렸다고 한다. Cloud TPU v2로 가정했을 때, 16(TPU 수) * 4(학습 일) * 24(시간) * 4.5(시간 당 US)=)=6,912로 계산된다. 즉, BERT-large의 사전학습 비용을 원화로 환산하면 대략 890만 원 정도라고 볼 수 있다.

 

 

생성 모델인 GPT-2는 어떨까? GPT-2는 약 1,500,000,000(15억)개의 파라미터를 가지고 있다. The Register에 따르면, GPT-2의 학습에는 시간 당 $256 만큼 소요되는 256개의 Google Cloud TPU v3 cores를 사용했다고 한다. 

이후 나온 GPT-3와 현재 화두가 되고 있는 ChatGPT의 경우, 175,000,000,000(1750억)개의 파라미터를 가지고 있다. 앞서 살펴본 모델과 단순 비교만 하더라도 실로 엄청난 규모이다. 참고로 GPT-3를 학습하는 데 약 1200만 달러, 즉, 150억 원 정도가 소요되었다.

 

 

DistilBERT의 등장 배경

위와 같이 DistilBERT의 등장 이후에도 여전히 언어모델은 일반 기업들은 손대지도 못 할 정도의 초거대 모델로 개발되고 있다. 지금도 여전한 문제이기도 한 언어모델의 거대화는 DistilBERT의 등장 배경이기도 하다. 이와 함께 논문에서 언급한 내용들은 아래와 같다.

• 거대 언어모델을 학습하기 위한 연산에는 상당한 carbon footprint, 즉, 환경 비용(environmental cost)이 소요됨
• 거대 언어모델의 연산과 메모리 소요 등을 고려했을 때, 디바이스 상에서 실시간으로 사용되기 어려움

따라서 본 논문은 지식 증류(knowledge distillation) 기법을 활용하여 언어모델인 BERT를 경량화 하는 방식을 제시한다.

 

 

지식 증류 (Knowledge Distillation)

---

지식 증류(knowledge distillation)란, 이미 사전학습 되어있는 대규모 모델인 teacher로부터 경량화된 압축 모델인 student로 AI의 지식을 나누어 주는 개념이다. 증류(distillation)라는 단어가 액체 혼합물을 가열하여 액체 혼합물을 분리하는 과정을 의미한다는 점을 생각하면 그 뜻이 쉽게 와닿는다.

 

02-1. Softmax Temperature

기본 Softmax의 한계

일반적인 지도학습의 분류 모델은 모델의 예측 결과(logit)에 softmax를 취한 분포와 원 핫 인코딩 된 정답의 분포 간의 크로스 엔트로피 로스를 최소화 한다. 이에 따라 잘 학습된 모델은 정답 클래스에 대해서는 높은 확률로 분류해 낼 것이고, 오답 클래스에 대해서는 0에 가까운 값(near-zero)을 보일 것이다.

하지만, 같은 non-zero의 오답 클래스라 하더라도 어떤 클래스는 조금이나마 더 정답 클래스와 유사할 것이다. 극단적인 예를 들면, 바둑이를 분류하는 태스크에서 멍멍이라는 오답이 스파게티라는 오답보다 더 강아지와 유사하지 않은가? 여기서 ‘멍멍이’와 같은 지식을 암흑 지식(dark knowledge)라고 한다. softmax로 잘 학습된 분류기는 바둑이에 대해서는 0.99에 가까운 확률로 정답으로 분류를 하고, 멍멍이에 대해서는 0.000…1에 가까운 확률로 오답으로 분류할 것이다. 무언가 완벽해보이지 않는다. 즉, 단순한 softmax는 엔트로피가 낮다!

Softmax에 랜덤성을 더한 Temperature Scaling

 

이처럼 학습 시 temperature scaling을 적용한 확률을 soft target probability라고 한다. 학습을 마친 후 추론 시에는 scaling을 하지 않고, 기본 softmax를 사용한다. 이러한 방식은 generalization에 효과적이어서 테스트 데이터셋에 대해 더욱 좋은 성능을 보인다고 알려져 있다.

 

DistilBERT의 디테일

Next Sentence Prediction(NSP) 태스크를 수행하지 않음

DistilBERT의 student 모델은 기본적으로 BERT와 동일하나, 다음과 같은 차이가 있다.

• token-type embedding의 제거
  • token-type embedding이란, transformers 라이브러리의 token_type_ids에 대응되며, BERT에서 [SEP] 토큰으로 두 문장 간 관계를 학습할 때 각 문장을 구분함
• pooler의 제거
  • pooler란, BERT에서 얻은 contextual embedding 중 [CLS] 토큰의 임베딩에 해당됨

이러한 차이는 기존 BERT가 사전학습 할 때 수행하는 Next Sentence Prediction(NSP) 태스크를 수행하지 않음을 의미한다.

 

 

distilBERT는 distillation 기법으로 빠르게 학습하였음에도 불구하고,

기존 BERT에 비해 40%나 가벼워졌고 60%나 빨라졌으며 97% 성능을 유지했다.