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

🎯 Pitch

For the first time, we show that independently trained specialist LLMs (for math, code, and knowledge) can be woven into a single, sparsely-activated model that beats all matched baselines—even a larger dense model—with no catastrophic forgetting. The trick: brief joint fine-tuning with load-balancing, which forces the router to actually use the experts rather than collapse to one.

1. Executive Summary

This paper introduces Branch-Train-MiX (BTX), a continued pretraining method that combines embarrassingly parallel domain-specialist training with a Mixture-of-Experts (MoE) architecture to produce a single unified LLM with broad capabilities. Starting from a Llama-2 7B seed model, the approach branches into independently trained experts on math, code, and Wikipedia data, then mixes their feedforward sublayers as experts within MoE layers while averaging the remaining parameters, followed by a brief MoE-finetuning stage that learns token-level routing across the combined experts. BTX improves the seed model's average performance across math, code, knowledge, reasoning, and MMLU benchmarks from 40.7 to 47.9, outperforming both a data-matched dense continued-pretraining baseline (44.5) and a compute-matched sparse upcycling baseline (47.3), while using less than half the additional training compute of a Llama-2 13B model whose overall score it surpasses (47.9 vs. 45.4). The embarrassingly parallel expert-training stage provides higher training throughput and reduced communication cost compared to synchronized MoE training, establishing that domain-specialist knowledge acquired through independent continued pretraining can be effectively integrated into a single sparsely-activated model through a short joint finetuning phase without catastrophic forgetting, though the benefits manifest only when load-balancing regularizes the router to prevent expert collapse and preserve the specialists' complementary capabilities.

2. Context and Motivation

The Core Problem: We Know How to Build Specialists, but Not How to Unify Them

The fundamental tension this paper tackles is one that has become increasingly acute as the field of large language models matures: specialization and generalization are in direct conflict during continued pretraining. On one hand, the most dramatic performance improvements in technical domains like mathematics and code generation come from taking a general-purpose pretrained model and continuing to train it on domain-specific data. This is the recipe behind Llemma (Azerbayev et al., 2023), CodeLlama (Rozière et al., 2023), and numerous other specialist LLMs—and it works remarkably well within each domain. The math expert in this paper nearly triples the seed model's GSM8K performance (from 14.7% to 39.5%) and improves MATH more than sevenfold (from 2.5% to 18.8%), as shown in Table 1.

On the other hand, these specialist models suffer from catastrophic forgetting of capabilities they previously possessed. Table 1 quantifies this vividly: the math expert's TriviaQA accuracy drops from the seed model's 58.5% to just 37.1%—a ~21 percentage point degradation. The code expert similarly drops on world knowledge tasks (Natural Questions: 16.4% → 11.5%; TriviaQA: 58.5% → 29.9%). Even the Wikipedia expert, which one might expect to retain knowledge capabilities, loses ground on reasoning benchmarks (MMLU: 46.1% → 43.1%). Each specialist model becomes a single-domain savant at the expense of its general competence.

This creates a practical dilemma. Organizations deploying LLMs need a single model that performs well across all domains—they cannot realistically ship three separate specialist models and route between them at inference time, at least not without introducing additional complexity (domain classifiers, separate serving infrastructure) and losing the ability to perform subsequent fine-tuning (SFT, RLHF) on a unified model. The paper frames this explicitly in Section 1:

"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... both of which can boost performance further, and are crucial steps in building aligned LLMs."

This is not merely a convenience issue—it is a blocker for the standard LLM alignment pipeline. Without a single unified model, post-training improvements that are now considered essential (instruction tuning, preference optimization) cannot be applied.

Why This Problem Matters Now

The problem's importance has grown in lockstep with several converging trends in the field:

The death of "one big pretraining run." The era where a single monolithic pretraining run on a general web corpus produced adequately capable models is ending. As LLMs asymptotically approach the limits of what general web text can teach them, continued pretraining on specialized, high-quality data becomes the primary lever for capability improvements. The Llama-2 seed model used in this paper was pretrained on ~2 trillion tokens of general data—yet its MATH accuracy is a paltry 2.5%. Specialist continued training adds another ~200 billion tokens of math data but yields an 18.8% MATH score. The incremental pound of specialized data is dramatically more valuable than the last pound of general data. But this creates a fragmentation problem: each domain improvement risks the degradation of others, and we lack principled methods for combining these gains.

The communication bottleneck in synchronized training. The standard approach to training large models—data parallelism with synchronized gradient updates—faces an increasingly severe scaling wall. As Section 1 notes:

"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."

This is not just a theoretical concern. Training runs involving thousands of GPUs spend a substantial fraction of their wall-clock time waiting for gradient synchronization. Hardware failures, which are statistically inevitable at scale, cascade into full-training halts rather than isolated retries. Any method that reduces synchronization requirements directly translates to faster, cheaper, and more reliable training.

The inference cost imperative. Models with strong performance across all domains tend to be very large (e.g., 70B+ parameters), making inference expensive. MoE architectures address this by decoupling total parameters from active (computed) parameters—a model can have many experts but only activate a subset per token. This paper's BTX models activate roughly 6.7–11.1 billion parameters during inference (Table 4) while containing expert knowledge equivalent to four separately trained 7B models plus the original seed weights. The efficiency of sparse activation is what makes domain consolidation practically deployable.

The self-improvement and specialization flywheel. If organizations can efficiently train domain experts and then merge them back into a generalist model, this enables a continuous improvement cycle: identify a new domain, train an expert, integrate it into the main model via BTX, evaluate, and repeat. The embarrassingly parallel nature of expert training means this cycle can be fast—multiple domains can be pursued simultaneously without coordination overhead.

Where Existing Approaches Fall Short

The paper identifies three broad families of prior work and diagnoses their limitations with respect to the unification problem:

Branch-Train-Merge (BTM): Independence Without Integration

