Train Separately, Merge Together: Modular Post-Training with Mixture-of-Experts

summary

Video file (mp4)

The gist

The gist: BAR presents a modular approach to post-training that extends fully trained language models with new domain capabilities via independently trained mixture-of-experts.

In short

BAR introduces a modular approach to extend fully trained language models by training independent domain experts for specific tasks like math or code. These experts are composed into a unified model using a Mixture-of-Experts architecture and a lightweight router. This allows for adding, upgrading, or removing domain capabilities independently without degrading performance on other domains.

Key concepts

Mixture-of-Experts (MoE)
A model architecture where each transformer layer has multiple parallel Feed-Forward Networks (experts). A learned 'router' decides which expert processes a given input token, allowing the model to specialize in different tasks simultaneously.
Modular Training
The design isolates the training pipeline for each domain expert. This isolation prevents 'catastrophic forgetting,' meaning updates or training on one domain do not negatively affect the performance of other domains already trained.
Anchor Expert
A specific expert initialized from the original post-trained model's parameters and frozen. This anchor preserves the base model's general capabilities, ensuring that specialized experts only add new knowledge without erasing existing skills.

Terminology used across episodes

This episode discusses

The paper

Train Separately, Merge Together: Modular Post-Training with Mixture-of-Experts · Read on arXiv

Jacob Morrison, Sanjay Adhikesaven, Akshita Bhagia, Matei Zaharia, Noah A. Smith, Sewon Min

University of Washington · University of California, Berkeley · Allen Institute for AI

Extending a fully post-trained language model with new domain capabilities is fundamentally limited by monolithic training paradigms: retraining from scratch is expensive and scales poorly, while continued training often degrades existing capabilities. We present BAR (Branch-Adapt-Route), which trains independent domain experts, each through its own mid-training, supervised finetuning, and reinforcement learning pipeline, and composes them via a Mixture-of-Experts architecture with lightweight router training. Unlike retraining approaches that mix all domains and require full reprocessing for any update (with cost scaling quadratically), BAR enables updating individual experts independently with linear cost scaling and no degradation to existing domains. At the 7B scale, with experts for math, code, tool use, and safety, BAR achieves an overall score of 49.1 (averaged across 7 evaluation categories), matching or exceeding re-training baselines (47.8 without mid-training, 50.5 with). We further show that modular training provides a structural advantage: by isolating each domain, it avoids the catastrophic forgetting that occurs when late-stage RL degrades capabilities from earlier training stages, while significantly reducing the cost and complexity of updating or adding a domain. Together, these results suggest that decoupled, expert-based training is a scalable alternative to monolithic retraining for extending language models.

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Train Separately, Merge Together".

Jane: Here's the gist of the paper. The gist: BAR presents a modular approach to post-training that extends fully trained language models with new domain capabilities via independently trained mixture-of-experts.

