MobileBERT
스마트폰에서도 빠르게 돌아가는 BERT
자연어 처리(Natural Language Processing)는 이전부터 많이 연구되어 온 분야이다. 번역부터 시작해서 글 요약, 문장의 감정 분석 등 일상생활에서 볼 수 있는 많은 문제들이 자연어 처리에 해당된다. 이러한 자연어 처리는 구글이 2018년도에 발표한 BERT를 기점으로 급격하게 발전하기 시작했다. BERT는 기존의 자연어 처리 문제들에서 SOTA(State-of-the-Art)를 달성했고, 비교적 간단한 fine-tuning을 통해 다양한 자연어 처리 문제들에 빠르게 적용(task-agnostic)할 수 있다는 장점이 있다. 그래서 이후의 연구들(RoBERTa, XLM 등)은 대부분 BERT 또는 BERT에 사용된 transformer의 응용이다. BERT와 관련된 정보들은 구글링 해보면 쉽게 찾을 수 있다. 당장 brunch에서도 BERT를 소개한 글을 찾아볼 수 있다.
BERT에도 치명적인 단점이?
BERT가 자연어 처리 연구의 발전에 엄청난 기여를 한 것은 명백한 사실이다. 인터넷에서 쏟아지는 BERT와 관련된 정보들이 이를 보여준다. 하지만 BERT에도 치명적인 단점이 존재한다. 바로, 너무 크다는 것이다. 성능은 좋지만 parameter 수가 너무 많아서, inference 하는 것이 오래 걸린다는 단점이 있다. BERT-base 모델의 경우, Pixel 4 스마트폰에서 inference 하는데 0.3 초 정도의 시간이 걸린다. 물론 GPU가 있으면 어느 정도 완화할 수 있다. 하지만 실시간 번역이나 자동완성과 같은 문제들은 GPU가 없는 스마트폰에서도 자주 사용되는 기능들이다. 타자를 누르는 속도를 생각해보면, inference 하는데 걸리는 0.3초만으로 만족스러운 UX(User Experience)를 만들 수 없다는 것을 알 수 있다. 결과적으로 BERT는 GPU가 없는 스마트폰과 같은 환경, 이른바 resource-limited mobile device에서는 제대로 활용할 수 없다.
그래서 현재는 성능이 다소 떨어지더라도, inference가 빠른 모델들을 스마트폰에서 사용하고 있다. 그렇다면 우리는 언제까지 성능이 떨어지는 자연어 처리 기술들을 스마트폰에서 경험해야 할까? 최근, 자연어 처리 학회 ACL 2020에서 이 문제에 대해 연구한 논문을 공개하였다. 그것이 바로 MobileBERT(https://arxiv.org/pdf/2004.02984.pdf)이다. MobileBERT는 BERT와 비슷한 구조를 가지면서 BERT-large보다 4.3배 작고, 5.5배 빠른 inference time을 보여줬다. 그러면서 BERT-base와 유사한 성능을 보여줬다. 즉, 기존의 BERT와 비슷한 성능을 내면서 스마트폰에서도 사용할 수 있는 모델이라고 보면 된다. 이외에도 BERT의 장점 중 하나인 task-agnostic도 그대로 이어받았다. 그래서 논문에서는 MobileBERT를 compact task-agnostic 모델이라고 소개하였다. 다른 말로 BERT의 유의미한 부분만 남기면서(compact), 여전히 task-agnostic 하다는 것이다. 과연, MobileBERT는 어떻게 BERT의 장점들을 남기면서 훨씬 빠른 모델이 되었을까?
MobileBERT
가벼운 BERT를 만드는 가장 쉬운 방법은, 말 그대로 BERT를 가볍게 만든 다음 학습하는 것이다. 하지만 여러 연구들에서 이 방법은 좋은 성능을 보여주지 못했다. 이는 가벼운 모델이 다양한 자연어 처리 문제들에 대해서, 텍스트의 의미 있는 정보들을 모두 담을 만큼(representation power) 충분히 크지 않기 때문이다. 다른 방법은 BERT-large 만큼 복잡한 모델을 학습한 뒤, 이 모델에서 학습한 개념들을 가벼운 모델로 전이하는 것이다. MobileBERT는 이 방법을 사용하였고, knowledge transfer라고 부른다. 여기서 말하는 개념과 전이하는 방법은 뒤에서 자세히 설명할 것이다. 아무튼 이 방법은 마치 선생이 학생을 가르치는 것처럼 보여서 복잡한 모델은 teacher model, 가벼운 모델은 student model이라고 부른다.
IB-BERT(b)와 MobileBERT의 개략적인 구조위 그림의 왼쪽이 teacher model, 오른쪽이 student model이다. 여기서 student model이 바로 BERT를 가볍게 만든 MobileBERT이다. Teacher model은 BERT-large와 거의 유사하지만, 완전히 같지 않아 IB-BERT(Inverted-Bottleneck BERT)라고 부른다.
모델 설계
IB-BERT와 MobileBERT의 구조, 그리고 그 구조를 만들어내기 위해 진행한 실험들을 살펴보자.
주요 모델들의 parameter 수를 나타낸 도표Layer 크기/Head 수 Parameter 수를 보여주는 도표를 보면 MobileBERT의 구조는 BERT와 거의 똑같고, 각 layer의 크기만 다른 것을 알 수 있다. BERT-large에서는 feature map의 차원이 항상 1024로 유지되고 있는 반면, MobileBERT는 128로 유지하고 있다. 128로 맞춘 명확한 이유는 밝히지 않았지만, 아마 embedding 차원과 동일하게 맞추기 위해 그런 것이라고 추측한다. 실제로 BERT도 그렇기 때문이다. 그리고 IB-BERT, MobileBERT의 MHA에서 head의 수가 16에서 4로 줄은 것을 확인할 수 있다. 이는 IB-BERT를 학습하는 과정에서 head 수를 줄여도 성능이 거의 변하지 않는다는 것을 확인하였기 때문이다. 반면에 IB-BERT의 feature map 차원은 BERT-large와 동일하다. 이는 학습하는 과정에서 feature map 차원을 줄이면 급격한 성능 저하가 일어나는 것을 확인하였기 때문이다. 논문에서는 이 feature map들이 non-linear 한 관계들을 규명하고 있기 때문에, 성능에 많은 영향을 미친다고 해석했다.
Bottleneck IB-BERT와 MobileBERT는 기존의 BERT와 다르게 transformer가 두 개의 linear layer로 둘러싸여 있다. 이 두 개의 linear layer는 차원을 축소/확장함에 따라 bottleneck/inverted-bottleneck이라고 부른다. 이 두 개의 bottleneck을 추가한 이유는, transformer의 input/output 차원을 변경하기 위해서이다. 각 transformer에서 나오는 feature map들이 가장 의미 있는 정보를 가지고 있기 때문에, 얘네를 줄일 수 있을 만큼 줄여야 진정한 의미의 compact 한 모델을 얻을 수 있다고 생각한 것 같다. 그래서 transformer의 input/output 차원을 다르게 하면서, IB-BERT를 train 하는 실험을 진행하였다. 결과적으로 기존의 1024에서 512까지 줄여도 성능에 문제가 없다는 것을 확인할 수 있었다. 또 살펴보아야 할 점은 MobileBERT도 똑같이 transformer의 intput/output 차원이 512가 되도록 bottleneck을 추가했다는 것이다. 이를 똑같이 맞춘 이유는 뒤에 나올 knowledge transfer에서 설명하도록 하겠다.
Embedding 논문에서는 embedding 차원이 모델의 전체 parameter 수에 많은 영향을 미친다고 설명했다. 그래서 IB-BERT와 MobileBERT는 기존의 embedding 차원을 1024에서 128로 확 줄였다(128로 선정한 특별한 이유는 나와있지 않다). 그리고 bottleneck 실험을 통해 얻어진 최적의 feature map 차원인 512로 차원을 늘리기 위해 여러 개의 1D convolution을 거친다.
Stacked FFNs 기존의 BERT에서는 FFN의 parameter 수가 MHA의 2배이다. 하지만 MobileBERT에 위에서 언급한 방법들을 모두 적용하면 MHA가 FFN보다 parameter 수가 많아지는 문제가 생긴다. 논문에서는 기존의 parameter 수 균형을 맞추기 위해 stacked feed-forward network라는 것을 사용하였다. 어려운 용어인 것 같지만, FFN의 parameter 수가 부족하니 FFN을 더 쌓는다는 단순 무식한 해결법이다. 단순 무식하지만, 실제 실험 결과는 좋게 나왔다. 정확히는 MHA와 stacked FFN의 parameter 수 비율이 0.4~0.6이 되도록 FFN을 쌓으면 최적의 결과가 나왔다. 그러면서 일반적인 transformer에서 이 비율이 중요하다는 사실을 증명했다고 첨언하였다.
LayerNorm / GeLU 기존의 BERT는 각 MHA, FFN 뒤에 layer normalization를 거쳤다. 그리고 activation으로 GeLU를 사용하였다. 논문에서는 이 두 단계가 inference time에 많은 영향을 준다고 생각하여, layer normalization은 아예 제외했고 GeLU는 ReLU로 변경하였다. 실험 결과 192ms의 inference time이 62ms, 약 3배 정도 빨라진 것을 확인할 수 있었다.
Knowledge Transfer
Knowledge transfer에 대해서 다루기 전에, MobileBERT로 특정 자연어 처리 문제를 해결하기 위해 학습하는 과정을 정리해보자.
1. IB-BERT를 BERT와 같은 방식으로 pre-train 한다.
2. Pre-train 된 IB-BERT를 이용하여 MobileBERT를 학습한다(Knowledge Transfer).
3. 풀고자 하는 자연어 처리 문제에 맞춰서 MobileBERT를 fine-tune 한다.
1과 3은 기존의 BERT 학습 방법과 동일하기 때문에, 2번 과정에 대해서만 살펴보면 된다. Knowledge transfer를 직역하면 개념을 전이한다는 의미이다. 그렇다면 여기서 말하는 '개념'이란 것이 무엇이고, '전이'하는 방법은 무엇인가?
Transformer's Output Teacher model이 학습한 개념 중 하나는, 각 transformer의 output에 해당하는 feature map이다. 위에서도 언급했듯이 논문에서는 각 transformer의 output이 가장 의미 있는 정보들을 지니고 있다고 생각했다. 그래서 이를 줄이기 위한 실험도 진행하였다. 그만큼 transformer의 output은 중요하고, student model에 전이하면 좋은 개념이다. 그렇다면 각 transformer의 output들을 어떻게 student model로 전이할까? Bottleneck을 적용하는 부분을 다시 읽어보면, 논문에서는 IB-BERT와 MobileBERT의 transformer output 차원을 동일하게 만들었다. 이러면 전이하는 것이 간단해진다. 말 그대로, student model이 teacher model과 똑같은 feature map이 나올 수 있도록 학습하면 되는 것이다. 그래서 전이하는 과정은 student, teacher model의 각 transformer에 대해서, 서로의 output의 MSE loss를 줄이는 방향으로 학습하는 것이다. 둘이 동일한 output 차원을 가지고 있다는 것을 알고 있으면, 가장 먼저 생각할 수 있는 학습 방법일 것이다.
Self-Attention 이외에도 제시한 개념이 하나 더 존재한다. 바로, MHA의 각 head에서 학습하는 self-attention이다. BERT에서 사용한 transformer는 self-attention 문제를 획기적으로 해결한 방법 중에 하나로, 그만큼 attention이 BERT에서 중요하다는 것을 보여준다. 그래서 논문에서는 self-attention도 teacher model이 배우는 중요한 개념으로 보고, 전이하였다. 마찬가지로 head 수가 동일하기 때문에 전이하는 방법이 쉽다. 각 head에서 학습한 self-attention의 분포가 동일해지도록 만들면 된다. 이는 두 분포의 KL divergence를 최소화하는 방법으로 해결할 수 있다.
Distillation 사실 knowledge transfer에 대한 설명은 위에서 끝났다. 하지만 논문에서는, knowledge transfer 이외에 정제(distillation)하는 과정을 거친다. 이는 기존의 BERT pre-training 방법인 MLM(Masked Language Modeling)과 NSP(Next Sentence Prediction)을 knowledge transfer가 끝난 MobileBERT에 적용하는 것이다. 물론 이 두 개의 loss에 MLM KD(Knowledge Distillation) loss까지 추가하지만, 기존의 BERT 학습 과정과 동일하기 때문에 자세한 설명은 생략한다. 아무튼 MobileBERT의 정제는 MLM, MLM KD, NSP loss를 linear combination 한 loss를 최소화하는 방향으로 학습함으로써 이루어진다.
Training Strategy
Knowledge transfer과 distillation에 대해 살펴보았다. 여기서 중요한 것은 knowledge transfer는 각 transformer에 대해서 loss가 나오기 때문에, 이들을 학습하는 방법이 여러 개가 나올 수 있다는 것이다. 가장 쉽게 생각할 수 있는 방법은 각 transformer에서 나온 loss들을 모두 linear combination 하여, 이를 최소화하는 방향으로 학습하는 것이다. 하지만 MobileBERT의 base가 되는 BERT의 경우, 매우 깊기 때문에 제대로 학습이 안될 가능성이 높다.
그래서 논문에서는 가장 얕은(input과 가장 가까운) transformer부터 학습해서 어느 정도 수렴하면 freeze 하고 다음 transformer를, 이 과정을 반복해서 모든 transformer를 학습하는 방법을 채택하였다. 즉, 각 단계에서 하나의 transformer에 대해서만 knowledge transfer가 이루어진다. 이를 progressive knowledge transfer라고 부른다. 모든 transformer에 대해서 knowledge transfer가 완료되면, distillation 과정을 거쳐 마무리한다.
MobileBERT를 pre-train 하는 과정실험 결과
실험은 language modeling 문제에서 자주 쓰이는 GLUE와 SQuAD 데이터를 사용하였다. 자세한 결과는 논문에서 확인할 수 있으며, 중요한 결과들을 요약하면 다음과 같다.
MobileBERT는 BERT보다 빠르면서, BERT-base와 비슷한 성능을 보인다.
Bottleneck은 효과적인 knowledge transfer를 가능하게 만들었다.
Progressive Knowledge Transfer를 사용하여, MobileBERT를 효율적으로 학습할 수 있었다.
MobileBERT가 ACL 2020에서 공개된 지 1달도 되지 않았다. 그만큼 검증이 덜 된 상태이다. 연구 자체도 영어로 진행되었고, 당장 다른 영어 자연어 처리 문제들에 적용한 사례들도 많이 나오지 않았다. 그래서 한글에서 비슷한 성능이 나올지 의문이다. 하루빨리 다양한 task들에 대해서 검증이 되어, 우리들이 사용되는 스마트폰에서도 사용할 수 있는 날이 왔으면 좋겠다.
