Paper

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

gyuilLim 2025. 5. 16. 10:34

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

Increasing model size when pretraining natural language representations often results in improved performance on downstream tasks. However, at some point further model increases become harder due to GPU/TPU memory limitations and longer training times. To

arxiv.org

Author : Lan, Zhenzhong, et al.
Journal : ICLR 2020
Keyword : ALBERT
Published Date : 2019년 9월 26일

Abstract

자연어 표현을 사전학습할 때 모델의 크기를 키우면, 일반적으로 다운스트림 태스크에서의 성능이 증가한다.
하지만 TPU/GPU 메모리는 제한되어 있기 때문에, 무작정 모델의 크기를 키우는 것은 불가능하다.
이러한 문제를 해결하기 위해 BERT의 학습 속도를 증가시키고 메모리 사용량을 줄이는 두 가지 방법을 제안한다.
또한 Inter-setence coherence를 모델링하는 Self-supervised Loss를 사용한다.

Introduction

NLP의 다양한 벤치마크에서의 성능 향상 사례들은 보통 SOTA 성능을 달성하기 위해서는 큰 규모의 네트워크가 중요하다는 점을 보여준다.
또한 사전학습된 대규모 모델을 더 작은 모델로 증류(Distill)하여 실제로 사용하는 것이 관행이 되었다.
하지만 모델 크기의 중요성을 고려할 때, 단순한 크기 증가만으로 성능이 충분히 보장될까? 라는 의문이 생긴다.
최신 SOTA 모델들은 수억, 수십억개의 파라미터를 갖고 있기 때문에, 메모리 한계에 쉽게 부딪히고 Multi GPU를 사용하기 때문에 GPU 간 Communication overhead로 인해 학습 속도 또한 크게 느려질 수 있다.

이런 문제들의 기존 해결책으로는 Model Parallelization, Clever Memory Management 등이 있다.

하지만 이 방법들은 메모리 한계 문제는 해결해도 Communication overhead는 해결하지 못한다. 이 논문에서는 이러한 문제들을 해결하기 위해 BERT 구조에서 파라미터 수를 크게 줄인 ALBERT(A Lite BERT)를 제안한다.
ALBERT는 두 가지 파라미터 감소 기법을 도입했다.
- Factorized Embedding Parameterization
  - Vocabulary 행렬을 두 개의 작은 행렬로 분해한다.
  - Hidden Layer의 크기를 키워도 Vocabulary Embedding의 크기는 크게 증가하지 않는다.
- Cross-Layer Parameter Sharing
  - 네트워크의 깊이가 깊어져도 전체 파라미터는 증가하지 않도록함.
이러한 방법들을 적용하여 ALBERT는 BERT와 비교했을 때 약 18배 적은 파라미터, 1.7배 빠르게 학습된다.
성능을 더 향상시키기 위해 문장 순서 예측, SOP(Sentence Order Prediction)이라는 손실 함수도 도입한다.

SOP는 문장 간 일관성에 집중할 수 있도록 하고 기존 BERT의 NSP 손실 함수의 비효율성을 극복하기 위해 설계되었다.

따라서 SOP에서는 같은 문단에서 뽑힌 두 문장에 대해 순서를 바꾼 문장 쌍을 Negative로 사용한다.

The Elements of ALBERT

Model Architecture Choices

ALBERT의 모델 아키텍처는 BERT와 유사하다. GELU를 활성화 함수로 사용하는 Transformer 인코더 기반 네트워크이다.
아래와 같이 BERT의 표기법을 따른다.
- $E$ : Vocabulary의 Embedding Size
- $L$ : Encoder Layer 개수
- $H$ : Hidden size
BERT와 같이 FFN의 크기는 $4H$, 어텐션 헤드 수는 $H/64$ 이다.
ALBERT 아키텍쳐에서 세 가지 Contributions는 아래와 같다.
- Factorized Embedding Parameterization
- Cross-Layer Parameter Sharing
- Inter-Sentence Coherence Loss

Factorized Embedding Parameterization

