ToMoE: Converting Dense Large Language Models to Mixture-of-Experts through Dynamic Structural Pruning
Shangqian Gao, Ting Hua, Reza Shirkavand, Chi-Heng Lin, Zheng Tang, Zhengao Li, Longge Yuan, Fangyi Li, Zeyu Zhang, Alireza Ganjdanesh, Lou Qian, Xu Jie, Yen-Chang Hsu
cs.LG, cs.CL
2025-01-26
ToMoE turns dense LLMs into MoE via differentiable dynamic pruning and frozen weights. On LLaMA-2 7B at 50% active params, WikiText-2 perplexity is 8.36 versus 9.84 for DISP-LLM.
Structural pruning deletes whole channels, heads, or layers once they look unimportant. Capacity leaves with them, and recovering it usually means another round of fine-tuning that looks like continued pretraining. Dynamic pruning keeps structures input-dependent, but the compute budget then varies per token, which is awkward for batched decode and prefilling.
Turning a dense model into a sparse Mixture-of-Experts (MoE: each token activates only a few expert subnetworks) is the other common fix. Prior dense-to-MoE recipes such as LLaMA-MoE and CMoE split expert construction from router training, then spend billions to trillions of tokens on continued pretraining. ToMoE claims something stricter: usable experts already sit inside dense MLPs, and you can learn the split and the router together without touching the original weights.
The target is decoder-only Transformers. Each MLP is sliced along the intermediate dimension into N experts (8 by default) with top-1 routing per token. Attention is handled along the head dimension: Query and Key get a static mask shared across RoPE's paired subspaces; Value and Output get token-wise top-K. Every token therefore keeps the same head width, so prefilling and batched serving stay parallel.
Discrete choices are made differentiable with Straight-Through Gumbel-Sigmoid and Gumbel-Softmax. A Bi-GRU hypernetwork emits expert embeddings that couple layers, then small projections turn those embeddings into masks. Original weights stay frozen. Only the router, the projections, and the hypernetwork train, for 10,000 AdamW steps on 1 to 4 A100s. The task loss is KL between the dense logits and the converted model's logits. The teacher is the same checkpoint with ToMoE modules switched off, so a second copy of the model is not loaded.
Three regularizers pin the architecture down. A union loss pushes the experts to cover almost every original neuron. A parameter loss penalizes the widest expert so the active-parameter budget is a hard cap. A Switch-Transformer load-balancing term spreads tokens. After training, the hypernetwork and the MLP projection can be dropped. On LLaMA-2 7B the leftover extras are about 0.0184B parameters, 0.27% of the dense model. Calibration mixes WikiText, Alpaca, and Code-Alpaca one third each; WikiText alone hurts zero-shot scores.
On WikiText-2, dense LLaMA-2 7B sits at 5.12 perplexity. At 70%/60%/50% active parameters ToMoE scores 6.41/7.17/8.36, against DISP-LLM at 6.85/8.11/9.84 and ModeGPT at 7.51/8.41/11.88. The 50% result of 8.36 already beats most baselines at 70% active. The 13B copy follows: 6.78 at 50% versus DISP-LLM's 7.11. Against 2:4 semi-structured pruning on 7B, SparseGPT is 10.17 and ToMoE is 8.36.
Zero-shot average over ARC-e/c, PIQA, WinoGrande, and HellaSwag:
| Model | Active | ToMoE | Strongest baseline | Dense |
| LLaMA-2 7B | 60% | 60.72 | ModeGPT-Alpaca 57.58 | 69.00 |
| LLaMA-3 8B | 70% | 66.67 | ModeGPT-Alpaca (75%) 63.96 | 72.75 |
| Qwen-2.5 7B | 50%, N=16 | 59.90 | DISP-LLM 53.69 | 71.77 |
| Phi-2 | 70% | 64.16 | DISP-LLM 59.00 | 72.17 |
The gap versus conversion MoEs is wider. On LLaMA-2 7B at 50% active, ToMoE averages 56.07; LLaMA-MoE E8A2 after fine-tuning reaches 42.31. On a broader suite (BoolQ, SciQ, and related tasks) ToMoE hits 61.39 using 0.02B tokens; CMoE fine-tuned on 1.2B tokens reaches 51.68. Throughput at batch 1536 is 2919 tokens/s, a bit above LLaMA-MoE's 2888, against about 1858 for the dense model.
Raising N from 8 to 16 helps; 24 barely moves the needle. Whole-head pruning in the ablation collapses the average to 44.99. Dropping dynamic routing on attention V/O costs about one point.
This is a deployment shortcut for a dense model you already paid to train: leave the weights alone, spend about as much as learned pruning, and serve through an existing MoE stack. It fits teams that want roughly half the active parameters out of a 7B/8B/13B checkpoint and do not want another continued-pretraining run. It patches the capacity hole that static pruning leaves. It does not train an MoE from scratch.
On LLaMA-3 8B at 50% active, with MMLU in a wider suite, ToMoE averages 57.40. LLaMA-MoE-v2 trained on 7B tokens reaches 59.61. That is about 350 times fewer tokens and a slightly lower score. If the budget for continued pretraining exists, the latter path can still win.
The paper has no dedicated limitations section. Query and Key cannot be pruned dynamically per token: mismatched masks shrink the effective attention width below the budget K. Expert visualizations mostly track syntactic spans rather than topics; a numeric expert shows up mainly on math. Regularization also pushes expert widths together, so capacity specialization is flattened.
Most numbers are commonsense zero-shot. MMLU appears only on the LLaMA-3 8B accuracy-vs-sparsity curve. There is no comparison against a from-scratch native MoE such as DeepSeekMoE at matched active parameters. "No fine-tuning" means the original weights stay frozen; the router still trains for 10,000 steps. The gap to dense remains large: 56.07 versus 69.00 on LLaMA-2 7B zero-shot at 50% active. Throughput gains show up clearly only at large batch sizes.