NL-052, Patient Knowledge Distillation for BERT Model Compression (2019-EMNLP)

*Distilling
1.Task data로 Fine-tuned BERT을 student에 distilling을 한다.
2.이 때는 data labels와 teacher의 predicted prob을 동시에 이용한다. (Loss 식이 두 개임)
3.초기화는 pre-trained BERT을 이용한다.

■ Comment
  • TinyBERT보다 늦게 읽었지만, arxiv 기준으로 TinyBERT보다 먼저 나온 논문이다.
  • PKD의 핵심은 사실 TinyBERT처럼 출력층뿐 아니라 중간층을 학습하면 일반화가 잘 된다는 의미이다.
  • TinyBERT도 그랬지만 PKD-Last보다는 PKD-skip이 성능이 좀 더 괜찮다.
  • 또한 실험 결과에서 4.5 부분에서 실제 distilling할 때 도움될만한 착안점이 많다고 생각한다.

0. Abstract

  • BERT와 같은 pre-trained LM은 NLP에 효과적임이 증명 되었다.
  • 하지만 학습할 때 높은 계산 리소스가 들고 application에 적용하기 어렵다.
  • 여기서는 Patient Knowledge Distillation 접근법을 제안하여 얇은 네트워크(student)에 효과적으로 모델을 압축한다.
  • 이전의 distillation 방법들은 teacher 네워크의 마지막 출력만을 사용했지만, 여기서는 incremental KD을 통하여 student 모델이 multiple 중간의 층을 배움으로써 끈기있는 학습을 하려고 하였다.
    • 방법으로는 다음의 두 가지를 제시한다.
    • PKD-Last: 마지막 k 층을 배우는 것
    • PKD-Skip: 매 k 층마다 학습을 배우는 것(TinyBERT 처럼)
  • 이러한 두 가지 patient distillation 방법은 teacher의 hidden layers으로부터 풍부한 정보를 전달할 수 있으며 끈기있게 student model이 배울 수 있다.
  • 실험을 해보면, 이 방법이 학습 효율성이 좋고 모델의 성능또한 해치지 않는다고 한다.

1. Introduction

  • Pre-trained LM이 language representation에 효과적임은 증명이 되었다.
  • 대표적으로 ELMo, GPT, BERT 등이 NLP tasks에서 좋은 성능을 거두고 있다.
  • 하지만 BERT는 계산량이 효과적이지 않고 학습할 때는 4~16 TPU로 4일 정도 학습이 걸린다.
  • fine-tuning도 1 epoch에 몇 시간씩 걸리게 된다.
  • 따라서 모델들의 계산 cost을 줄이는 것은 중요하다.
  • 이 논문에서는 Patient-KD을 제안하는데, teacher 모델이 출력한 logit 확률과 예측 label을 student 모델이 학습하게 한다.
    • 실제 label이 있으면 그것을 써도 되는데 추가적인 unannotated samples들도 학습을 확장하기 위해서라고 함
    • 자세한 것은 뒷 부분을 봐야할 듯
  • 이전의 KD와 다른 것은 teacher의 마지막 층만을 이용해서 knowledge을 뽑아내지 않고 이전의 층도 이용한다는 것이다.
  • 이 방법을 Patient KD라고 부르고 이 방법은 teacher 모델의 여러 층의 풍부한 정보를 전달하는 것이다.
  • 여기서 두 가지 전략을 제시하는데
    • PKD-Last
      • student 모델이 teacher의 마지막 k 층을 배우도록 학습하는 것
      • 이것은 top layers들이 기존의 네트워크에서 중요한 정보, 지식을 가진다고 가정한 것이다.
    • PKD-Skip
      • student 모델이 teacher의 k 층마다 학습하도록 하는 것(즉 중간 중간마다)
      • 이 방법은 중요한 정보가 lower layers에서도 포함하고 있다고 생각하는 것
      • 사실 이 방법이 TinyBERT와 똑같음. 실제로 그 논문과 비슷한 시기에 나옴
  • 논문의 접근법을 Sentiment Classification, Paraphrase Similarity Matching, Natural Language Inference, and Machine Reading Comprehension에서 실험을 하여서 효과적임을 보여준다고 한다.