BERT와 BERT 이후 개선 모델들은 $E$와 $H$가 동일하게 설정되어 있다. 하지만 이는 최적이 아니다.
모델링 관점에서, Wordpiece 임베딩은 문맥에 독립적인 표현을, Hidden Layer는 문맥에 의존적인 표현을 학습한다. 즉 Hidden Layer가 Wordpiece 임베딩에 비해 문맥을 포함한 더 많은 정보를 학습하게 된다. 따라서 $H \gg E$로 설정하여야 효율적으로 학습할 수 있다.
실용적인 관점에서 일반적으로 Vocabulary의 크기인 $V$가 크기 때문에, $E \equiv H$ 인 경우, Hidden Layer의 차원을 증가시키면 임베딩 행렬의 크기 $V \times E$ 도 마찬가지로 증가하게 된다. 각각의 임베딩은 대부분 학습 중 드물게 업데이트되기 때문에 비효율적이다.
그래서 ALBERT에서는 임베딩 파라미터를 분해하여, 두 개의 작은 행렬로 나눈다.
$V\times H$를 $V \times E + E \times H$로 분해하는 것이다. $H \gg E$ 일 때 파라미터 감소량이 더 커진다.

Cross-Layer Parameter Sharing

파라미터 효율성을 증가시키기 위해 또 다른 방법으로는 Cross-layer parameter sharing이 있다.
ALBERT에서는 모든 레이어에서 모든 파라미터를 공유한다.
이런 방법을 사용한 선행 연구에서는, Layer 간 파라미터를 공유하면 어느 순간 입력과 출력이 같아지는 지점인 평형 지점에 도달한다는 결과가 있었지만 ALBERT에서는 다른 결과가 관측되었다.

위 그래프는 각 레이어의 입력 임베딩과 출력 임베딩 사이의 L2 Distance와 코사인 유사도를 시각화 한 것이다.
ALBERT의 경우 BERT에 비해 레이어간 전이가 훨씬 부드럽게 진행된다는 것을 확인할 수 있다.
각 그래프에 대해 분석해보면
- L2 Distance가 점점 감소한다. → 값 자체가 비슷해진다. 하지만 0에 수렴하진 않기 때문에, 학습하는 의미가 있다는 것을 시사한다.
- 코사인 유사도가 감소한다. → 입/출력 임베딩의 방향이 변화한다. 마찬가지로 학습하는 의미가 있다는 것을 시사한다.

Inter-Sentence Coherence Loss

BERT에서는 MLM(Masked Language Modeling) 이외에도 NSP(Next Sentence Prediction)라는 손실함수를 사용한다.
NSP는 문장 두 개를 입력으로 받고, 두 문장이 실제로 연속된 문장인지를 예측하는 이진 분류이다. Postive 샘플과 Negative 샘플은 아래와 같이 설정한다.
- Positive : Training Corpus에서 실제로 연속된 문장 쌍
- Negative : 서로 다른 Corpus에서 임의로 뽑은 문장 쌍
- 이 때, Positive와 Negative는 같은 확률로 샘플링된다.
NSP는 NLI(Natural Language Inference) 등 문장 간 관계 정보가 필요한 다운스트림 태스크를 위한 것이었으나 RoBERTa 등 후속 연구에서 NSP를 제거한 결과 오히려 성능이 향상된다는 결과가 있었다.
저자들은 NSP의 효과가 없었던 이유를 MLM에 비해 너무 쉽기 때문이라고 추정한다. 즉 모델은 문장 간 맥락(Context)나 일관성(Coherence)이 아니라 단순히 주제(Topic)의 차이만으로 정답을 맞출 수 있다는 것이다.
따라서 문장 간 일관성에 집중할 수 있는 SOP(Sentence-Order Prediction) Loss를 사용한다. SOP는 주제 예측은 피하고 일관성 학습에 집중할 수 있도록 설계했다.
- Positive : Training Corpus에서 실제로 연속된 문장 쌍
- Negative : Positive 샘플의 순서를 바꾼 경우
이 학습 방식은 더 세밀한 일관성 차이를 학습하도록 한다.

NSP로 학습된 모델은 SOP 태스크에 대한 성능이 낮게 측정된다. 이는 NSP가 단순히 주제 예측만으로 학습할 수 있다는 근거가 된다. 하지만 SOP로 학습한 모델은 NSP 태스크를 어느정도 잘 해결한다.

Experimental Setup

의미있는 비교를 위해 BERT를 따라 BOOKCORPUS, English Wikipedia 데이터셋으로 사전학습한다.
또한 입력 형식도 BERT와 같은 “$[\text{CLS}]\ x_1\ [\text{SEP}]\ x_2\ [\text{SEP}]$”이고 Vocabluary의 크기도 30,000개이다.
최대 3개의 그램까지 마스킹하고 배치는 4096, 학습률은 0.00176, 125,000 스텝 동안 학습한다.

Evaluation Benchmarks

