BASE Layers: Simplifying Training of Large, Sparse Models

In plain words

Large language models can be made bigger without making every step more expensive, by splitting one layer into many expert sub-networks and sending each token to just one of them. The catch is that experts easily end up unevenly loaded: a few get almost all the tokens while others sit idle. Existing fixes add extra penalty terms that have to be tuned.

BASE treats the choice of expert as a fair matching problem. During training, every expert is given the same number of tokens, and an optimal-assignment solver picks which token goes where, based on how well each token fits each expert. Each expert's output is blended into the token's representation, so experts gradually specialise in particular kinds of tokens. The result is one extra layer that can be dropped into a language model without new loss terms.

TL;DR

• BASE is a sparse expert layer. Each token is routed to exactly one expert, and there is one expert per worker (GPU).
• During training the routing is a balanced linear assignment (each expert gets T/E tokens), solved with an auction algorithm. There is no balancing loss and no capacity factor.
• At test time, each token simply goes to its highest-scoring expert.
• The authors report that it matches data-parallel training on 8 GPUs, gains at larger compute budgets, and is comparable to Sparsely Gated MoE and better than Switch on validation perplexity. No perplexity numbers are given in the text.
• Pitfalls: the paper's mixing equation differs from the released code, the solver's pseudocode differs from its C++ implementation, and each BASE layer adds four all-to-all communication steps.

Problem

Sparse expert layers add capacity at roughly constant compute per token, but routing tokens to experts is a hard discrete problem. Routing by learned gates needs auxiliary balancing losses with weights that must be tuned, and even with them balance is not guaranteed. Stable training also needs hard per-expert capacity limits, after which surplus tokens are ignored §2.2. Each sparse layer costs four routing communication steps during training (before and after the expert, in forward and backward), and workers can idle while waiting §2.2. The authors' stated goal is an expert layer that balances load without new hyperparameters or losses §1.

Method

BASE has E experts, one per worker, each with a position-wise network f_e and an expert centroid w_e in R^D §3.1 §4.1. A token h_t that is assigned to expert a_t leaves the layer as sigma(h_t · w_{a_t}) f_{a_t}(h_t) + h_t eq. 1 §3.1. The gate is high when the expert helps the token, and gradient descent then pulls h · w towards the tokens that expert serves §3.1.

Training assignment. The assignment maximises the sum of token-expert affinities h_t · w_{a_t} under the constraint that each expert receives T/E tokens eq. 2 §3.2.1. The paper solves it with the auction algorithm rather than the Hungarian method, because the auction is easier to parallelise on GPUs §3.2.1.

Sharding. Each worker solves the assignment over its own T tokens. It then sends T/E tokens to each other worker with one all-to-all. Globally this decomposes the ET-token problem into E smaller problems §3.2.1.

Shuffling. Tokens within a sequence are correlated, so before assignment each worker first sends an equal number of its tokens to every other worker at random. This breaks the sequence correlation so that experts can specialise across domains §3.2.1.

Forward pass. The full training procedure is given in the implementation listing fig. 2. The pseudocode below follows it. It uses the residual form of eq. 1 that the listing shows, which is not what the released code does (see Paper vs code).

# BASE layer, training forward  [fig. 2] [§3.2.1]
# Input:  features (T, D) on this worker; centroids (E, D); this expert_id
# Output: features (T, D) in the original order

def base_layer_train(features, centroids, expert_id, expert_network):
    # Randomly send tokens to all workers to break sequence correlation
    shuffle_sort = random_permutation(T)
    shuffled = all_to_all(features[shuffle_sort])        # (T, D)

    # Score each token against each expert centroid  [fig. 2]
    affinities = shuffled @ centroids.T                  # (T, E)

    # Balanced assignment: each expert gets T/E tokens  [eq. 2] [fig. 9]
    sort_by_expert = balanced_assignment(affinities)     # (T,)

    # Send each token to the worker that hosts its expert  [fig. 2]
    routed = all_to_all(shuffled[sort_by_expert])        # (T, D)

    # Blend in this worker's expert, gated per token  [eq. 1] [fig. 2]
    gate = sigmoid(routed @ centroids[expert_id])        # (T,)
    routed = gate[:, None] * expert_network(routed) + routed

    # Undo the expert sort and send results back  [fig. 2]
    shuffled = all_to_all(routed)[inverse_sort(sort_by_expert)]

    # Undo the random shuffle and return to the original worker  [fig. 2]
    return all_to_all(shuffled)[inverse_sort(shuffle_sort)]

Solver. The paper's listing of the auction algorithm is given in the appendix (it is typeset as a figure) fig. 9. It is pseudocode for the iterative bidding below. eps is not defined in the paper (see Pitfalls).

