# Diversifying the Expert Knowledge for Task-Agnostic Pruning
in Sparse Mixture-of-Experts

Zeliang Zhang1111This research was conducted during the author’s internship at
Microsoft Research. Xiaodong Liu2 Hao Cheng2 Chenliang Xu1 Jianfeng Gao2
1University of Rochester 2Microsoft Research
{zeliang.zhang<EMAIL_ADDRESS>{xiaodl, cheng.hao,
<EMAIL_ADDRESS>
###### Abstract

By increasing model parameters but activating them sparsely when performing a
task, the use of Mixture-of-Experts (MoE) architecture significantly improves
the performance of Large Language Models (LLMs) without increasing the
inference cost. However, the memory consumption due to the growing number of
experts presents a challenge to the deployment of these models in many real
world settings.

Our empirical study reveals that some experts encode redundant knowledge
during pre-training. We thus propose a method of grouping and pruning similar
experts to improve model’s parameter efficiency. We validate the effectiveness
of our method by pruning two state-of-the-art MoE models, Mixtral-8x7B and
Mixtral-8x22B. Evaluation shows that our method outperforms other model
pruning methods on a range of natural language tasks. To facilitate future
research, we will release our code and the pruned MoE models.

Diversifying the Expert Knowledge for Task-Agnostic Pruning
in Sparse Mixture-of-Experts

Zeliang Zhang1111This research was conducted during the author’s internship at
Microsoft Research. Xiaodong Liu2 Hao Cheng2 Chenliang Xu1 Jianfeng Gao2
1University of Rochester 2Microsoft Research {zeliang.zhang,
<EMAIL_ADDRESS>{xiaodl, cheng.hao<EMAIL_ADDRESS>

##  1 Introduction

Large Language Models (LLMs) have achieved outstanding performance across
various tasks by learning a large number of model parameters on large amounts
of data, as shown by the scaling laws [16]. In addition to increasing the
depth of neural network models, widening neural networks by using the
sparsely-activated mixture-of-experts (MoE) architecture is also proved
effective. MoE widens the feed-forward network (FFN) layer (one expert) by
having multiple parallel FFNs (experts). During forward propagation, only a
subset of these experts is activated. Thus, compared to dense models, MoE
models achieve better end-task performance and generalize better to new tasks
without increasing computation costs. Notable examples of MoE models include
Switch Transformer [9], Mixtral-MoE [14], and Uni-MoE [22].

Despite significant progress in developing wider and deeper MoE LLMs, the
increased memory consumption due to larger model sizes (i.e., increased number
of experts) poses a substantial challenge to the deployment of these models in
real-world settings. For example, storing and loading Mixtral-8x7B, which has
8 experts in each of its 32 layers, requires approximately $88$ GB. The MoE
layers constitute the majority of the parameters. Adding or removing even one
expert in each layer can significantly impact overall memory cost and model
performance. For example, [27] shows that randomly dropping $2$ experts in
each MoE layer reduces the memory cost by $21$ GB, and decreases model
performance by $7\%$ on the MMLU benchmark [12]. In this study, we strive to
seek the best trade-off between memory efficiency and task performance by
identifying an optimal set of experts in each MoE layer to prune.

There are many recent studies on pruning MoE models. One line of work utilizes
task-specific information to prune irrelevant experts. For example, Chen et
al. [3] prune the less frequently visited experts based on experiments on a
range of tasks. Chowdhury et al. [5] find that less important experts usually
exhibit smaller changes in routing weights during the fine-tuning stage. Li et
al. [21] merge experts that are frequently visited by tokens of a fine-tuned
dataset for pruning. Compared with task-specific pruning methods, task-
agnostic pruning is more challenging due to the lack of explicit criteria to
evaluate the redundancy of experts on unknown tasks. He et al. [11] explore
pruning experts with less visited frequency in a task-agnostic calibration
dataset but report a significant performance drop. In comparison, Lu et al.
[27] enumerate all the combinations of experts and prune some to achieve a
minimum loss of reconstruction, which significantly improves performance. We
illustrate the difference among these works in fig. 1. Although pruning MoE
models in task-agnostic settings is of great practical value, this area has
not been fully explored.

Figure 1: Removing several experts from the original MoE layer would not cause
model collapse but improve efficiency. Prior works often utilize expert access
information to filter out unimportant experts. In our work, we first group
different experts with similar knowledge in the feature space, then merge them
along with the routers to prune the MoE layer. This post-processing approach
allows us to diversify the features of each MoE layer, thereby preserving the
knowledge of the original large models as much as possible while reducing
computation and storage consumption.

In this work, we explore how to prune MoE models in a task-agnostic fashion.
Our study is motivated by the finding that, given the same input, many experts
respond similarly, indicating that these experts likely encode similar
knowledge, and thus are somewhat redundant. We propose a method to improve
model parameter efficiency by pruning redundant experts in two stages. As
shown in fig. 1, we first identify and group similar experts in the feature
space. Then, for each group, we merge experts in the weight space to diversify
the knowledge in different MoE layers. We validate the effectiveness of our
method by pruning experts for two state of the art MoE models, Mixtral-8x7B
and Mixtral-8x22B. Evaluation shows that our method outperforms other model
pruning methods on a range of natural language tasks.

##  2 Related Work

Sparse MoEs. The Mixture-of-Experts (MoE) structure is firstly applied in
classical machine learning models by Jacobs et al. [13] and Jordan and Jacobs
[15], then widely used in various deep learning models [38, 30, 39]. Recently,
some works employ the MoE to scale the capacity of transformer-based models,
especially the large language models [33, 19, 42]. It adapts the original
large feed-forward network (FFN) in each transformer block into multiple
smaller FFNs, forming an expert layer with a router that computes the weighted
output of each MoE layer. Sparse MoEs were first proposed by Fedus et al. [9].
In this approach, only a few experts are activated in each layer, accelerating
training and inference while significantly increasing the number of parameters
for greater model capacity. Many sparse MoEs have been developed and open-
sourced within the AI community, such as Switch Transformer [9], Mixtral-8B
[14], and Uni-MoE [22]. Recent studies also indicate that neural networks with
the MoE structure exhibit better generalization ability compared to dense
models [40, 20].

Model Pruning. Model pruning involves removing unimportant parameters from a
well-trained neural network to balance task performance and computational
efficiency [25, 36, 24]. Pruning techniques can be categorized into
unstructured pruning [23, 35, 29], which introduces sparsity in the weight
matrix by setting some parameters to zero, and structured pruning [18, 8, 34],
which removes entire neurons, layers, or blocks, reducing redundancy and being
more suitable for acceleration on GPUs [4]. Many efforts have been made to
leverage model pruning techniques to reduce the memory consumption of neural
networks, spanning a range of models from conventional architectures like CNNs
[28], RNNs [41], and LSTMs [7] to modern large models such as Llama [37] and
Stable-Diffusion [2].

While the large amount of parameters in sparse MoEs benefits the model’s
capacity to achieve good performance at the pre-training stage, the increasing
memory consumption causes great challenges to fine-tuning different downstream
tasks. In this paper, we work on pruning the sparse MoEs to reduce redundant
experts in the task-agnostic setting, which enhances the computational and
memory efficiency throughout the fine-tuning process, and scalability of
deploying these models.

##  3 Methodology

Notations. We denote $F(\cdot;\Theta,W,K)$ as an MoE layer, where
$\Theta=\\{\theta_{1},\theta_{2},...,\theta_{N}\\}$ indicates the parameter
set for $N$ experts $\\{f_{n}(\cdot;\theta_{n})\\}_{n=1}^{N}$,
$W\in\mathbb{R}^{N\times d}$ indicates the routing weights for each expert and
$K$ indicates the number of experts selected for each token.

For a new token $x\in\mathbb{R}^{d}$, we first compute the routing logits to
score the match between $x$ and the $n$-th expert:

$\displaystyle p_{n}(x)=\frac{e^{W_{n}x}}{\sum_{t=1}^{N}e^{W_{t}x}},$ (1)

where $W_{t}\in\mathbb{R}^{d}$ is the routing policy of $t$-th expert. Then,
we take the experts $\\{f_{n}(\cdot;\theta_{n})\\}_{n=i_{1}}^{i_{K}}$ with
top-$K$ logits $\\{l_{n}\\}_{n=i_{1}}^{i_{K}}$ as the selected target experts.
Last, we compute the weighted average of the outputs from selected experts as
the final output of the MoE layer,

$\displaystyle y=\sum\limits_{n=i_{1}}^{i_{K}}p_{n}(x)\cdot
f_{n}(x;\theta_{n}).$ (2)

###  3.1 Task Definition

Several prior works [27, 11, 21] have identified that removing some experts in
different MoE layers does not incur model collapse but improves efficiency in
inference and fine-tuning on downstream tasks. The remaining question is which
expert(s) should be pruned in the MoE layer.

Specifically, for a MoE LLM $\mathcal{M}(\cdot;\mathcal{F})$, which consists
of multiple MoE layers
$\mathcal{F}=\\{F^{l}(\cdot;\Theta^{l},W^{l},K)\\}_{l=1}^{L}$, we want to
prune $r$ experts from the total $N$. This leads to the expert modules
$\Theta$ being reduced to
$\hat{\Theta}^{l}=\Theta^{l}\setminus\\{\theta^{l}_{s^{l}_{1}},\theta^{l}_{s^{l}_{2}},\ldots,\theta^{l}_{s^{l}_{r}}\\}$,
and the routing matrix $W$ being reduced to
$\hat{W}^{l}=W^{l}\setminus\\{W^{l}_{s^{l}_{1}},W^{l}_{s^{l}_{2}},\ldots,W^{l}_{s^{l}_{r}}\\}$,
where $s^{l}=\\{s^{l}_{1},s^{l}_{2},\ldots,s^{l}_{r}\\}$ is the set of indices
to delete. We aim to find the pruned $\hat{\mathcal{M}}$ by deleting the
optimal indices set $s^{l}$ in different MoE layers $l$ that satisfy the
following optimization objective:

$\displaystyle\min\limits_{(x,y)\sim\mathcal{D}}\mathcal{L}(\hat{\mathcal{M}}(x;\hat{\mathcal{F}}),y),$
(3)

where $\mathcal{L}$ is the loss function,
$\hat{\mathcal{M}}(\cdot;\hat{\mathcal{F}})$ is the pruned LLM by deleting $r$
experts from each MoE layer $F$ in $\mathcal{F}$, and $\mathcal{D}$ is the
dataset.

This task is particularly challenging as this is an NP-hard search problem
with $(C^{r}_{N})^{L}$ possible solutions. In this work, we propose to
estimate expert similarity in the feature space for greedy structured pruning
within each MoE layer to address this problem.

Figure 2: Evaluation of the expert similarity for different MoE layers in
Mixtral-8x7B under two kernel-based CKA criteria (Linear and RBF). A darker
color indicates a greater similarity between experts.

###  3.2 Expert Similarity in MoE layer

Motivation: Intuitively, the more independent the experts are within the same
MoE layer, the more diverse the representations learned by each expert. This
diversity makes it harder to replace or drop any of the experts. Conversely,
the redundancy of shared knowledge learned by different experts provides
opportunities for model pruning.

We employ the Centered Kernel Alignment (CKA) [17] as the criteria to evaluate
the similarity between experts in each MoE layer. Specifically, given a batch
of inputs $\\{x_{1},x_{2},...,x_{s}\\}$, the similarity $\rho_{ij}$ between
two experts $f_{i}$ and $f_{j}$ is computed as follows,

$\displaystyle\rho_{ij}=\frac{\text{HSIC}(K^{i},K^{j})}{\text{HSIC}(K^{i},K^{i})\cdot\text{HSIC}(K^{j},K^{j})},$
(4)

where $\text{HSIC}(K^{i},K^{j})=\frac{1}{(s-1)^{2}}\text{tr}(K^{i}HK^{j}H)$,
$K^{i}_{mn}=k(f_{i}(x_{m}),f_{i}(x_{n}))$, $H=I-\frac{1}{s}11^{T}$, and
$k(\cdot,\cdot)$ is the kernel function.

Here, we compute the expert similarity of Mixtral-8Bx7B using randomly
selected $32$ samples from the pre-trained dataset C4 [32]. We respectively
use the linear kernel and RBF kernel in similarity computation of eq. 4, where
the use of linear kernel actually corresponds to the dot product similarity
between the expert outputs.

As illustrated in Fig. 2, the darker the cell, the more similar the two
related experts. The two kernel-based distance measures demonstrate similar
changes in expert similarity across different layers. There are varying
moderate ratios of similar paired experts in different MoE layers, suggesting
the plausibility of pruning the MoE layer with many experts down to fewer
experts. For example, in the first layer, both the linear kernel and the RBF
kernel indicate that the $2$-nd to $5$-th experts exhibit great similarity
(greater than 0.7) with each other, suggesting that one of them can be pruned
without significantly degrading model performance.

#### Remark.

We notice that a concurrent work [26] also studies expert similarity, but
focuses on the cosine similarity between expert weights. Our work examines the
differences between experts using a more generalized metric based on the
correlation between inputs and outputs, and obtains a different observation:
even without considering the last layer, expert similarity does not gradually
decrease with increasing depth of layers. For example, we observe higher
pairwise expert similarity at the $11$-th layer compared to the $6$-th layer,
which is cross-verified by both distance metrics. Additionally, the experts in
the $26$-th layer also demonstrate a higher similarity with others compared
with shallow layers. We provide the full evaluation results on visualizing the
expert similarity of MoE layers in Mixtral-8x7B and Mixtral-8x22B in appendix
A.

###  3.3 Discover and Merge Similar Experts on Feature Space

Figure 3: Overview of our method. We first leverage different metrics to
obtain the expert representation, then compute the pairwise distance to get
the expert similarity matrix. Based on the expert similarity matrix, we can
group the experts with shared knowledge in the same cluster, which can be
merged on the weight space for MoE pruning.

Having identified the existence of moderate expert similarity in each MoE
layer, we propose pruning the experts within each MoE layer by merging experts
with shared knowledge. We present an overview of our method in fig. 3. Suppose
we want to reduce the number of experts from $N$ to $r$, the details of our
approach are outlined as follows.

Discovering Similar Experts from the Expert Graph. We construct a
unidirectional graph $\mathcal{G}=\\{\mathcal{V},\mathcal{E}\\}$ to group
similar experts into the same clusters, where we denote $\mathcal{V}$ as the
set of nodes and $\mathcal{E}$ as the connections between nodes.

First, for each expert $f_{i}$, we compute its feature representation
$\mathcal{M}(f_{i})$, where $\mathcal{M}(\cdot)$ is a specialized function to
represent $f_{i}$.

Second, we compute the pairwise distance based on expert representation and
obtain the disparity matrix $D\in\mathbb{R}^{N\times N}$, where
$D_{ij}=\text{DistanceFunction}(\mathcal{M}(f_{i}),\mathcal{M}(f_{j}))$.

Then, we initialize the $\mathcal{G}$ as a complete graph, namely the expert
graph, where $\mathcal{V}$ is set as all expert indices, and $\mathcal{E}$ is
set as the distance matrix $D$.

Last, we leverage the graph partition algorithm to split $\mathcal{G}$ into
$r$ subgraphs $\\{\mathcal{G}_{i}\\}_{i=1}^{r}$. The experts indexed by
$\mathcal{V}_{i}$ of $\mathcal{G}_{i}$ shares the most similar knowledge with
each other, and shows much difference with experts indexed by
$\mathcal{V}_{t}$ of $\mathcal{G}_{t}$, $t\neq i$. To achieve this goal, we
adopt a two-step strategy. At the first step, we partition the $\mathcal{V}$
of $\mathcal{G}$ into two subsets $\mathcal{V}_{1}$ and $\mathcal{V}_{2}$,
where $\mathcal{V}_{1}$ has $r-1$ nodes and $\mathcal{V}_{2}$ has $G-r+1$
nodes. Then, we split $\mathcal{V}_{1}$ into $r-1$ subsets,
$\mathcal{V}_{1,1},...,\mathcal{V}_{1,r-1}$. The splitting process can be
formulated as follows,

$\displaystyle\min\sum_{i,j\in\mathcal{V}}x_{i}x_{j}D_{ij}-\sum_{i,j\in\mathcal{V}}x_{i}(1-x_{j})D_{ij}$
(5) $\displaystyle-\sum_{i,j\in\mathcal{V}}(1-x_{i})(1-x_{j})D_{ij},$
$\displaystyle\textit{s.t.}~{}\sum_{i\in\mathcal{V}}x_{i}=r-1,\sum_{i\in\mathcal{V}}(1-x_{i})=G-r+1,$

where $x_{i}=1$ if node $i$ in group $\mathcal{V}_{1}$, and $x_{i}=0$
otherwise. The first term is used to minimize the distance between nodes in
$\mathcal{V}_{1}$, the second term is usd to maximize the distance of nodes
between $\mathcal{V}_{1}$ and $\mathcal{V}_{2}$, and the third term is ued to
maximize the distance of nodes in $\mathcal{V}_{2}$.

When pruning $G$ experts to $r$ experts, there are $r-1$ experts keeping
unaffected in the entire process. Our graph partition strategy can preserve
the most knowledge while diversifying the expert features in each MoE layers.

Merging Similar Experts with Routers. To preserve the different knowledge
learned by different models clustered in the same group, we merge these
experts as well as their routers on the weight space as follows,

$\displaystyle\hat{\theta}_{n}$
$\displaystyle\leftarrow\sum_{i=1}^{|\mathcal{V}_{n}|}\alpha_{i}\theta_{\mathcal{V}_{n}(i)},$
(6) $\displaystyle\hat{W}_{n}$
$\displaystyle\leftarrow\sum_{i=1}^{|\mathcal{V}_{n}|}\alpha_{i}W_{\mathcal{V}_{n}(i)},$

where $\mathcal{V}_{n}$ is the set of similar experts in the $n$-th cluster
$\mathcal{G}_{n}$, and we have $\sum_{i=1}^{|\mathcal{V}_{n}|}\alpha_{i}=1$.
We update the MoE layer $F$ by respectively replacing all the experts grouped
in $\mathcal{G}_{n}$ by a single FFN expert layer $f(\cdot;\hat{\theta}_{n})$
and corresponding routing weights $\hat{W}_{n}$, $n=1,2,...,N-r$.

###  3.4 Practical Implementation

While we have developed a general strategy to prune sparse MoEs, several
practical questions arise during implementation:

Q1: How to decide $\mathcal{M}(\cdot)$ and perform task-agnostic pruning on
sparse MoEs?

Intuitively, given the same input, similar experts should generate similar
outputs. Thus, we propose to use the mean squared error on the outputs of two
experts to evaluate the similarity between them.

We follow previous work [27] to select few samples from pre-trained dataset C4
[32] as the calibration dataset to compute the expert feature representation
at the third step for similar experts discovery. The C4 dataset comprises
multiple domains of contents, which can be a good surrogate to simulate the
performance of the task-agnostic pruning.

However, with the limited size of the calibration dataset, the estimation of
expert similarity based on the aforementioned expert representation could be
inaccurate due to biased sampling. To improve the generalization ability at
the discovery stage, we randomly mix up the token embedding for data
augmentation and then use the augmented data for computing the expert
representation and similarity. In practice, we iteratively update the expert
similarity matrix by continuously shuffling and mixing up the data for full
augmentation.

Q2: Is it possible to prune sparse MoEs without using any data?

Experts in the MoE layer of current models typically have the same structure;
for instance, in Mixtral-8x7B, each expert is a two-layer FFN with SiLU
activation. By training on different routed tokens, the experts usually learn
different distributions during the training process. This results in a
disparity between the learned feature spaces of different experts, which is
also potentially reflected by their weights.

Based on this analysis, we can use the weights as the expert representation to
discover similar experts for pruning. In our implementation, we consider two
strategies: (1) directly using the vectorized weights, which does not account
for the structure of experts, and (2) leveraging the local linearity of neural
networks to compute the approximated weight matrix, namely the surrogate
weights in our paper.

Taking the FNN $f(\cdot;\theta)$ in Mixtral-8x7B as an example, it consists of
three linear layers, i.e., $\theta=\\{\theta_{1},\theta_{2},\theta_{3}\\}$.
For input $x$, the output is computed as
$f(x;\theta)=\theta_{2}(\sigma(\theta_{1}x)\cdot\theta_{3}x)$. Then, the
vectorized weight is $\text{concat}\\{\theta_{1},\theta_{2},\theta_{3}\\}$,
while the surrogate weights can be obtained by
$\theta_{2}(\theta_{1}\cdot\theta_{3})$.

Q3: How do we determine the value of $\alpha$ for merging similar experts?

Two strategies are explored in our implementation to merge similar experts.

One strategy is to simply merge experts and gates with the same coefficients,
i.e.,
$\alpha_{1}=\alpha_{2}=...=\alpha_{|\mathcal{V}_{n}|}=\frac{1}{|\mathcal{V}_{n}|}$,
which encourages experts contributing equally.

The other strategy is to learn the merging coefficients for different experts
towards minimizing the following loss function,

$\displaystyle\mathcal{L}(\\{\alpha_{n}\\}_{n=1}^{|\mathcal{V}_{n}|})=\|y-F(x;\hat{\Theta},\hat{W},K)\|,$
(7)
$\displaystyle\textit{s.t.}~{}\hat{\Theta}(n)=\lambda(\sum_{i=1}^{|\mathcal{V}_{n}|}\alpha_{i}\theta_{\mathcal{V}_{n}(i)}),$
$\displaystyle\hat{W}_{n}=\sum_{i=1}^{|\mathcal{V}_{n}|}\alpha_{i}W_{\mathcal{V}_{n}(i)},$

where the ground-truth is the output of original MoE layer
$y=F(x;\Theta,W,K)$, and we jointly optimize $\lambda$ and $\alpha$ for
different merging groups.

##  4 Evaluation

Table 1: Results on pruning the Mixtral-8x7B from 8 experts to 6 and 4 experts
in each MoE layer. We present the results of our four strategies, namely 1)
Vectorized and surrogate $\theta$: prunable experts discovery using vectorized
or surrogate weights and merging with uniform coefficients; 2) Augmented:
prunable experts discover using augmented data and merging with uniform
coefficients; 3) Learn: prunable experts discovery using vanilla data and
learn coefficients to merge based on eq. 7. The first and second columns
respectively indicate the results on pruned model with $6$ and $4$ experts.

Dataset | MMLU | BoolQ | OpenBookQA | RTE | Average
---|---|---|---|---|---
Method | humanities | social science | stem | Other
Mixtral-8x7B [14] | 60.5 | 77.8 | 58.9 | 74.2 | 85.4 | 34.4 | 71.1 | 66.0
Router-guided [21] | 51.8/24.8 | 60.5/26.5 | 46.9/24.7 | 60.5/25.0 | 82.6/39.9 | 32.0/11.6 | 70.4/50.9 | 57.8/29.1
Count-guided [11] | 49.2/36.9 | 59.7/45.6 | 45.0/35.1 | 58.2/43.4 | 77.2/76.6 |  33.0/26.4 | 56.6/55.9 | 54.1/45.7
Enumerate [27] | 52.4/43.5 | 66.4/52.7 | 49.0/40.4 | 63.7/43.5 | 84.0/80.8 | 32.6/28.8 |  71.1/66.4 | 59.9/50.8
Ours |  Vectorized $\bf{\theta}$ | 54.4/48.1 | 70.2/58.5 | 51.8/45.2 | 66.8/55.2 | 85.6/83.7 | 31.4/26.2 | 68.9/62.4 | 61.3/54.2
Surrogate $\bf{\theta}$ | 56.8/47.1 | 69.7/56.4 | 50.9/42.2 | 66.0/55.0 | 86.9/83.8 | 32.6/27.0 | 68.5/64.6 | 61.6/53.7
Augmented |  56.2/47.8 | 72.7/57.2 | 52.1/42.7 | 68.1/54.0 | 85.6/83.2 | 32.8/28.6 | 68.6/66.5 | 62.3/54.3
Learn | 56.0/48.0 |  73.1/57.0 |  52.4/43.3 |  68.2/54.6 |  86.4/83.3 | 31.4/28.5 |  69.3/67.1 |  62.4/54.5

###  4.1 Experiment Setup

Studied Models. We take pruning Mixtral-8x7B and Mixtral-8x22B [14] as the
example. The Mixtral-8x7B has $32$ sparse MoE-involved layers, in each there
are $8$ experts. The hidden dimension for token embedding is $4,096$, and the
intermediate feature dimension in each FFN of MoEs is $14,336$. The
Mixtral-8x22B is similar with Mixtral-8x7B, but with $56$ sparse MoE layers.
It has the token embedding dimension of $6,144$ and the intermediate feature
dimension of $16,384$. During the inference, each token will select $2$
experts in each MoE layer.

Compared Baselines. We take three advanced MoE pruning methods [11, 21, 27] as
our baseline for comparison. Among them, router-guided merging [21] is
initially designed for task-specific MoE pruning, where we set the target
dataset as the samples from the pre-training dataset. For task-agnostic
methods, we select the frequency-based pruning method Expert Trimming [27],
also the count-guided strategy and loss-based pruning method [27], namely the
Enumerate in our paper. Under the task-agnostic pruning setting, we disable
all the fine-tuning stage of all methods for fair comparison.

Our Approaches. We use four metrics to prune the model, including 1)
Vectorized $\theta$: concact the vectorized model parameter $\theta$ together
as the expert representation for grouping, and apply the uniform merging on
the experts and routers. 2) Surrogate $\theta$: multiply the model parameters
together as the weight for one-layer surrogate FFN and use it as the expert
representation. The merging strategy is uniform merging. 3) Augmented: compute
the expert representation based on the expert outputs on the augmented inputs.
The merging strategy is uniform merging. 4) Learn: Use the same grouping
strategy with Augmented but learn to merge for experts and routers.

