[论文解读] Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy

Pingzhi Li, Zhenyu Zhang+5 more|arXiv (Cornell University)|Oct 2, 2023

Domain Adaptation and Few-Shot LearningComputer Science被引用 3

一句话总结

本文提出 MC-SMoE，一种新颖的框架，首先通过路由策略引导的专家整合来合并稀疏 Mixture-of-Experts（SMoE），然后利用低秩和结构稀疏性对合并后的专家进行压缩。通过利用激活频率和路由相似性，M-SMoE 减少了冗余和内存占用，在 8 个基准测试中实现了高达 80% 的内存减少和 20% 的 FLOPs 降低，且性能损失可忽略不计。

ABSTRACT

Sparsely activated Mixture-of-Experts (SMoE) has shown promise to scale up the learning capacity of neural networks, however, they have issues like (a) High Memory Usage, due to duplication of the network layers into multiple copies as experts; and (b) Redundancy in Experts, as common learning-based routing policies suffer from representational collapse. Therefore, vanilla SMoE models are memory inefficient and non-scalable, especially for resource-constrained downstream scenarios. In this paper, we ask: Can we craft a compact SMoE model by consolidating expert information? What is the best recipe to merge multiple experts into fewer but more knowledgeable experts? Our pilot investigation reveals that conventional model merging methods fail to be effective in such expert merging for SMoE. The potential reasons are: (1) redundant information overshadows critical experts; (2) appropriate neuron permutation for each expert is missing to bring all of them in alignment. To address this, we propose M-SMoE, which leverages routing statistics to guide expert merging. Specifically, it starts with neuron permutation alignment for experts; then, dominant experts and their "group members" are formed; lastly, every expert group is merged into a single expert by utilizing each expert's activation frequency as their weight for merging, thus diminishing the impact of insignificant experts. Moreover, we observed that our proposed merging promotes a low dimensionality in the merged expert's weight space, naturally paving the way for additional compression. Hence, our final method, MC-SMoE (i.e., Merge, then Compress SMoE), further decomposes the merged experts into low-rank and structural sparse alternatives. Extensive experiments across 8 benchmarks validate the effectiveness of MC-SMoE. For instance, our MC-SMoE achieves up to 80% memory and a 20% FLOPs reduction, with virtually no loss in performance.

研究动机与目标

• 为解决稀疏 Mixture-of-Experts（SMoE）模型中高内存占用和专家冗余的问题，特别是在资源受限的部署场景中。
• 探究是否可以将专家信息有效整合为更少、更具知识性的专家，而不会导致性能下降。
• 开发一种在 SMoE 设置下优于传统模型合并技术的合并策略。
• 探索合并后的专家是否天然具有更高的可压缩性，从而实现进一步的效率提升。

提出的方法

• M-SMoE 利用路由统计信息指导专家合并：首先通过排列对齐专家之间的神经元，以确保特征对应。
• 基于路由策略相似性识别主导专家及其‘组内成员’，形成专家组。
• 利用激活频率作为加权平均值，对每组内的专家进行合并，从而降低次要专家的影响。
• 合并后的专家被发现在更低维的权重空间中，从而实现进一步压缩。
• MC-SMoE 对合并后的专家应用低秩和结构稀疏性，实现额外的参数和 FLOPs 减少。
• 该方法在 8 个 NLP 基准测试上，使用 switch-base-32 SMoE 模型进行了端到端评估。

实验结果

研究问题

• RQ1 在 SMoE 中，是否可以将专家信息有效整合为更少、更具知识性的专家，而不会造成性能损失？
• RQ2 为何传统模型合并方法在 SMoE 设置下会失效，其关键失败模式是什么？
• RQ3 如何利用路由策略指导专家合并，并识别冗余专家与主导专家？
• RQ4 合并专家是否自然导致更低维的权重空间，从而实现进一步压缩？
• RQ5 在实现内存和 FLOPs 效率方面，合并与压缩的相对贡献分别是什么？

主要发现

• MC-SMoE 在 switch-base-32 SMoE 模型上实现了高达 80% 的内存减少和 20% 的 FLOPs 减少，且在 8 个基准测试中性能损失可忽略不计。
• M-SMoE 在所有任务上均优于 C-SMoE（仅压缩），证明在保持性能方面，合并比单纯压缩更有效。
• 基于频率的加权合并始终优于均匀加权和 Fisher 加权合并，在所有任务上平均提升 1.5–2% 的准确率。
• M-SMoE 合并后的专家表现出更低的内在维度，使得通过低秩和结构稀疏性进行下游压缩更加有效。
• 后期 SMoE 层表现出更高的可压缩性，这从热力图可视化中剩余参数的分布可得。
• 消融实验证实，基于频率感知的合并能够抑制冗余专家，防止关键知识被掩盖。