Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy

TL;DR

This paper addresses the high memory and redundancy issues of sparse Mixture-of-Experts (SMoE) by introducing MC-SMoE, a Merge-Then-Compress pipeline that uses routing statistics to guide expert merging and subsequent compression.The core innovation, M-SMoE, aligns expert permutations, groups similar experts with dominant representatives, and merges groups using frequency-weighted averaging informed by routing activation, yielding compact but effective experts.A further post-merge compression step (MC-SMoE) decomposes merged weights into a low-rank factorization plus a pruned sparse component, enabling substantial memory and FLOPs reductions with minimal performance loss across eight benchmarks.Empirical results show up to 80% memory savings and 20% FLOPs reductions for MC-SMoE, with robust performance on several tasks, and substantial memory savings even for M-SMoE, supported by extensive ablations and analyses of merging strategies, permutation alignment, and KD effects.

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.

Figures (13)

• Figure 1: Accuracy ($\%$) on the COPA with the switch-base-32 SMoE. MC-SMoE reaches up to an $80\%$ memory saving with only a negligible compromise in performance.
• Figure 2: The overview of our proposed MC-SMoE pipeline. (a) In the conventional SMoE, each token embedding is directed to a small number of relevant experts. (b) The routing policy inspires expert merging. Across all SMoE layers, M-SMoE identifies the most frequently activated experts as dominant ones, groups the other non-dominant experts, and then merges them within each group in a frequency-weighted fashion. (c) After merging, the weight space of resulted experts tends to exhibit lower dimensionality, paving the way for additional compression. It clarifies the design of our MC-SMoE.
• Figure 3: Distribution of expert activation frequencies in the switch-base-32 model, encompassing $12$ SMoE layers with $32$ experts per layer. The top of the heatmap is the first MoE layer while the bottom is the last. The left two tasks, COPA and SQuAD, are characterized by answer-generation prompts. The right two tasks, WikiQA and SST2, are typified by answer-selection prompts. SMoE models fine-tuned on answer-selection tasks demonstrate a more skewed distribution in their transformer decoder layers, wherein a significant portion of experts remain inactivated all the time.
• Figure 4: Experts are more compressible after merging. We calculate the average stable-rank change ratio ($\frac{\text{after}-\text{before}}{\text{before}}$) of all dominant experts within each layer of the switch-base-32 SMoE model, reflecting the difference before and after merging. These mostly negative values throughout the SMoE layers emphasize a lower dimensionality achieved through the merging process.
• Figure 5: Ratio of remaining parameters after further compressing the dominant experts from MC-SMoE.