Skip to main content

DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter

· 4 min read
Victor Sanh
연구원
Lysandre Debut
연구원
Julien Chaumond
연구원
Thomas Wolf
연구원

논문 정보​

  • 제목: DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter
  • 저자: Victor Sanh, Lysandre Debut, Julien Chaumond, Thomas Wolf (Hugging Face)
  • 학회/저널: arXiv
  • 발행일: 2020-03-01
  • DOI: arXiv:1910.01108
  • 주요 연구 내용: BERT와 같은 대규모 사전 훈련 모델을 경량화하기 위해 지식 증류(Knowledge Distillation) 기법을 사전 훈련 단계에 적용. 학생 모델(DistilBERT)이 교사 모델(BERT)의 동작을 학습하도록 언어 모델링, 증류, 코사인 거리 손실을 결합한 삼중 손실(triple loss)을 도입하여 훈련.
  • 주요 결과 및 결론: DistilBERT는 BERT보다 파라미터 수가 40% 적고, 추론 속도는 60% 빠르면서도 GLUE 벤치마크에서 BERT 성능의 97%를 유지함. 이를 통해 제한된 컴퓨팅 환경이나 모바일 기기에서도 효율적으로 동작할 수 있는 범용 언어 모델의 가능성을 입증.
  • 기여점: 기존의 태스크 특화적 증류 방식에서 벗어나, 사전 훈련 단계에 지식 증류를 적용하여 범용적으로 사용 가능한 작고 빠른 언어 모델을 성공적으로 개발함. 또한 삼중 손실 함수와 교사 모델의 가중치를 활용한 초기화 방법을 제안하여 효과적인 지식 전수를 이룸.

요약​

초록​

대규모 사전 훈련 모델을 활용한 전이 학습이 NLP 분야에서 보편화되었지만, 이러한 모델들은 컴퓨팅 자원이 제한된 환경에서 운영하기 어렵다. 이 연구는 사전 훈련 단계에서 지식 증류를 활용하여 더 작고 빠른 범용 언어 표현 모델인 DistilBERT를 제안한다. DistilBERT는 BERT 모델의 크기를 40% 줄이면서도 언어 이해 능력의 97%를 유지하고, 60% 더 빠르다. 이를 위해 언어 모델링, 증류, 코사인 거리 손실을 결합한 삼중 손실을 도입하여 교사 모델이 학습한 귀납적 편향을 효과적으로 활용한다. 이 작고, 빠르고, 가벼운 모델은 사전 훈련 비용이 저렴하며, 온디바이스(on-device) 환경에서의 활용 가능성을 실험으로 입증했다.

서론​

최근 NLP 분야에서는 BERT, RoBERTa와 같은 대규모 사전 훈련 모델이 표준 도구로 자리 잡았지만, 이 모델들은 수억 개의 파라미터를 가지며 점점 더 커지는 추세다. 이러한 경향은 모델 훈련에 필요한 막대한 계산 비용으로 인한 환경 문제와, 모바일 기기와 같은 엣지 디바이스에서의 활용을 저해하는 문제를 야기한다. 이 연구는 지식 증류를 통해 훨씬 작은 계산 비용으로 사전 훈련된 경량 언어 모델이 다운스트림 태스크에서 대규모 모델과 유사한 성능을 낼 수 있음을 보여주고자 한다.

배경​

지식 증류는 크기가 작은 '학생(student)' 모델이 크기가 큰 '교사(teacher)' 모델의 동작을 모방하도록 훈련하는 모델 압축 기술이다. 학생 모델은 교사 모델의 최종 출력 확률 분포(soft targets)를 학습함으로써, 단순히 정답 레이블(hard targets)만 학습할 때보다 풍부한 정보를 얻는다. 예를 들어, 교사 모델의 예측에서 정답 클래스 외 다른 클래스들이 갖는 미세한 확률 값들은 모델의 일반화 능력에 대한 정보를 담고 있다.

