No Need to Talk: Asynchronous Mixture of Language Models

## Abstract

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. 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 % percent 75 75\% 75 % of the tasks.

## 1 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., ; Brown et al., ; Henighan et al., ; Hoffmann et al., ; Dubey et al., ). Scaling introduces substantial operating and engineering costs for both inference and training. In general, training is achieved on a large number of nodes via synchronous gradient descent techniques, which relies on high-bandwidth communication to scale. Inference of large models may require multiple compute nodes to distribute the model, which relies on low-latency communication. In both cases, state-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 on the scaling side, access to a large number of well interconnected nodes remains a privilege in the machine learning community.

In this paper, we explore strategies to 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 at mitigating reliance on high-bandwidth interconnects, researchers have developed algorithms that reduce the need for frequent or comprehensive gradient synchronizations. Such techniques include asynchronous training (Douillard et al., ; Aji & Heafield, ; Zhang et al., ; Liu et al., ) and gradient compression (Lin et al., ; Dettmers, ) which effectively decrease communication overhead. By performing updates less frequently or by communicating less data, these methods sustain high training throughput while diminishing dependence on high-speed interconnects. However, these algorithms still require some level of gradient synchronization, and resulting models often under perform (in terms of perplexity) compared to training approaches synchronizing at every step (Diskin et al., ).

Regarding efficient inference, a number of sparse parameter activation techniques have recently become popular (Shazeer et al., ; Fedus et al., ; Artetxe et al., ; Lewis et al., ; Du et al., ), in particular the Switch Transformer mixture of experts (MoE). These approaches are 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 et al., ; Li et al., ; Gururangan et al., ). In these methods, the input space is clustered using unsupervised techniques like K-Means, with each model in the mixture trained on an assigned cluster. These approaches reduce communication overhead and have demonstrated improvements in model perplexity, while maintaining similar inference costs to baseline models (Gross et al., ; Gururangan et al., ).
However, in the context of language processing, their practical applicability on standard downstream tasks with relatively short input prefixes remains a subject of debate, as they rely on full corpus clustering or access to metadata (Li et al., ; Gururangan et al., ).

In this paper, we introduce SmallTalk LM, 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 of the overall architecture 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 are as follows:

- •
We introduce an algorithm that trains a mixture of independent language models, significantly 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 a given sequence based on a 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.
- •
We evaluate our method on a set of downstream tasks, showing that it achieves better or similar accuracy compared to a dense baseline on $75\%$ of the tasks (see § 3.3 and App. B).

## 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. Subsequently, in § 2.2, we show how we can use independent language models to perform routing to different experts in the mixture. We also present the training procedure and implementation details of our method. Finally, we discuss its marginal computational overhead and minimal distributed communication costs.

### 2.1 Background

##### Language modeling

Let $\mathcal{V}$ denote a fixed vocabulary, and consider a sequence of tokens $\mathbf{x}=(x_{1},x_{2},\dots,x_{S})$, where each $x_{s}\in\mathcal{V}$ for $s=1,2,\dots,S$. Language modeling is about modeling the data distribution $p(\mathbf{x})$, which in practice is cast as training a language model $p(x_{s+1}\mid\mathbf{x}_{1:s}\,;\,\theta)$ generating a token $x_{s+1}$ given a context sequence $\mathbf{x}_{1:s}=(x_{1},x_{2},\dots,x_{s})$. Typically, the model parameterized by $\theta$ is a neural network, and parameters are obtained by minimizing the negative log-likelihood:

$$ $\mathcal{L}(\mathbf{x};\theta)=-\,p(\mathbf{x}\,;\,\theta)=-\sum_{s=1}^{S-1} \log p(x_{s+1}\mid\mathbf{x}_{1:s}\,;\,\theta)\,.$ (1) $$

##### Mixture of Experts

There exist many variants of the mixture of experts model (Jacobs et al., ; Jordan & Jacobs, ; Tresp, ; Collobert et al., ; Shahbaba & Neal, ). Following the original formulation and applying it to language modeling, we factorize the conditional distribution for predicting a given token by summing over a set of $E$ experts:

$$ $\mathcal{L}(\mathbf{x};\theta)=-\sum_{s=1}^{S-1}\log\sum_{e=1}^{E}p(x_{s+1} \mid\mathbf{x}_{1:s},e\,;\,\theta^{e})\,p(e\mid\mathbf{x}_{1:s}\,;\,\theta^{r}),$ (2) $$

where each expert $p(x_{s+1}\mid\mathbf{x}_{1:s},e\,;\,\theta^{e})$ is intended to focus on different parts of the data distribution, and the router $p(e\mid\mathbf{x}_{1:s}\,;\,\theta^{r})$ assigns weights to the expert predictions. Each expert $e$ has its own set of independent parameters $\theta^{e}$, and the router is parameterized by $\theta^{r}$.

While this type of mixture of experts partitions the input space across different models, all experts are still involved when computing the conditional likelihood in Equation (2), both at training and inference time. Hard mixtures of experts circumvent this issue by hard-assigning a sequence $\mathbf{x}_{1:s}$ to a *single* expert $e^{\star}$, minimizing an upper bound of Equation (2):

$$ $\mathcal{L}(\mathbf{x};\theta)\leq-\sum_{s=1}^{S-1}\log p(x_{s+1}\mid\mathbf{x }_{1:s},e^{\star}\,;\,\theta^{e^{\star}})\,p(e^{\star}\mid\mathbf{x}_{1:s}\,; \,\theta^{r}),$ (3) $$

