跳转至

Pruning and Distilling Mixture-of-Experts into Dense Language Models

会议: NeurIPS 2026
arXiv: 2605.28207
OpenReview: https://openreview.net/forum?id=hmvlyb1yQu
代码: https://github.com/krafton-ai/moe-to-dense
领域: 模型压缩
关键词: 混合专家模型 (MoE), 模型剪枝, 知识蒸馏, 密集语言模型, D-最优选择

一句话总结

首个将训练好的 MoE 语言模型系统化转换为标准密集(Dense)架构的压缩与蒸馏框架,通过基于 Gram 对数行列式的 D-最优多样性专家打分与无平均纯剪枝拼接,搭配前向 KL 蒸馏,在相同参数量与显存占用下显著超越传统密集到密集剪枝。

研究背景与动机

混合专家模型(MoE)凭借稀疏路由机制,在大幅扩展总参数量与表征容量的同时维持了恒定的单 token 激活计算量,已成为前沿基础大模型(如 DeepSeek-V2/V3、Qwen-MoE、Grok 等)的标准架构。然而,MoE 的高推理性能在边缘端或单卡显存受限场景下遭遇了严峻的工程瓶颈:即便每个 token 仅激活极少数专家(例如 128 个中选 8 个),系统在推理前仍必须将全部专家参数完整加载至 GPU 显存中,导致极高的显存占用与内存带宽压力。为了适应端侧部署,工业界通常不得不额外从头预训练同等规模的小型密集模型,或者对专门预训练的超大密集教师模型实施逐层结构化剪枝与蒸馏,这不仅消耗了数百万 GPU 时的庞大算力,更彻底浪费了既有 MoE 旗舰模型所蕴含的高阶表征。

现有针对 MoE 的压缩研究(如 MC-SMoE、HC-SMoE、Sub-MoE、REAP 等)主要聚焦于削减专家总数或将相似专家聚合,但其压缩产物仍然保持着稀疏路由的 MoE 结构,不仅未能根除“全量专家必须常驻显存”的物理限制,还因保留门控网络而引入了不规则访存与动态路由开销。更为关键的是,以往方法普遍依赖“激活频率(Selection Frequency)”来挑选保留专家,极易选出功能高度重叠的“通用型专家(Generalists)”,导致压缩后模型在专业任务上的表征退化严重;而若直接采用简单的参数平均或粗暴降维,专家间非线性的权重干涉又会诱发灾难性能力崩塌。

本文的核心切入点在于重构 MoE 到稠密模型的映射通路:MoE 每个专家的中间维度通常经过专门设计,使得所有激活专家的特征维度之和天然匹配标准密集 FFN 的隐藏层宽度;因此完全可以通过智能挑选一组信息正交且高激活的专家集合,直接在通道维度完成无损拼接。核心 idea:将 MoE 转换为标准密集架构的关键不在于参数平均,而在于利用 Gram 矩阵对数行列式最大化(D-最优设计)联合评估专家重要度与互补多样性,以“纯剪枝直接拼接”保留互补特化专家的完整表征,并通过 MoE 教师模型的前向 KL 散度蒸馏快速恢复任务性能。

方法详解

整体框架

MoE-to-Dense 转换框架旨在将一个包含 \(L\) 层、每层 \(E\) 个专家且每个 token 激活 \(k\) 个专家的 MoE 教师模型,无缝转换为主干参数量严格等于该 MoE 激活参数量的标准稠密 Transformer 学生模型。整体流程分为四个连续阶段:首先在少量校准数据上执行前向传播,收集各专家的门控激活概率与输出隐藏状态统计量;接着利用 D-最优准则从 \(E\) 个候选中选出 \(K\) 个兼顾重要性与特征多样性的核心专家集合;随后对所选专家执行分组与矩阵拼接,对下投影权重应用静态幅度补偿因子构成标准 Dense FFN,并直接继承教师模型的自注意力层、Embedding 与 LayerNorm;最后以原 MoE 教师为监督信号,执行高效的知识蒸馏微调。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["校准数据前向采样<br/>收集专家路由置信度与输出特征Gram矩阵"] --> B["D-最优多样性感知专家选择<br/>DO-ACP贪心最大化特征内核对数行列式"]
    B --> C["纯剪枝无重合特征拼接<br/>保留K=k个专家沿维度直接拼装为稠密FFN"]
    C --> D["静态下投影幅度缩放<br/>引入静态因子α_g补偿门控路由缺失的激活能量"]
    D --> E["前向KL散度知识蒸馏<br/>由MoE教师监督快速对齐Dense学生输出分布"]
    E --> F["全稠密结构目标语言模型<br/>参数完全匹配激活量且彻底摆脱MoE显存驻留限制"]

