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

📄 arXiv: 2409.06211v2 📥 PDF

作者: Jaeseong Lee, seung-won hwang, Aurick Qiao, Daniel F Campos, Zhewei Yao, Yuxiong He

分类: cs.LG, cs.CL

发布日期: 2024-09-10 (更新: 2025-07-21)

备注: ACL 2025 main

💡 一句话要点

提出结构化后无结构化剪枝以提升MoE模型的可扩展性

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 剪枝技术 深度学习 模型压缩 推理优化 结构化剪枝 无结构化剪枝

📋 核心要点

  1. 现有的专家剪枝方法在处理大规模MoE时面临计算复杂度高的问题，无法有效扩展。
  2. 本文提出了一种新的剪枝方法，先进行结构化的专家剪枝，再进行无结构化剪枝，以提高整体性能。
  3. 在Snowflake Arctic模型上，采用该方法仅需一台H100显卡和两小时即可实现40%稀疏度，几乎无性能损失。

📝 摘要（中文）

混合专家模型（MoE）通过稀疏激活专家来降低大型语言模型的推理成本。然而，专家数量庞大仍使得其服务成本高昂。本文研究了如何通过剪枝来解决这一问题。尽管无结构化剪枝在给定剪枝比例下性能最佳，但我们发现，专家剪枝作为一种结构化剪枝形式，实际上可以在无结构化剪枝之前进行，从而超越仅采用无结构化剪枝的方法。我们提出了一种可扩展的剪枝方法，其复杂度为O(1)，且在性能上优于现有的高成本方法。实验表明，在Snowflake Arctic模型上，我们的方法在40%稀疏度下几乎没有性能损失，尤其在生成任务中表现优异。

🔬 方法详解

问题定义：本文旨在解决混合专家模型（MoE）在推理时的高成本问题，现有的专家剪枝方法由于计算复杂度高，无法有效扩展到大规模模型。

核心思路：我们提出了一种新的剪枝策略，先进行结构化的专家剪枝，再进行无结构化剪枝。通过利用专家之间的潜在结构和行为相似性，优化剪枝决策，从而提高剪枝效果。

技术框架：整体方法包括两个主要阶段：第一阶段是专家剪枝，利用行为相似性进行结构化剪枝；第二阶段是无结构化剪枝，进一步优化剩余专家的选择。

关键创新：本文的创新在于提出了先进行结构化剪枝的策略，挑战了传统无结构化剪枝的主导地位，显示出结构化剪枝在某些情况下的优势。

关键设计：在剪枝过程中，我们设计了基于行为相似性的贪婪决策机制，确保剪枝决策能够有效捕捉到联合剪枝的效果，且在复杂度上优化至O(1)。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在Snowflake Arctic模型上，我们的方法在达到40%稀疏度的情况下，仅需两小时的训练时间，几乎没有性能损失。相比于传统的无结构化剪枝方法，我们的策略在生成任务中表现出色，显示出显著的性能提升。

🎯 应用场景

该研究具有广泛的应用潜力，尤其在大型语言模型和深度学习系统中，通过有效的剪枝策略，可以显著降低推理成本，提高模型的可扩展性。这将推动AI在实时应用和资源受限环境中的部署，具有重要的实际价值和未来影响。

📄 摘要（原文）

Mixture-of-experts (MoEs) have been adopted for reducing inference costs by sparsely activating experts in Large language models (LLMs). Despite this reduction, the massive number of experts in MoEs still makes them expensive to serve. In this paper, we study how to address this, by pruning MoEs. Among pruning methodologies, unstructured pruning has been known to achieve the highest performance for a given pruning ratio, compared to structured pruning, since the latter imposes constraints on the sparsification structure. This is intuitive, as the solution space of unstructured pruning subsumes that of structured pruning. However, our counterintuitive finding reveals that expert pruning, a form of structured pruning, can actually precede unstructured pruning to outperform unstructured-only pruning. As existing expert pruning, requiring $O(\frac{k^n}{\sqrt{n}})$ forward passes for $n$ experts, cannot scale for recent MoEs, we propose a scalable alternative with $O(1)$ complexity, yet outperforming the more expensive methods. The key idea is leveraging a latent structure between experts, based on behavior similarity, such that the greedy decision of whether to prune closely captures the joint pruning effect. Ours is highly effective -- for Snowflake Arctic, a 480B-sized MoE with 128 experts, our method needs only one H100 and two hours to achieve nearly no loss in performance with 40% sparsity, even in generative tasks such as GSM8K, where state-of-the-art unstructured pruning fails to. The code will be made publicly available.