[Paper Review] GLaM: Efficient Scaling of Language Models with Mixture-of-Experts

Nan Du, Yanping Huang+25 more|arXiv (Cornell University)|Dec 13, 2021

Topic Modeling168 citations

TL;DR

GLaM uses sparsely activated mixture-of-experts to scale language models up to 1.2 trillion parameters, achieving better zero/one/few-shot performance than GPT-3 on 29 tasks while using about 1/3 the training energy and half the FLOPs per token at inference.

ABSTRACT

Scaling language models with more data, compute and parameters has driven significant progress in natural language processing. For example, thanks to scaling, GPT-3 was able to achieve strong results on in-context learning tasks. However, training these large dense models requires significant amounts of computing resources. In this paper, we propose and develop a family of language models named GLaM (Generalist Language Model), which uses a sparsely activated mixture-of-experts architecture to scale the model capacity while also incurring substantially less training cost compared to dense variants. The largest GLaM has 1.2 trillion parameters, which is approximately 7x larger than GPT-3. It consumes only 1/3 of the energy used to train GPT-3 and requires half of the computation flops for inference, while still achieving better overall zero-shot and one-shot performance across 29 NLP tasks.

Motivation & Objective

• Motivate scalable language modeling with reduced training cost and energy usage compared to dense architectures.
• Explore sparsely activated MoE layers to increase model capacity without proportional compute.
• Evaluate zero-shot, one-shot, and few-shot performance across a broad NLU/NLG benchmark suite.
• Assess data quality, scaling behavior, and efficiency trade-offs of MoE versus dense models.

Proposed method

• Replace every other Transformer layer's feed-forward network with a Mixture-of-Experts (MoE) layer containing multiple experts.
• Use a gating network to select two best experts per token, yielding approximately O(E^2) combinatorial capacity with only two active experts per token.
• Adopt per-layer relative positional bias and GaLU/GeLU-like activations in non-MoE layers to improve efficiency and performance.
• Train on a 1.6 trillion-token high-quality dataset using a quality classifier to curate the web corpus and mix with Wikipedia, conversations, forums, books, and news with specified mixture weights.
• Optimize with Adafactor, apply MoE auxiliary load-balancing loss, and employ 2D tensor sharding for large models.

Experimental results

Research questions

• RQ1 Can sparsely activated MoE language models scale to trillion-parameter sizes while reducing training cost and energy compared to dense models?
• RQ2 Do MoE-based decoders provide superior zero/one/few-shot performance on a broad set of NLP tasks relative to dense equivalents at similar effective FLOPs?
• RQ3 How does data quality affect performance in large MoE language models?
• RQ4 What are the scalability and efficiency trade-offs when increasing the number of experts in MoE layers?
• RQ5 How do MoE models fare on knowledge-oriented tasks and open-domain QA compared to dense models?

Key findings

• GLaM with 1.2T parameters and 64 experts per MoE layer activates about 96.6B parameters per token, yet achieves competitive or better performance than GPT-3 across 29 benchmarks in zero/one/few-shot settings.
• GLaM requires only about one third the training energy of GPT-3 and about half the FLOPs per token for inference compared to GPT-3.
• GLaM (64B/64E) outperforms GPT-3 on average across 7 benchmark categories, with advantage in 6 of 7 categories.
• Data quality positively impacts performance; training on filtered, higher-quality data yields better downstream results than training on the unfiltered, larger dataset.
• MoE models demonstrate data efficiency, achieving stronger performance with less data at comparable effective FLOPs, and showing scaling benefits over dense counterparts at larger sizes.