Pea-KD: Parameter-efficient and accurate Knowledge Distillation on BERT

Posted by 날짜:

본 문서에서는 Pea-KD: Parameter-efficient and accurate Knowledge Distillation on BERT 논문을 소개합니다. 논문에 대한 상세한 정보는 다음과 같습니다.

BERT 압축 (BERT Compression)

최근 들어 자연어 처리 분야에서 사전학습 (pretrain) 후 특정 데이터셋에 미세조정 (fine-tuning)하는 모델들이 굉장히 좋은 성능을 보이고 있으며 그 중에 가장 대표적인 모델이 BERT입니다. BERT 모델은 하나의 임베딩 층과 트랜스포머 (transformer)의 인코더 (encoder) 부분을 여러 층 쌓은 구조로 이루어져 있습니다. 그리고 각 인코더는 멀티 헤드 셀프 어텐션 (multi-head self-attention) 구조를 사용합니다. 간단히 말해서 각 인코더에 들어오는 인풋 벡터를 n 등분한 후 각각에 셀프 어텐션을 적용한 후 다시 합치는 (concatenate) 과정을 사용합니다. BERT의 가장 기본적인 모델로 BERT-base를 흔히 사용하며, 이는 12개의 인코더 층과 벡터의 dimension으로 768을 사용합니다.  

하지만 이러한 BERT는 파라미터 수가 굉장히 많다는 단점이 있습니다. BERT-base의 파라미터 수는 대략 1.1억개로 상당히 많은 편이며 따라서 edge device에 적용하기가 어렵습니다. 따라서, BERT를 압축하는 여러 연구가 진행되고 있고 본 논문에서는 Knowledge Distillation과 Parameter Sharing 측면에서 BERT 압축을 기존보다 더 효율적으로 할 수 있는 방법을 제시합니다.    

지식 증류 (Knowledge Distillation (KD))

지식 증류란 실제 사용하고자 하는 작은 사이즈의 학생 모델 (student)을 학습시키는 데 있어서 미리 학습된 큰 사이즈의 교사 (teacher) 모델의 출력값을 활용하는 방법입니다. 가장 기본적으로 사용하는 지식 증류 방법은 학생 모델을 학습시킬 때 손실 함수 (loss function)에 학생 모델의 최종 출력값과 교사 모델의 최종 출력값 사이의 Cross Entropy Loss를 계산하여 추가하는 방법을 사용합니다. 일반적으로 지식 증류를 사용할 경우 학생 모델의 성능이 기존보다 좋아집니다.      

파라미터 공유 (Parameter Sharing)

파라미터 공유란 한 모델 내에서 일부 파라미터들 간에 공통된 파라미터를 사용하므로써 전체 파라미터 수를 줄이는 방법을 말합니다. 파라미터 공유는 모델의 전체 파라미터 수를 획기적으로 줄일 수 있다는 장점이 있지만 대부분 모델 성능이 현저히 떨어진다는 단점이 있습니다.   

Challenges

실제로 BERT 모델에 지식 증류와 파라미터 공유를 적용한 연구들은 여럿 있습니다. 예를 들어 Patient Knowledge Distillation과 ALBERT 논문이 대표적입니다. 그러나 근본적으로 지식 증류와 파라미터 공유 방법은 여전히 더 개선의 여지가 있으며 본 논문에서는 기존 방법들의 다음과 같은 한계점들에 착안하여 더 개선된 방법을 제안합니다. 
  • 학생 모델의 낮은 모델 복잡도: 기본적으로 학생 모델로 작은 사이즈의 모델을 사용하기 때문에 모델 복잡도가 낮을 수 밖에 없습니다. 이 낮은 모델 복잡도로 인해 학생 모델의 성능은 좋지 못 합니다. 
  • 초기화의 부재: 기존의 KD 방법들은 학생 모델 초기화를 전혀 고려하지 않습니다. 하지만 본 논문에서는 학생 모델을 적절하게 초기화 시켜줌으로써 KD의 결과가 더 좋아진다는 것을 실험적으로 보여줍니다.      

Proposed Method

Pea-KD는 SPS와 PTP라는 두 가지 새로운 메서드를 제안합니다. 

