Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy (2310.01334v2)

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.

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

This paper addresses the limitations associated with Sparsely activated Mixture-of-Experts (SMoE) models which, despite their potential to scale neural network capacity, suffer from inefficiencies in memory usage and expert redundancy. The paper poses a critical question about the possibility of creating a compact SMoE by judicious merging of expert information, and determining the best methodology for such a consolidation.

Key Contributions

The authors propose a novel algorithm termed as M-SMoE. This mechanism leverages the routing statistics within SMoEs to guide the merging of experts. The key steps in this process include:

• Neuron Permutation Alignment: Initiating the merger process with permutation alignment across neuron spaces of the experts, ensuring the neurons are accurately aligned.
• Group Formation and Expert Merging: The algorithm identifies dominant experts based on routing policies, forming groups with less significant "group members." This strategy allows for effective merging wherein each group is coalesced into a singular expert using a frequency-based weighting mechanism.
• Compression Beyond Merging: Upon merging, the authors observe a natural propensity for reduced dimensionality in the weight space of the resulting experts. Exploiting this feature, they introduce MC-SMoE (Merge, then Compress SMoE), which decomposes the consolidated experts using low-rank and structural sparseness methods to further amplify compression.

Experimental Validation

The paper provides empirical support across eight distinct benchmarks to validate the effectiveness of M-SMoE and MC-SMoE. For instance, results indicate that MC-SMoE achieves up to 80% memory savings and a 20% reduction in FLOPs with almost no performance loss. These strong numerical results demonstrate the method's efficacy in delivering memory and computational efficiency without sacrificing model performance.

Implications and Considerations

The research presents significant theoretical and practical implications. Theoretically, it offers a route to enhance the scalability of SMoE models via informed structural changes informed by underlying routing policies. Practically, the compression approach could become integral in resource-constrained scenarios, enabling deployment of larger, more capable models where computational resources are limited.

The finding related to dimensionality reduction post-merging prompting a secondary round of compression could pave the way for additional algorithmic innovations, particularly in tasks surrounding expert configuration and integration in SMoE architectures.

Speculations on Future Developments

Looking ahead, the approach of leveraging routing information could be expanded to other hybrid models beyond SMoE. Integrating such techniques with advancements in hardware specialization and parallelism could further optimize runtime efficiency, a current limitation due to routing and diverse expert implementation. The potential for inter-domain application - extending the algorithm from NLP to other fields like computer vision or multimodal processing - also presents an exciting avenue for future research.

In summary, the presented method of "Merge, Then Compress" for SMoE constitutes a profound step towards resolving key limitations in scaling mixture-of-expert systems. Its emphasis on efficiency and performance retention highlights an essential direction for developments in large-scale AI model deployment.