Train separately, merge together: Modular post-training with mixture-of-experts

There's a specific headache anyone who post-trains language models knows well: you fix or improve one capability, and something else quietly breaks. Teach a model better math reasoning through reinforcement learning, and its safety behavior can drift. Add tool-calling support, and instruction-following might slip. AI2's answer, detailed in a new technical report, is a recipe called BAR — Branch-Adapt-Route — that trains each skill as its own separate expert and only combines them at the very end.

The idea builds on AI2's earlier FlexOlmo work, which used a similar branch-and-merge trick during pretraining by freezing shared layers and only training feed-forward experts. But when the team tried the same trick during post-training, it flopped outright — during reinforcement learning with verified rewards, the reward curve stayed completely flat. Turns out post-training isn't just about stuffing in more knowledge, the way pretraining is. It reshapes behavior: output formats, reasoning style, safety guardrails. Those live in attention layers, embeddings, and the language modeling head, not just the feed-forward blocks that FlexOlmo left tunable.

So BAR unfreezes things gradually instead. During mid-training, shared layers stay frozen. During supervised fine-tuning, the embedding layer and LM head get unlocked — a change that pushed tool-use performance on the Berkeley Function Calling Leaderboard from 20.3 up to 46.4, since new function-calling tokens need somewhere to live. During RL, everything opens up, including attention, because RL apparently shoves the model around too much for frozen weights to absorb. Once each expert — math, code, tool use, safety — finishes its own pipeline, AI2 merges the diverged shared parameters by simple averaging, then trains a lightweight router on just 5% of the SFT data to decide which expert handles what.

The numbers back up the approach. Built on Olmo 2 at 7B scale, BAR averaged 49.1 across seven evaluation categories and 19 benchmarks, edging out standard full post-training retraining's 47.8, with math up 7.8 points and code up 4.7. Naive dense-model merging after mid-training, by contrast, basically destroyed the model — a score of 6.5. Full retraining from scratch, including mid-training, still wins outright at 50.5, but that requires the entire original pretraining setup, which most people simply don't have.

The more interesting payoff might be maintenance, not launch-day scores. AI2 swapped in a better code expert trained on newer data and RL, and code performance jumped 16.5 points in the merged model while everything else stayed put. They added an RL stage to an existing math expert and picked up 13 points in math, again without disturbing other domains. Only the swapped expert and the router needed retraining — a linear cost, versus the roughly quadratic cost of retraining a monolithic model every time one domain changes.