The way $e^{\star}$ is chosen depends of the goal of the approach. For example, in Jacobs et al. (), $e^{\star}$ is sampled at training time, according to the distribution $p(e\mid\mathbf{x}_{1:s}\,;\,\theta^{r})$. In other works, like Collobert et al. (), it is chosen as the most relevant expert (that is $e^{\star}=\operatorname*{arg\,max}_{e}p(e\mid\mathbf{x}_{1:s}\,;\,\theta^{r})$), both at training and inference. Regarding the training procedure, mixture models with both soft and hard assignments are typically trained with an Expectation-Minimization (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 to the most relevant expert $e^{\star}$, according to the router posterior $p(e\mid\mathbf{x}_{1:s}\,;\,\theta^{r})$, both at training and inference. However, instead of being a monolithic neural network model, the router is implemented as a language model per expert. More precisely, with Bayes’ rule we have:

$$ $\displaystyle e^{\star}$ $\displaystyle=\operatorname*{arg\,max}_{e}p(e\mid\mathbf{x}_{1:s}\,;\,\theta^{ r})$ (4) $\displaystyle=\operatorname*{arg\,max}_{e}\frac{p(\mathbf{x}_{1:s}\mid e\,;\, \theta^{r})\,p(e)}{\sum_{i=1}^{E}p(\mathbf{x}_{1:s}\mid i\,;\,\theta^{r})\,p(i)}$ (5) $\displaystyle=\operatorname*{arg\,max}_{e}p(\mathbf{x}_{1:s}\mid e\,;\,\theta^ {r})$ (6) $\displaystyle=\operatorname*{arg\,max}_{e}p(\mathbf{x}_{1:s}\mid\theta^{r,e})\,,$ (7) $$

where the likelihood $p(\mathbf{x}_{1:s}\mid\theta^{r,e})$ defines a routing language model attributed to the $e$-th expert with *independent* parameters $\theta^{r,e}$, and expert priors $p(e)$ are supposed to be uniform. So far, in our definition, there is nothing that prevents $\theta^{r,e}=\theta^{e}$, namely to use the experts themselves to implement the router. Interestingly, in that case selecting $e^{\star}$ as above is intuitive because it amounts to selecting the best expert for a given sequence (in terms of log-likelihood).

##### Deriving a Practical Model

In order for the above formulation to result in a model that is useful in practice, we have to solve the following main challenges. Firstly, we 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 full sequence (e.g. in question answering we do not have the answer). Secondly, using the experts 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$ to compute the score for each expert as follows:
$p(e^{\star}\mid\mathbf{x}_{1:s})\approx p(e^{\star}\mid\mathbf{x}_{1:M}).$
(8)
This enables us to use SmallTalk LM in real-world applications where only a very small prefix may be available. We show in § 3.4 that our method outperforms the dense baseline even when using as little as $32$ tokens for routing.
- 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 is good at predicting the beginning of a 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 is then used to 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 a 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 a sequence is similar to the training data or not. We show that even really small models ($4.4$M parameters) can perform that task very well and result in no discernible difference in performance compared to larger routers (§ 3.4).

Figure: Algorithm 1 SmallTalk LM training.

##### Training Procedure

We train our proposed hard mixture of experts with EM. Namely we alternate between a step of optimizing equation 1 and a step of assignments according to equation 4. In our framework, assignments $e^{\star}$ purposefully do not depend on the experts, and the routers are 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
$\mathcal{L}(\mathbf{x},\theta^{r})=-\log p(e^{\star}\mid\mathbf{x}_{1:M};
\theta^{r})=-\sum_{s=1}^{M-1}\log p\left(x_{s+1}|\mathbf{x}_{1:s};\theta^{r,e^
{\star}}\right),$
(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 selected 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 ($\theta^{e}=\theta^{r,e}$) our training algorithm comes down to a 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 as evenly as possible across all experts, or in other words that the model is utilizing the capacity of all experts efficiently instead of relying on a few good ones.
Common ways of tackling this problem are either introducing losses that penalize uneven assignments or introducing noise into the assignments namely by sometimes performing “sub-optimal” assignments to avoid local-minima. These techniques are used in all mixture models since their introduction in Jacobs et al. () until more recently in Fedus et al. ().

Figure: (a) Naive assignments
Refer to caption: x1.png

In our setup, we can solve this problem by ensuring that each expert is assigned equally sized distinct portions of the dataset.
However, a challenge arises when assigning data sequentially based solely on Equation 4. In 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 in Figure 1(a).

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_{e}\log p(\mathbf{x}_{1:M}\mid e;\theta^{r})$ and perform the assignments in that order. This 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 1 we refer to this procedure as *balanced assignments* and it is demonstrated in Figure 1(b).

During inference, no balancing is performed, and the expert is selected solely based on equation 4.

##### Computational and Communication Cost

It is 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 on (see Appendix A.3 for further details). As a result, the additional computational overhead from training and inference with the routers is negligible.

Regarding communication requirements for a distributed implementation, in Algorithm 1 all operations except for the balanced assignments are completely independent and embarrassingly parallelizable without any communication. Namely, each expert or each router can be trained on its 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 in $2$MB per $1$M sequences if we represent the scores in 16-bit floats. A detailed analysis is provided in the Appendix A.4.

## 3 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 techniques.

Second, in § 3.3, we evaluate our method on downstream tasks, demonstrating that it 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 our method, the routing. We show that the performance of the mixture does not depend 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  Gururangan et al. (). Moreover, we demonstrate that the perplexity gains are not driven by a 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, ), which is a 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 ().

Both the experts and routers utilize a Transformer architecture (Vaswani, ) with rotary positional encoding (Su et al., ). All experiments use sequences of $1,024$ tokens. We train the models using the AdamW optimizer (Loshchilov & Hutter, ) with $\beta_{1}=0.9$, $\beta_{2}=0.99$, and a 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\times 10^{-4}$ over the first $3,000$ steps, followed by a cosine decay for the remainder of the training period. We experiment with two model sizes: a $335$M-parameter model with $4$, $8$, $16$, and $32$ experts and a $1.3$B-parameter model with $4$, $16$, and $32$ experts. Tables 1 and 2 describe the architectures and the training parameters in more detail.

For the routers, we use models with $4.4$M parameters. The routers are trained for $128,000$ steps using a constant learning rate of $1\times 10^{-4}$, following a linear warm-up over the first $1,000$ steps. Routers are trained with a 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 of 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 as 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 FLOPs.

### 3.2 Language Modeling Results

Figure: (a) FLOPs vs perplexity, $335$M.
Refer to caption: x3.png

Across all experimental configurations, our method *consistently outperforms the baseline* while utilizing the same volume of data, with only minor increases in computational costs. Specifically, for the 1.3B parameter model, our approach results in less than a 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 are modest, with less than a 4% increase in training FLOPs and less than a 10% increase in inference FLOPs, alongside performance gains of up to nearly 14% in perplexity (see Appendix A.3 for details on FLOPs calculation). Figures 2(a) and 2(b) illustrate the perplexity on a held-out test set relative to the total training cost, measured in petaFLOPs, across different model sizes. Additionally, Figure 2(c) shows the test perplexity as a function of the number of training tokens for the 1.3B parameter model.

Moreover, it 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 LM 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 in perplexity increases with the number of experts. This suggests that our method more efficiently utilizes the training data to improve performance, likely due to the specialization of experts on different segments of the data distribution.

##### Communication Overhead

Our approach performs minimal communication and does not require a fast interconnect. As detailed in § 2, we first train the routers which are then used to efficiently 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 $6$ 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 in real-world scenarios, we performed zero-shot evaluation on a set of natural language processing (NLP) tasks, including ARC Challenge and ARC Easy (Clark et al., ), HellaSwag (Zellers et al., ), SciQ (Welbl et al., ), and MMLU (Hendrycks et al., ). Figure 3 shows both accuracy and perplexity for each task (see Appendix B for details regarding perplexity and accuracy computation).

Specifically, our largest configuration, a model with $32$ $1.3$B parameter experts, demonstrates better performance on four out of five tasks, achieving gains of $3\%$, $2\%$, $3\%$, and $1\%$ on ARC Challenge (3(a)), ARC Easy (3(b)), HellaSwag (3(c)), and MMLU (3(d)), 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$), as shown in Table 4.

