MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices

tacorico·2021년 9월 9일
0

(작성중..)

논문 링크

Introduction

최근 NLP 연구는 큰 pre-trained 모델을 사용하여 성공적인 결과들을 이루어냈는데, 큰 모델은 필연적으로 high latency라는 근본적인 문제를 가지고 있다. 특히나 모바일처럼 한정된 자원을 사용하는 디바이스에서는 큰 모델이 더 부담이 될 수밖에 없다. 이런 이유들로, 큰 모델로 인한 성능 자체의 개선의 연구와 더불어 Model Compression 같은 연구들도 함께 진행되고 있다. 이 논문에서는 MobileBERT라는 모델을 제안하는데, Model Compression 중에서도 Knowledge Distillation에 관련된 내용이다. 기존 BERT 모델처럼 task-agnostic(특정 task에 한정되어 있지 않다는 말. 기존 BERT가 fine-tuning으로 이곳저곳 활용할 수 있는 것을 생각하면 됨)하다.

Structure

mobileBERT는 teacher 모델(여기서는 IB-BERT)을 먼저 디자인하고, IB-BERT에서 mobileBERT로 knowledge transfer를 수행하는 방식으로 학습한다. 결론적으로, BERT BASE 모델에 비하여 4.3배 작고, 5.5배 빠른 mobileBERT 모델을 구현할 수 있다. 이렇게 작아지고 빨라진 성능 대비, GLUE, SQuAD score의 하락폭은 낮은 편으로 합리적인 모델 사용을 할 수 있다. 또한, 이전 연구들과의 차이는 기존 연구들은 모델 사이즈를 줄이기 위해서 layer의 수를 줄였지만 mobileBERT는 width를 줄이고 layer를 그대로 가져갔다는 점이 차이점이다.

Bottleneck and Inverted-Bottleneck

bottle이라는 말 그대로 병목 현상을 만들어 주는 것이다. IB-BERT의 구조를 보면 초록색 박스로 Linear라는 부분이 있는데 단순하게 선형 변환을 해서 변수의 차원을 조절해주는 부분이다. 왜냐하면, 여기서는 IB-BERT를 teacher model로 사용하는데 나중에 나올 knowledge transfer가 용이하도록 input, output dimension을 맞춰준다.

Stacked Feed-Forward Networks

이렇게 bottleneck structure를 사용하다보면 문제가 되는 부분이 생긴다. 여기서 말하는 문제점은, Multi-Head Attention(MHA)와 Feed-Forward Network(FFN)의 parameter 수의 밸런스가 깨진다는 것이다. MHA는 모델이 서로 다른 하위 공간의 정보를 공동으로 확인할 수 있게 하는 반면, FFN는 모델의 비선형성을 증가시킨다. Original BERT에서는 MHA와 FFN의 비율이 1:2이다. 하지만, bottleneck structure에서는 더 넓은 MHA feature map, 더 좁은 FFN feature map을 가지기 때문에 비율이 깨지게 된다. 이 문제를 해결하기 위해서 FFN을 더 쌓는 방식으로 해결했다.

Training Objectives

두 가지의 knowledge transfer를 수행한다.

  • Feature map transfer
  • Attention transfer

Feature Map Transfer(FMT)

각각의 layer가 teacher model과 비슷해지는 것이 목적이다. 특히, mobileBERT와 IB-BERT의 MSE가 knowledge transfer objective에 사용된다.
LFMTl=1TNt=1Ta=1ADKL(at,l,atrat,l,ast)L^l_{FMT} = {1 \over TN} \sum_{t=1}^T \sum_{a=1}^A D_{KL}(a_{t,l,a}^{tr}||a_{t,l,a}^{st})
ll: index of layers
TT: sequence length
NN: featurer map size

Attention Transfer(AT)

중요한 아이디어는, 잘 최적화 된 teacher에서 self-attention map을 사용해서, mobileBERT를 학습하면 되지 않을까?라는 생각이다. 특히, mobileBERT와 IB-BERT의 헤드 별 self-attention의 분포들 사이의 KL-divergence를 최소화하는 것으로 학습한다.
LATl=1TAt=1Ta=1ADKL(at,l,atrat,l,ast)L^l_{AT} = {1\over TA} \sum_{t=1}^T \sum_{a=1}^A D_{KL}(a_{t,l,a}^{tr}||a_{t,l,a}^{st})
AA: number of attention heads

Datasets & Evaluations

Conclusion

profile
우당탕탕 자연어 일기

0개의 댓글