PaperlayerView exact revision on arXiv

Source-linked AI summary

Unified Scaling Laws for Routed Language Models

Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew Johnson, Katie Millican, Albin Cassirer, Chris Jones, Elena Buchatskaya, David Budden, Laurent Sifre, Simon Osindero, Oriol Vinyals, Jack Rae, Erich Elsen, Koray Kavukcuoglu, Karen Simonyan

arXiv:2202.01169v2cs.CLcs.LG

TL;DR

Routing Networks challenge parameter-only scaling because each input uses only a subset of parameters, separating total capacity from computational cost. The paper derives shared scaling laws across three routing techniques and finds that routing improves performance while enabling a unified effective parameter scale for dense and routed models.

  • Problem

    Routing Networks decouple parameter count from per-input computation, but their general scaling behavior is not well understood.

  • Method

    The paper evaluates three routing-training techniques across routed language models and derives scaling laws in model size, experts, parameters, and inference compute.

  • Results

    Routing improves language-model performance across all tested sizes and variants, and reinforcement-learning routing is comparably effective to state-of-the-art techniques.

  • Takeaways & Limitations

    The scaling laws provide an empirical framework for comparing routing architectures and map dense and routed models onto a shared Effective Parameter Count.

  • Takeaways & Limitations

    The reported numerical values apply specifically to models trained on 130 billion tokens, although the analysis is described as robust to token count.

Abstract

from arXiv · show

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.

1 Introduction

Routing Networks separate parameter capacity from computation because each input uses only a subset of parameters. The paper studies their scaling and finds laws that unify routed and dense models.

  • Motivation: Routing Networks conditionally activate only a subset of parameters for each input, making parameter count and computation independent size axes.This prevents scaling laws based on parameter count alone from fully describing the architecture class.
  • Key findings: Models improve with routing across all tested sizes and variants, while reinforcement-learning routing matches the effectiveness of state-of-the-art techniques.The evaluation includes models with up to 200 billion parameters.
  • Scaling laws: Performance is accurately described by scaling laws in both the underlying dense model size and the number of experts.These laws generalize earlier scaling laws for standard language models.
  • Unified representation: The laws can be restated using parameter count and inference compute, allowing a shared fit across a wider range of routing architectures.This reframes routing along two independent resource axes.
  • Applications: An Effective Parameter Count maps dense and routed networks onto equivalent performance and scaling.The mapping provides a single scale for comparing the two architecture families.

2 Background

The background motivates routing as conditional computation: dense scaling improves performance but also increases compute, whereas routing separates total capacity from per-input cost. The paper studies sparse mixture-of-experts Transformers under fixed training conditions.

  • Language modeling: Language modeling autoregressively predicts natural-language tokens, with validation negative log-likelihood as the primary performance metric.Zero-shot transfer to other tasks is considered only briefly.
  • Existing scaling laws: Power-law scaling links larger parameter counts to lower loss but also to greater training computation, motivating architectures that disentangle size and compute.Routing Networks are presented as one such class.
  • Routing Networks: A Routing Network dynamically selects a fixed subset of parameters for each input, while the studied layout uses sparsely activated mixtures of experts.Several feedforward components are converted to routed equivalents in decoder-only Transformers.
  • Routing a single layer: For a routed layer, E expert parameter versions are created and a router selects which expert processes each input.The routed output is the selected expert’s transformation, with router computation adding only a small overhead.
  • Model size and inference cost: N denotes the parameters any input interacts with, whereas P denotes total parameters; dense models have P = N, while routed models have P roughly proportional to N · E.Inference cost F is approximately that of the dense equivalent, aside from router overhead.
  • Training setup: All models were trained for 130 billion tokens regardless of N, making token count an important condition on interpreting the scaling results.Training used fixed sequence length, batch size, and optimization setup.

3 Routing Techniques

