MoE-Pruner: Pruning Mixture-of-Experts Large Language Model using the Hints from Its Router (2410.12013v1)

Abstract: Mixture-of-Experts (MoE) architectures face challenges such as high memory consumption and redundancy in experts. Pruning MoE can reduce network weights while maintaining model performance. Motivated by the recent observation of emergent large magnitude features in LLMs (LLM) and MoE routing policy, we propose MoE-Pruner, a method that prunes weights with the smallest magnitudes multiplied by the corresponding input activations and router weights, on each output neuron. Our pruning method is one-shot, requiring no retraining or weight updates. We evaluate our method on Mixtral-8x7B and Mixtral-8x22B across multiple language benchmarks. Experimental results show that our pruning method significantly outperforms state-of-the-art LLM pruning methods. Furthermore, our pruned MoE models can benefit from a pretrained teacher model through expert-wise knowledge distillation, improving performance post-pruning. Experimental results demonstrate that the Mixtral-8x7B model with 50% sparsity maintains 99% of the performance of the original model after the expert-wise knowledge distillation.

• The paper introduces MoE-Pruner, a one-shot pruning strategy that uses router-informed metrics to identify and remove redundant weights in MoE models.
• The method achieves up to 50% sparsity in the Mixtral-8x7B model while maintaining 99% of its original performance through expert-wise knowledge distillation.
• This approach significantly reduces computational overhead and enhances deployment efficiency across nine language benchmarks without requiring retraining.

Overview of "MoE-Pruner: Pruning Mixture-of-Experts LLM using the Hints from Its Router"

The paper "MoE-Pruner: Pruning Mixture-of-Experts LLM using the Hints from Its Router" investigates techniques for compressing Mixture-of-Experts (MoE) architectures, which are often associated with high memory usage and redundancy. By introducing MoE-Pruner, the authors aim to enhance efficiency by reducing weights without significantly compromising the model's performance. This work offers a compelling approach to addressing the inefficiencies in MoE architectures, especially in the context of managing large-scale LLMs.

Contributions of the Paper

The proposed pruning strategy, MoE-Pruner, leverages a novel metric based on the magnitude of weights, input activations, and router weights. This method distinguishes itself as a one-shot pruning technique that circumvents the need for any form of retraining or weight updating, thus saving on computational resources. Furthermore, the authors enhance the MoE model's post-pruning performance through expert-wise knowledge distillation from a larger pre-trained teacher model. The results confirm that with up to 50% sparsity, their method allows the Mixtral-8x7B model to maintain 99% of its original performance post-distillation.

Key Numerical Results

The numerical evaluation of MoE-Pruner emphasizes its superiority over existing pruning methods such as SparseGPT and Wanda. The paper demonstrates that the 50% sparsely pruned Mixtral-8x7B model maintains a level of performance very close to the original, evidencing a minimal drop-off when subjected to their expert-wise knowledge distillation process. These experiments, conducted across nine language benchmarks, signify a substantive advancement in achieving an effective balance between computational efficiency and model integrity in MoE architectures.

Pruning Metrics and Techniques

The paper outlines that traditional methods such as SparseGPT and Wanda consider heuristics like magnitude or inverse Hessian information for pruning. In contrast, MoE-Pruner employs router weight-informed magnitude metrics, providing a more nuanced understanding of weight importance. The computational efficiency of MoE-Pruner, achieved through simplified complexity compared to SparseGPT, also facilitates faster deployment and iteration on large-scale MoE architectures.

Implications and Speculations on Future Research

The results and methodologies of this paper have substantial implications. Practically, the research presents a pathway for effectively reducing memory and computation overheads in LLMs, which could foster broader accessibility and deployment of sophisticated linguistic models in resource-constrained environments. Theoretically, the observations around router weight utilization suggest further exploration into adaptive model architectures that dynamically choose expert pathways, optimizing both efficiency and performance.

Future research could explore extending MoE-Pruner's capabilities to structured sparsity types, such as channel or expert pruning, which could offer additional hardware acceleration benefits. The exploration of more complex knowledge distillation strategies, considering not just expert layers but also interactions across layers, could propel additional gains in model compactness and efficacy.

In summary, this paper contributes significantly to the discourse on optimizing MoE architectures, presenting a robust method that enhances the practicality and scalability of LLMs. The work highlights the potential of efficient pruning techniques as pivotal in the evolution of AI architectures, aiming for greater utility without compromising on the capabilities of expansive LLMs.