Source-linked AI summary

SlimQwen: Exploring the Pruning and Distillation in Large MoE Model Pre-training

Shengkun Tang, Zekun Wang, Bo Zheng, Liangyu Wang, Rui Men, Siqi Zhang, Xiulong Yuan, Zihan Qiu, Zhiqiang Shen, Dayiheng Liu

arXiv:2605.08738v2cs.LGcs.AIcs.CL

TL;DR

MoE compression at pretraining scale lacks clear guidance on pruning, expert reduction, and distillation. This paper systematically evaluates these choices and finds that structured pruning, partial-preservation expert merging, complementary LM loss, MTP distillation, and progressive schedules support competitive compressed models.

  • Problem

    It remains unclear how structured pruning and knowledge distillation should compress MoE language models effectively at pretraining scale.

  • Method

    The paper systematically studies depth, width, and expert compression alongside continual-pretraining objectives and progressive pruning schedules.

  • Results

    Pruned models outperform scratch-trained targets, partial-preservation merging improves benchmarks, LM loss plus KD and MTP distillation yield gains, and progressive pruning beats one-shot compression.

  • Takeaways & Limitations

    Structured pruning provides a strong initialization and the studied recipe compresses Qwen3-Next-80A3B to 23A2B with competitive performance across major benchmarks.

Abstract

from arXiv · show

Structured pruning and knowledge distillation (KD) are typical techniques for compressing large language models, but it remains unclear how they should be applied at pretraining scale, especially to recent mixture-of-experts (MoE) models. In this work, we systematically study MoE compression in large-scale pretraining, focusing on three key questions: whether pruning provides a better initialization than training from scratch, how expert compression choices affect the final model after continued training, and which training strategy is most effective. We have the following findings: First, across depth, width, and expert compression, pruning a pretrained MoE consistently outperforms training the target architecture from scratch under the same training budget. Second, different one-shot expert compression methods converge to similar final performance after large-scale continual pretraining. Motivated by this, we introduce a simple partial-preservation expert merging strategy that improves downstream performance across most benchmarks. Third, combining KD with the language modeling loss outperforms KD alone, particularly on knowledge-intensive tasks. We further propose multi-token prediction (MTP) distillation, which yields consistent gains. Finally, given the same training tokens, progressive pruning schedules outperform one-shot compression, suggesting that gradual architecture transitions lead to better optimization trajectories. Putting it all together, we compress Qwen3-Next-80A3B to a 23A2B model that retains competitive performance. These results offer practical guidance for efficient MoE compression at scale.

1 Introduction

This section presents a systematic study of structured pruning, expert compression, distillation, and progressive training for large-scale MoE compression. It finds that pruning provides strong initialization, compression methods converge similarly after continual pretraining, and the proposed recipe compresses Qwen3-Next-80A3B to 23A2B with competitive performance.

  • Key findings: Pruning a pretrained MoE provides a significantly better initialization than training the identical target architecture from scratch under matched training tokens.The improvement holds across depth, width, and expert compression, benefiting both reasoning and generation performance.
  • Key findings: After 400B-token continual pretraining, one-shot expert pruning and merging methods achieve marginally different final performance, with no single approach dominating.The study compares criteria based on routing frequency, routing scores, and expert activations, motivating partial-preservation expert merging.
  • Key findings: LM loss improves knowledge-intensive benchmark performance when combined with knowledge distillation, while multi-token knowledge distillation yields consistent gains across major benchmarks.The work also reports improvements to backbone model training and speculative decoding from multi-token distillation.
  • Key findings: Progressive pruning strategies consistently outperform one-shot compression at the same final sparsity and total training tokens.The result supports gradual architecture transitions during post-compression training.

2 Related Work

Prior work studies structured pruning across width, depth, and MoE expert dimensions, while post-compression training uses language-modeling loss and knowledge distillation to recover degraded performance. Existing approaches differ in model scope and training strategy, motivating SlimQwen’s combination of pruning, progressive schedules, and distillation.

  • Structured Pruning in LLMs: Structured pruning for MoE LLMs targets width, depth, or experts through removal or merging.Width pruning changes hidden or FFN intermediate sizes; depth pruning removes transformer blocks; expert pruning or merging modifies MoE modules.
  • Post-Compression Training for Recovery: Post-compression training is generally required because structured pruning causes non-negligible performance degradation.Prior work uses continued training to recover the performance of pruned models.
  • SlimQwen Overview: SlimQwen combines structured pruning, partial-preservation expert pruning or merging, progressive pruning, and distillation to transform a teacher MoE model into a target architecture.The overview includes width, depth, and expert compression, staged depth-first, width-first, or joint schedules, and multi-token prediction distillation.
  • Post-Compression Training for Recovery: Prior recovery methods combine distillation or language-modeling loss, but their applicability differs across dense and MoE models.Minitron and Slim apply distillation to pruned dense models, while DarwinLM and SlimMoE use conventional LM loss and KD respectively; Minitron is non-MoE only.

