No Need to Talk: Asynchronous Mixture of Language Models
Authors: Anastasiia Filippova
Angelos Katharopoulos
David Grangier
Ronan Collobert
Venue: ICLR 2025 Spotlight (2025)
Keywords: language models
distributed learning
divide and conquer
efficient inference
Primary area: foundation or frontier models, including LLMs

Abstract: We introduce SMALLTALK LM, an innovative method for training a mixture of language models in an almost asynchronous manner. Each
model of the mixture specializes in distinct parts of the data distribution, without the need of high-bandwidth communication between the nodes training each model. At inference, a lightweight router directs a given sequence to a single expert, according to a short prefix. This inference scheme naturally uses a fraction of the parameters from the overall mixture model. Unlike prior works on asynchronous LLM training, our routing method does not rely on full corpus clustering or access to metadata, making it more suitable for real-world applications.  Our experiments on language modeling demonstrate that SMALLTALK LM achieves significantly lower perplexity than dense model baselines for the same total training FLOPs and an almost identical inference cost. Finally, in our downstream evaluations we outperform the dense baseline on 75% of the tasks.

Published as a conference paper at ICLR 2025
No NEED TO TALK: ASYNCHRONOUS MIXTURE OF LANGUAGE MODELS
Anastasiia Filippova EPFL
Angelos Katharopoulos Apple
David Grangier Apple
Ronan Collobert Apple
ABSTRACT
We introduce SMALLTALK LM, an innovative method for training a mixture language models in almost asynchronous manner. Each model of the mix- ture specializes in distinct parts of the data distribution, without the need of high- bandwidth communication between the nodes training each model. At inference, a lightweight router directs a given sequence to a single expert, according to a short prefix. This inference scheme naturally uses fraction of the parameters from the overall mixture model. Unlike prior works on asynchronous LLM train- ing, our routing method does not rely on full corpus clustering or access meta- data, making it more suitable for real-world applications. Our experiments on lan- guage modeling demonstrate that SMALLTALK LM achieves significantly lower perplexity than dense model baselines for the same total training FLOPs and an almost identical inference cost. Finally, our downstream evaluations we outper- form the dense baseline on 75% of the tasks.
INTRODUCTION
Recent research has demonstrated that scaling large language models (LLMs) by increasing model capacity and expanding training data consistently leads to significant performance improvements on a wide range of downstream tasks (Kaplan et al., 2020; Brown et al., 2020; Henighan et al., 2020; Hoffmann et al., 2022; Dubey et al., 2024). Scaling introduces substantial operating and engineering costs for both inference and training. In general, training is achieved on large number of nodes via synchronous gradient descent techniques, which relies on high-bandwe Inference of large models may require multiple compute nodes to distribute the model, which relies on low-latency communication. In both cases, state-of- of-the-art interconnect hardware is critical, and careful engineering is required to scale, and maintain a large compute cluster. While mainstream machine learning frameworks have eased the engineering work the scaling side, access large number of well interconnected nodes remains a privilege in the machine learning community.
communication to scale.
In this paper, explore strategies mitigate the communication cost of large language models, both at training and inference, while keeping the inference efficient. We show that efficient training and inference can be achieved without relying on fast interconnects, and without compromising model performance, both in terms of perplexity or downstream task accuracy.
In recent studies aimed mitigating reliance on binwihhh interconnects, researchers have developed algorithms that reduce the need for frequent or comprehensive gradient synchronizations. Such techniques include asynchronous training (Douillard et al., 2023: Aji & Heafield, 2017: Zhang et al., 2015; Liu et al., 2024) and gradient compression (Lin et al., 2017; Dettmers, 2016) which effectively decrease communication overhead. By performing updates less frequently or by commu- nicating less data, these methods sustain high training throughput while diminishing dependence on high-speed interconnects. However, these algorithms still require some level of gradient synchro- nization, and resulting models often under perform (in terms of perplexity) compared to training approaches synchronizing at every step (Diskin et al., 2021).
Regarding efficient inference, number of sparse parameter activation techniques have recently become popular (Shazeer et al., 2017; Fedus et al., 2022; Artetxe et al., 2022; Lewis et al., 2021; Du al., 2022), in partieular the Switch Transformer mixture of experts (MoE). These approaches
Work done while interning Apple.

---

Published as a conference paper at ICLR 2025
effective at reducing significantly the number of active model parameters at inference time. Yet, they require routing decisions to be made for each token, demanding rapid interconnections and requiring all parameters to be accessible in RAM.
To achieve both low-bandwidth training and efficient inference, several prior works have explored asynchronously training mixtures of models (Gross al., 2017; Li et al., 2022; Gururangan al., 2023). In these methods, the input space is clustered using unsupervised techniques like K-Means, with each model in the mixture trained on assigned cluster. These approaches reduce communica- tion overhead and have demonstrated improvements in model perplexity, while maintaining similar inference costs to baseline models (Gross et al., 2017: Gururangan al., 2023). However, in the context of language processing, their practical applicability on standard downstream tasks with rela- tively short input prefixes remains subject of dehate, as they rely on full corpus clustering or access to metadata (Li et al., 2022; Gururangan et al., 2023).
In this paper, we introduce SMALLTALK LMM, an asynchronous mixture of language models, which combines the advantage of asynchronous LM training methods (significantly reducing interconnect requirements), and sparse activation methods (it naturally uses only a subset the overall architec- ture parameters during inference). SMALLTALK LM achieves better perplexity, and better accuracy on a majority of downstream tasks, compared to a regular dense language model trained with the same amount of FLOPs. Our main contributions
follows:
We introduce an algorithm that trains a mixture of independent language models, signif- icantly reducing bandwidth requirements during training compared to regular distributed training methods. Our approach leverages a lightweight router (which accounts for less than 1.5% of the size of each expert) to efficiently route sequences. This router determines the most suitable expert for given sequence based on short prefix (see 2).
We empirically demonstrate that our method achieves significantly lower perplexity than a dense baseline for near-identical training and inference FLOPs and identical training data volume (see Fig. 2). Furthermore, we show that for a constant mixture model size, the improvement in perplexity increases with the number of experts.
of downstream tasks, showing that it achieves better or similar accuracy compared to dense baseline on 75% the tasks (see 3.3 and App. B).
We evaluate method
2
METHOD
In this section, we formalize our proposed method: SMALLTALK LM. In $ 2.1, we introduce general language modeling background, and present a mixture of experts framework in this context. Sub- sequently, in 2.2, show how different experts in the mixture. We also present the training procedure and implementation details of method. Finally, we discuss its marginal computational overhead and minimal distributed communication costs.
use independent language models to perform routing to
2.1 BACKGROUND
Language modeling Let V denote a fixed vocabulary, and consider a sequence of tokens x =
(11, 22, 2g), where each F, E for s 1,2..... Language modeling is about modeling the data distribution p(x), which in practice is cast as training a language model p(r,+1 X1:> 0) generating a token 2441 given a context sequence X1:s (I1, .e..,,....,.2 ....). Typically, the model parameterized by 0 is neural network, and parameters are obtained by minimizing the negative log-likelihood:
8-1
[(x;0) = = --xx; 0) = log p( X1:s ; 0)-
Σ
Mixture of Experts There exist many variants of the mixture of experts model (Jacobs et al., 1991; Jordan & Jacobs, 1994; Tresp, 2000; Collobert et al., 2001: Shahbaba & Neal, 2009). Following the original formulation and applying it to language modeling, we factorize the conditional distribution for predicting given token by summing over of F. experts:
2

---

