SlimQwen：探索大型 MoE 模型预训练中的剪枝与蒸馏 SlimQwen: Exploring the Pruning and Distillation in Large MoE Model Pre-training

SlimQwen 的完整流程分为三大模块。模块一：三维结构化剪枝。输入是 Qwen3-Next-80A3B 教师和 1024 条校准样本 $D$。（a）深度剪枝直接删去最后 12 个 Transformer 块（3 个完整 + 9 个线性注意力），保留前 36 层，得到 $\tilde{L}=36$；（b）宽度剪枝按 $I_k^{norm}=\sum_{\ell=0}^L \mathrm{Mean}[\mathrm{RMSNorm}(X^{(\ell)})]_k$ 在校准集上聚合 RMSNorm 输出的逐维 mean absolute activation，保留 top-1536 个维度；（c）专家剪枝/合并用三类 importance 指标——Frequency $I_i^{Freq}=\mathbb{E}[\mathbb{1}[i\in A(x)]]$、Soft Logits $I_i^{Soft}=\mathbb{E}[\frac{\mathbb{1}[i\in A(x)]\cdot z_i(x)}{\sum_{j\in A(x)} z_j(x)}]$、REAP $I_i^{REAP}=\frac{1}{|X_i|}\sum_{x\in X_i} z_i(x)\|E_i(x)\|_2$——评估 512 个专家，合并到 256 个，且 top-k 由 10 降为 8。模块二：部分保留合并策略（Algorithm 1）：取 $\tilde{N}=256$，令 $|S_{keep}|=\lfloor 256/2\rfloor=128$ 保留 importance 最高的 128 个专家，从余下专家中再选 128 个作为 merge base $S_{base}$，把剩余 $512-128-128=256$ 个被丢弃专家按 CosineSim 找最近 base 并按 importance 加权合并进去，得到最终 $\tilde{E}=\{E_i: i\in S_{keep}\}\cup\{\tilde{E}_j: j\in S_{base}\}$。模块三：持续预训练与蒸馏。输入是剪枝后的学生模型和 1024 样本训练流。损失为 $\mathcal{L}=(1-\lambda)\mathcal{L}_{LM}+\lambda\mathcal{L}_{KD}+\beta\big((1-\lambda)\mathcal{L}_{MTP-LM}+\lambda\mathcal{L}_{MTP-KD}\big)$，其中 $\lambda$ 由 1.0 线性衰减到 0.75，$\beta$ 由 0.3 余弦衰减到 0.1。$\mathcal{L}_{LM}=-\frac{1}{T}\sum_i \log p_i[t_i]$，$\mathcal{L}_{KD}=-\sum_v q_i[v]\log p_i[v]$，MTP 损失形式见方法 key_idea。模块四：渐进式两阶段调度（depth-first）：Stage-1 把深度先剪到目标的一半（深度从 48 → 42），用 40B tokens 训练；Stage-2 再剪到最终 36 层并完成 width（2048→1536）和 expert（512→256）压缩，用 360B tokens 训练。学习率全局余弦衰减 4e-4 → 3e-5，前 2000 步 warmup，batch size 120B 设置下 512，400B 设置下 1024。

从技术新颖性看，本文「集成性」远大于「单点新颖性」，但每一点都对应着对 MoE 压缩的新认识。第一，把「剪枝作为强初始化」提升到工业可用的实证证据层面：Random Init + KD 平均仅 61.66，而 Pruned + KD 达 73.45，差距 11.79 分，且在 MMLU、GSM-8K、EvalPlus 上全部单调增益，这一证据规模在 MoE 领域前所未有。第二，部分保留专家合并解决了 MoE 专家压缩的一个内在张力：既要保留教师已学到的专家特化（纯剪枝风险低），又要整合被丢弃专家的互补知识（纯合并风险同质化），half-half 是一个简洁的对称设计，且在 400B tokens 训练后仍能保持一致提升（表 2 中 MMLU 69.05→69.28、MMLU-Pro 42.62→44.05、GSM8K 71.08→74.18）。第三，MTP 蒸馏是首个把 MTP 模块的软标签作为蒸馏目标的 pretraining-scale 工作，在 backbone（MMLU 75.13 vs 74.16、EvalPlus 67.32 vs 69.32）和推测解码接受率（GSM8K acc_2 24.35% → 37.06%、acc_4 4.09% → 8.24%）两端都带来提升。第四，渐进式 depth-first 的简单两阶段调度在所有 8 个基准上稳定优于一次性剪枝（MMLU 75.86 → 77.39），证明「先减一半再减完」比「一步到位」的优化轨迹更平滑。整体上，这些设计都不是孤立的技巧，而是同一个假设——「MoE 压缩的本质是把教师的能力以更平滑的轨迹转移到学生」——的不同侧面。

