최근 NLP 연구는 큰 pre-trained 모델을 사용하여 성공적인 결과들을 이루어냈는데, 큰 모델은 필연적으로 high latency라는 근본적인 문제를 가지고 있다. 특히나 모바일처럼 한정된 자원을 사용하는 디바이스에서는 큰 모델이 더 부담이 될 수밖에 없다. 이런 이유들로, 큰 모델로 인한 성능 자체의 개선의 연구와 더불어 Model Compression 같은 연구들도 함께 진행되고 있다. 이 논문에서는 MobileBERT라는 모델을 제안하는데, Model Compression 중에서도 Knowledge Distillation에 관련된 내용이다. 기존 BERT 모델처럼 task-agnostic(특정 task에 한정되어 있지 않다는 말. 기존 BERT가 fine-tuning으로 이곳저곳 활용할 수 있는 것을 생각하면 됨)하다.
bottle이라는 말 그대로 병목 현상을 만들어 주는 것이다. IB-BERT의 구조를 보면 초록색 박스로 Linear라는 부분이 있는데 단순하게 선형 변환을 해서 변수의 차원을 조절해주는 부분이다. 왜냐하면, 여기서는 IB-BERT를 teacher model로 사용하는데 나중에 나올 knowledge transfer가 용이하도록 input, output dimension을 맞춰준다.
이렇게 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을 더 쌓는 방식으로 해결했다.
두 가지의 knowledge transfer를 수행한다.
각각의 layer가 teacher model과 비슷해지는 것이 목적이다. 특히, mobileBERT와 IB-BERT의 MSE가 knowledge transfer objective에 사용된다.
: index of layers
: sequence length
: featurer map size
중요한 아이디어는, 잘 최적화 된 teacher에서 self-attention map을 사용해서, mobileBERT를 학습하면 되지 않을까?라는 생각이다. 특히, mobileBERT와 IB-BERT의 헤드 별 self-attention의 분포들 사이의 KL-divergence를 최소화하는 것으로 학습한다.
: number of attention heads