Unified Scaling Laws for Routed Language Models

Introduction

It is a commonly held belief that increasing the size of a neural network leads to better performance, especially when training on large and diverse real-world datasets. This vague and debated notion has become increasingly justified as large empirical studies have shown that the performance of models on many interesting classes of problems are well understood as power-laws; where a multiplicative increase in model size leads to an additive reduction in the model’s loss (Kaplan et al., 2020; Hernandez et al., 2021; Henighan et al., 2020; Rosenfeld et al., 2019). These relationships are not well understood, but a key implication is that a sequence of smallMeasured as training or inference floating point operations, devices or time required, financial cost, carbon emissions, etc. models can be used both to infer the performance of models many times more powerful, but also to provide global information about the scalability of an architecture.

Enter Routing Networks: models with the unusual property that each input interacts with only a subset of the network’s parameters — chosen independently for each datapoint (Bengio et al., 2016, 2013; Denoyer and Gallinari, 2014). For a Routing Network, the number of parameters is nearly independent from the computational cost of processing a datapoint. This bifurcates the definition of size and prevents a scaling law in parameters alone from fully describing the model class. Specific Routing Networks have been trained successfully at large scales (Fedus et al., 2021; Du et al., 2021; Artetxe et al., 2021), but the general scaling behavior is not well understood. In this work we analyze the behavior of routed language models so that we might infer the scaling laws that describe their performance.

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). With models up to 200 billion parameters, we observe the following:

Routing improves the performance of language models across all sizes and variants attempted (see Fig. 1).

Training a Routing Network with RL (§3.3), a technique used in early routing work (Bengio et al., 2013), is of comparable effectiveness to state-of-the-art techniques.

The performance of all Routing Networks is accurately described by scaling laws in the number of experts and in the underlying dense model size (§4) which generalize those from Kaplan et al. (2020).

These laws can be restated in terms of parameter count and inference compute, capturing an even wider set of routing architectures under a shared fit (§4.4).

They further imply an Effective Parameter Count: a mapping equating the performance and scaling for both dense and routed networks (§5).

Background

We first review the language modelling problem and existing scaling laws before discussing the process of routing a neural network and how it is applied to language models.

We consider the problem of autoregressively predicting natural language, a task with consistent and predictable scaling characteristics across many orders of magnitude (Henighan et al., 2020; Kaplan et al., 2020). The objective is to maximize the likelihood of a sequence of tokens $P(x_{1},\dotsc,x_{T})$ factored auto-regressively as $p\left(x_{1},\dotsc,x_{T}\right)=\prod_{i}^{T}p\left(x_{i}|x_{j<i}\right)$. Our primary metric of performance is the negative log-likelihood of a validation dataset whose statistics match the training distribution. We focus on this validation loss, but briefly consider zero-shot transfer to other tasks in App. E.

We train on a multi-trillion-token compendium of English language text comprising documents from the internet alongside open-source text datasets, details of which are given in Rae et al. (2021). In this setting Kaplan et al. (2020) argue that the converged performance of a model trained on a dataset of infinite size is a power-law in the model’s parameter count $N$. Our dataset is not infinite, but its size – and the lack of any observed overfitting – make this a reasonable approximation. We consider the final (and best) evaluation value as the converged value, though this is also an approximation which is discussed further in App. F.

1 Routing Networks

Power-law scaling implies the performance of a language model increases with size, but so too does the compute needed to train the model. This undesirable connection between size and computation motivates a search for architectures wherein the two are disentangled. Routing Networks are one such class of model: a type of neural network that incorporates a specific flavor of conditional computation. In a Routing Network, each input (e.g., a token of text) is transformed into an output while only interacting with a fixed subset of the network’s parameters – dynamically selected based on the input itself. Many sparsely-activated networks have this property, but here we exclusively study the layout based on Sparse Mixtures of Experts (Shazeer et al., 2017) where multiple sub-components of a deep neural network (i.e., several layers) are independently converted to routed equivalents and jointly trained with the rest of the network.

We apply routing to a decoder-only Transformer (Vaswani et al., 2017) to measure the scaling properties that result: an architecture chosen due to its state-of-the-art performance. Details of the baseline architecture we use are in App. A. We will refer to non-routed Transformers as dense models, in opposition to Routed Transformers which sparsely activate some of their parameters. Our conversion to a Routed Transformer is the same as is used in prior work (Lepikhin et al., 2020; Fedus et al., 2021). Namely, we apply routing to every other set of feedforward components (FFWs) of the Transformer, sub-components that act on each timestep independently. Though different layers can have different numbers of experts, here all routed layers share the same number of experts $E$, and we will refer to the network as being routed $E$ ways.

We use $N$ to indicate a network’s dense model size: the number of parameters any one input interacts with. This is in opposition to $P$: the total number of parameters. For a dense model, $P=N$, whereas for a Routing Network $P$ is roughly proportional to $N\cdot E$, with factors that depend on details of the routing architecture (§4.4). Except for a small overhead due to running the routers, the cost $F$ (in TeraFLOPs) of executing a Routed Transformer is the same as its dense equivalent.

All models are trained on TPUs with JAX (Bradbury et al., 2018) using a combination of data, expert (see App. C) and sharding parallelism (Shoeybi et al., 2019). Models were trained with a sequence length of $2048$ and batch size of $256$ for 250,000 steps, i.e. 130 billion tokens, regardless of $N$. This is an important detail, and we discuss some of the implications in App. F. All were optimized with AdamW (Loshchilov and Hutter, 2018) and ZeRO Stage 1 was used to shard the optimizer state (Rajbhandari et al., 2020). Appendix A contains further details.

Routing Techniques

If the benefit of Routing Networks is the decoupling of parameter capacity from network cost, the fundamental difficulty is in effectively learning the parameters $\varphi$ of the router given the non-differentiability of its output. Much research in Routing Networks has therefore focused on techniques for learning $\varphi$. A major finding of this work is that three notably different techniques of training Routing Networks are effectively described by the same scaling laws. We now introduce and contextualize these three methods.

Sparse Mixture-of-Experts (smoe) methods (Shazeer et al., 2017) solve the problem of non-differentiability by reusing the probability of expert selection as a scalar multiplier on that expert’s output, guaranteeing a gradient passed to the logits of selected experts despite the the non-differentiability of sampling from those logits. Formally, the router is given as $\rho(x)=\operatorname{topk}(Wx+b)$, where $Wx+b$ is an unnormalized distribution over $\left[1,E\right]$ from which the experts corresponding to the top $K$ values are selected. In the final output of the routed layer, the normalized logits are reused as gating weights, i.e. the final output of the routed layer is $h(x)=\sum_{i\in\rho(x)}g_{i}(x)f_{i}(x)$ where $g(x)=\operatorname{softmax}(Wx+b)$.

Though this formulation supplies a gradient to $\varphi=(W,b)$, it represents changes to the scalar multiplier and does not directly correspond to optimizing expert selection. This method is nevertheless effective, and can be seen as a sparse approximation to dense mixture of experts models (Eigen et al., 2014; Jacobs et al., 1991) where the likelihood of skipping an expert is inversely proportional to the value of its scalar gate $g_{i}$.

It was conjectured that smoes require $(K\geqslant 2)$-way routing to produce effective gradients in the routers (Shazeer et al., 2017), and many attempts at incorporating routing into large Transformers use $K=2$ (Lepikhin et al., 2020; Du et al., 2021). However recently this has been challenged, and stable modifications have been proposed for $K=1$; namely the Switch Transformer (Fedus et al., 2021). Most smoes, including Switch, are reliant on auxiliary balancing losses which encourage the router output $\rho(x)$ to be more uniform across minibatches of inputs. To improve on this, base (Lewis et al., 2021) post-processes the router output with a Hungarian Matching algorithm that re-assigns expert selections to ensure that all experts are selected evenly.

