Condense, Don't Just Prune: Enhancing Efficiency and Performance in MoE Layer Pruning
作者: Mingyu Cao, Gen Li, Jie Ji, Jiaqi Zhang, Xiaolong Ma, Shiwei Liu, Lu Yin
分类: cs.LG, cs.CL
发布日期: 2024-11-26 (更新: 2025-02-16)
🔗 代码/项目: GITHUB
💡 一句话要点
提出ConDense-MoE，通过压缩而非剪枝MoE层，提升效率与性能
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 混合专家模型 模型压缩 模型剪枝 大型语言模型 推理加速
📋 核心要点
- MoE模型内存需求大，直接移除MoE层会导致显著的性能下降，难以在实际应用中落地。
- ConDense-MoE通过将大型稀疏MoE层压缩为小型稠密层，减少内存占用，同时保持硬件友好性。
- 在DeepSeekMoE-16B上，CD-MoE在降低27.5%内存的同时，保持90%准确率，推理速度提升1.26倍。
📝 摘要（中文）
混合专家模型(MoE)因其在利用相同甚至更少活跃参数的情况下扩展神经网络的能力而备受关注。然而，MoE并未减轻网络的大量内存需求，这限制了它们在实际应用中的可行性，尤其是在大型语言模型(LLM)时代。现有工作探索了移除整个MoE层以减少内存的可能性，但性能下降仍然显著。本文提出了ConDense-MoE (CD-MoE)，它不是删除整个MoE层，而是将大型稀疏MoE层压缩成一个更小、更密集的层，其中只有少数专家为所有token激活，同时保持硬件友好性。我们的方法专门为具有共享专家的细粒度MoE设计，其中前馈网络被分成许多小专家，某些专家被隔离用作始终激活的共享专家，例如DeepSeekMoE和QwenMoE。实验结果表明了我们方法的有效性。具体而言，对于DeepSeekMoE-16B模型，我们的方法在保持90%平均准确率的同时，将内存使用量减少了27.5%，并将推理速度提高了1.26倍。此外，我们表明，通过对压缩层进行轻量级专家微调（仅在单个80G A100 GPU上使用5小时），我们可以成功恢复98%的原始性能。
🔬 方法详解
问题定义：现有MoE模型虽然参数效率高，但内存占用仍然很大，限制了其在资源受限环境中的应用。直接剪枝整个MoE层虽然能减少内存，但会导致严重的性能损失。因此，如何在减少MoE层内存占用的同时，尽可能保持模型性能是一个关键问题。
核心思路：ConDense-MoE的核心思想是“压缩而非剪枝”。与其直接移除整个MoE层，不如将其压缩成一个更小、更密集的层。通过减少激活专家的数量，降低计算和内存需求，同时保留MoE层的部分能力。这种方法旨在在内存效率和性能之间取得更好的平衡。
技术框架：ConDense-MoE主要针对具有共享专家的细粒度MoE结构，如DeepSeekMoE和QwenMoE。其流程包括：1) 确定需要压缩的MoE层；2) 分析各专家在不同token上的激活情况；3) 选择少量专家，确保它们能够覆盖尽可能多的token；4) 将原始MoE层替换为仅激活这些选定专家的更小、更密集的层；5) (可选) 对压缩后的层进行微调，以恢复性能。
关键创新：ConDense-MoE的关键创新在于它改变了MoE层优化的思路，从简单的“移除”转变为更精细的“压缩”。通过保留部分专家并使其稠密化，可以在减少内存占用的同时，更好地保留原始模型的知识。此外，针对具有共享专家的MoE结构进行优化，使其更具实用性。
关键设计：关键设计包括：1) 专家选择策略：选择哪些专家进行保留，需要考虑专家对不同token的覆盖率，以及专家之间的互补性。可以使用贪心算法或其他优化方法来选择最佳的专家子集。2) 压缩后的层结构：压缩后的层可以使用更少的专家，并且可以调整专家的大小，以进一步减少内存占用。3) 微调策略：对压缩后的层进行微调可以恢复性能。可以使用轻量级的微调方法，例如只微调压缩层的参数，或者使用较小的学习率。
🖼️ 关键图片
📊 实验亮点
ConDense-MoE在DeepSeekMoE-16B模型上取得了显著的成果。在保持90%平均准确率的情况下，内存使用量减少了27.5%，推理速度提高了1.26倍。通过在单个80G A100 GPU上进行5小时的轻量级专家微调，可以成功恢复98%的原始性能。这些结果表明，ConDense-MoE是一种有效的MoE层压缩方法。
🎯 应用场景
ConDense-MoE可应用于各种需要部署大型语言模型的场景，尤其是在资源受限的环境中，如移动设备、边缘计算设备等。通过降低模型内存占用，可以使这些设备能够运行更大规模的MoE模型，从而提升其智能水平。此外，该方法还可以用于加速模型推理，提高用户体验。
📄 摘要（原文）
Mixture-of-Experts (MoE) has garnered significant attention for its ability to scale up neural networks while utilizing the same or even fewer active parameters. However, MoE does not alleviate the massive memory requirements of networks, which limits their practicality in real-world applications, especially in the era of large language models (LLMs). While recent work explores the possibility of removing entire layers of MoE to reduce memory, the performance degradation is still notable. In this paper, we propose ConDense-MoE (CD-MoE), which, instead of dropping the entire MoE layer, condenses the large, sparse MoE layer into a smaller, denser layer with only a few experts activated for all tokens, while maintaining hardware friendliness. Our approach is specifically designed for fine-grained MoE with shared experts, where Feed-Forward Networks are split into many small experts, with certain experts isolated to serve as shared experts that are always activated, such as DeepSeekMoE and QwenMoE. We demonstrate the effectiveness of our method. Specifically, for the DeepSeekMoE-16B model, our approach maintains 90% of the average accuracy while reducing memory usage by 27.5% and increasing inference speed by 1.26 times. Moreover, we show that by applying lightweight expert fine-tuning -- only to the condensed layers -- and using 5 hours on a single 80G A100 GPU, we can successfully recover 98% of the original performance. Our code is available at: https://github.com/duterscmy/CD-MoE/tree/main.