BERT의 등장 이후로 많은 연구들이 BERT를 기반으로 여러 nlp 태스크에서 좋은 결과를 만들어내고 있다. ALBERT 또한 그러한 논문 중 하나인데, 이름으로는 BERT모델을 가볍게 만드는 용도로 보이지만, 여러 트릭들을 통해 성능들도 많이 개선되었다. SOTA 모델 중, ALBERT를 기반으로 한 모델이 많이 있기에 이 논문을 선정하게 되었다.

논문의 내용이 전체적으로 BERT에서 세부적인 내용들을 개선한 것이기 때문에, 큰 골자는 BERT에서 크게 벗어나지 않는다. 따라서 BERT로부터 개선된 내용을 중심으로 간략하게 정리하도록 한다.

ALBERT에서 해결하고자하는 문제는 바로 파라매터의 개수를 감소시키는 것이다. 기존 BERT는 파라매터 수가 매우 많아서, 모델이 커지면 일반적인 GPU/TPU 메모리 한계로는 감당하기 어렵다. 또한 학습 시간도 그만큼 길어지게 된다.

논문에서 사용한 테크닉

ALBERT에서는 파라매터 감소와 학습속도를 올리면서, 성능을 향상시키기 위해 크게 다음과 같은 테크닉을 사용한다.

1. Factorized Embedding Parameterization

   큰 단어 임베딩 행렬을 작은 두 행렬로 분해하고, 히든레이어의 크기와 임베딩 크기를 각각 설정하도록 분리한다.

2. Cross-Layer Parameter Sharing

   파라매터가 네트워크 깊이의 증가에 따라 계속 커지는 것을 막기 위해서 이를 적용한다.

3. Sentence-Order Prediction (SOP)

   BERT에서 사용하던 NSP 대신에 새로 적용하는 학습이다. 문장 사이의 순서를 self-supervised loss 를 통해 학습시켜 문장간의 coherence(일관성)을 더 효율적으로 학습시키도록 한다.

기존의 관련 연구

Representation의 크기와 성능

지난 2년간 NLP 분야에서 가장 중요한 변화중 하나는 바로, 기존에 단어 임베딩부분만 pre-training을 진행했던 것이 전체 네트워크를 pre-training으로 학습시키고, 이후에 태스크에 대한 fine-tuning 단계를 진행하는 방식으로 바뀐 것이다.

여기에서 BERT를 포함한 기존의 연구들은 이 모델의 크기가 커질수록 성능이 좋아진다는 결과가 많았다. 그러나, 이 논문에서는 BERT가 단지 hidden size를 1024까지만 확인했다는 점을 지적하고 있다. SOTA 결과를 낸 모델들은 파라매터 개수가 너무 많아서 GPU/TPU 메모리 제한에 금방 도달하는 경우가 많다.

Cross-Layer Parameter Sharing

기존 transformer 구조에서도 레이어간에 파라매터를 공유하는 아이디어를 찾아볼 수 있다. 인코더의 output을 디코더의 multi-head attention 의 Key, Value 입력으로 넘기는 것이다. 하지만 transformer에서는 고전적인 인코더-디코더 구조 학습에만 이를 사용하고, pretraining/finetuning 단계에는 활용하지 않았다.

Universal Transformers라는 논문에서는 이 cross-layer parameter sharing이 기존 transformer보다 좋은 성능을 낸다는 결과를 얻었다.

또한, Deep Equilibrium Model(DQE) 논문에서는 이 DQE 모델이 input 임베딩과 output 임베딩이 특정 레이어에서 평형을 이룰 수 있다는 연구가 있다.

Sentence Ordering

위에서도 언급했듯이 ALBERT는 pre-training에서 SOP(Sentence Ordering Prediction)을 사용하는데, 이 것이 BERT에서 사용했던 NSP보다 더 어려운 pre-training이 되어, 특정 downstream task 에서는 더 유용하다고 한다.

ALBERT 모델

BERT 구조로부터 파생된 ALBERT에 대해서 알아보도록 한다.

모델 구조

ALBERT의 구조는 BERT와 마찬가지로 transformer의 인코더부분을 GELU¹와 함께 사용한다. 여기서 사용하는 각 크기에 대해 다음과 같이 정의한다.

• $E$ : 단어 임베딩 크기
• $L$ : 인코더 레이어 개수
• $H$ : 히든 레이어 크기
• 기타
  • Feed-forward 와 필터의 크기는 $4H$로 한다.
  • Attention head의 개수는 $\frac H{64}$로 한다.

여기까지는 BERT와 동일하며, 아래 BERT에 없는 3가지 주요 요소가 추가되었다.

Factorized embedding parameterization

BERT에서는 임베딩 크기 $E$ 와 히든 레이어 크기 $H$가 서로 동등하다. 즉, $E\equiv H$가 된다. 이러한 결정은 다음과 같은 관점에서 좋지 못하다.