Shuffled Parameter Sharing (SPS) 
SPS의 motivation은 학생 모델의 전체 파라미터 수는 동일하게 유지하되 모델 복잡도를 증가시켜 더 높은 성능을 가지게끔 하는 것입니다. SPS는 먼저 (step 1) 파라미터 공유된 층들을 기존 모델에 추가적으로 쌓고 (step 2) 공유된 파라미터들에 추가적으로 shuffling이라는 과정을 도입하여 모델의 복잡도를 더욱 높이는 방법입니다. 구체적으로, shuffling이란 아래 그림과 같이 공유된 layer들 사이에서 Query와 Key 파라미터의 위치를 바꾸어주는 과정으로 정의합니다. Shuffling 과정을 도입할 경우 각 파라미터들은 기존보다 더 많은 feature들을 학습하게 되어 일종의 정규화 (regularization) 효과 혹은 다중 작업 학습 (multi-task learning) 효과처럼 작용하여 모델의 일반화 성능을 높일 수 있습니다. 예를 들어, 아래 그림의 (b)의 경우 첫 번째 층의 Query 파라미터는 첫 번째 층에서는 Query 파라미터로, 네 번째 공유된 층에서는 shuffling에 의해 Key 파라미터로 쓰이게 됩니다. 따라서 이 파라미터는 Query의 feature도 배우고 Key의 feature도 배우게 되므로 정규화 효과 혹은 다중 작업 학습 효과로 인해 더 좋은 일반화 성능을 가지게 됩니다.    

[그림1] SPS 절차를 도식화한 그림

Pretraining with Teacher's Predictions (PTP)
PTP의 motivation은 교사 모델의 출력값에 들어 있는 중요한 정보들 (교사 모델이 정답을 맞추는지 못 맞추는지, 어느 정도 확신 (confidence)를 가지고 맞추는지)를 label로 노골적으로 나타낸 후 그 label들을 가지고 학생 모델을 먼저 사전 학습 시키는 것입니다. 그럴 경우 학생 모델은 교사 모델의 출력값에 담겨 있는 중요한 일반화된 정보 (generalized knowledge)를 사전에 학습할 수 있고 따라서 본격적인 KD에 앞서 더 준비된 초기 상태를 가지게 됩니다. 이는 결과적으로 KD를 하였을 때 기존보다 더 좋은 성능으로 이어지게 됩니다. 
  
[그림2] PTP label을 정의하는 방법

PeaBERT: SPS and PTP Combined 
본 논문에서 제안한 두 가지 메서드 SPS와 PTP를 동시에 BERT에 적용한 모델을 PeaBERT라고 합니다. 논문에서는 실험적으로 PeaBERT가 기존의 BERT보다 뛰어난 성능을 가진다는 것을 보여주고 ablation study들을 통해서 SPS와 PTP가 각각 효과가 있음을 보여줍니다.   

[그림3] PeaBERT의 전체적인 절차. 먼저 기존 BERT 모델 (student)에 SPS를 적용한 후 PTP로 사전학습 시킨 다음 KD를 적용합니다. 

Experiments

본 논문에서는 총 네 가지의 GLUE task (RTE, MRPC, SST-2, QNLI)에 대한 성능을 측정한 실험 결과를 제시합니다. 그 결과를 정리한 것은 아래 그림과 같습니다. PeaBERT는 모든 task에 대해 기존 BERT보다 월등히 좋은 성능을 내며 다른 경쟁 모델들 (DistilBERT, TinyBERT)보다도 같은 파라미터 수일 때 더 좋은 성능을 보임을 확인할 수 있습니다. 


그와 더불어 ablation study에서는 SPS와 PTP 각각의 효과를 확인할 수 있는 실험을 진행하였습니다. 그 결과는 아래와 같습니다. 결과를 보면 SPS를 적용했을 때 모델의 성능이 평균 3.0%, 추가로 PTP까지 적용했을 때 평균 4.5%까지 증가하는 것을 확인할 수 있습니다. 이를 통해 SPS와 PTP가 각각 모델 성능 향상에 효과가 있음을 확인할 수 있습니다. 


Conclusion

본 문서에서는 Pea-KD: Parameter-efficient and accurate Knowledge Distillation 논문을 소개하였습니다. 해당 논문은 기존의 KD에 새로운 초기화 방법을 제안하고 파라미터 공유 방법에 shuffling이라는 새로운 메커니즘을 도입하였습니다. 또한, 실험을 통해 기존의 방법들보다 월등히 뛰어난 성능을 보임을 확인하였습니다. 본 논문은 기존에 없던 새로운 방법들을 제안하였고 그 가능성을 실험을 통해 입증하였습니다. 지식 증류 방법과 파라미터 공유 방법은 자연어 처리뿐만 아니라 컴퓨터 비전, 그래프 모델 등 다양한 분야에서 사용되기 때문에 본 논문의 방법은 다양한 종류의 모델들을 (e.g., CNN, GCN 등) 효율적으로 압축하는데 유용하게 사용될 수 있습니다. 해당 논문에 대한 더 자세한 정보는 링크를 통해 확인할 수 있습니다.