Unified Scaling Laws for Routed Language Models

By Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake A. Hechtman, Trevor Cai, Sebastian Borgeaud and 16 more. 2022

285 citations recorded by the source corpus.

The performance of a language model has been shown to be effectively modeled as a power-law in its parameter count. Here we study the scaling behaviors of Routing Networks: architectures that conditionally use only a subset of their parameters while processing an input. For these models, parameter count and computational requirement form two independent axes along which an increase leads to better performance. In this work we derive and justify scaling laws defined on these two variables which generalize those known for standard language models and describe the performance of a wide range of routing architectures trained via three different techniques. Afterwards we provide two applications of these laws: first deriving an Effective Parameter Count along which all models scale at the same rate, and then using the scaling coefficients to give a quantitative comparison of the three routing techniques considered. Our analysis derives from an extensive evaluation of Routing Networks across five orders of magnitude of size, including models with hundreds of experts and hundreds of billions of parameters.

Subfields: Artificial Intelligence, Scaling & Learning Theory

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 Scaling & Learning Theory

Citations in the corpus

Training Compute-Optimal Large Language Models (2022, neutral): “However, for very large models the computational benefits of routed models seems to diminish (Clark et al., 2022).”
Mixtral of Experts (2024, methodology): “There are multiple alternative ways of implementing G ( x ) [6, 15, 35], but a simple and performant one is implemented by taking the softmax over the Top-K logits of a linear layer [28].”
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).”
Memorization Without Overfitting: Analyzing the Training Dynamics of Large Language Models (2022, unknown): “Recent efforts have focused on trying to model the scaling laws for language models, including data and model size [46, 47], applications to transfer learning [48], routing networks [49], and various autoregressive generative tasks [50].”
QMoE: Practical Sub-1-Bit Compression of Trillion-Parameter Models (2023, background): “Mixture-of-expert models are a popular research direction aimed at creating significantly more efficient large-scale models (Fedus et al., 2022; Artetxe et al., 2022; Clark et al., 2022).”
Jamba: A Hybrid Transformer-Mamba Language Model (2024, methodology): “These choices were inspired by prior work on MoE [7, 49] and verified in preliminary experiments.”
A Survey of Large Language Models (2023, unknown): “It has been shown that substantial performance improvement can be observed by increasing either the number of experts or the total parameter size [158].”
Designing Effective Sparse Expert Models (2022, background): “Concurrently, Clark et al. (2022) provides additional design recommendations including higher layer frequency and top-1 routing as per Fedus et al. (2021).”
From Sparse to Soft Mixtures of Experts (2023, background): “Other works suggest more advanced (and often costly) algorithms to compute the assignments, such as approaches based on Linear Programming algorithms (Lewis et al., 2021), Optimal Transport (Liu et al., 2022; Clark et al., 2022) or Reinforcement Learning (Clark et al., 2022).”
ST-MoE: Designing Stable and Transferable Sparse Expert Models (2022, background): “Concurrently, Clark et al. (2022) provides additional design recommendations including higher layer frequency and top-1 routing as per Fedus et al. (2021).”
A Survey on Mixture of Experts in Large Language Models (2024, unknown): “4(a)) can also yield competitive results, a ﬁnding that has been substantiated and adopted by later studies [63].”
Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints (2022, methodology): “A number of routing algorithms have been developed, for example Top-K (Shazeer et al., 2017), BASE and Sinkhorn-BASE layers (Lewis et al., 2021; Clark et al., 2022), Hash layers (Roller et al., 2021), and Expert Choice routing (Zhou et al., 2022).”
A Survey on Dynamic Neural Networks for Natural Language Processing (2022, methodology): “Clark et al. (2022) examine the scaling law of BASE Layer (Lewis et al., 2021), Hash Layer (Roller et al., 2021) and earlier Reinforcement Learning-based routing algorithms providing suggestions for best-practices in training MoE models.”
On the Representation Collapse of Sparse Mixture of Experts (2022, methodology): “Besides, the scaling behaviors of SMoE are also studied (Clark et al., 2022; Du et al., 2021).”