모델링 관점에서 보면, WordPiece 임베딩은 context-independent 한 representation을 학습시키는게 목적인 반면, hidden-layer 임베딩은 context-dependent 한 representation을 학습시키는 것이다. Context의 길이 지표에 따른 실험 결과에 따르면 BERT와 같은 representation의 힘은 context-dependent representation과 같은 신호를 제공하는데 context를 사용하는데서 나온다고 한다. 따라서 중요한 것은 히든 레이어의 크기 $H$가 되기 때문에, $H\gg E$가 되어도 문제없다는 뜻이 된다.

활용 측면에서 보면, NLP는 보통 vocabulary size가 매우 크다. 만약에 $E\equiv H$라면, $H$를 증가시키면 임베딩 행렬의 크기 $V\times E$또한 같이 증가하게 된다. 이는 결국 쉽게 모델의 파라매터가 너무 많아지는 결과를 낳게 된다.

따라서 ALBERT에서는 임베딩 파라매터에 대해서 두개의 작은 행렬로 분해한다. 직접 one-hot 벡터를 $H$ 크기의 hidden space에 projection 하는 대신, 작은 차원의 $E$크기에 대해 projection시키고, 이를 다시 hidden space 에 projection 시키는 방식이다. 이 방식은 결국 임베딩 파라매터를 $O(V\times H)$에서 $O(V\times E+E\times H)$로 감소시킬 수 있고, $H\gg E$일 경우에 이는 매우 효과적인 결과가 된다. 논문에서는 모든 word piece에 대해서 동일한 $E$값을 사용했다.

Cross-layer parameter sharing

ALBERT에서는 파라매터 효율을 위해 cross-layer parameter sharing을 사용했다. 파라매터 공유에는 Feed-forward network 파라매터만 레이어에서 공유하거나 Attention 파라매터만 공유하는 등 일부 파라매터들만 공유하는 방법이 있다. 하지만 ALBERT에서는 기본적으로 모든 파라매터를 레이어에 공유한다. 논문의 뒷부분에서 이러한 공유방법들에 대한 실험결과가 있다.

Universal Transformer와 Deep Equilibrium Models에서 비슷한 방법이 사용되었는데, 차이점은 다음과 같다.

• UT는 기본적인 transformer구조를 사용했다.
• DQE는 특정 레이어의 입력과 출력 임베딩이 equilibrium point에 도달하는 것이 목적이다.

BERT와 ALBERT의 레이어에 따른 L2거리값과 Cosine 유사도 값은 다음과 같다. 이 측정은 ALBERT의 임베딩이 수렴하지 않고 진동하고 있는 것을 보여준다.

이 그래프는 각 레이어의 입력과 출력 임베딩의 L2거리와 코사인 유사도를 나타낸다. 이를 통해 ALBERT에서 레이어간의 전환이 BERT보다 훨씬 자연스러운 것을 알 수 있다. 이 결과는 가중치 공유가 네트워크 파라매터의 안정화에 효과가 있다는 것을 나타낸다. 초반 레이어에서 급격한 꺾임이 있지만, 이후로는 24번째 레이어까지도 0으로 수렴하지 않고 계속 값을 유지한다. 이는 ALBERT 파라매터의 solution space가 DQE와는 확실히 다르다는 것을 보여준다.

Inter-sentence coherence loss

BERT에서는 MLM(Masked Language Modeling) loss와 더불어 NSP(Next-Sentence Prediction) loss를 사용했다. NSP는 두개의 세그먼트가 원래의 문장과 연속적인지를 예측하는 일종의 binary classification loss이라고 볼 수 있으며, 참인 경우와 거짓인 경우가 동일한 확률로 들어있다. 하지만 BERT의 후속 논문들에서 NSP의 효과에 대해 신뢰하기 힘들고, 이 부분을 제거하면서 몇몇 downstream task에서 더 좋은 결과를 얻었다는 결과가 있었다. ²

이 논문에서는 이러한 NSP의 가장 큰 문제점을 MLM과 비교했을 때 하나의 태스크로서의 어려움이 너무 없기 때문인 것으로 추정했다. 즉, 더 정확히 표현하면, NSP는 topic prediction과 coherence prediction을 단일 태스크에 동시에 포함하고 있다. 하지만 topic prediction은 coherence prediction에 대해 학습이 매우 쉽고, MLM과 학습하는 대상이 겹치는 것이다.

ALBERT에서는 inter-sentence 모델링도 언어 이해에서 중요하다고 생각했기 때문에, 위와 같이 태스크를 제거하기보다는 유지하는 방향을 제시했다. 대신에 위와 같은 단점을 보완하기 위해 coherence를 학습시키는데 중점을 두었고, 여기서 등장한 것이 바로 SOP(Sentence-Order Prediction) loss다.

SOP는 topic prediction을 제외하고 inter-sentence coherence 모델링에 초점을 맞춘 loss다. SOP의 postive 경우는 BERT와 마찬가지로 같은 문서에서 두개의 연속적인 세그먼트이고, negative 경우는 마찬가지로 두개의 연속적인 세그먼트이지만, 순서가 뒤바뀐 경우를 나타낸다.