Figure: (a) ARC Challenge
Refer to caption: x6.png

### 3.4 Routing analysis

In this section we analyse the most critical component of SmallTalk LM, the routing. In particular, 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 of the mixture (Figure 4(a)). 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 $335$M parameter model configured with $4$ experts, using four different router sizes for data partitioning: a $335$M parameter router (where the model routes data for itself, optimizing for the data it handles best), and routers with $110$M, $64$M, and $4.4$M parameters.
All models were trained with a batch size of $32$ for $256,000$ steps, using a routing prefix length of $256$ tokens, following the same experimental setup as in § 3.2 (see Appendix A
for more details about model architectures and training parameters). Figure 4(a) illustrates the test perplexity wrt the number of training tokens for these configurations. We observe that all router sizes perform practically identically.

Figure: (a) Router size
Refer to caption: x10.png

##### 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 in Figure 4(b) that SmallTalk LM outperforms the dense baseline even when routing with shorter prefixes. In particular, we show that even with prefixes of just $32$ tokens ($\frac{1}{8}$th of the length during training) the model still outperforms the dense baseline.

##### Comparison With Routing Using TF-IDF

An alternative approach might suggest that if a tiny LLM suffices for routing, simpler methods such as TF-IDF encoding combined with balanced K-Means clustering, as proposed by Gururangan et al. (), could be adequate. To test this hypothesis, we trained a 335M parameter model using the routing strategy outlined in (Gururangan et al., ): 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.

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 4(c), our proposed routing method significantly outperforms the TF-IDF-based routing. This demonstrates that our 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 under perform on their assigned data. To address this concern, we conducted an analysis of each expert’s performance individually. Figures 5(a), 5(b), and 5(c) compare the perplexity achieved by our method and the dense baseline on the routed dataset segments for the $1.3$B parameter model, utilizing mixtures of $4$, $16$, and $32$ experts, respectively.

Our findings demonstrate that all experts contribute positively, with consistent perplexity improvements observed across all routed dataset segments. Moreover, the performance gap between our method and the baseline widens as the number of experts increases, suggesting enhanced specialization 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 a few dominant experts.

Figure: (a) $4$ experts vs dense, $266$B
Refer to caption: x13.png

## 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, we discuss existing methods that aim to make data parallelism more communication-efficient via reducing the frequency of the synchronization. Since our paper also explores methods for efficient inference without performance degradation, we discuss this branch of research as well.

##### Distributed optimization

To reduce the frequency of synchronisation during training of large models, researchers have explored distributed optimization techniques both theoretically and empirically
(Mcdonald et al., ; Stich, ; Lian et al., ; Wang & Joshi, ; McMahan et al., ; Lin et al., ; Zinkevich et al., ). The core principle of these methods is to 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-training (Diskin et al., ; Douillard et al., ) and fine-tuning (Borzunov et al., ; Liu et al., ; Hilmkil et al., ; Ro et al., ).

While theses approaches reduces the frequency of communication, thereby decreasing communication overhead, they often result in degraded model performance compared to every step synchronization (Diskin et al., ), or they lack direct performance comparisons in the context of pre-training (Douillard et al., ). The degradation in performance is often attributed to the staleness of gradients and the divergence between local and global models due to delayed synchronization (Zhang et al., ; Yu et al., ; Mitra et al., ). Moreover, to our knowledge, all experiments have been conducted at a small scale ($\leq 400$M parameters) (Douillard et al., ; Liu et al., ; Ro et al., ). Therefore, scalable solutions that can reduce communication frequency for LLM training without compromising performance at larger scales are still needed.