아래와 같은 세 가지 벤치마크를 사용하여 모델을 평가한다.
- GLUE(General Lanuage Understanding Evaluation) : 총 9개의 NLU(Natural Language Understanding) 태스크로 구성된 벤치마크(CoLA, SST 등)로, 자연어 이해 능력 평가.
- SQuAD(Stanford Question Answering Dataset) : v1.1, v2.0 두 가지 버전이 있고 Wikipedia 기반 QA 데이터셋이다.
- RACE(ReAding Comprehension from Examinations) : 중국의 영어 시험에서 수집된 4지선다 독해 데이터셋. 4개의 선택지 중 하나를 선택한다.

Overall Comparison

BERT와 ALBERT 모델의 구성을 나타낸 것이다.
ALBERT-large(18M)가 BERT-large(334M) 보다 파라미터가 18배 작다. ALBERT-xxlarge(235M)은 BERT-large(334M)의 약 70%이다.

ALBERT-xxlarge는 BERT-large의 약 70% 수준의 파라미터로 유의미한 성능 향상을 보여준다.
BERT-large와 ALBERT-xlarge의 성능이 비슷한데, 학습 속도는 ALBERT가 더 느리다. 이는 Layer의 개수가 24개, Hidden 차원이 2048로 연산량이 증가했기 때문이다.
또한 ALBERT-base는 BERT-base와 비교했을 때 BERT-large보다 학습이 5.6배 더 빠르다.

Factorized Embedding Parameterization

ALBERT-base 모델로 Embedding size $E$ 를 변경했을 때 성능을 측정한 결과이다.
Not-shared에서 $E$ 가 커질수록 성능이 증가한다. 하지만 All-shared에서는 $E=128$ 일 때의 성능이 가장 좋다.
따라서 모든 실험에서 $E=128$ 를 사용한다.

Cross-Layer Parameter Sharing

ALBERT-base 모델로 $E=768, 128$ 일 때 각각에 대하여 네 가지 공유 방식의 실험 결과를 보여준다.
1. All-shared : 레이어의 모든 파라미터 공유
2. Not-shared : 공유하지 않음(BERT)
3. Shared-attention : 어텐션 파라미터만 공유
4. Shared-FFN : FFN 파라미터만 공유
All-shared는 E=128 보다 E=768에서 성능 저하가 더 크다. Shared-FFN도 마찬가지로 모두 성능 저하가 크다.
이외에 Shared-attention은 다른 방법에 비해 성능 저하가 작다.
결과적으로, All-shared 방식을 선택한다.

Sentence Order Prediction(SOP)

ALBERT-base 모델로 Inter-sentence loss에 대해 세 가지 실험 조건을 직접 비교한다.
- None : 추가 손실 없음
- NSP : BERT의 Next Setence Prediction 방식
- SOP : ALBERT의 Sentence Order Predcition 방식
실험 결과에 따르면, NSP로 학습한 모델은 SOP를 거의 수행하지 못하나, SOP로 학습한 모델은 NSP도 비교적 잘 수행하는 것을 확인할 수 있다.
또한 SOP로 학습한 모델은 Downstream Task에서 전반적으로 더 좋은 성능을 낸다.

What if we train for the same amount of time?

이전 비교에서 ALBERT-xxlarge가 BERT-large에 비해 약 3배 느린 것으로 측정이 됐는데, 이에 대해 학습 스텝을 고정하는 대신 시간을 고정하여 측정한다.
약 32시간 동안 125K 스텝을 학습한 ALBERT-xxlarge 모델이 34시간 동안 400K 스텝 학습한 BERT-large 모델에 비해 성능이 전반적으로 높게 측정된다.
이는 ALBERT 모델의 파라미터가 적기 때문에 수렴이 빠르다는 것을 나타낸다.

Discussion

ALBERT-xxlarge는 BERT-large보다 파라미터 수는 적지만 구조가 크기 때문에 계산 비용이 더 많이 든다.
따라서 sparse attention이나 block attention과 같은 기법을 통해 학습 및 추론 속도를 높이는 것이 다음 과제다.
표현력을 향상시키기 위한 방향으로는 hard example mining이나 효율적인 언어 모델 학습 기법이 있다.
또한 SOP가 효과적인 학습 과제임은 입증되었지만, 현재의 자기지도 학습 손실만으로는 포착하지 못한 표현 차원이 더 있을 수 있다고 가정한다.
이를 통해 더 강력한 언어 표현을 얻을 가능성이 있다.

명사: 토론, 논문, 변론