Tom: How it works BAR trains independent domain experts, each through its own midtraining, supervised finetuning (SFT), and reinforcement learning (RL) pipeline, and composes them into a unified model via a Mixture-of-Experts architecture with lightweight router training<ref:2604.18473#pg2> The modular design decouples learning for each domain, enabling experts to be developed, added, or upgraded independently without re-training other components<ref:2604.18473#pg3> This approach enables updating individual experts independently with linear cost scaling and no degradation to existing domains<ref:2604.18473#pg4> The model architecture involves replacing each layer’s single FFN with multiple parallel FFN “experts” and a learned router that determines which experts process each input token<ref:2604.18473#pg6> In the final, fully trained BAR architecture, each transformer layer contains shared layers and expert FFNs, including an anchor expert that preserves M’s capabilities<ref:2604.18473#pg7> The anchor expert’s FFN weights are initialized from M’s post-trained FFN parameters and frozen, preserving M’s original capabilities<ref:2604.18473#pg8> Each domain expert’s FFN weights are initialized from the pre-trained (not post-trained) base model<ref:2604.18473#pg9> Training Algorithms BAR training proceeds in three stages<ref:2604.18473#pg10> Stage 1 involves independent expert training, where each domain expert Ei is trained independently as a two-expert MoE Each expert follows the full training pipeline as applicable to its domain, including mid-training on large-scale corpora, supervised finetuning (SFT), and reinforcement learning with verifiable rewards (RLVR)<ref:2604.18473#pg10> A key difference from pre-training is how shared layers are unfrozen during these stages<ref:2604.18473#pg10> For mid-training, they follow prior work and freeze all shared layers<ref:2604.18473#pg10> During SFT, embeddings and the LM head are unfrozen to support new tokens<ref:2604.18473#pg10> During RLVR, all shared parameters are unfrozen for the behavioral shifts RL requires<ref:2604.18473#pg10> Stage 2 involves expert merging, where the k domain experts are combined with the anchor expert into a single (k + 1)-expert MoE<ref:2604.18473#pg10> For shared parameters that may have diverged across expert training runs, BAR averages the diverged parameters across all domain expert models<ref:2604.18473#pg10> Stage 3 is router training, where only the router parameters are trained while all expert and shared weights are frozen<ref:2604.18473#pg10> The router is a small linear layer at each transformer block that produces a probability distribution over experts for each input token<ref:2604.18473#pg10> Training uses a stratified sample of SFT data from all domains, ensuring the router sees examples from every domain<ref:2604.18473#pg10> Experimental Setup and Results The setup starts with an existing fully post-trained model M, which is a 7B-parameter dense transformer Four new domains—math, code, tool use, and safety—are added using newer datasets that were not available when the initial model was trained The evaluation spans 19 benchmarks across seven categories: chat (AlpacaEval, IFEval), knowledge (MMLU, PopQA, SimpleQA), reasoning (BBH, GPQA, ZebraLogic), math (MATH, GSM8K), code (HumanEval+, MBPP+; Liu et al., 2023), tool use (BFCL; Patil et al., 2025), and safety (HarmBench; Mazeika et al., 2024) The overall score for the full BAR model is 49.1, exceeding the retraining (post-train only) baseline of 47.8 The results show that adding new experts improves overall performance without degrading math or code, demonstrating that experts do not interfere with each other For instance, the combined BAR 5×7B model achieves a score of 38.7 in the Math category compared to 31.3 for the initial model M Furthermore, upgrading individual experts demonstrates significant gains; for example, upgrading math from v1 (SFT only) to v2 (SFT + RL) improves math performance by +13.0 points in the combined model while other domains remain stable Discussion and Limitations The structural advantage of modular training is that it avoids catastrophic forgetting because each domain’s pipeline is isolated<ref:2604.18473#pg10> This advantage becomes apparent when extending modular training to post-training, where behavioral shifts during SFT and RL require adapting shared layers<ref:2604.18473#pg10> Modularity also provides a cost advantage: adding or upgrading domains scales linearly (one expert per domain), while re-training requires reprocessing all domains each time, scaling quadratically<ref:2604.18473#pg10> Limitations include that total parameters grow linearly with the number of experts, increasing inference cost as domains are added<ref:2604.18473#pg10> While activating 4 of 5 experts achieves nearly equivalent performance to the full model, better routing strategies are needed to reduce cost without sacrificing quality<ref:2604.18473#pg10> Upgrading the anchor model currently requires retraining all domain experts<ref:2604.18473#pg10> The paper concludes that BAR offers a practical path toward scalable, modular LM development by enabling domain experts to be upgraded independently, composed without interference, and integrated via lightweight router training<ref:2604.18473#pg10> The overall results position decoupled, expert-based training as a scalable alternative to monolithic retraining for extending language models<ref:2604.18473#pg10> Dataset Details The math mid-training uses the math mid-training mix from FlexOlmo (Shi et al., 2025), which combines the Dolmino Math Mix (OLMo et al., 2025) and FineMath4+ (Allal et al., 2025)<ref:2604.18473#pg9> The code v1 uses StarCoder-based (Li et al., 2023) mid-training data and code SFT data from Tulu 3 (Lambert et al., 2025), while code v2 uses the OLMo 3 (Olmo et al., 2025) code mid-training data<ref:2604.18473#pg9> Tool use training involves SFT only, using the Dolci tool use instruction data from OLMo 3 (Olmo et al., 2025)<ref:2604.18473#pg9> Safety training uses SFT only, drawing from CoCoNot (Brahman et al., 2024), WildGuardMix (Han et al., 2024), and WildJailbreak (Jiang et al., 2024b)<ref:2604.18473#pg9> Training Details & Hyperparameters Mid-training uses a decaying learning rate schedule on large-scale domain corpora with a learning rate of 9e-4<ref:2604.18473#pg9> Supervised Finetuning (SFT) each domain expert is fine-tuned for 2 epochs on a mixture of domain-specific and general SFT data with a learning rate of 1e-4<ref:2604.18473#pg9> Reinforcement Learning (RLVR) uses GRPO with verifiable rewards for math and code, utilizing a learning rate of 6e-7<ref:2604.18473#pg9> Router Training involves training only the router parameters while all expert and shared weights are frozen, using a learning rate of 1e-4 on a 5% stratified SFT sample<ref:2604.