##### Mixture of Experts

Mixture of Experts (MoE) models have emerged as a prominent approach to increase model capacity without a proportional increase in inference costs (Jacobs et al., ; Jordan & Jacobs, ). The key idea behind MoE is to 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 to enhance performance while reducing computational overhead during inference.

Early implementations of MoE employed algorithms such as Gaussian Processes (Tresp, ; Theis & Bethge, ; Deisenroth & Ng, ), Dirichlet Processes (Shahbaba & Neal, ), and Support Vector Machines (SVMs) (Collobert et al., ) to model these specialized functions. Recent advancements have integrated deep neural networks as experts, providing scalable and flexible implementations 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. () 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 so the number of active parameters is decreased. Similar approaches have been effectively employed in models like the Switch Transformer (Fedus et al., ), GShard (Lepikhin et al., ), GLaM (Du et al., ) and Mixtral of Experts (Jiang et al., ) which leverage MoE layers to achieve state-of-the-art 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., ; Fedus et al., ; Riquelme et al., ).
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., ; Li et al., ; Gururangan et al., ).
For instance, Gross et al. () proposed to train a 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. () and Gururangan et al. () 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-time applications involving shorter inputs. For instance, Gururangan et al. () focus on performance in text classification tasks, which may not fully reflect real-world scenarios.

## 5 Conclusion

In this work, we present SmallTalk LM – a method that significantly reduces communication overhead during LLM training. It 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 we 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.

## 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.

## Appendix A Language modeling experiments

### A.1 Architecture and training details

Our architecture is based on the transformer decoder (Vaswani, ). We use rotary positional encoding (Su et al., ). We use a SentencePiece (Kudo & Richardson, ) tokenizer with a vocabulary of $32000$ tokens. The model specific parameters for the different sizes are summarized in Table 1. We implemented our EM scheme for training the routers in PyTorch (Paszke et al., ). After segmenting the training set, the experts were trained independently using Jax (Bradbury et al., ). All training was done in bfloat16. The optimizer states and operations are in float32. The training parameters for specific models are summarised in Table 2.

**Table 1: Model parameters.**
| Model parameters | Role | Hidden size | FFW expansion factor | Layers | Attention heads |
| --- | --- | --- | --- | --- | --- |
| $335$M | expert | $1024$ | $4$ | $24$ | $16$ |
| $1.3$B | expert | $2048$ | $4$ | $24$ | $16$ |
| $4.4$M | router | $96$ | $4$ | $12$ | $12$ |
| $64$M | router | $416$ | $4$ | $12$ | $12$ |
| $110$M | router | $768$ | $4$ | $12$ | $12$ |

**Table 2: Training parameters.**
| Model | Steps (k) | Tokens (B) | Batch size | # GPUs |
| --- | --- | --- | --- | --- |
| $335$M (dense) | $256$ | $133$ | $512$ | $8$ |
| $335$M ($4$ experts) | $256$ | $133$ | $128$ | $8$ |
| $335$M (dense) | $512$ | $266$ | $512$ | $32$ |
| $335$M ($8$ experts) | $256$ | $266$ | $128$ | $8$ |
| $335$M (dense) | $1024$ | $532$ | $512$ | $32$ |
| $335$M ($16$ experts) | $256$ | $532$ | $128$ | $8$ |
| $1.3$B (dense) | $512$ | $266$ | $512$ | $32$ |
| $1.3$B ($4$ experts) | $512$ | $266$ | $128$ | $8$ |
| $1.3$B (dense) | $1024$ | $1000$ | $1024$ | $64$ |
| $1.3$B ($16$ experts) | $512$ | $1000$ | $128$ | $8$ |
| $1.3$B (dense) | $1024$ | $2000$ | $2048$ | $128$ |
| $1.3$B ($32$ experts) | $512$ | $2000$ | $128$ | $8$ |
| $4.4$M (router) | $128$ | $4$ | $32$ | $1$ |

### 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_{\text{ff}}$)

For the router LLM:

- •
Prefix length for routing ($M$)
- •
Number of layers ($L_{r}$)
- •
Hidden dimension size ($H_{r}$)
- •
Number of attention heads ($A_{r}$)
- •
Batch size ($B_{r}$)
- •
Feedforward network dimension ($D_{r_{\text{ff}}}$) – parameters specific to the router model
- •
Number of training tokens per router between communications $T$

We also define $N_{\text{steps}}$, $N_{\text{steps_expert}}$, and $N_{\text{steps_router}}$ as the number of training steps for the dense baseline, the expert in the mixture, and the router, respectively. The sequence length ($S$) and vocabulary size ($V$) are consistent across the dense baseline, 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 calculated 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 as follows:

