DivMoE: Fine-Grained MoE Upcycling via Cross-Domain Expert Composition

AI summary

Overview

Research area: Efficient training and architecture design for large language models, specifically Mixture-of-Experts (MoE) upcycling — converting pre-trained dense models into sparse MoE models.

Technical level: Intermediate. The paper assumes familiarity with transformer FFN layers, top-k routing, and continual pre-training, but its core ideas are explained largely through the interaction between expert initialization and routing behavior.

Scope: The paper introduces DivMoE, a framework for fine-grained MoE upcycling that combines domain-specialized expert initialization with a hard diversity-constrained routing rule, and evaluates it across two dense base models, 15 benchmarks, and six upcycling baselines.

What This Paper Is About

Sparse upcycling takes an already-trained dense language model and turns it into an MoE by splitting its feed-forward layers into multiple experts, which is far cheaper than training an MoE from scratch. The authors identify a failure they call fine-grained upcycling collapse: when fine-grained experts are all derived from a single source model, routing collapses and accuracy drops to near-random. DivMoE's goal is to make fine-grained upcycling work by building experts from several domain-specialized dense models and structurally forcing each token to draw experts from different domains.

Key Contributions

1. Identification of a fine-grained upcycling collapse pathology. The authors show that combining fine-grained expert partitioning with single-source initialization produces routing collapse and near-random downstream accuracy, independent of the specific partitioning scheme. On Qwen3-1.7B-Base, Drop-Upcycling-fine-grained reaches 23.2% average accuracy across 15 benchmarks, essentially matching from-scratch training at 22.2%, while the same method's coarse-grained variant reaches 50.2%.

2. The first fine-grained MoE upcycling framework. DivMoE combines (a) domain-specialized fine-grained expert initialization — vertically slicing each of n domain specialists' FFNs into m shards to yield N = n × m experts in n domain groups — with (b) diversity-constrained routing, a hard structural rule requiring each token to activate at most one expert per domain group.

3. A public-corpus training pipeline that avoids post-CPT regression. Across three stages (domain-specialist initialization, joint MoE continual pre-training, supervised fine-tuning) using only public corpora, DivMoE strictly improves over the dense base model on every benchmark after Stage 2, closing a regression gap that prior upcycling pipelines exhibit.

4. Comprehensive controlled experiments. Results span two base models plus a third scale, 15 benchmarks, three model scales (1B, 1.7B, 4B base to 3.2B, 5.4B, 12B MoE), and six upcycling baselines sharing identical Stage 2 corpus, optimizer, and step count.

Main Findings

• DivMoE outperforms six upcycling baselines. On Qwen3-1.7B-Base, DivMoE reaches 55.6% average accuracy over 15 benchmarks versus 51.6% for the strongest baseline (BTX) and 51.1% for NVIDIA Upcycling. Over the 11-benchmark main-paper subset it scores 54.3% versus 50.8% for BTX and 50.5% for NVIDIA Upcycling. On Llama3.2-1B it reaches 37.4% versus 35.5% for BTX on the 11-benchmark subset.

• Fine-grained upcycling collapses without domain diversity. From-Scratch (22.2%) and Drop-Upcycling-fine-grained (23.2%) both collapse to near-random performance on Qwen3-1.7B-Base over 15 benchmarks, despite very different initialization schemes. Coarse-grained Drop-Upcycling does not collapse (50.2%), isolating the failure to the combination of fine-grained partitioning with non-domain-specialized initialization.

• Stage 2 continual pre-training removes the regression below the dense base. DivMoE surpasses Qwen3-1.7B-Base on every benchmark evaluated — for example MMLU 64.5 vs. 62.6, MATH 40.5 vs. 38.4, and BBH 55.5 vs. 53.5. It also beats a data-matched dense baseline (53.4% average, continually pre-trained on the same 540B tokens without MoE conversion) by roughly 2 points while activating only about 2.4B of its 5.4B parameters per token.

