STUN: Structured-Then-Unstructured Pruning for Scalable MoE Pruning

会议: ACL 2025
arXiv: 2409.06211
代码: https://github.com/thnkinbtfly/STUN
领域: Model Compression
关键词: MoE剪枝, 结构化剪枝, 非结构化剪枝, 专家剪枝, 模型压缩

一句话总结

STUN 提出了结构化→非结构化的两阶段 MoE 剪枝方法，第一阶段用 O(1) 前向传播实现可扩展的专家级剪枝，第二阶段在剩余专家内做非结构化剪枝，在 480B 参数的 Snowflake Arctic 上以 40% 稀疏度几乎无性能损失。

研究背景与动机

1. 领域现状：MoE 通过稀疏激活减少推理成本，但参数总量巨大（Mixtral 56B、DBRX 132B、Arctic 480B）仍需大量 GPU 内存
2. 现有痛点：
3. 非结构化剪枝无法利用 MoE 天然的专家结构
4. 结构化剪枝（行/列级）破坏模型能力较大
5. 现有专家剪枝方法需要 \(O(k^n/\sqrt{n})\) 次前向传播穷举组合，不适用于 128+ 专家的大规模 MoE
6. 核心矛盾：专家级剪枝效果好但搜索空间指数增长，非结构化剪枝可扩展但效果差
7. 本文要解决什么：找到结构化和非结构化之间的最优插值，同时保证可扩展性
8. 切入角度：先用 O(1) 的方法做专家级结构化剪枝，再在剩余专家内做非结构化剪枝
9. 核心idea一句话：利用专家行为相似性做 O(1) 可扩展的专家剪枝，再叠加非结构化剪枝

方法详解

整体框架

MoE 模型 → 第一阶段：基于行为相似性的专家级结构化剪枝（O(1) 前向传播）→ 第二阶段：在剩余专家内做非结构化权重剪枝（如 Wanda/SparseGPT）→ 最终压缩模型。

关键设计

1. O(1) 可扩展专家剪枝:
2. 做什么：只需一次前向传播就能确定哪些专家可以剪掉
3. 核心思路：利用专家之间的行为相似性（基于激活频率和输出相似度）构建潜在结构，使得贪心逐个剪除决策能近似捕获联合剪枝效果

4. 设计动机：现有方法需要穷举专家组合，128 个专家时根本不可行

5. 两阶段插值 (Structured-Then-Unstructured):

6. 做什么：在专家剪枝后继续做细粒度权重剪枝
7. 核心思路：专家剪枝后的 MoE 对非结构化剪枝更鲁棒——因为冗余专家已被去除，剩余专家的重要性更均匀
8. 实验验证：纯结构化或纯非结构化都不如两者组合（图1中峰值在中间）

损失函数 / 训练策略

剪枝后无需微调（training-free）。整个过程在单张 H100 上两小时完成。

实验关键数据

主实验

模型	方法	稀疏度	GSM8K	说明
Snowflake Arctic (480B)	STUN	40%	几乎不降	单卡H100, 2小时
Mixtral-8x7B	纯非结构化	50%	大幅下降	现有方法失败
Mixtral-8x7B	纯结构化	50%	大幅下降	现有方法失败
Mixtral-8x7B	STUN	50%	最佳	插值点最优

关键发现

• 结构化和非结构化的最优插值显著优于两端（图1峰值在中间）
• O(1) 专家剪枝在 128 专家的 Arctic 上可行，而穷举方法不可能
• STUN 在生成任务（GSM8K）上表现尤其好——这是对抗剪枝的最难场景
• 40% 稀疏度几乎无损说明大型 MoE 确实存在大量冗余专家

亮点与洞察

• O(1) 可扩展专家剪枝解决了 MoE 剪枝的关键瓶颈——未来 MoE 专家数会继续增长
• "先结构化后非结构化"的两阶段范式可以推广到其他混合剪枝场景
• 在 480B Arctic 上单 H100 两小时完成剪枝，实用性极高

局限性 / 可改进方向

• 目前只在 MoE 上验证，dense 模型不适用
• 专家相似性度量可能对某些任务不够精确
• 40% 以上稀疏度的性能退化情况未详细分析

相关工作与启发

• vs Wanda/SparseGPT: 纯非结构化方法，STUN 在 MoE 上叠加结构化后效果更好
• vs MC-SMoE: 需要穷举专家组合，STUN 的 O(1) 方法更可扩展
• vs LLM Surgeon: 通用剪枝方法，STUN 专门利用 MoE 结构

评分

• 新颖性: ⭐⭐⭐⭐ 两阶段插值思路简洁有效，O(1)专家剪枝的可扩展性是关键贡献
• 实验充分度: ⭐⭐⭐⭐ 480B Arctic + Mixtral + 多种对比
• 写作质量: ⭐⭐⭐⭐ 动机图（图1）直观地展示了插值优势
• 价值: ⭐⭐⭐⭐⭐ 对大规模MoE部署有直接实用价值