# Balanced assignment by auction  [fig. 9]
# Input:  scores (E, T): affinity of each token for each expert
# Output: for each expert, the indices of its T/E tokens

def balanced_assignment(scores, max_iterations=100):
    # Start from raw affinities; no token has a price yet
    adjusted = scores.clone()
    price = zeros(1, T)
    is_won = zeros(T, dtype=bool)
    round_no = 0
    while not all(is_won):
        # Each expert looks at its T/E + 1 best tokens  [fig. 9]
        best_vals, best_ids = topk(adjusted, k=T / E + 1, dim=1)
        # Bid: margin over the (T/E + 1)-th token, plus eps  [fig. 9]
        raise_by = best_vals[:, :-1] - best_vals[:, -1:] + eps
        bids = scatter(zeros(E, T), best_ids[:, :-1], raise_by)
        # Last round's holders keep their token with a minimal bid  [fig. 9]
        if 0 < round_no < max_iterations:
            bids[holders, is_won] = eps
        # The highest bid per token wins it  [fig. 9]
        top_bid, holders = bids.max(dim=0)
        is_won = top_bid > 0
        holders = holders[is_won]
        # Raise the price of contested tokens  [fig. 9]
        price += top_bid
        adjusted = scores - price
        # Winners stay in their expert's list  [fig. 9]
        if round_no < max_iterations:
            adjusted[holders, is_won] = infinity
        else:
            adjusted[holders, is_won] = scores[holders, is_won]
        round_no += 1
    return best_ids[:, :-1].reshape(-1)

Test time. Balancing would leak information about future tokens, so at inference each token goes to its single best expert §3.2.2. Memory is lower at inference because no gradients or optimizer states are kept, and the paper notes that the trained model still ends up reasonably balanced §3.2.2 §5.1.

Gradient clipping. Gradients are clipped by global l2 norm. The norm is computed only over shared parameters, so no extra communication is needed, but all gradients, including the experts', are rescaled by the same factor §3.3.

Placement and shape. A BASE layer is a stack of feedforward blocks, each with layer norm, a projection to 4× the input width, ReLU, a projection back, and a residual connection §4.1. In BASE×N, N layers are inserted after the layers with indices floor(L/(N+1)) to floor(NL/(N+1)), and each has about floor(10/N) sublayers, so that the parameter count stays roughly fixed §4.1. The default single BASE layer is inserted after the L/2-th layer §4.4.

Hyperparameters

Results

Author-reported. Setup: about 100B tokens, 128 × 32 GB V100 GPUs for about 2.5 days in the main comparison, validation perplexity against wall-clock time §4.1 §4.2. The text gives no perplexity values; the comparisons are plotted curves fig. 3 fig. 4.

The efficiency table caption also says a 43B-parameter model trains at 90% of the speed of the 1.5B data-parallel baseline tab. 2. Balance of routing (top-1 BASE vs Sparsely Gated MoE and Switch, across 128 experts) is plotted in fig. 8. A perfectly uniform split over 128 experts would be 0.78% of tokens (derived, not labelled in the figure). BASE is less balanced than the baselines that use a balancing loss, but more balanced than the second expert of MoE §5.1 fig. 8.

Paper vs code

The official code is the fairseq repository linked in the abstract, at the commit above. The BASE layer is fairseq/modules/base_layer.py; the solver is fairseq/clib/libbase/balanced_assignment.cpp.