Evaluation Datasets. The open-sourced Language Model Evaluation Harness
library [10] is used to evaluate the performance. We select four zero-shot
tasks, including MMLU [12], BoolQ [6], OpenBookQA [31], and RTE [1]. Among
these tasks, MMLU consists of $57$ subtasks, where we present four groups,
namely the humanities, social science, stem and other.

Implementations. We randomly select $128$ samples from the pre-training
dataset C4 [32] as the calibration dataset. For the prunable candidate group
discovery, we use the mean square error to evaluate the similarity between
experts. For the learnable merging, we initialize the merging coefficients
$\alpha$ as the uniform and set $\lambda$ as $1$. In our experiments, we
utilize four NVIDIA A100 GPUs, each with 80GB of memory.

###  4.2 Pruning the Mixtral Architecture

We present the results of pruning the Mixtral architecture, including
Mixtral-8x7B and Mixtral-8x22B, from 8 experts to 6 and 4 experts in each
layer, respectively. The results on the two models are depicted in table 1 and
table 2 respectively.

Results on Mixtral-8x7B. We can see that all our proposed four strategies
surpass the related works, with a clear margin performance improvement of
$1.5\%$ on average. Compared with count-guided strategy [11] which just drops
the experts less visited, router-guided strategy [21] has a large improvement
of $3.7\%$ on average by merging these experts, showing that the merging
operation plays a crucial role in preserving the expert knowledge. Besides,
compared with count-guided and enumerate strategies [27] which all adopt the
dropping strategy, we can see that directly leveraging the expert feedback
rather than the routing frequency is more suitable for task-agnostic pruning
in MoE layers.

