GLaM: Efficient Scaling of Language Models with Mixture-of-Experts

In plain words

Most large language models are dense: every weight takes part in every prediction, so a bigger model costs more to run and to train. This paper asks whether a model can hold far more knowledge without paying for all of it on each word. Its answer is a family of models called GLaM, which swap some of the feed-forward blocks for a bank of specialist sub-networks, called experts. A small gating step picks a couple of experts for each token, and only those run.

The largest GLaM has over a trillion weights, yet any single token uses only a small fraction of them. The authors trained it on a large, quality-filtered text mix and compared it with GPT-3 on 29 language tasks, where it does better on average without any task-specific fine-tuning. The authors also report that it cost much less energy to train than GPT-3.

TL;DR

GLaM replaces every other feed-forward sub-layer of a decoder-only Transformer with a mixture of 64 experts, and routes each token to its top two §4 tab. 4. The 64B/64E model has 1.2T total parameters and 96.6B active per token tab. 4. It beats GPT-3 (175B) on the average of 29 zero-, one- and few-shot tasks, with about half of GPT-3's inference FLOPs per token and a third of its training energy tab. 1. All numbers are author-reported.

Problem

Scaling a dense model raises the cost of every prediction in proportion to its size §5.3. Sparse mixture-of-experts models can grow parameter count without growing per-token compute, but earlier sparse models were mostly evaluated on fine-tuning benchmarks, not on in-context (zero-, one- and few-shot) learning, which is the setting GPT-3 made standard §2 §1. The paper asks whether a sparse decoder-only model can match or beat a dense model of similar per-token compute, and what it costs to train §1.

Method

Architecture. GLaM is a decoder-only Transformer. Following the GShard MoE Transformer, the feed-forward sub-layer of every other Transformer block is replaced by an MoE layer §4 fig. 2. Each MoE layer holds E independent feed-forward networks (the experts) and a gating network §4. A softmax over the experts gives each token a distribution over them §4. Each token is sent to the two experts with the highest scores, and its output is the weighted combination of those two experts' outputs §4. Because the two chosen experts can differ from token to token, one MoE layer offers O(E²) pairs of feed-forward networks instead of one §4.

Other changes to the Transformer. The standard positional embedding is replaced by a per-layer relative positional bias from a 2019 work the paper cites §4. In the dense (non-MoE) feed-forward sub-layers, the first linear projection and activation are replaced by a Gated Linear Unit whose output passes through a GELU §4.

Training data. The paper builds a corpus of 1.6T tokens §3. Web pages are filtered with a feature-hashed linear classifier trained to separate curated text (Wikipedia, books, some selected websites) from other web pages §3. The classifier score is turned into a sampling probability through a Pareto distribution, so that some low-scoring pages still get in §3. The filtered web data is then mixed with books, Wikipedia, forums, news and public-domain conversations, using weights chosen from the performance of each source in a smaller model §3 tab. 3.

Training objective. Standard next-token cross-entropy, plus the GShard load-balancing auxiliary loss with coefficient 0.01 §5.2.

Parallelism. The weights and computation of the large models are partitioned with the 2D sharding scheme of GSPMD over TPU-v4 devices §4 app. C. Experts with the same index in different MoE layers are placed on the same device, so that the repeated layers share one computation graph app. C.

# MoE feed-forward sub-layer with top-2 routing  [§4]
# Input and output: tokens (B, S, M)

def moe_layer(tokens, gate, experts):
    # 1. Softmax over the E experts gives a distribution per token  [§4]
    probs = softmax(gate(tokens))                 # (B, S, E)

    # 2. Keep the two highest-probability experts per token  [§4]
    weights, chosen = top_k(probs, k=2)           # (B, S, 2) each

    # 3. Weighted sum of the two chosen experts' outputs  [§4]
    #    (whether the two weights are renormalised is not stated)
    output = zeros_like(tokens)
    for rank in range(2):
        output += weights[..., rank] * expert_apply(
            experts, chosen[..., rank], tokens)
    return output

# Training loss  [§5.2]