Our implementation of base replaces the Hungarian Matching with a regularized Optimal Transport formulation (Cuturi, 2013) using the Sinkhorn algorithm as an approximate matching step during expert selection. This substantially improves routing efficiency on accelerated hardware (details in §B.2.1). We call the resulting method Sinkhorn-base (s-base), and use it as the representative of smoe methods, as early tests showed the benefit of its balancing mechanism.

2 Input-based Deterministic Hash Routing

An alternative approach eschews extra parameters completely and represents $\rho$ as a fixed function of the input. This is the concept pioneered by hash Layers (Roller et al., 2021) which circumvents the need to simultaneously learn $\varphi$ and $\theta$. Our implementation takes the token ID assigned to the input by the SentencePiece tokenizer (Kudo and Richardson, 2018) and uses the remainder of it divided by $E$ as the expert selection. See §B.4 for details.

3 Routing via Reinforcement Learning

Finally, we re-analyze a technique that optimizes the router via Reinforcement Learning (a class of methods we call rl-r), which was proposed in early work on neural conditional computation (Bengio et al., 2013, 2016; Bengio, 2017; Denoyer and Gallinari, 2014). In this approach each router is seen as a policy whose actions are the selection of an expert in each routed layer and whose observations are the activations passed to that router. After completing the forward pass, the probability the Routed Transformer assigns to the correct output token can be used as a reward, maximization of which is equivalent to minimization of NLL. To jointly train the experts and the router, we minimize a composite loss formed with the language modelling loss and a policy-gradient term (Sutton et al., 2000) using the selected set of experts as actions. We highlight that the optimal expert selection is dependent not only on the input activations but on the parameters of the rest of the network. This disrupts the theoretical underpinning, crucial to RL, that this is a Markov Decision Process. Nevertheless, it has been observed that this theoretical issue does not affect the practicality of the method (Rosenbaum et al., 2019).

Relative to smoe, rl-r benefits from directly optimizing actions to improve the language modelling loss. However this absence of bias comes with complications, especially the high variance of the gradient (Rosenbaum et al., 2019; Denoyer and Gallinari, 2014). We use reinforce with a learned baseline (Williams, 1992; Sutton and Barto, 2018) to address this issue, so that improving the policy means increasing the likelihood of selecting experts which lead to a better than average next token prediction. As with smoe, we find it useful to add a balancing term. To our knowledge, we are the first to experiment routing with Reinforcement Learning on large Transformer-based language models—we therefore explore key ablations in Appendix B.3.

Scaling Behavior at Convergence

Our main hypothesis is that the converged log-loss of a Routing Network is bilinear in the terms $\log N$ and $\log\widehat{E}$, where $\hat{E}$ is a saturating transformation of $E$. Specifically, we fit the 6-parameter scaling law:

We can generalize this law across a wider range of routing architectures by a change of variables, using the model inference cost $F$ and the total number of parameters $P$, as:

where $B\triangleq\frac{P}{F}$ and $B\to\hat{B}$ is the same saturating transform as $E\to\hat{E}$. Before justifying Equation (1), we validate its candidacy by fitting it to empirical data obtained on a large sweep of models. This sweep consists of a Routing Network trained for each of the three techniques described in §3: across six model sizes (described in Table 4) while varying $E$ across $$. This totals 168 different models, including dense baselines.

The observed losses for each model are shown in Fig. 2(a-c). We fit Eq. (1) to each routing method and plot predictions for fixed values of $N$ as dotted lines. The goodness-of-fit across all methods is apparent, as is the clear behavior that increasing $E$ leads to a reduction in validation loss. Fig. 2(d) plots the relative predictions for all three techniques, clearly showing that s-base performs best across all model sizes, followed by rl-r, followed by hash (see §5.3). The remainder of this section justifies the chosen functional forms (1) and (2); first supposing independent power laws in $N$ and $E$ (§4.1), then introducing a multiplicative interaction (§4.2) and saturation in the second term (§4.3), followed by a change of variables (§4.4). The benefit gained by this progression of fits can be seen in Table 2. Notations are recalled in Fig. 3.

Kaplan et al. (2020) argue that the converged performance of a dense model with $N$ parameters can be modelled accurately as the two-parameter power law

where $\alpha_{N}\triangleq-a$ and $N_{c}\triangleq 10^{d/-a}$. We can re-estimate these coefficients from the performance of our own dense models, leading to estimations in Table 2. The similarity of $\alpha_{N}$ is a reassuring sanity check (there are differences in dataset, vocabulary, tokenization and model which effect $N_{c}$).

An immediate hypothesis is that for all values of $N$, scaling in $E$ obeys a similar power law:

Because $L_{N}(1)=L(N)$ (a fact we will call dense equivalence), (3) and (4) can be combined into:

corresponding to the multiplicative separated power law:

If Eq. (4) fits observed data for any $N$ we can proceed with an assumption that scaling in $E$ obeys a power-law for fixed $N$. Observing a constant $b$ across $N$ would allow to fit Eq. (5) to models ranging across $N$ and $E$ simultaneously.

The first hypothesis is easily tested and confirmed to a reasonable degree. We fit Eq. (4) for each technique and value of $N$ separately, plotted as colored lines in Fig. 4. The values of $b$ are shown in Fig. 3.

We observe that $b(N)$ is increasing with $N$ (values listed in Table 6), corresponding to a reduction in benefit from routing as size increases, with a slope that is approximately linear in $\log N$ (Fig. 3). Eq. (5) requires that $b$ remains fixed across $N$; therefore we expect it to poorly predict model performance. We can attempt a fit nevertheless: plotted in grey in Fig. 4. Qualitatively, this mis-predicts some validation losses by over 0.2, particularly overestimating the performance at large $N$ and $E$. As reported in Table 2, the fit has held-out RMSLE values greater than 80e-4.

2 Quadratic Interaction in N𝑁N and E𝐸E

This motivates us to introduce a simple extension: that of a multiplicative interaction between $\log N$ and $\log E$. This is conveniently the exact function which leads to $b$ scaling with $\log N$ and takes the following form:

This function has the property that the log-log slope in both $N$ and $E$ are affine in the logarithm of the other variable. In other words, with $E$ or $N$ fixed, the performance $L$ scales with $N$ or $E$ following (3) and (4) with slopes given by:

$b(N)$ matches the behavior reported in Table 6. A transposed table, fitting sets of models with fixed $E$ and changing $N$, can be found to match the behavior predicted by $a(E)$ (see Table 8). There are two symmetric non-logarithmic representations of (7), useful for comparison to (6):

Fitting the bilinear (7) instead of (5) substantially reduces the prediction error for large $N$ (Table 2, Eq. (5) vs Eq. (7)), as displayed in Fig. 4 (dotted lines match the dashed ones, where the grey separable fit doesn’t). We verify dense equivalence: $\alpha_{N}\approx a$, while $N_{c}\approx\exp(d/a)$, and thus the law (7) gives similar prediction to the reference law (3) for dense models. Predictions for fixed $N$ are visualized as grey lines in Fig. 2.