2. Related Work

  • Language Model Pre-training
    • LM모델 pre-training은 두 가지 카테고리로 나눠져 연구되고 있다.
      • feature-based 접근법
      • fine-tunning 접근법
    • Feature-based
      • Context-independent word representation
        • Word2vec
        • FastText
      • Sentence-level representation (문장 표현을 의미하는 연구)
        • Kiros
        • Conneau
        • Logeswaran
      • Contextualized word representation
        • Cove
        • ELMo
    • Fine-tunning
      • GPT, BERT에서 쓰이는 것들
      • downstream task에 labeled data를 학습하는 것
    • 그러나 이러한 pre-train 모델은 수백많의 파라미터들을 가지고 있기 때문에 이것들의 실제 계산량때문에 실제로 어플리케이션에 적용하기는 어렵다.
      • 이러한 문제점을 해결하기 위해 성능을 해치지 않고 얇은 모델에 압축시키는 연구를 여기서 보여준다.
  • Model Compression & Knowledge Distillation
    • Network pruning
      • Han, 2015, He, 2017
      • 네트워크의 연결을 감소시키거나 희소하게 만드는 방법의 모델 압축 방법
    • Weight quantization
      • Gong 2014, Polino, 2018
      • 연결하는 weight에 제약을 줘서 discrete 값을 가지게 하거나 더 적은 비트로 표현을 하는 방법
    • 하지만 이러한 pruning과 quantization 방법론들은 CNN에서 대부분 쓰인다.
    • 소수의 연구만 언어 모델에 있다고 한다. (Changpinyo, 2017)
    • Knowledge distillation
      • 힌튼이 2015년에 보인 것으로 큰 파라미터를 빠르게 실행할 수 있는 모델에 압축시키는 법이다.
      • 기존 모델의 soft ouput을 모방하도록 compact 모델이 모방을 한다.
      • Romero 2015에서는 작은 모델이 큰 모델의 중간 표현법을 학습하는 것도 학습 과정중에서 성능을 향상시켜주는 방법이라고 보여준다.
      • 그 외에 Chen 2015, Liu 2019a, Tan 2019 등의 연구도 있다.
    • 이 논문에서는 Patient knowledge distillation이라는 방법으로 효과적으로 teacher 모델을 student 모델한테 압축시킨다.
      • 사실상 핵심은 TinyBERT과 비슷한 것 같고 위에서 언급한 Romero의 아이디어에서 착안해서 한 것이 아닐까 싶음

3. Patient Knowledge Distillation

  • Problem Definition
    • Teacher 모델인 을 Student 모델인 가 모방하도록 loss 을 이용하여 학습한다.

3.1 Distillation Objective

  • teacher 은 deep bidirectional encoder, BERT을 의미하고 은 얇은 student 모델을 의미한다.
  • BERT-Base, BERT-Large는 BERT(12), BERT(24)로 표기하고 BERT(k)은 k layers of Transformers을 의미한다.
  • 은 N개의 학습 데이터로 가 입력이고 가 ground-truth label 이다.
  • 는 contextualized embedding이고 는 classification의 예측값이다.
  • W는 weight matrix으로 학습되는 파라미터다.
  • 예를 들어 학습은 다음과 같은 식으로 된다.
    • 진짜 학습되는 loss는 식 3,4이다.
    • 이 식은 단순히 참고만 하면 되고 t는 여기서 teacher model을 의미한다.
  • 입력 에 대한 output 확률은 다음과 같다.
    • 는 soft labels으로 사용되고 T는 temperature으로 teacher의 soft prediction을 얼마나 신뢰하냐를 결정해준다.
  • teacher의 예측과(soft labels) student의 예측사이의 distance을 다음과 같이 정의한다.
    • 식을 보면 cross entropy인데 gt를 soft label로 사용한 것이다.
  • student 모델을 target task에 fine-tune 하기 위해 target specific cross-entropy을 다음과 같이 사용한다.
  • 최종 KD의 objective 함수는 다음과 같다.

3.2 Patient Teacher for Model Compression

  • 이제 이렇게 마지막 층을 이용하여 KD framework을 학습하면 되는데 이 때 figure2에서보면 학습 데이터세트에서는 teacher 모델을 잘 모방하는데 test 세트에서는 빠르게 saturation이 된다고 한다.
  • 따라서 여기서 세운 가정이 KD가 일반화가 약해서 오버피팅이 일어난다!
  • 이러한 이슈를 해결하기 위해서 student을 마지막 층의 logit으로만 학습하지 않고 중간 층에서도 학습하고 이를 "patient"라고 부른다.
  • 따라서 두 개의 patient distillation 전략이 있다.
    • PKD-Skip: the student learns from every k layers of the teacher (Figure 1: Left)
    • PKD-Last: the student learns from the last k layers of the teacher
  • BERT에서 classificaion할 때 [CLS] 토큰을 쓰는 경우가 많다. (오리지널 BERT도 그렇고)
    • SDNet에서는 [CLS] embedding이 적용이 될 때, 중간 층의 hidden layer vector을 가중치 합을 사용하여서 예측을 하는 방식이다.
    • 여기서 착안을 얻어 모델 압축을 할 때도 [CLS]의 중간 층의 표현을 학습하는 것이 일반화 능력을 얻을 가능성이 높을 것이라고 생각했다고 한다.
    • 근데 개인적인 생각은 당연한거 아닌가? 싶긴 하다..
  • [CLS] 토큰의 모든 층은 다음과 같이 표기를 한다.

  • 는 KD를 할 중간 층을 의미한다.
    • PKD-Skip: ={2,4,6,8,10}
    • PKD-Last: ={7,8,9,10,11}
  • 여기서 student는 k=5이다.
    • BERT의 경우 마지막 층은 hidden state가 softmax을 포함하기에 보통 뺀다고 함.
    • 즉, 일반적으로 BERT가 n layers이면 k=n-1임
  • 추가적인 학습을 위하여 PKD loss는 MSE로 다음과 같이 정의함
    • 여기서 M은 student의 layers 수를 의미하고 s, t는 각각 student와 teacher을 의미한다.
  • 전체의 loss는 따라서 다음과 같이 변경 된다.