关键设计

1. D-最优多样性感知专家选择(DO-ACP):联合优化专家输出重要性与空间多样性以避免冗余

传统 MoE 专家打分依赖全局选择频率(Selection Frequency),倾向于保留频次高但信息冗余的通用专家,丢失了低频但置信度极高的特化专家(Specialists)。本文首先提出激活加权条件概率(Activation-weighted Conditional Probability, ACP),将门控置信度从频率中剥离,并与专家输出向量模长解耦结合:

\[s_{\ell}^{(e)} = \bar{p}_{\ell}^{(e)} \cdot \sqrt{\mathbb{E}_{t}[\|f_{e}(t)\|^{2}]}\]

其中 \(\bar{p}_{\ell}^{(e)}\) 为专家 \(e\) 在被激活 token 子集上的平均 softmax 条件概率,\(f_{e}(t)\) 为输出特征。为解决高分专家之间表征高度重叠、侵占密集模型有限容量的问题,本文引入试验设计中的 D-最优准则(D-Optimality)。定义输出特征 Gram 矩阵 \(\mathbf{G}_{ij} = \mathbb{E}_{t}[\langle f_{i}(t), f_{j}(t)\rangle]\),构建重要度加权核矩阵 \(\boldsymbol{\mathcal{K}}_{ij} = \sqrt{s^{(i)}s^{(j)}} \cdot \mathbf{G}_{ij}\),专家子集选择问题形式化为对数行列式最大化:

\[S^{*} = \argmax_{|S|=K} \log\det\left(\boldsymbol{\mathcal{K}}_{S} + \lambda_{\text{reg}}\,\mathbf{I}\right)\]

该目标函数在数学上具有单调次模性(Monotone Submodular)。本文采用基于 Schur 补的贪心算法在 \(\mathcal{O}(K^3 E)\) 时间复杂度内求解,保证获得至少 \((1 - 1/e)\) 的理论近似比。对数行列式天然惩罚核矩阵中非对角线的互相关项,使得入选专家在保持单体重要性的同时,最大化张成子空间的有效秩(Effective Rank),彻底避免了高相关冗余专家的扎堆入选。

2. 纯剪枝无重合特征拼接(Pure Pruning & FFN Concatenation):保留未受权值平均污染的高质量专家表征

在决定候选专家数 \(K\) 与目标分组数 \(k\) 时,传统融合思路主张取 \(K > k\) 并通过加权平均聚合专家权重。然而实验发现,非同源训练的专家在参数空间中存在复杂的非线性干涉,直接参数平均不仅破坏激活流形,反而在蒸馏初期引入难以消除的混淆噪声。因此,框架主选 \(K = k\) 的“纯剪枝(Pure Pruning)”配置,令每个槽位独占一个完整特化专家。

由于现代 MoE 架构(如 Qwen3、DeepSeek、GPT-OSS)通常设置单专家中间维度 \(d_{\text{expert}}\) 满足 \(k \times d_{\text{expert}} \approx (3 \sim 5) \times d\),正好落入标准密集 Transformer 的 FFN 隐藏维度区间 \(d_{\text{dense}}\)。对于 SwiGLU 专家结构,将入选的 \(k\) 个专家的上投影和门控投影按行拼接,下投影矩阵按列拼接:

\[\mathbf{W}_{\text{gate}} = \begin{bmatrix} \mathbf{W}_{\text{gate}}^{(1)} \\ \vdots \\ \mathbf{W}_{\text{gate}}^{(k)} \end{bmatrix} \in \mathbb{R}^{d_{\text{dense}} \times d}, \quad \mathbf{W}_{\text{up}} = \begin{bmatrix} \mathbf{W}_{\text{up}}^{(1)} \\ \vdots \\ \mathbf{W}_{\text{up}}^{(k)} \end{bmatrix} \in \mathbb{R}^{d_{\text{dense}} \times d}\]
\[\mathbf{W}_{\text{down}} = \begin{bmatrix} \tilde{\mathbf{W}}_{\text{down}}^{(1)}, \dots, \tilde{\mathbf{W}}_{\text{down}}^{(k)} \end{bmatrix} \in \mathbb{R}^{d \times d_{\text{dense}}}\]

这种分块结构能够严格无损地保留被选专家对输入状态的内部激活响应,消除了参数平均导致的特征失真。

3. 静态下投影幅度缩放(Static Down-Projection Magnitude Scaling):在去除门控路由后保全中间层激活尺度一致性

原 MoE 网络中各个专家的输出在残差连接前均乘以动态路由归一化权重 \(p_{\ell}^{(e)}(t)\)(各激活专家权重之和通常为 1)。密集模型舍弃了动态门控网络,若直接将 \(k\) 个专家的输出逐项相加,会导致送入残差连接的信号能量膨胀约 \(k\) 倍,引发深层数值溢出与梯度爆炸。

为此,对每个专家分块在下投影权重处乘以静态缩放常数 \(\alpha_g\):\(\tilde{\mathbf{W}}_{\text{down}}^{(g)} = \alpha_g \mathbf{W}_{\text{down}}^{(g)}\)。针对 top-\(k\) 经 Softmax 严格重归一化的模型(如 Qwen3、GPT-OSS),采用均匀缩放因子 \(\alpha_g = 1/k\);而针对无重归一化或具备常驻共享专家的模型(如 DeepSeek-V2-Lite),共享专家缩放系数固定为 1.0,路由专家则采用校准集上的平均条件概率进行匹配补偿,确保了初始化阶段隐藏层表征幅度的数值稳定性。

损失函数 / 训练策略

初始化后的密集学生模型通过知识蒸馏(KD)从原 MoE 教师处恢复语言与推理能力。训练目标采用标准的前向 KL 散度(Forward KL Divergence),迫使学生模型覆盖教师完整的输出概率分布形态:

\[\mathcal{L}_{\text{KD}} = \frac{1}{|\mathcal{T}|}\sum_{t\in\mathcal{T}} D_{\text{KL}}\!\left(p_{\text{teacher}}(\cdot \mid x_t) \parallel p_{\text{student}}(\cdot \mid x_t)\right)\]

在优化策略上,相比于计算密集层隐藏状态均方误差(Hidden MSE)或反向 KL 散度,纯 Logit-level 前向 KL 表现出最强的收敛性。此外,蒸馏过程中 MoE 教师模型直接保持原生 \(k\) 个专家的稀疏激活推理,使得前向传递算力显著低于密集教师模型(相比同参数量稠密教师实现 1.6 倍训练墙钟加速)。

实验关键数据

主实验

主实验在 Qwen3-30B-A3B(\(E=128, k=8\),激活参数 3.3B)上进行转换评估,基线包括主流 MoE 压缩指标在密集化下的复刻(MC-SMoE、HC-SMoE)、对同量级稠密模型 Qwen3-32B 实施剪枝与蒸馏(D2D Pruning)、随机 FFN 初始化以及从头随机训练,均在统一的 0.3B token 与 ~4B token 扩展蒸馏预算下评估下游评测集(Winogrande、HellaSwag、ARC-Easy、ARC-Challenge、MMLU)。