In Eq. (7), when $c$ is positive, the expert improvement slope $b(N)$ reduces with model size $N$. All three routing techniques considered therefore predict diminishing improvements from routing when increasing scale. However, the scaling of s-base is predicted (and seen) to be substantially better. When designing a new technique, we can fit (7) and predict a better scaling behavior if the fitted $c$ is lower than with other techniques. A clear goal for future work in routing techniques should be to find a method with scaling coefficient $c\approx 0$.

3 Bounded Scaling in E𝐸E

Equation (5) models scaling in $E$ as a power law. For both small and large values of $E$, there are reasons to expect some deviation. If a routing technique degrades with $E$ (for instance, the variance of gradients in rl-r will increase), performance for large $E$ might be worse than predicted. On the other hand, fixed overhead (e.g., interference from auxiliary losses) might worsen scaling for low values of $E$, counter-intuitively leading to better than expected performance. Both phenomena appear clearly in Fig. 2. We seek to model this saturation such that the limit behavior in $E$ is bounded on both sides. We choose the following transformation, but discuss in §5.1 a number of implications which are independent of the specific saturating form used:

This is constructed so that we have $\hat{E}(E_{\min})=E_{\text{start}}$, while $\hat{E}\to E_{\max}$ as $E\to\infty$. We fix $E_{\min}=1$, indicating the lower bound of meaningful expert counts. $\hat{E}$ can be seen as a thresholded version of $E$: increasing past $E_{\max}$ will give improvement, but not following a power law. Similarly, when $E_{\operatorname{start}}>1$, $\hat{E}>E$ for small values of $E$. Practically, the fit is the same over a wide range of different thresholding functions.

Solving Equation (1), equal to Eq. (7) with $E\to\hat{E}$, is complicated by its non-convexity. We find the coefficients $(a,b,c,d,E_{\operatorname{start}},E_{\max})$ as the best of repeated solutions provided by the L-BFGS-B algorithm (Byrd et al., 1995). Fig. 2 shows fitted curves from these equations; coefficients are reported in Table 3.

Relative to using the simple bilinear law (7), fitting Eq. (1) improves prediction for the lowest and highest values of $E$ considered. Crucially, while the deviation from a power-law (and therefore improvement in RMSLE) is relatively minor for the values of $E$ considered, the deviation is nonetheless clear (seen best looking at the raw losses in Fig. 21). We believe it is important to model this saturation because (as argued in §5.2) the limit behavior of model performance as $N$ increases is substantially different when bounded, with important properties that are independent of $E_{\max}$. We further hypothesize that future work, able to test still larger values of $E$, will see a more quantitative benefit from including these terms. This can be already observed in Fig. 20 when noting that the law (7) does not over and under estimate the performance for $E=\{2,4,256,512\}$ as it does in Fig. 4. Level curves of Eq. (1) enumerate the $\{(N,E)\}$ which are predicted to achieve fixed performance, as visualized in Fig 1(b). This demonstrates of the power of routing: a model with $N=5M$ and $E=128$ equals the performance of a model with $N=55M$ and $E=1$,which requires over ten times more compute per inference.

4 Generalizing Across Architecture Variants

The models trained so far use fixed choices for two key details of routing: the number of experts executed per-datapoint $K$ and the frequency of routed layers across depth $R$ (previously set at 1 and $0.5$, respectively). For any selected value of $K$ and $R$ we may fit Eq. (1) to observed performance, but since these variables are independent of $N$ and $E$, we do not expect the same coefficients to remain valid across values of $K$ and $R$. To allow for a unified scaling law, we modify Eq. (1) to use terms in $F$, the TeraFLOPs required per forward pass, and in the ratio $B\triangleq\frac{P}{F}$ where $P$ is the total number of parameters. Specifically, $F$ is motivated by the approximation from Kaplan et al. (2020) that $F=2N$. $B$, the parameter utilization ratio, is an affine function of $E$, close to linear when most parameters lie in the routed components of the model.

Using $(F,B)$ instead of $(N,E)$ (and setting $E_{\min}$ to $\frac{1}{2}$) results in Eq. (2). To show the advantage of this change of variables we conduct two experiments: varying $K$ across $\{1,2,4\}$ and $R$ across $\{0.25,0.5,1.0\}$. In both cases, we vary $E\in\{8,64,256\}$ and $N\in\{15M,370M,870M\}$.

Eq. (2) predicts the scaling behavior of models as well as Eq. (1) for a given routing architecture, as indicated in Fig. 24. The benefit of the change of variables is seen most clearly in Fig. 5, which plots contours of fixed loss value as functions of $(N,E)$ and of $(F,B)$. For varying $(K,R)$, the loss surface as a function of $N$ and $E$ changes: meaning a joint fit would be inaccurate. Plotted as functions of $(F,B)$, the loss surface is almost the same, suggesting a shared fit between all three methods (see Fig. 25 and Fig. 26 for joint fits for $K$ and $R$ respectively). We highlight that $R=0.25$ deviates slightly. Plausible explanations are discussed in §D.4. The possibility to use a shared fit indicates a singular takeaway: the architectural details $K$ and $R$ little affect the scaling behavior of a Routing Network. The loss of the network can thus be predicted based only on inference flops $F$ and total number of parameters $P$.

Scaling Law Applications

Next we provide two applications of the scaling laws presented. We re-emphasize that all values are only valid at the specific token count all models were trained at: 130B. App. F provides evidence that our analysis, if not the numerical values, are nevertheless robust to token count.

We leverage Eq. (1) to compute the size $\bar{N}$ of a dense model giving the same performance as a Routing Network. Specifically, we solve for $L(\bar{N},1)=L(N,E)$, yielding

Here $\alpha(E)=a+c\log E$. Given a model with $N$ and $E$, we call $\bar{N}$ that model’s Effective Parameter Count (or epc). Eq. (1) predicts that the performance of all models increases as a power law in this variable

The result of plotting all models as a function of $\bar{N}$ is shown in Fig. 1(c): a good fit across four orders of magnitude. Scaling in terms of $\bar{N}$ results in a unifying power law: valid for dense and routed language models alike.

2 Routing Behavior for Large N𝑁N

epc leads to a better grasp of the behavior of routing as $N$ increases. Of immediate interest is $N_{\textrm{cutoff}}$: the value of $N$ where $\bar{N}(N,E)\leqslant N$. For larger $N$, routing will not improve performance. This is easily found to obey $\log N_{\textrm{cutoff}}=\frac{b}{c}$. $N_{\textrm{cutoff}}$ equals $937\textrm{B}$, $85\textrm{B}$ and $83\textrm{B}$ for s-base, rl-r and hash respectively. These values are highly dependent on the number of tokens seen, and $N_{\textrm{cutoff}}$ is expected to increase with increased numbers of tokens.

Next we consider $\bar{N}_{\max}(N)\triangleq\max_{E}\bar{N}(N,E)$, i.e. the maximal effective parameter count that a routing network can reach. Eq. (11) predicts that $\log\bar{N}$ is an affine function of $\log N$ for any fixed $E$, and $\bar{N}_{\max}(N)=N$ for $N>N_{\textrm{cutoff}}$. Therefore $\log{\bar{N}}_{\max}$ is piecewise-affine in $\log N$, as displayed in Fig. 6:

