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

Tasker_Jang·2024년 8월 7일

DistilBERT: 더 작고 빠르며 가벼운 모델

대규모 사전 학습된 모델을 활용한 전이 학습이 자연어 처리(NLP) 분야에서 점점 더 중요해지고 있다. 하지만 이런 큰 모델을 엣지 디바이스에서 사용하거나 제한된 컴퓨팅 자원 내에서 운영하는 것은 여전히 도전 과제로 남아 있다. 그래서 우리는 DistilBERT라는 더 작은 범용 언어 모델을 제안한다. 이 모델은 다양한 작업에서도 대형 모델처럼 훌륭한 성능을 발휘할 수 있도록 설계되어 있다.

기존 연구들은 주로 작업별 모델을 구축하기 위해 지식 증류를 사용했지만, 우리는 사전 학습 단계에서 지식 증류를 활용했다. 그 결과, BERT 모델의 크기를 40% 줄이면서도 언어 이해 능력의 97%를 유지하고 속도를 60%나 향상시킬 수 있었다.

또한, 더 큰 모델이 사전 학습에서 배운 귀납적 편향을 잘 활용하기 위해, 언어 모델링, 증류, 코사인 거리 손실을 결합한 삼중 손실(triple loss)을 도입했다. 이로 인해 DistilBERT는 사전 학습 비용이 적고, 개념 증명 실험과 비교 온디바이스 연구를 통해 온디바이스 연산의 가능성을 보여주고 있다.

도입부

최근 몇 년 동안 자연어 처리(NLP) 분야에서는 전이 학습이 대세가 되었다. 대규모 사전 학습된 언어 모델들이 많은 NLP 작업에서 기본 도구로 자리 잡았다. 이 모델들은 성능을 크게 향상시키지만, 보통 수억 개의 매개변수를 가지고 있다. 연구에 따르면, 더 큰 모델을 학습시키면 여전히 더 나은 성능을 보인다고 한다.

하지만, 이렇게 큰 모델을 사용하는 데는 몇 가지 문제가 있다. 첫째, 모델의 계산 요구가 기하급수적으로 늘어나면서 환경적 비용이 커지고 있다. 둘째, 이러한 모델을 디바이스에서 실시간으로 운영하면 새로운 언어 처리 애플리케이션을 만들 수 있지만, 계산 및 메모리 요구 사항이 증가하면서 널리 사용되기 어렵다.

물론입니다! 주어진 내용을 "~이다" 어투로 번역해 드리겠습니다.

Knowledge Distillation

지식 증류(Knowledge Distillation)는 작은 모델(학생)이 큰 모델(교사)이나 여러 모델을 조합한 것의 동작을 재현하도록 훈련하는 압축 기술이다.

지도 학습에서의 역할

지도 학습에서는 일반적으로 분류 모델이 올바른 클래스를 예측하도록 훈련하여 정답 레이블의 확률을 최대화한다. 표준 훈련 목표는 모델의 예측 분포와 원-핫(One-hot) 형태의 훈련 레이블 분포 간의 교차 엔트로피를 최소화하는 것이다. 훈련 세트에서 성능이 좋은 모델은 정답 클래스의 확률이 높고 다른 클래스의 확률은 거의 0에 가까운 분포를 예측한다. 하지만 이러한 "거의 0"인 확률 중 일부는 다른 것보다 크며, 이는 모델의 일반화 능력과 테스트 세트에서의 성능을 어느 정도 반영한다.

훈련 손실

학생 모델은 교사의 소프트 타겟 확률을 사용한 증류 손실로 훈련한다:

[ L{ce} = \sum{i} t{i} \times \log(s{i}) ]

여기서 ( t{i} )와 ( s{i} )는 각각 교사와 학생이 추정한 확률이다. 이 목표는 전체 교사 분포를 활용하여 풍부한 학습 신호를 제공한다.

[ p{i} = \frac{\exp(z{i}/T)}{\sum{j}\exp(z{j}/T)} ]

여기서 ( T )는 출력 분포의 부드러움을 조절하고, ( z_{i} )는 클래스 i에 대한 모델 점수이다. 훈련 시에는 학생과 교사 모두에게 동일한 온도 ( T )를 적용하며, 추론 시에는 ( T = 1 )로 설정하여 표준 소프트맥스를 사용한다.

최종 훈련 목표는 증류 손실 ( L{ce} )와 지도 학습 손실(우리의 경우 마스크드 언어 모델링 손실 ( L{mlm} ) [Devlin et al., 2018])의 선형 결합이다. 학생과 교사 은닉 상태 벡터의 방향을 정렬하려는 코사인 임베딩 손실 ( L_{cos} )를 추가하는 것이 유익하다는 것을 발견하였다.

"학생 모델"이라는 용어는 지식 증류(Knowledge Distillation)라는 기법에서 사용되는 개념입니다. 지식 증류는 큰 모델(교사 모델)에서 작은 모델(학생 모델)로 지식을 전이하는 기술입니다. 이 기술은 원래 모델의 성능을 최대한 유지하면서도 모델의 크기와 복잡성을 줄이기 위해 개발되었습니다.

지식 증류에서의 학생 모델