From the results of our strategies, we can also notice two facts: 1) Both
Vectorized $\theta$ and Surrogate $\theta$ strategies surpass all the other
data-involved pruning baseline methods. This suggests that weights already
encode fruitful data information and can be deployed to group experts for
pruning. 2) Compared with only using $\theta$, the involvement of $data$ in
the pruning process can help better candidates for merging and pruning for
improving the performance. Besides, though the learning strategy can bring a
further improvement in the average performance, the computation cost is larger
compared with using uniform coefficients for merging.

Results on Mixtral-8x22B. With more parameters introduced in Mixtral-8x22B,
the search for optimal experts to drop and merge becomes more challenging.
Unlike the results on pruning Mixtral-8x7B, the strategy using the surrogate
$\theta$ achieves state-of-the-art results, with only a minor performance drop
of $2.8\%$ on average compared to the full model.

At the same time, it should be noted that the use of data doesn’t help boost
the pruning performance as well, where the augmented and learning-based method
has an average performance drop of $1.05\%$ and $1.0\%$, respectively. Also,
the routing similarity-based pruning suffers model collapse under both two
settings. We argue it is due to the heavy imbalance between the used data and
parameter space, leading to the computation of expert similarity over-fitted
on the small calibration dataset.

Table 2: Results on pruning the Mixtral-8x22B from 8 experts to 6 and 4
experts in each MoE layer. We present the results of our four strategies,
namely 1) Vectorized and surrogate $\theta$, 2) Augmented, and 3) Learn. The
first column indicate the results on pruned model with $6$ experts, while the
second column present the result of $4$ experts.