• The diversity constraint helps domain-specific tasks rather than hurting them. Adding the constraint improves HumanEval by +1.8, MBPP by +2.3, GSM8K by +2.1, and MATH by +1.4, refuting the intuition that a coding token should only use coding experts.

• Component synergy is real, not additive. In the 15-benchmark ablation on Qwen3-1.7B-Base, removing the diversity constraint drops DivMoE from 55.6% to 52.5%; removing domain experts (NVIDIA Upcycling) gives 51.1%; removing fine-grained partitioning (BTX) gives 51.6%; removing diversity and fine granularity (Drop-Upcycling-fine-grained) collapses to 23.2%.

• The 12B DivMoE matches a 16B model. Upcycled from Qwen3-4B-Base, DivMoE reaches 64.5% average accuracy on the full 8-benchmark set, matching Moonlight-MoE (16B total) at 64.5%, and reaches 63.0 vs. 62.3 on the 7-benchmark subset in Table 4, where it leads on MMLU (+1.5), TriviaQA (+7.0), and MATH (+2.2). The abstract reports it outperforming a controlled NVIDIA-Upcycling baseline by 6.1 percentage points; Table 4's controlled comparison shows 63.0 versus 56.2 (a +6.8-point gap) and +9.3 over OLMoE (53.7) on that subset.

• Routing behavior. Sparse Upcycling shows routing collapse in deeper layers, with over 60% of tokens routed to two experts at layer 36. DivMoE's domain-specialized initialization already improves utilization, and the constraint balances routing across groups while preserving dataset-specific preferences — GSM8K activates more math experts, HumanEval more code experts.

• Granularity. Varying m = k over {1, 2, 4} shows progressively lower training loss with higher granularity, but m = k = 4 exhibits mild overfitting on MMLU and HellaSwag. The default is m = k = 2.

• Compute cost. DivMoE and BTX share identical total compute of 540B tokens (40B dense Stage 1 plus 500B MoE Stage 2), about 28,800 GPU-hours on NVIDIA H20 (80GB), versus about 26,500 GPU-hours for methods using 500B tokens only. Stage 1 adds roughly 8% overhead relative to from-scratch training.

Methodology in Plain English

The researchers start with a dense base model and make four copies of it. Each copy is continually pre-trained on a different public domain subset of Dolmino-Mix — math, code, academic, and web — for 10B tokens each (40B total), producing four domain specialists that keep general ability while absorbing domain-specific knowledge.

Each specialist's feed-forward network is then vertically sliced into two halves, producing eight fine-grained experts organized into four domain groups. Because picking one shard from any specialist reconstructs a complete FFN, the assembled model is dense-equivalent at initialization: its output distribution matches the base model, but the specialists' differing weights inject diversity into the router scores.

For routing, instead of a standard top-k over all experts — which can pick several experts from the same domain — the router first picks the highest-scoring expert inside each of the four groups, then selects the top-k groups by those representative scores. With k = 2, a token always activates two experts from two distinct domains. This is a hard structural constraint, not a gradient-based auxiliary loss, so it holds regardless of gradient magnitudes; the authors keep only a small auxiliary load-balancing loss (α = 0.01) to encourage balance within groups. Expert weights are also scaled by a magnitude-correction factor to compensate for softmax attenuation, and non-expert components (embeddings, attention, normalization) are initialized as the mean of the base model and all specialists.

The assembled MoE is then continually pre-trained on Dolmino-Mix for 500B tokens, using the same corpus, optimizer, and step count for every baseline. A final supervised fine-tuning stage on Tülu-3-SFT-Mixture is applied only for the open-source MoE comparison, and is applied identically to DivMoE, NVIDIA Upcycling, and OLMoE to isolate architecture from data.

Why This Matters