3 Method

The method systematically compresses MoE models along depth, width, and expert dimensions, then transfers pretrained capabilities through knowledge distillation and progressive pruning schedules. It combines standard and multi-token prediction losses in a unified training objective.

  • Compression dimensions: MoE compression is studied across three dimensions: depth, width, and experts.The method introduces separate strategies for each compression dimension.
  • Depth pruning: Depth pruning drops the last N layers, with experiments pruning the last 25% of layers.The retained layers are Lkeep = {1, . . . , L − N}, with ˜L = L − N.
  • Width pruning: Width pruning estimates hidden-dimension importance from activation statistics and retains the dt dimensions with the highest scores.Importance is computed from a sampled calibration dataset across hybrid attention, MoE, and normalization modules.
  • Expert compression: Expert compression compares pruning and merging using frequency, soft-logits, and router-weighted output-activation importance criteria.The partial-preservation strategy retains half of the target experts intact and merges discarded experts into selected bases.
  • Distillation and progressive pruning: The training objective combines backbone language-modeling and KD losses with MTP language-modeling and MTP distillation losses.Progressive schedules interleave structural pruning with fixed-token distillation, using depth-first, width-first, or joint reduction priorities.

4 Experiments

Experiments show that pruning provides a strong initialization for large-scale MoE pretraining, while continual training reduces differences among one-shot expert-compression methods. Combining language modeling with distillation and using progressive pruning further improves compressed-model training across benchmarks.

  • Q1: Pruning Initialization: 73.45 average score versus 61.66 for random initialization, a +11.79-point gain, demonstrates the benefit of pruning-based initialization under a 120B-token budget.Both setups use knowledge distillation from the Qwen3-Next teacher.
  • Q2: Expert Compression: No one-shot pruning or merging method is uniformly superior across downstream tasks after 400B-token continual pretraining.The frequency-based router-logits grouping method reaches 60.17 on BBH, but methods differ across benchmarks; partial-preservation expert merging improves major benchmarks.
  • Q3: Training Recipe: Combining next-token prediction knowledge distillation with language modeling loss outperforms pure distillation, particularly on knowledge-intensive tasks.Adding LM loss improves knowledge benchmarks such as MMLU and MMLU-Pro, while MTP distillation yields consistent gains.
  • Q3: Training Recipe: MTP distillation consistently improves multi-token generation acceptance rates over MTP loss across most benchmarks in both pretraining and supervised-finetuning stages.The comparison uses speculative decoding acceptance rates.
  • Progressive Pruning and Distillation: Progressive pruning and distillation consistently outperform one-shot pruning on most benchmarks under the same 400B-token total budget.Progressive training uses two stages of 40B and 360B tokens, whereas one-shot pruning trains directly for 400B tokens.

5 Conclusion

The paper finds that structured pruning remains a strong initialization for continual MoE pretraining, even at high compression ratios. Expert pruning and merging metrics differ only slightly after large-scale pretraining, while partial-preservation merging consistently improves benchmark performance.

  • Pruning and merging: Structured pruning provides a strong initialization for continual MoE pretraining, including at high compression ratios.The conclusion frames pruning as an effective starting point for continued training rather than requiring training the compressed architecture from scratch.
  • Pruning and merging: Different expert pruning and merging metrics show only minor differences after large-scale pretraining.This suggests that continued pretraining substantially reduces the final-performance gap between one-shot expert compression choices.
  • Pruning and merging: Partial-preservation expert merging delivers consistent performance improvements across major benchmarks.The strategy is presented as a simple expert-merging approach that improves downstream results across the evaluated benchmark suite.

A Appendix · A.1 Architecture Details

The appendix specifies the teacher and pruned student architectures, including Gated Attention, Grouped-Query Attention, and Gated DeltaNet components. It also identifies the model configurations and parameter counts summarized for different MoE variants.

  • A.1 Architecture Details: Gated Attention operates on input hidden states X ∈ R^n×d, where d is the model hidden size and hq is the number of query heads.The architecture description introduces the formulation using these dimensions and head counts.
  • A.1 Architecture Details: Its output projection is WO ∈ R^(hqdhead)×d, and its gate uses the sigmoid function σ(z) = 1/(1+e^-z).The sigmoid is explicitly defined in the architecture description.
  • A.1 Architecture Details: The gating mechanism uses a learnable weight g ∈ R^d×1, while each attention head is computed by scaled dot-product attention.The passages separately specify the learnable gate weight and scaled dot-product computation.
  • A.1 Architecture Details: The value projection uses learnable parameters V ∈ R^d×dhead, and the model applies Grouped-Query Attention with hq query heads and hkv key/value heads.This describes the attention parameterization and query/key-value head grouping.
  • A.1 Architecture Details: Gated DeltaNet maintains St ∈ R^dv×dk, qt ∈ R^dk, kt ∈ R^dk, and vt ∈ R^dv, then reads out yt = Stqt ∈ R^dv.The state matrix and vector dimensions define the recurrent token-mixing representation.
  • A.1 Architecture Details: The Gated DeltaNet output is mapped back to model dimension d as Yt = ytWout ∈ R^d with Wout ∈ R^dv×d.The implementation maps the token-mixing output through Wout.
  • A.1 Architecture Details: In the implementation, dk corresponds to the Q/K hidden size and dv corresponds to the V hidden size.These correspondences clarify the roles of the Gated DeltaNet dimensions.
  • A.1 Architecture Details: Table 6 summarizes model configurations and parameter counts for different MoE variants.The table covers both the original teacher model and pruned student models.