The paper compares three routing-training techniques—SMOE, HASH, and reinforcement learning—and argues that their behavior can be analyzed within shared scaling laws. The methods differ in how expert selection is learned or fixed.

  • Methods: The study examines three techniques: Sinkhorn-BASE, HASH Layers, and routing via Reinforcement Learning.These methods represent sparse mixture-of-experts, fixed non-parametric routing, and learned policy routing.
  • Sparse Mixture-of-Experts: SMOE methods use expert-selection probabilities as output multipliers, providing gradients to router logits despite non-differentiable expert sampling.The representative implementation, Sinkhorn-BASE, replaces Hungarian matching with regularized optimal transport and Sinkhorn approximation.
  • Observed performance: Increasing the number of experts reduces validation loss, with S-BASE performing best across model sizes, followed by RL-R and HASH.The figure compares fitted validation losses and overlapping prediction curves across techniques and model sizes.
  • HASH Layers: HASH Layers use no additional router parameters and select experts by taking the token ID modulo the number of experts.This avoids jointly learning router and expert parameters.
  • Routing via Reinforcement Learning: RL-R treats routing as a policy whose actions select experts and trains it with a policy-gradient term combined with language-modeling loss.The reward is based on the probability assigned to the correct output token.
  • Routing via Reinforcement Learning: RL-R directly optimizes expert-selection actions for language-modeling loss but faces high gradient variance addressed here with REINFORCE and a learned baseline.The method’s theoretical Markov assumption is also disrupted because optimal selection depends on the rest of the network.

4 Scaling Behavior at Convergence

Routing-network performance is modeled jointly through base-model size and expert count, with extensions that capture their interaction, expert-count saturation, and architectural differences. The resulting laws fit the evaluated routing methods and support comparison across parameter count and inference compute.

  • Scaling-law form: The proposed six-parameter law models converged loss as bilinear in log N and log Ẽ, where Ẽ is a saturating transformation of expert count E.The model is log L(N, E) = a log N + b log Ẽ + c log N log Ẽ + d.
  • Empirical validation: Across 168 models spanning six sizes and expert counts from 2 to 512, increasing E reduces validation loss and the fitted law shows good agreement across all three routing methods.The sweep includes models trained with Sinkhorn-BASE, HASH, and RL-R, plus dense baselines.
  • Separable scaling: The separable power law performs poorly because b(N) increases with N, indicating that routing benefits diminish as the base model grows.The separable fit can mis-predict validation losses by over 0.2 and has held-out RMSLE values greater than 80e-4.
  • Quadratic interaction: The bilinear interaction improves prediction for large N and is consistent with dense-model scaling through α_N ≈ -a and N_c ≈ exp(d/a).Fitting the bilinear law substantially reduces prediction error for large N while retaining similar predictions for dense models.
  • Interpretation: Positive c implies that expert-improvement slopes decrease with model size, and all three routing techniques therefore predict diminishing gains from routing at larger scale.The paper identifies c ≈ 0 as a goal for future routing methods and reports S-BASE as having substantially better scaling behavior.
  • Bounded and shared scaling: Replacing E with a bounded transform captures deviations from power-law scaling at low and high expert counts, while changing variables to inference compute F and parameter count P enables nearly shared fits across architectures.Level curves in (N, E) differ across routing settings, whereas curves in (F, P) almost overlap; the bounded transform approaches E_max as E increases.

5 Scaling Law Applications

The scaling laws yield an Effective Parameter Count that unifies dense and routed models, clarifies routing’s cutoff with increasing base-model size, and supports comparisons and practical recommendations across routing techniques.

  • 5.1 Effective Parameter Equivalence: Effective Parameter Count maps a routed model to a dense model with equivalent predicted performance.It is defined by solving L(N̄, 1) = L(N, E).
  • 5.1 Effective Parameter Equivalence: A single power law fits dense and routed models across four orders of magnitude when performance is plotted against Effective Parameter Count.
  • 5.2 Routing Behavior for Large N: Routing stops improving performance above a method-dependent cutoff: 937B for S-BASE, 85B for RL-R, and 83B for HASH.These cutoff values depend strongly on the number of training tokens and are expected to increase with more tokens.
  • 5.3 Comparative Analysis: S-BASE consistently outperforms RL-R and HASH, while RL-R remains competitive at smaller base-model sizes.
  • 5.3 Comparative Analysis: All routing methods lose efficacy as N increases, with S-BASE scaling best because it has the lowest fitted interaction coefficient c.
  • 5.3 Comparative Analysis: HASH and RL-R preserve power-law behavior to larger Emax, but greater interference c generally produces worse performance than S-BASE.
  • Recommendations: Practitioner guidance recommends routing for N ≤ 1.3B, using S-BASE by default, targeting 64–128 experts, and routing layers at frequency 0.5 ≤ R ≤ 1.Larger expert counts bring diminishing returns, while lower routing frequency reduces performance.
  • Recommendations: Future comparisons should evaluate multiple values of N and E because results at a single size cannot be extrapolated.

6 Related Work