Impact on research. The paper reframes what makes upcycling work. Rather than treating routing collapse as a tuning problem solved by auxiliary losses, it argues that the reachable selection set itself must be constrained, and that expert diversity must come from the initialization rather than from random perturbation. It also documents a reproducible failure mode — fine-grained upcycling collapse — that others can now test against.

Real-world applications:

• Cheaper domain-capable models. Organizations with an existing dense checkpoint can upcycle it into an MoE with domain specialists instead of training an MoE from scratch, using only public corpora.
• Reasoning- and code-heavy assistants. DivMoE's largest measured gains are on GSM8K, MATH, HumanEval, and MBPP, the tasks where multi-step cross-domain composition appears most useful.
• Deployment under memory constraints. DivMoE's Qwen3-1.7B variant activates roughly 2.4B of 5.4B parameters per token, so capacity grows while per-token compute stays limited.
• Open-model development. The full pipeline uses Dolmino-Mix-1124, Tülu-3-SFT-Mixture, and public base checkpoints, and the authors state they release the slicing and routing code with per-domain Dolmino-Mix splits.

Industry relevance. The controlled comparison is the practical headline: a 12B-parameter upcycled model matches a 16B model trained from scratch, on 540B Stage-2 CPT tokens. That offers a concrete cost argument for upcycling over from-scratch MoE training, though the paper does not report a direct dollar comparison.

Future Directions

• Scaling to larger base models. The authors note their experiments use base models up to 4B parameters, yielding MoE models up to 12B, and that scaling to 7B–13B base models is an important direction.
• Learning the domain partition. The four-domain partition (math, code, science, commonsense) is fixed and hand-chosen; learning it end-to-end is left to future work.
• Amortizing the Stage-1 overhead. DivMoE requires an additional 40B Stage-1 token budget per base model — an 8% overhead that the authors describe as small in absolute terms but non-trivial at very large scale.
• Understanding token-level domain roles. The paper argues the diversity constraint helps because individual tokens within a coding prompt serve heterogeneous functional roles, but it does not measure those roles directly.

Target Audience

This paper is most useful to ML engineers and researchers working on MoE architectures, model upcycling, and efficient pre-training pipelines; to practitioners deciding between training an MoE from scratch and converting an existing dense checkpoint; and to readers interested in routing collapse and load-balancing mechanisms. It will also interest teams with domain-specific dense checkpoints who want a recipe for combining them into a single sparse model using public data.

Authors’ abstract

Mixture-of-Experts (MoE) architectures have become essential for scaling large language models, with recent work demonstrating the benefits of fine-grained expert designs. Training such models from scratch is expensive, and sparse upcycling from pre-trained dense models is an attractive alternative. However, we identify a structural pathology of fine-grained upcycling: when fine-grained experts are derived from a single source model, naive routing collapses and downstream accuracy drops to near-random (e.g., on Qwen3-1.7B, Drop-Upcycling-fine-grained reaches only 23.2% average accuracy across 15 benchmarks, essentially matching from-scratch training at 22.2%, while the same method's coarse-grained variant reaches 50.2%). We propose DivMoE, the first framework achieving fine-grained MoE Upcycling with structurally-balanced routing. DivMoE introduces domain-specialized fine-grained expert initialization, deriving experts from dense models that have undergone domain-adaptive continual pre-training, and diversity-constrained routing, a hard structural constraint guaranteeing that each token activates experts from distinct domain groups. Across two base models and 15 benchmarks, DivMoE consistently outperforms six upcycling baselines (55.6% vs. 51.6% for the strongest baseline on Qwen3-1.7B) and strictly improves over the dense base model on every benchmark after Stage 2 continual pre-training -- closing the regression gap that has plagued prior fine-grained upcycling. After supervised fine-tuning on a public reasoning mixture, our 12B-parameter DivMoE model matches Moonlight-MoE (16B) at 64.5% average accuracy while outperforming a controlled NVIDIA-Upcycling baseline by 6.1 percentage points.