ALBERT: A Lite BERT for Self-supervised Learning of Language Representation (ALBERT) (Zhenzhong Lan et al.) - [2020 ICLR]

 

 

본 논문에서 제안한 ALBERT는 BERT의 한계를 극복하기 위해 제안된 경량화 언어모델로, 핵심은 파라미터 효율성과 문맥 이해력 개선에 있다. 기존 BERT는 임베딩 차원과 은닉 차원이 강하게 묶여 파라미터 수가 폭발적으로 증가했는데, ALBERT는 Factorized Embedding Parameterization를 통해 어휘 임베딩 차원과 은닉 차원을 분리하여 파라미터를 크게 줄였다. 또한 Cross-layer Parameter Sharing를 도입해 모델 깊이가 늘어나도 파라미터 수가 증가하지 않도록 했으며, BERT의 NSP 대신 SOP를 설계해 문장 간 일관성(coherence)을 학습하게 했다. 이 결과, ALBERT는 BERT-large보다 훨씬 적은 파라미터로도 GLUE, SQuAD, RACE 등에서 새로운 SOTA를 달성하며, 특히 다문장 이해가 필요한 Task에서 큰 성능 향상을 보였다. 즉, ALBERT의 핵심 기여는 더 적은 파라미터로 더 크고 강력한 모델을 학습할 수 있도록 하는 설계 혁신에 있다.

 

 

1. Introduction

본 논문은 대규모 사전학습 언어모델의 한계에 주목한다. BERT 이후, 모델 크기를 늘리면 일반적으로 성능이 향상되지만, GPU/TPU 메모리 한계와 Communication Overhead(분산 학습이나 병렬 학습에서 GPU가 서로 정보를 주고받는 비용)로 인해 더 큰 모델을 학습하기 어려워진다. 기존 방법들은 모델 병렬화나 메모리 관리 기법을 통해 일부 문제를 완화했지만, Communication Overhead는 여전히 해결되지 않았다.

이에 따라 저자들은 ALBERT(A Lite BERT) 를 제안한다. ALBERT는 다음 두 가지 기법으로 파라미터 효율성을 획기적으로 개선하였다.

1. Factorized Embedding Parameterization: 어휘 임베딩을 저차원 공간으로 Projection한 뒤 은닉 차원으로 다시 사상하는 방식으로, 어휘 임베딩 크기와 은닉 차원을 분리하였다.
2. Cross-layer Parameter Sharing: Transformer 계층 간 가중치를 공유하여 깊이가 깊어져도 파라미터 수가 선형적으로 증가하지 않도록 하였다.

또한 BERT의 Next Sentence Prediction(NSP) 대신 Sentence Order Prediction(SOP) 을 제안하여 문장 간 일관성을 더 잘 학습할 수 있도록 했다. 이 조합은 파라미터 수를 크게 줄이면서도 성능을 유지하거나 오히려 향상시켰다.

결과적으로, ALBERT는 GLUE, RACE, SQuAD 등 주요 벤치마크에서 새로운 SOTA를 달성했다. 예컨대 RACE 데이터셋에서 정확도를 89.4%까지 끌어올리며 기존 대비 큰 폭으로 개선하였다.

 

 

2. Related Work

2.1 Scaling Up Representation Learning For Natural Language

Word2Vec, GloVe 같은 정적 임베딩에서 시작해, ELMo·BERT 같은 맥락적 언어모델로 발전하면서 모델 크기와 성능이 비례 관계를 보여왔다. 그러나 모델 크기가 커질수록 학습과 추론이 느려지고 자원 소모가 급격히 증가한다. 기존에는 Gradient Checkpointing이나 Model Parallelism을 통해 메모리 문제를 해결하려 했지만, 속도 저하가 불가피했다.

2.2 Cross-Layer Parameter Sharing

Universal Transformer(UT), Deep Equilibrium Model(DQE) 등에서 시도된 바 있으나, 이들은 주로 언어모델링에 초점이 있었다. ALBERT는 이를 사전학습+미세조정(Pretraining + Finetuning) 구조에 적용해, 효율성과 성능을 동시에 확보했다.

2.3 Sentence Ordering Objectives

기존 BERT의 NSP는 효과가 불안정했고 주로 Topic 예측에 치우쳐 있었다. ALBERT는 대신 SOP를 도입하여 문서 내 연속된 두 세그먼트의 순서를 맞히도록 했다. 이는 Coherence 학습에 집중하게 만들며, 다문장 이해가 필요한 Task에서 효과적임을 보였다.

 

 

3. The Elements of ALBERT

3.1 Model Architecture Choices