Note that $\bar{N}_{\max}$ is continuous near $N_{\textrm{cutoff}}$, since for all $E$, $\bar{N}(N_{\textrm{cutoff}},E)=N_{\textrm{cutoff}}$. Moreover, the slope of $\bar{N}_{\max}(\cdot)$ for $N\leqslant N_{\textrm{cutoff}}$ is positive whenever $E_{\max}\leqslant E_{\operatorname{start}}10^{-a/c}$, which is true for our coefficients. In this setting $\bar{N}_{\max}(\cdot)$ is a non-decreasing function of $N$. Therefore for any routing network where $N<N_{\textrm{cutoff}}$, $N\leqslant\bar{N}_{\max}(N)\leqslant N_{\textrm{cutoff}}$, meaning routing will never let you train a model more powerful than $N_{\textrm{cutoff}}$. Note that despite this value not depending on $E_{\max}$, its existence crucially depends on the saturating transformation: without it ${\bar{N}}_{\max}$ is unbounded.

3 Comparative Analysis

Kaplan et al. (2020) use scaling laws to encapsulate and contrast the behavior of entire model classes. Here we mirror this analysis by using the scaling laws we have proposed to summarize the relative behavior of the three routing techniques considered. We make four concrete observations:

s-base consistently outperforms rl-r and hash, though rl-r is very competitive at smaller $N$.

All routing techniques suffer from reducing efficacy as $N$ increases. Amongst the three techniques, s-base scales best: the fitted parameter $c$ is lowest.

For small $N$, rl-r and s-base scale similarly with expert count and better than hash (as indicated by computing the effective expert slope $b(N)=b+c\log N$).

hash and rl-r maintain power-law behavior for longer than s-base (larger $E_{\max}$). However they suffer from more interference ($c$); leading to worse performance for most model sizes.

hash has large initial overhead (bigger $E_{\text{start}}$), clearly visible as a more obvious curvature at small $E$.

For a practitioner interested in applying routing techniques, we conclude with some recommendations:

Use routing when training any model with $N\leqslant\text{1.3B}$.

s-base is a good default routing algorithm. rl-r will sometimes match s-base in performance but is less robust and scalable (§D.1).

Target using $E\in\{64,128\}$ experts. Larger values will continue to improve, but with diminishing returns.

Use $K{=}1$ experts. Route layers at frequency $0.5\leqslant R\leqslant 1$; lower frequency reduces performance.

Future routing research should focus on the terms $c$ and $E_{\max}$; indicative of limits to arbitrary scaling.

New routing techniques must be validated at multiple values of $N$ and $E$ when comparing with prior work. Results on single sizes cannot be extrapolated.

Related Work

In studying the empirical aspects of scaling, this work follows Kaplan et al. (2020); which triggered much research including Henighan et al. (2020), Hernandez et al. (2021) and Ghorbani et al. (2021). The underlying theory is less understood, but there is some exploration of this space including Hutter (2021) and Bahri et al. (2021).

These studies, and ours, are mutually reliant on a large corpus of work improving the scalability of Transformers. This includes models like GPT-2 (Radford et al., 2019), GPT-3 (Brown et al., 2020), Jurassic-1 (Lieber et al., 2021) and Gopher (Rae et al., 2021), as well as work improving the ability of these models to be efficiently parallelized across multiple devices, including Shoeybi et al. (2019), Narayanan et al. (2019), Kim et al. (2021) and Xu et al. (2021).

Parallel to all this has been a long study of Routing Networks; a term introduced by Rosenbaum et al. (2018) but developed extensively in the literature as Conditional Computation (Bengio et al., 2013, 2016; Bengio, 2017; Denoyer and Gallinari, 2014) and Mixture of Experts (Jacobs et al., 1991; Collobert et al., 2003; Eigen et al., 2014). The framework is sometimes further generalized, seen as per-example architecture search in Ramachandran and Le (2018) or as a graph problem in Denoyer and Gallinari (2014). Routing was popularized for large scale training by Shazeer et al. (2017), and furthered by work including GShard (Lepikhin et al., 2020), Switch Transformer (Fedus et al., 2021) and GLaM (Du et al., 2021). In this vein, Artetxe et al. (2021) undertake a comparative analysis of dense networks and smoes with $E=512$ that aligns with our results. Finally, the core routing architecture is still being improved. Nie et al. (2021) adapt $K$ through training where Hazimeh et al. (2021) learn it via a differentiable loss. Ramachandran and Le (2018) increase $K$ through depth and encourage architectural diversity across experts. Caccia et al. (2021) grows $E$ throughout training and Rajbhandari et al. (2022) propose networks where $E$ changes with depth.

Conclusion

Using conditional computation to scale neural networks has long been a research goal, and methods based on Routing Networks have been increasing in popularity. Here we have introduced a scaling law (Eq. (1)) that models the behavior of these networks. This scaling law predicts that, for all models considered, introducing routing into a language model improves performance. That improvement follows a power-law in the number of experts $E$ that diminishes with model size $N$, and can be further generalized across routing architectures with Eq. (2). These scaling laws quantify the differences between three different routing techniques and lead to a single scalar (Eq. (11)) that simultaneously describes the performance of routed and dense models alike.

This work provides an empirical framework with which to analyze future innovations in routing. We hope the overwhelming evidence we provide towards the benefits of routing encourage it to be more rapidly adopted as a powerful tool for model improvement, whose scaling characteristics align with traditional methods of scaling (in depth and width) and which will remain beneficial up to models with base model size greater than 900 billion parameters.

Acknowledgments

We would like to thank Marc’Aurelio Ranzato, Nando de Freitas, Jacob Menick and Andy Brock for useful comments and feedback on early drafts of this paper. The infrastructure needed to train these models wouldn’t have been possible without the dedicated work of the JAX and XLA teams, especially Peter Hawkins, Roy Frostig and James Bradbury who all were crucial in the development of the routing software.

References

Appendix A Architecture

Our Transformer [Vaswani et al., 2017] is based on the architecture in [Radford et al., 2019] with relative positional encodings [Dai et al., 2019]. Text is tokenized via SentencePiece [Kudo and Richardson, 2018] with $32,000$ tokens and a byte-level backoff. We use Megatron-style FFW sharding [Shoeybi et al., 2019] where useful. Parameters are stored in bfloat16 but all optimizer statistics are kept in float32. As a result, the activations of the language models are calculated in bfloat16 (though we explicitly upcast to perform all operations involving a softmax, including the Attention Block and Router, in full float32 precision). This is crucial to maintain stability on larger models [Fedus et al., 2021, Rae et al., 2021]. The learning rate starts at 1e-7 and decays to 2e-5 with a cosine decay rate over the entire $250,000$ steps, after an initial warmup phase ramping up to 2e-4 in the first $1500$ steps.

We use seven different model sizes, with names and architectures specified in the following table. The width of the hidden layer $d_{\mathit{ffw}}$ is fixed at four times the width of the activations $d_{\mathit{model}}$, and we use the same dimension for keys and values.

The number of models we trained was too large to practically include multiple runs of each model with different seeds. To give an idea of the potential error introduced by random chance, we trained all three routing techniques with 3 different seeds on a 130M model for $100,000$ steps with $8$ and $256$ experts (along with a dense baseline). Results are shown in Fig. 7. Different seeds (which influence not only parameter initialization but Expert Parallelism – see Appendix C) lead to extremely minimal model divergence after an initial transitory period, with different seeds diverging by no more than $0.01$ before $100,000$ steps. This is a close match to the $0.02$ error mentioned in [Kaplan et al., 2020]Anecdotally, throughout the development of this work we used $0.02$ as the cutoff to denote statistical significance..

Appendix B Detailed Routing Techniques

Here we detail aspects of the routing techniques crucial to their implementation and provide comparisons to key alternatives.