def training_loss(logits, targets, gate_stats):
    # 1. Standard cross-entropy on next tokens  [§5.2]
    lm_loss = cross_entropy(logits, targets)

    # 2. GShard-style load-balancing term, weight 0.01  [§5.2]
    #    (the exact form is defined in the GShard paper, not here)
    balance_loss = load_balance_loss(gate_stats)
    return lm_loss + 0.01 * balance_loss

# Learning-rate schedule  [§5.2]

def learning_rate(step, base=0.01, warm_steps=10_000):
    # 1. Constant initial rate for the first 10K steps  [§5.2]
    if step < warm_steps:
        return base
    # 2. Then inverse-square-root decay; the anchor at 10K is not stated
    return base * (warm_steps / step) ** 0.5

Evaluation protocol. Zero-shot: each development example is scored directly §5.3. One- and few-shot: one or a few random examples from the task's training set are concatenated with two newlines before the test example §5.3. Few-shot runs tune the number of shots up to the count GPT-3 used for that task tab. 14. Generative tasks use beam search with width 4 §5.3. For multiple choice, each option is scored by its log-likelihood normalised by token length, except on ReCoRD and COPA, where the unnormalised score is used §5.3.

Hyperparameters

Shared by all GLaM sizes §5.2:

Architecture per model size tab. 4. Naming: GLaM (B/E) uses the dense base size B, with every other layer's feed-forward replaced by an E-expert MoE layer §5.1. GLaM (B) is the dense model of that size §5.1.

Symbols as defined in the paper: L layers, M model dimension, H feed-forward hidden dimension §5.1 tab. 4. The paper does not say whether H for MoE models is the hidden size of each expert (see Pitfalls).

Results

Author-reported. Setup: 29 public benchmarks (8 generative, 21 understanding), grouped into 7 categories §5.3 §A; mainly development-set scores §6. GPT-3 (175B) is the main baseline; Gopher (280B) and Megatron-NLG (530B) are shown for reference in the appendix tab. 11.

Headline comparison, GLaM 64B/64E against GPT-3 175B tab. 1:

Average NLG and NLU scores per setting, from the appendix main table tab. 11. Columns are GPT-3 (175B) and GLaM (64B/64E) for each setting:

Open-domain TriviaQA tab. 5: GLaM one-shot scores 75.0 on the test server and 75.8 on dev. Reference rows in the same table are KG-FiD (large, fine-tuned, test) at 69.8, Switch-C (fine-tuned, dev) at 47.5, GPT-3 one-shot (dev) at 68.0 and GPT-3 64-shot (test) at 71.2. The paper states that the test-server result beats the few-shot GPT-3 test-server score by 5.3% §6.1.

Scaling and data results (no tables; figures only, described in words):

• At similar effective FLOPs per token, MoE models score better than dense models on NLG and, at larger scales, on NLU §6.3 fig. 3.
• Training on filtered web data (143B tokens) beats training on unfiltered web data (about 7T tokens) for the 1.7B/64E model, with a larger gap on NLG than on NLU §6.2 fig. 3.
• Adding experts at a fixed compute budget per prediction improves average scores, as the expert count grows from 1 to 256 on the 1.7B base app. B fig. 6.

Efficiency §6.4 fig. 4 app. F:

• The 64B/64E model trained on 280B tokens beats GPT-3 trained on 300B tokens on 4 of the 6 learning settings and matches it on the remaining one §6.4.
• Reaching GPT-3-level scores took 574 hours on 1,024 TPU-v4 chips and 213 MWh §6.4 app. F.
• A 600B-token run used 456 MWh tab. 1 app. F.
• The paper estimates 18.7 t CO2e for the 280B-token run, against 552 t for GPT-3 app. F.

Bias and social dimensions §7.2 §1: the paper reports that GLaM closes the gap between stereotypical and anti-stereotypical examples on WinoGender. It gives no WinoGender figures in the main text.

Paper vs code

No official code.

No link to code appears in the paper's LaTeX source, and the arXiv record has no comments. A GitHub repository search for "glam mixture-of-experts", "GLaM language model" and "glam" restricted to the google-research and google organisations returned no matching repositories, so no code is cited. The Acruxos/GLAME result is unrelated (knowledge-graph model editing) and is not used.