模型 / 压缩配置 蒸馏预算 Wino Hella ARC-E ARC-C MMLU 平均准确率 (%)
DO-ACP, \(K=8\) (本文) 0.3B tokens 57.0 41.1 57.4 29.9 31.7 43.41
SF \(\times\) AB (MC-SMoE 准则) 0.3B tokens 52.2 31.6 47.1 23.1 27.5 36.31
SF \(\times\) OC (HC-SMoE 准则) 0.3B tokens 53.1 33.1 49.3 24.7 27.3 37.52
D2D 密集剪枝 (Qwen3-32B \(\to\) 3.4B) 0.3B tokens 51.7 27.7 38.9 22.4 25.7 33.28
Random FFN + 教师 Attention 0.3B tokens 53.4 28.0 34.7 22.2 25.2 32.70
Random Init (从头训练) 0.3B tokens 51.2 25.3 28.9 22.4 23.0 30.15
DO-ACP, \(K=8\) (本文) ~4B tokens 63.1 60.3 75.6 45.4 46.1 58.10
SF, \(K=16\) (最佳频率融合) ~4B tokens 61.2 56.3 74.0 43.1 32.7 53.46
D2D 密集剪枝 (Qwen3-32B \(\to\) 3.4B) ~4B tokens 60.5 57.5 73.1 41.5 26.6 51.84
Random FFN + 教师 Attention ~4B tokens 54.4 45.4 66.0 34.2 27.1 45.44
教师模型 (Qwen3-30B-A3B) 上限参考 82.6 82.0 85.3 65.3 83.2 79.65
稠密教师 (Qwen3-32B) 上限参考 76.8 84.0 89.9 73.1 81.9 81.14

消融实验