We encourage uniform routing in both our SMoE and RL-R methods with the differentiable load balancing loss adapted from the mean square auxiliary loss in Shazeer et al. and introduced in Lepikhin et al. , Fedus et al. .

Where $m_{e}$ is the mean gate per expert:

And $g_{e}$ is the gating decision per expert:

For $x$ in batch $B$ of size $N$ and policy $p(x)=\operatorname{softmax}(W_{p}x+b_{p})$. There are two cases where the selected experts may not be the ones used: in s-base after the Sinkhorn redistribution step (see §B.2.1) and when experts are skipped due to load-balancing (see §C.2). In both cases, the balancing loss is applied to the original gating decisions made by the policy. We found that the auxiliary loss is less effective if post-balancing experts were considered.

B.2 SMoE with Sinkhorn redistribution (s-base)

Our implementation of s-base differs from that proposed in Lewis et al. in two ways. First, we replace the auction algorithm for re-assigning expect selections with a continuous rebalancing process implemented via a Sinkhorn algorithm [Cuturi, 2013, Peyré and Cuturi, 2019]. Second, we add a shuffling step, similar to Lewis et al. , before computing the optimal assignment via Sinkhorn per-device (as opposed to across all devices as done in Lewis et al. ). In addition, we did not use any input jitter on the activations sent to $\rho$ as we did not see a noticeable effect. This is in line with base but differs from recommendations in other smoe papers [Lepikhin et al., 2020, Fedus et al., 2021].

Mathematically, $\Pi$ is obtained by solving a simple problem with constraints: each input must, on average, prefer exactly one expert. This is made clear by the variational formulation of the softmax:

where $H$ is the Shannon entropy of the matrix $\Pi$, i.e. $H(\Pi)\triangleq\sum_{i=1}^{T}\sum_{j=1}^{E}p_{ij}\log p_{ij}$, and $[\cdot]$ denotes horizontal stacking. This variational formulation offers a natural alternative to incorporate extra constraints. For ideal performance, each expert should be assigned the same number of tokens on average $B=\frac{T}{E}$. We therefore add $E$ additional constraints:

which yields the doubly constrained regularized linear problem

that we recognize as the regularized Kantorovich problem of optimal transport [Kantorovitch, 1958, Cuturi, 2013].

These updates converge towards an optimal couple $(f,g)$, such that

is the solution to Eq. (19), where $(f\oplus g)_{ij}\triangleq f_{i}+g_{j}$ for all $i,j\in[T]\times[E]$. As detailed below, we early stop the iterations (24) by measuring the primal violation of constraints in $L_{1}$ norm, i.e. when

Once the plan is computed, we greedily select, for each token, the device with highest device-selection probability, effectively applying an $\operatorname*{argmax}$ operation on top of the Sinkhorn logits to form a transportation plan projection.

Compared to using an exact (early-stopped) auction algorithm as Lewis et al. , the complexity of the Sinkhorn algorithm is in $\mathcal{O}(N\times E)$ versus $\mathcal{O}((N\times E)^{3/2})$, and its update are well adapted to batch computations on TPU/GPU. In contrast, the auction algorithm must be run on CPU as it is a greedy per-coordinate algorithm; it becomes a computational bottleneck applied to models with many routed layers. Replacing the softmax output by an regularized optimal transport plan is very naturally interpreted as adding a balancing distribution constraint to the softmax operator. Using an auction algorithm on top of the softmax assignment does not have this property.

Moreover, the Sinkhorn algorithm can be halted before it has fully converged with a proper tolerance parameter (26) where Lewis et al. uses a hard number of iterations. We find an error tolerance of $e_{\textrm{tol}}=10^{-2}$ gives consistently good performance. In practice we observe an end-to-end model overhead of $1\%$ to $3\%$ compared to Switch (the same routing technique without this reassignment). This computational offset is negligible compared to the per-step performance gain. Without the rebalancing step, Switch is very sensitive to balancing loss hyperparameters (as noted in Lewis et al. ) whereas s-base maintains uniform routing decisions with improved performance and robustness while varying $E$ and $N$.

B.2.2 Shuffling Tokens

Similar to Lewis et al. , we shuffle router inputs across workers by first computing a random permutation of the inputs and sending the $t$th row of the batch to the $\lfloor\frac{tE}{T}\rfloor$th worker. We found that this shuffling stage was necessary to prevent training from becoming unstable at larger scales. Our hypothesis is that the re-assignment provides a subtle side channel through which information can be propagated backwards in time, and this can be abused by larger models resulting in the validation loss diverging during training. Adding a shuffling stage ameliorates this issue by introducing a large number of irrelevant elements to the rebalancing process, making it harder to infer behavior of future inputs. Further work is needed to confirm this theory, but the introduction of the shuffling step does eliminate this performance degradation.

B.3 Routing with Reinforcement Learning (rl-r)

We will first describe a naive REINFORCE [Williams, 1992] implementation of routing, then describe possible extensions and improvements which lead to the form used in the main text as rl-r.

Our implementation of REINFORCE uses the balancing loss in Equation 14 and a policy gradient loss:

Where $R_{i}$ is the reward for each sequence in the batch of size N and $\pi$ is the normalized expert preferences output by a linear transformation as in smoe. The proper thing is for $\rho$, the selected experts, to be samples from the distribution $\pi$, but we found that this substantially degraded performance at larger scales. This phenomenon can be attributed towards unwanted interference, where exploratory steps for $\rho$ which turn out to be unnecessary lead to bad gradient updates to the rest of the network [Rosenbaum et al., 2019]. We therefore consider a greedy selection method, where router outputs are selected as $\rho(x)=\textrm{TopK}(\textrm{softmax}(W_{p}x+b_{p}))$.

While sampling (even when tuning softmax temperature) decreased the performance of the model, we would nevertheless like to regain some of its exploratory power. To ameliorate this, we can use Nucleus Sampling [Holtzman et al., 2019], which samples from the top-$p$ set of experts $E^{(p)}$.

Where $E^{(p)}$ is the smallest set of experts such that:

This eliminates the possibility of selecting experts with very low likelihood, while still introducing some randomness. It is important to emphasize that this introduces a distributional shift to the samples, which can be corrected with off-policy correction methods such as Importance Sampling.

An alternative improvement is to learn an additional baseline function for each router. This method has an additional entropy regularization loss and computes advantages $A_{i}=R_{i}-b_{i}$ for the learned baseline $b_{i}$:

Where we use the Huber Loss to calculate the value loss $v_{i}$.

Greedy REINFORCE (RLR-G). REINFORCE selecting the top-$k$ experts and no additional auxiliary losses.

Nucleus-sampled REINFORCE (RLR-S). REINFORCE using nucleus sampling to eliminate less reliable expert selections and reduce noise in the policy gradient update. In this method we sample from the top-$p$ truncated distribution. Nucleus sampling at a fixed top-$p$ scales well with increasing the number of experts.

REINFORCE with baseline (RLR-B). Our RL method which stabilizes training with a learned baseline and a policy entropy regularization loss. We learn a baseline with a value function that has a single hidden layer of size $\frac{d_{\textit{model}}}{8}$.

Table 5 details the hyperparameters chosen for each rl-r variant and Fig. 8 contains validation losses across a number of models. Note that the entropy loss is negative to encourage a more concentrated policy, and the weight must be tuned jointly with the load balancing loss to keep routing balanced. This is in line with Bengio et al. , who also use two loss terms to both encourage early specialization and expert diversity. Additionally, since the policy entropy loss has a similar effect to nucleus sampling, we did not see an improvement from including both regularization methods. RLR-B consistently performed the best, especially with regards to scalability in $E$ and $N$. For that reason we selected it as our prime example, and refer to it as rl-r elsewhere.