Dataset | MMLU | BoolQ | OpenBookQA | RTE | Average
---|---|---|---|---|---
Method | humanities | social science | stem | Other
Mixtral-8x22B [14] | 68.6 | 84.1 | 67.1 | 78.7 | 87.9 | 0.358 | 71.2 | 70.4
Router-guided [21] | 27.3/22.7 | 25.4/25.8 | 24.4/24.0 | 27.9/23.4 | 62.8/62.7 | 12.8/13.0 | 54.2/49.5 | 33.5/31.6
Count-guided [11] | 58.0/45.7 | 74.9/57.7 | 54.1/42.0 | 70.2/45.7 | 81.5/74.4 | 35.2/27.0 | 69.3/57.4 | 63.3/50.0
Enumerate [27] | 60.4/53.9 | 78.0/67.2 | 59.5/52.3 | 73.0/64.2 | 87.4/80.5 | 35.0/31.1 | 70.1/67.9 | 66.2/59.6
Ours |  Vectorized $\bf{\theta}$ | 26.7/24.2 | 28.5/21.7 | 26.9/21.3 | 31.7/23.8 | 62.0/53.6 | 19.6/11.4 | 52.0/53.1 | 35.3/29.9
Surrogate $\bf{\theta}$ |  63.7/58.1 |  80.0/72.5 |  62.1/54.3 |  75.6/68.3 | 88.0/85.2 | 34.6/31.2 | 69.0/68.6 |  67.6/62.6
Augmented | 62.3/57.8 | 78.5/69.7 | 60.2/51.3 | 73.4/64.2 | 87.6/83.1 |  35.8/33.2 |  71.1/68.1 | 67.0/61.1
Learn | 61.4/57.9 | 78.3/70.1 | 61.2/51.4 | 72.8/65.0 |  88.2/84.9 | 35.6/32.7 | 70.5/67.3 | 66.9/61.3

