BASE Layers: Simplifying Training of Large, Sparse Models

By Michael Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, Luke S. Zettlemoyer. 2021

411 citations recorded by the source corpus.

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/

Subfields: Artificial Intelligence, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Efficiency & Compression

Citations in the corpus

Unified Scaling Laws for Routed Language Models (2022, methodology): “We analyze three different techniques for training Routing Networks, detailed in §3: Sinkhorn-BASE, a sparse mixture-of-experts (SMOE) approach modifying BASE [Lewis et al., 2021]; non-parametric HASH Layers [Roller et al., 2021]; and routing via Reinforcement Learning (RL-R).”
Efficient Large Scale Language Modeling with Mixtures of Experts (2021, methodology): “Recent work (Lewis et al., 2021; Lepikhin et al., 2021; Fedus et al., 2021; Fan et al., 2021) has studied different conditional compute strategies that work well with”
GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (2021, background): “Various routing strategies (Gross et al., 2017; Lewis et al., 2021; Dua et al., 2021) have been investigated such that transfer learning can be achieved with task dependencies.”
A Review of Sparse Expert Models in Deep Learning (2022, methodology): “Load balance has been accomplished by auxiliary losses (Shazeer et al., 2017) as well as through treating this as a linear assignment problem (Lewis et al., 2021; Clark et al., 2022).”
Mixture-of-Experts with Expert Choice Routing (2022, background): “Base Layers [22] formulate a linear assignment that maximizes token-expert afﬁnities while ensuring each expert receives an equal number of tokens.”
Language Models that Seek for Knowledge: Modular Search & Generation for Dialogue and Prompt Completion (2022, methodology): “RoBERTa+CC100en We use the same data used to train the BASE language model (Lewis et al., 2021), which consists of approximately 100B tokens, combining corpora used in RoBERTa (Liu”
BlenderBot 3: a deployed conversational agent that continually learns to responsibly engage (2022, methodology): “R2C2 uses RoBERTa+cc100en Data – the same data used to train Lewis et al. (2021), which consists of approximately 100B tokens, combining the corpora used in RoBERTa (Liu et al., 2019) with the English subset of the CC100 corpus (Conneau et al., 2020).”
8-bit Optimizers via Block-wise Quantization (2021, methodology): “, 2016) and evaluated on en-de WMT’14 (Macháček and Bojar, 2014), large-scale language modeling (Lewis et al., 2021; Brown et al., 2020) and RoBERTa pretraining (Liu et al.”
Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (2021, methodology): “They are frequently used in transformer language modeling (Baevski & Auli, 2018; Lewis et al., 2021) and machine translation (Vaswani et al., 2017; Ott et al., 2018) models.”
Efficient Transformers: A Survey (2020, background): “However, we believe that the emergence of an entire research direction (Roller et al., 2021; Lewis et al., 2021; Lepikhin et al., 2020; Du et al., 2021) based on sparse efficient should warrant a new category of efficient Transformers.”
Hash Layers For Large Sparse Models (2021, methodology): “We use the architecture, data (RoBERTa+cc100en), and hyperparameters directly from [10], using either a single sparse routing layer consisting of 3 stacked FFNs (D = 8192) on the middle layer of a 25 layer network, or 3 routing layers evenly spaced in the network.”
DEMix Layers: Disentangling Domains for Modular Language Modeling (2021, methodology): “These settings are inspired by those reported by Lewis et al. (2021).”
Director: Generator-Classifiers For Supervised Language Modeling (2022, positive): “We use GPT2-Medium (Radford et al., 2019) as our baseline model, ﬁne-tuning on the BASE data of (Lewis et al., 2021) to predict the next sentence, and using greedy decoding during generation.”
Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts (2022, background): “Many works aim to improve the gating mechanism itself, by making it differentiable [38], reformulating as a linear assignment task [39] or even swapping it out for a simple hashing algorithm [40].”
Taming Sparsely Activated Transformer with Stochastic Experts (2021, neutral): “1 (noisy gating, Shazeer et al. 2017), limiting the maximum number of inputs that can be routed to an expert (expert capacity, Lepikhin et al. 2020), imposing a load balancing loss (Lepikhin et al., 2020; Fedus et al., 2021), and using linear assignment (Lewis et al., 2021).”
Tricks for Training Sparse Translation Models (2021, methodology): “BASELay-ers (Lewis et al., 2021) employ a simple mechanism that learns a balanced routing without the need for additional auxiliary losses.”