B.4 Hash layers (hash)

hash is simple compared to rl-r or s-base, but is highly reliant on the particular choice of hashing function. Many functions rely on knowing the integer ID which the tokenizer assigns to each unique token (characters, bytes, subwords, etc.). Roller et al. describe multiple alternative functions, including pre-computing expert assignments for each token using a greedy assignment based on the frequency counts of the token on the training set. They do not observe any improvement in terms of perplexity relative to simpler random assignments of token to expert, but argue that balanced hashing has better properties for distributed training.

Our implementation uses a simple modular hashing function, namely the token index modulo the number of experts. Tokens are indexed by our tokenizer in an order that is roughly ordered by their underlying frequencies in the training dataset, which means this strategy will be more balanced than an arbitrarily random assignment, while simpler to implement than fully balanced hashing. We note that poor balancing with increasing expert count is to some extent inevitable for any routing technique that defines one-to-one mappings between tokens and experts, assuming a bounded Expert Capacity (see Section C.2), as it becomes progressively harder to assign high frequency tokens into a bigger number of smaller buckets due to the tokens’ heavy-tailed distribution. This can be seen in Fig. 9.

Appendix C Distributed Routing Details

Here we describe the key aspects of Routing relevant to training on large clusters. We note there are several libraries available for supporting large-scale Routing, including DeepSpeed [Kim et al., 2021, Rajbhandari et al., 2022] and GSPMD [Xu et al., 2021]. Unfortunately these were incompatible with our preexisting infrastructure.

We briefly review parallelism techniques, building up to Expert Parallelism, a technique for efficiently distributing parameters over an accelerator cluster. For a more in-depth exposition we recommend Lewis et al. , Lepikhin et al. or Rajbhandari et al. . In a fully data-parallel world, every device has an identical copy of all parameters $\Theta$ and a different input batch $X$. Each device executes a forward and backward pass on $X$ and (usually) does a synchronous all-reduce across all devices on the gradients to $\Theta$. This is effective, but requires one copy of $\Theta$ for each device, wasteful when $|\Theta|$ is large.

The general class of techniques known as Model Parallelism reduce this duplication by having any individual device store only a subset of the entire model parameters. This reduction in memory comes with a cost: no longer can a single device take an input and produce the model’s output; that device no longer contains all of $\Theta$. Most techniques therefore require some additional synchronization or data exchange.

Sharding Parallelism [Shoeybi et al., 2019] takes advantage of a mathematical property present both in 2-layer-MLPs and a Transformer’s attention blocks: namely, that the output can be represented as the sum of $N$ components, where each component applies the same functional form with independent weights on the same input. Shoeybi et al. contains more details, but a simplified example can be given for a matrix multiplication where we observe the effect of splitting a matrix into columnwise sub-matrices: $Wx=[W_{1},...,W_{N}]x=\sum_{i}^{N}W_{i}x$. The effect of applying this technique such that each device has a separate subcolumn is to prevent the duplication of the weight matrices (which consist of the vast majority of $\Theta$). The disadvantage is that all devices must see the same input, meaning the total throughput of data on the cluster has been reduced $N$-fold. In addition, the sum described above is actually now a sum across devices, which introduces additional communication overhead.

Expert Parallelism takes further advantage of the structure of a routed layer to similarly reduce the necessity of parameter duplication while avoiding the need to duplicate data between devices. In particular, rather than duplicating experts across all devices, each device contains only a subset of the experts which are not replicated anywhere else. Different devices still see different inputs. The key motivation is that a given input $x$ never needs to interact with the parameters corresponding to experts which the router did not send $x$ to. Therefore, a single input $x$ need only be present on a single device (the one which contains the experts which the router selected for $x$) to produce the correct output. In order to produce an output, the router selects an expert for all inputs and an additional data-exchange is introduced which sends all inputs to the device which contains the requested experts. Each device then processes the inputs it was sent, then returns all inputs to their original devices. Crucially, a roughly uniform router distribution leads to an evenly balanced computation across devices. This allows routed layers to be stored across a cluster with no duplicated data and without a reduction in data throughput. The downside is that this data exchange required across devices is generally more costly than the cross-device-sum required by sharding. More details are given in Lewis et al. . Previous work [Fedus et al., 2021] suggests using one expert per device. We believe this to be an implementation detail dependent on many aspects of the infrastructure in use. For us, typically using $4$ or $8$ local experts per device gave good performance.

All of Data, Sharding and Expert parallelism can be applied simultaneously. We use all three methods at will, selecting the combination which works fastest for a given cluster structure and model size. There are still more variations of model parallelism, notably Pipeline Parallelism [Narayanan et al., 2019, Huang et al., 2019], which we do not use.

C.2 Load Balancing

This at-will changing of parallelism techniques is dependent on the parallelism not affecting the output of the model. This is generally true, but expert parallelism brings in one complicating factor: load balancing. In the description above, we emphasized that a roughly-uniform router (averaged over a minibatch) will send the same number of inputs to each device (we will call the expected value $BS_{\text{avg}}$). However, in the worst case all inputs on all devices might select the same expert, and therefore need to be sent to a single device. If memory is pre-allocated to accommodate this worse case, then each device must have enough free memory to potentially store the entire global batch size: prohibitive for large clusters.

The most common solution is to specify a capacity factor $C$, and only allocate space for $BS_{\text{avg}}\times C$ tokens. When an expert is oversubscribed tokens are dropped at random until no experts are exceeding capacity. Having $C>1$ is useful during training to prevent unnecessarily large numbers of tokens from being dropped. We set $C=2$ for all experiments (though during evaluation we always allow all tokens to be routed to the desired expert). This strategy works well for the Transformer architecture due to its residual connections – dropping a token means skipping that transformer block. As long as the amount of dropped tokens is kept at a reasonable bound, it does not impact learning.

Appendix D Architectural Variations

Throughout this work we have focused on a narrow subset of possible Routing Net architectures, which we believe are representative of recent work on large scale Routing Nets [Roller et al., 2021, Fedus et al., 2021, Lewis et al., 2021, Shazeer et al., 2017, Artetxe et al., 2021, Lepikhin et al., 2020]. However, we also experimented with many variations of these architectures, some of which we highlight now in more depth.

We evaluated the robustness of s-base and rl-r to changes in hyperparameters in Fig. 10. We focus on $E=512$ due to anecdotal experience that the largest performance variance occurred at this scale. rl-r is found to be highly sensitive to the hyperparameters in Table 5, especially the choice of balancing weight. In addition, changes to the policy entropy weight can lead to unbalanced routers when the balancing weight is not tuned jointly.

D.2 Varying Routing Frequencies

All of our models thus far have been routed every other layer with experts which are single FFWs [Lepikhin et al., 2020, Fedus et al., 2021]. However, Lewis et al. , Roller et al. explored stacking FFWs in the experts and placing $N$ routed layers at $\frac{L}{N+1}...\frac{NL}{N+1}$. We consider the performance impact of alternative routing frequencies, varying the frequency $R=\frac{N}{L}$ and placing routed layers at $\frac{L}{N}...\frac{NL}{N}$.

We compare routing every layer to routing at frequencies $R\in\{\frac{1}{2},\frac{1}{4},\frac{1}{L}\}$. For routing a single layer we chose the second to last layer [Roller et al., 2021], but consider routing at $\frac{L}{2}$ in subsection D.4. s-base scales well with routing frequency, but hash degrades in performance as shown in Fig. 11. At a single routed layer, hash has the lowest validation loss across model sizes.

