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  rozière
jacob  kahn
daniel  li
wen-tau  yih
jason  weston
xian  li
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-21 20:04:38

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.

Main paper: 
 introduction: 
 In recent years, Large Language Models (LLMs) have shown impressive performance in a wide-range of tasks (Brown et al., 2020; Touvron et al., 2023; Achiam et al., 2023) , including code generation (Li et al., 2022b; Rozière et al., 2023) , solving math problems (Azerbayev et al., 2023) , multilinguality (Zhao et al., 2024), etc. Training such LLMs requires a 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 as a single failed GPU can cause the whole training to halt (Zhang et al., 2022; Gemini Team, 2023) . Li et al. (2022a) proposed the Branch-Train-Merge (BTM) method for embarrassingly parallel training of LLMs without any synchronization for improving the throughput of pretraining. It starts by creating multiple copies of a seed LLM, then separately 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 a 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 Mixtureof-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 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 Branch-Train-Merge, Mixture-of-Experts are often trained in a fully synchronized fashion, and the communication cost increases with the number of experts due to all-to-all communication.In this paper, we aim for the best of both worlds, combining the advantages of Branch-Train-Merge and Mixture-of-Experts, while mitigating their disadvantages. We achieve this by training multiple expert LLMs separately as in the Branch-Train-Merge method, but subsequently combine those experts into a single model using an MoE architecture. More specifically, the feedforward sublayers from all the expert LLMs are brought together into a single MoE module at each layer, and a 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 Branch-Train-MiX (BTX).The main advantage of BTX compared to MoE is that expert training is embarrassingly parallel and asynchronous, reducing communication cost and increasing training throughput.Compared to Branch-Train-Merge, the final BTX model is a unified neural network that can be finetuned or used like any other standard LLM. The final BTX model will not significantly increase inference FLOPs compared to the seed model since it 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 a relatively short period compared to 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 all tasks demonstrating the benefits of learnt routing through MoE finetuning. Compared to purely MoE training such as sparse upcycling, BTX is more compute efficient with higher training throughput and more balanced performance across tasks in different domains.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 Branch-Train-Merge 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.MoE is used to scale deep networks in Shazeer et al. (2017) using a 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 Transformer 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 a random mapping based on input tokens. In larger scale experiments with recent LLMs, Jiang et al. (2024) demonstrated that the MoE approach can match the performance of dense LLM counterparts using a 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. For training a math expert, starting from a code expert rather than a general LLM was shown to be more beneficial (Shao et al., 2024; Azerbayev et al., 2023) . 
