[리뷰] GLaM(Generalist Language Model): Efficient Scaling of Language Models with Mixture-of-Experts
Model Size가 늘어날수록 dense 모델의 경우, 더 많은 에너지와 컴퓨팅 리소스가 필요하다. 최근 발표되는 1T 이상의 parameters을 가진 모델은 MoE와 Sparsity를 활용하여 에너지 사용 및 컴퓨팅 리소스의 사용을 줄여 학습 효율성을 증가시키고 있다.
구글은 올해 초 MoE기반의 Switch Transformer에 이어 21년 12 월 새로운 MoE기반 GLaM(Gerneralist Language Model)을 발표하였다. GLaM은 Large Sparsely Activated Network로 SOTA dense 모델에 비해 더 나은 컴퓨팅 효율과 few-shot task에서 더 우수한 경쟁력있는 결과를 달성하였다.
Introduction
GLaM의 가장 큰 버전은 1.2 Trillion 파라미터와 MoE 레이어 당 64개의 experts를 가지며 입력 batch에 있는 각 token은 95B parameters(8% of 1.2T)의 subnetwork만을 활성화한다. 희소적으로 활성화되는 모델 아키텍처와 모델 parallelism 알고리즘의 효율적인 구현덕분에, 학습기간 동안 에너지 총 에너지 소비는 GPT-3의 1/3수준이다.
MoE 모델이 NLP에 아직 일반적이지 않지만, 이 연구 결과로부터 MoE의 기본 버전조차도 대규모의 언어 이해에 굉장히 잘 동작함을 확인할 수 있었다. 또한 Sparsity가 에너지를 절감하면서 고품질 NLP 모델을 달성하기 위해 가장 약속된 방향 중 하나임을 입증한다.
아래 Table.2는 GLaM과 관련 모델의 주요 차이점을 비교하였다. n_params는 학습된 모델 parameters의 총 개수이며 n_act_params는 입력 토큰당 활성화된 모델 parameters의 개수이다.
Training Datasets
구글은 GLaM을 학습하기 위해 다양한 자연어 사용 사례를 대표하는 1.6조개의 토큰으로 구성된 고품질 데이터 세트를 만들었다. 원본 코퍼스로부터 고품질 웹 코퍼스를 생성하기 위해 구글은 자체 텍스트 품질 분류기를 개발하였다.
이 분류기는 선별된 텍스트 모음(Wikipedia, 책 및 일부 선택된 웹사이트)과 다른 웹페이지를 분류하도록 학습되었다. 이 분류기는 웹페이지의 콘텐츠 품질을 평가하며 그들의 점수에 따라 웹 페이지를 샘플링하기 위해 Pareto 분포를 사용한다. 이를 통해 분류기의 체계적인 편향을 방지하기 위해 일부 낮은 품질의 웹페이지를 포함할 수 있습니다
이 프로세스를 사용하여 필터링된 고품질 웹 페이지 하위 집합을 생성하고 이를 책, Wikipedia 페이지 및 기타 데이터 소스와 결합하여 Table 3와 같은 최종 GLaM 데이터 세트를 생성하였다. 또한 Adiwardana et al. (2020) [3]에서 사용하는 공개 도메인 소셜 미디어 대화의 데이터를 통합하였다. 이 데이터 필터링 단계가 모델의 품질에 큰 영향을 미친다.
Model Architecture
GLaM은 GShard MoE Transformer와 유사하게 Transformer Layer의 Feed-Forward 컴포넌트를 MoE Transformer로 대체하였다. GLaM의 아키텍처의 특징은 다음과 같다.
GLaM의 아키텍처 특징
- 각 MoE 레이어(bottom block)는 Transformer 레이어(upper block)가 반복적으로 교차 연결되어 있다. (e.g., MoE layer-Transformer layer-MoE layer-Transformer layer- …)
- 각 MoE 레이어는 더 많은 매개변수가 있지만 experts는 대부분 활성화되지 않는다. 주어진 입력 토큰에 대해서 제한된 experts의 subset만 사용되어 계산을 제한하여 더 많은 용량을 모델에 제공한다. (Subset의 크기는 2이다. 즉, 한번에 2개의 experts를 활성화한다.)
- 학습 중에 각 MoE 레이어의 게이팅 네트워크는 입력 시퀀스의 각 토큰에 대해 64개의 experts 중 최고의 experts만을 활성화하도록 학습된다.
- 이렇게 학습된 게이팅 네트워크는 추론 중에 2개의 최고의 experts를 동적으로 선택한다.
- E개의 experts를 가진 MoE 레이어는 O(E²)의 다양한 FFN의 조합을 제공한다. 토큰의 최종 학습된 표현은 선택된 experts의 출력의 weighted sum이다.
GLaM은 original Transformer 아키텍처에 몇가지 수정을 추가하였다.
- 기본 positional embedding을 per-layer relative positional bias로 교체함
- Non-MoE Transformer Feed-Forward sub-layer에서 첫번째 linear projection & activation 함수를 Gated Linear Unit로 교체함
- 기본 LayerNorm 대신 RMSNorm을 사용함
GLaM의 Model Hyper-parameters
GLaM은 MoE 언어모델의 scaling 효과를 이해하기 위해 Model 하이퍼파라미터를 변경한 여러개의 variant 버전이 존재한다. 또한 동일한 학습 데이터로 MoE와 dense 모델을 비교할 수 있도록 토큰당 유사한 유효 FLOP로 dense 모델셋을 학습한다.
GLaM의 학습방법
GSPMD[4]에서 서술한 것처럼, TPU 클러스터의 디바이스 네트워크의 2D 토폴로지를 활용한 2D sharding 알고리즘을 사용하여 GLaM 모델의 weights와 computation을 분할한다. GLaM 논문에서는 서로 다른 MoE 레이어에 대해 동일한 계산 그래프를 생성하기 위해 동일한 디바이스에 서로 다른 MoE 레이어에 걸쳐 동일한 인덱스를 가진 experts를 배치합니다.
결과적으로 MoE Transformer 아키텍처의 반복적인 모듈을 while_loop control flow 문으로 래핑하여 컴파일 시간을 줄일 수 있다. 실험은 고품질 모델을 얻기 위해 experts의 규모를 키워야 한다는 것을 보여준다. 그러므로 각 experts가 충분히 커지면 N/E 디바이스 집합에 각 expert를 할당해야 한다.
예를 들어, Expert dimension(E)와 Hidden dimension(H)를 따라 MoE 레이어에서 [E, M, H] 모양의 expert weight 텐서를 분할하고 Batch dimension(B)와 Model dimension(M)을 따라 [B, S, M] 모양으로 입력 activation 텐서를 분할한다. 이 2D sharding 알고리즘을 사용하면 큰 weight와 activation 텐서를 더 작은 조각으로 완전히 나눌 수 있으므로 모든 디바이스에서 데이터나 계산이 중복되지 않는다.
Experiment Setup
Training Setting
GLaM 모델은 130M 부터 1.2T 모델까지 다양한 규모로 학습되었으며 GLaM의 variants를 다음과 같은 표기한다.
- GLaM(8B/64E): Transformer 레이어:8B parameter을 가진 dense 모델, Non-transformer 레이어: 64 expert MoE 레이어인 GLaM 모델을 의미한다.
- GLaM(137B): 같은 데이터셋으로 학습된 dense 137B 파라미터 모델을 의미한다.
Model Training
모든 GLaM 모델에 대해서 동일한 학습 하이퍼파라미터를 사용한다.
- Max. sequence length: 1024 tokens
- Up to 1M tokens per batch
- Dropout rate: 0
- Optimizer: Adafactor (1st-moment decay beta1=0, 2nd-moment beta2=0.99 with 1- t^-0.8)
- Initial lr=0.01 (first 10K training steps)
- FP32 for weights/BF16 for activations
- 1024 Cloud TPU-V4 chips (for the largest GLaM 64B/64E)
결과
MoE와 Dense 모델간 비교
29개 benchmark에 대해서 GLaM(64B/64E)과 dense 모델(GPT3,175B)의 성능을 비교하면 GLaM이 zero-shot 태스크 중 80%, one-shot 태스크 중 90% 우수하다.
게다가 GLaM(64B/64E)은 총 1.2T parameters를 가진 반면, 추론시 token 당 97B parameters(1.2T 중 8%)의 subnetwork만이 활성화된다.
Figure.3은 Open-domain question answer task인 TriviaQA을 대상으로 prediction당 유사한 compute을 투입했을 때 MoE와 dense 모델의 성능을 비교한 결과이다. GLaM이 유사한 compute FLOPS에서 Dense 모델(GPT3)에 비해 더 좋은 정확도를 보인다.
또한 Table. 6에서 보는 것과 같이 TriviaQA Zero-shot/One-shot에서 이전 SOTA 모델에 비해 더 좋은 성능을 보인다. 이 결과는 GLaM(64B/64E)의 n_act-params가 GPT-3의 절반임에도 불구하고 GLaM의 추가용량이 성능 향상에 중요한 역할을 담당하고 있음을 제안합니다.
레퍼런스
[1] More Efficient In-Context Learning with GLaM
[2] GLaM: Efficient Scaling of Language Models with Mixture-of-Experts
[3] Towards a human-like open-domain chatbot.
[4] General and scalable parallelization for ml computation graph
