[Paper Review] BASE Layers: Simplifying Training of Large, Sparse Models

Mike Lewis, Shruti Bhosale+3 more|arXiv (Cornell University)|Mar 30, 2021

Topic ModelingComputer Science29 references63 citations

TL;DR

BASE layers introduce an optimal, balanced token-to-expert assignment using a linear assignment problem, enabling a single-expert per token sparse routing with no extra losses or hyperparameters, improving compute efficiency.

ABSTRACT

We introduce a new balanced assignment of experts (BASE) layer for large language models that greatly simplifies existing high capacity sparse layers. Sparse layers can dramatically improve the efficiency of training and inference by routing each token to specialized expert modules that contain only a small fraction of the model parameters. However, it can be difficult to learn balanced routing functions that make full use of the available experts; existing approaches typically use routing heuristics or auxiliary expert-balancing loss functions. In contrast, we formulate token-to-expert allocation as a linear assignment problem, allowing an optimal assignment in which each expert receives an equal number of tokens. This optimal assignment scheme improves efficiency by guaranteeing balanced compute loads, and also simplifies training by not requiring any new hyperparameters or auxiliary losses. Code is publicly released at https://github.com/pytorch/fairseq/

Motivation & Objective

• Motivate the use of sparse expert models to scale language models while reducing training costs.
• Propose a balanced, linearly-assigned BASE layer that eliminates balancing losses and capacity factors.
• Demonstrate that a single-expert per token routing with balanced assignment yields strong compute efficiency across large models.

Proposed method

• Formulate token-to-expert assignment as a linear assignment problem to ensure each expert handles an equal number of tokens.
• Use an auction algorithm to solve the assignment in parallel across workers and edge cases with token shuffling for robustness.
• Implement a BASE layer where each token is routed to a single expert, followed by a soft mixing of the expert output via a residual connection.
• Train with balanced token allocation to maximize throughput and avoid extra balancing losses, while testing with greedy expert selection at inference.
• Measure compute efficiency as perplexity given fixed GPU runtimes and compare against dense and prior sparse approaches.

Experimental results

Research questions

• RQ1 Can a linear assignment-based routing scheme balance expert usage without additional loss terms?
• RQ2 Does a single-expert BASE layer match or exceed the efficiency of existing dense and sparse methods across large models?
• RQ3 How does BASE placement and size affect performance and compute efficiency?
• RQ4 What is the impact of balanced routing on expert specialization and load balance during training?

Key findings

• BASE layers achieve balanced token-to-expert routing via linear assignment without Balancing losses or capacity factors.
• A single BASE layer can substantially outperform dense data-parallel and model-parallel baselines in higher compute budgets.
• BASE compares favorably with Sparsely Gated MoE and Switch transformers on compute efficiency, sometimes matching or exceeding their performance at similar budgets.
• Interleaving multiple BASE layers yields increasing performance gains while keeping total parameter count roughly constant.
• BASE layers show robust performance across different placements and numbers of BASE sublayers, indicating architectural robustness.