Unveiling Hidden Collaboration within Mixture-of-Experts in Large Language Models

TL;DR

The paper tackles hidden collaboration among experts in Mixture-of-Experts LLMs by extracting an expert activation matrix $X$ and applying Hierarchical Sparse Dictionary Learning (HSDL) to mine cross-layer patterns, followed by a Contribtion-Aware Expert Pruning (CAEP) to compress the model. It formalizes the mining task as $X ≈ DR$ with $D ∈ R^{N_e×N_p}$ and $R ∈ R^{N_p×N_s}$, then extends to a hierarchical decomposition $D_k ≈ D_{k+1} R_{k+1}$ with sparsity, inter-layer consistency, and reconstruction constraints, optimized via $L_{total} = L_{sparse} + λ_1 L_{hier} + λ_2 L_{rec}$. Empirically, the approach uncovers domain-specific collaboration patterns on MMLU-pro, achieves an average improvement of about 2.5% after pruning, and can prune up to 50% of experts with modest performance loss in many tasks. This work enhances interpretability and efficiency of MoE LLMs and suggests domain-aware pruning strategies for scalable deployment.

Abstract

Mixture-of-Experts based large language models (MoE LLMs) have shown significant promise in multitask adaptability by dynamically routing inputs to specialized experts. Despite their success, the collaborative mechanisms among experts are still not well understood, limiting both the interpretability and optimization of these models. In this paper, we focus on two critical issues: (1) identifying expert collaboration patterns, and (2) optimizing MoE LLMs through expert pruning. To address the first issue, we propose a hierarchical sparse dictionary learning (HSDL) method that uncovers the collaboration patterns among experts. For the second issue, we introduce the Contribution-Aware Expert Pruning (CAEP) algorithm, which effectively prunes low-contribution experts. Our extensive experiments demonstrate that expert collaboration patterns are closely linked to specific input types and exhibit semantic significance across various tasks. Moreover, pruning experiments show that our approach improves overall performance by 2.5\% on average, outperforming existing methods. These findings offer valuable insights into enhancing the efficiency and interpretability of MoE LLMs, offering a clearer understanding of expert interactions and improving model optimization.

Figures (9)

• Figure 1: In MoE LLMs, a group of experts often collaborate to analyze a certain type of tokens, and they are not necessarily in the same layer.
• Figure 2: Here ($x,y$) refers to the $y$-th expert in $x$-th layer. By selecting any two experts from the MoE, we can calculate the probability of their co-activation. It can be observed that Expert 21 from the layer 5 and Expert 3 from the layer6 frequently activate simultaneously, forming an expert collaboration pattern.
• Figure 3: Overview of Our Study's Pipeline.
• Figure 4: Hierarchical Sparse Dictionary Learning.
• Figure 5: Hierarchical Semantic Annotation of Dictionary Elements on MMLU.