Published as a conference paper at ICLR 2025
[(x:
log p(as+1 X1:s; 0*) p(e xi:s; 0*),
where each expert p( X1:s; C 0*) is intended focus different parts of the data distribu- tion, and the router p(e X1:s; ; (*)) assigns weights to the expert predictions. Each expert has its of independent parameters 8 and the router is parameterized by 6".
While this type of mixture of experts partitions the input space across different models, all experts still involved when computing the conditional likelihood in Equation (2), both training and inference time. Hard mixtures of experts circumvent this issue by hard-assigning a sequence X1:s to single expert ex, minimizing an upper bound Equation (2):
C(x: 0) S
Σ log p( X1:a,e;; 0*) p(e" Xi:s (*),
The way ex is chosen depends of the goal of the approach. For example, in Jacobs et al. (1991), e* is sampled at training time, according the distribution p(e X1:s i 0*). In other works, like Collobert et al. (2001), it is chosen as the most relevant expert (that e* = arg max, p(e x1:s; (")), both at training and inference. Regarding the training procedure, mixture models with both soft and hard assignments are typically trained with an Expectation-Mininizatie (EM) scheme, alternating between experts and router optimization.
2.2 SMALLTALK LM
Routing With Independent Language Models SMALLTALK LM can be viewed as an instance of hard mixture of experts, where a sequence is routed the most relevant expert e*, according to the router posterior p(e X1:s 0"), both at training and inference. However, instead of being monolithic neural network model, the router implemented language model per expert. More precisely, with Bayes' rule we have:
arg max p(e
; 0*)
(4)
p(xi: e; p(e)
arg
Σ E
Di=1 p(x i;; 0r) p(i)
(5)
arg max p( 11: e; 0")
(6)
arg max p(11= @me).
where the likelihood P(X1: (re)) defines routing language model attributed to the e-th expert with independent parameters gr and expert priors p(e) are supposed to be uniform. So far, in our definition, there is nothing that prevents 6" = 8" namely to use the experts themselves to implement the router. Interestingly, in that case selecting e as above is intuitive because it amounts to selecting the best expert for given sequence (in terms of log-likelihood).
Deriving a Practical Model In order for the above formulation to result model that useful in practice, we have solve the following main challenges. Firstly, need to be able to route using a small prefix of the full sequence since in most real world applications we do not have access to a do not have the answer). Secondly, using the experts
full sequence (e.g. in question answering to implement the router results in a large computational overhead since we need to evaluate every single one of them, largely negating the benefits of the hard mixture of experts.
We solve both of the above problems with the following two key ideas:
1. We use a short prefix of length M compute the score for each expert as follows:
p(e" X1:s) R p(e" xI.M).
This enables to SMALLTALK LM in worl-world applications where only very small prefix may be available. We show in $ 3.4 that our method outperforms the dense baseline even when using little as 32 tokens for routing.

---

Published as a conference paper at ICLR 2025
2. We implement the router with E language models that are orders of magnitude smaller than the experts. This results in marginal computational overhead, as little as 1% during training and 3% during inference ($ 3.2).
Using the log-likelihood of a prefix to route a sequence is an intuitive solution since a model that good at predicting the beginning of sequence is likely to be good at the full sequence. Another interpretation is that the prefix can be used to perform some sort of thematic identification of the whole sequence that then used route to the appropriate expert.
On the other hand, the efficacy of smaller language models as routers is less intuitive. Even though one can expect large model to be able to perform well on sequences that are beyond the capabilities of smaller models, the task of routing boils down to identifying whether sequence is similar to the training data or not. We show that even really small models (4.4M parameters) can perform that task very well and result in no discernible difference in performance compared to larger routers ($ 3.4).
Algorithm SMALLTALK LM training.
Train the routers
2: X <- N new sequences from the dataset
3: XI.E = raandaaaaasipppp saridomnsssinmeetss/
4: for i = 1 do
5:
for = 1... E do
6:
gr arg mingr.e L(X,; 0*:))
Optimize Equation 9 with SGD for the e-th router
7: end for
8: X N new sequences from the dataset
9:
X = balanced b
assignments(X,, { 0*)
# Segment the data according to Equation 4
end for
11: Train the experts
12: X < new sequences from the dataset comprising the total number of training tokens
X1.E balanced. (x.areeeeesssigmmeessX(.., 0")
14: for == 1... do
l arg ming- ((
# Optimize Equation with SGD for the e-th expert
16: end for
Training Procedure We train our proposed hard mixture of experts with EM. Namely we alter- nate between step of optimizing equation and step of assignments according equation 4. In framework, assignments e* purposefully do not depend on the experts, and the routers independent language models. This allows us to split the training in two stages:
1. We train the routers using EM without using the experts at all, namely we minimize
M
C(x,0 log p(c"
XI:M; 0*)
Σ
log (Fe41|Xi Orie")
(9)
in the likelihood maximization step and perform the assignments using equation 4 in the expectation step.
2. We train the experts as E independent language models on disjoint dataset segments se- lected by the trained routers.
This procedure is described in pseudocode in Algorithm 1. In the specific case where the experts and the routers are the same language models (0" fr our training algorithm comes down regular EM algorithm where the sequences are routed to the best performing expert. In that light, we can think of each router as lightweight approximation of the corresponding expert. See $ 3.4 for experiments where the experts themselves are used for routing.
Balancing the Assignments to Experts A common hurdle to overcome when training mixture of experts models is to ensure that the modeling load is shared evenly as possible across all experts, or in other words that the model utilizing the capacity of all experts efficiently instead relying on few good ones. Common ways of tackling this problem are either introducing losses that penalize
4

---

Published as a conference paper at ICLR 2025
scores matrix
assignments
scores matrix
sorted matrix assignments
1.1
0.9
1
1.1
0.9
1
0.1
0.4
0.8
0.1
0.4
0.8
0.1
0.4
0.8
0.8
0.2
0.7
K
0.8
0.2
0.7
0.8
0.2
0.7
1.1
0.9
(a) Naive assignments
(b) Assignment with respect to the minimum in chunk.
Figure 1: Balanced assignments. In this scenario we have 3 experts (columns with different colors) and 3 sequences to assign (rows) under the constraint that each expert should get 1 sequence. In (a) we assign sequentially each row, by negative log-likelihood, leading to sub-optimal assignment because the first expert is full when we try assign the last row. On the other hand, in (b), we first sort wrt to the minimum log- ligliikli which results in optimal assignments.
uneven assignments or introducing noise into the assignments namely by sometimes performing "sub-optimal"" assignments avoid local-minima. These techniques are used in all mixture models since their introduction in Jacobs et al. (1991) until more recently in Fedus et al. (2022).
In our setup, we can solve this problem by ensuring that each expert is assigned equally sized distinct portions of the dataset. However, challenge arises when assigning data sequentially based solely such cases, an expert might already have more than its fair share of the dataset assigned yet we continue to find sequences for which it has the lowest loss among all experts, as illustrated Figure la. To overcome this issue, during training we perform the assignments considering the whole set of sequences rather than one sequence at a time. In particular, we sort the sequences based on max, log log e; 8") and perform the assignments ensures that sequences with high likelihood will be assigned first and avoids the edge-case where a sequence with low-loss cannot be assigned to the appropriate expert because it is at capacity. In Algorithm we refer this procedure as balanced assignments and it is demonstrated in Figure 1b.
on Equation
that order. This
During inference, no balancing is performed, and the expert is selected solely based on equation 4.
Computational and Communication Cost It important to note that the computational cost of training and inference with the routers is significantly lower than that of training and inference with the experts. This difference arises due to the routers' much smaller size, smaller batch sizes, and the fewer tokens they train (see Appendix A.3 for further details). As result, the additional computational overhead from training and inference with the routers is negligible.
Regarding communication requirements for distributed implementation, in Algorithm all oper- ations except for the balanced assignments are completely independent and embarrassingly paral- lelizable without any communication. Namely, each expert or each router can be trained on own on a different node or set of nodes. The assignments, however, require that every node has access to the scores of every router in order to perform the assignment according to equation 4. This requires a small amount of communication as each node needs to share 1 score per sequence which results 2MB per 1M sequences if we represent the scores 16-bit floats. A detailed analysis is provided in the Appendix A.4.
EXPERIMENTS
In this section, we experimentally analyze the performance of our proposed method. First, in 3.2 we benchmark our approach against a dense LLM trained using regular distributed training tech- niques. Second, in 3.3, we evaluate our method on downstream tasks, demonstrating that achieves performance comparable to baseline models at the same level of perplexity, improving upon a dense model given the total training FLOPs spent. Finally, in $ 3.4, we analyse the core component of method, the routing. We show that the performance of the mixture does not de- pend significantly on the size of the router. Furthermore, we show that our approach outperforms less complex routing methods, such as the TF-IDF encoding with K-Means clustering proposed by
5

---

Published as a conference paper at ICLR 2025
Gururangan et al. (2023). Moreover, demonstrate that the perplexity gains not driven by few dominant experts but result from each expert contributing positively to the overall performance.
3.1 EXPERIMENTAL SETUP
We use the RedPajama- V2 dataset (Computer, 2023), which is large-scale collection of text data designed for training language models. The dataset is built from the ground up based on publicly available web data, consisting of 84 crawls provided by Common Crawl (2024).
Both the experts and routers utilize a Transformer architecture (Vaswani, 2017) with rotary posi- tional encoding (Su al., 2024). All experiments use sequences of 1, 024 tokens. We train the models using the Adam W optimizer (Loshchilov & Hutter, 2017) with 31 = 0.9, B2 = 0.99, and weight decay of 0.1. Gradient clipping is applied with a maximum norm of 0.1.
For the experts, we employ a learning rate schedule featuring a linear warm-up to 5 x 10 over the first 3, 3,000 steps, followed by cosine decay for the remainder of the training period. We experiment with two model sizes: 335M-parameter model with 4, 8, 16, and 32 experts and -ara-marre model with 4, 16, and 32 experts. Tables and describe the architectures and the training parameters in more detail.
For the routers, we use models with 4.4M parameters. The routers are trained for 128, 000 steps using a constant learning rate of x 10-4, following a linear warm-up over the first 1, 000 steps. Scheduler choice described in details in App. A. Routers are trained with batch size of 32 and perform all-to-all communication of the loss on the dataset chunk approximately every 45 million training tokens (see App. A.4). The length the prefix used for routing is set to 256 tokens, which is 25% of the context length.
Comparison to the Dense Model To ensure a fair comparison to the dense baseline, we designed our experiments such that the computational cost during inference and the total computational cost during training are the same. We achieve this by comparing the performance of models where each expert has the same architecture and number of parameters the dense model, resulting in the same inference cost. In addition, we train SMALLTALK LM for the same number of total tokens which means we spend the same number of total training FL.OPs.
3.2 LANGUAGE MODELING RESULTS
Across all experimental configurations, our method consistently oulperforms the buseline while uti- lizing the same volume of data, with only minor increases in computational costs. Specifically, for the 1.3B parameter model, approach results in less than 1% increase in training FLOPs and less than a 3% increase in inference FLOPs, while achieving performance gains of up to almost 18% in perplexity. Similarly, for the 335M parameter model, the cost increases than a 4% increase in training FLOPs and less than a 10% increase in inference FLOPs, alongside performance gains of up nearly 14% in perplexity (see Appendix A.3 for details on FLOPs calcu- lation). Figures 2a and 2b illustrate the perplexity on a held-out test set relative to the total training cost, measured in petaFLOPs, across different model sizes. Additionally, Figure 2c shows the test perplexity as function of the number of training tokens for the 1.3B parameter model.
modest, with less
Moreover, is interesting to compare our model with the smaller 335M parameter experts to the 1.3B dense baseline. When using 32 experts and roughly the same training budget in FLOPs, SMALLTALK LMM achieves slightly better perplexity (9.07 vs 9.11) but requires three times less compute during inference.
Finally, we observed that for a constant model size, the improvement perplexity increases with the number of experts. This suggests that our method more efficiently utilizes the training data to improve performance, likely due the specialization of experts on different segments of the data distribution.
Our approach performs minimal communication and does not re- quire fast interconnect. As detailed in 2, first train the routers which are then used to effi- ciently shard the dataset and each expert is trained completely independently. During their training, the routers communicate approximately 100 times with each transmission involving less than
Communication Overhead

---

Published as a conference paper at ICLR 2025
11
Respinalle
39L4d ddee tilee, SCeeprite
200
0.5M
0.MM
1.2M
2:4M
2,31
334
1034
2003
Vriting out (PELO2S)
Trasing (PELOPS)
(a) FLOPs vs perplexity, 335M.
(b) FLOPsvs perplexity, 1.3B.
(c) Tokens vs perplexity, 1.3B.
Figure 2: Better perplexity for the same price. Test perplexity comparison between our approach and the dense baseline, as a function of training cost measured in PFLOPs. In (a), we report results for models with 335M parameters using 4, 8, 16, and 32 experts and in (b) for models with 1.3B parameters using 4, 16, and 32 experts. addition, (c) shows the perplexity comparison between approach and the dense baseline, plotted against the cumulative number of tokens processed throughout the training for the 1.3B parameter models. We observe that our method significantly out- performs the baseline across all experimental configurations. Notably, our 335M parameter model with 32 experts achieves a perplexity of 9.07, outperforming the 1.3B dense baseline's perplexity of 9.1. This improvement is achieved with training budget of 2.5 x 1021 FLOPs, which is compara- ble to the baseline's s 2.2 x 10 FLOPs, while requiring three times less computational cost during inference (0.87 1022 FLOPs compared 2.81 10 FLOPs). See $ 3.2 and App. A for detailed description of our experimental setup.
Loce
0.33
031
cmae:
0 dean
Amis 2011
0 dmnn S0n
Depada
(a) ARC Challenge
(b) ARC Easy
(c) HellaSwag
(d) MMLU
Figure 3: Downstream evaluation. Accuracy with respect to perplexity on (a) ARC Challenge,
(b) ARC Easy, (c) HellaSwag and (d) MMLU, for 1.3B parameter dense baselines trained on 266B, 1T and 2T tokens (empty symbols) and mixture models with 1.3B parameter experts and 4, 16 and 32 experts respectively (filled symbols). The models that have the same symbol shape have near identical training and inference FLOPs.
megabytes per router resulting in truly minimal communication requirements. See A.4 for details about the calculation of these numbers.
3.3 DOWNSTREAM EVALUATION
To demonstrate the applicability of our approach real-world scenarios, we performed zero-shot evaluation on set of natural language processing (NLP) tasks, including ARC Challenge and ARC Easy (Clark et al., 2018), HellaSwag (Zellers et al., 2019), SciQ (Welbl et al., 2017), and MMLU (Hendrycks et al.. 2020). Figure 3 shows both accuracy and perplexity for each task (see Appendix B for details regarding perplexity and accuracy computation).
Specifically, our largest configuration, model with 32 1.3B parameter experts, demonstrates better performance on four out of five tasks, achieving gains of 3%. 2%. 3%, and 1% on ARC Challenge (3a), ARC Easy (3b), HellaSwag (3c), and MMLU (3d), respectively, under the same inference cost.
Moreover, on the MMLU benchmark, our approach either beats the baseline or achieves the same performance on 75% of the tasks (42 out of 56), shown in Table 4.
7

---

Published as a conference paper at ICLR 2025
23.3
Deylerdy
Poglestee
13-
1800
12-
denan 11F
13/
300
2015
64
125
506
32
0.1
125
355
Tokena
Preftc keggth
Fiwfic length
(a) Router size
(b) Prefix length
(c) Ours TF-IDF
Figure 4: Routing analysis. (a) Test perplexity over training steps for different router sizes using a 335M parameter model with 4 experts. We compare routers of sizes 335M (where the model routes data for itself), I 110M, 65M, and 4.4M parameters. (b) Test perplexity as function of routing prefix length during inference for 1.3B parameter model with 4, 16 and 32 experts. We examine how reducing the prefix length M used during inference affects performance when the data partitioned during training using a prefix size 2 M. (c) Test perplexity over training steps for a 335M parameter model with 16 experts, comparing our proposed routing using TF-IDF document encoding followed by SVD projection and balanced K-Means clustering.
3.4 ROUTING ANALYSIS
In this section we analyse the most critical component of SMALLTALK LM, the routing. In particu- lar, we perform experiments to investigate the impact of the size of the router to the mixture model, the impact of the size of the prefix and finally whether the EM algorithm presented in $ 2 is critical or it could be replaced by simple content based clustering.
Impact of Router Size
One of the most critical findings of our work is the fact that the size of the model used for routing does not impact the performance the mixture (Figure 4a). This allows us to utilize small router models which reduce dramatically the computational cost of routing during training and inference without compromising the model's performance.
In detail, we investigated the effect of router size on the performance of our method by experimenting with a 335M parameter model configured with experts, using four different router sizes for data partitioning: 335M parameter router (where the model routes data for itself, optimizing for the data handles best), and routers with 110M, 64M, and 4.4M parameters. All models were trained with batch size of 32 for 256, 000 steps, using routing prefix length of 256 tokens, following the same experimental setup as $ 3.2 (see Appendix A for more details about model architectures and training parameters). Figure 4a illustrates the test perplexity wrt the number of training tokens for these configurations. We observe that all router sizes perform practically identically.
Impact of Prefix Length An important hyperparameter for our mixture model is the length of the context used for routing to the expert model. In our experiments, we use a prefix length of 256 which may be too large for some real-world applications like conversational AI. However, we show Figure 4b that SMALLTALK LM outperforms the dense baseline even when routing with shorter show that with prefixes of just 32 tokens (2 th of the length during
prefixes. In particular, training) the model still outperforms the dense baseline.
Comparison With Routing Using TF-IDF An alternative approach might suggest that if tiny LLM suffices for routing, simpler methods such as TF-IDF encoding combined with balanced K- Means clustering, esis, we trained a 335M parameter model using the routing strategy outlined in (Gururangan et al., 2023): applying TF-IDF transformation to the text, followed by Singular Value Decomposition (SVD) projection into a low-dimensional space, and then clustering using balanced K-Means.
proposed by Gururangan et al. (2023), could be adequate. To test this hypoth-
We then trained 16 experts on these clustered results. Our findings reveal that routing with TF- IDF encoding under performs when the prefix is short, indicating limitations of this simple routing approach. As shown in Figure 4c, proposed routing method significantly outperforms the TF-
8

---

Published as a conference paper at ICLR 2025
IDF-based routing. This demonstrates that approach is more effective in capturing the nuances of the data distribution, even when using limited context for routing decisions.
Experts Do Specialize A potential concern is that the observed performance gains might be driven primarily by a few experts among the E, while the remaining contribute minimally or may even un- der perform on their assigned data. To address this concern, we conducted an analysis of each expert's performance individually. Figures 5a, 5b, and 5c compare the perplexity achieved by our method and the dense baseline the routed dataset segments for the 1.3B parameter model, utiliz- ing mixtures of 4, 16, and 32 experts, respectively.
Our findings demonstrate that all experts contribute positively, with consistent perplexity improve- ments observed across all routed dataset segments. Moreover, the performance gap between our method and the baseline widens the number of experts increases, suggesting enhanced special- ization due to finer partitioning of the data distribution.
Notably, despite the absence of capacity constraints on experts during inference, unlike the enforced capacities during training, all experts are actively utilized and receive substantial portions of the data. The consistent improvements across all routed dataset segments confirm that the experts do specialize and they all contribute to the overall improvement instead of only few dominant experts.
Rapert
(a) 4 experts vs dense, 266B
(b) experts vs dense, 1T
(c) 32 experts vs dense, 2T
Figure 5: Experts Do specialize. Test perplexity comparison between our method and the dense baseline on the routed dataset segments for the 1.3B parameter model, using mixtures of experts in (a) (trained 266B tokens), 16 experts in (b) (IT tokens), and 32 experts in (c) (2T tokens). Each bar represents a dataset segment, with the color intensity indicating the percentage of data assigned to that expert darker shades correspond to higher proportion of data. Overlapping bars depict the perplexity achieved by the dense baseline (translucent) and our proposed mixture model (opaque). The results demonstrate that all experts specialize effectively on their assigned segments of the data distribution, leading to consistent improvements over the baseline. While the data distribution among experts is not perfectly even with some experts receiving more data than others all experts receive a substantial portion of the data. This shows that each expert contributes meaningfully to the overall performance gains.
4 RELATED WORK
Data-parallel training requires synchronization of the gradients after each backward pass, which can be slow if the network bandwidth is limited. In this section, discuss existing methods that aim to make data parallelism more communication-cfiiee tion. Since our paper also explores methods for efficient inference without performance degradation, we discuss this branch research as well.
via reducing the frequency of the synchroniza-
Distributed optimization To reduce the frequency of synchronisation during training of large models, researchers have explored distributed optimization techniques both theoretically and empir- ically (Medonald et al., 2009; Stich, 2019; Lian et al., 2018: Wang Joshi, 2019; McMahan al., 2017: Lin et al., 2018; Zinkevich et al., 2010). The core principle of these methods is allow each worker to execute several local training iterations before engaging in global synchronization - each device performs multiple optimizer updates per communication round. This technique also has been studied for LLM pre- pre-training (Diskin et al., 2021; Douillard et al., 2023) and fine-tuning (Borzunov et al., 2022; Liu et al., 2024; Hilmkil et al., 2021; Ro et al., 2022).
While theses approaches reduces the frequency of communication, thereby decreasing communica- tion overhead, they often result in degraded model performance compared to every step synchroniza-

---

Published as a conference paper at ICLR 2025
tion (Diskin et al., 2021). or they lack direct performance comparisons in the context of pre-training (Douillard et al., 2023). The degradation in performance is often attributed to the staleness of gra- dients and the divergence between local and global models due to delayed synchronization (Zhang et al., 2016; Yu et al., 2019; Mitra et al., 2021). Moreover, to our knowledge, all experiments have small scale (w 400M parameters) (Douillard
been conducted Ro et al., 2022). Therefore, scalable solutions that can reduce communication frequency for LLM training without compromising performance larger scales are still needed.
al., 2023; Liu al., 2024;
Mixture of Experts Mixture of Experts (MoE) models have emerged prominent approach to increase model capacity without a proportional increase in inference costs (Jacobs et al., 1991; Jordan & Jacobs, 1994). The key idea behind MoE is partition the input space into subsets corresponding to specific subtasks, each managed by a specialized expert. By training different experts on distinct tasks, MoE models aim enhance performance while reducing computational overhead during inference.
Early implementations of MoE employed algorithms such as Gaussian Processes (Tresp. 2000; Theis & Bethge, 2015; Deisenroth & Ng. 2015), Dirichlet Processes (Shahbaba & Neal, 2009), and Sup- port Vector Machines (SVMs) (Collobert al., 2001) to model these specialized functions. Recent advancements have integrated deep neural networks as experts, providing scalable and flexible im- plementations suitable for large-scale models.
In natural language processing, researchers have generalized the original MoE idea to enhance large- scale language models. For instance, Shazeer et al. (2017) introduced a Mixture of Experts feed forward layer in which, in contrast to prior works, routing decision is made for every token. This layer takes a token representation as an input and routes it to the best-matched top-k experts. By doing the number of active parameters is decreased. Similar approaches have been effectively employed in models like the Switch Transformer (Fedus et al., 2022), GShard (Lepikhin et al., 2021), GLaM (Du et al., 2022) and Mixtral of Experts (Jiang et al., 2024) which leverage MoE layers to achieve state-of-the-arr results with reduced number of active parameters costs. However, since the routing decision in models like Switch Transformer MoE is made on the token level: (1) all experts must be retained in RAM during both training and inference phases, necessitating substantial memory resources, and (2) a high communication overhead is required for data transfer and gradient synchronization across experts (Lepikhin et al., 2021; Fedus et al., 2022; Riquelme et al., 2021). In that respect this approach is completely orthogonal to our proposed SMALLTALK LM. In fact, each of the experts in our method could be such a token-level mixture model which would result in even smaller inference FLOPs for a specific achieved perplexity albeit with a significantly more complex implementation and higher RAM requirements during inference.
Independent Training of Mixture of Experts Models To train experts fully in parallel, recent research has also explored independent training of the mixture of experts models (Gross et al., 2017; Li et al., 2022: Gururangan et al., 2023). For instance, Gross et al. (2017) proposed to train gater independently from the experts and use the output of the gater to hard assign the input image to experts. In language processing, Li et al. (2022) and Gururangan et al. (2023) proposed methods that train expert models separately and merge them post-training. However, these approaches rely on full document context or metadata for routing, which can be impractical in real-tine applica- tions involving shorter inputs. For instance, Gururangan et al. (2023) focus on performance in text classification tasks, which may not fully reflect real-world scenarios.
5 CONCLUSION
In this work, we present SMALLTALK LM - method that significantly reduces communication overhead during L.M training- also outperforms dense model baselines trained with a similar amount of FLOPs and the same data volume and comparable inference costs. In addition, our proposed lightweight routing with short prefixes makes SMALLTALK LM practical for real-world downstream tasks, where show it largely outperforms the dense baselines.
These findings open up multiple directions for future research, including exploring different routing strategies for data partitioning among experts and investigating further the impact on downstream task performance. Another promising avenue is scaling the mixture to hundreds or even thousands of experts, which could further enhance model performance while benefiting from sparse inference.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Jagrid Digani and Awni Hannun for their insightful discussions on the algorithm and experiments. We thank Justin Deschenaux for his contributions to the early stages of the project.
REFERENCES
Alham Fikri Aji and Kenneth Heafield. Sparse communication for distributed gradient descent. In Martha Palmer, Rebecca Hwa, and Sebastian Riedel (eds.), Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pp. 440-445. Association for Computa- tional Linguistics, September 2017.
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Out, Sam Shleifer, Xi Victo- ria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, Giridharan Anantharaman, Xian Li, Shuohui Chen, Halil Akin, Mandeep Baines, Louis Martin, Xing Zhou, Punit Singh Koura, Brian O'Horo, Jeffrey Wang, Luke Zettlemoyer, Mona Diab, Zornitsa Kozareva, and Veselin Stoyanov. Efficient large scale language modeling with mixtures of experts. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 11699-111 117-1.17... Association for Computational Linguistics, December 2022.
Alexander Borzunov, Dmitry Baranchuk, Tim Dettmers, Max Ryabinin, Younes Belkada, Artem Chumachenko, Pavel Samygin, and Colin Raffel. Petals: Collaborative inference and fine-tuning of large models. Association for Computational Linguistics, 2022.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. 33: 13:1877-1 2020.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. 2018.
arXiv preprint
Ronan Collobert, Samy Bengio, and Yoshua Bengio. A parallel mixture syms for very large scale problems. Advances in Neural Information Processing Systems, 14, 2001.
Common Crawl. Common crawl: Open web data, 2024. Accessed: 2024-09-15.
Together Computer. Redpajama: An open source recipe to reproduce Ilama training dataset, 2023.
Mare Deisenroth and Jun Wei Ng. Distributed gaussian processes. In International conference on machine learning, pp. 1481-1490. PMLR, 2015.
Tim Dettmers. 8-bit approximations for parallelism in deep learning. International Conference Learning Representations, 2016.
Michael Diskin, Alexey Bukhtiyarov, Max Ryabinin, Lucile Saulnier, Anton Sinitsin, Dmitry Popov, Dmitry V Pyrkin, Maxim Kashirin, Alexander Borzunov, Albert Villanova del Moral, et al. Dis- tributed deep learning in open collaborations, Advances in Neural Information Processing Sys- tems, 34:7879-7897, 2021.
Arthur Douillard, Qixuan Feng. Andrei A Rusu, Rachita Chhaparia, Yani Donchev, Adhiguna Kuncoro, Mare Aurelio Ranzato, Arthur Szlam. and Jajuu Shen. Diloco: Distributed low- communication training of language models. arXiv preprint arXiv:2311.08100 2023.
11

---

Published as a conference paper at ICLR 2025
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-expere experts.. In International Conference Machine Learning, pp. 5547- 5569. PMLR, 2022.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1 2022.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite. Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024.
Sam Gross, Mare Aurelio Ranzato, and Arthur Szlam. Hard mixtures of experts for large scale weakly supervised vision. In Proceedings of the IEEE Conference Computer Vision and Pat- tern Recognition, pp. 6865-6876 2017.
Suchin Gururangan, Margaret Li, Mike Lewis, Weijia Shi, Tim Althoff, Noah A Smith, and Luke Zetllmmyer. Scaling expert language models with unsupervised domain discovery. arXiv preprint arXiv:2303.14417 2022.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt, Measuring massive multitask language understanding. arXiv preprint arXiv:2009.088 03300, 2020.
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, et al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010.144: 2020.
Agrin Hilmkil, Sebastian Callh, Matteo Barbieri, Leon René Sütfeld, Edvin Listo Zec, and Olof Mogren. Scaling federated learning for fine-tuning of large language models. In International Conference on Applications of Natural Language to Information Systems, pp. 15-23. Springer, 2021.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
Robert A Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Albert Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bam- ford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024.
Michael I Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2):181-2) 1994.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, 2020.
Taku Kudo and John Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Eduardo Blanco and Wei Lu (eds.), Pro- ceedings of the 2018 Conference Empirical Methods in Natural Language Processing: System Demonstrations, pp. 66-71, Brussels, Belgium, November 2018. Association for Computational Linguisties.
12

---

Published as a conference paper at ICLR 2025
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. In International Conference Learning Representations, 2021.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning. pp. 6265-6274, PMLR, 2021.
Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A Smith, and Luke Zenlemoyer. Branch-train-mergg:: Embarrassingly parallel training of expert language models. arXiv preprini arXiv:2208.08: 03306, 2022.
Xiangru Lian, Ce Huang. Yijun Li, and Ji Liu. Asynchronous decentralized parallel stochastic gradient descent. In Proceedings of the 35th International Conference Muchine Learning, pp. 3043-3052. PMLR, 2018.
Yujun Lin, Song Han, Huizi Mao, Yu Wang. and William J Dally. Deep gradient compression: Re- ducing the communication bandwidth for distributed training. arXiv preprint arXiv: 2017.
Yujun Lin, Song Han, Huizi Mao, Yu Wang, and William J Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training. International Conference on Learning Representations, 2018.
Bo Liu, Rachita Chhaparia, Arthur Douillard, Satyen Kale, Andrei Rusu, Jiajun Shen, Arthur Szlam, and Mare' Aurelio Ranzato. Asynchronous local-sgd training for language modeling. arXiv preprint 23X1v2240..01135, 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- ence on Learning Representations, 2017.
Ryan Medonald, Mehryar Mohri, Nathan Silberman, Dan Walker, and Gideon Mann. Efficient large-scale distributed training of conditional maximum entropy models. Advances in neural information processing systems, 22, 2009.
H. Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effiienn learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, pp. 1273-1282, PMLR, 2017.
Praneeth Mitra, Idan Bistritz, Ji Wang. Dan Alistarh, and Wotao Yin. Linear convergence in fed- erated learning: Tackling client heterogeneity and asynchronous updates. In Advances in Neural Information Processing Systems, volume 34, pp. 14606-14619, 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Pitch Patarasuk and Xin Yuan. Bandwidth optimal all-reduce algorithms for clusters of workstations. Journal of Parallel and Distributed Computing, 69(2): 17-124, 2009.
Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In The Tenth International Conference Learning Represen- tations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34:8583-8595, 2021.
Jac Hun Ro, Theresa Breiner, Lara McConnaughey, Mingqing Chen, Ananda Theertha Suresh, Shankar Kumar, and Rajiv Mathews. Scaling language model size in cross-device federated learn- ing. arXiv preprint ar 2rXiv:2204.097/1 2022.
13

---

Published as a conference paper at ICLR 2025
Babak Shahbaba and Radford Neal. Nonlinear models using dirichlet process mixtures. Journal of Machine Learning Research, 10(8), 2009.
Noam Shazeer, * Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017.
Sebastian U Stich. Local SGD converges fast and communicates little. In International Conference on Learning Representations, 2019.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputiing, 568: 167066, 2024.
Rajeev Thakur, Rolf Rabenseifner, and William Gropp. Optimization of collective communication operations in mpich. The International Journal of High Performance Computing Applications, 19 (1):49-66, 2005.
Lucas Theis and Matthias Bethge. Generative image modeling using spatial Istms. Advances in neural information processing systems, 28, 2015.
Volker Tresp. Mixtures of gaussian processes. Advances in neural information processing systems, 13, 2000.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Jianyu Wang and Gauri Joshi. Adaptive communication strategies to achieve the best error-runtime trade-off in local-update sgd. Proceedings of Machine Learning and Systems, 1:212 229, 2019.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science ques- tions. In Leon Derczynski, Wei Xu, Alan Ritter, and Tim Baldwin (eds.). Proceedings of the 3rd Workshop on Noisy User-generated Text, pp. 94-106. Association for Computational Linguistics, September 2017.
Hao Yu, Sen Yang, and Shenghuo Zhu. Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 5693-5700, 2019. doi: 10.1609/ aaai.v33i01 v33001.33015663,
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Annual Meeting of the Association for Computational Linguistics, 2019.
Sixin Zhang, Anna E Choromanska, and Yann LeCun. Deep learning with elastic averaging sgd. Advances in neural information processing systems, 28, 2015.
Wei Zhang, Suyog Gupta, Xiangru Lian, and Ji Liu. Staleness- aware async-sgd for distributed deep learning. Proceedings of the 25th International Joint Conference on Artificial Intelligence, pp- 2350-2356, 2016.
Martin Zinkevich, Markus Weimer, Lihong Li, and Alex Smola. Parallelized stochastic gradient descent. Advances in neural information processing systems, 23, 2010.
14

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
2
Method
2.1
Background
2.2
SMALLTALK LM
3
Experiments
3.1
Experimental Setup
3.2
Language Modeling Results
3.3
Downstream evaluation
3.4
Routing analysis
Related Work
5
Conclusion
10
Language modeling experiments
15
A.I
Architecture and training details
15
A.2 Notation
16
A.3 Computational Cost
17
A.3.1 Transformer.
17
A.3.2 Mixture of LLMs
19
A.4 Communication overhead
20
Downstream evaluation
22
C Additional Results
23
A LANGUAGE MODELING EXPERIMENTS
A.1 ARCHITECTURE AND TRAINING DETAILS
Our architecture is based on the transformer decoder (Vaswani, 2017). We use rotary positional encoding (Su et al., 2024). We use SentencePiece (Kudo & Richardson, 2018) tokenizer with vocabulary of 32000 tokens. The model specific parameters for the different sizes are summarized in Table We implemented our EM scheme for training the routers PyTorch (Paszke et al., 2019). After segmenting the training set, the experts were trained independently using Jax (Bradbury et al., 2018). All training was done in bf cat1 6. The optimizer states and operations are in fl float32. The training parameters for specific models are summarised in Table 2.
We employ different learning rate schedules for routers and experts due to their distinct roles. Ex- perts are trained for accurate next-token prediction, using a cosine decay learning rate to optimize their performance. In contrast, routers are used solely for routing purposes. We only need an esti- mation of how well router performs are concerned with relative performance rather than absolute accuracy, applying the same training strategy all routers is sufficient for effective routing decisions. Moreover, using constant learn-
prefix of sequence relative to other routers. Since
15

---

Published as a conference paper at ICLR 2025
ing rate for routers reduces computational overhead and eliminates the need to tune the number of training steps.
Model parameters
Role
Hidden size
FFW expansion factor
Layers
Attention heads
335M
expert
1024
4
24
16
1.3B
expert
2048
24
16
4.4M
router
96
12
12
64M
router
416
4
12
12
110M
router
768
12
12
Table 1: Model parameters.
Model
Steps (k)
Tokens (B)
Batch size
# GPUs
335M (dense)
256
133
512
335M (4 experts)
256
133
128
335M (dense)
512
266
512
32
335M (8 experts)
256
266
128
8
335M (dense)
1024
532
512
32
335M (16 experts)
256
532
128
1.3B (dense)
512
266
512
32
1.3B experts)
512
266
128
1.3B (dense)
1024
1000
1024
1.3B (16 experts)
512
1000
128
8
1.3B (dense)
1024
2000
2048
128
1.3B (32 experts)
512
2000
128
4.4M (router)
128
4
32
Table 2: Training parameters.
A.2 NOTATION
For the dense baseline LLM and expert LLM:
Number of layers (L)
Hidden dimension size (H)
Number of attention heads (A)
Sequence length (S)
Batch size (B)
Vocabulary size (V)
Feedforward network dimension (D#)
For the router LLM:
Prefix length for routing (M)
Number of layers (L,)
Hidden dimension size (fI,)
Number of attention heads (A,)
Batch size (B,)
Feedforward network dimension (Drr) parameters specific to the router model
16

---

Published as a conference paper at ICLR 2025
Number of training tokens per router between communications T
We also define Sttepe Expept and router as the number of training steps for the dense base- line, the expert the mixture, and the router, respectively. The sequence length (S) and vocabulary size (V) are consistent across the dense bascline, router, and expert models. Finally, we denote E as the number of experts (routers).
A.3 COMPUTATIONAL COST
A.3.1 TRANSFORMER
To estimate the computational cost of Transformer model during training and inference, we calcu- lated the total floating point operations (FLOPs) based on the model's architectural parameters.
For training, we computed the FLOPs per training step by summing the contributions from the embedding layer, multi-head attention (MHA), feedforward layers (FFN), and the output projection layer. The calculations are follows:
Although embeddings are typically implemented as lookups with neg- ligible computational cost, for completeness, we estimate the FLOPs as:
Embedding Layer:
Femp = BxSxE
Multi-Head Attention (MHA):
1. Linear Projections (Queries, Keys, Values):
FLOPSproj 2 S X
Scaled Dot-Product Attention:
FLOPSatn FLOPSox - -FLOPsV 2x 2 x
3. Output Projection:
FLOPSeat.pm Bx H
Total Multi-Head Attention (MHA):
FL.OPsMIA FL.OPs proj
FLOP Sat
FLOPS Sout Sont_ proj
B H2
B X H2
=8 B x 2 +4x 4 H
Feedforward Network (FFN):
FLOPSFIN 2 x
Bx
DT = 4 x Dr
Output projection:
FLOPS FLOP Soulyny
FLOP Syottma
2 B H X X
Total FLOPs per Layer:
Flaer (FLOPSMHA + FLOPSF L
Total Forward Pass FLOPs per Step:
FLOPStorward FLOPSemb FLOPs + FLOPsout
Total Training FLOPs per Step:
The backward pass is approximately twice as computationally intensive as the forward pass. There- fore:
muinstu = Sroward x FL.PP Sooward
FI.OPSforwart
Total Training FLOPs:
17

---

Published as a conference paper at ICLR 2025
Multiplying the FLOPs per training step by the total number of training steps:
Total Training FLOPs = F-OP strpin.sepp X Nsteps = 3 X FLOP Sfowwadd X N
steps
Total Training FLOPs:
Total Training FI.OPs Naeps x BxSxH
+ W
+>xBxSxHHVV3xXBKSKK (10)
Total Inference FLOPs:
Total Inference FLOPs
B
+Lx (
+2xSx
(11)
18

---

Published as a conference paper at ICLR 2025
A.3.2 MIXTURE OF LLMS
The total training FLOPs for mixture model consist four main components: training routers, shading data for routers, training experts, and shading data for the experts.
For model trained with batch size B over Naeps steps, the total number of sequences shaded can be estimated as N stcpps X B E. where cach model receives approximately Nuaps B sequenees during training.
Thus, the FLOPs required for shading can be calculated as the inference FLOPs with effective batch size Nteeps x B x (see Equation 11).
Therefore Total Training FLOPs:
Total Training FLOPs =Training FLOPs Routers
Training FLOPs Experts
(12)
+ 3 N Sstuperrouuer X
Bp S H,
+ T (8 Br Sx H + 4 B, 5 H, +4x B, S H, Drw)
+ +2 x B, x V 3x B, xSXV x E (Fq. 13: Training Routers)
(13)
+NNaqueoouter B, E
M H,
+ L, (8X M H + +4x H, + 4 M H, Dr.))
+ + M H, X V xV x E (Eq. 14: Shading data for routers)
(14)
3 N Nateps.oxport export X BxSxH
+L (8 B x S 2 +4x S + D+)
+ B x S +3x XE (Eq. 15: Training Experts)
(15)
+ Nauqu expen x BXE)
M H,
+ L, (8x M H 4 X M H, 4x M H, Dr
+2 M X V x E (Eq. 16: Shading data for experts)
x
(16)
In the above formula:
Equations 15 and 13 correspond to the total training FLOPs for E experts and E routers respectively.
Equations 16 and 14 represents the inference cost E routers on batches with a prefix length M used to partition the data for experts and routers respectively.
Table 3 presents the training and inference costs for our proposed mixture of language models com- pared the dense baseline L.LMs.
19

---

Published as a conference paper at ICLR 2025
The expert balancing during training introduces only negligible computational overhead relative to both training and inference costs. Therefore, it excluded from our FLOPs calculations.
Model
Perplexity
Training cost
Inference cost
335M (dense)
11.78
31.02
0.79
335M experts)
10.78 & 8.49%
31.02 + 0.22 t 0.07%
0.79 + 0.01 1.27%
335M (dense)
11.25
62.03
0.79
335M (8 experts)
10.20 & 9.33%
62.03 + 0.75 t 1.20%
0.79 + 0.02 2.53%
335M (dense)
10.80
124.06
0.79
335M (16 experts)
9.64 4 10.74%
124.06 + 2.71 t 2. 2.19%
0.79 + 0.04 t 5.06%
335M (dense)
10.5
248.
0.79
335M (32 experts)
9.07 4 13.62%
248. 12 + 10.28 t 4.14%
0.79 + 0.08 1 10.13%
1.3B (dense)
9.10
221.33
2.81
1.3B (4 experts)
8.75 4 3.85%
221.33 + 0.36 t 0.16%
2.81 + 0.01 0.36%
1.3B (dense)
8.48
885.32
2.81
1.3B (16 experts)
7.42 | 12.53%
885.32 + 4.87 1 0.55%
2.81 + 0.04 1.42%
1.3B (dense)
8.20
1770.65
2.81
1.3B (32 experts)
6.76 17.56%
1770.65 + 18.94 t 1.07%
2.81 + 0.08 2.85%
Table 3: Performance gain versus computational overhead. The table compares performance improvements and computational costs for models with 335M (trained with 4, 8. 16, and 32 experts) and 1.3B parameters (trained with 4, 16, and 32) experts. Training cost is measured in 10 FLOPs, inference cost measured in 10a2 FI.OPs. The dense baselines and corresponding mixture models were trained on the same data volume.
Regarding Computational Overhead As demonstrated in Table 3, the 1.3B parameter model with 32 experts achieves significant performance improvements almost cost-free, gaining nearly 18% in perplexity compared to the dense baseline while incurring only 1% additional cost during training and less than 3% during inference. Furthermore, 335M parameter mixture of 32 experts reaches the same performance level as the 1.3B parameter model, but with comparable training budget and three times lower inference cost. Despite the 335M parameter model incurring a higher percentage cost during inference, this computational overhead is not the definitive minimum neces- sary for effective routing: as detailed significantly reduced by utilizing smaller routers and shorter prefix sizes. This suggests that further reductions in computational costs, without compromising routing quality, are feasible and present a valuable direction for future research.
3.4, our results indicate that the routing overhead can be
A.4 COMMUNICATION OVERHEAD
In this section, we provide estimate of the communication overhead involved in model train- ing compared to a dense baseline trained with standard distributed data parallelism. We quantify both the volume of data transmitted and the frequency of communication required during training. For simplicity of computation, we assume that all collective communications are bandwidth-oppimall meaning that the total number of bytes transferred per node approximately 2K, where K is the size of the message (Thakur et al., 2005; Patarasuk & Yuan, 2009).
Every router performs all gather operation to share the loss on the dataset chunk approximately every T = 45M training tokens, as outlined in 3.2. This frequency was primarily determined by implementation specifics related data storage.
Mixture of LLMs
We can define the number of tokens processed per step by router as Nukknss perkens-prrssepp step. which can be calculated as:
N Nuokens.per.sngg step SSS x Br
20

---

Published as a conference paper at ICLR 2025
This leads to the minimum number of steps between communications, Nappaper.com given by:
T
N steps.per_.com
S B
The maximum number of communication events, Neomm, during the router's training can then be estimated as:
N Nitepssoouter
N Socppssouuerr x SSZ x B
N Noomm
N Netepspper_ccom
T
As shown in App. Table we train routers for 128, 000 steps with batch size 32. We use context size 1024. Therefore, Noomm 94 < 1 102.
The total amount of data sent and received by each node (router) can be calculated as the total number of sequences multiplied by the number of bytes per sequence. We consider that the loss is transferred in float16, which requires 2 bytes per element. The total number of sequences can be estimated as TxN
Each router needs to send and receive the loss values for each sequence to and from other routers, resulting in factor of 2 for the data sent and received. Therefore, the data per router can be approximated as:
Tx E
Data per router x
S
Given 7' = 45 x 10a training tokens per router between communication steps, E < 32 number of experts and 1,024 sequence length, we can estimated the data transferred per node as:
Data per router 2222 Tx
45
10t
32 5.625MB
4x
S
1,024
Once routers are fully trained, they communicate only when necessary, such when experts require new data training.
For instance, consider that the message size is K bytes, equivalent to loss array of K sequences in f1oat16 precision. The frequency communication after routers are fully trained is determined by the batch size of the expert B and the number of experts E. Specifically, the number of steps each expert makes between communications can be estimated as:
K
N stepps.per.comm
2 BxBXL
Thus, communication can be scheduled approximately every 4x expert's 's training steps. Assum- ing K 231 bytes, 128 as used in the majority of our experiments, and E (4,8, 16, 32) (see Appendix 2), the communication interval becomes:
K
231
223
211
VE {4,8, 16,
(17)
x 128
X
E
E
Comparison with Distributed Training Let us assume that we train a model with batch size of B on a single host, where W represents the model size in parameters. For simplicity, we consider only data parallelism and assume that the model weights can on one GPU. To compare commu- nication costs, we examine the additional expense incurred when increasing throughput by factor of E, i.e., utilizing hosts instead of one.
In regular distributed training, scaling to E hosts necessitates copying all weights to each host and performing gradient synchronization every training step. Assuming each parameter is represented by 32 bits (as optimizers typically operate in fll oatt3) the amount of data transferred to and from each node during gradient synchronization can be calculated using the bandwidth-optime algorithm:
Data per node 2 x W 4
Here, W x represent the size of the model gradients in bytes.
21

---

Published as a conference paper at ICLR 2025
If W = 1.3 x 10 parameters, then the amount of data transferred per node per training step is:
Data per node 2 x 1.3 x 10 x = 10.4GB
Therefore, regular distributed training, each node must send and receive gigabytes of data at every step. In contrast, our approach offers flexible communication strategies. We transfer large volumes of data only few times during training (see Equation 17), or optimise for more frequent communications with smaller data sizes. This flexibility allows us adapt based on network latency and other infrastructure constraints.
B DOWNSTREAM EVALUATION
For the evaluation, we utilized the Im-eval-harness software (Gao et al., 2024). For the computa- tion of perplexity, we compute the perplexity using the format: "Question: {question). Answer: (answer)."
Task
32 experts, 2T
dense, 2T
mmlu abstract algebra
0.1400
0.2300
mmlu astronomy
0.3553
0.2829
mmlu business ethics
0.4900
0.5300
mmlu college computer science
0.2500
0.2900
mmlu college mathematics
0.1500
0.1400
mmlu computer security
0.3500
0.3800
mmlu elementary mathematics
0.2672
0.2487
mmlu formal logic
0.3016
0.2937
mmlu global facts
0.3100
0.2800
mmlu high school geography
0.3586
0.3586
mmlu high school macroeconomics
0.3128
0.2872
mmlu high school microeconomics
0.3193
0.3193
mmlu high school statisties
0.2917
0.2593
mmlu human aging
0.4081
0.3857
mmlu human sexuality
0.3969
0.3893
mmlu international law
0.2314
0.2314
mmlu logical fallacies
0.3067
0.2945
mmlu machine learning
0.2589
0.2679
mmlu medical genetics
0.3900
0.3500
mmlu moral disputes
0.2803
0.2890
mmlu moral scenarios mmlu nutrition
0.2380
0.2380
0.3039
0.2516
mmlu prehistory
0.4167
0.3858
mmlu professional accounting
0.2872
0.2518
mmlu professional medicine
0.3566
0.2904
mmlu professional psychology
0.3203
0.3186
mmlu world religions mmlu anatomy
0,5731
0.4503
0.4667
0.3926
mmlu clinical knowledge
0.3434
0.2755
mmlu college biology
0.4028
0.3819
mmlu college chemistry
0.2800
0.2600
mmlu college medicine
0.2832
0.2659
mmlu college physics
0.1863
0.1961
mmlu conceptual physics
0.3787
0.3787
Table 4: Downstream Evaluation Results (Part 1). This table presents a comparison in accuracy on downstream tasks between 1.3B-parameter 32-expert model and 1.3B- --B-parameter dense model. Both models were trained using the 2T tokens and nearly the same training FLOPs, with only 1% difference. The inference cost of the expert model increased by less than 3% compared to the dense baseline FLOPs.
22

---

Published as a conference paper at ICLR 2025
Task
32 experts, 2T
dense, 2T
mmlu econometrics
0.2719
0.2018
mmlu electrical engineering
0.2552
0.2621
mmlu high school biology
0.3290
0.3097
mmlu high school chemistry
0.2167
0.1970
mmlu high school computer science
0.2400
0.2400
mmlu high school european history
0.2848
0.3212
mmlu high school government and polities
0.4093
0.3834
mmlu high school mathematics
0.1519
0.1444
mmlu high school physics
0.2980
0.2715
mmlu high school psychology
0.4642
0.4330
mmlu high school us history
0.3382
0.3137
mmlu high school world history mmlu jurisprudence mmlu management
0.2827
0.2996
0.2407
0.2222
0.3786
0.3981
mmlu marketing
0,4915
0.4915
mmlu miscellaneous
0.5121
0.4994
mmlu philosophy
0.2958
0.3119
mmlu professional law
0.2634
0.2536
mmlu public relations
0.3545
0.4273
mmlu security studies mmlu sociology
0.3224
0.3020
0. 0.2786
0.2786
mmlu us foreign policy are challenge are easy
0.3300
0.3400
0.3200
0,2892
0.6780
0,6549
hellaswag
0.4912
0,4663
sciq
0.8950
0.9130
Table 5: Downstream Evaluation Results (Part 2)
C ADDITIONAL RESULTS
Effect of Prefix Length on Routing Performance We also conducted experiments by partitioning the data using shorter prefix length M 32 tokens during training, whereas our main results (see & 3), we used a prefix length of M = 256 tokens.
We observed that training with smaller pre- fix length resulted in better expert performance when routing during inference was done using small number of tokens (see Figure 6). This suggests that using shorter prefix during train- ing enhances the router's ability to make effec- tive decisions with limited context.
10
Perplesity
We did not take any action to make that the router models could extrapolate to longer se- quences than seen in training (e.g. interpolating positional encodings, using ALiBi (Press et al., 2022), etc.). As expected, observe a perfor- mance degradation when making routing deci- sions for sequences longer than the ones used for training.
profii 155
128
Prefix length
Figure 6: Test perplexity function of routing prefix length during inference for L.3B parameter models with 16 experts. The orange curve repre- sents a prefix length of 32, while the blue curve corresponds to prefix length of 256.
23


Reviews:
[{"review_id": "iTjZBORNAM", "reviewer": "Reviewer_iShb", "summary": "Large language models at modern scale are typically trained by means of distributed, synchronous model-parallel algorithms that rely on bespoke compute clusters with expensive fast interconnect hardware. To make progress towards eliminating this requirement, the authors tap into recent literature in parameter-sparse language models. So-called \"mixture of experts\" (MoE) language models replace a monolithic language model with a number of smaller \"expert\" models that are allocated inputs according to some (typically co-trained) routing mechanism and therefore require many fewer active parameters per input at inference time. The authors propose efficient new routers that can be trained completely independently of the mixture of experts, permitting asynchronous training for each expert. They show that mixture-of-expert models trained with these routers are competitive with dense language models of the same size and also outperform MoE models trained with even more trivial routers previously explored in the literature.", "questions": "How were the number of tokens chosen for each configuration (see Weaknesses section above)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The method is simple (to its credit), intuitive, and effective, and it attacks an important problem. While information about the largest and most significant language models is, of course, sparse, they do not appear to be growing at the same rate as they have in previous years, and significant attention is now being devoted to the problem of improving the performance of comparatively small models (e.g. 2B models like Gemma). Good algorithms for asynchronous training could potentially unlock previously unattainable advantages of scale.\n\nThe paper is thorough and explicitly addresses most of the potential shortcomings of the proposed method that came to my mind as I read it (e.g. efficient allocation of inputs, the question of how each of the experts perform relative to each other). The authors ran a lot of very expensive experiments. It's also quite clearly written.\n\nI was also surprised by the effectiveness of small routing models. I wouldn't have guessed that the data partitions would have been so similar across model scales.", "weaknesses": "I'm not in love with the choice to benchmark the asynchronous MoE against a dense model with the same number of parameters as each expert but trained for as many tokens as all experts combined (see the para. starting \"Comparison to the Dense Model...\"). Depending on how the number of training tokens was chosen, it might unfairly bias the results towards the expert setup---training one really saturated model up to 32x longer than each expert might just be a waste of compute and not a realistic counterfactual. At inference time, it's also a little generous to the MoE approach. It implicitly assumes that all of the experts are loaded in GPU memory at inference time. While this is not completely unrealistic, especially if we're modeling the very plausible scenario where we're running some kind of chat service and the number of incoming requests is high enough that we can pre-load all experts and still keep them saturated, it ignores the very real strength of the dense model that it can more efficiently be run in situations where these assumptions do not hold. I generally prefer Gururangan et al., 2023's approach to evaluation, which pitches the MoE against a single larger but train-compute-matched dense model. Inference-time compute is matched by ensembling the predictions of the top k experts. I like it for the additional reason that it gestures toward the scenario where the dense model is so large that it does not fit on a single compute node and requires model-parallel training (whereas individual experts don't). If this isn't the case, we can always still train the dense model without fast interconnect hardware if we're just willing to train it for a longer time, which, relatively speaking, diminishes the argument for a sparse model. There are lines in this paper that hint at something similar by comparing 335M experts to the 1.3B dense model, so I assume it wouldn't require too much work/compute to add more evaluations with that flavor.\n\nThroughout, the paper also needs to do a better job positioning itself relative to Gururangan et al., 2023. The abstract is pretty representative; with the exception of the phrase \"according to a short prefix\" and the name of the method, the abstract of this paper could have been that for the prior paper, which differs from this one only in the choice of routing function. It needs to be clearer that this isn't the first paper to attempt asynchronous training using MoEs and that all the novelty lies in the routing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HdpRaw3lfy", "reviewer": "Reviewer_8sbF", "summary": "The study introduces a method for asynchronously training a mixture of language models (LLMs), called experts, which significantly reduces communication overhead by eliminating the need for high-volume interactions between experts during training. Each expert is specialized to process a distinct segment of data, directed by a small, specialized router—a lightweight language model with as few as 4 million parameters. These routers evaluate text sequences using a scoring system that determines the most suitable expert for each sequence based on a short prefix of 256 tokens. A novel balanced assignment approach is proposed to ensure that each expert receives a roughly equal volume of data, minimizing training imbalances. The only communication overhead is needed for the data partitioning phase and the conventional high-volume communication between experts is completely removed by the proposed method.", "questions": "Minor edits:\n- Line 300: glitch in the reference for Adam Optimizer\n- Figure 3 is difficult to read, especially the legends\n\n\nQuestions:\n- Line 309: a different schedule (warmup + stable) for the routers, it’d be great if authors could add a line of reasoning behind this decision\n- Line 388: Figure 4 (b): If the performance is sensitive to the prefix length, how would it affect the tasks where only a few tokens are provided as the prompt e.g., essay writing?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The authors show that the same perplexity can be achieved with three times less compute cost compared to the dense baseline\n\n- The final performance (test PPL) is not sensitive to the router size hence, making the approach succeed with even tiny routers (4M parameters)\n\n- The authors have conducted experiments that are highly expensive in nature. The community will benefit from such a study. Especially, the test PPL benefits (line 349) with three times less compute is a remarkable result, it can potentially open avenues for more research on large-scale LLM development with limited resources.", "weaknesses": "Performance (test PPL) is sensitive to the prefix length (256 tokens used to train routers) and the sensitivity increases with the number of experts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Cl23DlNG7Q", "reviewer": "Reviewer_GV45", "summary": "This paper introduces SmalltalkLM, where instead of having a single large dense connected network, we have multiple individual dense models (the experts) and a router, which predicts which expert a given input should be routed to. During training, the data is partitioned into separate groups using the router, and each expert is then trained with language modelling on the allocated subset of data. During inference, the router can assign a new sample to the relevant expert. This enables the model to be trained with more limited RAM requirements (as each expert can be trained independently on the partitioned data) and yields better performance while only having a marginal increase in cost at inference (with the increase only as the router is first used to route the current input).", "questions": "- Just to clarify, the performance shown in Figure 3 is the performance of the model, which was only trained using language modelling, when it is further prompted for the downstream tasks?\n\n- I want to raise the weaknesses above in case the authors have any further comments.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Unlike approaches like MoE, the experts here are independent, and therefore, the routing can be done before each training epoch. Since each system can be trained independently, this reduces the simultaneous RAM requirements, which, unlike MoE approaches, enables SmalltalkLM systems to be trained using the same computational hardware. Inference also has lower RAM requirements. \n\n- They demonstrate that this approach results in better performance across different model sizes and can yield improved downstream performance on different tasks beyond language modelling. \n\n- The paper is quite easy to follow, and the overall idea is quite simple and seems reasonable.", "weaknesses": "- As far as I understand, this work takes ideas of existing deep learning MoE but simplifies the process, and instead of using the top-k experts, operating at the token-level, or having experts per layer, have a single router to a single expert, which may be of limited novelty.\n\n- The savings in RAM might not be too substantial, as when deployed in live systems, you’ll need to load up all experts simultaneously anyway (to avoid the massive latency of continually reloading the next expert). Since the running inference costs are already more expensive than training, it seems to be quite a contrived advantage.\n\n- this point is a bit pedantic, but the base system was presented as a vanilla LLM (i.e. not actually helpful for any particular task), although it was demonstrated it yielded marginally better performance of MCQA tasks in an ‘emergent’ setting. If these systems are further tuned (which I assume is the purpose of having these base LLMs), the routing aspect may limit the practicality. The router balancing may no longer be applicable, and possibly, specialized systems may not generalize as well as a single network trained on all the data when adapted to particular tasks. It might be interesting to see how this model might operate in more practical settings, i.e. with instruction-tuning or fine-tuning to particular tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)