Title and authors: Tom: So, this paper is called Train Separately, Merge Together: Modular Post-Training with Mixture-of-Experts. I’m looking at the title now and it’s pretty direct about what they did here.

Jane: It sounds like they are talking about taking an already trained language model and adding new skills without having to retrain the whole thing from scratch, which is a big deal for cost and speed.

Tom: Exactly. They're focusing on this modular approach where you train these domain experts separately, each with its own training pipeline, and then they just merge them using a mixture-of-experts setup with some lightweight router training at the end.

Lu: What I find interesting is that they’re addressing a gap in the research. They point out that existing methods for pre-training modularity don't really extend to post-training well, which is what this paper tackles specifically.

Meng: So, it’s not just another way to fine-tune; it sounds like they are proposing a fundamentally different structure for how we attach new knowledge to these massive models.

Lalam: I think the title makes sense because the core idea is training separately and then merging together so you don't have to do all that reprocessing every time you want a new domain.

The paper's summary: Jane: So, what’s the actual mechanism behind this BAR system? It sounds complicated, so I need you to break it down for someone who doesn't read the deep technical stuff.

Tom: Okay, basically they train these independent domain experts—say one for math and one for code—through their own mid-training, supervised finetuning, and reinforcement learning pipelines. Then they combine all those experts into a single model using a mixture-of-experts architecture guided by this little router that decides which expert handles which part of the input token.

Lu: The key point is that this design decouples the learning for each domain, meaning you can develop or upgrade one expert without needing to retrain the other components in any other domain. That’s a huge structural advantage they are highlighting here.

Meng: If I'm running this at scale, that decoupling is important because it means if we want to improve our safety expert, we only touch that part, not the whole massive model.

Lalam: And they detail exactly how the experts are initialized too; for instance, an anchor expert preserves the original model’s capabilities while the domain expert starts from a pre-trained base model.

The paper's improvements: Tom: The main improvement they push is that this modular design enables updating individual experts independently with linear cost scaling and no degradation to existing domains, which is much better than what you get when you retrain everything monolithically, which scales quadratically.

Jane: So, the numbers are important here; they show that adding new experts doesn't mess up the performance on things the model already knows well, like math or code.

Lu: For example, they show that upgrading an individual expert can lead to significant gains; for instance, upgrading the math expert from version one only using supervised finetuning to version two with supervised finetuning and reinforcement learning improves its performance by thirteen point zero points in the combined model while other domains remain stable <ref:2604.18473#pg1,supervised finetuning and reinforcement learning>.

Meng: That's encouraging for deployment because it shows we can iteratively improve specific features without risking regressions on our established capabilities.

Lalam: And they also show that experts don't interfere with each other when things are routed correctly; for instance, the combined five times 7B model achieves a score of thirty-eight point seven in the Math category compared to a score of thirty-one point three for the initial model M, which is way higher than just retraining everything.

Conclusion: Jane: So, if we put it all together, what is the big picture implication here for how we build these next generation language models?

Tom: The big picture is that this modular post-training approach offers a practical path toward scalable, modular AI development by enabling domain experts to be upgraded independently and composed without interference.

Lu: They are showing that decoupled, expert-based training can be a scalable alternative to monolithic retraining for extending language models. It’s about making the extension of capabilities much more efficient than retraining from scratch every time you want something new.

Meng: The limitation they mention is that while adding experts scales linearly, the total number of parameters grows linearly with the number of experts, which increases inference cost as you add more domains.

Lalam: And they also flag that upgrading the anchor model currently requires retraining all domain experts, and they point out that better routing strategies are needed to reduce cost without sacrificing quality.

More episodes

← Home