The work builds on empirical scaling-law studies, advances in scalable Transformers, and the broader literature on Routing Networks, conditional computation, and Mixture of Experts.

  • The empirical analysis follows Kaplan et al. and related work on language-model scaling laws, while noting that the underlying theory remains less understood.
  • The study relies on advances that made Transformers scalable and efficiently parallelizable across multiple devices.
  • Routing Networks are connected to earlier work on conditional computation and Mixture of Experts, as well as per-example architecture search and graph formulations.

7 Conclusion

The paper introduces scaling laws for Routing Networks that quantify how experts and model size affect language-model performance. These laws unify routed and dense models and provide an empirical framework for comparing future routing methods.

  • The proposed scaling law predicts that routing improves performance, with gains following an expert-count power law that diminishes as base-model size increases.
  • The laws quantify differences among three routing techniques and produce a scalar that describes routed and dense models alike.
  • The framework is intended to support analysis of future routing innovations and indicates benefits up to base-model sizes greater than 900 billion parameters.

A Architecture

The architecture section specifies the Transformer implementation, model-size setup, routing experiments, and load-balancing procedures used to study Routing Networks.

  • Implementation: The models use a decoder-only Transformer with relative positional encodings, SentencePiece tokenization, Megatron-style feedforward sharding, and bfloat16 activations.
  • Model definitions: Seven model sizes are used, with feedforward width fixed at four times the activation width and matching key and value dimensions.
  • Random-seed evaluation: The routing study trains three random seeds on a 130M model with 8 and 256 experts; seed differences remain below 0.01 before 100,000 steps.The observed divergence closely matches the 0.02 error cited from prior scaling-law work.
  • Routing methods: The implementation details and comparisons cover routing techniques and key alternatives, including uniform-routing and load-balancing procedures.
  • Random-seed evaluation: The routing experiments use a figure that color-codes training curves by random seed for 1, 8, and 256 experts alongside maximum disagreement between runs.

B.2.1 Sinkhorn Redistribution

Sinkhorn-BASE redistributes router assignments through a constrained optimal-transport formulation solved efficiently with Sinkhorn iterations. The section also describes RL-R variants and HASH routing, highlighting their implementation trade-offs and balancing behavior.

  • Sinkhorn redistribution: Sinkhorn redistribution adds per-expert balancing constraints to router assignments, yielding a regularized optimal-transport problem.The formulation seeks approximately equal average token assignments across experts.
  • Sinkhorn redistribution: The Sinkhorn algorithm maps logits L ∈ R^T×E to a soft-assignment matrix Π ∈ R^T×E through alternating dual updates.Iterations are early-stopped by measuring primal constraint violation in L1 norm.
  • Sinkhorn redistribution: Greedy selection projects the Sinkhorn probabilities into a transportation plan by assigning each token to its highest-probability device.This follows computation of the balanced routing plan.
  • Efficiency: Sinkhorn has O(N × E) complexity versus O((N × E)3/2) for the auction algorithm and is better suited to batch computation on TPU/GPU.The auction method must run on CPU and can bottleneck models with many routed layers.
  • Efficiency: An error tolerance of etol = 10^-2 gives consistently good performance, with 1% to 3% end-to-end overhead relative to Switch.The reported computational offset is described as negligible compared with the per-step performance gain.
  • Other routing methods: RL-R includes greedy, nucleus-sampled, and learned-baseline variants, while HASH uses modular token-index assignment and becomes harder to balance as expert count increases.The RLR-B variant consistently outperforms the other RL-R variants, whereas HASH’s balancing depends strongly on its hashing function.

D.1 Robustness to hyper-parameter changes

S-BASE is comparatively robust to balancing-weight changes and improves with routing frequency, whereas RL-R is highly hyperparameter-sensitive and HASH degrades or saturates in several higher-frequency settings. Routing-layer placement and expert multiplicity create additional performance–efficiency trade-offs.

  • Hyperparameter robustness: RL-R is highly sensitive to hyperparameters, especially balancing weight, while policy-entropy changes can produce unbalanced routers unless jointly tuned.The evaluation focuses on E = 512, where performance variance was reportedly largest.
  • Hyperparameter robustness: S-BASE remains robust to balancing weights from 1e-3 to 1, though omitting the balancing loss makes training less stable.Switch shows higher expert oversubscription rates even after balancing-weight tuning.
  • Routing frequency: S-BASE performance improves with routing frequency, whereas HASH flattens or degrades as frequency increases across 8E, 64E, and 256E settings.With only one routed layer, HASH has the lowest validation loss across model sizes.
  • Shared routing decisions: When routing decisions are shared across layers, S-BASE and HASH have similar losses; S-BASE improves when routing is learned independently at each expert layer.Static HASH routing remains unchanged because it selects experts from the input token index.
  • Layer placement: Three evenly spaced routed layers slightly outperform three stacked expert feed-forward layers for a 32E 1.3B model, with S-BASE benefiting more from interleaving.This is consistent with the reported routing-frequency results.
  • Experts per datapoint: Higher K always improves performance per step, but K = 1 remains on the Pareto front for FLOP efficiency and reducing K by half nearly doubles practical speed.The speed observation includes both inference and training and is associated with reduced communication costs.