• Mixing equation. Paper: sigma(h·w) f(h) + h, residual eq. 1 and fig. 2. Code: alpha * expert(x) + (1 - alpha) * x, a convex blend without the identity term base_layer.py. Each expert block is residual internally base_layer.py.
• Training versus test switch. The paper switches on training versus test §3.2. The code uses whether the input requires gradients, not model.train() base_layer.py base_layer.py. Evaluation under no_grad therefore uses greedy top-1 routing.
• Shuffle. Fig. 2 always shuffles. The code shuffles only when training base_layer.py base_layer.py.
• Gradient into routing. The affinities used for assignment are computed under no_grad base_layer.py. Centroids only receive gradient through the sigmoid gate base_layer.py. The paper does not say whether the assignment is differentiable; this matches its description of learning through the gate §3.1.
• Retained winners. Fig. 9 sets winners' value to infinity before the cap. The C++ code sets it to the largest score in the whole matrix, computed once balanced_assignment.cpp balanced_assignment.cpp.
• After the cap. The fig. 9 caption says the solver switches to greedy after max_iterations. The C++ code does not: it changes the update rule and keeps looping until every token has a bid balanced_assignment.cpp balanced_assignment.cpp.
• Epsilon. The paper's listing uses an undefined eps. The code uses max((max − min)/50, 1e-4) balanced_assignment.cpp.
• Matrix orientation. Fig. 2 passes a T × E matrix. The fig. 9 listing takes (workers, jobs), i.e. E × T. The C++ code transposes its input balanced_assignment.cpp.
• Sublayers. The paper uses 10 sublayers in the standard layer §4.1. The code default is base_sublayers = 1 transformer_config.py; the 10 comes from the experiment config, which is not in the repository.
• Expert width. The paper specifies 4× expansion §4.1. The code takes the width from decoder_ffn_embed_dim base_layer.py base_layer.py.
• Initialisation. Centroids are orthogonal with gain 0.1 base_layer.py. The second feedforward matrix is zero-initialised base_layer.py. Neither is in the paper.
• Non-finite scores. The code replaces NaN and infinite affinities with the minimum finite score before solving base_layer.py. The paper says nothing about this.
• Clipping. Matches §3.3: the norm uses only non-expert gradients and the scale applies to all utils.py. The config default is 0 (off) configs.py; the 0.1 threshold is not in the repository.
• Expert gradient sync. Expert parameters are skipped by the legacy data-parallel sync legacy_distributed_data_parallel.py.
• Greedy test routing. Matches §3.2.2: top-1 by affinity base_layer.py.
• Placement. Matches: BASE layers are inserted at index ((i+1)·L) // (N+1) transformer_decoder.py.
• Training recipe. Warm-up, Adam settings and learning rates are not in the repository; no BASE example config was found in the clone.

Pitfalls

• Pick one mixing rule. Eq. 1 and fig. 2 use the residual form; the released code uses a convex blend. Choose deliberately and test it against a reference run. The code is what the authors released, so it is the safer default for matching results, but the paper and code do not agree.
• "No new hyperparameters" is not literally true. The paper says the method adds none §1. The implementation exposes base_layers, base_sublayers and base_shuffle transformer_config.py, and the solver has an epsilon and an iteration cap balanced_assignment.cpp. The centroid initialisation gain of 0.1 is also a setting base_layer.py.
• Solver epsilon is numerically sensitive. Eps is added to every bid, so it sets the minimum price step and the margin that keeps winners in their top list. The code scales it by the score range, so its size changes with the affinities balanced_assignment.cpp balanced_assignment.cpp. Use the same rule, or tune it.
• Solver termination. The paper promises a greedy fallback after 100 iterations fig. 9. The code does not use one. Make sure your implementation terminates when every token has a bid, and check the behaviour on hard inputs.
• Divisibility. The assignment assumes T is divisible by E eq. 2. The code uses integer division balanced_assignment.cpp and has no check base_layer.py. The paper does not say what happens otherwise.
• Hidden communication. Each BASE layer in training makes four all-to-all calls: shuffle, route, return and undo-shuffle base_layer.py base_layer.py base_layer.py base_layer.py. The paper concedes that balancing has non-trivial overhead and leaves faster solvers to future work §1 §7.
• Train-time and test-time routing differ. Training balances load; testing does not, so inference load can be uneven. The paper reports that the trained routing is reasonably balanced anyway §5.1 fig. 8.
• The solver needs a build step. The code loads a C++ extension and fails with a message to run python setup.py build_ext --inplace if it is missing base_layer.py.
• Checkpoints are per rank. When base_layers > 0, checkpoints are saved on every rank with a rank suffix, and loaded on every rank trainer.py trainer.py utils.py. The paper does not mention this.
• Grad mode silently selects the router. Any forward pass without gradients, including validation during training, uses greedy routing rather than the balanced assignment. Check the grad mode of each call site base_layer.py base_layer.py.
• Parameter count. The introduction says experiments reach 110B parameters §1; the ablation lists a 117B configuration §4.4. The paper does not reconcile these.
• Specialisation is mostly local. The paper's own analysis finds that many assignments track the previous token's identity, which it calls superficial §5.2 tab. 1. It does not claim semantic experts.
• Figure data is not in the text. Perplexity curves are only in the plots, so a reproduction needs the plot data from the LaTeX source or a rerun.

Lineage

• Builds on: Sparsely Gated Mixture of Experts (arXiv:1701.06538) §1; GShard (arXiv:2006.16668) §1; Switch Transformer (arXiv:2101.03961) §1. BASE drops the balancing loss and the capacity factor, and its authors show that a single expert layer is effective §1 §6.
• Model parallel training follows Megatron-LM (arXiv:1909.08053) §2.1.
• Compared against: data-parallel and model-parallel dense training §4.2; Sparsely Gated MoE and Switch re-implemented by the authors §4.3.
• Superseded by: not stated in the card's content.