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 用可微动态结构剪枝，把稠密 LLM 的 MLP 拆成专家、注意力做 top-K，原权重全程冻结。LLaMA-2 7B 只激活一半参数时 WikiText 困惑度 8.36，比 DISP-LLM 的 9.84 更低。
结构剪枝会把判定为不重要的通道、层、头直接删掉。删完容量就没了，补回来往往要再烧一轮接近继续预训练的微调。动态剪枝按输入决定哪些结构参与计算，但每个 token 的预算不固定，批处理和 prefill 很难对齐。
另一条路是把稠密模型改造成稀疏 Mixture-of-Experts（MoE，按 token 只激活一部分专家）。LLaMA-MoE、CMoE 这类已有做法先切专家再训路由，两段分离，还要喂十亿到数万亿 token。ToMoE 的主张更硬：稠密 MLP 里本来就藏着可拆的专家，路由和专家划分可以一次学完，原权重完全不动。
对象是 decoder-only 模型。MLP 沿中间维度切成 N 个专家（默认 8），每个 token 走 top-1。注意力沿 head 维度处理：Query/Key 做静态剪枝，并按 RoPE 的成对子空间共享掩码；Value/Output 做 token 级 top-K。所有 token 的 head 宽度因此一致，prefill 和 batch 服务可以并行。
可微离散用 Straight-Through Gumbel-Sigmoid 和 Gumbel-Softmax。一层 Bi-GRU 超网络生成跨层相关的专家嵌入，再投影成每层选择掩码。原权重冻结，只训路由、投影和超网络，1 到 4 张 A100，固定 1 万步。蒸馏目标是稠密模型与转换后模型 logits 的 KL，teacher 就是关掉 ToMoE 模块后的同一份权重，不必另载一份模型。
三项正则卡住结构。专家并集尽量覆盖原层全部位置，避免大量权重永远没人用；按最宽专家卡死激活量上限；再加 Switch Transformer 式负载均衡。训完超网络和 MLP 投影可丢掉。LLaMA-2 7B 上额外参数约 0.0184B，占原模型 0.27%。校准数据用 WikiText、Alpaca、Code-Alpaca 各三分之一；只用 WikiText 时零样本会掉。
WikiText-2 上 LLaMA-2 7B 稠密困惑度 5.12。激活 70%/60%/50% 时 ToMoE 为 6.41/7.17/8.36，DISP-LLM 为 6.85/8.11/9.84，ModeGPT 为 7.51/8.41/11.88。半激活的 8.36 已经低于多数方法在 70% 激活时的数字。13B 同样：半激活 6.78，对 DISP-LLM 的 7.11。对 2:4 半结构剪枝，7B 上 SparseGPT 是 10.17，ToMoE 是 8.36。
零样本五任务均值（ARC-e/c、PIQA、WinoGrande、HellaSwag）：
| 模型 | 激活 | ToMoE | 最强对照 | 稠密 |
| 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 |
对改造类 MoE 的差距更大。LLaMA-2 7B 半激活：ToMoE 56.07，LLaMA-MoE E8A2 微调后 42.31。扩大到 BoolQ、SciQ 等、只用 0.02B token 时 ToMoE 均值 61.39，CMoE 微调 1.2B token 后是 51.68。吞吐在 batch 1536 时 2919 token/s，略高于 LLaMA-MoE 的 2888，稠密基线约 1858。
专家数从 8 加到 16 有收益，24 几乎不再涨。消融里整头剪枝把均值打到 44.99；去掉注意力 VO 动态路由只掉约 1 个点。
这是给已经训好的稠密模型做部署压缩的一条捷径：不碰原权重，训练成本与学习式剪枝同级，系统侧可以走现成 MoE 推理栈。适合要从 7B/8B/13B 里抠一半激活量、又不想再烧一轮继续预训练的团队。它补的是剪枝留下的容量缺口，不是从零训一个 MoE。
对 LLaMA-3 8B 半激活、评完 MMLU 的更宽套件，ToMoE 均值 57.40，LLaMA-MoE-v2 用 7B token 训完是 59.61。token 少约 350 倍，分数略低。预算够去继续预训练，后一条路仍可能更高。
论文没有单独的局限节，缺口要自己看。Query/Key 不能按 token 动态剪：不同位置掩码一错位，注意力有效宽度就会小于预算 K。专家可视化里多数层对齐的是句法片段，不是语义主题，数学题里才出现数字专家。正则会把各专家宽度推向一致，容量上的分工被压扁。
评测以常识零样本为主，MMLU 只在 LLaMA-3 8B 的参数-精度曲线里出现。没有和 DeepSeekMoE 这类从零训的原生 MoE 按激活量对齐。所谓「无需微调」指原权重不动，路由本身仍要 1 万步。和稠密模型的绝对差距仍大：LLaMA-2 7B 半激活零样本 56.07，对稠密 69.00。大 batch 下吞吐优势才明显。