DIVE into MoE: Diversity-Enhanced Reconstruction of Large Language Models from Dense into Mixture-of-Experts¶
会议: ACL 2025
arXiv: 2506.09351
代码: https://github.com/yuchenblah/DIVE
领域: 模型压缩
关键词: Mixture-of-Experts, 结构化剪枝, Dense-to-MoE转换, 专家多样性, 领域亲和性挖掘
一句话总结¶
提出 DIVE,一种将 Dense LLM 重构为 MoE 架构的方法,核心洞察是不同领域的校准数据集会让结构化剪枝产生不同的剪枝结果,利用这种多样性构建领域特异的专家,配合高效的两阶段重训练(router dense训练 + expert LoRA稀疏训练),在仅调不到 1% 参数的情况下实现优于现有剪枝和 MoE 重构方法的效果。
研究背景与动机¶
- 领域现状:MoE 架构通过稀疏激活实现高效推理(如 Mixtral 8x7B),但 MoE 模型通常从头预训练,成本巨大。近年来 Dense-to-MoE 的转换方法受到关注——从已有的 Dense LLM 出发构建 MoE 模型,大幅节省训练预算。
- 现有痛点:
- 复制法(如 Sparse Upcycling):直接复制 FFN 作为多个专家,初始专家完全相同,缺乏多样性,需大量重训练才能分化
- 随机划分法(如 LLaMA-MoE):随机将 FFN 的中间维度切分为多个专家,但无法保证有意义的功能分化
- 两类方法都忽略了专家多样性,导致冗余问题
- 核心矛盾:Dense-to-MoE 转换的关键挑战在于如何在构建阶段就赋予专家初始多样性,而不是完全依赖后续训练来分化。
- 本文要解决什么?
- 如何利用 Dense LLM 自身的知识分布来构建具有初始多样性的 MoE 专家
- 如何设计高效的重训练流程,使转换后的模型快速恢复性能
- 切入角度:作者观察到结构化剪枝方法(如 FLAP)对校准数据集非常敏感——用数学数据做校准的剪枝模型在数学任务上表现好但在其他任务上差,不同领域的校准产生的剪枝 mask 截然不同。这个通常被认为是剪枝的"缺点"恰好可以用来构建多样化的专家。
- 核心 idea 一句话:利用结构化剪枝对校准数据的敏感性作为"特征"而非"缺陷",用不同领域的校准数据分别剪枝同一模型,得到领域特异的 FFN 子网络作为 MoE 专家。
方法详解¶
整体框架¶
Dense LLM \(\mathcal{M}\) → 领域亲和性挖掘(在 24 个数据集上交叉评估剪枝结果)→ 层次聚类(将数据集分成 N 个领域簇)→ 基于剪枝的专家重构(用每个簇的数据做校准分别剪枝 FFN,得到 N 个专家)→ 组装 MoE(FFN 替换为 MoE 层 + 随机初始化 router)→ 两阶段重训练(Stage 1: 全激活训练 router,Stage 2: 稀疏激活 LoRA 训练专家 + LayerNorm)→ 最终 MoE LLM
关键设计¶
- 领域亲和性挖掘 (Domain Affinity Mining):
- 做什么:发现数据集之间的领域关联,为专家构建提供分组依据
- 核心思路:
- 对 24 个数据集两两组合:用数据集 \(t_i\) 做校准剪枝,在数据集 \(t_j\) 上评估 → 得到 24×24 的 PPL 矩阵
- 用归一化 PPL \(\text{norm}(p)_{i,j} = \frac{\min(p)_{:,j}}{p_{i,j}}\) 消除不同评估集的量纲差异
- 用 Pearson 相关系数计算数据集间的关联度
- 层次聚类分为 N 组(如 8 组对应 8 个专家)
- 设计动机:不同领域的校准数据让 FLAP 剪枝保留不同的 FFN 通道——数学数据倾向保留数学相关的通道,NLI 数据保留推理相关的通道。这种敏感性是构建多样专家的天然工具
-
关键发现:
- 相似领域的校准集表现高度相关(MathQA ↔ GSM8K)
- 亲和性受数据来源影响大于任务类型(QNLI ≈ SQuAD2 > QNLI ≈ ANLI)
- C4(通用语料)作为校准集展现广泛泛化能力
-
基于剪枝的专家重构:
- 做什么:将 Dense FFN 剪枝为多个领域特异的瘦 FFN 作为 MoE 专家
- 核心思路:
- 使用 FLAP 方法,基于隐状态特征的波动方差评估通道重要性:\(\mathbf{S}_{:,j}^\ell = \frac{1}{N-1}\sum_{n=1}^{N}(\mathbf{X}_{n,j}^\ell - \bar{\mathbf{X}}_{:,j}^\ell)^2 \cdot \|\mathbf{W}_{:,j}^\ell\|_2^2\)
- 对每个领域簇,用该簇的数据做校准,剪枝 FFN 的中间维度(如 50% 或 75%)
- 8 个簇 → 剪枝得到 8 个 FFN → 组装为 MoE 层(加随机初始化的 noisy router)
- 除 FFN 外,embedding、LM head、MHA 模块保持不变
-
关键参数:记为 \(N/K\) 模型(如 1/8 = 激活 1/8 个专家,2/8 = 激活 2/8 个专家)
-
高效两阶段重训练:
- 做什么:用最少参数恢复重构后模型的性能
- 核心思路:
- Stage 1 (Dense训练):激活所有专家,只训练 router(引入温度系数 \(t\) 的 softmax),用少量数据(0.5B tokens)
- Stage 2 (Sparse训练):仅激活 top-k 专家,用 LoRA 训练专家 + 全量训练 LayerNorm,5B tokens
- MHA 模块完全冻结不训练(消融实验证明不需要)
- 设计动机:Stage 1 先让 router 学会分配,Stage 2 再让专家在稀疏激活下微调。不到 1% 的参数需要更新
损失函数 / 训练策略¶
- 标准语言建模目标(next token prediction)
- Stage 1: router 温度系数 \(t\)=0.05 (1/8) 或 0.5 (2/8),使 softmax 更 sharp 来模拟稀疏推理
- 重训练数据:SlimPajama(627B tokens),随机采样 5.5B tokens
- LoRA 应用于专家 FFN
实验关键数据¶
主实验(TinyLlama-1.1B,50% FFN激活)¶
| 方法 | 类型 | WikiText2 PPL ↓ | LAMBADA PPL ↓ | 下游任务平均 Acc ↑ |
|---|---|---|---|---|
| TinyLlama-1.1B | Dense 100% | — | — | 47.18 |
| LLM-Pruner | 剪枝 50% | 17.59 | 56.66 | 39.57 |
| FLAP | 剪枝 50% | 14.51 | 33.22 | 41.42 |
| LLaMA-MoE 1/8 | MoE 50%×1 | 19.57 | 87.27 | 39.34 |
| DIVE 1/8 | MoE 50%×1 | 13.52 | 24.84 | 42.17 |
DIVE 1/8 在 WikiText2 上比次优的 FLAP 低 0.99 PPL,LAMBADA 上低 8.38 PPL,下游任务平均准确率高 0.75%。
消融实验¶
| 配置 | WikiText2 PPL | LAMBADA PPL | ID任务 Acc | OOD任务 Acc |
|---|---|---|---|---|
| DIVE 2/8 (完整) | 18.09 | 63.45 | 47.88 | 27.67 |
| w/o 领域亲和性挖掘 | 20.02 (+1.93) | 76.03 (+12.58) | 47.08 (-0.80) | 26.15 (-1.52) |
| MHA训练 | DIVE 1/8 PPL变化 | DIVE 2/8 PPL变化 |
|---|---|---|
| 包含MHA | baseline | baseline |
| 不含MHA | <0.29 退化 | 反而提升 0.37 |
关键发现¶
- 领域亲和性挖掘至关重要:去掉 D.A.M. 后 LAMBADA PPL 恶化 12.58,证明多样性专家初始化的价值
- DIVE 仅用 1B tokens 就超过 FLAP 用 5B tokens 的表现,体现了高效性
- 路由分布验证了专家多样性:数学 tokens 正确路由到数学专家,NLI tokens 路由到 NLI 专家,C4 tokens 均匀激活所有专家
- MHA 模块不需要重训练,仅训练 router + expert LoRA + LayerNorm(<1% 参数)即可
- 高频分叉 token 类型与领域高度对应,说明 router 成功学到了领域分配
亮点与洞察¶
- "缺点变优点"的思路最为巧妙:结构化剪枝对校准数据的敏感性通常被视为弱点,但 DIVE 恰好利用这个特性来构建多样化专家。这种逆向思维值得借鉴——方法的局限性可能是另一个问题的解决方案
- 两阶段重训练设计很实用:先训 router 再训 expert,解耦了路由学习和专家适应两个问题。Stage 1 的温度系数技巧(让 dense training 模拟 sparse inference)是值得复用的技术细节
- 领域亲和性挖掘的系统性:用 24×24 的交叉评估矩阵量化了数据集间的关系,这个分析框架可以迁移到多任务学习、课程学习等场景
局限性 / 可改进方向¶
- 仅在 TinyLlama-1.1B 上验证,更大模型(7B+)的效果未知
- 领域亲和性挖掘需要 24×24 次剪枝+评估的前置开销,不适合即时使用
- 24 个数据集的选择和 8 个专家的数量是硬编码的,如何自动确定最优专家数和领域划分?
- 重训练仍需 5B tokens 的通用语料,对资源受限场景仍有门槛
- 未探索与 attention 层 MoE 化的结合(如 Switch Attention)
相关工作与启发¶
- vs LLaMA-MoE:LLaMA-MoE 用随机划分构建专家,缺乏多样性初始化;DIVE 在 PPL 和下游任务上全面领先
- vs Sparse Upcycling:复制法的专家初始完全相同,需要更长训练分化;DIVE 从一开始就有多样性
- vs FLAP (纯剪枝):FLAP 无法利用 MoE 的稀疏激活优势;DIVE 将剪枝与 MoE 结合,在相同激活参数下更优
评分¶
- 新颖性: ⭐⭐⭐⭐ 利用剪枝敏感性构建多样专家的 idea 有原创性,整体框架设计合理
- 实验充分度: ⭐⭐⭐⭐ 11 个下游任务、详细消融、路由分布分析,但仅 1.1B 模型稍显不足
- 写作质量: ⭐⭐⭐⭐ 图表设计出色(热力图清晰展示领域亲和性),算法描述完整
- 价值: ⭐⭐⭐⭐ 提供了一种高效的 Dense-to-MoE 转换方案,<1%参数微调的效率令人印象深刻