###  4.3 Empirical Analysis on Expert Hints

While expert pruning brings a significant speedup by reducing the parameters,
what changes are introduced to different MoE layers? To answer this question,
we analyze the change of expert hints on the calibration dataset C4.

Figure 4: Statistics of the visiting frequency for all experts in different
MoE layers.

The statistic results are shown in fig. 4. We can see that in most MoE layers
in Mixtral-8x7B, many pairs of experts have similar hints. In contrast,
pruning differentiates the hints of experts. Compared to directly using hints
as the pruning goal, using expert knowledge as the pruning criterion results
in more significant changes in hints. Additionally, the merging operation on
the gate aggregates the new expert (last column) more tokens, further
increasing the hint differences.

###  4.4 Ablation Study and Discussion

On the number of samples. In our work, we augment the data to improve the
accuracy of similarity estimation. While this technique is typically used in
scenarios with limited data, a natural question arises:

Would increasing the amount of data help preserve performance after pruning?

Table 3: Evaluation results of different pruning setting with ranging the
number of samples for expert similarity estimation. The augmentation is
disabled in this experiment.

# of samples | MMLU | BoolQ | OpenBookQA | RTE | Avg.
---|---|---|---|---|---
128 | 0.615 | 0.856 | 0.328 | 0.686 | 0.621
256 | 0.614 | 0.854 | 0.332 | 0.697 | 0.624
512 | 0.617 | 0.854 | 0.334 | 0.704 | 0.627
2,048 | 0.612 | 0.856 | 0.322 | 0.675 | 0.616