E Effects of scaling strategy on Zero-shot Transfer

Zero-shot scaling gains vary substantially across downstream tasks and routing techniques, so validation-set scaling does not determine a universal downstream outcome. Effective Parameter Count likewise differs by task because scaling coefficients and their interactions differ.

  • Task-dependent scaling: Scaling coefficients vary greatly across tasks, producing different relative gains from increasing expert count or base model size.The analysis estimates coefficients separately for each task and routing technique.
  • Task-dependent scaling: WikiText-103 has higher b and lower a than the validation set, illustrating that expert and base-size scaling can differ by task.The passage uses WikiText-103 as a concrete coefficient comparison.
  • Technique comparisons: S-BASE often has lower effective coefficients at larger E and N, but technique ordering varies across tasks.RL-R has better b(N) than S-BASE for much of LAMBADA before being overtaken at N = 410M, while S-BASE is always superior in C4.
  • Effective Parameter Count: Downstream tasks require careful analysis because their EPC values and slopes differ, potentially changing the apparent benefit of routing-based scaling.The paper frames this as a consequence of task-specific b and c coefficients.

F.1 Convergence on Large Datasets

The models do not reach convergence on the large datasets studied, so the reported scaling behavior is defined at fixed training steps or token counts rather than at an observed converged limit. Scaling coefficients change with training and depend on additional variables that were not fully modeled.

  • Convergence: Even the smallest models remain unsaturated after 130 billion tokens, with no observed overfitting or plateau in validation performance.The authors therefore question whether converged performance is empirically available for these settings.
  • Convergence: Training a 15M HASH model four times longer reduces loss by 0.05 without routing and 0.07 with E = 512.The longer runs use half a trillion tokens.
  • Token-count dependence: The reported scaling behavior is valid only as a function of a particular number of steps or tokens, with instantaneous coefficients evaluated at S steps.This avoids claiming analysis at an unobserved convergence point.
  • Token-count dependence: The expert-scaling coefficient b is non-constant and is not uniquely determined by token count alone.This motivates matching all possible variables when comparing scaling coefficients.
  • Limitations: Critical-batch-size effects are not analyzed because the models use a fixed batch size, and estimating them would require a substantially larger sweep.The authors identify this as a limitation for fitting scaling behavior.
  • Modeling implications: The tested training curves contradict the constant-scale prediction associated with the Kaplan et al. formulation, although the authors do not claim to disprove it generally.They conclude that the cited equation cannot model their specific training curves.

G Large Scale Routing Behavior, Coefficient Sensitivity, and Future Work

The analysis extrapolates routing behavior to larger model and expert counts while emphasizing diminishing returns, compute-limited regimes, and sensitivity to the scaling coefficient b. Supplementary figures and tables provide additional fits, visualizations, and downstream-task coefficients.

  • Large-scale behavior: Larger expert counts are predicted to improve performance, especially for small models, but at a diminishing rate.The analysis also predicts routing benefits as the base model size N increases by at least one, possibly two, orders of magnitude.
  • Limits and future work: Practical compute limitations prevented experiments from probing the predicted large-N and large-E regimes.The authors identify unanswered questions about behavior at the limits of the two scaling variables.
  • Large-scale behavior: A routed model with N > 7,000,000 substantially outperformed a same-size dense baseline despite differences preventing accurate prediction by the paper’s fits.The experiment was evaluated at 100B tokens and provides a rough upper bound on error in b for RL-R.
  • Limits and future work: The authors caution that concrete extrapolation is impossible for the larger RL-R model, while its competitive improvement suggests the measured scaling coefficients may not substantially overestimate scalability.They encourage future work to probe routing limits in both N and E and refine coefficient estimates.
  • Supplementary analyses: Supplementary material includes visualizations, affine-fit figures, joint fits, and tables of scaling coefficients, including downstream-task coefficients.The listed materials cover S-BASE, RL-R, and HashLayer models, along with fits varying K and R.
Loading 2202.01169v2…