branch-train-mix: 
 Given an existing LLM M which has been pretrained on a large corpora covering a 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 := {D 1 , . . . , D N }, each related to a specific knowledge domain such as math, code, etc. The proposed method contains three stages: Branch, Train, and MiX.Initializing from the seed model M, we train N expert LLMs {M 1 , . 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-Merge method (Li et al., 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 we will describe in the next section.We employ a Mixture-of-Experts approach to combine the domain expert models M i . However, instead of using the classical procedure of mixing the final outputs from M i , 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 a single MoE sublayer. If FF l i (x) is the feedforward sublayer at the l-th layer of the i-th domain expert M i , then the combined MoE layer for input representation x at layer l will compute:FF l MoE (x) = N ∑ i=1 g i (W l x)FF l i (x).Here W l is a linear transformation and g is a routing function, which usually has sparse output and hence switches on only some experts. Since we can skip computing FF l i (x) if the corresponding router output is zero, the actual computation of FF l MoE (x) will be much more efficient than computing all domain experts. However, routing decisions can change from token to token, so 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 g(W l x) = SoftMax(TopK(W l x)), unless otherwise stated.For the self-attention sublayers, we combine the different domain experts by simply averaging 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.) as well.Note that the only new parameters we introduce are the router's transformation parameters W l , 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 i to use. In addition, funetuning is helpful because the self-attention 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 a small amount of combined finetuning might make large improvements.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 a situation because a dead expert is never in the top-k selection, and therefore never receives a training signal. Load balancing offers a simple solution by adding an extra loss term that encourages the experts to be utilized equally. We use a loss term similar to (Fedus et al., 2022) :L LB = αN N ∑ i=1 u i p i where u i = 1 |B| ∑ x∈B g i (W l x) and p i = 1 |B| ∑ x∈B SoftMax i (W l x).Here B is the current data batch, and α is 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 a Top-1 routing method proposed by Fedus et al. (2022) .• Soft routing: We use softmax as 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.  
experimental setup: 
 We base our experiments on the setup used for LLAMA-2 pretraining (Touvron et al., 2023) .In particular, we use the LLAMA-2 7B model as our seed model.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 to 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 a "generalist" expert so that its general knowledge is transferred to 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 to 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.We compare to the following baselines:• LLAMA-2: We compare to the original LLAMA-2 7B that we use as a seed model, as well as LLAMA-2 13B. • Llemma 7B: A language model specializing in mathematics (Azerbayev et al., 2023) by continued training of CodeLlama 7B on math data.For evaluation, we use the zero-and few-shot performance on multiple benchmarks that test different skills: GSM8K and MATH for math reasoning; HumanEval 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. 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 days2 , 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 13B model.Domain experts excel at their respective tasks. We first analyze how expert LLMs specialize to specific domains. Results are summarized in Table 1 . As expected, individual expert LLMs achieve the best performance in their respective domain, where the math and code domains see especially large improvements. In addition, there are several interesting 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. BTX improves all tasks where experts specialize. Table 2 and Figure 2 (right) show aggregated 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 Top-2 corresponding to different number of active parameters) improve on all expert domains, such as math, coding and world knowledge without regressing on 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 are not their speciality such as world knowledge and commonsense reasoning. Compared to alternative data-matching (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 a large margin on average, indicating that MoE finetuning to learn token-level routing is beneficial. Overall, the results demonstrate that BTX is a better performing method for continued pretraining which is 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.We compare BTX with baselines in terms of compute efficiency in Figure 2 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 it lacks sparsity and has the highest number of active parameters. Overall, the Top-2 routing gives us a good balance between performance and efficiency.We also ablate additional design choices of BTX, with results summarized in Table 7 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 on 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 i .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 of the load between experts. Analyzing the token probability distributions, we observe a shift towards low probability scores across all experts with load balancing, especially closer to the final layers of the model, which contributes to the fair routing. Interestingly, all models without load balance heavily rely on the Math expert, with a 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, it not only becomes visible, but becomes the dominant expert in the math and code domains.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 we look at a more detailed token distribution (Appendix C, Figure 6 ), we find that the GSM8K 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. 
conclusion: 
 We introduced Branch-Train-MiX (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 a single Mixture-of-Experts (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 performance 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 a 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.  
 
 Appendix: Layer-by-layer comparison of the routing decision for different router designs and downstream 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 a 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 a 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 if 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 on Math and generalist LLM's expertise. The number of Floating-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 computematched 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 et al. (2022) . Indeed, BTX has higher throughput (see Table 6 ), i.e. given the same amount of compute budget, BTX can do more computations.Low-Rank Adaptation (LoRA) is a 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 are required to fit massive amounts of data. In Table 9 , we compared normal full-rank training against low-rank training (LoRA added to 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 a pre-training setup with large amounts of data.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 or using unsupervised domain discovery (Gururangan et al. , 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. In our experiments, we used a 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 an efficient MoE implementation could shorten the training time of BTX, and the sparse upcycling baseline as well.Compared to BTM, BTX provides an approach to finetune the combined experts, which can be directly applied in instruction finetuning or RLHF procedures. However, we leave that for future work as we focused on the pretraining stage in this paper.The question of whether experts in MoE are better off specializing in specific domains or not is an interesting one that is worth further investigation. Our approach explicitly Table 9 : Ablations on parameter-efficient 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 to two of its special variants, i. The GPU days of Llama-2 13B is an approximate measurement, calculated by doubling the training compute of a 7B model trained with the same amount of pretraining data (according toTouvron et al. (2023) Table2). Since Llama-2 13B is not trained from the seed model, we simply report their difference in GPU days.Sample Top-1 training is more efficient, so we trained for 160B tokens for the same compute. 


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

Reviews: review: paper_summary: paper_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.

Overall, 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.
strength_weakness: strength_weakness: - I don't see a good reason to reject this paper.
questions: questions: - 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.
- It will be also useful to have a baseline MoE model which has the same number of parameters (4 experts) but trained from scratch.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1. **The comparison to the upcycling baseline is not fair.**
   - 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.
   - 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.
   - A fair comparison should be under the same TFLOPs which is not included.

2. **Lack of comparison to the individual domain expert LLMs.** 
   - 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: questions: Please see the reasons to reject.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.

Comments: comment: We agree that unsuccessful experiments can be useful, and we will add our LoRA results in the appendix.

Below is our experimental result using LoRA. We compared normal full-rank training against low-rank training (LoRA added to both attention and FFN parameters) for the code expert LLM training. We found that with rank=64, LoRA’s performance lags behind 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 setups with large amounts of data. 

#tokens | Training ________ | Valid ppl | HumanEval (pass@1)

40B ____ | Full-rank _______ | 1.86 |  21.34

40B ____ | LoRa___________ | 1.99 |  15.24

40B ____ | LoRa (attn only) | 1.88 |  24.39

200B___ | Full-rank  _______ | 1.74 |  31.71

200B___ | LoRa (attn only) | 1.77 |  28.66
Thank you for your response.

Regarding the experiments that I found less convincing, I still have some concerns.

I appreciate the insights on the LoRA experts during pretraining. I've learned a lot from this discussion. I hope that the results of these seemingly unsuccessful experiments can be included in the appendix. **I will raise my rating** if you can provide this information (which will help more practitioners save unnecessary training costs).
Thanks for the rebuttal.

I agree on your points, and keep my score.
Thanks for your response.