모델 아키텍처 / 방법론​

  • 핵심 구조/방법: 학생 모델인 DistilBERT는 BERT와 동일한 트랜스포머 아키텍처를 기반으로 하지만, 전체 레이어 수를 절반으로 줄였다. 또한, 토큰 타입 임베딩과 풀러(pooler)를 제거하여 모델을 더욱 경량화했다. 학생 모델의 가중치는 교사 모델(BERT)의 레이어를 두 개 중 하나씩 가져와 초기화하여 안정적인 수렴을 유도했다.

  • 주요 구성 요소: 훈련은 교사 모델의 지식을 학생 모델에 효과적으로 전달하기 위해 세 가지 손실 함수를 결합한 삼중 손실(triple loss)을 사용한다.

    1. 증류 손실 (LceL_{ce}): 교사 모델의 출력 확률 분포와 학생 모델의 출력 분포 간의 크로스 엔트로피. 온도(T)를 적용한 소프트맥스를 사용한다.
    2. 마스크드 언어 모델링 손실 (LmlmL_{mlm}): BERT의 사전 훈련 목표와 동일하게, 마스킹된 토큰을 예측하는 손실.
    3. 코사인 임베딩 손실 (LcosL_{cos}): 학생 모델과 교사 모델의 은닉 상태 벡터 방향을 정렬하여 유사해지도록 만드는 손실.

  • 수식:

    • 소프트맥스-온도: pi=exp(zi/T)∑jexp(zj/T)p_{i}=\frac{exp(z_{i}/T)}{\sum_{j}exp(z_{j}/T)}
    • 증류 손실: Lce=∑iti∗log(si)L_{ce}=\sum_{i}t_{i}*log(s_{i}) (여기서 tit_i는 교사, sis_i는 학생의 확률)
    • 최종 손실 함수는 LceL_{ce}, LmlmL_{mlm}, LcosL_{cos}의 선형 결합으로 구성된다.

  • 알고리즘: DistilBERT는 BERT와 동일한 코퍼스(English Wikipedia, Toronto Book Corpus)를 사용하여 사전 훈련되었다. 훈련 시 동적 마스킹(dynamic masking)을 사용하고, 대규모 배치(최대 4K)를 활용했다.

실험 결과​

  • 주요 데이터셋: 범용 언어 이해 능력은 GLUE 벤치마크로 평가했으며, 다운스트림 태스크 성능은 IMDb 감성 분류와 SQuAD 1.1 질의응답 데이터셋으로 측정했다.

  • 핵심 성능 지표:

    • 성능: 논문의 Table 1에 따르면 GLUE 벤치마크 종합 점수에서 DistilBERT는 77.0점을 기록하여 BERT-base(79.5점) 성능의 97%를 달성했다.
    • 크기 및 속도: 논문의 Table 3에서 볼 수 있듯이 DistilBERT는 6,600만 개의 파라미터를 가져 BERT-base(1억 1,000만 개)보다 40% 작다. CPU 추론 시간은 60% 더 빠르며, 아이폰 7 플러스에서는 71% 더 빠른 속도를 보였다.

  • 비교 결과:

    • IMDb 데이터셋에서 DistilBERT는 92.82%의 정확도를 보여 BERT-base(93.46%)와 불과 0.6%p 차이를 보였다.
    • SQuAD 데이터셋에서는 BERT-base보다 약 3.9점 낮은 F1 점수를 기록했다.
    • Ablation study 결과(논문 Table 4), 증류 손실(LceL_{ce})과 코사인 임베딩 손실(LcosL_{cos})이 모델 성능에 가장 큰 영향을 미쳤으며, 교사 가중치로 초기화하지 않을 경우 성능이 크게 하락했다.

결론​

이 연구는 사전 훈련 단계에 지식 증류를 적용하여, 기존 대규모 모델의 성능을 상당 부분 유지하면서도 크기는 40% 줄이고 속도는 60% 향상시킨 DistilBERT 모델을 성공적으로 개발했다. DistilBERT는 제한된 자원을 가진 환경, 특히 온디바이스 애플리케이션에 매우 적합한 실용적인 대안임을 입증했으며, 이는 효율적인 NLP 모델 연구에 중요한 방향을 제시한다.