- •
Embedding Layer: Although embeddings are typically implemented as lookups with negligible computational cost, for completeness, we estimate the FLOPs as:
$\text{FLOPs}_{\text{emb}}=B\times S\times H$
- •
Multi-Head Attention (MHA):
1.
Linear Projections (Queries, Keys, Values):
$\text{FLOPs}_{\text{proj}}=3\times 2\times B\times S\times H\times H=6\times B
\times S\times H^{2}$
2.
Scaled Dot-Product Attention:
$\text{FLOPs}_{\text{attn}}=\text{FLOPs}_{\text{QK}}+\text{FLOPs}_{\text{V}}=2
\times B\times S^{2}\times H+2\times B\times S^{2}\times H=4\times B\times S^{
2}\times H$
3.
Output Projection:
$\text{FLOPs}_{\text{out_proj}}=2\times B\times S\times H\times H$
Total Multi-Head Attention (MHA):
$\displaystyle\text{FLOPs}_{\text{MHA}}$
$\displaystyle=\text{FLOPs}_{\text{proj}}+\text{FLOPs}_{\text{attn}}+\text{
FLOPs}_{\text{out_proj}}$
$\displaystyle=6\times B\times S\times H^{2}+4\times B\times S^{2}\times H+2
\times B\times S\times H^{2}$
$\displaystyle=8\times B\times S\times H^{2}+4\times B\times S^{2}\times H$
- •
Feedforward Network (FFN):
$\text{FLOPs}_{\text{FFN}}=2\times 2\times B\times S\times H\times D_{\text{ff}
}=4\times B\times S\times H\times D_{\text{ff}}$
- •
Output projection:
$\text{FLOPs}_{\text{out}}=\text{FLOPs}_{\text{out}_{\text{proj}}}+\text{FLOPs}
_{\text{softmax}}=2\times B\times S\times H\times V+3\times B\times S\times V$

Total FLOPs per Layer:

$$ $\text{FLOPs}_{\text{layer}}=\left(\text{FLOPs}_{\text{MHA}}+\text{FLOPs}_{ \text{FFN}}\right)\times L$ $$

Total Forward Pass FLOPs per Step:

$$ $\text{FLOPs}_{\text{forward}}=\text{FLOPs}_{\text{emb}}+\text{FLOPs}_{\text{ layer}}+\text{FLOPs}_{\text{out}}$ $$

Total Training FLOPs per Step:

The backward pass is approximately twice as computationally intensive as the forward pass. Therefore:

$$ $\text{FLOPs}_{\text{train_step}}=\text{FLOPs}_{\text{forward}}+2\times\text{ FLOPs}_{\text{forward}}=3\times\text{FLOPs}_{\text{forward}}$ $$

Total Training FLOPs:

Multiplying the FLOPs per training step by the total number of training steps:

$$ $\text{Total Training FLOPs}=\text{FLOPs}_{\text{train_step}}\times N_{\text{ steps}}=3\times\text{FLOPs}_{\text{forward}}\times N_{\text{steps}}$ $$

Total Training FLOPs:

$$ $\displaystyle\text{Total Training FLOPs}=3\times N_{\text{steps}}\times\Bigg{[}$ $\displaystyle B\times S\times H$ $\displaystyle+L\times\Big{(}8\times B\times S\times H^{2}+4\times B\times S^{2 }\times H+4\times B\times S\times H\times D_{\text{ff}}\Big{)}$ $\displaystyle+2\times B\times S\times H\times V+3\times B\times S\times V\Bigg {]}$ (10) $$

Total Inference FLOPs:

$$ $\displaystyle\text{Total Inference FLOPs}=$ $\displaystyle B\times\Bigg{[}\times S\times H$ $\displaystyle+L\times\Big{(}8\times S\times H^{2}+4\times S^{2}\times H+4 \times S\times H\times D_{\text{ff}}\Big{)}$ $\displaystyle+2\times S\times H\times V+3\times S\times V\Bigg{]}$ (11) $$

#### A.3.2 Mixture of LLMs

The total training FLOPs for the mixture model consist of four main components: training routers, sharding data for routers, training experts, and sharding data for the experts.

For a model trained with batch size $B$ over $N_{\text{steps}}$ steps, the total number of sequences shaded can be estimated as $N_{\text{steps}}\times B\times E$, where each model receives approximately $N_{\text{steps}}\times B$ sequences during training.

Thus, the FLOPs required for sharding can be calculated as the inference FLOPs with an effective batch size of $N_{\text{steps}}\times B\times E$ (see Equation 11).

Therefore
Total Training FLOPs:

