MINILM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers

[Taekyoon]

• link: https://arxiv.org/pdf/2002.10957.pdf
• code: https://github.com/microsoft/unilm/tree/master/minilm

[Taekyoon]

사전에 알면 좋은 내용들

• BERT 설명: http://jalammar.github.io/illustrated-bert
• UNILM 설명: https://eagle705.github.io/articles/2019-12/Unified-Language-Model-Pre-training-for-Natural-Language-Understanding-and-Generation
• Distillation 모델 관련
  • 모델 압축 방법 설명: https://blog.est.ai/2020/03/%EB%94%A5%EB%9F%AC%EB%8B%9D-%EB%AA%A8%EB%8D%B8-%EC%95%95%EC%B6%95-%EB%B0%A9%EB%B2%95%EB%A1%A0%EA%B3%BC-bert-%EC%95%95%EC%B6%95
  • Tiny Bert 설명: https://jeongukjae.github.io/posts/tiny-bert-distilling-bert-for-natural-language-understanding/
  • Mobile Bert 설명: https://brunch.co.kr/@choseunghyek/4

Abstract

• 모델 사이즈를 줄이기 위한 또다른 모델 distillation에 대한 내용
  • 기존에 Distill-BERT, Tiny-BERT, Mobile-BERT 등이 나와있었음
• Distillation을 Self-Attention(SA)을 활용하여 해결
  • 이 논문에서는 작은 모델이 Self-Attention을 따라하도록 하는도록 하는 것이 목적
• self-attention의 value를 scaled dot product 하여 knowledge distillation에 활용
  • key-query를 가지고 self-attention을 만들었다면 value-value를 가지고 self-attention을 만들었다고 보면 됨
  • query-key attention과는 어떤 다른 knowledge를 student에게 전달해주는지)
• teacher assistant라는 방식 도입
  • teacher 모델을 돕는 모델
  • distill을 하는데 어떤 도움을 주는지

논문에서 살펴볼 만한 포인트

• 이미 나온 모델들에 비해 학습 방식이 간단하다.
• 모델 사이즈를 조절하는 하이퍼파라메터의 제약이 적어졌다.

Main Contents

Knowledge Distillation

• Teacher 모델의 feature를 soft label (0~1 사이의 값으로 표현된 label)와 같이 Student 모델에 전이 시키는 학습
• 동일한 학습 데이터에 대해서 Teacher 모델과 Student 모델 feature 간의 로스를 표현하는데, 주로 MSE나 KLD를 많이 활용
• 주로 Masked Language Model prediction이나 임베딩 outputs, self-attention 분포 그리고 매 레이어에서 출력된 hidden-feature들을 활용하여 Distillation을 함

Key Ideas

• Self Attention Distribution Transfer
• Self-Attention Value-Relation Transfer
• Teacher Assistant

Self Attention Distribution Transfer

• Transformer 마지막 레이어의 Self-Attention을 따라하도록 학습
• Student 모델이 Teacher 모델의 Self-Attention 분포를 학습 (이때 KL Distribution을 활용)
• "Layer-by-Layer로 Self-Attention을 학습할 필요가 있는가" 라는 점을 시사
• Transformer 레이어 수를 하이퍼파라메터로 활용할 수 있다

Self-Attention Value-Relation Transfer

• Value에 대한 정보를 학습하기 위한 방법
• Transformer의 value-value에 대한 Self-Attention 맵을 만들고 분포를 위 방식과 동일하게 학습
• Value 간의 scaled dot product로 만든 분포(맵)를 가지고 Teacher 모델의 value 정보를 Student 모델에 학습하고자 함
• 임베딩 Hidden dimension을 하이퍼파라메터로 활용할 수 있다

• 학습하게 되는 Loss는 Key-Query에 대한 Attention과 Value-Value에 대한 Attention 분포들임

Teacher Assistant

• Student 모델을 Knowledge Distillation 돕기 위해 Teacher 모델보다는 작고 Student 모델보다는 큰 모델을 만듬 (intermediate-size student model)
• 여기서는 Teacher 모델과 비교했을 때 레이어는 같고 hidden dimension을 작게한 모델을 Teacher Assistant라 함
• Teacher 모델과 Student 모델 간의 Knowledge Distillation을 하는데 중간 다리 역할을 함
• 결과적으로 성능 향상에 도움을 줌