The most directly relevant prior work is BTM (Li et al., 2022a; Gururangan et al., 2023), which the paper positions as a special case of BTX with 100% of compute allocated to expert training and 0% to MoE finetuning. BTM solves the communication bottleneck: it trains $N$ expert models completely independently on different data domains, with zero inter-model communication. At inference time, BTM classifies each input prompt into one or more domains (using tf-idf similarity between the prompt and the experts' training data representations) and averages the output token distributions of the top- $k$ matching experts.

The paper's diagnosis of BTM's weaknesses is twofold. First and most critically, BTM does not produce a unified model. The experts remain separate models with their own parameters. The lack of a single model means:

No SFT or RLHF can be performed on the combined system—each expert would need to be separately fine-tuned, and the domain-classifier-based routing might not transfer to instruction-following settings.
Inference requires running multiple models for each token (when multiple experts are selected), which is computationally expensive. The BTM baselines in this paper select Top-2 experts, meaning two full forward passes compute token distributions that are then averaged.
The domain classifier itself is a heuristic (tf-idf similarity) that may not capture the nuanced token-level mixing that an MoE model with learned routing can achieve.

Second, BTM lacks a mechanism for the experts to learn from each other. The experts are frozen after their independent training, so there is no opportunity for transfer learning—the math expert cannot benefit from the code expert's knowledge despite the domains' overlap (visible in Table 1, where math training improves code performance from 16.8 to 33.6).

The paper's empirical results confirm BTM's limitations. From Table 2, BTM Top-2 achieves an average score of 43.4 compared to BTX Top-2's 47.9—a substantial gap despite using the same underlying expert models. BTM underperforms significantly on math (21.5 vs. 27.4 for BTX) and world knowledge (26.9 vs. 41.0), suggesting that the MoE finetuning stage provides benefits beyond simply having access to expert parameters.

Conventional MoE Training: Synchronized From the Start

Mixture-of-Experts architectures have emerged as a powerful paradigm for scaling model capacity without proportionally scaling inference cost. Standard MoE training (Fedus et al., 2022; Jiang et al., 2024) initializes expert feedforward sublayers randomly (or as identical copies of a dense model's FF layer in "sparse upcycling") and trains them jointly with the router from the beginning of pretraining or continued pretraining.

The paper identifies several shortcomings of this approach:

Synchronization overhead. Conventional MoE training is fully synchronized across all experts. While only a subset of experts is active per token, all experts receive gradient updates (either through the training tokens that route to them or through auxiliary losses like load balancing). The all-to-all communication patterns required by MoE routing add to the already-heavy communication burden of synchronized training. The paper's compute efficiency argument is that spending some fraction of the total compute budget on embarrassingly parallel expert training (with zero communication between experts) and only the remainder on synchronized MoE finetuning is more efficient than spending 100% of compute on synchronized MoE training.

Lack of domain specialization in naturally trained MoE experts. An important observation that the paper draws from prior work (Jiang et al., 2024) is that MoE experts trained in the standard way do not naturally specialize by domain. Instead, they tend to develop more abstract, non-domain-specific specializations (e.g., different experts might handle different syntactic patterns or semantic roles that cut across domains). While this can work effectively—Mixtral, for instance, achieves strong performance—it means that standard MoE training does not directly leverage the domain structure that continued pretraining on specialized data can provide. The paper hypothesizes that explicitly initializing experts with domain-specialized knowledge (as BTX does) and then fine-tuning the router to leverage that specialization could be more efficient.

Sparse upcycling as a special case. The sparse upcycling baseline (Komatsuzaki et al., 2022) initializes MoE experts as identical copies of the dense model's FF layers and then trains everything jointly. This is a special case of BTX where 0% of compute is allocated to expert training and 100% to MoE finetuning. Table 3 shows the comparison: sparse upcycling (CM) achieves 47.3 average score with 252B tokens of MoE training, while BTX achieves 47.9 with 533B total tokens but only 80B of those in the synchronized MoE finetuning stage—the remainder were trained in the higher-throughput parallel expert stage. The 7.8 days of BTX training time vs. 7.9 days of sparse upcycling training time (despite BTX processing more total tokens) demonstrates the throughput advantage of the embarrassingly parallel phase. The paper argues this advantage would grow with more domains, since each additional domain adds another parallel expert training process while adding proportionally less to the MoE finetuning compute.

Dense Continued Pretraining: The Simple Baseline That Doesn't Scale

The most straightforward approach to improving a seed model on multiple domains is to simply continue training it on a mixture of all the domain-specific datasets. This is the "Dense (DM)" baseline in the paper. While simple, this approach faces two fundamental problems:

Task interference. Training on multiple domains simultaneously can lead to negative transfer—gradients from one domain's data can hurt performance on another domain. This is the multi-task learning version of catastrophic forgetting: since all parameters are shared, improvements on one task can come at the expense of others. The paper's results bear this out: the Dense baseline achieves 44.5 average score (Table 2), significantly below BTX's 47.9, despite training on exactly the same data mixture. The gap is particularly visible in math (18.3 vs. 27.4) and MMLU (49.8 vs. 52.5).

Undifferentiated resource allocation. Dense continued pretraining treats all domains uniformly in the sense that every parameter is updated for every batch, regardless of whether the batch contains math, code, or Wikipedia data. This means that the model cannot allocate specialized capacity to different domains—all domain knowledge competes for the same parameter space. In BTX, domain experts have dedicated capacity (their own FF layers) that can specialize without interfering with other domains, while shared components (attention, embeddings) benefit from cross-domain exposure during MoE finetuning.

How BTX Positions Itself

The paper frames BTX as synthesizing the strengths of BTM and MoE while mitigating their respective weaknesses. This synthesis is not simply additive—the paper argues there is a genuine synergy between the two paradigms.

From BTM, BTX inherits embarrassingly parallel expert training with high throughput and low communication cost. The expert training phase requires no synchronization between experts, no all-to-all communication, and no joint optimization. Each expert can be trained with its own optimal hyperparameters, data schedule, and resource allocation. Hardware failures in one expert's training do not affect others. The paper quantifies this throughput advantage in Table 3: BTX processes 533B total tokens in 7.8 days compared to sparse upcycling's 252B tokens in 7.9 days—more than double the data throughput for the same wall-clock time.

From MoE, BTX inherits sparse activation with learned token-level routing that produces a single unified model. Unlike BTM's heuristic domain classifier, the learned router can make fine-grained decisions at every token and every layer, potentially routing mathematical notation tokens to the math expert while routing natural language scaffolding tokens to the generalist expert within the same sequence. The unified model can undergo SFT and RLHF—critical capabilities that BTM lacks entirely.

The novel contribution is the two-phase allocation of compute: the paper explicitly frames BTX as generalizing two extremes. BTM allocates 100% of compute to expert training and 0% to MoE finetuning (no routing, no joint optimization). Sparse upcycling allocates 0% to expert training and 100% to MoE finetuning (no specialization prior, all learning happens jointly). BTX sits between these extremes, allocating some fraction to each phase, and the paper's empirical results suggest that this intermediate allocation is superior to either extreme.

The paper is careful not to claim that domain specialization naturally emerges in MoE training—in fact, it acknowledges the opposite: "such specialization does not seem to emerge naturally during MoE training" (Section 6, citing Jiang et al., 2024). BTX therefore does not rely on emergent specialization; it injects domain specialization through the expert initialization and then preserves or refines it during MoE finetuning. The routing analysis in Section 4.3.2 and Figure 3 confirms this: after MoE finetuning, the domain experts still receive disproportionate activation on their respective domain tasks (the Code expert is dominant in code tasks, the Wikipedia expert is dominant in knowledge tasks), demonstrating that the injected specialization survives the joint finetuning phase.

The paper also positions itself relative to the continual learning literature, specifically parameter isolation methods (Section 2). The idea that different parameters should handle different domains to avoid catastrophic forgetting is well-established in continual learning, but applying this principle at the scale of 7B-parameter LLMs with MoE layers is novel. The key insight is that MoE architecture provides a natural parameter isolation mechanism: each expert's FF layers are dedicated to its domain, preventing interference during expert training, while the router learns to allocate capacity dynamically during finetuning.

An important but subtle aspect of the paper's positioning is its focus on continued pretraining rather than pretraining from scratch. All experiments start from Llama-2 7B, which has already been pretrained on 2 trillion tokens of general data. BTX is therefore not a method for training models from scratch but rather for efficiently extending the capabilities of an existing strong base model into new domains. This is a practically important distinction: most organizations will not train foundation models from scratch but will want to adapt existing ones to their specific domain needs. BTX provides a recipe for doing so while maintaining general capabilities and producing a single deployable model.

3. Technical Approach

3.1 Reader Orientation

Branch-Train-MiX (BTX) is a three-stage continued pretraining method that starts with an existing general-purpose LLM, trains multiple independent copies on different specialized data domains in parallel without any communication between them, and then fuses those copies into a single Mixture-of-Experts model through a brief joint finetuning stage. It solves the problem of how to combine the benefits of domain-specialist training (which produces strong performance in individual domains but causes catastrophic forgetting of other capabilities) with the deployability of a unified model that can be further fine-tuned or used as a standard LLM at inference time. The shape of the solution is a two-phase compute allocation: an embarrassingly parallel "Branch and Train" phase where experts acquire specialized knowledge independently, followed by a synchronized "MiX" phase where those experts are integrated into a single sparsely-activated architecture with learned token-level routing that preserves their complementary strengths.

3.2 Big-Picture Architecture (Diagram in Words)

The system has three sequential stages:

Stage 1 — Branch: Starting from a pretrained seed model $\mathcal{M}$ (Llama-2 7B), create $N$ identical copies, one for each target domain. In the paper's experiments, $N = 3$ (math, code, Wikipedia), plus the original model is retained as a fourth "generalist" expert.

Stage 2 — Train: Each copy $\mathcal{M}_i$ is trained independently on its domain-specific dataset $D_i$ using the standard language modeling objective (next-token prediction). There is zero synchronization between these training processes—they are embarrassingly parallel. At completion, this yields $N$ expert LLMs, each specialized in its domain but suffering from catastrophic forgetting in others.

Stage 3 — MiX: The feedforward (FF) sublayers from all expert models are combined into MoE layers, with one expert per domain per layer. Self-attention weights and other parameters (embeddings, layer norms, etc.) are averaged across all experts. A newly initialized learned router is added at each layer to select which FF expert(s) to activate per token. The entire composite model is then fine-tuned jointly on a mixture of all domain data for a relatively short period (~80B tokens) so the router learns token-level mixing and the averaged attention weights become coherent.

Information flows linearly through these stages: seed model → parallel experts → fused MoE model → final unified LLM. The key architectural insight is that only the FF layers are domain-specialized and kept separate; all other components are shared through averaging and subsequent joint optimization.

3.3 Roadmap for the Deep Dive

First, the Branch and Train stages—how expert models are created, what data they are trained on, and why this embarrassingly parallel design matters for throughput and fault tolerance.
Second, the MiX stage architecture—how the FF sublayers from separate experts become experts in MoE layers, how self-attention and other weights are averaged, and what the router mechanism looks like.
Third, the MoE finetuning objective and dynamics—what loss functions are used, how load balancing prevents expert collapse, and what the router learns during joint training.
Fourth, the routing mechanism in detail—Top-k routing versus alternatives, how the router's linear transformation works, and how sparse activation reduces inference cost.
Fifth, the variations explored—different routing methods (Switch, soft routing, Sample Top-1), load balancing, expert splitting, and expert blending, along with their motivations and empirical outcomes.
Sixth, the training hyperparameters and compute budget—exact numbers for tokens, batch sizes, learning rates, and how the total compute is split between the parallel and synchronized phases.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a methodology and systems paper whose core idea is that the optimal way to combine domain-specialist LLMs into a unified generalist model is to treat their feedforward layers as pre-trained experts in an MoE architecture and then jointly fine-tune only the router and averaged non-expert parameters while keeping the expert FF weights frozen or lightly updated.

The Branch Stage: Creating Independent Expert Initializations

The Branch stage is the simplest part of BTX but its design reflects specific assumptions about what makes expert training effective. The procedure is:

Begin with a fully pretrained seed model $\mathcal{M}$ , which in the paper's experiments is Llama-2 7B (Touvron et al., 2023), a dense Transformer with 7 billion parameters pretrained on approximately 2 trillion tokens of general web data.
Create $N$ identical copies of $\mathcal{M}$ , producing $\mathcal{M}_1, \mathcal{M}_2, \ldots, \mathcal{M}_N$ . In the paper, $N = 3$ domains are targeted, yielding three copies. Each copy inherits the full parameter set of the seed model—all weights, including attention, FF, embeddings, and layer norms, are exactly duplicated. There is no parameter randomization or re-initialization at this stage.
Crucially, the original seed model $\mathcal{M}$ is also retained as a fourth "generalist" expert. The paper states: "we also include the original seed LLM as a 'generalist' expert so that its general knowledge is transferred to the final model." This means the final MoE model contains $N + 1 = 4$ experts per layer: math, code, Wikipedia, and the frozen original Llama-2 7B weights. The motivation is that during MoE finetuning, the router can fall back on the generalist expert for tokens that do not require specialized knowledge, preserving the seed model's broad capabilities.

Why branching instead of training from scratch: The paper explicitly builds on the finding that "continual training of a seed LLM on a specific domain of code can produce a strong domain expert model, and this converges much faster than starting from scratch" (citing Rozière et al., 2023). Starting from a general pretrained model provides the experts with strong language modeling and reasoning foundations, so the domain-specific continued training only needs to inject specialized knowledge rather than teaching basic syntax and semantics from scratch. This dramatically reduces the compute required for each expert.

Why fix the seed model: The decision to keep the original seed model as a frozen generalist expert, rather than continuing to train it or letting it be updated during MoE finetuning, ensures that the model retains a stable reference point for general capabilities. If the generalist were also updated, it could drift and lose the broad knowledge that the specialized experts are meant to complement.

The Train Stage: Embarrassingly Parallel Domain Specialization

After branching, each copy $\mathcal{M}_i$ is trained independently on its domain-specific dataset $D_i$ using the standard causal language modeling objective. The paper emphasizes that "each expert model $\mathcal{M}_i$ can be trained in complete separation from the others" and that "the whole training process becomes $N$ -way embarrassingly parallel."

Training objective: Each expert is trained to minimize the next-token prediction cross-entropy loss on its assigned dataset. There is no auxiliary loss, no knowledge distillation, and no regularization beyond standard weight decay. The objective is:

$\mathcal{L}_{\text{expert}} = -\frac{1}{|\mathcal{B}|} \sum_{x \in \mathcal{B}} \log P_{\mathcal{M}_i}(x_{\text{next}} \mid x_{\text{context}})$

where $\mathcal{B}$ is the current batch of training sequences, $x_{\text{context}}$ is the preceding token context, and $x_{\text{next}}$ is the target token. $P_{\mathcal{M}_i}$ is the probability assigned by expert model $i$ to the correct next token.

What it computes: The standard autoregressive language modeling loss—for each sequence in the batch, take the model's predicted probability of the actual next token given all previous tokens, take the negative log, and average over the batch. Lower loss means the model assigns higher probability to the true continuations.

Why this form: This is the pretraining objective the seed model was originally trained with, so continued training with the same objective ensures compatibility. Using a different objective (e.g., masked language modeling) would require architectural changes and would not leverage the seed model's autoregressive pretraining.

Domain datasets and training scales (Section 4.1.1):

Math expert: Trained on the same data sources and mixture used in Llemma (Azerbayev et al., 2023). The data mixture is detailed in Table 7: AlgebraicStack (13.57%), OpenWebMath (54.27%), Arxiv (27.14%), Github (2.99%), and Commoncrawl (5.01%). Training runs for 48,000 steps with a total of 201 billion tokens processed.
Code expert: Trained on the same data sources and mixture used in CodeLlama pretraining (Rozière et al., 2023). The mixture (Table 7) is: Code (82.18%), Natural language related to code (9.90%), and Natural language (6.93%). Training runs for 50,000 steps with 210 billion tokens total.
Wikipedia expert: Trained on Wikipedia documents extracted between June and August 2022, preprocessed to remove hyperlinks, comments, and formatting boilerplate. The mixture (Table 7) is: Wikipedia (90.91%) and Commoncrawl (9.09%). Since Wikipedia is a smaller dataset, training is limited to 42 billion tokens total.

Optimization hyperparameters (Section 4.1.2): All experts use the same optimizer configuration as the MoE finetuning stage: AdamW optimizer with weight decay 0.1, learning rate warmed up over 100 steps to a peak of $1 \times 10^{-4}$ , then decayed to 10% of peak with a cosine schedule. Batch size is 4 million tokens with a sequence length of 4096.

Why embarrassingly parallel matters (Section 3.1): The paper identifies three concrete benefits:

Linear throughput scaling: "It allows linear scaling of overall training throughput when scaling up the size of compute, while joint training often faces uncertain performance from increasing batch size." In synchronized training, increasing the number of GPUs requires increasing the global batch size, which can hurt convergence. In embarrassingly parallel training, each expert runs with its own batch size, and adding more experts simply adds more independent processes.
Reduced communication cost: There is zero all-to-all communication between expert training processes. Each expert's GPUs only communicate among themselves for data parallelism within that expert's training, not across experts. This is the extreme version of the communication reduction that methods like DiLoCo (Douillard et al., 2023) approximate through periodic synchronization.
Fault tolerance: "It is also more resilient, as a single training failure will only affect one of the $N$ training processes instead of halting the entire training." In a synchronized training run spanning thousands of GPUs, one GPU failure cascades into a full training halt. With BTX, a failure in the math expert's training does not affect the code or Wikipedia experts.

Consequences of independent training (Table 1): After training, each expert shows dramatic improvements in its domain but substantial degradation elsewhere:

Math expert: GSM8K 14.7% → 39.5% (+24.8 points); MATH 2.5% → 18.8% (+16.3 points). But TriviaQA drops 58.5% → 37.1% (-21.4 points).
Code expert: HumanEval 12.8% → 31.7% (+18.9 points); MBPP 20.8% → 40.2% (+19.4 points). But TriviaQA drops 58.5% → 29.9% (-28.6 points).
Wikipedia expert: Natural Questions 16.4% → 21.8% (+5.4 points). But MMLU drops 46.1% → 43.1% (-3.0 points).

These divergent expert capabilities are the raw material that the MiX stage must integrate—the challenge is to preserve each expert's strengths while recovering the general capabilities that individual experts lost.

The MiX Stage Architecture: Building the MoE Model from Expert Parts

The MiX stage is where BTX diverges fundamentally from Branch-Train-Merge. Instead of keeping experts as separate models with a domain-classifier-based routing heuristic, BTX constructs a single Transformer model where each layer's feedforward sublayer is replaced by an MoE module containing all $N + 1$ expert FF layers (including the original seed model as a generalist). The self-attention and other parameters are merged by averaging.

Feedforward layers become MoE experts (Section 3.2): For a standard Transformer, the computation at layer $l$ with input representation $x$ includes a feedforward sublayer:

$y = \mathtt{FF}^l(x)$

In Llama-2 7B, this FF sublayer is a two-layer MLP with a hidden dimension $d_{\text{FF}}$ that is larger than the model dimension $d_{\text{model}}$ (typically $d_{\text{FF}} \approx 4 \times d_{\text{model}}$ for standard dense Transformers, though the exact dimensions for Llama-2 are not specified in this paper). The FF sublayer accounts for roughly two-thirds of the model's total parameters.

BTX replaces this single FF sublayer with a mixture of $N + 1$ expert FF sublayers:

$\mathtt{FF}_{\text{MoE}}^l(x) = \sum_{i=1}^{N+1} g_i(W_l x) \cdot \mathtt{FF}_i^l(x)$

where $\mathtt{FF}_i^l$ is the feedforward sublayer from the $i$ -th expert model at layer $l$ , $W_l$ is the learned linear transformation of the router, and $g$ is the routing function that produces a sparse output (most entries are zero).

What it computes: For each token at each layer, first compute the router's linear transformation $W_l x$ to produce a score for each expert, apply the routing function $g$ to select which experts to activate (typically 1 or 2 out of $N+1$ ), compute only the selected experts' FF outputs, weight them by the router's scores, and sum. Experts with zero routing weight are not computed at all, saving FLOPs.

Why this form: The key property is sparsity through $g$ . If all experts were activated, the computation would be $(N+1) \times$ the cost of a dense FF layer—prohibitively expensive. By activating only $k$ experts (where $k \ll N+1$ , typically $k = 2$ ), the inference cost remains roughly $k \times$ the cost of a single FF, comparable to a moderately larger dense model, while the total parameter count grows with $N$ . This decouples capacity from computation.

Self-attention and other parameters are averaged (Section 3.2): For all parameters that are not feedforward sublayers—self-attention projections (query, key, value, output), embeddings, layer norms, and any other parameters—BTX simply averages the weights across all $N + 1$ experts. If a parameter matrix in expert $i$ is $\theta_i$ , the combined parameter is:

$\theta_{\text{combined}} = \frac{1}{N+1} \sum_{i=1}^{N+1} \theta_i$

Motivation for averaging attention: The paper states this is based on "the assumption that the self-attention layers are less domain specialized than the feedforward layers." The reasoning is that attention mechanisms learn general patterns of token interaction (syntax, long-range dependencies, coreference) that are largely domain-invariant, while FF layers store more domain-specific factual and procedural knowledge. This assumption is partially validated by the routing analysis (Figure 3), which shows that expert routing decisions vary strongly by domain, suggesting domain specialization is concentrated in the FF layers that the router selects among.

Why not average FF layers too: If the FF layers were also averaged, the domain specialization acquired during expert training would be diluted or lost. The whole point of keeping FF layers separate is to preserve the specialized knowledge each expert acquired. Averaging would blend the math expert's FF weights (which encode mathematical reasoning patterns) with the Wikipedia expert's FF weights (which encode factual knowledge), likely destroying both specializations.

Router parameters are the only new learnable parameters: The router at each layer consists of a single linear transformation $W_l$ that maps from the model dimension $d_{\text{model}}$ to $N+1$ logits (one per expert). The paper notes these are "negligible in size compared to the rest of the network." For a model with hidden dimension ~4096 (typical for 7B models) and 32 layers with 4 experts each, the router parameters total $32 \times 4096 \times 4 = 524,288$ parameters—roughly 0.007% of 7 billion parameters.

The MoE Finetuning Objective and Dynamics

After constructing the MoE model by combining expert FF layers and averaging other parameters, BTX finetunes the entire model jointly on a mixture of all domain data. This is the stage where the router learns token-level mixing and the averaged attention weights become optimized for the combined model.

Finetuning data mixture (Section 4.1.1 and Table 7): The MoE finetuning stage uses a mixture of data from all domains, including the original Llama-2 pretraining data. The sampling ratios across domains are: math data (30.16%), code data (40.31%), Wikipedia data (10.30%), and original Llama-2 pretraining data (19.23%). The paper trains for 80 billion tokens in the default configuration.

Why this mixture ratio: The paper does not provide a detailed justification for these specific ratios, but the pattern suggests heavier weight on domains where continued pretraining caused the most catastrophic forgetting in individual experts. Code and math receive the largest shares, likely because they showed the largest improvements during expert training and the MoX stage needs sufficient signal to integrate those improvements while recovering general capabilities, while Wikipedia knowledge is relatively easier to recover given the generalist expert's retained capabilities.

The language modeling objective: During MoE finetuning, the model is still trained with the standard next-token prediction loss on the combined data mixture. There is no domain-conditional objective—the model does not know which domain a sequence comes from except through the tokens themselves. The router must learn to infer from the token content which expert(s) to activate.

Load balancing auxiliary loss (Section 3.3): A well-known failure mode of MoE training is "dead experts"—experts that the router never activates, and therefore never receive gradient updates, making them permanently useless. Once a router's initial random weights slightly favor some experts over others, the favored experts get more training signal and improve faster, widening the gap, while unfavored experts spiral into permanent irrelevance.

To prevent this, BTX adds a load balancing loss term similar to the one in Fedus et al. (2022):

$\mathcal{L}_{\text{LB}} = \alpha N \sum_{i=1}^{N} u_i p_i$

where:

$\alpha$ is a hyperparameter controlling the strength of the balancing penalty. The paper uses $\alpha = 0.01$ for default Top-2 routing.
$N$ is the number of experts ( $N = 4$ in the paper's main experiments, including the generalist).
$u_i$ is the fraction of tokens in the current batch $\mathcal{B}$ that are routed to expert $i$ : $u_i = \frac{1}{|\mathcal{B}|} \sum_{x \in \mathcal{B}} g_i(W_l x)$ , where $g_i(W_l x)$ is the routing weight assigned to expert $i$ after the Top-k operation.
$p_i$ is the average softmax probability assigned to expert $i$ before Top-k truncation: $p_i = \frac{1}{|\mathcal{B}|} \sum_{x \in \mathcal{B}} \text{SoftMax}_i(W_l x)$ .

What it computes: For each expert $i$ , multiply the fraction of tokens actually routed to it ( $u_i$ ) by the average probability the router assigns to it ( $p_i$ ), sum over all experts, and multiply by $\alpha N$ . This loss is computed at each MoE layer and added to the language modeling loss.

Why this form: The product $u_i p_i$ is minimized when the distribution of routing assignments is uniform—each expert receives an equal share of tokens. If one expert dominates ( $u_i$ is large), the term penalizes this proportionally. If an expert is unused ( $u_i = 0$ ), the term is zero for that expert, so the penalty doesn't force dead experts back to life directly; rather, it prevents the router from concentrating too heavily on a subset of experts in the first place. The $\alpha N$ scaling makes the penalty strength roughly independent of the number of experts and provides a tunable knob for the balance between load balancing and task performance.

The total training objective during MoE finetuning is:

$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{NLL}} + \mathcal{L}_{\text{LB}}$

where $\mathcal{L}_{\text{NLL}}$ is the standard negative log-likelihood language modeling loss and $\mathcal{L}_{\text{LB}}$ is summed over all MoE layers.

Freezing experts during MoE finetuning (Section 4.3.1, Table 5): The paper experimented with freezing the FF layers initialized from domain experts during MoE finetuning (training only the router, averaged attention weights, and other non-expert parameters). The result: "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 is a significant finding: it suggests that the domain experts already acquired sufficient specialization during the Branch-Train phase, and the MoE finetuning phase primarily serves to train the router to appropriately select among already-competent experts and to optimize the averaged attention weights. This justifies the paper's two-phase compute allocation—the expensive part (training experts on hundreds of billions of tokens) can be done independently, and only a relatively short joint phase is needed for integration.

Routing Mechanism: How the Router Selects Experts

The router is the central mechanism that makes the MoE model work—it decides, for each token at each layer, which expert FF sublayer(s) to activate. The paper explores several routing variants.

Top-k routing (default): The paper's primary routing method is Top-k with $k = 2$ , described in Section 3.2:

$g(W_l x) = \text{SoftMax}(\text{TopK}(W_l x))$

The computation proceeds step by step:

Linear projection: Compute $z = W_l x$ , where $W_l \in \mathbb{R}^{(N+1) \times d_{\text{model}}}$ is the router's learned weight matrix at layer $l$ , $x \in \mathbb{R}^{d_{\text{model}}}$ is the input representation at that layer, and $z \in \mathbb{R}^{N+1}$ is a vector of unnormalized scores (logits), one per expert.
TopK selection: Keep only the $k$ largest values in $z$ , setting all others to $-\infty$ . For $k = 2$ , this means only the two experts with the highest logits survive.
Softmax normalization: Apply softmax to the truncated logits: $g_i = \frac{\exp(z_i)}{\sum_{j \in \text{top-k}} \exp(z_j)}$ for $i$ in the top-k set, and $g_i = 0$ otherwise. This produces a probability distribution over only the selected experts.
Weighted combination: Compute the MoE output as $\sum_{i=1}^{N+1} g_i \cdot \mathtt{FF}_i^l(x)$ . Since $g_i = 0$ for all but $k$ experts, only $k$ FF forward passes are needed.

What it computes: For each token, select the $k$ most relevant domain experts based on the token's representation, compute their FF outputs, and form a weighted combination where the weights are the softmax of the router's scores. The output has the same dimensionality as a standard FF output ( $d_{\text{model}}$ ), so it feeds naturally into the next layer's residual stream.

Why this form: Top-k routing with $k > 1$ provides two benefits over Top-1: (1) it allows the model to combine knowledge from multiple domains for a single token (e.g., a math symbol might benefit from both the math expert's reasoning and the generalist expert's language understanding), and (2) it provides redundancy that makes training more stable—if one expert is temporarily suboptimal, the other can compensate, preventing the kind of routing collapse that plagues Top-1 methods.

Switch routing (Top-1): The paper also experiments with Switch routing (Fedus et al., 2022), which is Top-1 with a capacity factor. The capacity factor is a hard limit on how many tokens can be routed to each expert; tokens that would exceed an expert's capacity are dropped (their FF output is zeroed). The paper uses a capacity factor of 1.5, meaning each expert can handle up to $1.5 \times \frac{\text{total tokens}}{N}$ tokens per batch. Switch routing "was subpar in average performance" (Table 4: average score of 24.7 vs. 34.6 for Top-2 at 10B tokens of finetuning), which the paper attributes to the lack of redundancy and the difficulty of training with token dropping.

Soft routing: In soft routing, the function $g$ is simply $\text{SoftMax}(W_l x)$ without any Top-k truncation, meaning all $N+1$ experts are activated for every token. The paper notes this "is likely to provide the best performance, but it comes at the expense of increased compute" since all experts must be computed. Table 4 confirms this: soft routing achieves the highest average scores (35.8 at 10B tokens, 37.3 at 40B tokens) but requires $N+1 = 4$ active FF computations per token versus 2 for Top-2 or 1 for Switch/Sample Top-1. The active parameter count for soft routing is 19.7B versus 11.1B for Top-2 and 6.7B for Sample Top-1.

Sample Top-1 routing (Section 3.3): This is a stochastic routing method that attempts to combine the efficiency of Top-1 (only one expert computed per token) with the training stability of having a gradient signal for non-selected experts. It uses the Gumbel-Softmax reparameterization trick (Jang et al., 2016):

$g(W_l x) = \text{GumbelSoftMax}(W_l x, \tau)$

At training time, this produces a soft sample from the categorical distribution over experts, but then zeros out all but the largest value, so only one expert is actually computed. The non-zero Gumbel-Softmax probabilities still provide a gradient signal to all experts through the soft sample, even though only one is forward-computed. At inference time, hard sampling is used (no soft gradients needed).

The temperature $\tau$ is annealed during training to gradually reduce the discrepancy between the soft training-time distribution and the hard inference-time distribution. The paper uses $\tau = \max(0.5, \exp(-r t))$ where $r = 1 \times 10^{-4}$ and $t$ is the number of training steps. This starts at temperature 1.0 at $t=0$ and decays exponentially toward 0.5 over training.

What it computes: Similar to Top-1, but with stochastic expert selection during training rather than deterministic selection. The randomness during training helps explore different routing assignments and prevents premature convergence to suboptimal expert allocations.

Why this form: The key advantage of Sample Top-1 over Top-2 is efficiency at inference: only one expert is activated per token (6.7B active parameters), compared to two for Top-2 (11.1B active parameters). Table 4 shows Sample Top-1 achieves competitive performance with less inference compute: 36.9 average score at 160B tokens versus Top-2's 37.3 at 80B tokens. The Gumbel-Softmax training provides gradient signal to non-selected experts (unlike deterministic Top-1), which helps prevent dead experts.

Router weight initialization: The paper does not explicitly specify how router weights $W_l$ are initialized for BTX, but the sparse upcycling baseline notes that router parameters are "randomly initialized." It is reasonable to assume the same random initialization is used for BTX routers, since they are the only new parameters introduced in the MiX stage.

Load Balancing Analysis and Impact on Expert Specialization

The routing analysis in Section 4.3.2 and Figure 3 provides critical insight into how load balancing shapes the model's behavior.

Without load balancing (Figure 3, bottom): The distribution of routing decisions is highly skewed:

The Math expert dominates routing across almost all domains and layers. In code tasks, knowledge tasks, and reasoning tasks, the Math expert receives the majority of token assignments.
The Code expert is essentially dead—it receives negligible routing probability across all domains, including code tasks. The paper explicitly notes: "A dead Code expert comes 'back to life' with load balancing introduced in training."
This collapse occurs because the Math expert was trained on the most tokens (201B) and on data that overlaps with other domains (Table 1 shows math training improved code performance, indicating shared knowledge). The router, without balancing pressure, gravitates toward the expert that provides the strongest initial signal, starving others of training gradients.

With load balancing (Figure 3, top): The routing distribution becomes much more uniform:

The Code expert becomes the dominant expert in code domain tasks. For HumanEval and MBPP evaluations, Code receives the highest probability across most layers.
The Wikipedia expert is the dominant expert in world knowledge tasks (Natural Questions, TriviaQA).
The Math expert remains relevant for math tasks but no longer monopolizes other domains.
The Llama-2 generalist expert maintains a consistent baseline presence across all domains.

Why load balancing is necessary (Table 5): The ablation in Table 5 quantifies the impact. With load balancing, HumanEval performance is 27.4; without load balancing, it drops to 19.5. However, GSM8K performance improves without load balancing (34.6 vs. 29.8), because without load balancing the router heavily favors the Math expert, which benefits math tasks at the expense of code. This reveals a trade-off: load balancing forces the router to use all experts, which helps neglected domains (code) but can reduce performance on the domain that would naturally dominate (math). The default BTX configuration accepts this trade-off because the goal is balanced multi-domain performance.

Token-level routing examples (Table 6): The paper provides concrete routing examples that illustrate the router's behavior:

In a GSM8K math word problem, tokens are routed to a mix of experts: numerical tokens and operations route to Math and Llama-2, while natural language scaffolding routes to Llama-2 and Wikipedia. The Math expert handles the domain-specific computational reasoning, while the generalist and Wikipedia experts handle the common-sense language understanding.
In a HumanEval code generation task, Python keywords and structure tokens are routed to Code and Llama-2, demonstrating that even within code tasks, the generalist expert contributes to basic syntax while the Code expert handles domain-specific patterns.
In a Natural Questions factoid query, the key entity tokens are routed to Wikipedia (the domain expert for factual knowledge), while the surrounding query structure is handled by Llama-2.

The paper observes: "Tokens that were routed to the in-domain expert are underlined." This confirms that domain specialization partially survives the MoE finetuning stage—the experts initialized from domain-specific training are indeed preferentially activated on their domain's tokens, even though the router is free to assign any token to any expert.

Layer-wise routing patterns (Figure 4, Appendix Section 10): The per-layer routing analysis shows that routing distributions "slightly vary in the first few layers, but quickly become indistinguishable from layer to layer." This suggests that the router learns a relatively consistent strategy across most layers, rather than dramatically different routing at early vs. late layers. One exception is Switch routing, where "Math expert becomes dominant across tasks in the last model layer," indicating that Switch is particularly susceptible to expert collapse at late layers where the representations are most task-specific.

Expert Splitting and Blending Variations

The paper explores two variations on how domain experts are structured within the MoE layers (Section 3.3, results in Table 5).

Splitting experts: The number of MoE modules per layer can be increased by splitting each domain's FF sublayer into multiple chunks. Given $N$ domains and an FF activation size of $d_{\text{FF}}$ , each FF is split into $C$ chunks with dimension $d_{\text{FF}} / C$ . The resulting MoE layer has $N \times C$ modules. With $N = 4$ and $C = 2$ , this produces 8 modules.

The paper tested this with both Top-2 of 8 and Top-4 of 8 routing (Table 5). Top-2 of 8 (activating 2 out of 8 modules) performed poorly: 22.2 average score vs. 34.7 for standard BTX with 4 experts. Top-4 of 8 (activating 4 out of 8, matching the active parameter count of standard Top-2 of 4) achieved 34.5, essentially equal to the 4-expert version. The paper concludes that splitting experts "does not improve performance, even if Top-4 routing is used to match the active number of parameters."

Why splitting doesn't help: Splitting each domain expert into multiple chunks breaks the correspondence between one expert module and one domain. A single domain's knowledge is now fragmented across multiple modules, and the router must learn to co-activate all chunks from the same domain to reconstruct the original FF computation. This adds complexity without increasing total capacity or specialization. The fact that Top-4 of 8 matches but doesn't exceed Top-2 of 4 suggests the additional routing flexibility doesn't translate to better domain specialization.

Blending experts: Instead of each MoE module corresponding to exactly one domain's FF layer, the paper tries "including all domains in each MoE expert." Specifically, each domain expert's FF layers are split into $N$ chunks, and then the $n$ -th chunks from all domains are merged to build the $n$ -th MoE expert. This way, each MoE expert contains the same proportion of parameters from all domains.

The motivation is an observation from prior work (Jiang et al., 2024) that "MoE experts trained in a standard way do not show domain specialization, but rather are activated uniformly across different domains." The blending experiment tests whether BTX's domain experts should be explicitly de-specialized before MoE finetuning.

The result in Table 5 is stark: blending experts drops performance catastrophically, from 34.7 average to 22.2. The paper concludes: "domain FF layers cannot be mixed in this way." The catastrophic failure likely occurs because blending destroys the coherent domain specialization that each expert acquired during individual training. Each blended module contains a mixture of math, code, Wikipedia, and general FF parameters, which produces incoherent internal representations that the router cannot effectively leverage.

Why this is a significant negative result: The failure of blending confirms that domain specialization is genuinely encoded in the FF parameters and that this specialization is useful for the MoE model. The BTX approach works precisely because it preserves the domain structure of the experts, and the router learns to select experts based on token content, exploiting their complementary specializations. Destroying that structure by blending eliminates the benefit of the Branch-Train phase.

Compute Allocation and Training Budget

The paper's results depend heavily on the specific allocation of compute between the Branch-Train (parallel) and MiX (synchronized) phases. This allocation is what distinguishes BTX from its two special cases.

Total compute budget (Figure 2, Table 3): BTX training uses a total of 926.1 GPU-days of compute (measured from the seed model, not including the original Llama-2 pretraining). Within this, the parallel expert training phase accounts for approximately 77% of the total compute, and the MoE finetuning phase accounts for approximately 23%. The total number of tokens processed across all phases is 533 billion (201B + 210B + 42B for experts = 453B, plus 80B for MoE finetuning).

Training time vs. throughput (Table 3): Despite processing more than double the tokens (533B vs. 252B), BTX's total training time is slightly less than sparse upcycling's: 7.8 days vs. 7.9 days. This is possible because the parallel expert training phase has higher throughput per GPU-day than synchronized MoE training—the absence of all-to-all communication and the ability to use independent data pipelines and batch sizes per expert overcome the overhead of training more total tokens.

Special cases of BTX (Section 6): The paper explicitly frames BTX as generalizing two extremes:

"BTM with 100% compute allocated to expert training and 0% on MoE finetuning, and sparse upcycling with 0% compute allocated to expert training and 100% on MoE finetuning"

The paper acknowledges that it has not "performed a thorough sweep of the compute allocation ratio between expert training and MoE training" (Section 6). The chosen allocation (~77% expert training, ~23% MoE finetuning) was determined by pragmatic considerations rather than systematic optimization: the expert training volumes were chosen to match prior work (Llemma and CodeLlama data quantities), and the MoE finetuning volume was chosen to provide sufficient tokens for router learning.

Optimization hyperparameters for MoE finetuning (Section 4.1.2): The same optimizer configuration is used for both expert training and MoE finetuning. The paper states: "We use the AdamW optimizer with weight decay 0.1, and anneal the learning rate to the peak of $1 \times 10^{-4}$ with 100 steps of warmup, and decay to 10% of the peak with a cosine schedule. We use a batch size of 4M tokens with a sequence length of 4096."

This is important because it means the MoE model is fine-tuned with the same learning rate and schedule as the experts were trained with—there is no special treatment for the newly introduced router parameters. The router must learn its routing strategy under the same optimization regime as the expert FF layers.

First layer uses soft routing for Sample Top-1: The paper notes that for the Sample Top-1 configuration, "for the first layer only, we used soft-routing instead." This is because the first layer's input is the token embedding, which may not contain sufficient signal for hard routing decisions. Soft routing in the first layer ensures that all experts contribute to the initial representation, providing a richer input for subsequent layers' routing decisions.

Summary of Design Choices and Their Justifications

Branching from a pretrained seed model rather than training experts from scratch: Continued pretraining converges much faster than training from scratch because the seed model already possesses general language understanding and reasoning capabilities. This is the finding from CodeLlama and Llemma that BTX builds upon.
Embarrassingly parallel expert training with zero inter-expert communication: Eliminates the communication bottleneck that limits scaling of synchronized training. Provides linear throughput scaling, reduced all-to-all communication cost, and fault isolation between expert training processes.
Only FF layers become MoE experts; attention is averaged: Based on the hypothesis that domain specialization is concentrated in FF layers (which store factual and procedural knowledge) while attention mechanisms handle domain-invariant syntactic and structural patterns. The routing analysis (Figure 3) partially validates this: routing decisions vary by domain, suggesting domain-relevant information is concentrated in the expert FF layers that the router selects among.
Including the original seed model as a frozen generalist expert: Ensures that general capabilities are always available, preventing the MoE model from losing broad competence when no domain expert is clearly appropriate. The generalist provides a stable fallback for tokens that don't fit any specialized domain.
Top-2 routing rather than Top-1: Provides redundancy and smoother training dynamics. Top-2 allows combining knowledge from two domains per token, which is important because many tokens benefit from both domain-specific and general knowledge. The gap between Top-2 (34.6 average at 10B tokens) and Switch Top-1 (24.7) in Table 4 confirms the importance of $k > 1$ .
Load balancing with $\alpha = 0.01$ : Prevents dead experts, which the routing analysis shows is a real risk—without load balancing, the Code expert becomes completely unused. The trade-off (slightly reduced math performance, significantly improved code performance) is accepted because BTX targets balanced multi-domain capability.
Freezing expert FF weights during MoE finetuning works nearly as well as updating them: This finding (Table 5) suggests the Branch-Train phase is the primary source of domain knowledge, and the MiX phase mainly serves to train the router and optimize shared parameters. This validates the two-phase compute allocation.
Not blending domain experts across MoE modules: The catastrophic failure of blending (Table 5: 22.2 vs. 34.7 average) confirms that domain specialization is encoded coherently in each expert's FF parameters and that preserving the correspondence between experts and domains is essential for BTX to work.
Uniform sampling from all domains during MoE finetuning: While the paper uses weighted ratios (30% math, 40% code, 10% Wikipedia, 19% general), it acknowledges not experimenting with different data mixtures for MoE finetuning (Section 6). The chosen ratios reflect the relative volume of each domain's training data.
Short MoE finetuning relative to expert training (80B vs. 453B tokens): The router and averaged parameters can be optimized much faster than domain expertise can be acquired, so the majority of compute is allocated to the embarrassingly parallel phase where throughput is highest and communication cost is zero.

4. Key Insights and Innovations

Innovation 1: The Two-Phase Compute Allocation as a New Axis for Scaling Efficiency

The most intellectually distinctive move in BTX is not the architecture itself—Mixture-of-Experts models and embarrassingly parallel training are both established techniques—but the explicit treatment of the compute allocation ratio between parallel and synchronized training as a design dimension to be optimized. Prior work had explored the extremes: Branch-Train-Merge (Li et al., 2022a) allocates 100% of compute to independent expert training with zero joint optimization, while sparse upcycling (Komatsuzaki et al., 2022) allocates 0% to independent training and 100% to synchronized MoE training from randomly initialized or cloned experts. These were studied as alternative methods, not as points on a continuous spectrum.

BTX reframes them as special cases of a single procedure parameterized by one number: the fraction of compute spent on parallel expert training before synchronized integration. This is a conceptual shift analogous to how Hoffmann et al. (2022) reframed model size and data quantity from independent design choices into components of a joint compute-optimal allocation. The paper does not find the optimal allocation ratio—it acknowledges not sweeping this dimension (Section 6)—but the reframing itself is the contribution. It establishes that where you spend synchronization overhead matters independently of how much total compute you spend, a degree of freedom that was previously invisible in the method-vs-method comparisons of the literature.

The empirical evidence for this reframing's validity is in Table 3 and Figure 2 (left). BTX (77% parallel, 23% synchronized) achieves higher average performance (47.9) than sparse upcycling with compute-matching (47.3) despite both spending the same GPU-days, because the parallel phase has higher throughput—BTX processes 533 billion tokens versus 252 billion for the same wall-clock time. BTM (100% parallel, 0% synchronized) achieves only 43.4 despite using the same experts, because it lacks the integration stage. The extremes underperform, and an intermediate point outperforms both. This is the signature of a genuine optimization dimension, not just a binary method choice. The finding that freezing expert FF layers during the MoE finetuning phase has "little impact on performance across all tasks" (Section 4.3.1, Table 5) further supports the logic: the expensive parallel phase produces domain knowledge that is nearly sufficient on its own, and the synchronized phase primarily serves to organize it through the router and averaged attention weights.

This insight is fundamental rather than incremental because it changes how a practitioner should think about allocating a fixed training budget across model components. Before BTX, one might ask "should I train one generalist or several specialists?" After BTX, the question becomes "what fraction of my budget should I spend on independent specialization versus joint integration, and how does that fraction change with the number of domains?" The answer almost certainly depends on the number of domains, the data volume per domain, and the base model's starting capabilities—opening a new empirical scaling law to be characterized by future work.

Innovation 2: The Empirical Refutation of Natural Domain Specialization in MoE as a Sufficient Strategy

The paper provides a concrete and important negative result with positive implications: MoE experts trained jointly from random or cloned initialization do not naturally develop the kind of domain specialization that explicit pretraining on domain data produces, but you can inject that specialization through initialization and it partially survives joint finetuning.

The field had two competing intuitions about this. One, grounded in work like Mixtral (Jiang et al., 2024), observed that standard MoE training produces experts that specialize in abstract, non-domain-specific patterns (syntactic roles, semantic categories) and that this works well—domain-level specialization is unnecessary for strong performance. The other intuition, implicit in the Branch-Train-Merge and domain-conditional MoE work (Gururangan et al., 2021), was that domain specialization is valuable and should be explicitly engineered.

The BTX paper provides evidence that both intuitions are wrong in interesting ways. The failure of the "blending experts" variant (Table 5: 22.2 average score vs. 34.7 for standard BTX, a catastrophic 12.5-point drop) demonstrates that explicitly destroying domain structure before MoE finetuning is harmful—domain specialization matters. Yet the success of standard BTX demonstrates that you don't need to perfectly preserve domain specialization either. Figure 3 shows that after MoE finetuning, the Code expert is dominant in code tasks and the Wikipedia expert is dominant in knowledge tasks, but there is substantial cross-domain routing: the Math expert and generalist contribute across all domains. The router learns a soft, token-level version of domain specialization that is more nuanced than the hard, model-level specialization of Branch-Train-Merge.

This is a fundamental reframing of the relationship between architecture and domain knowledge. It suggests that MoE architectures should not be viewed as black-box capacity scalers that learn their own internal abstractions (the Mixtral view), nor as direct mirrors of human-specified domain ontologies (the domain-conditional view), but as structured prior injections that bias the model toward using specific parameters for specific types of knowledge while allowing the router to learn where those biases are useful and where they should be overridden. This explains why BTX outperforms sparse upcycling (which has no domain prior) and BTM (which has a rigid, non-learnable domain prior): the learnable router with load balancing finds the optimal middle ground between the two extremes.

The evidence is anchored in the routing analysis of Section 4.3.2 and Figure 3, where we see that without load balancing the Math expert consumes nearly all tokens across all domains—the prior overwhelms the learning—while with load balancing the domain-appropriate experts dominate but others contribute meaningfully. This balance is what neither extreme achieves.

Innovation 3: The Diagnostic Power of the Load Balancing Ablation

Load balancing in MoE training is typically understood as a practical necessity—an auxiliary loss that prevents expert collapse during training. The paper's routing analysis (Figure 3, Section 4.3.2) elevates load balancing from a technical workaround to a diagnostic tool that reveals the underlying structure of domain expertise in the fused model.

The comparison between BTX with and without load balancing (Figure 3, top vs. bottom) tells a clear story: without load balancing, the Math expert consumes 60-80% of token routing across code, knowledge, and reasoning tasks, while the Code expert effectively dies (its routing probability goes to zero in later layers). The paper identifies the mechanism: the Math expert was trained on the most data (201B tokens) and on data that overlaps with other domains (math training improves code performance, Table 1), so the router, unconstrained, gravitates toward it as the most generally useful expert.

This diagnostic reveals something non-obvious about the relationship between the experts. The Math expert doesn't dominate because math is more important—it dominates because its training data has the broadest coverage overlap with other domains. This would not be visible from per-task accuracy metrics alone. It implies that expert training order, data volume, and domain overlap are critical variables for BTX that the paper does not fully explore, and that load balancing serves not just as a training stabilizer but as a mechanism for enforcing that complementary expert knowledge is actually utilized.

The practical significance of this insight goes beyond the metric gains. Without load balancing, BTX would underperform on code (19.5 vs. 27.4 HumanEval, Table 5) but overperform on math (34.6 vs. 29.8 GSM8K), leading a practitioner to incorrectly conclude that BTX doesn't help code integration. The load balancing ablation thus functions as a controlled experiment that isolates whether the MoE integration is genuinely leveraging multiple experts or simply defaulting to the strongest one. This is a methodological contribution—a way to verify that MoE models are actually mixing expertise rather than degenerating to single-expert behavior—that future work in this area should adopt as a standard diagnostic.

5. Experimental Analysis

Evaluation Methodology

Dataset. The primary evaluation uses a diverse set of benchmarks covering five capability areas: Math (GSM8K 8-shot, MATH 4-shot), Code (HumanEval 0-shot, MBPP 3-shot), World Knowledge (Natural Questions 5-shot, TriviaQA 5-shot), Reasoning (ARC-Easy and ARC-Challenge 0-shot, SIQA 0-shot, PIQA 0-shot, WinoGrande 0-shot), and General (MMLU 5-shot). These are standard benchmarks drawn from the Llama-2 and CodeLlama evaluation suites. The paper does not use a single held-out test split from a unified dataset; rather, it evaluates on the standard test sets of each benchmark independently and reports aggregated domain averages.
Base model(s). All experiments start from Llama-2 7B (Touvron et al., 2023), a dense Transformer pretrained on approximately 2 trillion tokens of general web data. This model was chosen as the seed because it is a widely-used, open-weight model with strong general capabilities and published baselines. Expert models are derived from this seed through continued pretraining. For comparison, Llama-2 13B is also evaluated as a reference point for scaling pretraining compute rather than using BTX. CodeLlama 7B (Rozière et al., 2023) and Llemma 7B (Azerbayev et al., 2023) serve as domain-specialized reference models, both also derived from Llama-2 7B.
Metrics. Each benchmark uses its standard metric: for GSM8K and MATH, exact match accuracy; for HumanEval and MBPP, pass@1 with greedy decoding; for Natural Questions and TriviaQA, exact match; for MMLU, ARC-Easy, and ARC-Challenge, accuracy; for WinoGrande, SIQA, and PIQA, accuracy. All generations use greedy decoding. Domain-level scores are computed as the unweighted average of the component benchmarks (e.g., Math = average of GSM8K and MATH). The overall "Average" score reported in Table 2 is the unweighted average across the five aggregated domain scores (Math, Code, Knowledge, Reasoning, MMLU), which itself averages the underlying benchmark scores.
Baselines. The paper compares BTX against six distinct baselines: (1) Llama-2 7B — the seed model without any continued pretraining; (2) Llama-2 13B — a larger model trained from scratch with approximately 2× the pretraining FLOPs of the 7B model (according to Touvron et al., 2023, Table 2); (3) Dense (DM) — continued pretraining of the seed model on the same data mixture used for BTX expert training and MoE finetuning, in the same order (domain-specific data first, then the full mixture), with all parameters updated throughout; (4) Sparse upcycling (DM and CM) — initializing an MoE model from the seed model by making 4 identical copies of each feedforward layer as experts with randomly initialized Top-2 routers, then training on either the same data as BTX (DM, data-matching) or for the same GPU-days as BTX (CM, compute-matching); (5) Branch-Train-Merge (BTM) (Li et al., 2022a) — using the same expert models as BTX but keeping them separate, selecting Top-1 or Top-2 experts per input using tf-idf cosine similarity between the input and each expert's training data representations, then averaging their output token distributions; (6) CodeLlama 7B and Llemma 7B — published specialized models that serve as upper-bound references for code and math domains respectively, both derived from continued pretraining of Llama-2 7B.
Generation budget / compute accounting. Training compute is measured in GPU-days, calculated from the seed model onward (not including the original Llama-2 7B pretraining). For the FLOPs-matched comparison between BTX and sparse upcycling (CM), the metric is total GPU-days spent on continued pretraining: sparse upcycling (CM) uses 1,007.1 GPU-days entirely on MoE training (processing 252B tokens), while BTX uses 926.1 GPU-days split between expert training (processing 453B tokens in parallel) and MoE finetuning (80B tokens), totaling 533B tokens. The slightly lower GPU-days for BTX despite processing 2.1× more tokens reflects the higher throughput of embarrassingly parallel training. Inference efficiency is measured by active parameters: the number of parameters actually computed per token, which varies by routing method (6.7B for Sample Top-1, 11.1B for Top-2, 19.7B for soft routing). All models are evaluated with greedy decoding, so inference cost is deterministic per token given the routing decisions.
Cross-validation / statistical protocol. The paper does not report cross-validation, statistical significance tests, or confidence intervals. Results are reported as single-point estimates on the standard test sets of each benchmark. The MoE finetuning data mixture ratios (30.16% math, 40.31% code, 10.30% Wikipedia, 19.23% original Llama-2 data) were chosen based on the relative volumes of domain training data rather than through systematic hyperparameter sweeps, and the paper acknowledges not experimenting with alternative mixture ratios (Section 6). The compute allocation ratio between expert training and MoE finetuning was also not systematically swept, leaving open the question of whether the reported results represent the optimal allocation.

Main Quantitative Results

Individual Expert Performance: Dramatic Domain Gains at the Cost of Catastrophic Forgetting

Table 1 establishes the raw material that BTX must work with. Each expert model, trained independently on its domain data, shows striking improvements in its target domain and substantial degradation elsewhere:

Math expert: GSM8K improves from 14.7% to 39.5% (+24.8 percentage points) and MATH from 2.5% to 18.8% (+16.3 points). However, TriviaQA collapses from 58.5% to 37.1% (-21.4 points) and MMLU drops from 46.1% to 37.1% (-9.0 points). Notably, code performance improves from 16.8% to 33.6% on the code aggregate, suggesting genuine transfer between math and code domains.
Code expert: HumanEval improves from 12.8% to 31.7% (+18.9 points) and MBPP from 20.8% to 40.2% (+19.4 points). But TriviaQA drops from 58.5% to 29.9% (-28.6 points), MMLU drops from 46.1% to 39.6% (-6.5 points), and Natural Questions drops from 16.4% to 11.5% (-4.9 points). Math performance improves slightly (8.6% to 8.0% on GSM8K + MATH aggregate? Wait, Table 1 shows GSM8K 12.0%, MATH 4.0%, math average = 8.0% vs. seed 8.6%—actually a slight decrease on the math aggregate), indicating less cross-domain transfer from code to math than vice versa.
Wikipedia expert: Natural Questions improves from 16.4% to 21.8% (+5.4 points), and TriviaQA remains stable (58.5% → 57.2%, -1.3 points). Math and code performance degrade slightly (math: 8.6% → 7.4%; code: 16.8% → 13.1%). MMLU drops modestly (46.1% → 43.1%, -3.0 points).

The key observation is that catastrophic forgetting is severe for the math and code experts on knowledge tasks, while the Wikipedia expert shows milder forgetting—likely because Wikipedia training (42B tokens) is much shorter than math (201B) or code (210B) training, giving less opportunity for the model to overwrite previously learned knowledge. This asymmetry in forgetting severity is not discussed in the paper but is visible in the data and may influence the MoE finetuning dynamics.

BTX vs. Baselines: BTX Achieves the Best Average Performance with Balanced Domain Improvements

Table 2 presents the headline comparison. BTX with Top-2 routing achieves an average score of 47.9, compared to 40.7 for the seed model Llama-2 7B (+7.2 points) and 45.4 for Llama-2 13B (+2.5 points). The domain-level breakdown shows where the improvements come from:

Math: BTX Top-2 achieves 27.4 (average of GSM8K 37.1% and MATH 17.8% from Table 8), up from the seed model's 8.6. This approaches Llemma 7B's 28.0 but remains below sparse upcycling (CM)'s 28.2 (Table 3). Compared to Dense (DM)'s 18.3, BTX improves math by 9.1 points—the largest single-domain gain over the dense baseline.
Code: BTX Top-2 achieves 34.0 (HumanEval 28.7%, MBPP 39.4% from Table 8), up from the seed model's 16.8. This approaches CodeLlama 7B's 36.3 and substantially exceeds Dense (DM)'s 25.8 (+8.2 points). Sparse upcycling (CM) achieves only 30.7 (Table 3), suggesting BTX's expert initialization provides a meaningful code advantage.
World Knowledge: BTX Top-2 achieves 41.0 (Natural Questions 24.8%, TriviaQA 57.1% from Table 8), up from the seed model's 37.4. This is a domain where catastrophic forgetting was most severe in the math and code experts (Table 1: TriviaQA dropped to 37.1% and 29.9% respectively). BTX recovers and exceeds the seed model's knowledge performance (+3.6 points), while Dense (DM) achieves 39.6 (+2.2 points) and sparse upcycling (DM) achieves only 34.0. The large gap between BTX and sparse upcycling on knowledge (41.0 vs. 34.0) suggests the generalist expert (frozen seed model) provides critical knowledge retention that sparse upcycling cannot recover through joint training alone.
Reasoning: BTX Top-2 achieves 63.5, essentially flat compared to the seed model's 63.3 and Dense (DM)'s 63.3. This domain shows almost no variation across methods (range: 62.3–63.7 for all BTX and dense variants), suggesting that commonsense reasoning capabilities are already near the Llama-2 7B ceiling and are not significantly affected by continued pretraining or MoE integration.
MMLU: BTX Top-2 achieves 52.5, up from the seed model's 46.1 (+6.4 points). This is notable because individual experts varied dramatically on MMLU: the math expert outperformed the seed (52.0 vs. 46.1) while the code and Wikipedia experts degraded (39.6 and 43.1 respectively). BTX captures the math expert's MMLU improvement while recovering from the other experts' degradation.

The comparison with BTM is particularly informative about the value of the MiX finetuning stage. BTM Top-2 uses the exact same expert models as BTX but achieves only 43.4 average (vs. 47.9 for BTX Top-2), a gap of 4.5 points. The gap is largest in world knowledge (26.9 vs. 41.0, a 14.1-point difference) and math (21.5 vs. 27.4, a 5.9-point difference). This demonstrates that the MoE finetuning stage is not merely reorganizing already-available expertise—it is creating genuinely better performance through joint optimization of the router and shared parameters. The domain-classifier heuristic of BTM cannot replicate what learned token-level routing achieves.

Compute Efficiency: BTX Outperforms Sparse Upcycling with Higher Throughput

Table 3 and Figure 2 (left) address the core efficiency argument. Sparse upcycling with compute-matching (CM)—that is, allocating 100% of compute to MoE training without any embarrassingly parallel expert phase—achieves an average score of 47.3 using 1,007.1 GPU-days and 252B training tokens. BTX achieves 47.9 using 926.1 GPU-days and 533B total tokens. The efficiency advantage has two components:

Higher throughput in the parallel phase: BTX's expert training processed 453B tokens in parallel across three experts, benefiting from zero cross-expert communication and independent data pipelines. The 80B tokens of MoE finetuning were the only synchronized phase. Sparse upcycling's entire 252B tokens were trained with synchronized MoE communication. The result: BTX processes 2.1× more tokens in slightly less total time (7.8 vs. 7.9 days).
Better data efficiency during MoE finetuning: The BTX MoE model starts with domain-specialized experts, so the router only needs to learn which expert to use for each token—a relatively easy learning problem. Sparse upcycling starts with identical experts, so the model must simultaneously learn to specialize the FF layers and learn routing—a harder joint optimization problem. This manifests in BTX's slightly higher average score (47.9 vs. 47.3) despite using less total compute.

Figure 2 (left) visualizes this as a Pareto frontier. BTX (circle at 926 GPU-days, score 47.9) lies above and to the left of sparse upcycling (CM) (1,007 GPU-days, score 47.3) and substantially above Dense (DM) and BTM. Llama-2 13B, plotted for reference with approximate GPU-days calculated by doubling the 7B training compute, achieves only 45.4 with larger active parameters (13B vs. BTX's 11.1B for Top-2). The paper explicitly notes that "BTX uses less than half of the additional training compute compared to Llama-2 13B" (the additional compute being the difference between 7B and 13B pretraining costs), yet surpasses it on overall performance.

Figure 2 (right) normalizes the scores by dividing each method's domain performance by the best performance achieved by any method in that domain. BTX achieves normalized scores of approximately 0.95+ in math (matching Llemma) and code (approaching CodeLlama), while maintaining ~1.0 in knowledge and reasoning. This radar-chart visualization emphasizes BTX's balanced profile compared to specialized models, which spike in their domain but collapse in others.

Routing Method Comparison: Top-2 Balances Performance and Efficiency

Table 4 presents a systematic comparison of routing methods at different finetuning budgets (10B, 40B, 80/160B tokens). This ablation uses a subset of representative tasks (GSM8K, HumanEval, Natural Questions, ARC Challenge, MMLU) rather than the full evaluation suite, so the "Average Score" is not directly comparable to Table 2's overall average. The key patterns:

Switch Top-1 is substantially worse than all alternatives. At 10B tokens, Switch achieves 24.7 average vs. 34.6 for Top-2—a 9.9-point gap on this task subset. The paper attributes this to the difficulty of training with capacity-factor-based token dropping and the lack of redundancy from single-expert selection.
Sample Top-1 is competitive with Top-2 at higher finetuning budgets. At 10B tokens, Sample Top-1 achieves 33.0 vs. Top-2's 34.6. At 40B tokens, the gap narrows to 35.3 vs. 35.9. At 160B tokens, Sample Top-1 reaches 36.9, while Top-2 at 80B tokens reaches 37.3. Sample Top-1 requires only 6.7B active parameters vs. 11.1B for Top-2, making it more inference-efficient. The paper presents these as comparable: "Sample Top-1 achieves competitive performance with less inference compute."
Soft routing consistently outperforms sparse methods but at 3× the inference cost. At 10B tokens, soft routing achieves 35.8; at 40B, 37.3; both are the highest in their respective columns. However, soft routing activates all 4 experts (19.7B active parameters), making it 2.9× more expensive at inference than Sample Top-1 and 1.8× more expensive than Top-2. The paper treats soft routing as an upper bound on what MoE routing can achieve rather than a practical deployment option.
All methods improve with more finetuning tokens. Sample Top-1 improves from 33.0 (10B tokens) to 35.3 (40B) to 36.9 (160B). Top-2 improves from 34.6 (10B) to 35.9 (40B) to 37.3 (80B). The diminishing returns in later stages suggest that most of the routing learning happens in the first 40B tokens of MoE finetuning.

The finding that Sample Top-1 can nearly match Top-2 with significantly fewer active parameters is important for deployment: it suggests that the BTX approach is not inherently tied to Top-2 routing and can be adapted for more inference-constrained settings.

BTX vs. Specialized Models: Closing the Gap Without Sacrificing Generality

Table 2 and the per-task Table 8 allow comparison of BTX against the specialized reference models CodeLlama 7B and Llemma 7B:

Math: BTX Top-2 (27.4) is within 0.6 points of Llemma 7B (28.0), while drastically outperforming it on world knowledge (41.0 vs. 17.2, a 23.8-point gap) and reasoning (63.5 vs. 38.8). Llemma's catastrophic forgetting is most severe on reasoning tasks: ARC-Challenge drops from 43.8 (seed) to 26.8, and ARC-Easy drops from 76.4 to 28.7—the model essentially loses its commonsense reasoning capabilities. BTX recovers these to near-seed levels.
Code: BTX Top-2 (34.0) is within 2.3 points of CodeLlama 7B (36.3), while outperforming it on world knowledge (41.0 vs. 22.2, an 18.8-point gap), reasoning (63.5 vs. 56.6), and MMLU (52.5 vs. 38.6). The pattern mirrors math: specialized models achieve strong in-domain performance but collapse on out-of-domain tasks, while BTX preserves in-domain gains while recovering general capabilities.

These comparisons validate the paper's central claim that BTX produces a model that is competitive with dedicated specialists in their domains while maintaining the broad competence of a generalist. The gaps that remain (2-3 points in code, 0.6 points in math) represent the cost of unification—the MoE finetuning and load balancing slightly dilute domain performance relative to pure specialization—but the trade-off is massively favorable when considering performance across all domains.

Ablation Studies and Robustness Checks

Load balancing: Adding load balancing (α = 0.01) significantly improves code performance (HumanEval: 27.4 with LB vs. 19.5 without, Table 5) but slightly reduces math performance (GSM8K: 29.8 with LB vs. 34.6 without). The average score across the five representative tasks is nearly identical (34.7 vs. 34.6), but the distribution shifts dramatically. This is explained by the routing analysis in Figure 3: without load balancing, the Math expert absorbs ~60-80% of routing probability across all domains, starving the Code expert. With load balancing, the Code expert activates on ~40-50% of code tokens and contributes meaningfully across other domains. The load balancing loss forces the router to utilize all experts, preventing the collapse that would otherwise occur when one expert (Math, with its large and domain-overlapping training data) dominates routing.

Freezing expert FF weights during MoE finetuning: "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" (Table 5: no LB & freeze experts achieves 34.7 average vs. 34.6 for no LB with expert training). This is a significant finding: it means the Branch-Train phase produces FF weights that are already near-optimal for their domains, and the MiX phase primarily serves to train the router (which must learn to select among these frozen experts) and to optimize the averaged attention weights and other shared parameters. This finding is robust across tasks: GSM8K 34.8 vs. 34.6, HumanEval 18.3 vs. 19.5, Natural Questions 24.1 vs. 23.2, MMLU 51.4 vs. 51.6. The near-identical performance validates the paper's two-phase compute allocation—the expensive parallel phase is where domain knowledge is acquired, and the synchronized phase is primarily an integration step.

Blending experts across MoE modules: When domain FF layers are split into chunks and recombined such that each MoE module contains equal parameter shares from all domains, performance collapses catastrophically: average score drops from 34.7 to 22.2 (Table 5). Every task degrades—GSM8K: 34.6 → 13.9, HumanEval: 19.5 → 17.1, Natural Questions: 23.2 → 9.9. This negative result demonstrates that coherent domain specialization encoded in each expert's FF parameters is essential—the router cannot reconstruct domain knowledge from blended fragments. It also validates the paper's architectural choice to keep experts aligned with domains rather than allowing domain mixing at the parameter level.

Splitting experts into more granular modules: Using C = 2 chunks per domain expert to create 8 MoE modules instead of 4, then using Top-2 of 8 routing, reduces average performance from 34.7 to 28.0 (Table 5). Using Top-4 of 8 (matching the active parameter count of Top-2 of 4) recovers to 34.5, essentially equal to the 4-expert configuration. This suggests that splitting experts neither helps nor hurts when the total active parameter count is matched, but provides no benefit that would justify the added complexity. The router gains additional flexibility (it can select different chunks from different domains), but this flexibility does not translate into better performance, likely because the domain knowledge within each expert is not usefully decomposable into independently routable chunks.

Routing method comparison across finetuning budgets (Table 4): All routing methods improve with additional finetuning tokens, but with diminishing returns. Switch routing shows the steepest initial climb (24.7 at 10B tokens) but the paper does not report higher-budget Switch results, likely because it was abandoned as inferior. Sample Top-1 shows the most consistent improvement with budget (33.0 → 35.3 → 36.9), eventually nearly matching Top-2 (37.3 at 80B) with half the inference parameters. This suggests that for deployment scenarios where inference cost is the primary constraint, Sample Top-1 with extended finetuning is a viable alternative to Top-2.

First-layer soft routing for Sample Top-1 (Section 4.1.1): The paper mentions that for Sample Top-1 routing, "for the first layer only, we used soft-routing instead." This is a targeted design choice to address the cold-start problem: token embeddings at the first layer contain minimal context, making hard Top-1 routing decisions unreliable. Soft routing at layer 1 allows all experts to contribute to the initial representation, providing a richer signal for subsequent layers' routing decisions. The paper does not ablate this choice, so its quantitative importance is unknown, but the mention suggests it was empirically necessary for Sample Top-1 training stability.

Data mixture ratios during MoE finetuning: The paper uses fixed ratios (30.16% math, 40.31% code, 10.30% Wikipedia, 19.23% Llama-2 data) and acknowledges in Section 6 that it "did not perform experiments with different data mixtures for MoE finetuning other than uniform sampling." This is a significant gap: the mixture ratios almost certainly affect the balance of domain performance in the final model. A heavier weight on math data during MoE finetuning might recover the GSM8K gap with the no-load-balancing variant, while a heavier weight on code data might close the remaining gap with CodeLlama. The chosen ratios reflect the relative volumes of expert training data, but this is an implicit choice rather than an optimized one.

Number of domains and experts: All experiments use exactly four experts (math, code, Wikipedia, generalist). The paper does not ablate the number of domains, the choice of which domains to include, or the effect of including multiple generalist variants. Section 6 acknowledges this: "Training on more domains such as using unsupervised domain discovery (Gururangan et al., 2023) should amplify the benefit of the parallelization of experts training." This remains a forward-looking claim rather than an empirically validated one.

Top-k values for BTM: Table 2 reports BTM with both Top-1 and Top-2 expert selection. BTM Top-2 (43.4) slightly outperforms BTM Top-1 (43.1), but both are substantially below BTX. The small difference between Top-1 and Top-2 for BTM (0.3 points) compared to the large difference for BTX routing variants suggests that the BTM domain-classifier heuristic has low precision—adding a second expert provides minimal benefit because the classifier's second choice is often not the most relevant domain.

Sample Top-1 temperature annealing: The paper uses an exponential annealing schedule τ = max(0.5, exp(-r t)) with r = 1e-4. The minimum temperature of 0.5 (rather than 0.1 or 0.01) means the Gumbel-Softmax training distribution never becomes fully deterministic, maintaining some stochasticity throughout training. This likely helps prevent premature convergence to suboptimal routing assignments but is not ablated against alternative schedules.

Critical Assessment

Does BTX genuinely outperform the baselines, or are the comparisons confounded?

The comparison between BTX and Dense (DM) is the cleanest test of whether the BTX architecture adds value beyond simply training on more data. Both methods use exactly the same total data in the same order (domain-specific data first, then the full mixture). BTX achieves 47.9 vs. Dense's 44.5, a 3.4-point average improvement. This gap is driven primarily by math (27.4 vs. 18.3) and code (34.0 vs. 25.8), with smaller contributions from world knowledge (41.0 vs. 39.6) and MMLU (52.5 vs. 49.8). The dense model's underperformance suggests that joint training on mixed-domain data causes negative interference between domains—a known failure mode of multi-task learning that BTX's parameter isolation (separate FF layers per domain) successfully mitigates.

However, the Dense baseline has a confound: it processes all data sequentially with synchronized updates, while BTX trains domain experts in parallel. The Dense model sees each domain's data interleaved with others during training, which might produce different optimization dynamics independent of the architectural differences. A fairer comparison would train the dense model on each domain's data sequentially (as in continual learning) to match the exposure pattern of BTX experts, but the paper's "Dense (DM)" baseline mixes data during the domain-specific phase as well. The paper does not specify the exact interleaving, so it's unclear whether the dense model's underperformance is due to architectural limitations or training order effects.

The comparison with sparse upcycling (CM) is the paper's strongest efficiency claim. BTX achieves 47.9 in 926 GPU-days vs. sparse upcycling's 47.3 in 1,007 GPU-days. The 0.6-point difference is small, and without statistical significance testing, it's unclear whether this difference is reliable or within noise. The more defensible claim is that BTX is competitive with sparse upcycling while offering higher training throughput—the 533B vs. 252B token processing advantage is substantial and unambiguous. However, this throughput advantage depends on the assumption that expert training can proceed in parallel with perfect linear scaling, which the paper's small cluster of three experts cannot fully validate. At larger scales with more domains, Amdahl's law would apply: synchronization points, data loading, and hardware allocation overhead would eat into the theoretical speedup.

Does BTX genuinely "bridge the gap with specialized models"?

This claim requires careful qualification. BTX Top-2 achieves a math score of 27.4 vs. Llemma 7B's 28.0—a 0.6-point gap. In code, BTX achieves 34.0 vs. CodeLlama 7B's 36.3—a 2.3-point gap. These are small differences, and BTX's broader performance is dramatically better (e.g., +23.8 points in world knowledge over Llemma). So the claim is well-supported in the sense that BTX is competitive with specialists while being vastly more general.

But the specialists used for comparison (Llemma 7B, CodeLlama 7B) are themselves derived from Llama-2 7B with continued pretraining—they are not the strongest possible specialists. Would BTX close the gap with a specialist trained on more data, or with a larger specialist model? The paper cannot answer this because it only tests one specialist scale per domain. Furthermore, the specialists were trained on exactly the same data as the BTX experts, making the comparison fair but also narrow. A stronger test would compare BTX against a specialist trained with the full compute budget that BTX uses (expert training + MoE finetuning) allocated entirely to that domain—this would test whether the MoE integration imposes a fundamental ceiling on domain performance.

The load balancing dependence is a genuine fragility.

The paper's own data shows that without load balancing, code performance collapses (HumanEval 19.5 vs. 27.4 with LB, Table 5). The load balancing coefficient α = 0.01 is not systematically tuned—the paper does not show ablations across α values or demonstrate that 0.01 is near-optimal. This matters because load balancing is a known hack in MoE training that trades off between expert utilization and routing optimality. Too much load balancing forces the router to use suboptimal experts for some tokens; too little causes expert collapse. The paper's chosen α was likely found to work through trial and error, but without documentation of the tuning process or sensitivity analysis, a practitioner attempting to replicate BTX on a different set of domains could easily choose a suboptimal α and get substantially worse results.

The deeper issue is that load balancing is needed precisely because the experts are not equally useful across all domains—the Math expert trained on 201B tokens with broad domain overlap is naturally more generally useful than the Wikipedia expert trained on 42B tokens. Load balancing papers over this asymmetry rather than addressing it. A more principled approach might equalize the "general usefulness" of experts through more balanced training data or through curriculum learning during expert training, but the paper does not explore this.

The routing analysis is descriptive, not causal.

Section 4.3.2 and Figure 3 show that routing decisions correlate with domain-appropriate experts, but do not demonstrate that these routing decisions cause the performance improvements. The routing analysis shows the Code expert is dominant in code tasks after load balancing—but is this because the Code expert actually provides better representations for code tokens, or because the router learned to associate code-like token patterns with the Code expert index without those representations being genuinely useful? The paper cannot distinguish these possibilities because it does not perform causal interventions (e.g., forcing the router to use non-domain experts for domain tokens and measuring the performance impact).

The token-level routing examples in Table 6 are anecdotal and selected post-hoc. They demonstrate that the model can route domain-appropriate tokens to the corresponding experts, but not how consistently it does so across all tokens in all tasks. A quantitative measure of domain-routing alignment (e.g., what fraction of code dataset tokens are routed to the Code expert vs. others) would strengthen the claim, but the paper only provides aggregate per-layer distributions (Figures 3, 4, 5, 6).

Missing experiments that would strengthen the paper.

Several experiments are conspicuously absent given the paper's claims:

Scaling the number of domains. The paper argues that BTX's parallel training advantage grows with more domains, but tests only three. An experiment with 5, 10, or 20 domains (perhaps using unsupervised domain discovery as suggested in Section 6) would validate the scaling claim.
SFT/RLHF on the BTX model. The paper argues BTX's key advantage over BTM is that it produces a unified model that can be fine-tuned, but never demonstrates this capability. Running instruction tuning or RLHF on the BTX model and comparing with BTM would close this gap.
Systematic sweep of compute allocation ratio. The paper frames BTX as generalizing the two extremes of BTM (100% expert training) and sparse upcycling (0% expert training), but tests only one intermediate ratio (~77% expert training). Sweeping this ratio would reveal the shape of the trade-off and whether the paper's chosen allocation is near-optimal.
Data mixture ratio sensitivity during MoE finetuning. The 30/40/10/19 split across domains is likely suboptimal, and the paper acknowledges this, but provides no guidance on how mixture ratios affect final performance. An experiment varying the ratio would help practitioners adapt BTX to their own domain distributions.
Larger seed models. All experiments use Llama-2 7B. Would BTX provide similar relative gains starting from a 13B or 34B seed model? The paper's efficiency argument—that BTX is more effective for late-stage pretraining—would be strengthened by showing it scales with base model size.
Statistical significance. With 500-question MATH and 164-question HumanEval test sets, differences of 1-2 points between methods may not be statistically significant. The paper reports no confidence intervals, no bootstrap estimates, and no significance tests, making it difficult to assess the reliability of the comparative claims.

Where do the claims hold conditionally?

"BTX outperforms BTM on all tasks" — This holds unconditionally in the reported results (Table 2), but the gap varies dramatically by domain: 14.1 points in world knowledge, 5.9 points in math, 1.5 points in MMLU, and essentially zero in reasoning. BTM's primary weakness is in domains where catastrophic forgetting was most severe (knowledge), suggesting that the MoE finetuning stage is most valuable for recovering lost general capabilities rather than improving domain-specialized performance.
"BTX is more compute-efficient than sparse upcycling" — This holds in terms of throughput (more tokens per GPU-day), but the final model quality difference (47.9 vs. 47.3) is small and unquantified for reliability. The claim is better stated as "BTX achieves comparable or slightly better performance with higher training throughput and more total data processed."
"BTX bridges the gap with specialized models" — This holds for Llemma 7B (math: 27.4 vs. 28.0) but the gap is larger for CodeLlama 7B (code: 34.0 vs. 36.3). It holds only for specialists of comparable scale; the paper does not compare against larger specialists (e.g., CodeLlama 13B, Llemma 34B) or against specialists trained for longer.
"Freezing experts during MoE finetuning has little impact" — This holds for the specific set of five tasks in Table 5 at 10B tokens of finetuning. Whether it holds for all tasks, at larger finetuning budgets, or with different expert training recipes is unknown. The paper does not report the full evaluation suite for this ablation.

6. Limitations and Trade-offs

The Difficulty Estimation Prohibitive Cost Is Unaddressed

The assumption or constraint. BTX relies on a two-phase allocation of compute between embarrassingly parallel expert training and synchronized MoE finetuning, but the paper makes no attempt to determine the optimal allocation ratio. Section 6 acknowledges this explicitly:

"We only compared BTX to two of its special variants, i.e. BTM with 100% compute allocated to expert training and 0% on MoE finetuning, and sparse upcycling with 0% compute allocated to expert training and 100% on MoE finetuning. Future work could perform a thorough sweep of the compute allocation ratio between expert training and MoE training."

The specific allocation used in the experiments (~77% expert training, ~23% MoE finetuning) was determined by matching prior work's training volumes for individual experts (Llemma's 201B tokens for math, CodeLlama's 210B tokens for code) and choosing a MoE finetuning budget (80B tokens) that felt sufficient for router learning. It was not optimized.

The consequence. A practitioner cannot determine, without re-running the experiments, whether the reported performance represents a local optimum along the allocation-ratio axis or whether substantially better results could be achieved with different ratios. The paper's central framing—that BTX generalizes the two extremes and that an intermediate point outperforms both—establishes that the allocation ratio matters, but provides no guidance on how to choose it for a new set of domains, data volumes, or model scales. If an organization has 10 domains with varying data volumes rather than 3 domains with roughly balanced data, the optimal ratio almost certainly shifts, and the paper provides no transfer function to estimate it.

What evidence exists in the paper. The only evidence is the single-point comparison: BTX (~77/23 split) outperforms both BTM (100/0) and sparse upcycling (0/100). Figure 2 (left) and Table 3 provide the empirical comparison, and the special-case framing in Section 6 acknowledges this as a continuous dimension worth sweeping, but no ablation varies the ratio even coarsely (e.g., 50/50 or 90/10). The finding that freezing expert FF weights during MoE finetuning has "little impact on performance across all tasks" (Table 5) hints that the allocation could be shifted even further toward expert training without much loss, but this is not tested beyond the no-load-balancing setting at 10B finetuning tokens.

Mitigation status. Not addressed. The paper flags this as future work but provides no characterization of the trade-off surface, no theoretical model for predicting the optimal ratio, and no heuristics for practitioners.

Load Balancing Is a Brittle Necessity, Not a Robust Design Principle

The assumption or constraint. BTX's strong multi-domain performance depends critically on load balancing during MoE finetuning, with a specific coefficient α = 0.01 chosen without systematic tuning. Without load balancing, the MoE router collapses to predominantly using the Math expert across nearly all domains and layers, with the Code expert becoming effectively dead (routing probability near zero in later layers). The paper's routing analysis (Section 4.3.2, Figure 3) documents this vividly but treats it as an interesting phenomenon rather than a structural fragility.

The consequence. Load balancing is a known hack in MoE training—an auxiliary loss that penalizes uneven expert utilization rather than a principled mechanism for learning optimal routing. The trade-off it introduces is measurable in the paper's own results (Table 5): with load balancing, HumanEval is 27.4 but GSM8K is 29.8; without load balancing, HumanEval collapses to 19.5 while GSM8K improves to 34.6. The routing analysis (Figure 3) explains why: load balancing forces the router to use the Code expert for code tokens (improving code performance) but also forces it to use non-Math experts for some math tokens (degrading math performance). The average score across these tasks is nearly identical (34.7 vs. 34.6), but the distribution shifts dramatically.

This means the α = 0.01 value is effectively controlling a domain-performance trade-off knob without the paper characterizing it as such. A practitioner deploying BTX who cares more about code than math would want a lower α (or no load balancing); one who cares primarily about math would want the opposite. The paper provides no sensitivity analysis and no principled way to set α based on domain priorities—the chosen value was found empirically and locked in for all experiments. For a new set of domains with different degrees of expert quality asymmetry, this brittle dependence on α could produce substantially different results.

What evidence exists in the paper. Table 5 provides the direct ablation comparing with and without load balancing at 10B finetuning tokens. Figure 3 provides the qualitative routing distribution difference, showing the Code expert "coming back to life" with load balancing. The paper also notes the dependence in Section 4.3.2: "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." Figure 5 in the appendix further documents the dead-expert phenomenon with per-layer routing probability histograms.

Mitigation status. The paper does not treat this as a limitation requiring mitigation. Load balancing is presented as a standard MoE training technique, and the chosen α is reported without ablations across different values. Section 6 does not mention load balancing as an open issue. The failure mode (dead experts without load balancing) is treated as an empirical observation that justifies using load balancing, rather than as evidence that the BTX integration approach requires a fragile balancing act that practitioners will need to tune per setup.

No Demonstration of the Claimed Unified Fine-Tuning Advantage

The assumption or constraint. The paper's primary justification for BTX over BTM is that BTX produces a single unified model that can undergo SFT and RLHF, while BTM's collection of separate experts cannot. Section 1 states this as a key motivation:

"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, both of which can boost performance further, and are crucial steps in building aligned LLMs."

And Section 6 restates it as a direction for future work:

"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 consequence. The paper never tests whether a BTX model actually benefits from SFT or RLHF, or whether such post-training would preserve the multi-domain balance that the MoE finetuning achieved. There are several plausible failure modes that remain unexamined:

SFT on instruction-following data might cause the router to shift its allocation patterns, potentially re-collapsing experts in ways that undo the load balancing benefits.
RLHF reward models are typically trained on general preference data; a BTX model's domain-specialized routing might interact with RLHF optimization in ways that degrade domain-specific performance in favor of general conversational ability.
If SFT or RLHF disproportionately updates certain experts (e.g., if instruction-following data routes primarily to the generalist expert), the model could lose specialized capabilities even if the expert FF weights themselves are frozen.

The paper's central value proposition—that BTX enables a complete post-training pipeline—is asserted but not demonstrated. A practitioner choosing between BTX and BTM cannot know from this paper whether the downstream SFT/RLHF advantage is real or hypothetical.

What evidence exists in the paper. Zero. No SFT or RLHF experiments are reported. The paper evaluates only base-model (pretrained) performance on standard benchmarks. The claim about SFT/RLHF compatibility is purely architectural and motivational. The finding that freezing expert FF weights during MoE finetuning has little performance impact (Table 5) provides circumstantial evidence that the experts might remain stable during further training, but this was tested only in the context of continued language modeling, not instruction tuning or RL optimization.

Mitigation status. Explicitly deferred to future work. The paper is transparent about this gap in Section 6, but the limitation is consequential because the SFT/RLHF pipeline is, by the paper's own argument, the primary reason to prefer BTX over BTM. Without this evidence, the comparison between BTX and BTM in Table 2 (47.9 vs. 43.4) must be interpreted as a comparison of base models only, and the downstream advantage after full alignment remains unknown.

Single Seed Model, Single Scale: Generalization Across Architectures and Sizes Is Untested

The assumption or constraint. All experiments use Llama-2 7B as the seed model. The domain experts, the MoE architecture, the routing strategies, and the baselines are all derived from this single starting point. The paper's claims about BTX being "more compute efficient than training a larger generalist LLM or several separately specialized LLMs" (Section 5) are based entirely on the 7B starting scale. Section 4.1 justifies the choice:

"We base our experiments on the setup used for Llama-2 pretraining... we use the Llama-2 7B model as our seed model."

Section 6 acknowledges the scale limitation indirectly:

"Due to compute limitations, we only experimented with three domains and four experts in this paper."

The consequence. Several dimensions of generalization are untested, and there are plausible reasons why results might differ:

Seed model scale: At 7B parameters, Llama-2 7B has limited general knowledge and reasoning capacity. The relative benefit of domain-specialized continued pretraining might be larger at this scale (where the seed model has clear gaps that specialists can fill) than at 34B or 70B (where the seed model already has stronger domain performance, leaving less room for improvement). Alternatively, larger seed models might benefit more from BTX because they have stronger representations that make expert specialization more sample-efficient. The paper provides no data to distinguish these possibilities.
Seed model family: Llama-2 is a specific architecture (dense Transformer with SwiGLU activations, rotary position embeddings, specific pretraining data distribution). The finding that FF layers are more domain-specialized than attention layers (the justification for keeping FF layers separate while averaging attention) might be specific to Llama-2's training recipe. Models with different architecture choices (e.g., different FF-to-attention parameter ratios, grouped-query attention, different activation functions) might show different specialization patterns that break BTX's design assumptions.
Domain composition: The three domains tested (math, code, Wikipedia) have clear semantic boundaries and distinct data distributions. BTX's routing benefits likely depend on domains being sufficiently distinct that a router can learn to differentiate them from token content. For domains with overlapping vocabulary and structure (e.g., different programming languages within code, or different scientific disciplines), the domain-level routing might be less effective, and the benefit of BTX over dense continued pretraining might shrink.

What evidence exists in the paper. None beyond the single configuration. The paper does not report experiments with Llama-2 13B as a seed model, with non-Llama models, or with alternative domain compositions. The qualitative routing analysis (Figures 3 and 6) provides some evidence that domains are distinguishable by token content even within the Llama-2 representation space (GSM8K vs. MATH tokens route differently), but this is within-family evidence that does not test generalization.

Mitigation status. Not addressed. The paper makes no claim that BTX generalizes across seed model scales or architectures, but also provides no caveats. The practical consequence is that an organization considering BTX for a different base model has no transfer guidance.

Three Domains Are Insufficient to Validate the Scaling Argument

The assumption or constraint. The paper's strongest efficiency argument is that embarrassingly parallel expert training has higher throughput than synchronized MoE training, and that this advantage grows with the number of domains. Section 3.1 states:

"It allows linear scaling of overall training throughput when scaling up the size of compute."

And Section 6 projects forward:

"Training on more domains such as using unsupervised domain discovery (Gururangan et al., 2023) 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."

The consequence. All experiments use exactly three domains plus the generalist (four experts total). At this small scale, the throughput advantage is measurable but modest: BTX processes 2.1× more tokens than sparse upcycling in slightly less wall-clock time (Table 3). However, the paper's scaling argument projects that this advantage should grow linearly with the number of domains. Several failure modes could prevent this:

Hardware allocation overhead: Training 10 domain experts in parallel requires 10× the GPU allocation of training one expert, even though each expert training is independent. Cluster management overhead, data pipeline contention, and I/O bottlenecks could eat into the theoretical throughput gains at larger scales.
Load balancing with many experts: The router's job becomes harder as the number of experts grows. With 4 experts, each token routes to 2 (Top-2), giving the router a relatively constrained allocation problem. With 20 experts and still Top-2 routing, the router must learn to select among 20 options for each token, and 18 experts go unused per token. The load balancing constraint (equal utilization across 20 experts) becomes more restrictive, potentially forcing the router to route tokens to suboptimal experts more frequently.
Diminishing returns from additional domains: Not all domains benefit equally from specialized experts. The three domains tested (math, code, Wikipedia) are high-value domains with large, high-quality training datasets and clear performance improvements from specialization. If a practitioner adds domains with smaller datasets or less clear specialization gains (e.g., "academic papers" as separate from "Wikipedia"), the throughput advantage of parallel training might be offset by reduced per-expert quality gains.
MoE finetuning data mixture complexity: With three domains, the MoE finetuning data mixture has three tunable weights plus the generalist data weight. With 10 domains, this becomes an 11-dimensional optimization problem that the paper provides no method to solve.

What evidence exists in the paper. The three-domain experiment provides a single data point. Table 3 establishes the throughput advantage at this scale, and Table 2 establishes the performance advantage over baselines. But the scaling projection—that more domains amplify the benefit—is purely theoretical, with no empirical evidence. The expert splitting experiment (Table 5), which tested 8 modules instead of 4 with Top-2 and Top-4 routing, provides weak indirect evidence: increasing the number of modules without adding genuinely new domain knowledge did not improve performance, suggesting that simply adding more experts is not automatically beneficial.

Mitigation status. Acknowledged as a limitation in Section 6 with a forward-looking suggestion to use unsupervised domain discovery. The paper does not claim to have solved the scaling problem, but its core value proposition—that BTX is attractive because it scales linearly with domains—is undersupported by the three-domain experiment.

No Accounting for Inference Latency or Expert Parallelism in Deployment

The assumption or constraint. The paper evaluates inference efficiency solely in terms of active parameter count (6.7B for Sample Top-1, 11.1B for Top-2) and treats total FLOPs as the primary cost metric. Section 3.2 notes the computational savings:

"Since we can skip computing FF_i^l(x) if the corresponding router output is zero, the actual computation of FF_MoE^l(x) will be much more efficient than computing all domain experts."

And Section 6 projects efficiency gains from more experts:

"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."

The consequence. MoE models present a fundamental latency-vs-throughput tradeoff that the paper does not address. Even though only k experts are active per token, the inactive experts' parameters still occupy GPU memory. In a typical deployment, all expert parameters must be loaded onto the same device or communicated across devices. This has several implications:

Memory footprint: A BTX model with four 7B-parameter experts has a total parameter count of ~28B even though only 11.1B are active per token (for Top-2). This means the model requires memory proportional to the total expert count, not just the active subset. Deploying BTX Top-2 requires roughly 4× the GPU memory of the original Llama-2 7B, even though inference FLOPs are only ~1.6× higher.
Expert parallelism vs. latency: The paper mentions "placing different experts on different GPUs to run them in parallel" (Section 6) as a future optimization. Without this, experts on a single GPU are computed sequentially, adding latency proportional to k (the active expert count). With expert parallelism across GPUs, the all-to-all communication between GPUs for routing tokens to their assigned experts adds network latency. Neither regime matches the latency characteristics of a dense model with equivalent active parameters—the dense model has lower communication overhead and can be optimally sharded across GPUs without the dynamic routing bottleneck.
Batch size sensitivity: MoE models are typically more efficient at large batch sizes, where the token distribution across experts can be balanced. At small batch sizes (common in interactive inference with low query rates), some experts may be underutilized, wasting GPU memory and potentially causing load imbalance that the training-time load balancing doesn't address at inference time.

What evidence exists in the paper. None. All inference cost discussions are in terms of active parameter counts. The paper does not report inference latency, memory usage, throughput at different batch sizes, or comparison with a dense model of equivalent active parameters on any inference metric. The compute-matching comparisons in Table 3 and Figure 2 are entirely about training compute, not inference characteristics.

Mitigation status. Section 6 acknowledges that "an efficient MoE implementation could shorten the training time of BTX," referring to expert parallelism during training, but does not address inference latency. The paper focuses on training efficiency and leaves inference deployment characteristics entirely to future work, making it difficult for a practitioner to assess whether BTX's parameter efficiency translates to latency-competitive deployments or whether the ~4× memory overhead (for four 7B experts) negates the active-parameter advantage.

7. Implications and Future Directions

How This Work Changes the Landscape

Branch-Train-MiX introduces a conceptual reframing rather than a paradigm shift. It does not invent either embarrassingly parallel expert training or MoE architectures—both are well-established. What it changes is how the field should think about the relationship between these two techniques: not as alternative methods to choose between, but as complementary phases along a continuous spectrum of compute allocation.

The magnitude of this reframing becomes clear when you examine how the paper's own baselines sit on this spectrum. Branch-Train-Merge (Li et al., 2022a) allocates 100% of compute to parallel expert training with 0% joint optimization and achieves 43.4 average performance on the evaluation suite. Sparse upcycling (Komatsuzaki et al., 2022) allocates 0% to independent training and 100% to synchronized MoE training, achieving 47.3 in compute-matched comparison. BTX, at an intermediate allocation of roughly 77% expert training and 23% MoE finetuning, achieves 47.9—outperforming both extremes. This U-shaped (or more precisely, inverted-U-shaped) relationship between the allocation ratio and final performance is the signature of a genuine optimization dimension. Before this paper, the question "how should I allocate my continued pretraining budget between independent expert training and joint integration?" was not a question the literature recognized as askable. After this paper, it is.

There is also a methodological shift in what counts as evidence that MoE routing is working. The paper's routing analysis in Figure 3, particularly the dead-expert phenomenon visible without load balancing, establishes a diagnostic standard that goes beyond aggregate accuracy metrics. Showing that the Code expert collapses to near-zero routing probability in later layers without load balancing, and that load balancing restores it to dominance on code tasks, is a stronger form of evidence than simply reporting that BTX outperforms a dense baseline. It demonstrates how the architecture achieves its gains rather than just that it does. This kind of per-layer, per-expert routing distribution analysis should become standard practice for MoE papers going forward, analogous to how attention visualization became standard for understanding dense Transformers. The paper's own per-task routing breakdown (Figure 6: GSM8K preferring Code and Llama-2 experts while MATH prefers the in-domain Math expert) provides a template for this kind of analysis.

Another shift concerns the resolution of the "do MoE experts specialize by domain?" question. Prior work (Jiang et al., 2024) had shown that naturally trained MoE experts do not develop domain-level specialization, instead developing more abstract, non-domain-specific functional specializations. This paper's negative result with expert blending (Table 5: performance collapses from 34.7 to 22.2 when domain coherence is destroyed) provides complementary evidence from the opposite direction: when you inject domain specialization through initialization, it is useful and the model leverages it, but it requires explicit engineering—it does not emerge spontaneously. The natural-specialization hypothesis, at least at the granularity of domains like math/code/Wikipedia, is effectively dead after this paper and the Mixtral findings taken together. The productive question shifts from "will experts specialize?" to "what kind of initialization bias produces useful specialization, and how do we preserve it during joint training?"

The paper also subtly shifts the narrative around how to deal with catastrophic forgetting in continued pretraining. The standard framing is that catastrophic forgetting is a problem to be solved—through replay buffers, elastic weight consolidation, progressive networks, or other continual learning techniques. BTX reframes it differently: the forgetting that occurs during individual expert training (Table 1: the math expert's TriviaQA drops from 58.5% to 37.1%) is temporary and recoverable through architectural integration. You don't need to prevent forgetting during the expert phase; you just need to recover the lost capabilities during the MiX phase. This is a fundamentally different way of thinking about the problem. The paper's evidence for this is in the world knowledge scores: individual experts lost 21-29 points on TriviaQA, but the final BTX model recovers to 57.1%, within 1.4 points of the seed model's 58.5%. The forgetting happened, and then it was undone through joint finetuning of the router and averaged attention weights, without needing to replay the original pretraining data to each expert individually.

This reframing makes certain research directions more attractive and others less so:

More attractive: Work on the optimal allocation ratio for BTX-style training (what fraction of compute should be parallel vs. joint?), on methods for equalizing expert "general usefulness" so load balancing is less necessary, on domain discovery for BTX (how do you decide what domains to train experts on?), and on understanding the minimum finetuning budget needed to recover from catastrophic forgetting given the expert configuration.
Less attractive: Work on preventing catastrophic forgetting during continued pretraining through architectural regularizers or replay mechanisms, at least for the use case where experts will be integrated via MoE. If forgetting is recoverable through a short joint phase, preventing it during expert training may be an unnecessary cost. Also less attractive: work on sophisticated domain classifiers for Branch-Train-Merge-style inference, since learned token-level routing substantially outperforms heuristic classifiers (47.9 vs. 43.4 for BTX vs. BTM).

Follow-Up Research This Work Enables

Characterizing the compute-allocation trade-off surface. The paper frames BTX as generalizing BTM (100% expert training) and sparse upcycling (0% expert training) but tests only one intermediate ratio. A systematic sweep of the allocation fraction—say, at 0%, 25%, 50%, 75%, 90%, and 100%—would reveal the shape of the performance-vs-allocation curve. Is the optimum broad or sharp? Does it depend on the number of domains? On the data volume per domain? The paper's finding that freezing expert FF weights has "little impact" (Table 5) hints that the curve might be relatively flat beyond some minimum finetuning budget, but verifying this requires varying the allocation ratio while keeping total compute fixed. A strong study would hold total GPU-days constant while sweeping the ratio across 6-8 values for 5, 10, and 20 domains, producing scaling laws for BTX analogous to how Hoffmann et al. (2022) produced scaling laws for pretraining data-vs-parameters allocation.

Causal verification of domain-expert routing benefits. The paper's routing analysis is correlational: it shows that code-domain tokens are preferentially routed to the Code expert, but not that this routing causes the performance improvements. A causal experiment would force the router to use specific experts for specific token types and measure performance. For example, on the HumanEval evaluation set, run inference three times: (1) normal learned routing, (2) forced routing to the Code expert for all tokens, and (3) forced routing to the Math expert for all tokens. If normal routing outperforms both forced configurations, that demonstrates the router is making genuinely useful expert selections rather than simply associating expert indices with token patterns. If forced Code-expert routing matches or exceeds learned routing, that suggests the domain specialization is so strong that a simple domain detector could replace the learned router. The paper's anecdotal routing examples in Table 6 are insufficient to distinguish these possibilities.

BTX followed by SFT/RLHF: does the multi-domain balance survive alignment? The paper's primary motivation for BTX over BTM is that BTX produces a single model capable of downstream fine-tuning, but this capability is never tested. A direct experiment would take the BTX Top-2 model (the one achieving 47.9 average in Table 2), apply standard instruction tuning (e.g., on the same data mixture used for Llama-2-Chat), and evaluate both the aligned model's multi-domain performance and its routing distributions. Key questions: Does SFT cause the router to collapse toward the generalist expert (since instruction data is typically not domain-specialized)? Does RLHF preserve the domain-appropriate routing patterns visible in Figure 3, or does it shift all routing toward the expert that produces outputs the reward model prefers? Does the aligned BTX model still outperform an aligned dense baseline on domain tasks? The absence of this experiment is the single largest gap between the paper's stated motivation and its empirical validation.

Scaling the number of domains with unsupervised domain discovery. The paper's three-domain experiment (math, code, Wikipedia) is the minimal interesting case. Gururangan et al. (2023) demonstrated unsupervised domain discovery for Branch-Train-Merge, clustering training data into domains based on data distribution similarity. Applying this to BTX would test the central scaling claim: does the throughput advantage of parallel expert training grow linearly with the number of domains? A 20-domain experiment, with each domain expert trained on its cluster's data and then combined via MoE, would reveal several dynamics the three-domain case cannot: whether load balancing becomes harder with many experts (since the router must distribute tokens across 20 experts while still making useful specialization decisions), whether the performance gain per added domain diminishes, and whether the MoE finetuning data mixture becomes the bottleneck for balancing performance across many domains rather than the expert training itself.

Stress-testing the load balancing dependence with systematic α sweeps. The paper's α = 0.01 load balancing coefficient was chosen without documented tuning. A systematic sweep of α across orders of magnitude (0, 0.001, 0.003, 0.01, 0.03, 0.1) would map the domain-performance trade-off that load balancing controls. The current data (Table 5) shows that α = 0.01 trades math performance (GSM8K 29.8 vs. 34.6 without LB) for code performance (HumanEval 27.4 vs. 19.5). Sweeping α would reveal whether there exists an intermediate value that captures most of the code gain with less math loss, or whether the trade-off is inherently sharp. It would also test whether the dead-expert phenomenon (Code expert routing probability collapsing to zero without LB, visible in Figure 3) is a threshold effect—does a small load balancing coefficient like α = 0.001 prevent the collapse entirely, or do dead experts emerge below some critical α?

Expert training order and domain overlap effects. Table 1 reveals that the math expert improved code performance (from 16.8 to 33.6 on the code aggregate) even without code training, while the code expert barely improved math (8.6 to 8.0). This asymmetry in cross-domain transfer suggests that the order in which experts are trained matters—a math expert trained first might produce better representations for subsequent code specialization than vice versa. An experiment varying expert training order (train math then code vs. code then math vs. both from a common seed) and measuring both individual expert quality and final BTX performance would characterize transfer effects. More broadly, the degree of domain overlap—measured by cross-domain improvement in Table 1—might predict which experts will dominate routing without load balancing and therefore how large α needs to be. If domain overlap is high (as between math and code), load balancing becomes more critical because the strongest domain-overlapping expert will naturally attract most tokens.

Practical Applications and Downstream Use Cases

Multi-domain model consolidation for organizations with separate specialist teams. A common pattern in large AI organizations is that different teams independently train specialist models for their domains—a code team produces a code model, a math team produces a math model, a knowledge team produces a retrieval-augmented model. These specialists excel in their domains but cannot be shipped as a single product. BTX provides a turnkey recipe for consolidating them: branch from a shared seed, let each team train independently on their data with their preferred hyperparameters, then run the MiX stage to fuse them into a single deployable model. The paper's throughput numbers (Table 3: 533B tokens processed in 7.8 days, with the expert phase being embarrassingly parallel) mean the consolidation step adds minimal overhead relative to the expert training that already happened. The final model would activate roughly 11.1B parameters per token (for Top-2 routing with four experts) while carrying the specialized knowledge of all teams' work, dramatically simplifying deployment and enabling post-training improvements that are impossible with separate models.

Cost-efficient continued pretraining for domain-specific enterprise deployments. An enterprise wanting to adapt a general-purpose LLM to their domain (say, legal documents, financial reports, or medical literature) traditionally faces a choice: fine-tune the full model on domain data (risking catastrophic forgetting of general capabilities) or train a separate domain model and route between it and a general model at inference time (increasing serving complexity and cost). BTX offers a third path: train a domain expert on the enterprise's specialized corpus in parallel with retaining the original model as a generalist expert, then fuse them via a short MoE finetuning stage. The finding that freezing expert FF weights during finetuning "has little impact on performance" (Table 5) means the expensive domain-specific training is a one-time cost, and adding new domains later requires only training the new expert and re-running the MiX stage. An enterprise with 5-10 specialized data domains could scale their model's expertise linearly with the number of domain experts while keeping inference cost roughly constant (always activating k experts, regardless of total expert count), achieving the paper's projected efficiency: "the number of active experts can remain the same while its overall capacity increases."

Late-stage pretraining augmentation for open-source model releases. The paper's comparison with Llama-2 13B is particularly instructive for organizations releasing model families at multiple scales. Llama-2 13B required roughly 2x the pretraining compute of Llama-2 7B (Touvron et al., 2023, Table 2) and achieved 45.4 average score. BTX, starting from the already-released Llama-2 7B, achieves 47.9 with "less than half of the additional training compute compared to Llama-2 13B." This suggests a release strategy: instead of pretraining a larger dense model from scratch, release a base model (7B), then periodically release BTX-augmented versions that incorporate new domain expertise without requiring a full retraining from scratch. Each new domain (e.g., a new programming language, a new scientific discipline) would be trained as an independent expert in the embarrassingly parallel phase and integrated via a brief MiX update, producing a model that is both more capable and more up-to-date than a monolithic pretrained model frozen at its training cutoff. This turns model releases from one-time events into a continuous integration pipeline where new capabilities are added without regressing on existing ones.

When to Prefer This Method

The paper explicitly positions BTX against three alternatives: dense continued pretraining, Branch-Train-Merge, and sparse upcycling (conventional MoE training without the parallel expert phase). The empirical comparisons in Tables 2 and 3 support the following decision rules grounded in the paper's own numbers:

Prefer BTX over dense continued pretraining when: You have multiple clearly distinct data domains (math, code, factual knowledge, etc.) and the seed model already has non-trivial performance on each domain. The paper's dense baseline achieved 44.5 vs. BTX's 47.9 on the same total data (Table 2), with the gap driven primarily by domains where multi-task interference is most severe (math: 18.3 vs. 27.4; code: 25.8 vs. 34.0). If your domains overlap heavily or you lack sufficient data per domain to train meaningful expert specializations (the Wikipedia expert improved Natural Questions by only 5.4 points, Table 1), the advantage of parameter isolation through separate FF layers diminishes and the dense baseline becomes more competitive. The dense approach also wins on implementation simplicity—if the 3.4-point average gap (47.9 vs. 44.5) is not worth the additional engineering complexity of MoE training with load balancing, routing, and expert management, dense continued pretraining remains a reasonable default.

Prefer BTX over Branch-Train-Merge when: You need a single unified model for downstream fine-tuning (SFT, RLHF) or for simpler deployment. BTM keeps experts as separate models and requires running a domain classifier plus multiple forward passes at inference time, and it cannot undergo post-training as a unified system. BTX outperforms BTM by 4.5 points on average (47.9 vs. 43.4, Table 2), with the largest gaps in domains where catastrophic forgetting was most severe during expert training (world knowledge: 41.0 vs. 26.9, a 14.1-point gap). However, BTM's training is simpler (no MoE finetuning phase, no router training, no load balancing to tune) and the expert models remain independently usable. If your deployment can accommodate separate expert models with a domain classifier, and you never need to run SFT or RLHF on the combined system, BTM's 4.5-point performance penalty may be acceptable given its reduced complexity—though the paper provides no inference latency or memory comparison to quantify the deployment trade-off.

Prefer BTX over sparse upcycling when: Training throughput and data volume matter more than conceptual simplicity. BTX and sparse upcycling achieve comparable final performance (47.9 vs. 47.3, Table 3) with similar total training time (7.8 vs. 7.9 days), but BTX processes 2.1x more tokens (533B vs. 252B) because the embarrassingly parallel expert phase has higher throughput than synchronized MoE training. If you have large volumes of domain-specific data that would benefit from processing (the paper's experts trained on 201B, 210B, and 42B tokens respectively), BTX lets you process all of it without paying the synchronization cost for every token. If data volume per domain is small, the throughput advantage shrinks and sparse upcycling's simplicity—just one training phase with one data mixture—becomes more attractive. Sparse upcycling also avoids load balancing fragility (it was tested with load balancing in the paper's DM and CM baselines, but the optimal α value may differ from BTX's).

Prefer none of the above (use individual specialists) when: Domain performance is the sole metric that matters and general capabilities are irrelevant. The paper's individual experts (Table 1) achieve stronger single-domain performance than any unified model: the math expert's GSM8K of 39.5 vs. BTX's 37.1, the code expert's HumanEval of 31.7 vs. BTX's 28.7 (Table 8). BTX leaves 1-4 points of domain performance on the table in exchange for recovering 20+ points of general performance. If your application only needs math capability and will never see code, knowledge, or reasoning queries, training a pure math specialist and deploying it alone will outperform BTX on the metrics you care about. The BTX architecture adds 4x memory overhead (storing four experts) for multi-domain capabilities that the single-domain use case never invokes.