Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM

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.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

- BTX's parallel and asynchronous training of experts reduces communication costs and increases throughput. - 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.

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.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

1. The idea of merging finetuned LLMs to MoE is intuitive and interesting. 2. The introduced pipeline seems acceptable considering the trade-off between the performance and efficiency. 3. 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.** - 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 to authors

Please see the reasons to reject.

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.

Rating

8

Confidence

5

Ethics flag

1

Reasons to accept

- Paper is very well written - The idea is simple and intuitive, yet it gives some good gains in quality and efficiency. - 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. - It will be also useful to have a baseline MoE model which has the same number of parameters (4 experts) but trained from scratch.