Comparison with previous Work

• 얼마나 Knowledge Distillation을 하면서 제약을 받지 않을 수 있는지를 보여줌

• Layer-to-Layer Distillation
  • Teacher 모델에 어느 Layer의 정보를 선택하여 학습할 것인지에 대한 선택이 필요함
• Requirements on the number of layers of students
  • 레이어 수를 조절하는데 제약을 받음
• Requirements on the hidden size of students
  • Hidden dimension에 대한 제약을 받음
• MiniLM은 이전에 활용했던 방식과 비슷하게 진행하되 더 간편하고 학습을 하는데 자유도를 준 모델임을 보임

Experiments

Distillation Setup

• BERT Base 모델(12-layers, 768 hidden-size, 109M parameters)을 teacher 모델로 설정

  • vocab size: 30,522
  • sequence length: 512

• Adam Optimizer

  • Beta 1: 0.9
  • Beta 2: 0.999

• For 6-layers 768 hidden-size

  • learning rate: 5e-4 (for 400,000, peak: 3e-4)
  • batch size: 256
  • linear warmup (first 4,000 steps) and linear decay
  • dropout: 0.1
  • linear decay: 0.01

• Distill에서 사용한 머신

  • V100 x 8 for distillation
  • P100 for inference test

• 6-layers 384 hidden-size의 경우 in-house BERT를 활용

• 학습 파라메터 설정이 일부 다르지만 큰 차이는 없음

Downstream Task Results

• 전반적으로 downstream task에서 teacher 모델과 비슷한 성능을 보이고 이전 distilled model과 비교했을 때 좋은 성능을 보임
• 제안한 논문 모델이 TinyBERT의 학습 방식보다 단순하지만 전반적으로 좋은 점수를 보임 (Teacher 모델의 Knowledge를 어떻게 다양하게 넣느냐 보다는 임팩트 있게 넣는지를 고민한 부분이라 봄)



• MLM-KD는 DistillBert를 대신해서 평가한 것으로 보임
• MINILM은 Teacher Assistant를 적용한 것과 안한 것 둘 다 기존 Distilled 모델보다 나음

Ablation Studies

• 주요 포인트
  • 과연 Value-Relation에 대한 Attention 학습이 필요했는가?
  • Attention Distribution에 대해서 KL Divergence를 활용했는데 이전 모델들에서 활용한 MSE보다 얼마나 나은가?
  • Layer-to-layer 보다 더 성능이 떨어지지 않고 효과를 잘 보이는가?

Discussion (아직 이야기 끝나지 않았습니다...)

• UNILM을 Distillation 하여 언어 생성 테스크 수행을 진행
• UNILM은 MLM, NLG, Seq2Seq이 세 방식을 모두 적용하여 학습한 모델 => NLU, NLG, Seq2Seq에 관한 모든 테스크에 finetune이 가능

• 여기서 진행한 NLG finetune은 UNILM과 동일하게 seq2seq 방식으로 Fine-tune을 했을 것으로 보임

Generation tasks

• 파라메터 수가 많이 줄어들었지만 Generation 성능에는 Teacher 모델과 비교했을 때 큰 차이를 보이지 않음

• Question Generation

• Abstractive Summarization

• 그 외 Multilingual 여러 실험들이 있지만 여기까지만 설명.

Personal Opinion

• Task agnostic 모델을 Distilation 하는데 효율적인 방법으로 학습 하는 방법을 제안
• hidden dimmension이나 token embedding의 제약을 받지 않고 학습 할 수 있는 점이 긍정적
• fine-tune이나 generation에서 distillation 모델 중에서는 좋은 성능을 보이는 것으로 확인
• 모델 성능까지 전부 검증을 했지만 attention map 학습만으로 마지막 layer의 teacher의 token representation vector들 간의 관계를 어떻게 재현가능한 것인지 의문
• 실용적으로 좋아보이는 방식이지만 어텐션 맵만으로도 어떻게 Knowledge distillation이 가능한지에 대해서는 잘 모르겠음