We conduct experiments with different numbers of samples to answer this
question. As shown in table 3, when we increase the number of samples from
$128$ to $512$, we can see a noticeable performance improvement with up to
$0.6\%$. It should be noted that when we scale up the number of samples to
$2,048$, the average performance dramatically drops to $61.6\%$.

We attribute this phenomenon to the use of mean operation in our expert
similarity estimation. On the one hand, the mean operation keeps the feature
dimension for efficient pairwise distance computation. On the other hand, it
naturally downgrades the precision of function similarity computation by
“round-off” error compared with other per sample-based distance metrics like
CKA used in section 3.2.

On the choice of distance function. We also considered other distance
functions, such as cosine similarity, to evaluate expert similarity during the
discovery stage. However, it performed worse compared to our mean-squared-
error function, showing a drop of $0.9\%$ on MMLU, $3.3\%$ on RTE, and $0.2\%$
on BoolQ.

On merging the routing policy. One relevant work to ours is [21], which
identifies and merges similar experts based on routing policy, setting the
routing policy of the merged expert to that of the most frequently visited
expert. While we identify similar experts based on expert knowledge, such as
weights and features, we merge not only the experts but also the routing
policies. To demonstrate the importance of merging routers, we evaluate the
performance of our pruning method without merging the routing policies. As
shown in table 4, this results in performance degradation across all tasks and
different numbers of experts, highlighting the crucial role of simultaneously
merging the routers and experts.

