Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy
Authors: Pingzhi Li
Zhenyu Zhang
Prateek Yadav
Yi-Lin Sung
Yu Cheng
Mohit Bansal
Tianlong Chen
Venue: ICLR 2024 spotlight (2024)
Keywords: Sparse Mixture-of-Experts
Efficiency
Merging
Compression
Primary area: general machine learning (i.e., none of the above)
TL;DR: We propose an SMoE merging and compression framework, which leverages routing statistics as guidance, achieveing impressive results of both memory- and parameter- efficiency.

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$) $\textit{High Memory Usage,}$ due to duplication of the network layers into multiple copies as experts; and ($b$) $\textit{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 these challenges, we propose a novel merging algorithm for SMoE, $\textit{i.e.}$, $\texttt{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 based on routing policies; 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 draw an interesting observation 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, $\texttt{MC-SMoE}$ ($\textit{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 our proposals. For instance, our $\texttt{MC-SMoE}$ achieves up to $80\%$ memory and a $20\%$ FLOPs reduction, with virtually no loss in performance. Our code is provided as supplementary material.

Published as a conference paper at ICLR 2024
MERGE, THEN COMPRESS: DEMYSTIFY EFFICIENT SMOE WITH HINTS FROM ITS ROUTING POLICY
Pingzhi Li' Zhenyu Zhang² Prateek Yadav!
Yi-Lin Sung'
Yu Chengè
Mohit Bansal' Tianlong Chen 1A5
'The University of North Carolina at Chapel Hill
The University of Texas at Austin
The Chinese University of Hong Kong
*MIT
*Harvard University
Mlangzhi,praazi y ylsung, mbansal, Srttanoogg@eeee unc. edd zhenyu. zhangautexas. edu chengyu& cse. cuhk edu. hk
ABSTRACT
Sparsely activated of (SMoE) has shown promise scale up the learning capacity of neural networks, however, they have issues like: (a) High Memory Usage, due duplication of the network layers into multiple copies experts; and (b) Redundancy in Experts, as common learning-basee routing poli- cies suffer from representational collapse. Therefore, vanilla SMoE models are memory inefficient and non-scalable, especially for resource-constraing stream 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 ex- perts into fewer but more knowledgeable experts? Our pilot investigation reveals that conventional model merging methods fail to be effective in such expert merg- ing for SMoE. The potential reasons are: (1) redundant information overshadows critical experts; (2) appropriate neuron permutation for each expert is missing bring all of them in alignment. To address these challenges, we propose a novel merging algorithm for SMoE, i.e., M-SMOE, which leverages routing statistics to guide expert merging. Specifically, it starts with neuron permutation align- ment for experts: then, dominant experts and their "group members" are formed based on routing policies; 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 draw an interest- the
down-
ing observation that our proposed merging promotes a low dimensionality merged expert's 's weight space, naturally paving the way for additional compres- sion. Hence, our final method, MC- DMCS (i.e., Merge, then Compress SMoE), further decomposes the merged experts into low-rank and structural sparse alter- natives. Extensive experiments across 8 benchmarks validate the effectiveness of proposals. For instance, MC-SMoE achieves up 80% memory and 20% FLOPs reduction, with virtually no loss in performance.
INTRODUCTION
Transformers (Vaswani et al., 2023) have become the de facto network architecture in various natural language processing (NLP) scenarios (Devlin et al., 2019; Yang et al., 2019; Liu al., 2019; Raffel et al., 2020; Fedus et al., 2022; Wei et al., 2022), and even for computer vision applications (Dosovitskiy et al., 2021; Touvron et al., 2021; Mao et al., 2022: Zheng al., 2021; Liu et al., 2021). Nowadays, the parameter counts of such models are commonly measured in billions rather than millions. It is mainly because certain empirical scaling laws (Kaplan et al., 2020) reveal a power-law relationship be- tween the final model quality and the amount of (data, model capacity, and computing time), ). Unfortunately, it poses infea- sible requirements for computational resources, e.g., training a GaT-based model (Brown et al., 2020) typically leads to thou-
Accuracy (%)
MC-SMoE M.SMoE SMoE
Baselines
0.5
1.0
1.5
2.0
Model Size (B)
Figure 1:
Accuracy (%) the
COPA with the switch-base-33 SMoE. MC-SMoE reaches up to an 80% mem- saving with only a negligible com- promisc in performance.
Our code is provided http:://ghhuh.coo//////S/AAAAA.s

---

Published as a conference paper at ICLR 2024
sands of GPU days. Sparse Mixture-of-Eaie of-xeeertss (SMoE) (Shazeer et al., 2017) was then proposed to trim down the computing cost while enabling efficient scaling of network capacity. For pre- dictions of given input, it leverages input-dependent conditional computation to sparsely activate (i.e,, routing) the relevant model pieces (ie., experts). Hence, the network parameter counts/capac- ity can be amplified with minimal extra training cost. For instance, Fedus et al. (2022) scales the TS-Base (Raffel et al., 2020) dense model to 35x x larger Switch-Base SMoE model, with roughly the same training FLOPS.
However, several crucial limitations persist in SMoE for expanding the capacity of large language models. Firstly. SMoE trades space for FLOPs', which introduces substantial memory overheads and constrains its practical usage in real-world resource-restriceed platforms, especially for down- stream deployment and inference. Secondly, SMoE has a poor utilization of its capacity. The prevalent learning-basee based routing policy in SMoE suffers from representation collapse issues, since encourages token embeddings to be clustered around expert centroids (Chi et al., 2022) and re- sults in redundant experts (Mittal et al., 2022; Chen et al., 2022). A recent investigation (Chen al., 2023) also points out a similar observation that the "effective capacity" in conventional SMoEs is low. To address these drawbacks and fully unleash the power of SMoE, one possible solution is con- solidating information from insignificant experts, aiming to establish a more compact SMoE without hurting performance. Nevertheless, naively combining existing model merging mechanisms leads to substandard results in the SMoE scenarios, as demonstrated in our pilot studies in Section 4.2. The potential reasons could be: ! Critical experts are prone be overshadowed by redundant information during merging, @ Experts are usually initialized and trained along with diverse opti- mization trajectories, thus an expert permutation can play an essential role in bringing them into alignment (A (Ainsworth et al., 2022). These primary challenges drive us to ask:
(Q) How to effectively consolidate the redundant experts of SMoE into selected few ones without sacrificing vital knowledge?
In this paper, we systematically investigate the above research question (Q), and target a compact and high-quality SMoE downstream fine-tuning/ifference scenarios. We discover that the rout- ing policies from SMoE contain the "clues" for effective expert merging. To be specific, (1) the activation frequency of experts indicates its utilization and can be regarded as great proxy for its importance. It enables an automatic way to determine how many and which experts should be kept in each SMoE layer; (2) The routing decision measures how similar are the experts to each other, in terms of the relevance to given input samples. It helps in associating redundant experts with dif- ferent dominant experts. Based on these insights, we proposed a novel M-SMoE method for SMoE merging. Furthermore, we find that the merged experts from M-SMoE lie in a low dimensional parameter space, which seems to suggest that an appropriate merging reduces the potential noisy weight signals (Han et al., 2016). We utilize this additional benefit of expert merging to design our MC-SM4 DM-SMoE (Merge, then Compress SMoE) method that organically integrates low-rank decomposi- tion techniques for further expert compression. Our main contributions are as follows:
We propose novel framework MC MC-SMoE, i.e. Merge, then Compress SMoE, for SMoE efficiency at the downstream scenarios, including fine-tuning and zero-shot evaluation.
We design an innovative merging approach (M-SMoE) based on the guidance from routing policies. Specifically, it begins with a customized permutation alignment for experts, then identifies the dominant experts globally along with their "group members" within SMoE layers, and concludes with a weighted averaging according to their activated frequency.
We observe that resultant experts from M-SMoE inherently exhibit a lower weight dimen- sionality. This interesting phenomenon paves the way for additional compression, enabling our MC-SMoE method to further boost memory and parameter efficiency.
Extensive experiments across eight benchmarks validate the effectiveness of our MC-SMoE. An example is presented in Figure 1. Notably, M-SMoE yields up to 60% reduction in memory overhead with even slightly improved performance. MC-SMoF achieves up to 80% memory and 20% FLOPSs reduction, with only marginal performance drops.
FLOPP means the floating point operations per second. Note that the vanilla design of SMoE does not necessarily bring running time benefits. Instead. to mitigate the extra latency costs from routing and diverse experts, it usually requires specialized parallelism (Rajbhandari et al., 2022; Fedus et al., 2022; He et al., 2021; 2022) and hardware designs (Fan et al., 2022).
2

---

Published as a conference paper at ICLR 2024
Roused Frequency
Dominant Expert Dominant Expert
Expert
Expert 2
...
Expert k
...
Expert N
Grouping
Expert Index
Layer
(b) Routing Hints of SMoE Merging
Routing Policy
SMoF Layer
Token Embeddings
SMoF Layer
SMoF! Layer M-1
(a) Expert Routing of SMoE
SMeB Layer
lixpert
Merged Expert k
Compressing
Frequency-Awee Aware Expert Merging
Merged Layer
Merged Layer
Expert
Merged Layer
(MC-SMoE) Merge, Then Compress
Merged Layer M-1
(c) Merging Encourages Low Dimensionality.
Mergod Layer
The overview of proposed MC-SMoE pipeline. (a) In the conventional SMoE, each token (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 frequency-weighter fashion. (c) After merging, the weight space of resulted experts tends exhibit lower dimensionality, paving the way for additional compression. It clarifies the design of our MC-SMoE.
Figure 2:
embedding is directed a small number of relevant experts.
RELATED WORKS
Sparse Mixture-of-Experte (SMoE). The benefits of scaling model size are widely acknowledged, which usually offers increased learning capacity and enhanced generalization (Brown et al., 2020; Kaplan et al., 2020; Chung al., 2022: Chowdhery al., 2022). SMoE is efficient approach to train larger models with negligible additional overhead, which has been broadly studied in Shazeer al. (2017); Lepikhin et al. (2021); Fedus et al. (2022). SMoE models activate different pieces of the model for different input tokens opposed to utilizing the full network parameters, For instance, GShard (Lepikhin et al., 2021), an SMoE model scales up a Transformer- 600B parameters with training cost being lower than 100B dense model. Recently, Fedus et al. (2022) created a T5 (Raffel al., 2020) based SMoE model with trillion parameters.
model from 2B to
Efficiency Concerns in SMoE and Existing Solutions. SMoE models require huge memory to host experts, moreover, many experts have low utilization during inference. To address this, Chen et al. (2022); Kim et al. (2021); Koishekenov et al. (2023) prune experts based on their utilization to save memory, however, this leads to lower performance. In contrast, Gao et al. (2022) uses a tensor decomposition method to share the central tensor's parameters across experts and keep dif- ferent auxiliary tensors for each expert. Moreover, some works employ knowledge distillation (KD) (Rajbhandari al., 2022; Artetxe al., 2022; Fedus et al., 2022) to create either a smaller dense model or SMoE model with fewer layers. However, they also overlook the existing redundancy within SMoE layers. Moreover, Yadav et al. (2023a) show that experts can be compressed to huge degree without any performance loss.
Model Merging Language Models. The abundance of open-source models necessitates har- nessing these existing models to create superior ones. Network ensembling (Zhu et al., 2019; Ortega al., 2022) emerges as an intuitive solution, however, its computational burden during inference increases proportionally with the inclusion of more models. Recent literature has increasingly em- phasized the concept model merging (Yadav et al., 2023b; Cai et al., 2023; IIharco et al., 2022b; Matena & Raffel, 2022; Jin et al., 2022; Don-Yehiya et al., 2022; Rame et al., 2023). Yet, most of these studies assume that the merged models originate from the same initialization (Yadav et al., 2023b; Ilharco et al., 2022a; Wortsman et al., 2022), narrowing the pool of potential source models suitable for merging. However, this assumption might not be applicable to SMoE models. Typically, different experts within SMoE start with distinet random parameter initializations, and each expert
3

---

Published as a conference paper at ICLR 2024
COPA
SQuAD
WikIQA
SST2
11
533a2 Lones
122
62
12
32
10
10
Kepprt
Repert
Repert
Repprtt
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 the first MoE layer while the bottom is the last. The left tasks, COPA and SQuAD, characterized by answer-generation prompts. The right two tasks, WikiQA
Figure 3:
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.
is optimized with only subset of the training data, as determined by the routing networks. These characteristics make the task of merging experts in SMoE more challenging.
To tackle these challenges, numerous investigations resort to mode connectivity (Draxler et al., 2018; Frankle et al., 2020; Freeman Bruna, 2016; Garipov al., 2018) as a metric measure the intri- cacy of merging between two experts. The underlying premise is that models within the same loss basin are mergeable. Additionally, some works employ permutation invariance (Ainsworth et al., 2022; Jordan et al., 2022; Peña et al., 2023) to transfer models in different error basins into the same one without affecting their functionality. Jolicoeur-Martinee terms during training to enhance the mergeability of models, and Gueta al. (2023) systematically analyzes how training tasks, datasets, and recipes influence the difficulty of merging. work, SMEAR (Muqeeth al., 2023) dynamically merges various experts into single one during the training process to avoid discrete routing. Note that this approach doesn't 't offer any memory reduction and necessitates retaining the whole SMoE during inference.
et al. (2023) applies regularization
concurrent
METHODOLOGY
In this section, we present the details of our proposed MC-SMot method. Section 3.1 introduces the expert merging technique M-SMoE and how it is guided by the routing policy. In Section 3.2, we illustrate the extra benefit of merged experts and how it leads further compression. The whole procedure of MC-SMoE is provided at the end in Algorithm 1.
3.1 ROUTING POLICY GUIDES EXPERTS MERGING
Experts Permutation Alignment. Our M-SMoE method begins with the alignment of expert weight permutations since merging without it could potentially lead to the inferior fusion of mis- matched neurons. In our case, the target experts operate in the same input-output space, which makes the merging more feasible. The experts 2-layer feed-forward networks, where Win and Wout denote two weight matrices of input and output layers, respectively. æ is the input vector and (-t(-) represents the activation function. Then, a feed-forward network is defined as mapping F + (act )). Ainsworth et al. (2022) tells us that for any arbitrary permutation matrix P. the following equation (au( (Wn,.))) Wam (act..PW always holds. In other words, preserves the function F.
We follow the weight matching optimization in Ainsworth al. (2022) to align experts without altering their functionalities. For example, given two experts E; and E, with weight matrices W; and Wj. it try to locate the optimal P; and P., by minimizing the &2 distance between their corresponding permutated weights W¹ and Wy. Details are included in A2. This process provides beneficial first step for merging.
Routing Policies Reflect the Expert Similarity. One of the main challenges in SMoE expert merging comes from the expert specialization (Mittal et al., 2022) cultivated during the joint train- ing of experts and routers. Although representation collapse happens (Chi et al., 2022) and massive redundancies exist among experts, Figure 3 demonstrates that the utilization of several (more than one) experts is significantly larger compared to the rest. Therefore, it is challenging to merge all ex- perts within an SMoF layer into a single dense expert. Instead, we divide them into multiple groups based on their similarity, and keep all dominant (most used) experts preserve the performance. To
4

---

Published as a conference paper at ICLR 2024
COPA
SQuAD
SST2
Rabs 30
16
-X0
-15
I1N6-
30
12
Layer
Lager
Figure 4: Experts are more compressible after merging. We calculate the average stable-rand change ratio aafler-belirg of all dominant experts within each layer of the switch-base-32 SMoB model, reflecting the before
difference hefore and after merging. These mostly negative values throughout the SMoF layers emphasize a lower dimensionality achieved through the merging process.
meet the goal, our MM-MMEE method exploits the implicit guidance from SMoE's routing policy: (1) Similar rows (output channel) in router weight matrix tend to feed similar input tokens to their cor- responding experts, pushing these experts to be trained in a similar fashion; (2) Intuitively, experts that are similar tend to exhibit similar router logits across the majority of input tokens. Based on this, we can either use the rows in a router weight matrix or the router logits vector derived from a batch of input tokens, to measure expert similarity. Detailed comparisons are provided in Section 4.3 and we describe the superior one here, Le., router logits, and leave the other to Appendix A2, Specifically, the similarity Sim(.,-) between experts E, and E, an SMoE layer computed by:
= (x Si (E,, By) cos ne (R.,,, H,,-);
where X is an input embedding, W, is the router weight, H and Hy,, are row vectors in logits H.
Dominant Experts, Expert Grouping, and Frequency-Based Merging, Based on the expert utilization as depicted in Figure 3, we first treat the most commonly active experts as dominant experts. Such expert utilization is calculated by inputting and routing randomly picked subset of training data. Then, as demonstrated in Figure 2 (b), each non domminatt expert gravitates toward and joins the group led by its most similar dominant expert, using the similarity function defined by Equation 1. After grouping, each group consists of a few non-dominant and one dominant expert. Lastly, for group of k experts {E,,-.. E. a frequency-based merging is performed as follows:
>x 1 O.E,
Emerged
(2)
-i=1 O4
where a; is the usage frequency of expert E, The superiority of emphasizing the dominant experts is detailed and validated in our ablation study (Section 4.3).
Adaptive Layer- Wise Merging Ratio. As shown in Figure 3, the activated frequency of each expert varies across different SMoE layers, suggesting diverse number of dominant experts and corresponding groups. To consider this phenomenon, we normalize the frequencies within each SMoE layer and select the dominant experts in global manner across all layers', Take extreme case as an example, if the expert routing is uniform in one SMoE layer, then all experts will be treated dominant ones, echoing our intuitions.
3.2 MERGING ENCOURAGES EXPERT DECOMPOSITION
Merging Encourages Low-Rank Weights. We observe that M-SMoF promotes lower dimen- sionality in the weight space of merged experts, naturally facilitating additional compression. We adopt the metric from Wang et al. (2023) to measure the rank of weight spaces. This metric has proved to be practical as it primarily remains unswayed by minuscule singular values, providing rank estimation for the weight matrix W from network layer. It is defined below:
E
stable-rank(o)=
max of
where o denotes the singular value vector of Figure 4 showeases several stable-rank change ratio instances of SMoEs fine-tuned on various tasks. We measured the stable-ranks change after merging by calculating the ratio of its difference to its initial value. We see that the aver- aged stable-rank change ratio of all experts is consistently non-positive, i.e. stable-rank decreases, over most the SMoE layers, after merging. It inspires us to conduct post-merging compression, illustrated in Figure 2 (c).
'To ensure computational stability, we adjust the frequency of the most active expert in each SMoE layer to 1.0. In this way, at least one expert will be labeled as dominant. However, our experiments show that there are always least two dominant experts in each SMoF layer.
5

---

Published as a conference paper at ICLR 2024
Post-Merging Compression of SM-SMOE. To enjoy the extra benefits from merging, tailor the previous SoTA decomposition methods (Chen et al., 2021; ; Li et al., 2023) for SMoE, and propose an upgraded algorithm MC-SMoE for further memory and parameter efficiency. To be specific, the weight matrix W of a merged expert is decomposed into UV + S. Here, the product of U & Rdd x and V R represents low-rank approximation, where is much smaller rank compared to the full dimensionality of W. S contains the incoherent part of weights in W, and will be further pruned in structural manner. An importance score of weight $..j is computed as [(s.))
| V V.., C), where L indicates the training objective of SMoEs. To trim down S, the weight columns with the lowest cumulative scores Σ; I(s.g) will be removed, which is determined across all S weights and naturally leads to layer- wise adaptive compression ratio. As summary, Algorithm presents the full procedures of our proposed CMME framework.
Algorithm The Overall Procedures of MC-SMcE.
I: Initialize: model M with SMoE layers, training dataset with & tokens, the total number original experts and the number of the remaining experts k.
Let H E xxxx and A E R denote the ruter logits and frequeenc respectively
3: Let D represents the set of dominant experts
4: H, A + forward(M, 7): D + top (k, rox-normallzo(A)
5: for layer 1 1, .....,
6: for expert i = ......... do
7: E we we math
Expert Permutation Alignment
8:
end for
2(1) == arggaaxee.ccciinneeeee. (Hr.,,,, Hy H.)
Group Label Assignment
10:
for d E D do
G G f @(i) == dH: + Exeoh t8
11:
D. Merging based on Activated Frequencies Then compress
2.eg Ae,
12:
E + U +
13: end for
for do
15: Dropping E from M
16: end for
17:
end for
18: Return: A compact SMoE produced from MC- BC-SMoE.
4 EXPERIMENTS
4.1 IMPLEMENTATION DETAILS
Datasets and Network Backbones. Our ex- periments adopt the two open-source large lan- guage model families with their SMoE vari- ants: (a) the Switch Transformers (Fedus al., 2022) and (b) Meta's GPT--ased SMoE mod- els (Artetxe
Table I: Two SMoE models and their correponddng dense model checkpoints. act-size: number of activated parameters for each token, size: total number of param- eters, the number of transformer layers, h: hidden dimension, e: the number of number of experts, arch: the type of transformer architecture.
al., 2022). A summary of the specific model configurations is provided in Ta- ble 1. We use eight popular NLP tasks for supervised fine-tuning and evaluation: SST- (Socher et al., 2013) for sentiment classifi- cation, MRPC (Dolan & Brockett, 2005) for paraphrase identification, MultiRC (Khashabi et al., 2018) for multiple-choice QA, COPA (Gor- don et al., 2012) for sentence completion. WinoGrande (Sakaguchi et al., 2019) for conference resolution, SQuAD v1.1 (Rajpurkar et al., 2016) for extractive QA, WikiQA (Yang et al., 2015) and HotpotQA (Yang et al., 2018) for closed-book QA. For zero-shot evaluation, pick three representative benchmarks: MRPC in GLUE (Wang et al., 2019), WinoGrande for reasoning, and OpenBookQA (Mihaylov et al., 2018) for QA.
Model Identifier
act-size
size
arch
15 5-bare switch-base--2
220M
220M
12
768
1
enc-dev
220M
2.0B
12
708
cac-dec
Kuirseq-dense Saaneee-15mm 125m
125M
125M
12
708
1
dec
fairseq-moe 15b
15B
768
512
Comparison Baselines. We compare our proposals to six baselines including two pruning and four merging methods. Firstly, we consider the "task-specifie" expert pruning method from Chen

---

Published as a conference paper at ICLR 2024
Table 2: Performance evaluations on the switch-base-32 model with 32 experts in each SMoE layer, as well as its comparative dense model 15-base. We found the first SMoE layer has a profound impact on the model's performance, and merging it results in more significant performance degradation compared to other layers. Thus for all merging/compression mechanisms, the first SMoE layer is skipped following Ma et al. (2023), and it maintains average of 8 experts in other SMoE layers. We report exact-match/F/ score for SQuAD and HotpotQA, Fl-score for MultiRC, and accuracy for other tasks. For each task, we highlight the best performance over all baselines in blue and mark the performance no worse than full SMoE in bold.
Methods
Model Size
TFLOPs
SST-2
MRPC
MultiRC
COPA
WinoGrande
SQuAD
WikiQA
HotpotQA
Dense
220M
4.65
94.61
98.97
74.25
58.00
08.72
63 65/83 76
96.12
66. 13/33-55
Full SMcE
2.0B
4.65
95.75
90.20
76.19
68.00
61.80
65.39/85.81
96.45
67.55/8.1.00
Pruning
733M
4.65
91.30
88.97
73.13
63.00
61.04
64.80855 13
96.27
67.39/8 1.56
Task- Sokk-Speifiie
733M
4.65
91:28
82.04
53.63
52.00
08.56
54. 71.10778.00 00
96.24
64. 7117070.770
Averaging
733M
4.65
92.66
68.73
74.04
62.00
59.50
64.49/94.77 75
96.19
67.36/94.61
Ziplt
733M
1.65
93.12
91 18
73.26
65.00
60.38
65.01/85.06
96.05
67. 17/59/844 4.70
REPAIR
733M
4.65
92.80
90 444
74.44
65.00
61.48
64.67/94.48 84
96.27
67. 07.66784.77 77
93.35
74.25
65.00
59.25
61.61/8-1.92
96.20
67.23/8-L.66
M-SMoF
733M
4.65
94.20
90. 69
75.57
68.00
61.
65 77..66855.0
96:34
67.91/84.88 883
MC-SMoE
381M
3.33
93.33
89.2 22
73.98
67.00
69.52
65. 05.41/85..5
96.08
67.64/84.77 77
et al. (2022), which gradually drops non-active experts during fine-tuning. Additionally, we evaluate the one-shot pruning of non-dominant experts as a sanity check. Secondly, given the absence of prior work on expert merging, we directly adapt Averaging (Choshen et al., 2022), Ziplt (Stoica al., 2023), REPAIR (Jordan et al., 2022) and Git Re-basin (Ainsworth et al., 2022) merging methods to our SMoE scenarios as strong baselines for comparison.
Training and Evaluation Details. For the encoder-decoder models, including the switch-base-32 SMoE model and the 15-base dense model, we report supervised fine-tuning results. For each task, we first undertake a comprehensive hyper-parameter search. This encompasses batch sizes from {8, 16, 32, 64), learning rates from (3 10 4, 1 10-4 3 10-5, 1 x 10-5), "), and epoch counts spanning {3, 5. 10, 20), to pinpoint the optimal fine-tuned models. Further fine-tuning hyper-pann parameters are fixed, as shown in Appendix Table A15. After merging and compression, we proceed to fine-tune the condensed model to restore its performance. Further, compel the F and MC-SMoF models to imitate the outputs generated by the full SMoE model on the training dataset. The sappee-prrmeterss in the added KD loss are fixed for all tasks, please refer to Appendix A2 for more details. As for the decoder-only models, including the fairseq-moe- 15b SMoE model and the fairseq-dense- 125m dense model, we report zero-shot results, i.e. without undergoing any further training. For the compression phase in MC we set the sparse ratio to 0.1 and the low-rank factor 32, following Li et al. (2023). The model size and the number of tera floating point operations (TFLOPs) are reported to measure the efficiency. The TFLOPs is evaluated by batch of the first 64 samples in the SQuAD dataset, with the input sequence length of 329 and the target sequence length of 13. All experiments are conducted with Py PyTorch and DeepSpeed on NVIDIA A100 and A6000.
apply knowledge distillation (KD) to
4.2 COMPETITIVE PERFORMANCE AND SUPERIOR EFFICIENCY OF MC-SMOE
Table 2 presents the performance comparisons among MMSMEE, MC-SMoE, and eight baselines in supervised fine-tuning manner on {SST2, MRPC, MultiRC, COPA, WinoGrande, SQuaD, WikiQA, HotpotQA} } datasets. Note that all the compared methods activate the same number of parameters. From Table 2, the following observations can be drawn: ! M-SMoE achieves 60% memory reduc- tion while retaining performance on (MRPC, COPA, WinoGrande, SQuAD, HotpotQA). and even obtains {0.49. 0.25, 0.41} (%) extra performance improvement on (MRPC, SQuAD, HotpotQA} over the full SMoE model, respectively. Although M-SMOE shows marginal drop in performance for the memory efficiency on (SST2, MultiRC, WikiQA} benchmarks, however, it still outperforms all other pruning and merging baselines. These impressive results validate the superiority of our M in consolidating the redundant experts. @ MC-SMcB is performed on top of the expert merging from M... The resulting model achieves up to 80% in memory and 20% in FLOPs sav- ing, while the performance degradation remains less than 1% on (MRPC, COPA, SQuAD, WikiQA, HotpotQA). @ In addition, the zero-shot learning comparisons between ours and baselines with the fairseq-moe- 15b SMoE and fairseq-dense 125m dense models included in Appendix A
7

---

Published as a conference paper at ICLR 2024
4.3 ABLATION STUDY AND EXTRA INVESTIGATION
Ablation on Different Merging Ratio Designs. To tes- tify whether our adaptive merging ratio is effective not, we conduct an ablation study on different merging ratios, i.e., uniform (constant ratio per layer) v.s. adap- five (ours). Experimental results are produced with the switch-base-32 backbone on four datasets, as shown in Table 3. Our adaptive ratio presents a consistent advan- tage in terms of merging performance, compared to the uniform ratio. is within expectation since the pilot study in Figure 3 reveals that the number of frequently utilized experts is different across different transformer blocks.
Table 3: Comparison between Uniform and Adaptive (ours) merging ratio with the switch batt--aae-22 model on four datasets.
Merging Ratio
Uniform
Adaptive
MultiRC COPA
74.48
75.57
63.00
68.00
MRPC SQUAD
90.44
90.69
64.30/84.56
65. 65.66/85.44 49
Ablation on Different Grouping Methods. A pivotal component of our M-SMcE frame- work is to compute the similarity among ex- perts by router output logits, i.e. router-logits, which directly determines their grouping sta- tuses. Here, we carry out an ablation study for comparing our router-logits with seven other similarity functions: (i) random, which generates random vector for each expert; (iz) expert-weight, using the flattened weight of each expert's ss feed-forward network; (iii) expert-weeight-feuturr, leveraging the product of the ex- pert's weight and the L2 norm of its associated features; (iv) expert-gradient, utilizing the flattened gradients of each expert's feed-forward network; (1) expert-feature, adopting the average input hid- den states of each expert; (vi)
Table 4: Comparison between router-logits (ours) and seven other similarity functions for grouping experts.
Representations
MultiRC
COPA
MRPC
SQuAD
Random
74.69
62.00
80.95
64.97/84. 96
weight
75.29
63.00
89.46
64. 64.98/85.18
74.96
62.00
89.95
64. 14..8895.19
Expert Expert-grodient Expert-feature
59.00
74.74
60.00
89.95
65.03/85.22 21
Expert- feature.at
75.20
65.00
89.22
64. (44.5995.155
Ronter-weigha
88.73
64.9W95..
Router logits (Ours)
75.57
68:00
90.69
65. 65.68/85.49
using the average of absolute values of each ex- pert's 's input hidden states; (vit) youter-weight, adopting the corresponding row vector from the router weight matrix; and our (vili) router-logits, which uses the router output logits vector corresponding to the expert after feeding a batch to the SMoE model. Experimental results with the switch-base-32 model across four datasets are presented in Table 4. We observe that our router-logits consistently outperforms all other similarity variants. The strength reflect the routing decision distribution of input samples. During the training, experts with similar routing decision are optimized with a similar subset of data, leading to potential redundancy.
router-logits lies its ability to directly
Contribution from Knowledge Distillation. Knowledge distillation (KD) has been proven to be effective in inherit- ing information from large models. Therefore, we by default use KD for all merged and compressed SMoEs, including our MMMMOE, MC-SMoE, and all baselines. show its contribu- tion, we perform an ablation study comparing M-SMoF and w.o. the inclusion of KD loss during fine-tuning. Experimen- tal results presented in Table 5, with the switch-base-32 SMoE model across four datasets, underscore the advantages derived from the application of KD.
Table 5:
Comparison between fine- tuning M- EMoE wao. and w: (ours) KD with the switch--ase-32 model.
Methods
u.O. kD
w: kD
MultiRC COPA
74.77
75.57
64.00
68.00
MRPC SQuAD
90.69
63. 03.25/84.003
65.66/85.0 499
Contribution from Expert Permutation Alignment. Consider an expert with two feed-forward layers with an intermediate dimension of d, there are d! kinds of per- mutation possibilities to match and merge two experts. Next, we present an ablation study to compare M -MoE w. and w.o. alignment to assess the effectiveness of ex- pert permutation alignment. In Table 6, we present results with the switch-bas base-32 SMoE model on four datasets. It demonstrates clear performance improvement when ap- plying the expert permutation alignment before merging. alignment, expert merging could result in an inferior fusion of mismatched neurons.
Table 6: Comparison between M-SMoE W.O. and w. permutation alignment (PA) with the switch-base-32 32 model.
Methods M-SMoR w.o. M-SMoR w.
MultiCC COPA MRPC SQUAD
74.84
75.57
66.00
68.00
89.95
90 69
64.73/841.73
65.665 86.66/85.44
Therefore, without proper permutation
8

---

Published as a conference paper at ICLR 2024
Impact of Merging vs. Decomposition. To quantify the extra benefit of the low dimensionality arising from M-SMoE, E, we look at the effects of merging experts and compressing SMoEs separately. We con- sider the evaluation of three tasks using the switch-base-32 SMoE model and com- pare M-SMoE that only merges experts, C-SMoE that only compresses, and with MC-SMoE that does both merging and compression. From Table 7, we observe: 1 M-SMOR reduces the model size while maintaining or boosting performance. In contrast, C-SMoE.
Table 7: Comparison among M-SMoE that only merges, C-SMoE that only compresses, and MC-SMoE that merges and then compresses. Experiments are conducted with the switch-base-32 model. We highlight the better performance hetween C- SMOR and MC-SMOR, bold for each task.
Methods
SMoF
M-SMo7
C-SMoT
MC-SMcF
Model Size TFLOPs
2.0B
733M
570M
381M
4.65
4.65
COPA
68.00
68.00
61.00
67 .0
MRPC
90,20
69
88.97
89.22
SQUAD
65.39/85.81
65.66/85.47
6L 18/88999
65.41/85.30
(i.e., compression only) leads to a significant performance drop. It suggests that merging is a su- perior option to pursue memory efficiency and maintain model quality. @ The success of M SMoE paves the way for further compression. This is supported by MM-SMOE outperforming C- SMoo with even fewer parameter counts.
Ablation on Different Merging Strate- gies. To examine the effectiveness of our proposed frequency-aware expert merg- ing, an ablation study on different merg- ing strategies is needed. Specifically, we investigate uniform (Wortsman al., 2022), fisher-weighted (Matena & Raf- fel, 2022), and frequency-weiighted (ours) merging methods with the switch-base-32
Table 8: Comparison among different averaging strategies of Uniform, Fisher-weighted and Frequency-weightee (ours), evaluated with the switch-base-32 SMoE models.
Methods
Uniform
wishhr-weighted
Prequency-wrightee
MultiRC
75.11
73.77
75.57
COPA
61,00
65.00
68.00
MRPC
89.95
89.46
90.63 60
SQuAD
64.55/84.885
63.99/84.44
35.66/85.45
model across four datasets. As detailed Table 8, we see that our frequency-weighttd merging consistently reaches the best performance. A possible reason is that merging based on activation frequencies suppresses the impact of less significant experts. In contrast, the uniform approach redundant information, overshadowing critical experts during the merging process. As for the fisher-weighted merging strategy, which relies on gradient magnitude for expert reighigtingg, does not quite hit the mark, since our case, the experts have already been well pre-trained before merging.
tends give inappropriate prominence
Visualization of Compaet SMoEs from MC-SMoE. We visualize the distribution of dominant experts in the switch-buse-32 SMoE model produced by M-SMoE, and their compressed ver- sions from MC-SMoE in Figure 5. Each grid box denotes a dominant expert, and the darker color indicates more remaining parameters in that expert. Later SMoE layers, at the bottom of the heatmap, seem to be more mergeable and compressible.
WikiQA
SQuAD
66
SMOE Layer
q
Export
Expert
Figure 5: Ratio of remaining parameters after further compress- ing the dominant experts from MC-SMoB.
5 CONCLUSIONS
Sparse Mixture-of-Eii -x-eeerss (SMoE) is promising framework to scale up the model capacity, which enjoys roughly unchanged training and inference FLOPPs at the cost of significantly increased mem- ory overheads. The memory requirements and expert redundancy highly limit its practical usage. In this work, we propose an innovative SMoE merging approach, i.e., M M-SMcE, based on the hints from routing policies, to consolidate expert information into fewer but more knowledgeable ones. Moreover, such merged experts are demonstrated to be more compressible. our proposed, MC- SMoE methods pursue superior memory and parameter efficiency with competitive performance. We con- duct comprehensive experiments to support the effectiveness of our proposals. Future works mainly lie in the extension of multi-modality scenarios and co-designs with hardware platforms.

---

Published as a conference paper at ICLR 2024
6 REPRODUCIBILITY STATEMENT
To encourage reproducibility, we have made our source code available at our GitHub repository. Stwp:/bbbbboo//NNTTSSSabMbMM
including the data pre-processing. SMoE merging/- compression/pruning, and evaluation scripts. The hyperparameter details are provided in Ap- pendix A2 and the detailed pseudo-code about SMoE expert merging is provided in Appendix A3. We also provide clear and concise Algorithm for our MC-SMcE pipeline.
REFERENCES
Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.04836 2022.
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, Giri Anantharaman, Xian Li, Shuohui Chen, Halil Akin, Mandeep Baines, Louis Martin, Xing Zhou, Punit Singh Koura, Brian O' O'Horo, Jeff Wang, Luke Zettlemoyer, Mona Diab, Zornitsa Kozareva, and Ves Stoyanov. Efficient large scale language modeling with mixtures of experts. arXiv preprint arXiv:21/2.106644 2022.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Ruisi Cai, Zhenyu Zhang, and Zhangyang Wang. Robust weight signatures: Gaining robustness as easy as patching weights? arXiv preprint arXiv:2302 12480, 2023.
Tianlong Chen, Zhenyu Zhang, Ajay Jaiswal, Shiwei Liu, and Zhangyang Wang. Sparse moe as the new dropout: Scaling dense and selm-slimmable transformers. arXiv preprint arXiv:2303.01610, 2023.
Tianyu Chen, Shaohan Huang, Yuan Xie, Binxing Jiao, Daxin Jiang, Haoyi Zhou, Jianxin Li, and Furu Wei. Task-specific expert pruning for sparse mixture-of-exxeets. arXiv preprint TrXiv:2206.00777 2022.
Xuxi Chen, Tianlong Chen, Yu Cheng, Weizhu Chen, Zhangyang Wang, and Ahmed Hassan Awadallah. Dsee: Dually sparsity-embedded efficient tuning of pre-trained language models. arXiv preprint arXiv:2111.001660 2021.
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, and Furu Wei. On the representation collapse of sparse mixture of experts. arXiv preprint arXiv:2204.09/ /9999999999.99.. 2022.
Leshem Choshen, Elad Venezian, Noam Slonim, and Yoav Katz. Fusing finetuned models for better pretraining. arXiv preprint arXiv:2204.03044 2022.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay. Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard. Guy Gur-Ari, Pengeheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito. David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowyez, Erica Moreira, Rewon Child, Oleksandr Polozov. Katherine Lee, Zongwei Zhou, Xuezhi Wang, Bren- nan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. arXiv preprint ar 224:02204. 2022.
Douglas
10

---

Published as a conference paper at ICLR 2024
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros, Marie Pel- lat, Kevin Robinson, Dasha Valter, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei. Scaling crruution-finetune language models. arXiv preprint arXiv:2210.
2022.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. ArXiv, abs/ abs/1810004805,, 2019.
William B. Dolan and Chris Brockett. Automatically constructing corpus sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (IWP2005), 2005. URL https: htt////aaaaatthhceee
org/ 115555022
Shachar Don-Yehiya, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. Cold fusion: Collaborative descent for distributed multitask finetuning. arXiv preprint
arXiv:2212.01378, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recognition at scale. ArXiv, 1bb//2000.1199, 2021.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference on machine learning, pp. 1309- 1318. PMLR, 2018.
Zhiwen Fan, Rishov Sarkar, Ziyu Jiang, Tianlong Chen, Kai Zou, Yu Cheng, Cong Hao, Zhangyang Wang, et al. M Mixture-of-exppet vision transformer for efficient multi-task learning with model-accelerator co-design. Advances in Neural Information Processing Systems, 35:28441- 28457, 2022.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-99 2022. URL http://///. / jmlr. org/ papers/ /ttee//////////// 2 0998. html.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode con- nectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pp. 3259-3269, PMLR, 2020.
C Daniel Freeman and Joan Bruna. Topology and geometry of half-rectified network optimization. arXiv preprint arXiv: 161 1611.01540 2016.
Ze-Feng Gao, Peiyu Liu, Wayne Xin Zhao, Zhong-Yi Lu, and Ji-Rong Wen. Parameter- efficient mixture-of-experts architecture for pre-trained language models. arXiv preprint arXiv:2203.01104, 2022.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in neural information pro- cessing systems, 31, 2018
Andrew Gordon, Zornitsa Kozareva, and Melissa Roemmele. SemEval-2012 task 7: Choice of plausible alternatives: An evaluation commonsense causal reasoning. In *SEM 2012: The First Joint Conference on Lexical and Computational Semanties Volume 1: Proceedings of the main conference and the shared lask, and Volume 2: Proceedings of the Sixth International Workshop on Semantic Evaluation (SemEval 2012), pp. 394-398, Montréal, Canada, 2012. Association for Computational Linguistics. URL https:/ / /aaaaaaaanaooogoooooeo www.ollyyoooooo.oooooooooooooore /122112--22 1052.
Almog Gueta, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. Knowledge is a region in weight space for fine-tuned language models. arXiv preprint arXiv:2302. 04863, 2023.
11

---

Published as a conference paper at ICLR 2024
Song Han, Huizi Mao, and William J. Dally. Deep compression: Compressing deep neural networks with pruning. trained quantization and huffman coding. arXiv preprint arXiv:1510000149 2016.
Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang, Jidong Zhai, and Jie Tang. Fastmoe: A fast mixture-of-expert training system. arXiv preprint arXiv:2103. 13262, 2021.
Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang, Fuwen Luo, Shangfeng Shi, and Qin Li. Fastermoe: Modeling and optimizing training of large-scale dynamic pre-trained models. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programminng PPoPP '22, pp. 120-134, New York. NY, USA, 2022. Association for Com- puting Machinery. ISBN 9781450392044. doi: 00.1155/503221.350444 URL https: //doi ort.oorg//ii 10. 1145/ 3503221.3508418.
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. arXiv preprint arXiv:2212.0400 04089, 2022a.
Gabriel Ilharco, Mitchell Wortsman, Samir Yitzhak Gadre, Shuran Song, Hannaneh Hajishirzi, Si- mon Kornblith, Ali Farhadi, and Ludwig Schmidt. Patching open- vocabuocyyy models by interpo- lating weights. Advances in Neural Information Processing Systems, 35:29262-29277, 2022b.
Xisen Jin, Xiang Ren, Daniel Preotiuc-Pietro, and Pengxiang Cheng. Dataless knowledge fusion by merging weights of language models. arXiv preprint arXiv:22 02.999 2022.
Alexia Salicoeur-Martinees Emy Gervais, Kilian Fatras, Yan Zhang, and Simon Lacoste-Julien. Population parameter averaging (papa), arXiv preprint arXiv:2304.0300) 2023.
Keller Jordan, Hanie Sedghi, Olga Saukh, Rahim Entezari, and Behnam Neyshabur. Repair: Renor- malizing permuted activations for interpolation repair. arXiv preprint urXiv:2211.08400 2022.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08366 2020.
Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Looking be- yond the surface: A challenge set for reading comprehension over multiple sentences. In Proceed- ings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lunguage Technologies, Volume (Long Pupers), pp. 252-262 New Or- leans, Louisiana, 2018. Association for Computational Linguistics. doi: 10. 10.18653/v1/118 1023. URL /attps://aclaae http:///aaaaathhoooo
Young Jin Kim, Ammar Ahmad Awan, Alexandre Muzio, Andres Felipe Cruz Salinas, Liyang Lu, Amr Hendy, Samyam Rajbhandari, Yuxiong He, and Hany Hassan Awadalla. Scalable and effi- cient moe training for multitask multilingual models. arXiv preprint arXiv:2109. 10465, 2021.
Yeskendir Koishekenov, Alexandre Berard, and Vassilina Nikoulina. Memory-efficiee NLLB-200. Language-speeifiic expert pruning of massively multilingual machine translation model. In Anna Rogers, Jordan Brad-rraber, and Naoaki Okazaki (eds.), Proceedings of the 61 Annual Meet- ing of the Association for Computational Linguistics (Volume I: Long Pupers), pp. 3567-3585, Toronto, Canada, July 2023. Association for Computational Linguisties. doi: 18.55/ 1/2023. 198.
acl-long. 198. URL ht htaptt:://acaaatthlleo /aaaaaaaannttoooooooooooo org/ org/202/ 2022
Harold W. Kuhn, The Hungarian Method for the Assignment Problem. Naval Research Logistics Quarterly, 2(1-2):83-97 March 1955. doi: 10. 1002/nav.38000000000
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. {GS) hhard: Scaling giant models with condi- tional computation and automatic sharding. In International Conference on Learning Represen- fations, 2021. URL https: //wpnreviieeeeeeeeemoemeeeeeeeee net/ ftret//orummaa dewrrummddigweeeee XHTTbb..
Yixiao Li, Yifan Yu. Qingru Zhang, Chen Liang, Pengcheng He, Weizhu Chen, and Tuo Zhao Losparse: Structured compression of large language models based on low-rank and sparse ap- proximation. arXiv preprint arXiv:2306. 1222, 2023.
12

---

Published as a conference paper at ICLR 2024
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zertlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 19X.::197..119999 2019.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 10012-10022, October 2021.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. LIm-pruner: On the structural pruning of large language models. arXiv preprint vrXiv:2305.11622 2023.
Zhiyuan Mao, Ajay Jaiswal, Zhangyang Wang, and Stanley H. Chan. Single frame atmospheric turbulence mitigation: A benchmark study and a new physics-inspiret transformer model. ArXiv, abs/2207.100400, 2022.
Michael S Matena and Colin Raffel. Merging models with fisher-weighted averaging. Advances in Neural Information Processing Systems, 11117700-17716 2022.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP, 2018.
Sarthak Mittal, Yoshua Bengio, and Guillaume Lajoie. Is a modular architecture enough? Advances in Neural Information Processing Systems, 35:28747-28760, 2022.
Mohammed Muqeeth, Haokun Liu, and Colin Raffel. Soft merging of experts with adaptive routing. arXiv preprint arXiv:2306.03745 2023.
Xiaonan Nie, Pinxue Zhao, Xupeng Miao, Tong Zhao, and Bin Cui. Hetumoe: efficient trillion- scale mixture-of-expert distributed training system. arXiv preprint arXiv:2203.14685, 2022.
Luis A Ortega, Rafael Cabañas, and Andres Masegosa. Diversity and generalization in neural net- work ensembles. In International Conference on Artificial Intelligence and Statistics, pp. 11720- 11743. PMLR, 2022.
Fidel A Guerrero Peña, Heitor Rapela Medeiros, Thomas Dubail, Masih Aminbeidokhti, Eric Granger, and Marco Pedersoli. Re-basin via implicit sinkhorn differentiation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 20237-20246, 2023.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 10683, 2020.
Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Am- mar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-mos -me: Advancing mixture-of-experrs inference and training to power next-gemeration ai scale. arXiv preprint arXiv:2201.055/1 2022.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ ques- tions for machine comprehension of text. Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 2383-2399,, Austin, Texas, 2016. Association for Computational Linguistics. doi: 10.18653/v1/D16 1/44.1 URL https: http::s:///caaahhlogg..s org/D16- 1264.
Alexandre Rame, Kartik Ahuja, Jianyu Zhang, Matthicu Cord, Léon Bottou, and David Lopez-Pat Model ratatouille: Recycling diverse models for out-of-distribution generalization. arXiv preprint 2rXiv:22/2.10444 2023.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale. arXiv preprint arXiv: 100.//907/4.4.. 10641, 2019.
Victor Sanh, Thomas Wolf, and Alexander M. Rush. Movement pruning: Adaptive sparsity by fine-tuning, 2020.
13

---

Published as a conference paper at ICLR 2024
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The spansely-gate mixture-of ex-eeertss layer. 2017.
arXiv preprint
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing. pp. 1631-1642, Seattle, Washington, USA, 2013. Association for Computational Lin- guistics. URL https: : /tt//aaaaaaatloooeeeeeeaeaa 1tt.Looyy.or//1-11117710
George Stoica, Daniel Bolya, Jakob Bjorner, Taylor Hear, and Judy Hoffman. Zipit! merging models from different tasks without training. arXiv preprint arXiv:2305.03053 2023.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herv'e J'egou. Training data-efficient image transformers & distillation through attention. In ICML, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. arXiv preprint (766:07706 2023.
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy. David Courna- peau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. Nel- son, Eric Jones, Robert Kern, Eric Larson, C Carey, ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mul- bregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-274 2020. doi: 11000///41112-019-000000
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXivv:8044 07441, 2019.
Hongyi Wang, Saurabh Agarwal, Pongsakorn U-chupala, Yoshiki Tanaka, Eric Xing, and Dim- itris Papailiopoulos. Cuttlefish: Low-rank model training without all the tuning. arXiv preprint arXiv:2305.022 02538, 2023.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoe Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11900 2022.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning, pp. 23965-23998, PMLR, 2022.
Prateek Yadav, Leshem Choshen, Colin Raffel, and Mohit Bansal. Compeft: Compression for communicating parameter efficient updates via sparsification and quantization, 2023a.
Prateek Yadav, Derek Tam, Leshem Choshen, Colin Raffel, and Mohit Bansal. Resolving interfer- ence when merging models. In NeurIPS, New Orleans, USA, 2023b. Proceedings of Machine Learning Research.
Yi Yang, Wen-tau Yih, and Christopher Meek. WikiQA: challenge dataset for open-domain ques- tion answering. In Proceedings of the 2015 Conference on Empirical Methods in Natural Lan- guage Processing, pp. 2013-2018, Lisbon, Portugal, 2015. Association for Computational Lin- guistics. doi: 10.18653/v 100165/////115--11 1237. URL https: / acl.ntoooontoyyycccccoo org/D15- 1237.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. HotpotQA: A dataset for diverse, explainable multi-hop hop question an- swering. In Proceedings of the 2018 Conference Empirical Methods in Natural Language
14

---

Published as a conference paper at ICLR 2024
Processing, pp. 2369, Brussels, Belgium, October-November 2018. Association for Com- putational Linguistics. doi: 10.18653/11/D18-12555 URL https:://caaattholggg/.ooooo///s .rrg/ D18-1259.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32. 2019.
Ofir Zafrir, Ariel Larey, Guy Boudoukh, Haihao Shen, and Moshe Wasserblat. Prune once for all: Sparse pre-trained language models, 2021.
Qingru Zhang, Simiao Zuo, Chen Liang, Alexander Bukharin, Pengcheng He, Weizhu Chen, and Tuo Zhao. Platon: Pruning large transformer models with upper confidence bound of weight importance, 2022.
Minghang Zheng, Peng Gao, Renrui Zhang, Xiaogang Wang, Hongsheng Li, and Hao Dong. End- to-end object detection with adaptive clustering transformer. ArXiv, 02s/2011.093111 2021.
Michael Zhu and Suyog Gupta. To prune, or not to prune: exploring the efficacy of pruning for model compression, 2017.
Shilin Zhu, Xin Dong, and Hao Su. Binary ensemble neural network: More bits per network or more networks per bit? In Proceedings the IEEE/CVF Conference on Computer Vision and Patlern Recognition (CVPR), June 2019.
15

---

Published as a conference paper at ICLR 2024
APPENDIX
AI MORE EXPERIMENTAL RESULTS
A1.1 ZERO-SHOT EVALUATION RESULTS
We compare our proposed M-SMcE, MC-SMoE with one-shot pruning of non-dominant experts and the "task-specific" expert pruning method, in zero-shot learning manner. Our MMM consis- tently outperforms the baseline methods, as shown in Table A9, The performance might be further improved if only we can fine-tune the routers, given that our M-SMOE highly leverages routing information during the merging phase.
Table A9: Performance evaluation on the fuirseq-moe-130 model with 512 experts in each SMoE layer. as well as its comparative dense model fuirseq-dense-125mr Different from the fine-tuned switch-buse-32 model, we apply muning/merging methods on every SMoE layer here and maintain an average of 16 experts. We highlight the best performance over all baselines bold.
Methods
Model Size
TFLOPs
MRPC
OpenBookQA
WinoGrande
Dense
125M
5.08
37.75
34.00
49.49
Full SMoE
15B
5.08
60.54
36,80
51.78
Pruning
562M
5.08
52:20
30.60
48.46
Task Task-Specific
552M
5.08
40.19
23.60
48.38
M-SMoR
552M
5.08
52.69
34.40
50.43
MC-SMoE
160M
4.45
47.35
34.60
49.09
A1.2 EFFICIENCY DISCUSSIONS AND LIMITATIONS
Latency Limitations Despite the (dense, SMoE, M-SMoE, MC-SMoE} models sharing the same theoretical TFLOPs, they do not necessarily produce the same latency. This is because the vanilla de- sign of SMoE in the real world suffers from significant extra latency costs introduced by routing (Nie et al., 2022). Our proposed M-SMcE and WC-SMMB. achieve impressive memory and TFLOPs effi- ciency for SMoE. However, they do not improve latency. Ideally, the merging process is supposed to reduce the number of classes managed by the router classifier due to the reduction in the number of experts in each layer. However, in practical implementation, we face challenge: explicitly creating a new router for the merged experts is non-trivial. To address this issue, we adopt the following strategy as shown in Appendix A3: within each group, we retain representative expert and let other routers point towards this representative. Yet, all such routing decisions into this group will now be directed towards a single new merged expert. This implies that, although the count of ex- perts reduces, the number of classes managed by the router remains constant, i.e. the routing latency costs remain constant. Thus, if we manage to prune the router output channels without affecting its functionality, we can realize notable improvement in latency efficiency.
To examine the potential efficiency from router pruning upon M- SMoE, we conduct experiments with the switch-base-32 backbone batch size {32, 256, 512) and compare inference latency of these four models: 4D dense, @ SMoE, @ M-SMOE, M-SMoE w. pruning router. Notably, results in Table A10 across three batch size settings demonstrate a latency ordering of 2>>3> ①> O. This indicates the latency limitation and encourages future work for router pruning.
Table A10: Latency analysis of the switch-hase-32 model SQuAD task inference with BF16.
Bsz=32
BSZ=256
BsZ=512
Models
TFLOPs
Latency (8)
TFLOPs
Latency (8)
TFLOPs
Latency (8)
Dense
2.33
0.08
25.51
0.85
50.32
2.33
Full SMoE-32
2.33
0.18
25.51
1.02
59.32
2.50
M-SMoE-8
2.33
0.17
25.51
0.99
50.32
2.48
M-SMoF-8 w. pruning router
2.33
0.13
25.51
0.93
59.32
2.38
Potential Specialization for Inference Implementation We first present a comprehensive inves- tigation of the inference cost of our full SMoE, M-SMoE, and MC-SMoE models, focusing on both computational and memory efficiency. Our investigation covers latency, throughput, and FLOPPs for
A16

---

Published as a conference paper at ICLR 2024
computational aspects, along with model size and memory cost for memory aspects. As shown in Table A11, the underscored results demonstrate the marginal inference gain from M-SMoE, which confirms our analysis in the first paragraph of Appendix A1.2. On the other hand, the throughput of MC-SMoE, is lower than that of M M-SMoF, despite it consuming less memory and FLOPs. This is due our lack of specialized sparse matrix support software hardware for MC-SMoE, which encourages our future work.
Table A11:
Computational and memory efficiency evaluation our full SMoE, M-SMcE, and MC-SMoE models without specialized implementation. All performance is produced using the same input size, including (throughput (token per ms), latency (ms), GFLOPs, memory, model size (number of parameters))
Models
Throughput
Latency
GFL.Ps
Memory
Model Size
Full SMoE
4.47
114.3
232
3893MB
2.0B
M-SMoF
4.82
106.2
232
1456MB
733M
MC-SMoR
2.71
189.0
186
715MB
381M
However, theoretical speedup exists. This is because, in conventional SMoE implementation, the routing process involves two drawbacks of throughput: (1) a layout transform of the tensors (to arrange tokens targeting the same experts into a continuous memory buffer) and its reverse opera- tion (Nie et al., 2022), and (2) breaking down one large matrix block GEMM operation into many smaller matrix block GEMM operations (each corresponding to an individual expert), leading to less efficient utilization of modern computational hardware's s advantages. These factors lead to decrease in real throughput for the sparsely activated computation in SMoE when the number of experts rises, a topic that remains an open area for research (Nie et al., 2022) and is carmarked for exploration in our future studies. While our M-5McE confronts the first challenge due to the diffi- culty of pruning the router's output channels, we capable of optimizing the inference speed from the second challenge.
We conduct an extended evaluation of computational and memory costs for a specialized inference design. Our approach involves gathering tokens routed to all experts of one group and processing them through one single expert, leveraging the shared weights within the group. This strategy is designed to take advantage of the parallel processing capabilities of hardware accelerators, typically GPUs. The underscored results presented in Table A12 clearly illustrate the enhanced throughput and latency performance of our M-SMoE and MC-SMoE models post-implemmnnation of this opti- mization technique. We believe these promising initial results will catalyze additional exploration and research.
Table A12: Computational and memory efficiency evaluation on our full SMoE, M-SMoE, and MC-SMoE models with specialized implementation. All performance is produced using the same input size, including (throughput (token per ms), latency (ms), GFL.OPs, memory, model size (numher of parameters))
Models
Throughput
Latency
GFLOPs
Memory
Model Size
Full SMoF
4.47
114.3
232
3895MB
2.0B
M-SMoE
7.91
64.7
232
1456MB
733M
MC-SMoE
6.27
81.6
186
715MB
381M
A1.3 COMPUTATIONAL COST DISCUSSION OF M-SMoE
We present a detailed computational cost analysis for each stage of our merging procedure. The M- MMSMoE merging approach encompasses three principal stages: Daligning expert permutations, @grouping experts, and @merging expert weights. To begin with, aligning expert permutations is performed separately in each SMoE layer, which results in the computational costs being linearly correlated with the number of SMoE layers. Secondly, expert grouping involves model inference to assess activation frequencies and router logits, followed by calculating pair-wise similarity among experts. Owing to the sparse activation computations inherent in SMOE, the model's inference costs remain unchanged regardless of the number of SMoE layers, leading to the similarity computa- tions within each SMoE layer being the main contributors to linear increase in computational costs. The final stage, merging expert weights within each SMoE layer, also adds this linear increase in computational demands. To sum up. while some aspects of our approach maintain a constant computational load, our overall cost analysis indicates trend of linear growth in these demands.
A17

---

Published as a conference paper at ICLR 2024
To validate analysis, conduct extra experiments for the computational costs of merging. We evaluate the switch-base-32 model's computation time costs of Dexpert permutation alignment, @expert grouping, and @expert weight merging respectively. We maintained a constant (24) total number of Transformer layers while varying the number of SMoE layers. The results shown in Table A13 confirm our analysis, indicating that the primary bottleneck in terms of time cost is rooted in the expert permutation alignment, while the bulk of memory cost is attributed to model inference.
Table A13: Computational costs of our M- SMoE merging method, evaluated with the wwtch--base-32 the COPA task. We maintain a constant total number Transformer layers of and vary number of SMoE layers from 2 to 12. The three principal stages of M-SMcE are evaluated separately, including Permutation Alignment (PA), Expert Grouping (EG), and Weight Merging (WM).
Stage
Metrie
SMoE-2
SMoE-4
SMoE-6
SMoE-8
SMoE-10
SMoF-12
Time Cust
2.35 min
4.61 min
6.54 min
8.40 min
10.30 min
12.30 min
PA
Memory Cost
1.23 GB
2.36 GB
3.48 GB
4.61 GB
5.73
6.86 GB
Time Cost
8.0 s
8.2
8.2
8.3
8.2
8.2
EG
Memory Cost
4.19 GB
5.29 GB
6.39 GB
7.48 GB
8.58
9.68 GB
Time Cost
23 ms
44 ms
66 ms
87 ms
109 ms
139 ms
WM
Memory Cost
1.33 GB
1.83 GB
2.32 GB
2.82 GB
3.31 GB
3.81 GB
A1.4 COMPARISON BETWEEN DIFFERENT PRUNING RATIO SCHEDULES
Our compression method for MC-SMoE uses a cubic schedule of pruning ratio, which is widely applied in many existing methods (Zhang et al., 2022; Zhu & Gupta, 2017: Sanh et al., 2020; Zafrir al., 2021). We con- duct extended comparison experiments with two other pruning ratio schedules, including linear and quadratic schedules, on the COPA and MultiRC tasks. The out- comes, shown in A A14, illustrate a performance ordering of cubic quurs)>=qadrrati>>> >ineer schedules. This is po- tentially because, in the early stages of pruning, an ag- gressive pruning schedule is less likely to lose useful information in the weights; while it is the opposite in the later stages of pruning.
Table A14: MC-SMoE performance eval- uation the switch-base-32 model with (linear, quadratic, cubic (ours)} schedules pruning ratio. We highlight the best per- formance over all baselines in bold.
Methods
COPA
MultiRC
Linear
59.00
73.83
Quadratic
61.00
73.92
Cubic (ours)
67.00
73.98
A2 MORE TECHNIQUE DETAILS
Supervised Fine- Tine-Tuning Hyper- Paramettrss Besides {batch size, learning rate, epoch counts} which vary for each task, we keep other tasks. These are shown in Table A15.
of supervised fine-tuning fixed for all
Table A15: Fine-tuning
of the switch-buse-32 model.
Hyper- Param--arssssrss
Values
Optimizer
ADAMW
Adam
1e-6
Adam
(0.9, 0.98)
Warm-up steps
16
Weight decay
0.01
LR scheduler
LINEAR DECAY
KD KD T
0.2
2.0
Details in Zero-Shot Learning We evaluate our approaches and baselines with the fairseq-moe 15b model in the zero-shot learning setting. Specifically, We use the language model to separately score each label choice, and pick the one with the highest score as the prediction. Although we utilize the training sets, they are only incorporated when essential in merging/compression, such as
A18

---

Published as a conference paper at ICLR 2024
when calculating the expert usage frequency. In short, optimization occurs any stage of the process, i.e. no fine-tuning at all.
Compression Hyper- ParameParseetess For M-SMoE, randomly pick 256 samples from training data to calculate both expert usage frequency and router-logits similarity for all tasks. For the com- pression phase in MC-SMoE, following Li et al. (2023), we adopt the cubic pruning ratio scheduler to control the S pruning process.
0st< To
P PT +1 (1 (1 LTT 3 (1-
T TT
PT
O.W.
where T is the total steps. Ti is the number of initial warm-up steps. Ty is the number of final cold-down steps. We set T to 10000, Tr to 400 and 7, to 1600 for tasks.
Knowledge Distillation In this paragraph we illustrate the detail of knowledge distillation (KD) applied in the supervised fine-tuning setting on all merged and compressed SMoE models for per- formance recovery, including our M SMoE, SMoS and all baselines. The goal is to force them, iee. the students, to imitate the outputs from the full SMoE model, i.e. the teacher. Specifically, the training objective can be formulated
min E (x.))~D [xx:; θ) + <<<u(<;; θ)], &
where the value of a is fixed at 0.2 for all tasks. the cross-entropy loss between predictions and the given hard labels, LKD is the KL divergence loss between the predictions and the full SMoE model's s soft labels:
CKD = XL P y x; @(ral))
||
æ;; 0)
Moreover, employ temperature T in the KL divergence to control the smoothness of the output distribution for both student and teacher models, defined as:
ps exp(2//T)
where 2j the logit score for class j. and the fixed at for all tasks.
The Router- Weight Similarity Function We provide detailed description of the reuttr- similarity function in this paragraph, which inferior to our adopted router-logits Section 3.1. Specifically, the similarity Sim(.,)) between experts E, and E, in SMoE layer is computed by:
Sim(E,, Ey) cosine "ws"" WH")
where W, is the router weight, and w and wi are vectors in it.
Expert Permutation Alignment We provide a detailed description of our expert permutation alignment here.
First, we introduce the permutation matrix P. which a square matrix where each row and column has exactly one element of 1, with all other elements being 0. It perseveres the functionality of the expert, a feed-forward network consisting of two linear layers Win- Wouts and an activation function act(-). This is because the equation Wout (W,.)) M (acttPaaaaaaaa always holds.
Second, we minimize the L2 distance between two experts to align them. Consider the first layer weights, denoted Wm. each of its rows of this matrix are identical; in that case, they would generate the same feature, disregarding any bias for now. Furthermore, if we have [u['],,. similar to [ it logically follows that neurons i and j would have a connection or association. Applying this concept to the second layer, Wonts this observation leads us consider optimization approach:
corresponds to individual hidden feature. Suppose two
argmin, veo[ (1) w(ll))
vec([PW(2)
(2(2) (())PT) || argmaxp (W(!) PW(() )
(((2) PT
Wout
Finally, this optimization constitutes 'linear assignment problem" (LAP), which be efficiently and practically solved by the Hungarian Algorithm (Kuhn, 1955). The Python-style pseudo code is included in A3.
A19

---

Published as a conference paper at ICLR 2024
A3 MORE IMPLEMENTATION DETAILS
We show some pseudocode demonstrate the implementation of our proposed M SMoM in PyTorch-like style.
Details of Merging Experts in SMoE Feed-Forward Layer In experiments, the final step of merging involves replacing one expert in a group with the derived weight. Instead of pruning the other experts, redirect the remaining ones in that group the newly substituted expert. This implementation ensures that the routing functionality remains consistent. Below is the PyTorch-style pseudo code:
def
fin: Sutceatrssstrrrsspppsssrrrr
group_labels: torch. torch.longieenssee Longleensr, usage_frequencies: torch.Flcate Flcattensor, Switchrtaanfoommessspaeeeee Each expert has group label and usage frequency
assert (iggopu_labell__
len [nnuusgg_frequenccee
len (ffn.experts))
for label in group_labels unique
expert_indices = torch. (group_ == label) [0]
with serch.no_gradun
Step 1. Calculate usage- frageee-reueec-wweiiihe
averaging
fol. fct_weight surch.sum (torch, stack
[ffn. [asge__feqqueccies (expert_eeeepp__________e_e__ _dxx for expert_idk in
usage
expert inttiidddces,, dim=0
) dim=0) torch. (usage sreepeiesssseeeeeeeeeeeeeeeeisss (expert. htepeeeeeeeeiceeeeeeeeeeeeeeeeee indices), dim=0)
fc2 Fch_weight - torch. (torchh
no2.weight htt.:/freguecccesseeeetettttccee for expert_idx in expert_indices), dim=0 ) dim=0) torch. sum
[ffn.experts
(expert_ din=0)
# Stop 2. Copy weight to the first expert in the group first_expert ffn. wff.exxertst httpeepppretteeeeet_edeeeeeeeeee (fg_: first_expert. fc2. fel. chp.gghttttcccccaeeee (fp__
# Step 3. redi other merged experts to the first one for expert idx in expert expert_indices|1:11 indices|l: fin. experss
ff"xperrtt /fh.exettssf"expprrt ( (expert_idxxe
ffn.experts [f"expert____ pexxper_lexpprr_tiiee eexpert_indices ["]
return fin
Details of Solving Expert Permutation Alignment The optimal permutation matrix for align- ing two experts is computed by minimizing the L2 distance between the expert weight matrices, which constitutes linear assignment problem. We utilize SciPy (Virtanen et al., 2020) to solve this optimization problem, and the Python-style pseudo code is shown below:
def httpr/sssttcceeeetaaaaaeeeeeeeee
reference_mipe oxttecarranntoomeeeDssssaaaaaaa target_mlp: wwrtchrrrrasfffmerrreeeeeeeeeeee
) -> torch.Tensor.
1sa_cost_matrix torch.m
reference mlp.mpp... wi weight.datai target_mii wwwght.tw.we....mm.mm........... data.t
+ torch.mm
reference
pettrence_mll.eeee wwwpww.weight wwight.daaa.t()
- perm Sunaa__su_asssinnne
isa_cost_matria cpu () numpy ()) maximize=True)
return torch. from_nunpy (perm) .te (Lea.coot_matri..eeev..))
A20

---

Published as a conference paper at ICLR 2024
A4
SUPPLEMENTARY EXPERIMENT RESULTS
A4.1 GROUPING RESULTS OF M-SMOE
We provide expert grouping results of the switch-base-32 model on all eight tasks including {SST2, MRPC, MultiRC, COPA, WinoGrande, SQuAD, WikiQA, HotporQA} here, shown in Fig- ure A6 A7 A8 A9 A10 A11 A12 A13 respectively.
Expert Grouping Layers SST2
12
10
10
Layer Index
Figure A6: Expert grouping results of the switch-base model on the SST2 task.
Expert Grouping SMoE Layers MRPC
12
6
10
Layer Index
Figure A7: Expert grouping results of the switch-baxe-33 model the MRPC task.
Expert Grouping SMoE Layers MultiRC
Group Index
10
Layer Index
Figure A8: Expert grouping results of the switch-buse-31 32 model on the MultiRC task.
A21

---

Published as a conference paper at ICLR 2024
Expert Grouping SMoE Layers COPA
KODDo
Group
XXXD
000
O00
10
2
000
Layer Index
Figure A9: Expert grouping results of the switch-buse--2 model on the COPA task.
Expert Grouping in SMoE Layers
WinoGrande
14
Group Index
10
4
Layer Index
Figure A10: Expert grouping results of the switch-base-32 model on the WinoGrande task.
Expert Grouping in SMoE Layers SQuAD
8
6
GELO
10
Layer Index
Figure All: Expert grouping results of the switch-base- 32 model on the SQuAD task.
Expert Grouping in SMoE Layers WikiQA
14
12
Group Index
Layer Index
Figure A12: Expert grouping results of the switch-bass base-32 model on the WikiQA task.
A22

---

Published as a conference paper at ICLR 2024
Expert Grouping SMoE Layers HotpotQA
10
10
Group
6
COCO
O
00148
000
(000
CCX(
10
Layer Index
Figure A13: Expert grouping results of the switch-base model on the HotpotQA task.
A23


Reviews:
[{"review_id": "uJsSssSi3x", "reviewer": "Reviewer_5S2f", "summary": "This paper proposes MC-SMoE, which is a compression method for Mixture-of-Experts models. The idea is to split experts into several groups, where only the most important expert is kept in each group. The authors further propose an algorithm to compress the merged experts. Experiments are provided to demonstrate the effectiveness of the proposed method.", "questions": "See above", "limitations": "", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* In this work, the authors study how to consolidate experts in MoE models. This is a very important topic since one of the major bottlenecks of deploying MoE models is the memory usage.\n\n* The problem studied in the paper is well-motivated. It is well-known that there are redundancies in MoE models. This paper leverages this finding and propose algorithms to compress MoE models.", "weaknesses": "Concerns and questions about presentation:\n\n* What is the intuition behind experts permutation alignment? Specifically, how is “alignment” defined? I can understand from performance-wise (Table 6) that alignment is needed. However, I do not understand how two experts are “aligned” using the proposed alignment method. To make the paper self-contained, please include the detailed algorithm used for this.\n\n* How are the results in Figure 3 computed? From my understanding, a Switch-base-32 model is first fine-tuned on each individual task, and then the activation frequencies are computed. Are the models fine-tuned with the load balancing loss [1]? It seems in Figure 3, loads of different experts are extremely unbalanced.\n\n* How is the stable-rank computed? For example, in a specific layer, the 32 experts are compressed into 6 experts. Do you compute the average stable-rank of the 32 experts as “before” in Figure 4, and the average stable-rank of the 6 experts as “after”?\n\n* I do not fully understand how experts are grouped. It is mentioned that “each non-dominant expert gravitates toward and joins the group led by its most similar dominant expert“. For example, suppose we have two dominant experts $E_1$ and $E_2$, then for a non-dominant expert $E$, do you calculate the similarity of $E$ with $E_1$ and $E_2$, and then assign $E$ to the more similar one? If this is true, is it possible that nearly all non-dominant experts are assigned to the same group?\n\n* The pruning of $S$ needs more justification. It is mentioned that “the weight columns with the lowest cumulative scores will be removed”. Why are weights pruned according to cumulative importance scores instead of importance scores? The pruning procedure in Appendix A2 also seems ad-hoc. How is this particular pruning schedule chosen?\n\n\nConcerns about experiments:\n\n* I would like to further understand the role of knowledge distillation. The authors mention that all the models (including the baselines) in Table 2 use knowledge distillation. Could the authors provide some results of the dense and full-SMoE models without distillation?\n\n* From Table 2, it seems performance of the model considerably drops after applying the compression technique (M-SMoE vs. MC-SMoE). The authors should provide more detailed analysis on the design of the compression method. For example, will different pruning strategies/schedules work better?\n\n* The authors mention “further memory and parameter efficiency” in the paragraph above Algorithm 1. However, no experiments are conducted to evaluate the speed and memory of the MC-SMoE models. The latency results in Table A10 indicate that there is only marginal speed gain of M-SMoE compared with full-SMoE. The authors should benchmark the inference speed (throughput) and memory usage of M-SMoE and MC-MoE, and compare the metrics with the dense and the full MoE models.\n\n[1] https://arxiv.org/pdf/2101.03961.pdf", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aPKlLLxMBj", "reviewer": "Reviewer_FC53", "summary": "This paper presents a mechanism for merging multiple experts by merging redundant experts while preserving as much as knowledge as possible. It is achieved by:\n1. Aligning the weights of any two given experts at a time using a permutation matrix since two models that are merely weight permutations of each other are equivalent. \n2. The router policy is used to group experts into groups of similar experts. Based on activation frequency, weight permutation-corrected experts in a group are merged together.  \n3. The final merged expert is further compressed by a low-rank decomposition and pruning of the incoherent part.", "questions": "1. Could you provide some theoretical insights into why the merging of experts should be done before compression since compression might be able to get rid of irrelevant information and make it more convenient to compare experts later for merging?\n2. Could you provide some theoretical background for your claim that similar experts would show similar router logits?\n3. An interesting line of investigation is the long-term scalability of how one could add more experts later during the life-cycle after applying M-SMoE.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The authors perform an extensive experimental analysis of each of their design decisions for:\n1. Their averaging strategy (Tab. 8)\n2. Need for permutation alignment (Tab. 6)\n3. Similarity function and the superiority of use of router logits (Tab. 4)\n\nThe authors also provide comparisons on multiple text datasets", "weaknesses": "A theoretical proof of either 1. Optimality of their expert merging algorithm or 2. An error bound on either the information loss/performance degradation based on their proposed algorithm would have significantly helped this work.\nAn analysis of the computational cost as the number of SMoE layers increases would be helpful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NYVjdgowRk", "reviewer": "Reviewer_zDGs", "summary": "The paper distills a large body of experts (in a mixture-of-experts model) into a few experts. This saves memory and improves fine-tunability of the final model. The core idea is to repermute neurons, group based on neuron \"routes\", then merge models. The authors show memory and FLOP reductions with negligible quality loss.", "questions": "- nit: the changing underline baseline is visually unappealing. not sure if there's a way to glue underlines to the bottoms of letters. this is a super-nit of course, doesn't really matter at all", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The figures are clear, the motivation is well-written, and the paper is overall well put together. \n- The paper presents a large array of experimental results and ablations; the results are convincing.\n- Figure 1: Impressive results. (A tradeoff curve might be nice? e.g., lines connect your points. also nit: The legend kinda blends in. Maybe give it a strong border or clearer background?)\n- Figure 2 is likewise well done. The lighted dotted lines are and spacing clearly separate the three parts, and the illustration of \"highly frequent\" to \"cluster center\" is helpful. (nit: I wish the fonts and sketch-esque style was applied to everything)\n- Figure 3's insight and accompanying visualizations are clear and insightful. Is this used later on anywhere? e.g., model can be compressed more aggressively for SST2 than for COPA.", "weaknesses": "- Experts permutation alignment (then computing similarity based on routes) is a big part of the paper, but the details are a bit lost on me. It could be worth adding more to 3.1, covering the basics of how the different possible permutations are searched.\n- In 3.1, it *seems like there's a chicken and egg problem -- we need alignment to know which experts are more similar *but we need to know which expert is the \"reference\" to re-align. How is that resolved?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)