D.3 Varying the Routing Policy

Motivated by the improved scaling results for s-base, we investigate whether learning a routing policy becomes more beneficial as the frequency of routers increases.

In Fig. 12, the routing decisions are made at the first routed layer and shared across layers, which keeps the number of routers constant as $R$ increases. As hash selects experts based on the token index at the input layer, its routing function is unchanged for this variant. s-base and hash have similar losses for shared routing decisions, whereas s-base improves when learning to route at each expert layer.

Conversely, we tested a variant of hash where the hash function at each router uses a static permutation of the input tokens to select the experts. This allows tokens to be routed to the same expert at some layers without having the same hash. We found that performance was unchanged for this variant, suggesting that increasing the number of possible routing paths does not necessarily impact performance for static policies.

These router variants suggest that methods which can adapt to each expert layer will outperform static policies. Further work is needed in analyzing how policies can more effectively learn to route across layers.

D.4 Routing a Single Layer

We analyzed the scaling behavior of hash and s-base when only routing a single layer. We observed that the routing gains for $R=\frac{1}{L}$ deviated from higher frequencies, which also impacted $R=\frac{1}{4}$ to a lesser degree. We attribute this performance regression to the suboptimal behavior of the first routed layer. In both cases the total number of routers is low, and the first layer has a larger impact on overall performance than at higher routing frequencies. For $R=\frac{1}{L}$ the complexity of routing is reduced and a simpler routing method can reach competitive performance. hash and s-base have similar performance across expert counts in this case, as shown in Fig. 13.

We also compared routing a single layer at $\frac{L}{2}$ with three FFWs per expert to three evenly spaced routed layers in Fig. 13. Similar to the results shown in [Roller et al., 2021], three evenly spaced routed layers has slightly better performance than three stacked FFWs for a 32E 1.3B model. We also found that s-base benefits more from interleaving the routed and dense layers, which is consistent with our routing frequency results.

D.5 Varying number of experts per datapoint

In this work we have focused on routing each datapoint to a single expert at all routing layers, i.e. for the case where $K=1$. However, SMoE models have historically routed datapoints to more than one expert [Shazeer et al., 2017, Lepikhin et al., 2020, Ramachandran and Le, 2018]. Increasing $K$ incurs in extra computation on the experts, but this additional computation may be helpful for the end result, reflecting in better loss. Moreover, routing a datapoint through more experts means each expert gets to see more data for each forward pass, which may speed up training. For these reasons, it is not obvious that $K=1$ is the best setup. Section 4.4 investigated this and argued both that the generalized formula Equation (2) can accommodate such cases and also that the resulting fits show no substantial difference in performance for $K$. However we explore this variance more in Fig. 14: plotting both scaling curves for varying values of $K$ as well as plotting the loss in terms of $F$. Higher values of $K$ invariably yield better performance per step, but they are not necessarily more flop efficient. In fact, $K=1$ is always in the pareto front. We can verify that this holds for varying numbers of experts.

Note that this difference in flop-efficiency is not only theoretical, and is also followed by increased communication costs when using expert parallelism. We observed in practice that reducing K by half amounted to close to 2x speedup in inference and training.

Appendix E Effects of scaling strategy on Zero-shot Transfer

There is a strong relationship between the validation loss we have been discussing and the downstream performance of models and specific tasks [Kaplan et al., 2020]. However, recent work has shown that this relationship is not as straightforward for large Routing Networks, and individual tasks can benefit more or less from expert scaling. For example, Artetxe et al. show a narrowing performance gap between a smoe Routing Network with $E=512$ and its dense equivalent, with more marked improvement from routing in some tasks like HellaSwag and PIQA than in in tasks like Winogrande and ReCoRD. Likewise, Fedus et al. shows that Switch benefits more from scale better in TrivaQA than in SuperGlue.

A detailed analysis of the scaling properties of Routing Networks and how that transfers to downstream tasks merits dedicated work. Here we start the conversation by looking at zero-shot transfer on a set of well known downstream tasks: LAMBADA [Paperno et al., 2016], The Pile [Gao et al., 2020], Curation Corpus [Curation, 2020], WikiText-103 [Merity et al., 2016] and C4 [Raffel et al., 2020].

We estimate the scaling coefficients individually for each task and routing technique. For simplicity of interpretation we ignore the bounded scaling term and focus on the bilinear fit on Eq. 7. The coefficients can be seen in Table 9. We expect that scaling in both $N$ and $E$ will improve the downstream performance. The key question revolves around understanding changes in the relative magnitude of $a$, $b$ and $c$ as we move from task to task.

Viewing Table 9 it is immediately clear that the individual scaling coefficients vary greatly across tasks, i.e. different tasks have different relative gains at Zero-Shot performance as we move to larger scales. This can be better shown in Fig. 15, where all coefficients are displayed in a single plot. The variation across tasks are not the same for $a$ and $b$. e.g. WikiText-103 has higher values for $b$ and lower for $a$ when compared to the validation set. This means that even though tasks see monotonic improvement in performance by scaling through either adding more experts or increasing the base model size, some tasks benefit more and some less from which method is used.

For a more complete picture, we can account for the $N$ and $E$ interaction coefficient $c$ by incorporating it into one of the scaling coefficients – by holding the other quantity fixed – which leads to $a(E)$ and $b(N)$ (see Section 4.2). This can be seen in Fig. 16 for varying values of $N$ and $E$.

We see that s-base tends to dominate with lower coefficients at higher values of $E$ and $N$ (due to its smaller interaction term relative to scaling terms), but this varies across tasks. For example, rl-r shows better $b(N)$ for most values of $N$ in LAMBADA, until it is overtaken by s-base at $N=410M$, but s-base is always superior in C4. Moreover, the ordering is not consistent between hash and rl-r across tasks, even though they often do not cross. This all means it is difficult to establish superior performance of a routing technique without looking at a variety of tasks and scales.

We often want to compare Routing Networks with a dense baseline with the same performance on the validation set, and see how this changes on downstream tasks. We can use these parameters in a simplified version of the Effective Parameter Count (EPC, Equation 11), by assuming $E_{start}=1$ and $E_{max}=\infty$, such that $\hat{E}=E$. First, we note that since the coefficients vary greatly across tasks, each task will have a different EPC for the same network configuration. Moreover, the effects of scaling by varying $E$ and $N$ will vary across tasks. Say we have a routing net of size $N$ with $E$ experts and we want to increase its base model size by a factor of $x$ while keeping the same number of experts. The effect on $\bar{N}$ in this case will be a multiplication by $x^{a(E)/a(1)}=x^{1+c\log E}$. Since $c$ varies per task, the improvement achieved by increasing the base model size will also be task dependent.

Say we have a routing net of size $N$ with $E$ experts and we want to increase its base model size by a factor of $x$ while keeping the same number of experts. The effect on $\bar{N}$ in this case will be a multiplication by $x^{a(E)/a(1)}=x^{1+\frac{c}{a}\log E}$. Since $\frac{c}{a}$ varies per task, the improvement achieved by increasing the base model size will also be task dependent.

For example, the epc validation for N=110M, E=32 is 370M, but epc lambada for the same model is 284M, while epc pile is 535M. The key implication here is not only do the values change, but their slopes are different. This means that downstream tasks must be analyzed carefully: a practitioner could scale a model via routing expecting some overarching improvement, but get a much diminished (or enhanced!) improvement on specific downstream tasks, depending on their specific values of $b$ and $c$.