이 부분이 모델을 담화수준의 coherence에 대해서 더 세분화하여 학습하도록 만든다. 이에 대한 실험 결과로, NSP로는 SOP 태스크를 전부 풀지는 못하지만, 반대로 SOP로는 NSP 태스크를 전부 풀 수 있다. 결국 ALBERT 모델이 여러문장에 관련한 downstream task의 성능을 일관되게 향상시키는 것이다.

모델 설정 비교

BERT와 ALBERT 모델의 하이퍼파라매터 차이는 다음과 같다.

모델	파라매터 개수	레이어	히든 크기	임베딩 크기	파라매터 공유
BERT-base	108M	12	768	768	X
BERT-large	334M	24	1024	1024	X
ALBERT-base	12M	12	768	128	O
ALBERT-large	18M	24	1024	128	O
ALBERT-xlarge	60M	24	2048	128	O
ALBERT-xxlarge	235M	12	4096	128	O

여기서 ALBERT-xxlarge의 레이어가 12개인 것은 24개로 설정했을 때의 결과와 별 차이가 없었기 때문이다.

위 표를 통해 BERT와 비교했을 때 다음과 같은 특징들이 있는 것을 알 수 있다.

• ALBERT-large는 BERT-large보다 18배 더 적은 파라매터를 사용한다.
• 히든 크기가 4096인 ALBERT-xxlarge도 BERT-large의 파라매터의 70%밖에 되지 않는다.

이와 같이 ALBERT가 기존의 BERT보다 파라매터 개수를 획기적으로 줄인 것을 알 수 있다.

실험

실험 설정

• 의미있는 결과 비교를 위해 BERT에서 사용했던 BookCorpus (2015) 와 English Wikipedia (2019) 데이터를 pretrain에 사용한다.

• 입력은 두개의 세그먼트의 경우 [CLS]x1[SEP]x2[SEP] 형식으로 구분한다.

• 최대 길이를 512로 설정하고, 입력 시퀀스를 10%확률로 512보다 짧게 만든다.

• BERT와 마찬가지로 vocab size는 30000으로 하고, SentencePiece를 통해 토큰화시킨다.

• $n$-gram 마스킹을 사용해서 MLM 타겟을 만든다. 여기서 길이 $n$은 다음과 같은 확률로 랜덤하게 선택된다. 여기서 최대 길이 $N=3$이다.

  \[p(n)=\frac{1/n}{\sum_{k=1}^N{1/k}}\]

• 배치 사이즈는 4096으로, 큰 배치 학습을 위해 LAMB optimizer를 사용한다. 여기서 learning rate는 0.00176이다.

• 특정 경우를 제외하고 모든 모델을 125000 step까지 학습시킨다.

벤치마크 측정

내부 평가

학습 진행을 확인하기 위해 SQuAD와 RACE를 개발 세트로 설정해서 학습을 진행했고, MLM과 sentence classification task들에 대한 정확도를 측정했다. 이는 단순히 모델이 제대로 수렴하는지를 측정하기 위함이고, 성능 측정이 목적이 아니다.

Downstream 평가

보편적인 벤치마크인 GLUE와, SQuAD, RACE 데이터셋을 사용하여 평가를 진행했다. GLUE 데이터의 경우 편차가 커서 5번의 실행 결과의 중간값을 사용했다.

BERT와 ALBERT의 최종 결과 비교

모델	파라매터	SQuAD1.1	SQuAD2.0	MNLI	SST-2	RACE	평균	속도
BERT-base	108M	90.4/83.2	80.4/77.6	84.5	92.8	68.2	82.3	4.7x
BERT-large	334M	92.2/85.5	85.0/82.2	86.6	93.0	73.9	85.2	1.0
ALBERT-base	12M	89.3/82.3	80.0/77.1	81.6	90.3	64.0	80.1	5.6x
ALBERT-large	18M	90.6/83.9	82.3/79.4	83.5	91.7	68.5	82.4	1.7x
ALBERT-xlarge	60M	92.5/86.1	86.1/83.1	86.4	92.4	74.8	85.5	0.6x
ALBERT-xxlarge	235M	94.1/88.3	88.1/85.1	88.0	95.2	82.3	88.7	0.3x

SQuAD 점수의 경우 F1/EM 점수 형태로 표기되었다. 또한, 속도는 BERT-large 를 기준으로 다른 모델들의 학습 속도를 나타낸다.

60M의 파라매터를 사용한 ALBERT-xlarge만 해도 BERT-large 와 거의 동등한 결과를 얻었고, ALBERT-xxlarge 는 모든 결과에서 기존 BERT보다 향상된 점수를 얻었다.

하지만 적은 파라매터로 훨씬 좋은 결과를 낸 반면에, 계산은 훨씬 많아지고, 속도는 BERT-large에 비해서 많이 줄어들었다는 점이 ALBERT 모델의 가장 아쉬운 부분이다.

논문에는 기타 여러 학습 설정에 대한 벤치마크 측정 결과들이 많다. 하지만 ALBERT 모델에 대한 이해 측면에 대해서는 중요하지 않아 나머지는 생략하도록 한다. 결과적으로 ALBERT가 대부분의 태스크에서 SOTA를 기록하였고, 이를 바탕으로 한 후속 연구들도 많이 진행되고 있다.