$$ $\displaystyle\text{Total Training FLOPs}=$ $\displaystyle\text{Training FLOPs Routers}+\text{Training FLOPs Experts}=$ (12) $\displaystyle+3\times N_{\text{steps_router}}\times\Bigg{[}B_{r}\times S \times H_{r}$ $\displaystyle+L_{r}\times\left(8\times B_{r}\times S\times H_{r}^{2}+4\times B _{r}\times S^{2}\times H_{r}+4\times B_{r}\times S\times H_{r}\times D_{r_{ \text{ff}}}\right)$ $\displaystyle+2\times B_{r}\times S\times H_{r}\times V+3\times B_{r}\times S \times V\Bigg{]}\times E\quad\text{(Eq.~{}\ref{eq:training_routers}: Training Routers)}$ (13) $\displaystyle+\bigg{(}N_{\text{steps_router}}\times B_{r}\times E\bigg{)} \times\Bigg{[}M\times H_{r}$ $\displaystyle+L_{r}\times\left(8\times M\times H_{r}^{2}+4\times M^{2}\times H _{r}+4\times M\times H_{r}\times D_{r_{\text{ff}}}\right)$ $\displaystyle+2\times M\times H_{r}\times V+3\times M\times V\Bigg{]}\times E \quad\text{(Eq.~{}\ref{eq:sharding_routers}: sharding data for routers)}$ (14) $\displaystyle\ 3\times N_{\text{steps_expert}}\times\Bigg{[}B\times S\times H$ $\displaystyle+L\times\left(8\times B\times S\times H^{2}+4\times B\times S^{2} \times H+4\times B\times S\times H\times D_{\text{ff}}\right)$ $\displaystyle+2\times B\times S\times H\times V+3\times B\times S\times V\Bigg {]}\times E\quad\text{(Eq.~{}\ref{eq:training_experts}: Training Experts)}$ (15) $\displaystyle+\bigg{(}N_{\text{steps_expert}}\times B\times E\bigg{)}\times \Bigg{[}M\times H_{r}$ $\displaystyle+L_{r}\times\left(8\times M\times H_{r}^{2}+4\times M^{2}\times H _{r}+4\times M\times H_{r}\times D_{r_{\text{ff}}}\right)$ $\displaystyle+2\times M\times H_{r}\times V+3\times M\times V\Bigg{]}\times E \quad\text{(Eq.~{}\ref{eq:sharding_experts}: sharding data for experts)}$ (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 of $E$ routers on $B$ 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 compared to the dense baseline LLMs.

The expert balancing during training introduces only a negligible computational overhead relative to both training and inference costs. Therefore, it is excluded from our FLOPs calculations.

**Table 3: Performance gain versus computational overhead. The table compares performance improvements and computational costs for models with $335$M (trained with $4$, $8$, $16$, and $32$ experts) and $1.3$B parameters (trained with $4$, $16$, and $32$) experts. Training cost is measured in $10^{19}$ FLOPs, inference cost is measured in $10^{12}$ FLOPs. The dense baselines and corresponding mixture models were trained on the same data volume.**
| Model | Perplexity | Training cost | Inference cost |
| --- | --- | --- | --- |
| $335$M (dense) | $11.78$ | $31.02$ | $0.79$ |
| $335$M (4 experts) | $10.78\,\color[rgb]{0.1,0.5,0.1}\definecolor[named]{pgfstrokecolor}{rgb}{<br>0.1,0.5,0.1}{\downarrow\,8.49\%}$ | $31.02+0.22\,\uparrow\,0.07\%$ | $0.79+0.01\,\uparrow\,1.27\%$ |
| $335$M (dense) | $11.25$ | $62.03$ | $0.79$ |
| 335M ($8$ experts) | $10.20\,\color[rgb]{0.1,0.5,0.1}\definecolor[named]{pgfstrokecolor}{rgb}{<br>0.1,0.5,0.1}{\downarrow\,9.33\%}$ | $62.03+0.75\,\uparrow\,1.20\%$ | $0.79+0.02\,\uparrow\,2.53\%$ |
| $335$M (dense) | $10.80$ | $124.06$ | $0.79$ |
| $335$M ($16$ experts) | $9.64\,\color[rgb]{0.1,0.5,0.1}\definecolor[named]{pgfstrokecolor}{rgb}{<br>0.1,0.5,0.1}{\downarrow\,10.74\%}$ | $124.06+2.71\,\uparrow\,2.19\%$ | $0.79+0.04\,\uparrow\,5.06\%$ |
| $335$M (dense) | $10.5$ | $248.12$ | $0.79$ |
| $335$M ($32$ experts) | $9.07\,\color[rgb]{0.1,0.5,0.1}\definecolor[named]{pgfstrokecolor}{rgb}{<br>0.1,0.5,0.1}{\downarrow\,13.62\%}$ | $248.12+10.28\,\uparrow\,4.14\%$ | $0.79+0.08\,\uparrow\,10.13\%$ |
| $1.3$B (dense) | $9.10$ | $221.33$ | $2.81$ |
| $1.3$B ($4$ experts) | $8.75\,\color[rgb]{0.1,0.5,0.1}\definecolor[named]{pgfstrokecolor}{rgb}{<br>0.1,0.5,0.1}{\downarrow\,3.85\%}$ | $221.33+0.36\,\uparrow\,0.16\%$ | $2.81+0.01\,\uparrow\,0.36\%$ |
| $1.3$B (dense) | $8.48$ | $885.32$ | $2.81$ |
| $1.3$B ($16$ experts) | $7.42\,\color[rgb]{0.1,0.5,0.1}\definecolor[named]{pgfstrokecolor}{rgb}{<br>0.1,0.5,0.1}{\downarrow\,12.53\%}$ | $885.32+4.87\,\uparrow\,0.55\%$ | $2.81+0.04\,\uparrow\,1.42\%$ |
| $1.3$B (dense) | $8.20$ | $1770.65$ | $2.81$ |
| $1.3$B ($32$ experts) | $6.76\,\color[rgb]{0.1,0.5,0.1}\definecolor[named]{pgfstrokecolor}{rgb}{<br>0.1,0.5,0.1}{\downarrow\,17.56\%}$ | $1770.65+18.94\,\uparrow\,1.07\%$ | $2.81+0.08\,\uparrow\,2.85\%$ |

##### Regarding Computational Overhead

As demonstrated in Table 3, the $1.3$B 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, a $335$M parameter mixture of $32$ experts reaches the same performance level as the $1.3$B parameter model, but with a comparable training budget and three times lower inference cost.
Despite the $335$M parameter model incurring a higher percentage cost during inference, this computational overhead is not the definitive minimum necessary for effective routing: as detailed in § 3.4, our results indicate that the routing overhead can be 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.

### A.4 Communication overhead

In this section, we provide an estimate of the communication overhead involved in our model training 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-optimal, meaning that the total number of bytes transferred per node is approximately $2K$, where $K$ is the size of the message (Thakur et al., ; Patarasuk & Yuan, ).

##### Mixture of LLMs

Every router performs all gather operation to share the loss on the dataset chunk approximately every $T=45$M training tokens, as outlined in 3.2. This frequency was primarily determined by implementation specifics related to data storage.

We can define the number of tokens processed per step by a router as $N_{\text{tokens_per_step}}$, which can be calculated as:

$$ $N_{\text{tokens_per_step}}\leq S\times B_{r}$ $$

This leads to the minimum number of steps between communications, $N_{\text{steps_per_com}}$, given by:

$$ $N_{\text{steps_per_com}}\geq\frac{T}{S\times B_{r}}$ $$

The maximum number of communication events, $N_{\text{comm}}$, during the router’s training can then be estimated as:

$$ $N_{\text{comm}}\leq\frac{N_{\text{steps_router}}}{N_{\text{steps_per_comm}} }=\frac{N_{\text{steps_router}}\times S\times B_{r}}{T}$ $$

As shown in App. Table 2 we train routers for $128,000$ steps with batch size $32$. We use context size $1024$. Therefore,
$N_{\text{comm}}\approx 94<1\times 10^{2}$.

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 $\frac{T\times N}{S}$.

Each router needs to send and receive the loss values for each sequence to and from other routers, resulting in a factor of 2 for the data sent and received. Therefore, the data per router can be approximated as:

$$ $\text{Data per router}\approx 2\times 2\times\frac{T\times E}{S}$ $$

Given $T=45\times 10^{6}$ – training tokens per router between communication steps, $E\leq 32$ – number of experts and $S=1{,}024$ – sequence length, we can estimated the data transferred per node as:

$$ $\text{Data per router}\leq 2\times 2\times\frac{T\times E}{S}=4\times\frac{45 \times 10^{6}\times 32}{1{,}024}=5.625\text{MB}$ $$

Once routers are fully trained, they communicate only when necessary, such as when experts require new data for training.

For instance, consider that the message size is $K$ bytes, equivalent to loss array of $K/2$ sequences in float16 precision. The frequency of 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:

$$ $N_{\text{steps_per_comm}}\approx\frac{K}{2\times B\times E}$ $$

Thus, communication can be scheduled approximately every $\frac{K}{4\times B\times E}$ expert’s training steps. Assuming $K=2^{31}-1$ bytes, $B=128$ as used in the majority of our experiments, and $E\in\{4,8,16,32\}$ (see Appendix 2), the communication interval becomes:

$$ $\frac{2^{31}-1}{2\times 128\times E}\approx\frac{2^{23}}{E}\leq 2^{18}\quad \forall E\in\{4,8,16,32\}$ (17) $$

##### Comparison with Distributed Training

Let us assume that we train a model with a 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 fit on one GPU. To compare communication costs, we examine the additional expense incurred when increasing throughput by a factor of $E$, i.e., utilizing $E$ hosts instead of one.

In regular distributed training, scaling to $E$ hosts necessitates copying all weights to each host and performing gradient synchronization at every training step. Assuming each parameter is represented by $32$ bits (as optimizers typically operate in float32), the amount of data transferred to and from each node during gradient synchronization can be calculated using the bandwidth-optimal algorithm:

$$ $\text{Data per node}=2\times W\times 4$ $$

Here, $W\times 4$ represent the size of the model gradients in bytes.

If $W=1.3\times 10^{9}$ parameters, then the amount of data transferred per node per training step is:

$$ $\text{Data per node}=2\times 1.3\times 10^{9}\times 4=10.4\text{GB}$ $$

Therefore, in regular distributed training, each node must send and receive gigabytes of data at every step. In contrast, our approach offers flexible communication strategies. We can transfer large volumes of data only a few times during training (see Equation 17), or optimise for more frequent communications with smaller data sizes. This flexibility allows us to adapt based on network latency and other infrastructure constraints.

## Appendix B Downstream evaluation

For the evaluation, we utilized the lm-eval-harness software (Gao et al., ). For the computation of perplexity, we compute the perplexity using the format: “Question: {question}. Answer: {answer}.”

**Table 4: Downstream Evaluation Results (Part 1). This table presents a comparison in accuracy on downstream tasks between a $1.3$B-parameter, $32$-expert model and a $1.3$B-parameter dense model. Both models were trained using the $2$T tokens and nearly the same training FLOPs, with only a $1\%$ difference. The inference cost of the expert model increased by less than $3\%$ compared to the dense baseline FLOPs.**
| Task | 32 experts, 2T | dense, 2T |
| --- | --- | --- |
| mmlu abstract algebra | $0.1400$ | $\mathbf{0.2300}$ |
| mmlu astronomy | $\mathbf{0.3553}$ | $0.2829$ |
| mmlu business ethics | $0.4900$ | $\mathbf{0.5300}$ |
| mmlu college computer science | $0.2500$ | $\mathbf{0.2900}$ |
| mmlu college mathematics | $\mathbf{0.1500}$ | $0.1400$ |
| mmlu computer security | $0.3500$ | $\mathbf{0.3800}$ |
| mmlu elementary mathematics | $\mathbf{0.2672}$ | $0.2487$ |
| mmlu formal logic | $\mathbf{0.3016}$ | $0.2937$ |
| mmlu global facts | $\mathbf{0.3100}$ | $0.2800$ |
| mmlu high school geography | $\mathbf{0.3586}$ | $\mathbf{0.3586}$ |
| mmlu high school macroeconomics | $\mathbf{0.3128}$ | $0.2872$ |
| mmlu high school microeconomics | $\mathbf{0.3193}$ | $\mathbf{0.3193}$ |
| mmlu high school statistics | $\mathbf{0.2917}$ | $0.2593$ |
| mmlu human aging | $\mathbf{0.4081}$ | $0.3857$ |
| mmlu human sexuality | $\mathbf{0.3969}$ | $0.3893$ |
| mmlu international law | $\mathbf{0.2314}$ | $\mathbf{0.2314}$ |
| mmlu logical fallacies | $\mathbf{0.3067}$ | $0.2945$ |
| mmlu machine learning | $0.2589$ | $\mathbf{0.2679}$ |
| mmlu medical genetics | $\mathbf{0.3900}$ | $0.3500$ |
| mmlu moral disputes | $0.2803$ | $\mathbf{0.2890}$ |
| mmlu moral scenarios | $\mathbf{0.2380}$ | $\mathbf{0.2380}$ |
| mmlu nutrition | $\mathbf{0.3039}$ | $0.2516$ |
| mmlu prehistory | $\mathbf{0.4167}$ | $0.3858$ |
| mmlu professional accounting | $\mathbf{0.2872}$ | $0.2518$ |
| mmlu professional medicine | $\mathbf{0.3566}$ | $0.2904$ |
| mmlu professional psychology | $\mathbf{0.3203}$ | $0.3186$ |
| mmlu world religions | $\mathbf{0.5731}$ | $0.4503$ |
| mmlu anatomy | $\mathbf{0.4667}$ | $0.3926$ |
| mmlu clinical knowledge | $\mathbf{0.3434}$ | $0.2755$ |
| mmlu college biology | $\mathbf{0.4028}$ | $0.3819$ |
| mmlu college chemistry | $\mathbf{0.2800}$ | $0.2600$ |
| mmlu college medicine | $\mathbf{0.2832}$ | $0.2659$ |
| mmlu college physics | $0.1863$ | $\mathbf{0.1961}$ |
| mmlu conceptual physics | $\mathbf{0.3787}$ | $\mathbf{0.3787}$ |

**Table 5: Downstream Evaluation Results (Part 2)**
| Task | 32 experts, 2T | dense, 2T |
| --- | --- | --- |
| mmlu econometrics | $\mathbf{0.2719}$ | $0.2018$ |
| mmlu electrical engineering | $0.2552$ | $\mathbf{0.2621}$ |
| mmlu high school biology | $\mathbf{0.3290}$ | $0.3097$ |
| mmlu high school chemistry | $\mathbf{0.2167}$ | $0.1970$ |
| mmlu high school computer science | $\mathbf{0.2400}$ | $\mathbf{0.2400}$ |
| mmlu high school european history | $0.2848$ | $\mathbf{0.3212}$ |
| mmlu high school government and politics | $\mathbf{0.4093}$ | $0.3834$ |
| mmlu high school mathematics | $\mathbf{0.1519}$ | $0.1444$ |
| mmlu high school physics | $\mathbf{0.2980}$ | $0.2715$ |
| mmlu high school psychology | $\mathbf{0.4642}$ | $0.4330$ |
| mmlu high school us history | $\mathbf{0.3382}$ | $0.3137$ |
| mmlu high school world history | $0.2827$ | $\mathbf{0.2996}$ |
| mmlu jurisprudence | $\mathbf{0.2407}$ | $0.2222$ |
| mmlu management | $0.3786$ | $\mathbf{0.3981}$ |
| mmlu marketing | $\mathbf{0.4915}$ | $\mathbf{0.4915}$ |
| mmlu miscellaneous | $\mathbf{0.5121}$ | $0.4994$ |
| mmlu philosophy | $0.2958$ | $\mathbf{0.3119}$ |
| mmlu professional law | $\mathbf{0.2634}$ | $0.2536$ |
| mmlu public relations | $0.3545$ | $\mathbf{0.4273}$ |
| mmlu security studies | $\mathbf{0.3224}$ | $0.3020$ |
| mmlu sociology | $\mathbf{0.2786}$ | $\mathbf{0.2786}$ |
| mmlu us foreign policy | $0.3300$ | $\mathbf{0.3400}$ |
| arc challenge | $\mathbf{0.3200}$ | $0.2892$ |
| arc easy | $\mathbf{0.6780}$ | $0.6549$ |
| hellaswag | $\mathbf{0.4912}$ | $0.4663$ |
| sciq | $0.8950$ | $\mathbf{0.9130}$ |

## Appendix C Additional Results

##### Effect of Prefix Length on Routing Performance

We also conducted experiments by partitioning the data using a shorter prefix length of $M=32$ tokens during training, whereas in our main results (see § 3), we used a prefix length of $M=256$ tokens.

Figure: Figure 6: Test perplexity as a function of routing prefix length during inference for $1.3$B parameter models with $16$ experts. The orange curve represents a prefix length of $32$, while the blue curve corresponds to a prefix length of $256$.
Refer to caption: x16.png

We observed that training with a smaller prefix length resulted in better expert performance when routing during inference was done using a small number of tokens (see Figure 6). This suggests that using a shorter prefix during training enhances the router’s ability to make effective decisions with limited context.

We did not take any action to make sure that the router models could extrapolate to longer sequences than seen in training (e.g. interpolating positional encodings, using ALiBi (Press et al., ), etc.). As expected, we observe a performance degradation when making routing decisions for sequences longer than the ones used for training.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: How were the number of tokens chosen for each configuration (see Weaknesses section above)?
soundness: 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.

The 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.

I 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.
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.
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.

Throughout, 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.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
questions: Minor edits:
- Line 300: glitch in the reference for Adam Optimizer
- Figure 3 is difficult to read, especially the legends


Questions:
- 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
- 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?
soundness: 4
strengths: - The authors show that the same perplexity can be achieved with three times less compute cost compared to the dense baseline

- The final performance (test PPL) is not sensitive to the router size hence, making the approach succeed with even tiny routers (4M parameters)

- 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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
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?

- I want to raise the weaknesses above in case the authors have any further comments.
soundness: 3
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. 

- 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. 

- The paper is quite easy to follow, and the overall idea is quite simple and seems reasonable.
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).
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.

- 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.

- 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.

Meta-review: This paper introduces a more parallelizable MoE system. The approach is very simple, introducing a single router to make asynchronous operation possible.

Pros: Appears to significant improve efficiency of an MoE. Very simple approach, can be further modified.

Cons: Longer prefix length makes it harder to reliably select the right expert. Baseline comparison is against a dense model with fewer total parameters than the entire mixture.

Decision: Accept (Spotlight)