ALBERT의 Backbone은 Transformer Encoder 구조를 유지한다. 그러나 다음 세 가지 주요 차별점을 도입했다.

1. Factorized Embedding Parameterization
   • BERT에서는 임베딩 차원(E) = 은닉 차원(H)으로 강하게 묶여 있어, H가 커질수록 어휘 임베딩 행렬(V×H)이 지나치게 커졌다.
   • ALBERT는 어휘를 먼저 낮은 차원(E)으로 임베딩한 후 은닉 차원(H)으로 Projection하는 2단계 구조를 사용했다.
   • ⇒ 이를 통해 파라미터 수가 O(V×H)에서 O(V×E + E×H)로 줄어들며, V=30k, E=128, H=4096일 때 효과가 극대화된다.
2. Cross-layer Parameter Sharing
   • Transformer 계층별로 독립된 파라미터를 두는 대신, 모든 층이 같은 가중치를 공유한다.
   • 실험적으로 FFN(Feed-Forward Network) 공유가 성능을 가장 저해하고, Attention 공유는 성능 저하가 거의 없음을 확인했다.
   • 최종적으로 All-Shared 전략을 채택하여 성능은 좀 낮춰지되, 파라미터 효율을 극대화했다.
3. Sentence Order Prediction (SOP)
   • NSP 대신 SOP를 도입. 양성 샘플: 문서 내 연속된 두 세그먼트, 음성 샘플: 같은 두 Segment의 순서만 바꾼 것
   • NSP는 서로 다른 문서에서의 두 Segment가 같은 문서에서 연속된 것인지 맞추는 것이기 때문에 Topic 예측에 특화될 수 밖에 없는데, 같은 문서에서 순서가 다른 두 Segment가 연속된 것인지 맞추는 SOP를 도입함으로써 Coherence를 더 잘 이해할 수 있게 된다. 이를 통해 모델은 Coherence 일관성을 학습하고, 다문장 태스크에서 성능이 개선된다.

3.2 Model Setup

• ALBERT-large: 18M 파라미터 (BERT-large 334M 대비 18배 적음)
• ALBERT-xlarge: 60M 파라미터
• ALBERT-xxlarge: 235M 파라미터 (여전히 BERT-large보다 적지만 성능은 더 높음)

 

 

4. Experimental Results

4.1 실험 설정

• 데이터: BOOKCORPUS(16GB), Wikipedia(16GB)
• 입력 최대 길이: 512
• MLM: n-gram 마스킹 (최대 3-gram)
• Optimizer: LAMB, lr=0.00176, batch=4096
• 학습 스텝: 125k (추가로 1M+까지도 수행)
• TPU V3, 64~512개 사용

4.2 평가 벤치마크

• Intrinsic: MLM 및 문장 분류 정확도
• Downstream: GLUE, SQuAD v1.1/v2.0, RACE

4.3 BERT vs ALBERT 비교

• ALBERT-xxlarge는 BERT-large보다 파라미터는 70% 수준이지만 성능은 크게 향상
• 예: SQuAD2.0 +3.1%, RACE +8.4%
• 속도: ALBERT-large는 1.7배 빠름, ALBERT-xxlarge는 구조가 커져 3배 느림

4.4 Factorized Embedding 효과

• FFN 공유 시 성능 저하 크지만, Attention 공유는 거의 영향 없음
• 모든 공유(all-shared)를 기본으로 선택

4.6 SOP 효과

• NSP는 Topic 전환만 학습해 SOP task를 풀지 못함
• SOP는 NSP도 어느 정도 해결 가능
• Downstream에서 SQuAD +1~2%, RACE +1.7% 개선

4.7 동일 시간 학습 비교

• ALBERT-xxlarge(125k step, 32h) vs BERT-large(400k step, 34h) → ALBERT가 여전히 우세

4.8 추가 데이터 및 Dropout

• 추가 데이터(XLNet, RoBERTa 수준)를 사용하면 성능 상승
• Dropout 제거 시 성능 개선 (대규모 Transformer에서 dropout이 오히려 해로울 수 있음)

4.9 SOTA

• GLUE 평균 89.4
• SQuAD v2.0 F1 92.2
• RACE 정확도 89.4 (이전 최고 대비 +6~8% 향상)

 

 

5. Discussion

• ALBERT는 BERT 대비 극단적으로 적은 파라미터로도 더 좋은 성능을 달성했다.
• 그러나 구조가 매우 크기 때문에 학습·추론 속도는 여전히 부담이 된다. 향후에는 Sparse Attention, Block Attention 등 효율적 구조가 필요하다.
• SOP는 NSP보다 효과적이지만, 언어 이해에 필요한 다른 차원의 self-supervised loss도 탐구할 여지가 많다.