ALBERT: A Lite BERT for Self-supervised Learning of Language Representations

Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut, [Google Research][Toyota Technological Institute at Chicago], (2019.09)

 

Abstract

• 자연어 표현을 사전학습 시킬 때 model size를 키우면 대체로 downstream task의 성능이 향상된다
• 하지만 model이 커짐에 따라 training time과 GPU, TPU memory 한계의 문제를 겪게 됨
• 이러한 문제를 해결하기 위해 BERT보다 더 적은 memory를 소모하고 학습 속도가 빠른 ALBERT를 제시
• inter-sentence coherence를 모델링하기 위해 self-supervised loss를 도입하였으며 이는 multi-sentences input을 가지는 downstream task에서 효과가 있었음
• ALBERT는 GLUE, RACE, SQuAD에서 BERT-large보다 더 적은 parameters를 사용하면서 SOTA를 달성
• 코드와 pretrained model은 github에 공개 (https://github.com/google-research/ALBERT)

---

1. Introduction

• language representations을 pretraining 할 때 large model을 사용해야 좋은 결과를 얻을 수 있었으며 large model을 학습시키고 비교적 작은 모델에 distillation 시키는 방식이 common practice
• 모델의 size는 물론 매우 중요하나 저자는 다음과 같은 질문을 던진다

Is having better NLP models as easy as having larger models?

• 위 질문에 대한 답을 하기 위해서 넘어야 할 가장 큰 장애물은 hardware의 memory 한계
• distributed training에서 communication overhead의 경우 model parameters에 비례하기 때문에 큰 모델의 경우 학습 속도가 크게 느려진다
• 이전 연구들은 model parallelization, memory management로 memory 문제를 해결하려 하였으나 communication overhead를 해결하지 않았다
• 기존 BERT 보다 적은 parameters를 가지는 A Lite BERT(ALBERT)를 제시하여 앞서 언급한 모든 문제를 해결
• ALBERT는 pre-trained model을 scaling 하는데 주요한 문제를 해결하기 위해 2개의 parameter reduction 기술을 사용
• 첫 번째 parameter reduction 기술은 factorized embedding parameterization
  • 큰 vocabulary embedding matrix를 작은 2개의 matrices로 decompose
  • hidden layer size를 vocabulary embedding size와 분리 (vocabulary embedding의 parameter를 크게 증가시키지 않으면서 hidden size를 증가시킬 수 있음)
• 두 번째 parameter reduction 기술은 cross-layer parameter sharing
  • 모델의 depth가 깊어짐에 따라 parameters 수가 늘어나는 것을 방지
• 두 parameter reduction 기술은 BERT의 성능을 크게 해치지 않으면서 parameters를 크게 줄였다
• ALBERT는 BERT-large와 비슷한 세팅으로 BERT-large보다 18배 적은 parameters로 1.7배 빠르게 학습
• parameter reduction 기술들은 일반화 성능, 학습 안정성에 도움이 되는 regularization의 역할도 한다
• ALBERT의 성능을 더 끌어올리기 위해 **sentence-order prediction(SOP)**를 위한 self-supervised loss를 도입
• SOP는 inter-sentence coherence에 집중하며 original BERT에서 제안된 next sentence prediction loss의 비효율을 해결하기 위해 디자인되었다
• 언급한 여러 design decisions으로 ALBERT는 BERT-large보다 적은 parameters를 가지며 모델을 scale up 할 수 있었으며 성능 또한 더 좋았다
• GLUE benchmark: 89.4, SQuAD 2.0 f1 score: 92.2, RACE accuracy: 89.4%로 SOTA를 달성하였다.

---

2. The Elements of ALBERT

Model Architecture Choices

• ALBERT는 transformer의 encoder + GELU activation을 사용하는 BERT와 유사한 백본 아키텍처를 가짐
• BERT의 notation convention을 그대로 사용 (E = embedding size, L = num of encoder layers, H = hidden size, filter size = 4H, num of attention head = H/64)

 

Factorized embedding parameterization

• BERT와 그 후속 논문들(XLNet, RoBERTa)에서 WordPiece embedding size E는 hidden layer size H와 묶여있는데(ex. E ≡ H) 이는 모델링 측면, practical 측면에서도 최적의 모델 디자인이 아니다.
• 모델링 측면에서 보면 WordPiece embedding은 context-independent representations을 학습하도록 설계되었고, hidden-layer embedding은 context-dependent representations을 학습하도록 설계되었다
• BERT와 같은 representations의 힘은 context-dependent representations을 학습하기 위한 신호를 제공하는 context의 사용에서 온다
• 그러므로 WordPiece embedding size E와 hidden layer size H가 비슷해야 한다는 제약에서 벗어나면 전체 모델 parameters를 모델링 목적에 맞게 더 효율적으로 사용할 수 있게 된다 ( H >> E )
• 현실적인 측면에서 보면 NLP에서 보통 큰 vocab size V를 요구하는데 만약 E가 H와 비슷하다면 H를 증가시키면 embedding matrix (V x E shape) 또한 커지게 된다
• embedding matrix가 커지다 보면 모델이 수십억의 parameters를 가지게 되는 결과를 불러오고 training 도중 sparsely update만 가능하게 되는 문제가 발생
• 그러므로 ALBERT는 embedding parameters를 2개의 작은 matrices로 나누는 factorization 사용
• one-hot vector를 바로 H size의 hidden space로 projection 하지 않고 더 낮은 차원의 size E를 가지는 embedding space로 projection을 거치고 hidden space로 projection
• 위 과정을 통해 원래 (V x H) 개의 embedding parameters를 (V x E + E x H) 개로 줄일 수 있게 되었다 (E가 작아야 효과적)
• 단어마다 다른 embedding size를 가지는 것이 중요한 whole-word embedding이 아닌 문서에 고르게 분포된 word pieces를 사용하기 때문에 E는 고정값으로 설정하였다

 

Cross-layer parameter sharing

• feed-forward network의 parameters만 share 하거나 attention parameters만 share 하는 등 많은 parameter sharing 방법이 있는데 ALBERT는 layer 간 모든 parameters를 share
• Transformer networks에서 비슷한 parameter sharing을 제안한 Universal Transformer (Dehghani et al. 2018), Deep Equilibrium Models (Bai et al. 2019)
• layer들의 input embedding, output embedding의 L2 distance와 cosine similarity는 아래와 같다

• BERT보다 ALBERT에서 layer 사이의 값이 더 smooth 한 것을 볼 수 있는데 이는 ALBERT의 sharing parameters가 network parameter 안정화 효과를 가짐을 뜻한다

 

Inter-sentence coherence loss

• BERT는 masked language modeling(MLM) loss와 next-sentence prediction(NSP)을 사용
• NSP는 sentence pairs의 관계를 추론해야 하는 natural language inference와 같은 downstream tasks의 성능을 향상시키기 위해 디자인된 loss
• 하지만 후속 연구들은(XLNet, RoBERTa) NSP의 효과에 의문을 가지며 이를 제거하였고 여러 downstream task에서 성능 향상에 성공하였음
• ALBERT의 저자는 NSP의 효과가 없는 이유를 MLM에 비해 task 자체의 난이도가 낮아서라고 추측하였다
• NSP는 topic prediction과 coherence prediction이 뒤섞인 single task이며 topic prediction은 coherence prediction 보다 더 쉽고 이미 MLM loss도 topic prediction의 역할을 어느 정도 수행하고 있음
• inter-sentence coherence에 집중하기 위해 ALBERT는 sentence-order prediction(SOP) loss를 사용
• SOP의 positive examples은 BERT와 동일하게 같은 document 내의 연속되는 2개의 segments로 구성하며, negative examples은 positive sample과 동일하게 2개의 segments를 뽑지만 둘의 순서를 바꾸어 구성한다
• SOP는 모델이 대화와 같은 고수준의 일관적인 속성에 대한 세분화된 차이를 학습하게 강제한다
• NSP는 SOP task를 풀 수 없지만 SOP는 적당하게 NSP task를 해결할 수 있음을 관찰하였음
• 결과적으로 ALBERT는 multi-sentence encoding tasks에서 일관적인 성능 향상을 보였다

---

3. Experimental Results

Model Setup

 

 

Experimental Setup

• BERT와 동일한 BookCorpus, English Wikipedia 데이터셋 사용
• input token: [CLS] segment1 [SEP] segment2 [SEP]
• XLNet과 같이 30,000 vocab size, SentencePiece tokenize 사용
• MLM의 mask를 n-gram masking을 사용하여 생성, n=3
• batch size = 4096, learning rate = 0.00176, LAMB optimizer 사용
• Cloud TPU V3에서 125,000 steps 학습

 

Evaluation Benchmarks

• ALBERT xxlarge는 더 적은 parameters로 BERT large 보다 좋은 성능을 보임

• 중간 embedding space size E에 따라 parameters가 줄어드는 정도가 다르며 parameter sharing을 적용하면 크게 parameters 수 감소, 성능도 약간 감소

• all-shared를 적용하면 성능이 하락, 성능은 E가 클 때 더 큰 폭으로 하락하는데 이는 FFN-layer paramters를 sharing 하는 것이 성능 하락의 주 원인임을 보여준다

• SOP는 여러 downstream task에서 일관적으로 좋은 결과를 이끌어냈다

• ALBERT-xxlarge는 BERT-large 보다 약 3배 이상 학습이 느림, 같은 시간 동안 학습하고 결과를 비교해도 ALBERT-xxlarge의 성능이 더 뛰어났다

• 1M step을 학습하여도 overfitting이 관찰되지 않았으며 CNN에서 batchnorm과 dropout의 조합이 성능을 해친다는 이전 연구의 결과와 비슷하게도 ALBERT에서 dropout을 제거하니 성능이 향상되었다

• GLUE, SQuAD, RACE에서 SOTA를 달성

---

4. Discussion

• ALBERT-xxlarge는 BERT-large 보다 적은 parameters로 더 좋은 결과를 기록하였으나, 거대한 structure로 인해 computation cost가 큼
• 중요한 다음 스텝은 ALBERT의 학습, 추론 속도를 sparse attention, block attention과 같은 method를 사용하여 개선하는 것
• 효율적인 언어모델 학습, 예시 추출도 추가적인 representation power를 제공할 수 있는 연구 방향
• SOP가 일관적으로 유용한 pretraining task인 증거를 보였지만 제시한 self-supervised training loss로 포착할 수 없는 representation power를 증가시킬 수 있는 다른 수단이 있을 수 있다고 생각 (pretraining task에 대한 future work 제시)

---

후기 & 정리

• Factorized embedding parameterization, Cross-layer parameter sharing로 parameters의 수를 줄이고 Inter-sentence coherence loss를 도입한 ALBERT를 제시
• ALBERT-xxlarge는 BERT-large 보다 적은 parameters로 GLUE, SQuAD, RACE의 여러 task에서 SOTA를 달성
• 중간 embedding size E, parameter sharing 방식(all-layer, only FFN, only attention)의 변화에 따른 결과를 모두 제시
• 후속 연구로는 ALBERT의 학습, 추론 속도 개선, representation power를 끌어올리기 위한 example mining, 언어모델 학습, pretraining task 선정을 제시
• NSP의 효과가 없는 이유에 대한 분석과 대안 SOP 제시, SOP로 학습한 모델로 NSP, NSP로 학습한 모델로 SOP를 수행하는 과정이 흥미로웠음

 

Reference

[0] Zhenzhong Lan et al. (2019). "ALBERT: A Lite BERT for Self-supervised Learning of Language Representations". https://arxiv.org/abs/1909.11942

ALBERT: A Lite BERT for Self-supervised Learning of Language Representations