BERT는 많은 NLP tasks에 대해 좋은 성능을 보인다. 하지만, 이러한 pre-training methods는 대체로 계산 비용이 많이 들기 때문에, 모바일과 같은 제한적인 기기에서는 작동시키기가 어렵다. 따라서 기존의 KD(Knowledge Distillation)을 적용한 Transformer distillation를 제안한다. 새로운 KD는 큰 'teacher' 모델인 BERT에서 작은 'student' 모델인 TinyBERT로의 지식 전이를 효과적으로 하기 위한 방법이다. Transformer distillation은 두 개의 과정으로 나뉘게 되는데, 이 두 과정은 기존의 pre-training과 fine-tuning에 해당하는 학습 과정이다. 이 방법을 사용함으로써 TinyBERT가 BERT로부터 효율적으로 지식을 전수받을 수 있다.
TinyBERT-4는 4개의 레이어로 선생 모델인 'BERT-base'를 7.5배 작은 모델 사이즈와 9.4배 빠른 추론 속도로 의 GLUE benchmark 성능인 96.8%을 뛰어넘는 성능을 달성했다.
Introduction
Pre-training된 모델을 fine-tuning하여 여러 tasks에 적용시키는 방법은 NLP의 새로운 패러다임이 되었다. 최근에 연구된 BERT, XLNet, RoBERTa, ALBERT, T5, 그리고 ELECTRA 등의 Pre-trained Language Models(PLM) 기반의 논문들은 좋은 성능을 보이지만 방대한 계산 비용을 가지고 있다는 단점과, 그로 인해 생기는 긴 추론 시간으로 실질적인 활용에 있어서 제한이 생길 수 있다(모바일 환경에서의 활용 등). 최근의 연구들은 PLM에 중복성이 있다고 주장하며, 이에 따라 PLM 모델의 사이즈를 축소시키면서 성능을 유지하는 것이 가능하다.
모델의 추론 속도를 높이고 사이즈를 줄이기 위해 quantization, KD, 등의 연구가 이뤄졌으며, 해당 논문은 KD에 집중한다. KD는 teacher-student 구조로 되어있으며, 학습된 teacher 모델의 지식을 student 모델에게 전이하는 학습 구조를 말한다.
teacher 모델과 student 모델의 behavior function을 각각 $f^T$, $f^S$라고 하자. 이때 behavior function은 입력을 특정 representation으로 변환하는 것에 중점을 두며, 이는 네트워크에 속하는 임의의 레이어의 출력이 된다. KD는 다음의 objective function을 최소화시키는 것을 목표로 학습한다:
$$L_{KD} = \sum_{x \in X} L(f^S(x), f^T(x))$$
Methods
TinyBERT는 transformer를 위한 transformer distillation과 two-stage learning을 통해 학습을 진행한다.
Transformer Distillation
학생 모델이 M개의 레이어를 가지고, 선생 모델이 N개의 레이어 수를 가진다고 했을 때, mapping function $n = g(m)$에 의해 M과 N이 결정된다. 선생 모델과 학생 모델의 크기가 다르기 때문에 mapping function을 통해 정보 전달을 용이하게 한다. Embedding layer의 index를 0으로 설정하고, $M + 1$과 $N + 1$을 예측 레이어로 설정하면 $N + 1 = g(M + 1)$으로 대응된다. Transformer distillation은 다음 objective를 최소화한다:
$$L_{model} = \sum_{x \in X} \sum_{m=0}^{M+1} \lambda_m L_{layer}(f_m^S(x), f_{g(m)}^T(x))$$
여기서 $L_{layer}$은 주어진 모델 레이어에 대한 loss function으로, transformer distillation은 세 가지 layer에 대한 loss를 정의한다.
Transformer-layer Distillation 은 hidden state와 attention 기반의 전이를 의미한다. BERT에 의해 학습된 attention은 언어적 지식을 다량 습득할 수 있기 때문에 이와 같은 distillation을 제안하며, object 수식은 다음과 같다:
$$L_{attn} = \frac{1}{h} \sum_{i=1}^h MSE(A_i^S, A_i^T)$$
$h$는 attention 헤드의 수를 의미하며, MSE loss를 사용하여 오차 수치를 계산한다. 모델은 transformer 레이어의 출력인 hidden state의 정보의 전이 또한 수행하며, objective function은 다음과 같다:
$$L_{hidden} = MSE(H^S W_h, H^T)$$
$W_h \in \mathbb{R}^{d' \times d}$은 학습 가능한 linear transformation matrix이며, 학생 네트워크의 hidden state $H^S \in \mathbb{R}^{l \times d'}$를 선생 네트워크의 hidden state $H^T \in \mathbb{R}^{l \times d}$와 같은 크기로 변환하는 기능을 가지고 있다.
Embedding-layer Distillation 은 hidden state 지식 전이와 유사한 embedding 레이어에서의 지식 전이를 일으킨다. objective function은 다음과 같다:
$$L_{embd} = MSE(E^S W_e, E^T)$$
$W_e$는 위의 hidden state distillation의 $W_h$와 유사한 역할로, 학생 모델과 선생 모델의 차원을 계산 가능한 크기로 변환하는 기능을 수행한다.
Prediction-layer Distillation 은 prediction 레이어에서 이뤄지는 지식 전이이다. $z_S$, $z_T$를 각각 학생 모델과 선생 모델이 예측한 값이라고 했을 때, objective function은 다음과 같다:
$$L_{pred} = CE(z^T / t, z^S / t)$$
$CE$는 cross entropy loss를 의미하며, $t$는 temperature value로, 논문의 실험에서 $t = 1$일 때 잘 작동하는 것을 확인했다.
세 가지 distillation을 정의함으로써 transformer distillation은 transformer 구조 전체 레이어에서 선생 모델과 학생 모델에 대응되는 distillation loss를 사용할 수 있게 되었다.
Learning Process
기존의 BERT는 pre-training과 fine-tuning 두 가지 학습 과정을 거친다. BERT 모델은 이 과정에서 자연어 처리에 대한 지식을 다량 습득하며, 이를 TinyBERT로 전이하기 위해 pre-training에 해당하는 general distillation과 fine-tuning과정에 해당되는 task-specific distillation 두가지 프로세스를 거친다. 또한 지식 전이에서 활용할 data augmentation 기법도 소개한다.
- General Distillation: 첫 번째 전이 과정에서는 선생 모델인 pre-trained된 BERT의 지식을 학생 모델인 TinyBERT가 학습한다. 이 과정에선 fine-tuning되지 않은 BERT 모델을 사용하며, 지식 전이가 완료된 TinyBERT는 fine-tuning 과정을 거쳐서 downstream tasks에 적용될 수 있지만, 축소된 모델 구조를 가지기 때문에 기존의 BERT보다 성능이 현저히 떨어진다.
- Task-specific Distillation: Pre-trained BERT는 fine-tuning을 통해 downstream task에 적용된다. Task-specific distillation은 general distillation을 거친 TinyBERT를 downstream tasks에 적용하기 위한 과정으로, 선생 모델인 fine-tuned BERT로부터 지식을 전이받는다. 이 때, augmented task-specific dataset을 사용한다. 데이터 증강 기법을 활용하여 모델의 일반화 성능을 향상시킬 수 있다.
- Data Augmentation: Pre-trained BERT와 GloVe word embedding을 활용하여 데이터를 증축시킨다. BERT는 마스킹된 토큰을 예측하기 때문에 '이 책은 파란색이야.'에서 파란색을 마스킹하면 BERT는 해당 토큰이 될 수 있는 분포를 계산한다. 이 분포에서 GloVe word embedding을 활용하여 비슷한 단어를 선별하여 데이터를 증축시킨다. 아래 [그림 1]은 데이터 증축을 위한 알고리즘이며, $p_t = 0.4$, $N_a = 20$, $K = 15$로 설정되었다.
학습 과정에서 general distillation과 task-specific distillation은 서로를 보완하며 완성도 있는 학습을 진행할 수 있게 된다. General distillation은 task-specific distillation의 학습 시작 시점에 도움을 주고, task-specific distillation은 최종 downstream tasks에 적용하기 위해 꼭 필요하다. 결론적으로 TinyBERT는 다양한 NLP task에서 좋은 성능을 보인다.
Experiments
- Datasets: 실험 과정에서 TinyBERT의 성능을 평가하기 위해 GLUE(General Language Understanding Evaluation) 벤치마크를 사용한다. GLUE는 2개의 단일-문장 문제, 3개의 문장 유사도 문제, 4개의 자연어 추론 문제 등이 포함되어 있다.
- Model Settings: 학생 모델인 TinyBERT의 사이즈는 { 레이어의 수 $M=4$, 히든 사이즈 $d'=312$, feed-forward/filter 사이즈 $d'_i=1200$, 헤드의 수 $h=12$ }로 총 14.5M개의 파라미터 수를 가진다. 선생 모델인 BERT-base는 { $N=12$, $d=768$, $d_i=3072$, $h=12$ }로 총 109M의 파라미터 수를 가진다. Mapping function은 $g(m) = 3 \times m$으로 설정되어 선생 모델의 레이어 3개 마다 지식 전이가 이뤄진다. TinyBERT$_6$은 { $M=6$, $d'=768$, $d'_i=3072$, $h=12$ }로 구성된다.
TinyBERT는 BERT-tiny, BERT-small, BERT-PKD(논문), DistillBERT 등과 비교되었다.
TinyBERT$_4$의 GLUE 벤치마크 성능은 BERT-base의 비교하면 조금 떨어지는 경향이 있지만, 9.4배 빠르며, 모델의 사이즈 또한 7배 가량 적다. 모델의 사이즈는 같지만 pre-training을 그대로 적용시킨 BERT$_{TINY}$는 TinyBERT와 평균 성능이 6.8%가 차이난다. 다른 4개의 레이어를 가진 KD baseline 모델들과 비교했을 때 TinyBERT$_4$가 가장 좋은 성능을 보인다.
모델의 레이어 수를 6개로 확장한 TinyBERT$_6$는 BERT-base에 비해 파라미터 수가 현저히 적고, 계산량과 추론 속도가 2배 가량 빠름에도 평균 성능이 거의 비슷하고, 다른 baseline 모델들보다 비교하여 2.6% 높은 벤치마크 성능을 달성한다.
Discussions
해당 논문은 transformer-based distillation으로 학습한 TinyBERT를 제안했다. TinyBERT는 학습된 BERT-base의 지식을 전이 받음으로써 제한된 환경에서 적용 가능하도록 추론 시간과 연산량이 적음에도 성능은 유지되는 결과를 보였다.
논문을 참고하고 싶으신 분은 아래 링크 참고해주세요
