Source-linked AI summary
Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM
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
TL;DR
Training LLMs across specialized domains is compute-intensive, and synchronized training can make communication a scaling bottleneck. BTX trains domain experts asynchronously, combines them in a unified MoE model, and finetunes routing; experiments report better performance than BTM across tasks and a more compute-efficient, balanced tradeoff than sparse upcycling.
Problem
Training LLMs across specialized domains requires substantial compute, while synchronized training's frequent communication is a scaling bottleneck and Branch-Train-Merge lacks a unified model for further finetuning.
Source: §1 · pp. 1–2Method
BTX trains domain experts separately, combines their feedforward parameters as MoE experts, averages other parameters, and finetunes the unified model to learn routing.
Source: §1 · p. 2Results
BTX outperforms BTM on all tasks and, compared with sparse upcycling, is more compute efficient with higher throughput and more balanced cross-domain performance.
Source: §1 · p. 2Takeaways & Limitations
The experiments suggest BTX can boost a generalist LLM through continued training on specialized datasets and inform compute allocation in late pretraining.
Source: §5 · p. 11Experiments cover only three domains and four experts, so the paper does not establish BTX's performance with broader domain or expert counts.
Source: §6 · p. 11
Abstract
from arXiv · showhide
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.
1 Introduction
Branch-Train-MiX (BTX) combines independently trained domain experts in a unified MoE model, retaining asynchronous expert training while enabling learned token-level routing. Experiments report improvements over the seed model and BTM, with greater compute efficiency and more balanced cross-domain performance than sparse upcycling.
1 Introduction
- Synchronized training's frequent communication is a major scaling bottleneck, and a single GPU failure can halt the entire training run.Training LLMs can require thousands of GPUs and trillions of tokens.Source: §1 · p. 1
- Branch-Train-Merge trains domain-specific copies independently but lacks a unified model for further SFT or RLHF finetuning.The paper notes that these finetuning stages can boost performance and are crucial to building aligned models.Source: §1 · p. 1
- MoE activates only a subset of parameters at a time, allowing total parameter count to grow without additional computation, but conventional MoE training can incur all-to-all communication that increases with expert count.MoE is applied to Transformer feedforward sublayers.Source: §1 · pp. 1–2
- BTX trains experts separately, combines their feedforward layers as MoE experts, averages other parameters, and finetunes the unified model so a router learns token-level selection.The router selects which feedforward expert to use at each token.Source: §1 · p. 2
- Expert training is embarrassingly parallel and asynchronous, reducing communication cost and increasing throughput; sparse activation avoids a significant inference-FLOP increase over the seed model.Unlike BTM, the final BTX model is a unified network that can be finetuned or used like a standard LLM.Source: §1 · p. 2
- BTX outperforms BTM on all tasks and is more compute efficient than sparse upcycling, with more balanced performance across domains.Experiments used Llama-2 7B experts trained on math, code, and Wikipedia data, plus the original model as a fourth expert.Source: §1 · p. 2
2 Related Work
The related work situates BTX among approaches to reduce synchronization, scale models with MoE, and preserve or specialize knowledge through continual learning. It distinguishes BTX by combining independent domain-expert training with MoE while retaining asynchronous training.
Asynchronous parallel training
- Branch-Train-Merge runs independent domain-specific training processes and averages selected experts' output distributions at inference; parameter averaging has also improved performance when models differ only in hyperparameters.Earlier asynchronous approaches instead let workers diverge and synchronize loosely or less frequently.Source: §2 · pp. 2–3
Mixture-of-Experts
- MoE scales networks through routing, and prior work found learned, fixed, or random routing can work; BTX is closest to domain-conditioned fixed routing but adds asynchronous training.Prior LLM studies report MoE can match dense-model performance with fewer active parameters, while other work studies fine-grained and shared experts.Source: §2 · p. 3
Continual learning
- BTX relates to continual learning because domain experts continue training from a seed model on differently distributed data, with separate parameters isolating domains.Prior work also reports that continued training can produce code experts faster than training from scratch, and that starting math training from a code expert can be beneficial.Source: §2 · p. 3
3 Branch-Train-MiX
BTX branches a seed LLM into independently trained domain experts, then combines their feedforward sublayers as MoE experts and averages other parameters. Its variations alter how experts are composed, how routing is trained, or how many expert modules are formed.
3.1 Branch & Train: Embarrassingly Parallel Expert Training
- Each expert starts from the seed LLM and trains independently on its domain dataset, enabling N-way parallel training with linear throughput scaling, lower all-to-all communication, and failures isolated to one process.After training, BTX differs from Branch-Train-Merge by combining experts into one model for further finetuning.Source: §3.1 · p. 3
3.2 MiX: Combining Separate Experts to be a Mixture-of-Experts
- BTX replaces output-level expert mixing with layer-level MoE over domain experts' feedforward sublayers; sparse routing activates only selected experts, while remaining parameters are averaged.The method assumes self-attention is less domain-specialized than feedforward layers; finetuning trains the router and can improve the averaged attention weights' fit.Source: §3.2 · p. 4
3.3 Variations
- Load balancing adds a loss term that encourages equal expert use, addressing dead experts that receive no router-selected training signal.The term is computed per layer and added to the negative log-likelihood loss.Source: §3.3 · p. 4
- Routing variations include Switch Top-1, soft routing that activates all experts at increased compute cost, and sampled Top-1 using Gumbel softmax with hard sampling at inference.Sampled Top-1 anneals temperature toward a sharp distribution to reduce the train–inference discrepancy.Source: §3.3 · pp. 4–5
- Splitting each domain's feedforward layer into C chunks of dimension dFF/C increases the MoE module count from N domains to MC modules.Source: §3.3 · p. 5
- Blending Experts splits each domain's feedforward layers into N chunks and combines same-index chunks across domains, so every resulting expert contains parameters from all domains.Source: §3.3 · p. 5
4 Experiments
Experiments branch a Llama-2 7B seed into domain experts, combine them with the seed as a generalist in an MoE, and compare against dense, sparse-upcycling, BTM, and specialized baselines across benchmark groups.
4.1 Experimental Setup
- Experiments use Llama-2 7B as the seed model, following the Llama-2 pretraining setup.Source: §4.1 · p. 5
4.1.1 BTX Training
- BTX makes three copies of pretrained Llama-2 7B and continues training them on domain datasets, then adds the original seed as a generalist and mixes all four experts into one MoE.It finetunes the MoE on all data sources used to train the four experts, including original Llama-2 7B pretraining data, for another 80B tokens.Source: §4.1.1 · p. 5
4.1.2 Baselines
- Baselines include Llama-2 models, dense training, sparse upcycling, BTM, CodeLlama, and Llemma.Dense training is data-matched; sparse upcycling is also compute-matched, while BTM uses the same experts without MoE finetuning.Source: §4.1.2 · p. 6
- Baselines, experts, and MoE models share AdamW and the same learning-rate schedule, with a 4M-token batch and 4096-token sequences.Source: §4.1.2 · p. 6
4.1.3 Evaluation
- Evaluation averages benchmark performance for math, code, world knowledge, and reasoning, and reports MMLU for general capability.Source: §4.1.3 · pp. 6–7
- Table 2 compares aggregated capabilities across generalist and specialized models; dense, sparse upcycling, BTM, and BTX use the same data mixture, except BTM omits finetuning.Source: Table 2 · p. 7
4.2 Main Results
BTX improves performance across expert domains while avoiding regressions on other tasks. Its compute-matched comparisons and routing ablations examine how parallel expert training and load balancing shape that performance.
4.2.1 Overall Performance
- BTX improves all expert domains without regressing on other tasks, approaches Llemma 7B and CodeLlama 7B in math and coding, and beats Llama-2 13B except on reasoning.The math and code experts individually improve in their domains but perform worse than the seed model on TriviaQA.Source: §4.2.1 · p. 7, Figure 2 · p. 8
- With the same compute budget, BTX trains more than 2× the data and achieves slightly higher average performance across domains than sparse upcycling.The compute-matched comparison reports higher throughput and more balanced performance for BTX.Source: Table 3 · p. 8, §4.2.1 · p. 8
4.3.1 Ablations of BTX training
- Top-2 routing offers a performance-efficiency balance; removing load balancing lowers HumanEval but raises GSM8K, and mixing domain feedforward layers lowers performance.Freezing expert feedforward modules has little impact; the authors suggest the mix stage mainly trains averaged self-attention weights and router transformations.Source: §4.3.1 · pp. 9–10
4.3.2 Routing Analysis
- Load balancing makes expert use more uniform and revives the Code expert, which becomes dominant in math and code; without it, models rely heavily on Math.With load balancing, math tasks often route tokens to Code and Llama-2 7B, while MATH relies more on Math; reasoning uses Math and the generalist similarly.Source: Figure 3 · p. 10, §4.3.2 · pp. 10–11
5 Conclusion
Branch-Train-MiX (BTX) continues pretraining by specializing seed-model copies in parallel and combining them into a finetuned Mixture-of-Experts model. The experiments suggest this improves generalist capabilities with greater compute efficiency than training a larger generalist or separate specialists.
5 Conclusion
- BTX trains seed-model copies asynchronously and in parallel, then merges them into a single Mixture-of-Experts model through finetuning.The parallel stage brings higher training throughput and scalability; the MoE finetuning stage makes the final LLM more performant.Source: §5 · p. 11
- The experiments suggest BTX is more compute efficient than training a larger generalist LLM or several separately specialized LLMs.The authors say these findings may inform compute allocation during late pretraining.Source: §5 · p. 11
- The authors conclude that continued training on specialized-knowledge and skills datasets can boost a generalist LLM's performance using BTX.They focus this conclusion on the method's supported experimental findings.Source: §5 · p. 11
6 Limitations & Future Work
The experiments cover only three domains and four experts, and use a simple, unoptimized MoE implementation. The authors also leave instruction finetuning and RLHF, expert specialization, broader compute allocations, and alternative MoE-finetuning mixtures for future investigation.
6 Limitations & Future Work
- Experiments cover only three domains and four experts, and the simple MoE implementation was not optimized with techniques such as placing experts on different GPUs.The authors say more domains could amplify parallelization benefits and more experts could increase capacity while keeping active experts constant.Source: §6 · p. 11
- Instruction finetuning and RLHF of the combined experts are left for future work because this study focuses on pretraining.The authors say BTX's finetuning approach could be applied to those procedures.Source: §6 · p. 12
- Whether MoE experts are better off specializing in specific domains remains an open question; the authors observed that domain specialization only partially remains after MoE finetuning.Some experts were used more on their corresponding domain tasks, but such specialization did not seem to emerge naturally during MoE training.Source: §6 · p. 12
- The comparison covers only BTM with all compute on expert training and sparse upcycling with all compute on MoE finetuning, leaving intermediate compute allocations untested.The authors suggest a thorough sweep of the allocation ratio as future work.Source: §6 · p. 12
- MoE finetuning was tested with uniform sampling only; other data mixtures were not evaluated.The authors explicitly note that they did not experiment with different data mixtures for MoE finetuning.Source: §6 · p. 12
A Data mixture
The appendix describes the data mixtures used for domain-expert training and the dataset sampling weights used to finetune the MoE model.
A Data mixture
- Table 7 reports the exact data-mixture ratios used to train each domain expert.Its caption identifies the table as covering data sources and weights for domain experts.Source: A Data mixture · p. 16, Table 7 · p. 16
- MoE finetuning samples the math-expert, code-expert, Wikipedia-expert, and original Llama-2 7B datasets with probabilities of 30.16%, 40.31%, 10.30%, and 19.23%, respectively.The original Llama-2 7B serves as the generalist expert.Source: A Data mixture · p. 16
B Evaluation
Evaluation uses task-specific metrics and greedy decoding; detailed individual-task results are reported in Table 8.
B Evaluation
- Code tasks HumanEval and MBPP are evaluated with pass@1.The reported metrics follow those used in Touvron et al. (2023) and Rozière et al. (2023).Source: B Evaluation · p. 16
- Math and knowledge tasks use exact match, while MMLU and ARC use accuracy.The math tasks are GSM8k and MATH; the knowledge tasks are Natural Questions and TriviaQA.Source: B Evaluation · p. 16
- All generations use greedy decoding, and Table 8 reports detailed results for individual tasks.Table 8 is captioned as individual task performance for BTX and baselines.Source: B Evaluation · p. 16, Table 8 · p. 16
C Routing analysis
Routing varies by task and router design: top-2 routing with load balancing distributes load more evenly, while expert preferences differ across math and reasoning tasks.
C Routing analysis
- Routing distributions vary slightly in early layers, then quickly become indistinguishable, except Switch routing makes the Math expert dominant in the final layer across tasks.Figure 4 compares routing decisions across layers and task domains.Source: C Routing analysis · p. 16, Figure 4 · p. 19
- GSM8K tokens prefer the Code and LLaMa-2 experts, MATH relies more on the Math expert, and reasoning tasks share reliance between Math and LLaMa-2 7B.The paper hypothesizes that GSM8K's grade-school word problems draw on common sense and basic arithmetic, unlike MATH's college-level content.Source: C Routing analysis · p. 17, Figure 6 · p. 21
- Top-2 routing with load balancing distributes load more uniformly across experts than other routing methods across layers.In the Code domain, the Code expert is dominant with load balancing; without it, the Math expert prevails across domains, and Code routing probabilities can fall to zero.Source: C Routing analysis · p. 17, Figure 4 · p. 19, Figure 5 · p. 20