Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM
Authors: Sainbayar Sukhbaatar
Olga Golovneva
Vasu Sharma
Hu Xu
Xi Victoria Lin
Baptiste Roziere
Jacob Kahn
Shang-Wen Li
Wen-tau Yih
Jason E Weston
Xian Li
Venue: COLM (2024)
Keywords: mixture of experts,continued pretraining
TL;DR: BTX improves a generalist LLM on multiple fronts. It trains expert LLMs in parallel for new skills in domains such as math, code & world knowledge, and later mix them together & finetune as a Mixture-of-Experts.

Abstract: We investigate efficient methods for training Large Language Models (LLMs) to possess capabilities in multiple specialized domains, such as coding, math reasoning and world knowledge. Our method, named Branch-Train-MiX (BTX), starts from a seed model, which is branched to train experts in embarrassingly parallel fashion with high throughput and reduced communication cost. After individual experts are asynchronously trained, BTX brings together their feedforward parameters as experts in  Mixture-of-Expert (MoE) layers and averages the remaining parameters, followed by an MoE-finetuning stage to learn token-level routing. BTX generalizes two special cases, the Branch-Train-Merge method, which does not have the MoE finetuning stage to learn routing, and sparse upcycling, which omits the stage of training experts asynchronously. Compared to alternative approaches, BTX achieves the best accuracy-efficiency tradeoff.

Published as a conference paper at COLM 2024
Branch-Train-MiX:
Mixing Expert LLMs into a Mixture-of-Experts LLM
Sainbayar Sukhbaatar Olga Golovneva
Vasu Sharma
Hu Xu Xi Victoria Lin Jason Weston Xian Li
Baptiste Rozière
Jacob Kahn
Daniel Li
Wen-tau Yih
FAIR at Meta (sainbar,xianl) cttbbaaxiannl@mtt.ccomm @meta.com
Abstract
We investigate efficient methods for training Large Language Models (LLMs) to possess capabilities in multiple specialized domains, such as cod- ing, math reasoning and world knowledge. Our method, named Branch- Train-MiX (BTX), starts from a seed model, which is branched to train experts in embarrassingly parallel fashion with high throughput and re- duced communication cost. After individual experts are asynchronously trained, BTX brings together their feedforward parameters as experts in Mixture-of-Expert (MoE) layers and averages the remaining parameters, followed by an MoE-finetuning stage learn token-level routing. BTX generalizes two special cases, the Branch-Train-Men not have the MoE finetuning stage learn routing, and sparse upcycling, which omits the stage of training experts asynchronously. Compared to alternative approaches, BTX achieves the best accuracy-efficieng tradeoff.
method, which does
Introduction
In recent years, Large Language Models (LLMs) have shown impressive performance in wide-range of tasks (Brown al., 2020; Touvron al., 2023; Achíam et al., 2023), including code generation (Li et al., 2022b; Rozière et al., 2023), solving math problems (Azerbayev et al., 2023), multilinguality (Zhao al., 2024), etc. Training such LLMs requires large amount of compute and data, exceeding thousands of GPUs and trillions of tokens. The training parallelization is typically done by maintaining multiple copies of the model on different GPUs and keeping them synchronized after each weight update. The cost of this frequent communication is the main bottleneck in scaling the training to more GPUs. Besides this issue, synchronized training is more vulnerable to hardware failures single failed GPU can cause the whole training to halt (Zhang et al., 2022; Gemini Team, 2023).
Recent work by Li et al. (2022a) proposed the Branch-Train-Merge (BTM) method for em- barrassingly parallel training of LLMs without any synchronization for improving the throughput of pretraining. It starts by creating multiple copies of seed LLM, then sepa- rately training each copy on different subsets of data. This results in multiple independent LLMs that do not share any parameters and each LLM is an expert specializing in its own data distribution, such as knowledge domains, languages or even modalities. At test time, an input prompt is classified into one or more of the domains, and then the final outputs are formed from the corresponding expert models which are combined to predict the next token. While this approach makes training more efficient, its main drawback is the lack of unified single model making it impossible to do further supervised finetuning (SFT) or reinforcement learning from human feedback (RLHF) finetuning (Ouyang et al., 2022), both of which can boost performance further, and are crucial steps in building aligned LLMs.
A separate line of work for reducing the computational footprint of LLMs is the Mixture- of-Experts (MoE) approach (Jacobs et al., 1991; Shazeer et al., 2017), where only a subset of parameteters are active at any given time. In particular, MoE is applied to the feedforward sublayer of Transformers (Fedus et al., 2022; Roller et al., 2021; Lewis et al., 2021), allowing

---

Published as a conference paper at COLM 2024
Data
Expert LLM
MoE
FF
Feedlorward
layer
Expert LLM
FF 2
LLM
router
Branch from a seed model
Expert
Dala
LLM
FF N
Train experts separately on their corresponding data
Mix the experts an unifled Misture-of-Expert model and finetune
Figure 1: The Branch-Train-Mix (BTX) method has three steps: 1) branch from a pretrained seed LLM by making multiple copies of it; train those copies separately on different subsets of data to obtain expert LLMs; 3) mix those expert LLMs by combining them into single LLM using mixture-of-exxert feedforward (FF) layers, and finetuning the overall unified model.
the total number of parameters to grow without additional computation. LLMs scaled in this way have shown impressive performance on downstream tasks (Jiang et al., 2024; Xue et al., 2024). Unlike synchronized fashion, and the communication cost increases with the number of experts due to all-to-all communication.
Mixture-of-Expe -ffexessss are often trained in fully
In this paper, we aim for the best of both worlds, combining the advantages of Branch- Train-Merge and Mixture-of-Expere by training multiple expert LLMs separately as in the Branch- MaanchhTrai---erge method, but subsequently combine those experts into single model using MoE architecture. More specifically, the feedforward sublayers from all the expert LLMs are brought together into single MoE module at each layer, and router network selects which feedforward expert to use at every token. We merge other modules of the expert LLMs, including self-attention layers, by simply averaging their weights. Then the resulting model is MoE-finetuned on all the combined data by continuing training, so that the router can learn to mix the expert feedforward (FF) modules. Figure 1 shows an overview of this method, which we call (BTX).
while mitigating their disadvantages. We achieve this
The main advantage BTX compared to MoE that expert training embarrassingly par- allel and asynchronous, reducing communication cost and increasing training throughput. Compared Branch-Train--e Train-Merge, the final BTX model a unified neural network that can be finetuned used like any other standard LLM. The final BTX model will not significantly increase inference FLOPs compared to the seed model since is sparsely activated, despite having a much larger number of parameters.
We conduct our experiments using LLAMA-2 7B (Touvron et al., 2023) as a seed model and train expert LLMs on different subsets of data corresponding to the domains of math, code and Wikipedia. With the original LLAMA-2 7B weights added as a fourth expert, we finetune the combined MoE model for relatively short period compared the pretraining process. The resulting BTX model brings significant improvements over the seed model on tasks across various domains, especially bridging the gap with specialized models on math and code related tasks, while retaining performance on the original capabilities where specialized models suffer from catastrophic forgetting, BTX outperforms BTM on tasks demonstrating the benefits of learnt routing through MoE finetuning. Compared to purely MoE training such as sparse upcycling, BTX more compute efficient with higher training throughput and more balanced performance across tasks in different domains.
Related Work
Asynchronous parallel training Reducing communication between training workers for computational efficiency a major topic of study for training deep learning systems. Zhang et al. (2015) introduced method that allows model instances different workers to diverge
2

---