A.2 Training Hyperparameters

The paper provides detailed pretraining hyperparameters in Table 7.

  • A.2 Training Hyperparameters: Detailed pretraining hyperparameters are provided in Table 7.The passage directs readers to Table 7 for the implementation details.

A.3 Implementation Detail

The implementation builds on Megatron-LM and follows Qwen3 MoE practices, using global-batch load balancing, calibration data from pretraining, and staged optimization for progressive pruning distillation.

  • Implementation Detail: The codebase is built upon Megatron-LM and applies global-batch load balancing loss for MoE, following Qwen3 MoE models.The load balancing loss follows Qiu et al. (2025a).
  • Implementation Detail: Calibration data is sampled from the pretraining data.
  • Implementation Detail: Progressive pruning distillation uses a single-stage learning-rate decay schedule, with the second stage starting from the first stage’s final-step learning rate.
  • Implementation Detail: All models use the AdamW optimizer.

A.4 Comparison of Different Depth Pruning Methods

The section compares activation-similarity pruning with directly removing the last layers in a pretrained 15A3B teacher. With four layers pruned one-shot, removing the last four layers causes only minor degradation, whereas activation-based pruning causes substantially larger drops.

  • Method comparison: The study compares activation-similarity pruning with directly pruning the last several layers.Activation similarity is computed from token-mean-pooled layer activations and adjacent-layer cosine similarity.
  • Experimental setup: The activation-based method identifies middle layers for pruning in the 24-layer teacher.Experiments use a pretrained 15A3B teacher, a 1024-example calibration dataset, and one-shot pruning of four layers.
  • Results: 75.62 to 73.86 on MMLU: pruning the last four layers causes only minor degradation, while activation-based pruning produces substantially larger drops.The comparison uses one-shot removal of four layers from the pretrained 15A3B teacher.

A.5 Results of Progressive Pruning and Distillation with More Stages · A.6 Evaluation on More Benchmarks · A.7 Efficiency Analysis

The appendix examines finer-grained progressive pruning, broader benchmark coverage, and SlimQwen’s inference efficiency. More stages do not improve over the two-stage setup, while evaluations compare pruned-weight initialization with scratch training and report standardized efficiency measurements.

  • A.5 Results of Progressive Pruning and Distillation with More Stages: Three-stage progressive pruning compares depth-first and width-first schedules, using staged pruning followed by continued training.The depth-first schedule trains for 20B tokens after each half-depth pruning, then trains for 360B tokens after width pruning.
  • A.5 Results of Progressive Pruning and Distillation with More Stages: After post-compression KD with 120B tokens, last-layer pruning recovers better performance than activation-based pruning.In one-shot pruning, last-layer pruning causes only minor degradation on benchmarks such as MMLU, whereas activation-based methods produce substantially larger drops.
  • A.5 Results of Progressive Pruning and Distillation with More Stages: More fine-grained three-stage partitions do not yield additional performance gains compared with the two-stage setup.
  • A.6 Evaluation on More Benchmarks: Additional evaluations cover Chinese and general knowledge, reasoning, in-context learning, coding, multilingual knowledge, and mathematical reasoning.Named benchmarks include CEval, SuperGPQA, KOR-Bench, ICLEval, MBPP, MMMLU, IncludeBase, and Mgsm.
  • A.6 Evaluation on More Benchmarks: Table 10 compares models trained from scratch with models initialized from pruned weights across the expanded benchmark suite.
  • A.7 Efficiency Analysis: SlimQwen’s speedup and memory are analyzed against the original model using HuggingFace and vLLM inference backends.Tests use prompt and generation lengths of 128, 10 runs with 3 warmups, the same two GPUs with tensor parallel size 2, and bfloat16 peak-memory monitoring.
  • A.7 Efficiency Analysis: Efficiency results are averaged across repeated runs under matched hardware and inference conditions.The setup uses a prompt length of 128, generation length limited to 128, 10 executions, 3 warmup runs, and tensor parallel size 2.
Loading 2605.08738v2…