4. Experiments

4.1 Datasets

  • SST-2, MRPC, QQP, MNLI, QNLI, RTE 등을 사용했다.
  • GLUE 데이터세트를 가지고 한 것이고 다른 논문들과는 다르게 데이터 증강은 사용안한 듯 하다.

4.2 Baselines and Training Details

  • 생략

4.3 Experimental Results

4.4 Analysis of Model Efficiency

  • PKD 방법이 위에서 성능이 잘 나온다는 것을 확인했고 여기서는 모델의 크기와 인퍼런스 시간을 확인해 본다.
  • Table 4를 보면 PKD는 layers 수에 따라 선형적으로 스피드가 빨라진다 (1.94x, 3.73x)
  • embedding layer은 24 million 파라미터를 가지며 30k 단어 사전수, 768 차원를 가지는 것으로 모든 모델이 동일하다.
  • 시간은 QNLI 학습 데이터의 105k 샘플로 측정하였으며 RTX GPU 1장을 사용하고 batch는 128, max seq length는 128, FP16을 사용했다.

4.5 Does a Better Teacher Help?

  • teacher 모델이 Base, Large에 따라 모델 압축의 성능 측정을 비교하였다.
  • Base와 Large는 embedding 차원이 다르기 때문에 teacher모델과 student 모델이 다른 경우 (#1, #2 같은 경우) 바로 loss (7)을 적용할 수가 없다.
  • 따라서 이런 상황에서는 KD만 적용해서 학습을 했다고 한다.
  • #1과 #2는 큰 차이가 없다.
  • #1과 #3을 비교하면 BERT(6)[Large]-KD가 BERT(6)[Base]-KD보다 성능이 확연히 안좋다.
    • Teacher 모델과 BERT(6)[Large]-KD가 파라미터는 더 많은 큰 모델인데 어떻게 된 것일까?
    • 이는 압출률 때문에 그럴 것이라고 생각한다.
    • #1은 teacher과 student의 파라미터수 비율이 2:1
    • #3은 teacher과 student의 파라미터수 비율이 4:1
    • 즉 #3이 모델은 더 크지만 더욱 큰 압축률을 가지므로 챌린지하기 때문이다.
  • #2와 #4를 비교하면 teacher 모델은 동일한데 Large 모델이 성능이 더 안좋다.
    • 이것도 상식적으로는 이상한데 이 논문에서는 weight initialization mismatch 때문이라고 언급한다.
    • weight을 초기화할 때, BERT(6)[Large]는 BERT(24)의 초기 6-layers로 weight 파라미터를 초기화하고 BERT(6)[Base]는 BERT(12)의 초기 6-layers로 weight 파라미터로 초기화를 하여 학습했다고 한다.
    • 근데 BERT(24)는 앞의 6-layers가 BERT(12)보다 비교적 앞의 층이므로 high-level features을 못담아냈기 때문이라고 말을 한다.
    • 이것을 증명한 것은 아니기 때문에 개인적인 생각으론 정말 이것 때문일까? 라는 생각이 들고 같은 모델구조를 가진 것에 단순히 distilling해서 그런게 아닐까 라는 생각도 든다.
    • 즉 초기화없이 단순히 KD만 해서 [Base]와 [Large] 모델 성능 차이를 비교했으면 좀 더 확실했을 것이라는 생각이 든다.
  • #3와 #4을 비교하면 #4의 PKD 방식이 성능이 더 좋다.
    • 따라서 teacher 모델에 상관없이 PKD는 일반적인 접근법이라고 주장한다.

5. Conclusion

  • 이 논문에서는 large BERT을 압축하는 novel 접근법을 소개했다. (PKD)
    • 사실 근데 TinyBERT비슷한 접근법이고 arxiv을 보니 TinyBERT가 늦게 나오긴 했다.
  • teacher 모델의 풍부한 정보를 담아내기 위해서는 student가 다층을 distilling을 해주는게 좋다.
  • 4개의 NLP tasks에서 우리 모델이 효과적임을 증명했다.
  • 앞으로는 pre-train할 때 BERT의 initialization mismatch 이슈를 해결하려고 한다.
  • 좀 더 정교한 distance metric을 설계하는 것도 하나의 방향이다.
  • PKD가 multi-task와 meta-learning과 같은 복잡한 세팅에 도움되길 원한다.

 Reference

댓글

댓글 쓰기