Table 4: Evaluation results of our augmentation-based pruning method without
merging the routers.

# of experts | MMLU | BoolQ | OpenBookQA | RTE | Avg.
---|---|---|---|---|---
4 | 0.491 | 0.833 | 0.274 | 0.664 | 0.565
6 | 0.604 | 0.847 | 0.312 | 0.675 | 0.609

##  5 Conclusion and Future Work

In this paper, we address the practical problem of reducing the number of
experts in MoE layers to accelerate the inference and fine-tuning process. We
propose discovering similar experts at the feature level and then merging them
in the weight space for MoE pruning while preserving as much original expert
knowledge as possible. This approach allows the MoE layer to maintain diverse
experts with different knowledge, thereby efficiently reducing redundancy.

However, current methods do not account for the uniqueness of different MoE
layers during the pruning process, often setting the same number of experts to
be reduced across all layers. Our analysis reveals that some experts have low
visitation frequencies, indicating significant potential for further pruning.
Future work will aim to push the compression of MoE layers to the limit.

##  6 Acknowledgments

We would like to thank Lifeng Li from Microsoft for setting up GPU machines.

## References

  * Bentivogli et al. [2009] Luisa Bentivogli, Peter Clark, Ido Dagan, and Danilo Giampiccolo. 2009.  The fifth pascal recognizing textual entailment challenge.  _TAC_ , 7(8):1.
  * Castells et al. [2024] Thibault Castells, Hyoung-Kyu Song, Bo-Kyeong Kim, and Shinkook Choi. 2024.  Ld-pruner: Efficient pruning of latent diffusion models using task-agnostic insights.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 821–830.
  * Chen et al. [2022] Tianyu Chen, Shaohan Huang, Yuan Xie, Binxing Jiao, Daxin Jiang, Haoyi Zhou, Jianxin Li, and Furu Wei. 2022.  Task-specific expert pruning for sparse mixture-of-experts.  _arXiv preprint arXiv:2206.00277_.
  * Choi and Yang [2021] Kyusik Choi and Hoeseok Yang. 2021.  A gpu architecture aware fine-grain pruning technique for deep neural networks.  In _European Conference on Parallel Processing_ , pages 217–231. Springer.
  * Chowdhury et al. [2024] Mohammed Nowaz Rabbani Chowdhury, Meng Wang, Kaoutar El Maghraoui, Naigang Wang, Pin-Yu Chen, and Christopher Carothers. 2024.  A provably effective method for pruning experts in fine-tuned sparse mixture-of-experts.  _arXiv preprint arXiv:2405.16646_.
  * Clark et al. [2019] Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. 2019.  Boolq: Exploring the surprising difficulty of natural yes/no questions.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 2924–2936.
  * Ding et al. [2020] Guiguang Ding, Shuo Zhang, Zizhou Jia, Jing Zhong, and Jungong Han. 2020.  Where to prune: Using lstm to guide data-dependent soft pruning.  _IEEE Transactions on Image Processing_ , 30:293–304.
  * Fang et al. [2023] Gongfan Fang, Xinyin Ma, Mingli Song, Michael Bi Mi, and Xinchao Wang. 2023.  Depgraph: Towards any structural pruning.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 16091–16101.
  * Fedus et al. [2022] William Fedus, Barret Zoph, and Noam Shazeer. 2022.  Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.  _Journal of Machine Learning Research_ , 23(120):1–39.
  * Gao et al. [2021] Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, et al. 2021.  A framework for few-shot language model evaluation.  _Version v0. 0.1. Sept_ , 10:8–9.
  * He et al. [2024] Shwai He, Daize Dong, Liang Ding, and Ang Li. 2024.  Demystifying the compression of mixture-of-experts through a unified framework.  _arXiv preprint arXiv:2406.02500_.
  * Hendrycks et al. [2020] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020.  Measuring massive multitask language understanding.  In _International Conference on Learning Representations_.
  * Jacobs et al. [1991] R Jacobs, MI Jordan, GE Hinton, et al. 1991.  Mixtures of expert networks.  _Neural Comput_ , 3:79–87.
  * Jiang et al. [2024] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. 2024.  Mixtral of experts.  _arXiv preprint arXiv:2401.04088_.
  * Jordan and Jacobs [1994] Michael I Jordan and Robert A Jacobs. 1994.  Hierarchical mixtures of experts and the em algorithm.  _Neural computation_ , 6(2):181–214.
  * Kaplan et al. [2020] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_.
  * Kornblith et al. [2019] Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. 2019.  Similarity of neural network representations revisited.  In _International conference on machine learning_ , pages 3519–3529. PMLR.
  * Lemaire et al. [2019] Carl Lemaire, Andrew Achkar, and Pierre-Marc Jodoin. 2019.  Structured pruning of neural networks with budget-aware regularization.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 9108–9116.
  * Lepikhin et al. [2020] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2020.  Gshard: Scaling giant models with conditional computation and automatic sharding.  _arXiv preprint arXiv:2006.16668_.
  * Li et al. [2022] Bo Li, Yifei Shen, Jingkang Yang, Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. 2022.  Sparse mixture-of-experts are domain generalizable learners.  _arXiv preprint arXiv:2206.04046_.
  * Li et al. [2024a] Pingzhi Li, Zhenyu Zhang, Prateek Yadav, Yi-Lin Sung, Yu Cheng, Mohit Bansal, and Tianlong Chen. 2024a.  Merge, then compress: Demystify efficient smoe with hints from its routing policy.  In _The Twelfth International Conference on Learning Representations_.
  * Li et al. [2024b] Yunxin Li, Shenyuan Jiang, Baotian Hu, Longyue Wang, Wanqi Zhong, Wenhan Luo, Lin Ma, and Min Zhang. 2024b.  Uni-moe: Scaling unified multimodal llms with mixture of experts.  _arXiv preprint arXiv:2405.11273_.
  * Liao et al. [2023] Zhu Liao, Victor Quétu, Van-Tam Nguyen, and Enzo Tartaglione. 2023.  Can unstructured pruning reduce the depth in deep neural networks?  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 1402–1406.
  * Liu et al. [2021] Jing Liu, Bohan Zhuang, Zhuangwei Zhuang, Yong Guo, Junzhou Huang, Jinhui Zhu, and Mingkui Tan. 2021.  Discrimination-aware network pruning for deep model compression.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 44(8):4035–4051.
  * Liu et al. [2018] Zhuang Liu, Mingjie Sun, Tinghui Zhou, Gao Huang, and Trevor Darrell. 2018.  Rethinking the value of network pruning.  In _International Conference on Learning Representations_.
  * Lo et al. [2024] Ka Man Lo, Zeyu Huang, Zihan Qiu, Zili Wang, and Jie Fu. 2024.  A closer look into mixture-of-experts in large language models.  _arXiv preprint arXiv:2406.18219_.
  * Lu et al. [2024] Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang, Bo Zhang, Junchi Yan, and Hongsheng Li. 2024.  Not all experts are equal: Efficient expert pruning and skipping for mixture-of-experts large language models.  _arXiv preprint arXiv:2402.14800_.
  * Luo et al. [2018] Jian-Hao Luo, Hao Zhang, Hong-Yu Zhou, Chen-Wei Xie, Jianxin Wu, and Weiyao Lin. 2018.  Thinet: Pruning cnn filters for a thinner net.  _IEEE transactions on pattern analysis and machine intelligence_ , 41(10):2525–2538.
  * Mason-Williams and Dahlqvist [2024] Gabryel Mason-Williams and Fredrik Dahlqvist. 2024.  What makes a good prune? maximal unstructured pruning for maximal cosine similarity.  In _The Twelfth International Conference on Learning Representations_.
  * Masoudnia and Ebrahimpour [2014] Saeed Masoudnia and Reza Ebrahimpour. 2014.  Mixture of experts: a literature survey.  _Artificial Intelligence Review_ , 42:275–293.
  * Mihaylov et al. [2018] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. 2018.  Can a suit of armor conduct electricity? a new dataset for open book question answering.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pages 2381–2391.
  * Raffel et al. [2020] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _Journal of machine learning research_ , 21(140):1–67.
  * Shazeer et al. [2017] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. 2017.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  _arXiv preprint arXiv:1701.06538_.
  * Shen et al. [2022] Maying Shen, Pavlo Molchanov, Hongxu Yin, and Jose M Alvarez. 2022.  When to prune? a policy towards early structural pruning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 12247–12256.
  * Shi et al. [2024] Xinyu Shi, Jianhao Ding, Zecheng Hao, and Zhaofei Yu. 2024.  Towards energy efficient spiking neural networks: An unstructured pruning framework.  In _The Twelfth International Conference on Learning Representations_.
  * Wang et al. [2020] Yulong Wang, Xiaolu Zhang, Lingxi Xie, Jun Zhou, Hang Su, Bo Zhang, and Xiaolin Hu. 2020.  Pruning from scratch.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 34, pages 12273–12280.
  * Xia et al. [2023] Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. 2023.  Sheared llama: Accelerating language model pre-training via structured pruning.  _arXiv preprint arXiv:2310.06694_.
  * Yuksel et al. [2012] Seniha Esen Yuksel, Joseph N Wilson, and Paul D Gader. 2012.  Twenty years of mixture of experts.  _IEEE transactions on neural networks and learning systems_ , 23(8):1177–1193.
  * Zhang et al. [2023] Yihua Zhang, Ruisi Cai, Tianlong Chen, Guanhua Zhang, Huan Zhang, Pin-Yu Chen, Shiyu Chang, Zhangyang Wang, and Sijia Liu. 2023.  Robust mixture-of-expert training for convolutional neural networks.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 90–101.
  * Zhu et al. [2022] Jinguo Zhu, Xizhou Zhu, Wenhai Wang, Xiaohua Wang, Hongsheng Li, Xiaogang Wang, and Jifeng Dai. 2022.  Uni-perceiver-moe: Learning sparse generalist models with conditional moes.  _Advances in Neural Information Processing Systems_ , 35:2664–2678.
  * Zhu and Gupta [2017] Michael Zhu and Suyog Gupta. 2017.  To prune, or not to prune: exploring the efficacy of pruning for model compression.  _arXiv preprint arXiv:1710.01878_.
  * Zoph et al. [2022] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. 2022.  St-moe: Designing stable and transferable sparse expert models.  _arXiv preprint arXiv:2202.08906_.

##  Appendix A Evaluation on the expert similarity

Following the same setting in section 3.2, we conduct experiments on
evaluating the expert similarity in all MoE layers of Mixtral-8x7B and
Mixtral-8x22B. The results are respectively depicted in fig. A1 and table 2.
We summarize the observation as follows,

  * •

Most MoE layers in the two Mixtral models contain significant expert
redundancy.

  * •

The most redundant MoE layers are located in the first and last several
layers, while the experts in the intermediate MoE layers learn more diverse
features.

Figure A1: Evaluation of the expert similarity for different MoE layers in
Mixtral-8x7B under the linear kernel-based CKA criteria. Figure A2: Evaluation
of the expert similarity for different MoE layers in Mixtral-8x22B under the
linear kernel-based CKA criteria.