独立分析本文的弱点有四条。第一，启发式选择缺乏消融：half-target-expert 是 partial-preservation 的核心超参，但仅与 0 比较了「是否保留」，没比较 1/4、1/3、2/3 等其他比例，读者无法判断 50% 是否真的最优；类似地，$|S_{keep}|$ 与 importance 排序之间是否对最终性能敏感也未讨论。改进方向：做 importance ratio grid search，并用更细粒度指标（如分桶 importance 差）指导比例选择。第二，校准集代表性：仅 1024 条样本做 importance 估计，但 SlimQwen 在 400B tokens 上训练，校准-训练分布错配的可能性被低估。改进方向：探索校准集大小（256/1024/4096/16384）和多样性（多语言、多领域）对最终性能的影响。第三，MTP 蒸馏的边际收益与开销不匹配：MTP KD 让 EvalPlus 涨 2 分左右，但 MTP 模块本身需要额外 forward 和显存（共享 embedding/head 不计），论文没有给出 MTP 蒸馏对训练 FLOPs 和显存的具体放大系数。改进方向：报告 training FLOPs ratio 与 accuracy gain 的 Pareto 曲线，并研究 selective MTP（只在后期开启）等加速技巧。第四，教师-学生架构绑定：SlimQwen 只能在 Qwen3-Next 内部完成压缩，因为 Gated DeltaNet + Gated Attention 的混合层结构被原样保留；若想压缩到纯 attention 或纯 Mamba 架构，本文方法不能直接复用。改进方向：探索 architecture-agnostic 的剪枝方案，例如把 Gated DeltaNet 块换成标准 attention 或反之，并研究块类型对齐对蒸馏的影响。第五，安全性与对齐未涉及：压缩可能放大偏见或降低安全对齐，论文没有任何 safety benchmark 的评估。改进方向：在 SlimQwen 上跑 BBQ、TruthfulQA、ToxiGen 等基准，验证压缩是否带来 alignment tax。

作者明确提出的方向：（1）把 partial-preservation 思想推广到 dense-to-MoE 或不同 MoE 拓扑之间的迁移；（2）研究自动化的 expert retention ratio 搜索；（3）把 MTP 蒸馏推广到 SFT 阶段的端到端训练。基于成果可延伸的方向：（a）多模态 MoE 压缩——本文仅在文本 LLM 上验证，但 Qwen2-VL、LLaVA 等多模态模型也使用 MoE，部分保留策略是否同样适用于视觉专家值得验证；（b）推测解码与蒸馏联合优化——表 4 显示 MTP KD 改善推测解码，可进一步把接受率作为训练 signal 做 RL 微调，而不是单纯用 KL 蒸馏；（c）更细粒度的 curriculum——表 9 显示 2 阶段已足够，但可以探索基于 validation loss 的自适应阶段切分；（d）小模型的 MTP 蒸馏 scaling law——当前结果只在 23B 和 6B 两个学生模型上验证，缺少 1B、3B 级别小模型的 MTP KD 是否仍带来一致收益的分析；（e）与稀疏 attention、量化、speculative decoding 的协同——例如把 4-bit 量化 + SlimQwen-23A2B + vLLM 推测解码组合，可能在消费级 GPU 上实现接近原 80A3B 的性能；（f）数据维度——本文训练 token 120B/400B 仍远小于现代 pretraining 规模（数万亿），后续可研究当 teacher 与 student 训练 token 数量级相近时，部分保留策略是否仍然有效。

开源情况：论文未明确给出代码仓库链接，但基于 Megatron-LM 与 Qwen3 MoE 训练框架（Maestro, Yuan et al. 2026）实现，这些底层框架已开源。数据：全部实验基于内部高质量、多样化文本语料，未公开——但训练设置、batch size、token 数都详细列出（表 7：120B 设置 batch=512，400B 设置 batch=1024，peak LR 4e-4，cosine 衰减到 3e-5，2000 步 warmup）。算力需求：在 Alibaba Cloud 上训练，400B tokens、batch size 1024、80B 教师的 KD——粗略估算需要数百到上千张 H800/A100 等效 GPU，中小机构难以复现最终 SlimQwen-23A2B 的训练，但部分消融（如 120B token、表 2、表 8）计算量可控。复现难度：中等到高。具体瓶颈在于：（1）1024 条校准数据的采样与 importance 计算虽然简单，但其与最终性能的关系缺乏验证，可能带来 1-2 分随机波动；（2）训练 loss 中 $\lambda$、$\beta$ 的衰减曲线虽列出，但 $\mathcal{L}_{KD}$ 实现细节（logit temperature、是否 masking padding 等）未完全给清；（3）MTP KD 的实现需要把教师的 MTP 模块完整加载，显存翻倍；（4）不同 seed 间方差未报告，单次实验的偶然性无法评估。降低复现门槛的建议：作者可在附录给出 importance 估计的 standalone 脚本，以及 partial-preservation merging 的伪代码（Algorithm 1 已给出）对应 50 行的 PyTorch 实现，使社区能在小模型（6B）上验证结论，再 scale 到 80B 教师。