Appendix F On Convergence, or Lack Thereof

Here we digress on two important details, both focusing on token count. First we argue that discussing converged performance of large transformers on modern and massive text datasets is probably a misnomer; scaling analyses should focus on optimal performance at a fixed number of tokens. Second, we provide evidence arguing against a proposed equation in Kaplan et al. (Eq. (1.6)).

There are two cases where the converged performance of a model can be clearly defined. The first is when continued training of the model produces no improved results (even analyzed at logarithmic scale in the number of tokens), the second is when continued training leads to reduced validation performance: overfitting.

Our models exhibit neither behavior. No overfitting is seen even for our largest models, likely due to the complexity and size of the dataset used. Furthermore, despite being trained for 130 billion tokens, not even our smallest models have saturated. We push this envelope even further: training two additional sets of 15M models with $1$, $64$, $128$ and $512$ experts. The first set is trained for just $75,000$ steps, and the second for $1,000,000$ steps: four times more data (half a trillion tokens). We highlight that this involves corresponding changes to the cosine cycle decay. We exclusively train hash models, both due to limits in the number of extra models we were able to train and also because it has the largest value of $E_{\max}$.

Results from these models are plotted in Fig. 17 (left). 15M with no routing, the smallest model we train as part of this work, is still far from having saturated its performance capability. Indeed, training for 4x longer further reduces the validation loss by $0.05$. This pattern continues, and is exacerbated, when increasing the expert count: the same model with $E=512$ gets a $0.07$ reduction in loss from 4x more tokens.

It is clear then that the very smallest models considered have yet to converge. The same is certainly true for larger ones, and probably more so. If 500 billion tokens is a lower bound to the convergence point of 15M, the analysis in Kaplan et al. would predict needing trillions of tokens to converge 1.3B: much more than what was used to train some of the largest language models yet created [Brown et al., 2020]. For large, complex text datasets of the scale used to train large language models, convergence is not a proper criteria.

F.2 Performance Qualified on Token Count

Rather than claiming analysis at a non-observed point of convergence, we emphasize that the scaling behavior we have described in this work is valid only as a function of a particular number of steps (or tokens). At each point, we can define instantaneous values of scaling coefficients, with the values from all models taken at $S$ stepsThis sidesteps the issue of critical batch size [McCandlish et al., 2018, Kaplan et al., 2020], consideration of which requires a substantially larger sweep of models. Future work estimating the critical batch size will likely lead to better model fits. .

In fact, the situation is more complicated that simply conditioning our scaling coefficients on token count. We can see this by plotting $b$, the scaling coefficient for changes in expert-count in Fig. 17(right). An immediate observation is that the values of $b$ are non-constant, supporting the need to qualify scaling on token count. A second, more substantial point, is that these values are not uniquely defined by token count. For a given number of tokens, the scaling behavior of three different sets of models is completely different, dependent on how far into the learning rate schedule those sets of models were. We note that this behavior is suggested by experiments in Kaplan et al. (App. D.6).

Attempting to find the full set of parameters on which these scaling terms depend is beyond the scope of this work. We highlight just the importance of insuring that all variables possible are matched when comparing values to calculate scaling coefficients.

F.3 Performance Modeled as L(N, S)

We conclude by highlighting one implication of the fact that scaling coefficients are dependent on token count. We analyze only the dense models trained as part of this work, and calculate values of $a$ in Equation (3) for all dense models trained as part of the primary sweep across all step counts; plotted in Fig. 18(a) with RMSE values plotted in Fig. 18(b). First, it is important to emphasize that the fits remain good throughout $S$ (after an initial period of transience). Namely, though the slope is different, the validation losses for a given intermediate $S$ follow a power law about as well as they do later in training (if anything, more so). Second, the estimated coefficients $a$ are clearly monotonically increasing with $S$.

[Kaplan et al., 2020] propose (Eq. 1.6) a unified prediction of the loss achieved by a model with size $N$ training for $S$ steps:

This comes with the subtlety that $S$ must be defined as the number of steps when training at the averaged critical batch size, where our models are trained with a fixed batch size. This means a proper analysis must use $S_{\mathit{min}}$ with $S=S(1+\frac{B_{c}}{BL^{-\alpha_{B}}})^{-1}$ for constants $B_{c}$ and $\alpha_{B}$. It is important to highlight however that $S_{\mathit{min}}$, as described in Kaplan et al. , should be independent of $N$. This implies that $\frac{\partial}{\partial N}\left(L(N,S)\right)$ is independent of S, or in log-log space:

This prediction of constant scale is in concrete opposition to the increasing value seen in Fig. 18(a). We can furthermore check that this functional form cannot be obviously fit to our learning curves, with examples show in Fig. 18(c).

There are subtle differences between training setups, and we do not want to claim our experiments wholly disprove the conjecture in [Kaplan et al., 2020]. However, the results in Fig. 18 motivate us to assume that Eq. (32) cannot be used to model our specific training curves. A consequence of this is that we can also no longer conclude Equation B.5 from [Kaplan et al., 2020], that:

With this equation, we might be able to lower-bound true converged performance (which we have not seen in our models) by inference from compute-efficient performance, which has been achieved by the majority of our models.

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

Our analysis predicts that larger values of $E$ will continue to improve performance, especially for small models, at a diminishing rate. §5.2 also predicts that routing will continue to help with increasing $N$ for at least one, if not two orders of magnitude larger base model size. Practical compute limitations prevented our sweep from exploring these regimes, and there are interesting unanswered questions in the limit of these two variables. In particular, exact predictions of $N_{\textrm{cutoff}}$ are highly dependent on the precise value of $b$, where error in the second decimal place shifts predicted values by orders of magnitude (not surprising, as it is the slope of a line in log-log space).

We believe exploring the limit behavior of $N$ and $E$, especially arriving at a more precise value of $b$, is crucial. Anecdotally, we can report the results of one experiment: a large rl-r model with $N>7,000,000$, providing a rough upper bound for error in $b$ for rl-r. In particular, we trained a model with $d_{model}=4096$, $n_{layers}=32$, $n_{heads}=32,E=64$ and K/V size of 128. There are some important eccentricities of this model which affect its match to the fits described in this work: it was trained with a batch size of 1024 for 100k steps with a policy gradient weight of 1e-1 and balancing weight of 1e-1. Other training details are consistent with Section 2.1.

The performance of this model, relative to a dense model of the same size and also to a number of smaller models, is plotted in Fig. 19 evaluated at 100B tokens. The changes described above prevent the analysis in this work from accurately predicting this model’s performance, but one key feature remains: the routed 7B model substantially outperforms the baseline. This is of particular interest since just a $0.01$ decrease in $b$ would predict an $N_{\textrm{cutoff}}$ at $9B$, meaning we would already be close to the regime where routing would cease to work. Nevertheless, at this value routing is clearly still a major improvement, and our estimate of $b$ is unlikely to be a substantial overshoot.

While the differences between this model and those analyzed in the paper make concrete extrapolation impossible, it shows that routing techniques still maintain competitive improvements at almost an order of magnitude larger value of $N$ than analyzed and it is unlikely the scaling coefficients measured in this work substantially overestimate the routing technique’s scalability. We encourage future work probing the limits of routing networks, both in $N$ and $E$, to better understand their properties and provide more accurate predictions of their scaling coefficients.

Appendix H Extra Plots and Tables

This section contains some helpful visualizations and data which are not included in the main text.