消融实验聚焦于知识蒸馏目标函数选择,以及蒸馏阶段扩大 MoE 教师路由宽度 \(k'\)(暴露未激活专家知识)对密集学生收敛质量的影响。

消融维度 / 实验配置 Wino Hella ARC-E ARC-C MMLU 平均准确率 (%) 相对变化 (\(\Delta\))
Forward KL (默认基准) 57.0 41.1 57.4 29.9 31.7 43.41 —
混合损失 (Logit KL + Hidden MSE) 55.5 39.1 53.3 28.7 30.9 41.50 -1.91
Reverse KL 53.0 33.8 42.4 25.6 31.0 37.17 -6.24
教师路由宽度 \(k'=8\) (原生 \(1\times k\)) 57.0 41.1 57.4 29.9 31.7 43.41 —
教师路由宽度 \(k'=16\) (\(2\times k\)) 57.9 40.9 57.8 31.1 32.8 44.11 +0.70
教师路由宽度 \(k'=32\) (\(4\times k\)) 57.1 40.6 57.0 31.0 32.5 43.63 +0.22
教师路由宽度 \(k'=64\) (\(8\times k\)) 55.0 39.5 56.1 30.8 32.1 42.71 -0.70
教师路由宽度 \(k'=128\) (全专家激活) 55.6 38.3 54.5 29.0 31.6 41.78 -1.63

关键发现

  • 打分机制是决定性因素,分组策略影响甚微:打分准则带来的性能方差高达 5.7 个百分点(DO-ACP 达 42.64% 对比 PP 的 36.94%),而分组方式对准确率的影响仅约 1.0 个百分点。D-最优准则使选出专家子集的有效秩(Effective Rank)从 6.07 提升至 7.37,避免了功能同质化。
  • 多样性打分下“纯剪枝”系统性优于“权重融合”:当挑选专家数 \(K=k=8\) 时,完全规避权重平均直接拼接表现最佳(43.41%)。相比之下,仅当使用传统频率打分(选出的专家冗余度极高)时,参数融合才能产生正向收益。
  • MoE 转换密集模型显著优于密集到密集(D2D)剪枝:在完全匹配参数量与算力下,MoE-to-Dense 无论在短期(0.3B tokens,领先 10.1 pp)还是扩展训练(4B tokens,领先 6.3 pp)均全面压制 D2D 剪枝。此外,由于 MoE 教师单步计算仅激活 3B 参数,训练墙钟速度比 32B 密集教师快 1.6 倍(73s/step vs 116s/step)。
  • 知识密集型任务收益尤甚:质性错误分析显示,DO-ACP 在 MMLU 人文类学科上准确率达到 49.2%,远超 D2D 的 19.2% 与 SF 的 25.4%,证明多样性专家保留了更完整的常识与事实型先验。

亮点与洞察

  • 将子模集选择(D-Optimality)优雅引入 MoE 压缩:将连续参数空间的专家选择映射为离散子模优化问题,利用 Schur 补贪心算法在多项式时间内找到兼顾重要度与表征正交性的专家组合,数学推导完备且具备理论近似比界限。
  • 打破了“专家合并必胜”的技术惯性:指出了在 MoE 转 Dense 场景下,由于激活 FFN 宽度天生可直接由 \(k\) 个专家拼接填充,参数平均反而会污染特征流形;“高多样性纯剪枝”才是保留特化能力的最佳范式。
  • MoE 教师是训练密集学生更优越的算力源泉:揭示了 MoE 模型不仅自身具备高容量,其稀疏激活特性更是高吞吐蒸馏教师的天然绝配,颠覆了以往必须专门训练超大稠密模型作为蒸馏母体的工业常规。

局限与展望

  • 合并权重与重要性打分的耦合:当 \(K > k\) 必须执行权重融合时,当前框架直接将选择得分作为融合权重,若解耦二者可能进一步释放融合潜能。
  • 蒸馏训练 token 规模仍有拓展空间:当前扩展实验仅验证至 4B tokens,而工业界完全收敛通常需要上百亿(10B+)甚至千亿 token 蒸馏,其渐近性能上限尚待建立。
  • 专家池容量对增益的敏感性:在小专家池模型(如 GPT-OSS-20B 仅 32 个专家)上,多样性选择相对于随机初始化的相对增益(+3.7 pp)明显小于大专家池模型 Qwen3(128 个专家,+13.3 pp),说明冗余空间越大的 MoE 从该方法中获益越显著。

相关工作与启发

  • vs REAP (Lasby et al., 2025):REAP 提出了门控权重与输出模长乘积的打分准则,但仅作独立标量排序,忽略专家间的表征共线性且最终输出仍为 MoE;本文将其因子化扩展后与 Gram 对数行列式结合,并首次落地于向完全密集架构的转换。
  • vs MC-SMoE / HC-SMoE (Li et al., 2024; Chen et al., 2025):此类方法基于选择频率聚集并平均专家,保留了稀疏路由开销且无法消除显存驻留瓶颈;本文证明在稠密转换场景下,其所依赖的频率打分落后 DO-ACP 达 5.9~7.1 pp。
  • vs Minitron (Muralidharan et al., 2024):Minitron 采用密集到密集的宽度/深度剪枝加蒸馏,但在相同参数量下其表征恢复能力大幅弱于 MoE-to-Dense,且其稠密教师训练吞吐落后 1.6 倍。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出将 MoE 转化为标准全密集大模型的系统化流水线,独创基于 D-最优 Gram 对数行列式的多样性感知专家选择。
  • 实验充分度: ⭐⭐⭐⭐⭐ 评估了 350 种配置组合,跨 Qwen3、DeepSeek-V2-Lite、GPT-OSS 三大主流架构验证,包含 4B token 扩展蒸馏与详细的质性错误分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论证明、实验对比与图表展示清晰严谨,深入分析了为什么“多样性纯剪枝”优于“冗余专家融合”。
  • 价值: ⭐⭐⭐⭐⭐ 具有极强的工业落地价值,为解决 MoE 模型显存驻留痛点、低成本高效率地产出高性能端侧密集模型提供了标准范本。