학생 모델(Student Model): 학생 모델은 교사 모델의 동작을 학습하도록 설계된 작은 모델입니다. 이 모델은 교사 모델보다 매개변수 수가 적고, 따라서 메모리 사용량이 적으며, 계산이 더 빠릅니다. 학생 모델은 교사 모델의 출력(예: 확률 분포)을 학습하여 원래 모델의 지식을 압축합니다.
교사 모델(Teacher Model): 교사 모델은 일반적으로 크고 복잡한 모델로, 높은 성능을 보여주지만 계산 자원과 시간이 많이 필요합니다. 교사 모델의 역할은 학생 모델이 학습할 수 있는 지식을 제공하는 것입니다. 이 모델은 데이터에 대한 복잡한 패턴을 학습하여 학생 모델이 따라야 할 목표를 설정합니다.

지식 증류의 과정

교사 모델 학습: 먼저, 교사 모델은 대량의 데이터를 사용하여 학습합니다. 이 모델은 복잡한 데이터 패턴을 이해하고 높은 정확도의 출력을 생성합니다.
학생 모델 초기화 및 학습: 학생 모델은 교사 모델의 예측(확률 분포)을 기반으로 학습됩니다. 학생 모델은 교사 모델이 예측한 출력값을 사용하여 자신의 출력을 조정합니다.
손실 함수: 학생 모델은 보통 소프트 타겟(교사 모델의 확률 분포)을 사용하여 손실 함수를 최소화하도록 학습됩니다. 일반적으로 소프트맥스 함수에 온도 매개변수를 사용하여 출력 확률 분포를 부드럽게 하고, 이를 통해 더 풍부한 학습 신호를 제공하게 됩니다.
성능 향상 및 경량화: 최종적으로 학생 모델은 교사 모델의 성능에 근접하면서도 더 작고 빠르게 동작하는 모델이 됩니다. 이는 실제 환경에서 모델을 운영할 때 효율성을 높이고, 메모리와 계산 자원의 소비를 줄이는 데 유리합니다.

물론입니다! 주어진 내용을 "~이다" 어투로 벨로그 스타일로 번역해 드리겠습니다.

DistilBERT: BERT의 증류 버전

학생 모델 구조

이번 연구에서 학생 모델인 DistilBERT는 BERT와 같은 일반 구조를 가지고 있다. 토큰 타입 임베딩과 풀러(pooler)가 제거되었으며, 레이어 수는 2배 줄어들었다. Transformer 아키텍처에서 사용되는 대부분의 연산(선형 레이어 및 레이어 정규화)은 최신 선형 대수 프레임워크에서 매우 최적화되어 있다. 조사를 통해 텐서의 마지막 차원(히든 사이즈 차원)에서의 변동이 레이어 수와 같은 다른 요소들의 변동보다 계산 효율성에 미치는 영향이 작다는 것을 알게 되었다. 따라서 우리는 레이어 수를 줄이는 데 집중한다.

학생 모델 초기화

앞서 설명한 최적화 및 아키텍처 선택 외에도, 훈련 절차에서 중요한 요소는 서브 네트워크가 수렴할 수 있는 적절한 초기화를 찾는 것이다. 교사와 학생 네트워크 간의 공통 차원을 활용하여, 우리는 두 레이어 중 하나를 선택해 학생을 교사에서 초기화한다.

Model	Score	CoLA	MNLI	MRPC	QNLI	QQP	RTE	SST-2	STS-B	WNLI
ELMo	68.7	44.1	68.6	76.6	71.1	86.2	53.4	91.5	70.4	56.3
BERT-base	79.5	56.3	86.7	88.6	91.8	89.6	69.3	92.7	89.0	53.5
DistilBERT	77.0	51.3	82.2	87.5	89.2	88.5	59.9	91.3	86.9	56.3

DistilBERT는 BERT 성능의 97%를 유지하고 있다. 이는 GLUE 벤치마크의 개발 세트에서의 비교이다. ELMo 결과는 저자가 보고한 바와 같다. BERT 및 DistilBERT 결과는 서로 다른 시드로 5번 실행한 중간값이다.

Model	IMDb (acc.)	SQuAD (EM/F1)
BERT-base	93.46	81.2/88.5
DistilBERT	92.82	77.7/85.8
DistilBERT (D)	-	79.1/86.9

DistilBERT는 다운스트림 작업에서도 유사한 성능을 발휘한다. IMDb (테스트 정확도)와 SQuAD 1.1 (개발 세트에서 EM/F1)에서의 비교이다. D는 미세 조정 중 두 번째 증류 단계가 있는 경우를 나타낸다.

Model	# param. (Millions)	Inf. time (seconds)
ELMo	180	895
BERT-base	110	668
DistilBERT	66	410

DistilBERT는 크기가 작으면서도 훨씬 빠르다. GLUE 작업 STS-B(감성 분석) 전체를 CPU에서 배치 크기 1로 처리할 때의 추론 시간이다.

증류 과정

BERT 모델을 훈련하기 위한 최근 제안된 최선의 방법을 적용했다. DistilBERT는 매우 큰 배치에서 그래디언트 누적(최대 4,000개의 예제를 한 배치로 활용) 및 동적 마스킹을 사용하여 다음 문장 예측 목표 없이 증류되었다.

데이터와 계산 자원

DistilBERT는 원래 BERT 모델과 동일한 코퍼스(영어 위키백과와 토론토 북 코퍼스 [Zhu et al., 2015]의 결합)에서 훈련되었다. DistilBERT는 8개의 16GB V100 GPU에서 약 90시간 동안 훈련되었다. 비교를 위해, RoBERTa 모델 [Liu et al., 2019]은 1,024개의 32GB V100에서 1일의 훈련이 필요했다.