Pitfalls

• Average scores do not reproduce from the appendix. The paper defines the overall average as the mean over all datasets §5.3. Weighting the appendix NLG and NLU averages by task count (8 and 21) gives about 62.1 for GLaM zero-shot (paper: 62.7) and about 57.1 for GPT-3 (paper: 56.9), by our arithmetic. The paper does not say how the headline averages were computed. Use per-task numbers in tab. 11–tab. 14 when precision matters.
• TriviaQA percentages mix splits. The "8.6%" over KG-FiD compares the GLaM development score (75.8) to KG-FiD's test score (69.8) §6.1 tab. 5. The test-server score (75.0) gives +7.4% over the same reference, by our arithmetic. The 5.3% over GPT-3 compares test-server scores §6.1 tab. 5.
• Two different training runs back the energy claim. The 456 MWh figure in tab. 1 belongs to the 600B-token run. The 213 MWh figure belongs to the 280B-token run that reaches GPT-3-level scores §6.4 app. F. The GPT-3 figure (1287 MWh) is cited to another paper app. F. The "one third" headline compares the 600B-token run (456 MWh) with that GPT-3 figure.
• Routing details are missing. The paper does not state the form of the gating network (linear layer, bias, input normalisation), whether the two selected weights are renormalised, the expert capacity factor, how tokens beyond capacity are handled, or whether a router z-loss is used §4 §5.2. Implementers must choose these, and the results depend on them.
• Learning-rate schedule has gaps. The decay constant after 10K steps is not stated §5.2. Warmup is not mentioned. The total number of training steps and the exact batch size in examples are not stated; only "up to 1M tokens per batch" is given §5.2.
• Expert hidden size is ambiguous. Table 4 lists H = 32,768 for 64B/64E and H = 65,536 for the dense 137B tab. 4. The paper does not say whether the MoE value is per expert or summed across experts.
• Model names do not give active size. "64B" names the dense base size, not the active parameters. The active count of 96.6B is 8% of the 1.2T total, as the paper states §1 §5.1. Do not read "64B" as a parameter count.
• Dataset construction is not reproducible from the text. The classifier's features, training data sizes, threshold and Pareto parameter are not stated §3. The mixture weights come from unspecified small-model experiments §3. The number of passes over each source during training is not stated, so the mixture in tab. 3 may not equal what the model saw.
• Data contamination is reported only as counts. The paper gives dirty and total example counts per dataset, for the subset also used by GPT-3, using an n-gram overlap test whose n varies by dataset app. D tab. 6. It gives no single contamination rate, and the text says the overall level roughly matches earlier work §3 app. D.
• Evaluation protocols differ across tasks. Beam width 4 for generation §5.3, the unnormalised scores for ReCoRD and COPA §5.3, F1 over answer options for MultiRC §5.3, and per-task shot counts for few-shot tab. 14. Whether the GPT-3 numbers were re-run or copied from the GPT-3 paper is not stated.
• Development and test scores are mixed. Most of the table uses development scores §6. Test-server results appear only for some tasks tab. 5. Compare like with like.
• Serving cost is discussed but not measured. The paper says sparse models need more devices and raise serving cost when traffic is low §8. The serving latency tables appear only in commented-out source, so the final text reports no serving measurements.
• Figure and table numbers are from the LaTeX source. Table numbers are counted in source order, including the appendix. Check them against the compiled PDF before citing externally.

Lineage

• Builds on: the sparsely gated mixture-of-experts layer §2 §4, and the GShard MoE Transformer, which GLaM follows for layer placement and the auxiliary loss §4 §5.2. Neither has a card in this wiki.
• Compared against: GPT-3 (175B), the main baseline §1 tab. 1; Gopher (280B) and Megatron-NLG (530B), dense models shown for reference tab. 11; Switch-C, a sparse encoder-decoder model, in the TriviaQA comparison §2 tab. 5. None of these has a card in this wiki.
• Superseded by: not known from content already in this wiki.