Published as a conference paper at COLM 2024
from each other, thus eliminating the constant need of synchronization, Instead, the workers are loosely synchronized to master weights using elastic averaging from time to time. A more recent work by Douillard et al. (2023) showed that less frequent synchronization of diverged workers by averaging their weight changes and applying Nesterov momentum works well in practice for training LLMs. The
method (Li et al., 2022a; Gururangan et al., 2023) takes parallel training to the extreme by running multiple training processes completely independently. Each training process uses specific domain data, thus the corresponding model becomes an expert in that domain. Finally, the output distributions of those expert models are averaged to make a next token prediction. Which experts to average is decided by classifying the input into one or more of the domains. Wortsman et al. (2022) showed simply averaging parameters of separately trained models improves performance, but the models only differed in their hyperparameters
Mixture-of-Experts MoE is used to scale deep networks in Shazeer et al. (2017) using simple Top-K routing scheme. Since the routing decisions are discrete and thus cannot be trained by gradient descent, various training methods have been explored for the Trans- former architecture (Fedus et al., 2022; Lewis et al., 2021). Surprisingly Roller et al. (2021) showed that even a fixed routing scheme without any learning works well, if the routing is done via random mapping based on input tokens. In larger scale experiments with recent LLMs, Jiang et al. (2024) demonstrated that the MoF approach can match the performance of dense LLM counterparts using much smaller number of active parameters. A study by Dai et al. (2024) showed the advantage of more fine-grained experts, as well as having a shared expert that always stay active. More similar to our work, Gururangan et al. (2021) makes experts in feedforward layers specialize to specific domains using a domain-conditioned fixed routing, but it lacks the asynchronous training of our approach.
Continual learning Our method relates to continual learning (Awasthi & Sarawagi, 2019) because domain experts are trained on datasets with different distributions from the initial data used for training the seed model, which is implemented by continued training after branching. Specifically, our approach is related to parameter isolation methods (Lange et al., 2019) as we have different parameters for different domains. Aljundi et al. (2016) also creates a new copy of a model to train on each domain. Rusu et al. (2016) adds a new model with a new domain, but connects it to the previous models so the previously learned features can be used. Rozière et al. (2023) showed continual training of a seed LLM on a specific domain of code can produce strong domain expert model, and this converges much faster than starting from scratch. For training math expert, starting from a code expert rather than a general LLM was shown to be more beneficial (Shao et al., 2024; Azerbayev al., 2023).
3 Branch-Train-Mix
Given an existing LLM M which has been pretrained on a large corpora covering wide variety of topics, we aim to improve its performance on N areas of expertise. This is achieved by continued pretraining with corresponding training datasets D := {Di,..., Dr J, each related to a specific knowledge domain such as math, code, etc. The proposed method contains three stages: Branch, Train, and MiX.
3.1 Branch & Train: Embarrassingly Parallel Expert Training
Initializing from the seed model M, we train N expert LLMs (MI...,,N'I) with each model M; being trained on the corresponding dataset D, in the same manner as during pretraining, using the usual LM objective. Since each expert model M; can be trained in complete separation from the others, the whole training process becomes N-way embarrass- ingly parallel. This training paradigm has several benefits in large-scale distributed training. It allows linear scaling of overall training throughput when scaling up the size of compute, while joint training often faces uncertain performance from increasing batch size. It has lower all-to-all communication cost. It is also more resilient, as a single training failure will only affect one of the N training processes instead of halting the entire training.
After all the expert training is finished, we will end up with N different LLMs, with each specializing in a specific distribution. At this point, the Branch-Train -eege- method (Li et al.,
3

---

Published as a conference paper at COLM 2024
2022a; Gururangan et al., 2023) uses these domain experts as is, choosing which expert to use by determining which domain the input belongs to at inference time. Usually multiple experts are chosen, and their final output distributions are simply averaged to generate the next token. Our BTX approach, in contrast, merges these domain experts back into a single LLM that is finetuned further, as will describe in the next section.
3.2 MiX: Combining Separate Experts to be Mixture-of-Expeers
We employ a Mixture-of- However, instead of using the classical procedure of mixing the final outputs from M,, we do a more fine-grained mixing by performing MoE within each layer of a Transformer. In particular, we combine the different feedforward sublayers from the domain experts into single MoE sublayer. (x) is the feedforward sublayer at the 1-th layer of the i-th domain expert Mi, then the combined MoE layer for input representation at layer I will compute:
approach to combine the domain expert models Mj-
FF FF (x)
E g: (x)
Here W/ a linear transformation and g a routing function, which usually has sparse output and hence switches on only some experts. Since we can skip computing Fx) if the corresponding router output is zero, the actual computation of FF MaME (x) will be much more efficient than computing all domain experts. However, routing decisions can change from token to token, one input sequence can employ all the domain expert FF layers if needed, even when only a few are accessed at any given token. In our experiments, we use Top-k (k=2) routing where gWix) = SoftMax CatiMaxxXoopp((/ (((op(((x)),, unless otherwise stated.
For the self-attention sublayers, we combine the different domain experts by simply aver- aging their weights. The motivation behind this is the assumption that the self-attention layers are less domain specialized than the feedforward layers. We do the same averaging for the remaining parameters (embeddings, etc.) well.
Note that the only new parameters we introduce are the router's transformation parameters Wi, which are negligible in size compared to the rest of the network. Nevertheless, those new parameters need to be finetuned, so the router can make optimal decisions in selecting which domain FF; to use. weights are constructed by averaging, and are likely not optimal. Overall, the entire system has not been optimized for working together at all in the embarrassingly parallel training framework, but our hypothesis is that even small amount of combined finetuning might make large improvements.
addition, funetuning is helpful because the self-attention
3.3 Variations
We also experimented with several variations of our method.
Load balancing A common problem with MoE is the emergence of dead experts, which do not get activated by the router at all. Common routing methods like Top-k are unlikely to escape from such situation because a dead expert is never in the top-k selection, and therefore never receives training signal. Load balancing offers simple solution by adding an extra loss term that encourages the experts to be utilized equally. We use loss term similar to (Fedus al., 2022):
1
LB = aN Σ u;P; where 11 = B
(Wix) and p; xEB
Σ
Σ xEB
SoftMax
B
Here the current data batch, and a a hyperparameter. This loss is computed in each layer and added to the NLL loss.
Routing method Besides Top-k routing, we also experiment with other routing methods:
Switch: It is Top-1 routing method proposed by Fedus et al. (2022).
Soft routing: We use softmax the routing function g, so all experts are activated both during training and inference. While it is likely to provide the best performance, it comes at the expense of increased compute.
4

---

Published as a conference paper at COLM 2024
Math
Code
General knowledge
GSMSK
MATH
HvaanBval
MBPP
NQ
TriviaQA
MMLU
LLAMA-2 7B
14.7
2.5
12.8
20.8
16.4
58.5
46.1
Math expert
39.5
18.8
25.0
33.6
14.4
37.1
52.0
Code expert
12.0
4.0
31.7
40.2
11.5
29.9
39.6
Wikipedia expert
11.7
3.1
11.0
15.2
21.8
57.2
43.1
BTX, 1op-2
37.1
17.8
28.7
39.4
24.8
57.1
52.5
Table 1: Individual domain expert LLM performance representative tasks, compared to the seed model LLAMA-2 7B. As expected, the code and math experts excel at their corresponding domain tasks. The Wikipedia expert performs better on Natural Questions (NQ), but the math expert has the best score on MMLU (MMLU contains many math subjects and math training is shown to help on this task (Shao et al., 2024)).
Sample Top-1: We use the gumbel softmax (Jang et al., 2016) for g. At training time, we generate soft sample from the gumbel softmax, but zero out all its values except the largest one. Then we compute only one expert corresponding to this largest value, omitting the other expert computations. At inference time, simply do hard sampling. We anneal the temperature to a sharp distribution at the end of training to gradually reduce the discrepancy between training and inference.
4 Experiments
4.1 Experimental Setup
We base our experiments on the setup used for LLAMA-2 pretraining (Touvron et al., 2023). In particular, use the LLAMA-2 7B model as seed model.
4.1.1 BTX Training
We use the pretrained Llama-2 (Touvron et al., 2023) with 7B parameters as our seed model. After making three copies of the seed model LLAMA-2 7B, we continue training them on the following domain datasets to derive three domain experts:
Math: The same data sources and mixture used in Llemma (Azerbayev et al., 2023) model training. To be comparable to Llemma, we train on the same amount of data as well, i.e. 48k steps with 201B tokens in total.
Code: The same data sources and mixture of code data used in CODELLAMA pretraining (Rozière et al., 2023). The code expert LLM is trained for 50k steps with 210B tokens in total to be comparable with the math expert.
Wikipedia: Wikipedia documents extracted between June to August 2022. The data was preprocessed remove hyperlinks, comments and other formatting boilerplate. Since this is a smaller dataset, we train on a total of 42B tokens.
While we can proceed with only these three domain experts, we also include the original seed LLM as "generalist" expert that its general knowledge is transferred the final model. Thus we mix these four expert models into a single MoE model as described in Section 3.2. Then we finetune this MoE model on all the data sources used train the four experts (including the original LLAMA-2 7B pretraining data for the generalist expert) and train for another 80B tokens. The detailed sampling ratio of datasets as well as other training details are described in Appendix A.
4.1.2 Baselines
We compare to the following baselines:
LLAMA-2: We compare the original LLAMA-2 7B that we use as a seed model, as well LLAMA-2 13B.

---

Published as a conference paper at COLM 2024
Math
Code
Knowledge
Reasoning
MMLU
Average
Specialized L.Mss
CODELLAMA 7B
8.1
36.3
22.2
56.6
38.6
37.9
LLEMMA 7B
28.0
33.5
17.2
38.8
33.5
32.1
Generalist LLMs
LLAMA-2 7B
8.6
16.8
37.4
63.3
46.1
40.7
LLAMA-2 13B
16.3
24.5
40.0
66.1
52.8
45.4
Dense (DM)
18.3
25.8
39.6
63.3
49.8
44.5
Sparse upcycling (DM), Top-2
28.1
34.7
34.0
62.3
51.1
46.3
BTM, Top-1
21.3
36.4
26.5
61.0
44.3
43.1
BTM, Top-2
21.5
36.6
26.9
61.2
44.3
43.4
BIX, Sample Top-1
26.4
31.5
40.1
63.7
53.2
47.3
BTX, Top-2
27.4
34.0
41.0
63.5
52.5
47.9
Table 2: Test performance of BTX compared various baselines, including both generalist and specialized pretrained models, on various capabilities aggregated across popular bench- marks. Dense, sparse upcycling, BTM and BTX are trained on exactly the same amount and mixture of data with the exception that BTM does not have the finetuning stage.
Dense: Instead of training separate LLMs on different domain datasets, the dense baseline continues to train the seed LLM with all the data. We use exactly the same training data as BTX, first training on the new domain-specific data used in the experts training stage, followed by the same data mixture that includes the LLAMA-2 pretraining data the MoE finetuning stage. We call this comparison data-matching (DM).
Sparse upcyeling: This baseline (Komatsuzaki et al., 2022) initializes MoE model from the seed model by making 4 identical copies of the feedforward module as experts. We use the Top-2 router with randomly initialized W, parameters. In addition training a data matching (DM) baseline with the same data as is used in BTX and the dense baseline, we also train sparse upcycling baseline with the same amount of GPU-days, ie. compute-matching (CM), using the MoE finetuning data mixture throughout training. This is equivalent to special case of BTX which does contain embarrassingly parallel expert training.
Branch- Teae--raaiiiiMM Merge (BTM): This baseline (Li et al., 2022a) uses the same expert LLMs as BTX (including the original seed model) but uses them directly without building a MoE model. For given context (input), it selects Top-k expert LLMs based on the similarity between the context and experts' training data. Following the efficient inference method used in Gururangan embedded via tf-idf. Top-k experts are selected based on cosine similarity to the mean tf-idf embedding of each expert.
al. (2023), both context and experts' training data
CodeLlama 7B: A language model specializing in code (Rozière et al., 2023) by continued training of the same seed model LLAMA-2 7B on code data. It also has other features such as long-context and infilling.
Llemma 7B: A language model specializing in mathematics (Azerbayev et al., 2023) by continued training of CodeLlama 7B on math data.
4.1.3 Evaluation
For evaluation, use the zero- and few-shot performance multiple benchmarks that test different skills: GSM8K and MATH for math reasoning; HumanEyal and MBPP for code generation; Natural Questions and TriviaQA for world knowledge; ARC-Easy, ARC- Challenge, SIQA, PIQA and WinoGrande for reasoning; and MMLU which covers multiple domains. See Appendix B for more evaluation details.

---

Published as a conference paper at COLM 2024
BTX
Math
Sparse Upcycling (CM)
46
0.4
Dense (DM)
Average
44
Llamal2 13B
8TM
42
Knowledge
Ressoning
L 7B
(seed)
40
640
1000
1600
2000
2500
3000
Liama-2 78
Codeliama 78
Train GPU days (from the seed madel)
138
The average performance vs training budget of BTX compared to various baselines, with different active parameters at inference time indicated by circle size. All the models except LLAMA-2 13B are trained starting from LLAMA-2 7B using the datasets described in Section 4.1.1. The X-axis shows the total training compute starting from the seed model measured in GPU days², and the Y-axis is the average score over all the tasks (as computed in Table 2). The BTX models outperform the baselines that started from the same seed model, as well as LLAMA-2 13B. Right: The normalized performance over different domains where the scores are divided by the highest one. We see large improvements for BTX in code (which matches the specialized model) and math tasks compared to the seed model LLAMA-2 7B, even outperforming the LLAMA-2 2 13B model.
Figure 2: Left:
4.2
Main Results
4.2.1
Overall Performance
Domain experts excel at their respective tasks. We first analyze how expert LLMs spe- cialize to specific domains. Results are summarized in Table 1. As expected, individual expert LLMMs achieve the best performance in their respective domain, where the math and code domains see especially large improvements. ing observations. We see that the math expert training improved its code performance as well, indicating a close relation of these domains. However, such single-domain continued training also suffers from catastrophic forgetting with significant performance drops on some tasks in other domains, For example, the math and code expert are much worse on TriviaQA than the seed model.
addition, there are several interest-
BTX improves all tasks where experts specialize. Table 2 and Figure 2 (right) show aggre- gated performance across multiple domains. More detailed per-task results are reported in Table 5 in the Appendix. Compared to the seed model LLAMA-2 7B, BTX models (both Sample Top-1 and T corresponding to different number of active parameters) improve on all expert domains, such as math, coding and world knowledge without regressing other tasks such as commonsense reasoning. BTX with Top-2 experts (our default) also approaches the best performance of the specialized models LLEMMA 7B and CODELLAMA 7B in the math and coding domains, while drastically improving over those models on domains that not their speciality such as world knowledge and commonsense reasoning. Compared to alternative data-mag (DM) methods for continued pretraining such as dense and sparse upcycling, BTX achieves better performance on average with small gaps in the math and coding domains. BTX outperforms BTM by large margin on average, indicating that MoE finetuning to learn token-level routing is beneficial. Overall, the results demonstrate that BTX is better performing method for continued pretraining which is
2The GPU days Llama-2 13B is an an approximate measurement, calculated by doubling the training compute of a 7B model trained with the same amount of pretraining data (according to louvron et al. (2023) Table 2). Since 1 13B is not trained from the seed model, we simply report their difference in GPU days.
7

---

Published as a conference paper at COLM 2024
Active parameters (B)
Average score
Routing method
Training
Inference
train=10B
train=40B
train=80B
Switch Top-1
6.7
6.7
24.7
Soft routing
19.7
19.7
35.8
37.3
Sample Top-1
6.7
6.7
33.0
35.3
36,9
Top-2
11.1
11.1
34.6
35.9
37.3
Table 3: Ablations on different routing methods during BTX training for increasing amount of train tokens. Average score is based performance on representative tasks including GSM8K, HumanBval, Natural Questions, ARC Challenge and MMLU.
robust to task interference from multi-task learning. BTX also outperforms LLAMA-2 13B on all tasks except reasoning, even though LLAMA-2 13B uses significantly more training compute and has slightly more active parameters.
We further compare BTX with the sparse upcycling baseline in the compute-matching (CM) scenario. Both train on the same data mixture during the MoE stage, but differ in terms of the percent of compute spent on MoE training. While sparse cycling performs close behind BTX, the parallel training of experts increases the training throughput of BTX, as is shown in Table 6 in the appendix. As a result, BTX can train with more than 2x the data than pure MoE given the same training compute budget, and achieves slightly higher average performance across all domains.
In Table 1, we directly compare BTX with the original expert LLMs. While we see some degradation in the math and code domains, the final BTX model has the best overall performance and competitive with the expert LLMs on their respective domain tasks.
4.2.2 Better compute-performen tradeoff
We compare BTX with baselines terms of compute efficiency in Figure (left). The X-axis shows the total training compute starting from the seed model measured in GPU days, which includes the domain expert training and finetuning of the MoE model. The Y-axis measures the overall performance reported in Table 2.
Better performance than dense and BTM. Despite that the MoE training stage uses fraction of the total training budget in pretraining (for example, LLAMA-2 pretraining uses 2T tokens), BTX brings steep improvements on general capabilities compared to alternative continued pretraining approaches such as multi-task learning of the dense model and Branch- Maanc--Tain-Mergee.
More efficient than sparse upcycling. As a special case of BTX, sparse upcycling without expert training outperforms dense and BTM but not BTX, given the same or larger compute budget. The compute efficiency gains of BTX are from the embarrassingly parallel training of experts before MoE finetuning.
In terms of the active number of parameters (shown as circle sizes in 2 (left)), the MoE models are similar to the LLAMA-2 13B model. BTX uses less than half of the additional training compute compared to LLAMA-2 13B, but demonstrates improved performance on expert domains (math, code, and knowledge) and achieves better overall performance. This indicates that BTX's training is more effective for the late stage of pretraining than using the same training protocol throughout the entire of pretraining.
4.3
Ablations & Analysis
4.3.1
Ablations of BTX training
First, we compare the different routing methods with varying amount of active parameters for different amounts of finetuning. For fair comparison, load balancing is not used in any
Sample Top-1 training is more efficient, so we trained for 160B tokens for the same compute.
8

---

Published as a conference paper at COLM 2024
Top-2 routing with load- balanding
Top-2 routing with no load- baddbbalancing
Layer 16
Layer
Layer
Layer
0.4
0.2
0.0
Wiki
Math
Code
LLaMa-2
Will
Math
Gode
LLaMa-2
Wild
Math
Codo
LLaMa-2
Wiki
Math
Code LLaMa-2
Math domain
Coda domain
World knowledge
Reasening
Figure 3: BTX token routing decisions at middle (16) and top (32) layers to different experts (Wiki, Math, Code, LLAMA-2 7B) for different downstream tasks. The tasks are aggregated by domain: Code (Human Eval, MBPP), Math (GSM8K, MATH), World knowledge (Natural Questions, IriviaQA), and Reasoning (ARC-Easy, ARC-Challenge, SIQA, PIQA, and WinoGrande). We observe that Top-2 routing with load balancing (left) ensures a more uniform distribution of the load between experts compared to Top-2 without load balancing (right).
of them. Results are shown in Table 3. For Switch routing, we set its capacity factor to 1.5 (a hard limit after which routed tokens will be dropped). We found the Switch router to be subpar in average performance. The soft routing performs the best, but that is expected since lacks sparsity and has the highest number of active parameters. Overall, the Top-2 routing gives us a good balance between performance and efficiencyy
We also ablate additional design choices of BTX, with results summarized in Table in the appendix. We found that MoE training without load balancing performs worse on the coding task (HumanEval), but has higher math (GSM8k) accuracy. The routing analysis in the next section will give more insight into this trade-off. Next, freezing the feedforward modules initialized from each expert, and only training the rest of the MoE model has little impact performance across all tasks. This suggests that individual experts already gained sufficient domain knowledge during the branch-train stage, while the mix (MoE finetuning) stage mainly trains the other parameters such as averaged weights in the self-attention and the router transformations W,-
4.3.2 Routing Analysis
To gain an in-depth understanding of the performance of BTX, we run model evaluations on downstream tasks and examine the routing decisions among the experts. The results are summarized in Figure 3, and we also report detailed ablation results for different BTX setups in Appendix C. Compared to other routing methods, Top-2 routing with load balancing ensures a more uniform distribution probability distributions, we observe a shift towards low probability scores across all experts with load balancing, especially closer to the final layers the model, which contributes to the fair routing. interestingly, all models without load balance heavily rely on the Math expert, with low overall contribution from other experts, especially the Code expert. A dead Code expert comes "back to life" with load balancing introduced in training. In fact, not only becomes visible, but becomes the dominant expert in the math and code domains.
the load between experts. Analyzing the token
Examples of the routing decisions for Top-2 with load balancing can be found in Table 8 in the appendix. Overall across math domain tasks, tokens are often routed to the Code and LLAMA-2 7B experts. If look at a more detailed token distribution (Appendix C, Figure 6), we find that the GSMSK task prefers Code and LLAMA-2 experts, while the MATH task relies more on the in-domain Math expert. We hypothesise that this happens because the GSM8K dataset consists of grade school math problems that require common sense knowledge and basic arithmetic operations. Both the Code and World knowledge tasks mostly route to the in-domain Code and Wikipedia experts respectively. As observed earlier in Section 4.3.1, when load balancing is introduced, there are improvements in coding tasks but degradation in math tasks, which can be explained with these changes in domain expert routing. The reasoning tasks in contrast exhibit similar behaviour, and rely equally on Math and generalist LLM's expertise.

---

Published as a conference paper at COLM 2024
5 Conclusion
We introduced Branch-Train-MiXX (BTX), a simple continued pretraining method to improve an LLM's capabilities. It trains multiple copies of a seed LLM to specialize in multiple domains in an asynchronous and parallel fashion and later merges them back into single Mixture-of-Experi (MoE) model via finetuning. While the initial parallel training stage brings higher training throughput and scalability, the second MoE. finetuning stage makes the final LLM more performant. Our experiments suggest that a generalist LLM's perfor- mance can be boosted by continued training on datasets with specialized knowledge and skills using our method. We find that the BTX approach is more compute efficient than training larger generalist LLM or several separately specialized LLMs. These insights can inform how to allocate compute in late pretraining to achieve a strong generalist model.
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint urXiv:2303.08774, 2023.
Rahaf Aljundi, Punarjay Chakravarty, and Tinne Tuytelaars. Expert gate; Lifelong learning with network of experts. 2017 TEEE Conference on Computer Vision and Pattern Recognition URL ht http:p///ppp.ssmmtniccoooooo/o// www.steescaoollaae.....ee http:iccchhlaaaoooo...oooo org/CorpusIDD 91 4027.
(CVPR), pp. 7120-7129, 2016.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie J. Cai, Michael Terry, Quoc V. Le, and Charles Sutton. Program synthesis with large language models. ArXio, abs/ /b5//218.07732, 2021. URL //tps://api.
semanticschol. wtt.cccaa...rr//oorrsoooooo.eeee www.or//rrosss:22500 2371 42385.
Abhijeet Awasthi and Sunita Sarawagi. Continual learning with neural networks: A review. In Proceedings of the ACM India Joinf International Conference on Data Science and Management of Data, Pp. 362-365, 2019.
Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster, Marco Dos Santos, Stephen McAleer, Albert Q Jiang, Jia Deng, Stella Biderman, and Sean Welleck. Llemma: An open language model for mathematics. ArXiv, abs/2310.106311 2023. URL https://api. semanticscholar.c htt.a//aa.rgg//ropsss00:0:0- 26:172333
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, T. J. Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, 1lya Sutskever, and Dario Amodei. Language models are few-shot learners. ArXiv, abs/ 14165,, 2020. URL https: / //pi.. semanticscholar. org/CorpusID: :ttp:puuII::18811111 218971783.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde, Jared Kaplan, Harrison Edwards, Yura Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, David W. Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss,, William H. Guss, Alex Nichol, Igor Babuschkin, Suchir Balaji, Shantanu Jain, Andrew Carr, Jan Leike, Joshua Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew M. Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating large language models
10

---

Published as a conference paper at COLM 2024
trained code. ArXiv, abs/2107.033744 2021. URL https: http::///pp..smmanccchhhlarrrrr Corpus1D:235755472.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? Try ARC, the A12 reasoning challenge. arXiv preprint arXie:1803.0545) 2018.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint WrXiv:2110.1 10.6616.,, 2021.
Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. Deepseekmoe: Towards ultimate expert specialization in mixture-of-expert language models. ArXiv, abs/2401 0ets/4400.00666, 2024. URL https://api. semanti cachooaar..ll......o..ooooooaa.a wtt.aii......//r/rrrousss....... httg/CoopuusD::26633333
Mostafa Dehghani, Yi Tay, Anurag Arnab, Lucas Beyer, and Ashish Vaswani. The efficiency misnomer. In International Conference on Learning Representations, 2022.
Arthur Douillard, Qixuang Feng, Andrei A. Rusu, Rachita Chhaparia, Yani Donchev, Adhi- guna Kuncoro, Marc' Aurelio Ranzato, Arthur Szlam, and Jiajun Shen. Diloco: Distributed low-communication training of language models. ArXiz, abs/2311.0810 2023. URL https://api. comanticscholan or org hwtg//oorpsI:::26:5 265158012..
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232-52700 2022.
Gemini Team. Gemini: family of highly capable multimodal models. arv preprint arXiv:2312.11805, 2023. Team, Gemini and Anil, Rohan and Borgeaud, Sebastian and Wu, Yonghui and Alayrac, Jean-Baptiste and Yu, Jiahui and Soricut, Radu and Schalkwyk, Johan and Dai, Andrew M and Hauth, Anja and others.
Suchin Gururangan, Michael Lewis, Ari Holtzman, Noah A. Smith, and Luke Zettle- moyer. Demix layers: Disentangling domains for modular language modeling. In North American Chapter of the Association for Computational Linguistics, 2021. URL wrgg/CrppsIID 236976189.
https://api
semantisschollaar wew.aantiiccoola..oooo/oroorrooo
Suchin Gururangan, Margaret Li, Mike Lewis, Weijia Shi, Tim Althoff, Noah A Smith, and Luke Zettlemoyer. Scaling expert language models with unsupervised domain discovery. ar Xiv preprint ar Xrx::22303.1411 1ett27.3..117,,, 2023.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021a. URL https::/ wtwp:e/wwwwwtooommieiiemmmem wtt.://oorrmmmmdddddddmmmm1mmm1
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Xiaodong Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. Arviv, abs/ /b00..03.,,874 2021b. URL https://api. semanticscholar.
org/CorpusID: 232134851
Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adap- tive mixtures of local experts. Neural Compulation, 3:79-87, 1991. URL https://api.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. Xiv preprint Xxx::1611.01144 2016.
11

---

Published as a conference paper at COLM 2024
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, L'elio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of experts. ArXio, htbs/240.00400 2024. URL https://api. swwanticscholaat.. wtt.eeccoola..roooooooooooo org/CorpusID: ID:
26:884447778888877 266844877
Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. Triviaqa: large scale distantly supervised challenge dataset for reading comprehension. ArXiv, abs/ 1005//755.3351, 2017. URL https: //ppp./pp...a// semanticscholas httpricchoooooooorrrrrosss/////e 26501111:265111 26501419.
Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of-experts from dense checkpoints. ArXiv, abs/2212..552 2022. URL https:
s/ap...sematticscho... http:/ooaa...orrrroooss::::::::0 254533552....
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Matthew Kelcey, Jacob Devlin, Kenton Lee, Kristina N. Toutanova, Llion Jones, Ming-Wei Chang, Andrew Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: benchmark for question answering research. Transactions of the Association of Computational Linguistics, 2019.
Matthias De Lange, Rahaf Aljundi, Mare Masana, Sarah Parisot, Xu Jia, Ales Leonardis, Gre- gory G. Slabaugh, and Tinne Tuytelaars. A continual learning survey: Defying forgetting in classification tasks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44: 3366-3385, 2019. URL https: //atp.://ap..... www.nticscoolllattt.oaoaaaaaaaaa org wrw.aaarrrrgr///oooooooor CorpusID 218889911
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, 2021. URL httpss:/aapp semanticscholar :otp::roooo/oo:::: /tsp://saaaaaiiiioooooooooooooo/ 2324283311
Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A. Smith, and Luke Zettlemoyer. Branch-train-merge Embarrassingly parallel training of expert language models. ArXiv, abs/ abs/2208.03306, 2022a. URL https://ppi..mmmaticcchhh. semanticscholar
org/CorpusID: /Bt/Corpus1D:25137133333
Yujia Li, David H. Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom, Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, Thomas Hubert, Peter Choy, Cyprien de, Masson d' d'Autume, Igor Babuschkin, Xinyun Chen, Po-Sen Huang, Johannes Welbl, Sven Gowal, Alexey, Cherepanov, James Molloy, Daniel Jaymin Mankowitz, Esme Sutherland Robson, Pushmeet Kohli, Nando de, Freitas, Koray Kavukcuoglu, and Oriol Vinyals. Competition-level code generation with alphacode. Science, 378:1092 1097, 2022b. URL https://api semanticscholar. org/ wrg/ /Cop: 246527904.
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke E. Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Francis Christiano, Jan Leike, and Ryan J. Lowe. Training language models to follow instructions with human feedback. ArXio, abs0/220.021 2022. URL https: seapi..smaaticcccool.. watplaa..rrgrrrreere/eeee org/CCrpusID: 246426909
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. Hash layers for large sparse models. In Neural Information Processing Systems, 2021. URL https: sew.iissmmtticcccchieee wtw.oolaar.rrr/opuusttto..s 235367626.
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Tan, Yossi Adi, Jingyu Liu, Tal Remez, Rémez,Jérémy Rapin, Artyom Kozhevnikov, Evtimov, Joanna Bitton, Manish P Bhatt, Cristian Cantón Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre D'efossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. Code llama: Open foundation models for
12

---

Published as a conference paper at COLM 2024
code. ArXiv, abs/2308.12955 2023. URL 261 100919.
Andrei A. Rusu, Neil C. Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirk- patrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. Progressive neural net- works. ArXiv, abs/ 0es/16606.0667,, 2016. URL https://apt semanticscholle wwweaarrrgg/ooooo .rgg/CrppssID: 15350923.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communicatiios of the ACM, 64(9):99-106, 2021.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions. arXiv preprint ar arXit:1904.09728, 2019.
Zhihong Shao, Peiyi Wang, Qihao Zhu, R. X. Xu, Jun-Mei Song, Mingchuan Zhang, Y. Li, Yu Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. ArXiv, abs 02///2222220000 2024. URL https: //api. semanticscholar. wrg//orppus1D::27415501
Noam M. Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Ge- offrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- gated mixture-of-experts layer. ArXiv, abs/1701.06538, 2017. URL https://api.
semanticscholar.co hra.a/arr.rrg//ooousii:1:100 12462234.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs, Raphael Gontijo- Lopes, Ari S. Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. ArXiv, abs/ 2203.05482 2022. URL https://api. commanticscholan or wrgg/orppssspC 247362888..
Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You. Openmoe: An early effort on open mixture-of-experts language models. arXiv preprint ar Xio:2440..017333 2024.
Sixin Zhang, Anna E Choromanska, and Yann LeCun. Deep learning with elastic averaging sgd. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 28. Curran Asso- ciates, Inc., 2015. URL https://proceeddin http://rrcceddinns...ieeerriiips httppsscccpppee.iiiieeeee httpppprr..lles/pppr///211.. 1etea441155ffBB
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models. ArXiv, abs/2205.01 068, 2022. URL https: /ttppp:://aa/ //ppi. semanticscholar hwt.oo/ar.rrrgooouusii////999990
13

---

Published as a conference paper at COLM 2024
Jun Zhao, Zhihao Zhang, Qi Zhang, Tao Gui, and Xuanjing Huang. Llama beyond english: An empirical study on language capability transfer. arX preprint arXiv:2401.00 01005,, 2024.
A Training Details
Table 4 shows the exact data mixture ratios used in training each domain expert. For finetuning the MoE model, sample datasets that used to train math expert, code expert, wikipedia expert and the original LLAMA-2 7B with probabilities 30.16%, 40.31%, 10.30% and 19.23%.
Domain
Dataset
Sampling ratio (%)
A AlgebraicStack
13.57
OpenWebMath
54.27
Math
Arxiv
27.14
Github
2.99
Commoncrawl
5.01
Code
82.18
Code
Natural language related to code
9.90
Natural language
6.93
Wikipedia
90.91
Wikipedia
Commoncrawl
9.09
Table 4: Data sources and weights for domain experts.
For BTX with default Top-2 routing, we use load balancing with a = 0.01, unless other- wise stated. For the Sample Top-1 routing, r=max(0.5, -rf) from Jang et al. (2016) with r = 1e where is the number of training steps. For the first layer only, used soft-routing instead of Top-1.
the temperature annealing schedule
We use the same optimization hyperparameters for training of the baselines, expert models and MoE models. We use the AdamW optimizer with weight decay 0.1, and anneal the learning rate to the peak of le 4 with 100 steps of warmup, and decay to 10% of the peak with cosine schedule. We use batch size of 4M tokens with sequence length of 4096.
B
Evaluation
We use the following tasks for evaluation:
Math: we report the average performance on GSM8K (8 shot) (Cobbe et al., 2021) and MATH (4 shot) (Hendrycks et al., 2021b) for math reasoning.
Code: we report the average performance of HumanEval (0 shot) (Chen et al., 2021) and MBPP (3 shot) (Austin et al., 2021) for code generation.
World knowledge: we report the average performance of Natural Questions (5 shot)(Kwiatkowski et al., 2019) and TriviaQA (5 shot) (Joshi et al., 2017).
Reasoning: we report the average 0-shot performance of ARC-Easy and ARC-Challenge
(Clark et al., 2018), SIQA (Sap et al., 2019), PIQA (Bisk et al., 2020) and WinoGrande
(Sakaguchi et al., 2021).
General: we report performance on MMLU (5 shot) (Hendrycks et al., 2021a) which covers multiple domains.
We use the same evaluation metrics as is used in Touvron et al. (2023) and Rozière et al. (2023): for code tasks (HumanEval and MBPP) we report pass@1, for math tasks (GSM8k and MATH) and knowledge tasks (Natural Questions and TriviaQA) we report exact match, report accuracy for MMLU and ARC. We use greedy decoding for all generations. Detailed results all tasks reported in Table 5.
14

---

Published as a conference paper at COLM 2024
GSM8K
Human MBPP Eval
NQ
Trivia
ARC-
SIQA PIQA
QA
Specialized LLMs
CODELLAMA 7B
13.0
3.3
31.1
41.4
11.5
32.8
67
34.0
62.7
46.1
72.9
38.6
LLEMMA 7B
39.3
16.7
25.6
41.4
9.4
24.9
28.7
26.8
50.1
37.3
51.0
33.5
Generalist LLMs
LLAMA-2 7B
14.7
2.5
12.8
20.8
16.4
58.5
76.4
43.8
69.2
48.3
78.8
46.1
LLAMA-2 13B
28.7
3.9
18.3
30.6
16.1
63.8
77.3
49.4
73.0
50.1
80.8
52.8
Dense (DM)
26.7
9.9
20.7
30.8
24.0
55.3
76.7
44.5
68.9
48.3
78.2
49.8
SU (DM), Top-2
37.3
18.9
29.3
40.2
18.8
49.2
76.3
43.4
66.4
47.3
77.9
51.1
SU (CM), Top-2
40.1
16.2
26.2
35.2
24.5
58.2
75.6
44.7
69.1
47.1
78.0
52.1
BTM, Top-1
27.4
15.2
30.8
41.9
15.0
38.0
72.8
38.1
68.4
47.8
77.9
44.3
8TM, Top-2
27.7
15.3
30.6
42.6
15.3
38.5
73.1
38.5
68.3
48.0
78.1
44.3
BTX, sample Top-1
36.9
15.8
25.6
37.4
23.7
56.4
76.7
45.0
70.6
48.0
78.2
53.2
BTX, Top-2
37.1
17.8
28.7
39.4
24.8
57.1
76.9
45.6
67.9
48.7
78.7
52.5
Table 5: Individual task performance of BTX and baselines. SU Sparse upcycling
MoE compute ratio
Training time (days)
Total compute (GPU-days)
num. tokens
Math Code Knowledge Reasoning MMLU Avg.
(B)
BTX
23%
7.8
926.1
533
27.4
34.0
41.0
63.5
52.5
47.9
SU (CM)
100%
7.9
1007.1
252
28.2
30.7
41.3
62.9
52.1
47.3
Table 6: Comparison between BTX and Sparse upcycling (SU) with compute-matching (CM), which special case of BTX without the expert training stage as is shown by the first column that 100% of compute is spent on MoE training. We also report total training time, compute and number of training tokens. Comparing both performance on individual domains as well as the average, we can see that BTX has more balanced performance, in addition to higher throughput.
C Routing analysis
Layer-by-layer comparison of the routing decision for different router designs and down- stream tasks aggregated by task domain is shown in Figure 4. Routing distributions slightly vary in the first few layers, but quickly become indistinguishable from layer to layer. One exception is in Switch routing where Math expert becomes dominant across tasks in the last model layer.
We observe that Code expert is dominant force in Code domain in Top-2 routing with load balancing. Note the difference with other models where load balancing is not added, and Math expert prevails across domains. We look at Code domain closer and compare routing probability distribution for models with and without load balancing in Figure 5, On the bottom three graphs of the picture we can observe phenomena of the dead expert, where routing probability to Code expert shifted to 0, while with load balancing added, probability distributions across experts look more similar, with slightly higher expectations for the Code expert.
To understand experts specialize in other domains, we look closer at per-task distribution. Routing decision of the tokens in Math and Reasoning domains are shown in Figure 6. We observe that GSM8K task prefers Code and LLAMA-2 experts, while Math task more relies on in-domain expert. We hypothesise that this happens because GSM8K dataset consists of grade school math word problems that require common sense knowledge and basic arithmetic operations, while Math task requires college-level math knowledge, and more aligned with Math expert's training data. In the Reasoning domain, all tasks exhibit similar behaviour and equally rely Math and generalist L.M''s expertise.
15

---

Published as a conference paper at COLM 2024
GSM8K
Human Eval
Natural Questions
ARC Challenge
MMLU
Average
Score
BIX
29.8
27.4
23.0
43.4
50.0
34.7
no load-balancing (LBB)
34.6
19.5
23.2
44.4
51.6
34.6
LB freeze experts
34.8
18.3
24.1
44.9
51.4
34.7
Table 7: Ablations on different BTX training strategies. All variants are initialized from the same experts and trained for total of 10B tokens during MoE finetuning.
Task
Question and generation
Q: Janet's ducks lay 16 eggs per day. She three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck ogg. How much in dollars does she make every day at the farmers' market?
GSM8K
A: Janee'ss ducks lay 16 eggs per day. She gats three for breakfast every morning So she has 16- 3 eggs left. She farmers' market daily for $2 per fresh duck egg. she 9.72 dollars.
eggs She bakes muffins for her friends every day with So she has
remainder The answer is 18.
Q: from typing import List
def has close elements(numbers: List(float), threshold: float)- bool:
w0th Check if in given numbers, are any two numbers closer to each other than
given threshold.
222 hws.close ceaiieeelmmett((... 2.0, 3.00 0.5)
False
Human Eval
227 has claseclss..elll aiineeelements(1111 (1.0, 2.8, 3.0, 5.0, 0.3)
True
A: for in rengellen(mumbers) -1): for in range(i Jen/mumbers]):
abas(numbers(] numbersfji threshold:
return
return False
Natural Questions
Q: who gut the first nobel prize in physics
Max
Planck
Table 8: Examples of the token routing decisions for the Top-2 routing with load balancing in the math (GSM8K), code (Human Eval), and knowledge (Natural Questions) domains. Tokens highlighted are routed to the following experts: Wikipedia and LLAMA-2 7B, Math and LLAMA-2 7B, Code and LLAMA-2 7B, Math and Code, Wikipedia and Math, Wikipedia and Code. Tokens that were routed to the in-domain expert are underlined.
D TFLOPs Comparison
The number of Floating-poit point operations (FLOPs) is often used for quantifying the efficiency of an algorithm. According to our approximate calculation (the active parameter count multiplied by the total number of training tokens), Sparse Upcycling (SU) in the compute- matched setting has 1.67E+10 TFLOPs while BTX has 2.39E+10. While the SU uses less FLOPs, looking only at FLOPs is flawed as it does not include communication cost. The main motivation of our work is to reduce the communication cost by making the training embarrassingly parallel. This improved efficiency cannot be measured in FLOPs alone and should be looked at together along with throughput. Such a comparison scheme has also been suggested in Dehghani al. (2022). Indeed, BTX has higher throughput (see Table 6), iee. given the same amount of compute budget, BTX can do more computations.
E LoRA Experts Training
Low-Rank Adaptation (LoRA) is popular technique for making further training efficient for neural networks. We looked into only training a LoRA adaptor for each expert, but early experiments show that LoRA experts do not work well in our pretraining setting, possibly because LoRA experts do not have enough capacity for when experts are trained on and required to fit massive amounts of data. In Table 9, compared normal full-rank
16

---

Published as a conference paper at COLM 2024
Top-2 routing with load-balanding
Jyerr
Layer
Leyer
Layor
Lager?
Layer
04
Layer 10
Luywr 11
Logw 12
Logw 15
Laysrr 14
Layer 15
Layw
03
Layw 17
Layer 18
Layar 1D
Laywr 20
Laye 21
Layar 22
Layar 23
Layer 24
02
90
01
Maut
Curde
Vild
Cliss
Ltable 2D
Wir
Cleta
LLabe 7B
Maat
Grad
Labid 70
Mal
Alata
Case
Ledde 7B
Man
Ceste
Leadin P9
Vili
Cond
Wild
Ladee
LLaade
ELeld
Top-2 routing
Layer
Layw
Layer
Loyer
Layor
Layer
Layer?
Layar
0.4
02
Leryar 10
Liwwer 11
Lawe 12
Layw 15
Layyrr 14
Luyer 15
Layerr
Laww 17
Luywe 18
Laywer 19
Loww 20
Layw 21
Layw 23
23
Lupprr
Laycr
Laycr
Layer
Layer
Layer
Layer 32
Layor
0.4
Wdu
Codee
Hald
Asse
Cheie
Wali
Cheis
LLaes 79
Hau
May
Cool
Mau
Medi
Code
vaki
Mani
Oode
Hdli
Ohde
Leeta 70
Hald
Mell
Choli
Laatee
Laatee22
L
LLaes
LLataa
Lesto
training against low-rank training (LoRA added both attention and FFN parameters) for the code expert LLM training. We found that with rank-64, LoRA's performance lags the full-rank training as early as training for 40B tokens. We further found that training FFN with full-rank updates (while only adding LoRA on attention) can bridge that gap at small-data regime (e-g. 40B tokens) but LoRA still lags behind full-rank training when scaling up the training data to 200B tokens. This indicates that LoRA is unsuitable for pre-training setup with large amounts of data.
F Limitations
Although our experimental results on BTX are promising, we have not fully explored its potential in this paper. Due to compute limitations, we only experimented with three domains and four experts in this paper. More domain experts can be added if data for more domains is available using unsupervised domain discovery (Gururangan et al.,
17

---

Published as a conference paper at COLM 2024
Soft routing
Layer
Layer
Leyer
Leyer
Layor
Layer
Layer?
Layer
Locpe 10
Luywr 11
Lagu 12
Layw 15
Laysr 14
Layer 15
Laywi
0.3
0.3
-ayar
Layer 18
Luynr 1D
Laywr 20
Layer 21
Layar 22
Layar 29
Layer 24
27
90
Mauth
Custe
Hilr
Akutt
clnsa
Ltable 2D
Wili
Alem
Cletr
LLabd 7B
Maat
Grud
Mili
Casle
Laddd TD
Hili
Man
Ceste
Ltadda 2 7D
Vili
Coot
Wild
Adats
Céslo
Leeee 78
LLabe
Labid
Liaade
Switch routing
Laya
Layar
Layen
Layar
Layal
Layor
Lagur7 7
Layer
0.75
0.50
0.5
0.20
Layer
Layw 12
Layer
Loyer 12
Layer 13
Leyer 14
Leyur 15
Loyor
0.59
0.25
0.20
Layer LT
Laye 1.5
Layer
Layer 20
Layer 21
Layer 22
Lagter 23
Layer 24
6/79
0.75
0:30
Layer
26
Lagec
Laytr 25
Layer
Lapor 31
Layer
6/9
0:0
Madi
Cone
LLade 7B
Meth
Codo
Madi
Coso
Leabd 73
Moui
Malh
Costo
Laabe 7B
Madt
Kode
Lokko 7B
Houi
Cosl
Ladda 7B
Vili
Mldt
Cole
Lalls 7B
Mat
Case
LLado TB
Laddee
Figure 4: BTX routing decisions of the tokens at various layers to different experts (Wiki, Math, Code, LLAMA-2 7B) for different downstream tasks. The tasks are aggregated by domain: Code (Human Eval, MBPP), Math (GSM8K, MATH), World knowledge (Natu- ral Questions, TriviaQA), and Reasoning (ARC-Easy, ARC-Challenge, SIQA, PIQA, and WinoGrande). We observe that top-2 routing with load balancing ensures more uniform distribution of the load between experts compared to the other routing methods across all layers.
2023). Then we can repeat the finetuning step using these new experts in addition to the previously trained experts. The MoE training itself should scale well to more experts as demonstrated in Fedus et al. (2022). Training on more domains should amplify the benefit of the parallelization of experts training. Having more experts will also make the final MoE model more efficient because the number of active experts can remain the same while its overall capacity increases.
18

---

Published as a conference paper at COLM 2024
Task Human Eval. kead-beleencing
MWU
Man
Lode
LLSMM
2000
2590
CA
Toatk Human Evall layer 16, top rousing with load balancing
Wil
Math
cosl
LLAV/ TB
3000
0/1
0.
0-4
26
0.0.
0.7
04
C/
6.6.
no
0.7
0.4
0.
om
0.25
0.50
0.55
1.00
Toak Human Evall layer 32, top rousing with lond belencing
Wili
wath
Code
LLAME /5
400
300
300
0.00
0.25
0.50
0.75
1.00
0.00
C25
0.50
0.75
02
14
0.5
04
0.00
0.25
0.50
0.75
1.00
Trask Human Eval, laper top routing
Widi
2000
1500
3000
61
Tesk Human Evail, Ityer 16 top-2 routing
500
Juman Cval, layer 32 top-2 routing
3500
Figure 5: Routing probabilities per expert across different layers for Human Eval task. We compare top-2 routing with (left) and without load balancing (right).
In our experiments, we used simple implementation of MoE and did not optimize it using more complex techniques such as placing different experts on different GPUs to run them in parallel. Such efficient MoE implementation could shorten the training time of BTX, and the sparse upeycling baseline as well.
Compared BTM, BTX provides approach to finetune the combined experts, which can be directly applied in instruction finetuning or RLHF procedures. However, we leave that for future work
focused on the pretraining stage in this paper.
The question of whether experts in MoE are better off specializing in specific domains not is interesting that is worth further investigation. Our approach explicitly
19

---

Published as a conference paper at COLM 2024
Top-2 routing with load balancing
Loyer
Layer
Layer
0.3
0.2
0.1
0.0
Wiki
Math
Code LLaMa-2 7B
Wiki
Math
Code LLaMa-2 78
Wiki
Math
Code LLaMa-2 2 7B
routing load baad-bllancing
Layer
Layer
Layer 32
0.4
0.3
0.1
0.0
WIR
Math
Code
ARC Easy
ARC Challlnge
SIOA
PIQA
Winogrande
Figure 6: Routing decision of the tokens in Math and Reasoning domains. We observe that GSM8K task prefers Code and LLAMA-2 experts, while MATH task relies more on in-domain expert. In the Reasoning domain, the load is distributed between Math and LLAMA-2 7B experts.
#tokens
Training
Valid ppl
HumanBval (pass@1)
40B
Full-rank
1.86
21.34
40B
LoRA (attn r=64, ffn r=64)
1.99
15.24
40B
LoRA (attn r-64, fin full-rank)
1.88
24.39
200B
Full-rank
1.74
31.71
200B
LoRA (attn r=64, ffn full-rank)
1.77
28.66
Table 9: Ablations parameter-efficiet experts with LoRA.
tied experts to certain domains, but such specialization does not seem to emerge naturally during MoE training (Jiang et al., 2024). We observed that some experts are used more in their corresponding domain tasks, showing that their domain specialization partially remains even after the MoE finetuning.
We only compared BTX two of its special variants, ie. BTM with 100% compute allocated to expert training and 0% MoE finetuning, and sparse upcycling with 0% compute allocated to expert training and 100% on MoË finetuning. Future work could perform thorough sweep of the compute allocation ratio between expert training and MoE training. Also, we did not perform experiments with different data mixtures for MoE finetuning other than uniform sampling.
20


Reviews:
[{"review_id": "bRpgBeZaUn", "reviewer": "Reviewer_EfwL", "summary": "In this paper, the authors propose a new method to train domain expert LLMs and mix them together to get a better model on all the tasks. Compared to the existing studies, it has two major differences - 1) compared to sparse upcycling (SU), it has a separate domain expert training stage which each individual expert trained on special tasks. 2) compared to branch-train-merge (BTM), combining different expertise is done as a single model rather than ensemble using MoE architecture. The authors show that the proposed method can outperform existing SU and BTM on average scores for various tasks. The authors also show the proposed method can be more efficient than other method thanks to the parallel training of different domains.\n\nOverall, the paper is very well-organized and easy to follow. The proposed idea looks interesting and intuitive, even though the gains in the quality are not huge (compared llama2-13b and SU). However, it will be a good contribution to the community to improve more in that direction.", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their time and careful consideration.\n\n> It will be nice to have a discussion how to extend this to the greater number of experts and larger models. Many open-source models use at least 8 experts.\n\nExtending to more experts should be relatively straightforward. First, one needs to identify more domains and collect data corresponding to those domains, then train corresponding experts starting from the seed model. The previously trained expert LLMs can be used as is. The mixing and fine-tuning steps can be applied to 8 experts or more in the same way. As you mentioned, there is evidence that MoE routing scales well to more experts. We expect this to be true in our MoE fine-tuning step as well. We will add this discussion to the paper. \n\n> It will be also useful to have a baseline MoE model which has the same number of parameters (4 experts) but trained from scratch.\n\nYes, it would be interesting to compare against such a baseline. However, in our setup, training any model from scratch will take too much computation, especially for a larger MoE model. Still, there is evidence that sparse-upcycling works well, even compared to MoE trained from scratch as shown in [1] (see their Fig 4). We quote from [1] “The language MoE model trained from scratch requires about 120% of the original dense checkpoint’s computation budget to catch up to the upcycled model”. Also, prior work on BTM [2] compared with training MoE from scratch (although with domain dependent routing) and they found the asynchronous expert training is more compute efficient. And in our work, we show that BTX is superior to BTM.\n\n[1] Komatsuzaki, Aran, et al. \"Sparse upcycling: Training mixture-of-experts from dense checkpoints.\" (2022) https://arxiv.org/pdf/2212.05055  \n\n[2] Li, Margaret, et al. \"Branch-train-merge: Embarrassingly parallel training of expert language models.\" (2022) https://arxiv.org/pdf/2208.03306", "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": "- Paper is very well written\n- The idea is simple and intuitive, yet it gives some good gains in quality and efficiency.\n- It has various ablation studies including routing analysis, compute efficiency and etc.", "reasons_to_reject": "- I don't see a good reason to reject this paper.", "questions_to_authors": "- It will be nice to have a discussion how to extend this to the greater number of experts and larger models. Many open-source models use at least 8 experts.\n- It will be also useful to have a baseline MoE model which has the same number of parameters (4 experts) but trained from scratch.", "extra_scores": {}}, {"review_id": "bxCvferkfc", "reviewer": "Reviewer_tFfH", "summary": "This paper introduces the paradigm of aggregating multiple dense models finetuned on different data into an intact MoE, which first separately finetunes $N$ domain LLMs with BTM, then aggregates these models into a MoE through further finetuning. It also proposes various variants with different choices of the gating networks and initializations of expert parameters, and conducts detailed experiments. The results show salient improvements of BTX on the downstream performance under Data Matching (DM) and Compute Matching (CM) settings.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> the compute budget of MoE highly depends on the library…\n\nMoE performance can vary depending on the implementation, so it is possible to increase the efficiency of MoE training using more complex training schemes. In contrast, the first step of BTX offers a simple recipe for training a large model while being less dependent on complex training infrastructures and implementations while being more robust to hardware failures. However, we still think MoE training is a very effective way of scaling, and we showed that it can be effectively used as the final step in BTX. This also means that an efficient MoE implementation will also make BTX more efficient, but of course to a lesser degree compared to full MoE training.\n\n> the Data Matching (DM) is not fair … BTX activates all the parameters in the first stage\n\nIn the data matching case, both models see the same amount of data during training. However, BTX does not activate all parameters on every data point in the first stage. For example, if the sample is from the math dataset, then only the math expert LLM will process this data. Therefore, only 1/4 of the total BTX parameters will be activated during the first stage of training. In contrast, sparse-upcycling activates 2 out of 4 FF experts.\n\n> A fair comparison should be under the same TFLOPs\n\nSparse upcycling (SU) in the compute-matched setting has 1.67E+22 FLOPS while BTX is 2.39E+22. The SU uses less FLOPS, but looking only at FLOPS is flawed as it doesn’t include communication cost. The main motivation of BTX was to reduce the communication cost by making the training embarrassingly parallel. This improved efficiency cannot be measured in FLOPS alone but should be looked at together along with throughput. Such a comparison scheme has also been suggested in https://arxiv.org/abs/2110.12894 . Indeed, BTX has higher throughput (926 GPU days vs SU’s 1007), i.e. given the same amount of compute budget, BTX can do more computations (higher TFLOPs). We will include FLOPS and discussion in the paper.\n\n> In Table 2, the specialized LLMs should include the individual domain expert LLMs \n\nWe will add a table comparing BTX vs.single-expert on tasks in each domain, but cannot show it here due to the space limit. In the math and code domains, the final model performance is little behind the expert LLMs, but overall it has balanced good performance. We will include this table in the paper. We can reply to you if you want to see the full table.", "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": "1. The idea of merging finetuned LLMs to MoE is intuitive and interesting.\n2. The introduced pipeline seems acceptable considering the trade-off between the performance and efficiency.\n3. Results are good under either the same amount of data (DM) or the same computation (CM).", "reasons_to_reject": "1. **The comparison to the upcycling baseline is not fair.**\n   - As far as I know, the compute budget of MoE highly depends on the library and hardwire environment used. So the Compute Matching (CM) may not be a fair comparison as MoE would encounter more communication bottlenecks as you have shown in Table 6, which could be enhanced by more painless designs of training frameworks for MoE.\n   - Also, the Data Matching (DM) is not fair for the Upcycling as its total parameter activation is less than BTX. BTX activates all the parameters in the first stage and switches to top-$k$ in the second finetuning stage. However, the upcycled MoE always uses top-$k$ activation.\n   - A fair comparison should be under the same TFLOPs which is not included.\n\n2. **Lack of comparison to the individual domain expert LLMs.** \n   - In Table 2, the specialized LLMs should include the individual domain expert LLMs in Table 1 for clarification of the gap between the general LLM and domain oracle LLMs. We are curious about whether the combined model can surpass the oracle model in the specified domains, rather than other standalone LLMs.", "questions_to_authors": "Please see the reasons to reject.", "extra_scores": {}}, {"review_id": "t75ZC4QWeh", "reviewer": "Reviewer_go34", "summary": "The paper introduces Branch-Train-MiX (BTX), a method for training Large Language Models (LLMs) with expertise in multiple domains. BTX branches a seed model to train experts in parallel, then integrates them into a Mixture-of-Experts (MoE) model and fine-tunes for token-level routing. This approach aims to enhance the seed model's performance across various tasks while maintaining efficiency.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper and providing insightful feedback.\n \n> The authors present only three domains and four experts, making the results less convincing.\n\nWe agree with the reviewer that further evaluation of the proposed model would be beneficial to strengthen the conclusions of our work, however due to compute and resource limitations, we were unable to include more experts. We tested our approach on the most common tasks - mathematical and code reasoning, and commonsense knowledge. However, even with only 4 experts, we see the benefit of having multiple experts from the large improvements compared to a dense baseline. \n\n> Admittedly, full model training is expensive, I suggest the authors report results with PEFT techniques on more domains and experts.\n\nWe looked into PEFT using LoRA experts, but from our experience we found that LoRA experts don't work well in the pretraining setting possibly because LoRA experts do not have enough capacity when experts are trained on massive amounts of data.", "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": "- BTX's parallel and asynchronous training of experts reduces communication costs and increases throughput.\n- It demonstrates better compute-performance trade-offs compared to training a larger generalist LLM or multiple specialized ones.", "reasons_to_reject": "- The authors present only three domains and four experts, making the results less convincing. Admittedly, full model training is expensive, I suggest the authors report results with PEFT techniques on more domains and experts.", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This work proposes a new way to train a mixture of expert model called Branch-Train-MiX (BTX). It starts with parallel finetuning of LLMs to train experts for different tasks. It then merges those experts into a single model as follows. The FeedForward layer of each model is combined together to create a single MoE layer for the end model. A router routes each token to select expert feedforward models. The self-attention layer is merged by averaging the weights across models. The end model is the finetuned on the aggregated data to train a generalist router. The experts are trained on math, coding, and general text (wikipedia). The method shows superior performance on several downstream tasks (including math, coding, knowledge-based questions, and reasoning) over prior approaches such as sparse upcycling and Branch-Train-Merge approaches. 

Overall, the independent and parallel finetuning of experts and then merging them using an MoE concept is an interesting and scalable way to train a generalist agent and provides a path for maintaining the benefits of finetuning while minimizing its adverse affects on the general capabilities of the model. This line of work can ignite new ideas on reusing and building on top of finetuned models in efficient ways towards a generalist model.
Decision: Accept