Multi-Head Mixture-of-Experts
Authors: Xun Wu
Shaohan Huang
Wenhui Wang
Shuming Ma
Li Dong
Furu Wei
Venue: NeurIPS 2024 poster (2024)
Keywords: Mixture-of-Experts; Multi-head Mechanism
Primary area: deep_learning_architectures
TL;DR: An novel Mixture-of-Experts model utilizing a multi-head mechanism to augment the expert utilization and refine the granularity of understanding capabilities.

Abstract: Sparse Mixtures of Experts (SMoE) scales model capacity without significant increases in computational costs. However, it exhibits the low expert activation issue, i.e., only a small subset of experts are activated for optimization, leading to suboptimal performance and limiting its effectiveness in learning a larger number of experts in complex tasks. In this paper, we propose Multi-Head Mixture-of-Experts (MH-MoE). MH-MoE split each input token into multiple sub-tokens, then these sub-tokens are assigned to and processed by a diverse set of experts in parallel, and seamlessly reintegrated into the original token form. The above operations enables MH-MoE to significantly enhance expert activation while collectively attend to information from various representation spaces within different experts to deepen context understanding. Besides, it's worth noting that our MH-MoE is straightforward to implement and decouples from other SMoE frameworks, making it easy to integrate with these frameworks for enhanced performance. Extensive experimental results across different parameter scales (300M to 7B) and three pre-training tasks—English-focused language modeling, multi-lingual language modeling and masked multi-modality modeling—along with multiple downstream validation tasks, demonstrate the effectiveness of MH-MoE.

Multi-Head Mixture-of-Experts
Xun Wu, Shaohan Huang b-
Wenhui Wang, Shuming Ma, Li Dong, Furu Wei
Microsoft Research Asia
xunmu@iccrosoft.com com, Mhachanh@microsoon com, fuwei wweiiiccrooootccom
https://aka. wac.aa.ma//CC m
Abstract
Sparse Mixtures of Experts (SMoE) scales model capacity without significant increases in computational costs. However, it exhibits the low expert activation issue, i.., only a small subset of experts are activated for optimization, leading suboptimal performance and limiting its effectiveness in learning larger number of experts in complex tasks. In this paper, we propose Multi-Head Mixture-of- Experts (MH-MoE). MH-MoE split each input token into multiple sub-tokens, then these sub-tokens are assigned to and processed by a diverse set of experts in parallel, and scamlessly reintegrated into the original token form. The above operations enables MH-MoE to significantly enhance expert activation while collectively attend to information from various representation spaces within different experts to deepen context understanding. Besides, it's worth noting that our MH-MoE is straightforward to implement and decouples from other SMoE frameworks, making it easy to integrate with these frameworks for enhanced performance. Extensive experimental results across different parameter scales (300M to 7B) and three pre-training tasks Sikss-Ennlish-focusse language modeling, multi-lingual language modeling and masked multi-modality modeling-along with multiple downstream validation tasks, demonstrate the effectiveness of MH-MoE.
Introduction
and Large Multi- modal Models (LMMs) 37 27 have demonstrated their efficacy across various domains and tasks. To further enhance performance, reliable approach involves scaling up these models by augmenting the parameter count [13]. But for most of these densely-activated large-capacity models (referred to as Dense models), which utilize all their parameters to process all inputs, the extremely large size of these models significantly reduces inference speed, further limiting their practicality.
Large capacity models, such as Large Language Models (LLMs) 39 28 6 25
A promising alternative, facilitating model scalability while mitigating the burdensome computational costs, resides in Sparse Mixtures of Experts (SMoE) [31 12 (7. In contrast Dense model, S.oE contains parallel feed-forward neural networks (referred to experts) within each building block, and strategically activates distinet experts for specific input tokens via router, thereby yielding noteworthy efficiency enhancements. For instance, GShard [1] scales a Dense model from 2B to 600B parameters with lower training costs than 100B Dense model. And recently, Mixtral 7B [66], SMoE model containing 8 experts is shown to outperform matches LLaMA-2 70B 34 and GPT-3.5.
Despite its success, SMoE exhibits the low experts activation issue, which means that only a small subset of experts are activated during optimization and inference, e.g., 8.33% activation ratid shown in Figure I (a). while the majority of them are not used at all (see the dark area). As a result, SMoE
'Experts activation ratio shown in Figre (a) is the ratio of each expert's selection frequency in each MoE layer to the total number of tokens, where those exceeding threshold (<1) considered activated.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

fails to utilize the full expressive power of these experts, especially when the number of experts is large, which significantly limits effectiveness and scalability of SMoE.
Our aim in this paper is to achieve denser expert activation (i.e., better utilization of "dead" ex- perts) without increase in computational cost. To achieve this, we propose Multi-Head Mixture-of- Experts (MH-MoE).. The workflow of MH-MoE is illustrated in Figure Inspired by the multi-head mechanism utilized in Multi-Head Self-Attention (MHSA) block, MH-MoE splits each input token into multiple sub-tokens and distribute them to different experts. After expert processing, sub- tokens are seamlessly reintegrated into the origi- nal token form, thereby achieving denser expert activation, c.g., 90.71 % activation in Figure | (a), while also circumventing any additional compu- tational burden in subsequent non-parallel layers, e.g., MHSA block. Specifically, as shown in Fig- ure 2] when provided with single input token, MH-MoE activates four experts by splitting it into four sub-tokens, whereas SMoE only activates single expert.
Layats
Export
Expert
SMoE Activation: 8.33%
MH-MoB
Activation: 90.71%
1
RGB Input
Sub-tokens Assign
(b)
Figure 1: (a) Expert activation distribution
Furthermore, we observe an interesting phe- nomenon: for tokens with richer semantic infor- mation, the sub-tokens split from these tokens are more likely to allocated to distinct experts. For example, refer to the bright area in Figure (b).. where sub-tokens split from these patches are located to greater number of different experts, facilitating the capture of semantically-rit mation (e.g., the eagle in the figure). We therefore speculate that the allocation of sub-tokens to dis-
on XNLI [IO] corpus, encompassing 6 par- allel expert layers with 32 experts per layer. (b) MH-MoE showcases finer-grained under- standing by distributing sub-tokens split from semantically-richh patches to more distinct ex- perts to capture semantic information. Brighter regions indicate that sub-tokens split from this patch are distributed to greater number of dif- ferent experts.
infor-
tinct experts enables MH-MoB simultaneously focus information from various representation spaces within different experts, ensuring a more granular understanding for subtle differences in both vision and language patterns, finally achieving better finer-grained understanding ability. See in Figure g sub-tokens assigned to Experts and 2 capture detailed understanding of each character's actions within an image patch, while those assigned Experts and explicitly model the semantics of the false cognate 'camera'.
MH-MoE maintains following strengths: (1) Higher experts activation & better scalability. MH- MoE can alleviate lower expert activation problem and significantly enhance the usage larger experts by enabling optimization of almost all of experts, allowing for efficient scaling of model capacity. (2) Finer-grained understanding ability. By adaptively assigning sub-tokens to different experts based the semantic richness the token, MH-MoE enabling to jointly attend information from different representation spaces at different experts, and finally achieving better finer-grained understanding ability. (3) Seamless integration. The implementation of MH-MoE is remarkably straightforward and decoupled from other SMoE optimization methods (e.g., GShard (1)). making very easy integrate them together achieve better performance.
We evaluate the proposed MH-MoE on three model pre- pre-training and fine-tuning setting: English- focused language modeling, multi-lingual language modeling and masked multi-modality modeling, across different parameter scales (300M to 7B). Extensive experimental among these three tasks demonstrate the effectiveness of MH-MoE.

---

Inputs
Embodding
Taken Ansign
Export
Expers
SHEE Assion
M44 Koogn
Kratati
valuel toukler
eaa orteed (quatr
mhoagraa upr snpflati andl bobba
Expert2
Experti
2h pwweeth buvrr not cirincoriiee SULT SUTTE FISS CW LE bn seceglaris a
Experti
Figure 2: Workflow of MH-MoE. For vision data, different heads routed to different experts try to capture different aspects of details within patches and relations between patches. For language data, different heads attend to capture the varying contexts of false cognates across different languages (e.g., Italian and English) or polysemous words within a single language.
2 Background
Sparse Mixture-of of Experts (SMoE) 31 12 5 7 enhances model capacity while maintaining constant computational demand, thus achieving better performance than aensely-activvtte models on various tasks 22
19 39 28
Different from censely-activated models, each MoE layer consists of independent Feed-Forwart Networks (FFN) ( ) the experts, along with gating function (-) to model probability distribution indicating the weights over these experts outputs. For the hidden representation h E R of each input token, the gating value of routing h to expert f JFFN is denoted as:
N
(f.TN) = exp (h c,)
Σ exp cc), 5=0
where e, denotes the trainable embedding of the i-th expert and Σ Ei=0 9 (f!TN) = Then, the corresponding & experts, according the top-k gated values, are activated and the output O of the MoE layer is O h + Eieo 9 (J.PN) f.PN (h), where & denote activated experts set and |+] k.
As described above, the most commonly used routing mechanism involves selecting the top-k experts from N experts, where k < N [32], e.g., k = 2 and N = 2048 in GShard [21]. Such a routing mechanism allows the combination of data parallelism and expert parallelism. Some works [38 21 suggest that larger values of l often contribute to better model performance. However, with the increase in the value of k training models with conventional top-k routing implementation becomes much less efficient [21]. In this paper, we introduce MH-MoE, simple but efficient manner to make denser expert activation without an increase in computational complexity.
3 Method
3.1 Multi-Head Mixture-of-Experts
Concretely, we denote sequence of inputs tokens by X Rixdd where the number of tokens and d represents the length of token dimension. In MH-MoE, each parallel layer contains a set of N experts, each presented as { R + R % 1 Fi=o- where denotes the number of heads (i.e.. the number of sub-tokens a single token is split into), which is decoupled from the head in the multi-head self-attention layer. For clarity, we describe the operation of single MH-MoE layer here only.
The full architecture of MH-MoE can be seen Figure with parameter matrices W head Rdddd
First, X is projected by a multi-head layer
X = X W
head
where X E Rixd After that, every token in X split into sub-tokens along the token dimensions, and these sub-tokens arranged in parallel according to the original token sequence, forming new
3

---

Transfomer Block
Transformer Block
Hoken Soks
Parge LaVe
ExpanErrbnddigg
RadAng
Zvt Protatt
FF34
Melkar
Nult-nee Syer
Drackio
Haker 5035 Sequaros
Transformer
(a) SMoE
(b) MH-MME
Figure 3: Illustration of typical SMoE layer and the proposed MH-MoE layer. (a) A SMoE layer consists of a router and expert networks, where the experts are sparsely activated according to dot-product token-expert routing scores. (b) MH-MoE introduces additional two MLP layers, namely the multi-head layer and merge layer to split and merge tokens, respectively.
feature space X E [(xxh)x
x =
1xh
where function / denotes the token splitting operation: glxd (lxh)x ft and each sub-token is presented as x) E x meaning it is the the jih sub-token split from the 7th token. Then all these sub-tokens are fed into the gating function g-)). The gating value of routing a certain sub-token xj into the ndh expert is computed
N Σ 8
(1 = exp (x) ep)
exp (x) eg).
where ep E R the learnable embedding of the nthh expert. In this paper, we mainly focus on top-k routing, i.e., only the experts with the largest top-k routing score is activated. = Top, (g (fEFR) denote the set of activated experts and |$| = k. Then x,. is processed by these activated experts as following.
oj = x + Σ
(J
FFRN (xx)
5
pEt
After that, all obtained of are rearranged in the original order of sub-tokens and concatenated together as2:
txh
where O E R After that, is transformed back the into original token form by token merging operation f : (lx & - lilxd and then projected by a merge layer with parameter ma- trices W merge dxd to effective integration of multiple features of capturing detailed information from different expert representation spaces:
X = (O)T
W T
merge
Then we get the final output X of the single MH-MoE layer.
By implementing the aforementioned operations, we effectively increase the average volume of data routed to specific expert by factor of h (as demonstrated in Eq. 3. thus achieving denser expert
Sub-tokeent within regions marked with the same color split from the same token.

---

activation. Besides, the shapes of the input and output in the MH-MoE layer remain unchanged, thus no additional computational cost is introduced in the subsequent block. Specifically, we introduce hyperparameter /3 to scale the inner dimensions of each expert, aiming to balance the parameters introduced by the multi-head layer and merge layer, aligning the model's parameters and computational complexity with the original SMoF. Furthermore, the allocation of sub-tokens to distinct experts within MH-MoE enables us to collectively capture semantic information from diverse feature spaces across these experts, thereby enhancing the model's ability to achieve finer-grained understanding.
As the Pytorch- like style pseudocode of MH-MoE shown in Appendix D MH-MoE is characterized by its overall simplicity of implementation and decoupled from other SMoE optimization strategies [2] ST. making it easy to integrate with other optimized SMoE frameworks to enhance performance.
3.2 Training Objectives
Load balancing loss. To mitigate the expert load imbalance issue, given the sub-token set X (depicted in Eq. 3) and the frequency tp of how many sub-tokens are routed the pth expert, we follow existing works 21 13 to compute the load balancing loss Chalance via:
N
balance
Σ x
1p . 9 (
(8)
X
p=1
where N denotes the number of experts, |X| is the number of sub-tokens contained in X. (ftN) denotes the gating value of routing certain sub-token x). into the pt expert (see in Eq. 4p.
Task specific loss. The term Lusk is dependent on the particular task that MH-MoE designed to learn. For instance, during pre-training in the English-focused Language Modeling task, we utilize the language modeling loss 29 whereas the model predicts the next word in a sequence.
So, the overall training objective of MH-MoE is to minimize:
L
Lrask + aLCallace:
(9)
where a is coefficient for load balancing.
Experiments
4.1 Experimental Setup
Compared Baselines. We include two baseline models for comparison: (1) Dense, which represents Transformer decoder without the incorporation of sparsely-activated parallel modules (i.e., SMoE layer). (2) X-MoE, which is our implementation based on the SMoE proposed by [5]. We build MH-MoE upon X-MoE and use identical settings. Note that the all models are pre-trained using the
same training data and loss (Eq. [ as MH-MoE, and we ensure that the parameter count of MH-MoE remains consistent with lower than that of X-MoE, ensuring a fair and equitable comparison. detailed comparison about parameter and computational complexity can be found in Table [1
Pre-training Data. We detail the pre-training data of MH-MoE in three areas: (1) English-focused experiments use the RedPajama dataset [8], which is an open-source pre-training dataset. (2) multi- lingual tasks follow XLM (20) and use the multilingual Wikipedia training data. (3) multimodal tasks use a masked multi-modality modeling task with a large dataset of images, documents, and image-text pairs. Further details available in the Appendix A
Model Architecture and Hy pyperaarameters. For all experiments, we use the X-MoE [5] as our backbone architecture to build our MH-MoE, which has shown better performance than prior SMoE models such as Switch Transformers [3] on cross-lingual understanding benchmarks. For English- focused Language Modeling and Multi-lingual Language Modeling, we construct Dense, X-Mol and MH-MoE using the Transformer 36 decoder (L. = 12, H = 768, = 12) with the GPT-4 vocabulary
5

---

X0AeE with exports
will copett
XMOE Reppetss
Mit-Adst upparts
hiskite with ncuett
WH NoE with expers
20le0: with roperts
Xiit with experts
25.
XORCE: with 32 seperts
Peppaddy
MotE with 32 experts
Purpaaee
MH-ME with 32 experts
M4H Nof with expents
37.5
100%
302
1080
(a) Foglish-focused language modeling
(b) Mutti-liguull language modeling
(c) Masked musti-moddal modeling
Figure 4: Perplexity on validation dataset during the training phase reported for Dense, X-MoE and MH-MoE across three pre-training tasks.
Table 2: Accuracy/ / normaazzoooniiiiiiioon scores for language understanding tasks using the LLM
Evaluation Harness [4] N denotes the number experts.
Model
ARC-Challenge
ARC-Easy
RTB
BookQA
Winogrande
PIQA
BoolQ
HellaSwag TrutifuQQA (mel/mc2)
Avg
Dense
18.1/23.3
44.9/39.7
51.5
17.1/29.0
48.2
66.6
55.0
29.7/34.1
24.1/39.3
37.2
X-MoB (N 8)
19.0/24.7
48.3/42.0
52.7
17.4/29.8
67.9
31.4/35.7
24.3/40.2
38.7
MH-MoE (N =8)
19.6/25.2
50.2/42.2
53.0
18.2/30.3
51.1
68.7
59.6
33,2/40.3
24.7/40.9
39.8
X-MoB (N 32)
19.4/24.8
50,4/42.5
52.7
17.8/30.0
51.3
68.8
52.8
33,4/40.1
24.3/39.1
39.1
MH-MoE (N- 32)
21.4/26.8
50.6/44.8
5344
18.8/31.6
53.8
69.3
56.6
35.0/42.1
24.8/39.5
40.6
as the backbone architecture. The pre-training procedure takes 14 days on 2 NVIDIA DGX-2 Stations.
For Masked Multi-modal Modeling, we build Dense, X-MoE and MH-MoE following the same
Transformer encoder architecture BEiT v3 (37). The pre-training procedure takes 4 days on 2 NVIDIA DGX-2 Stations. For all three pre-training tasks, we set the head number 4. More details about architecture and training hyperparameters can be found in Appendix and
4.2 Perplexity Evaluation
Table I: Results of upstream perplexity evaluation. We report the validation per- plexity cross two setting: 8 experts and 32 experts.
We examined the validation perplexity curves for all pre- trained models and pre-training tasks under two expert settings (8 experts and 32 experts). The perplexity trends are depicted Figure
4 with the final perplexity values listed in Table 1 We can observe that as training pro- gresses: 1) the perplexity of our MH-MoE remained lower in comparison to the compared baselines, indicating more effective learning; 2) MH-MoE achieved the lowest per- plexity across three distinct experimental setups; 3) increase in the number of experts led to a corresponding decrease in the perplexity of MH-MoE, suggesting that the model benefits from enhanced representation learning ca- pabilities more experts are incorporated. These results collectively demonstrate the superiority of MH-MoF in terms of learning efficiency and language representation across multiple pre-training paradigms.
Perplexity J
Model
Experts 32 Experts
English-focusel language modeling
Dense (without Experts)
16.23
16.23
X-MoE MH-MoE (Ours)
14.82
11.96
12.72
10.28
Multi-lingual language modeling
Dense (without Experts)
8.56
8.56
X-MoE
7.19
6.02
MH- MoE (Ours)
6.26
5.09
Masked multi-modd modeling
Dense (without Experts)
17.95
17.95
X-MoE
16.34
12.68
MH-MoE (Ours)
14.73
10.87
4.3 Downstream Evaluation
English-focused Language Modeling. We evaluate our models on a total of 9 different zero- shot benchmarks to assess their abilitics across various natural language tasks like common sense reasoning, general language understanding and knowledge understanding using the LLM Evaluation Hamess [4]. As shown in Table 2 comparing X-MoE with the Dense model, X-MoE show notable improvement, indicating that SMoE models (e.g.. X-MoE) benefit from the large model capacity. average Overall, for all benchmarks, our MH-MoE obtains the best performance, achieving performance gain of 1.1 for 8 experts setting and 1.5 for 32 experts setting compared to X-MoE, demonstrating the effectiveness of proposed MH-MoE on modeling English-focused language.
Multi-lingual Language Modeling. We evaluate our multi-lingual language models on the cross- lingual natural language inference (XNLI) corpus [10), which is the extension of the multi-genre NLI

---

Table 3: Accuracy / accuracy--normaliztton scores on multilingual understanding tasks using the
LLM Evaluation Harness 14] denotes the number of experts.
Model
bg
de
el
en
es
fr
hi
n
SW
th
tr
ur
vi
zh
Avg
Dense
33.1
33.3
33.0
35.1
32.8
34.4
33.6
34.2
33.3
33.1
33.3
33.9
33.5
32,9
33.5
X-MoE (N
33.9
33.4
33.4
37.3
33.3
35,9
34,5
35.0
33.5
33.6
33.4
34.2
33.3
33.2
34.1
MH-MoE (N =
34.4
33.2
33,9
40.1
34,0
36.4
34,6
35.2
33.8
34,4
33.3
34.7
34.6
33.5
34.7
X X-MoE (N = 32)
34.5
34.5
33.4
39.6
33.1
35.3
34.1
35.4
33.6
34.7
33.7
33.6
34.5
33.3
34.5
MH-MoE (N 32)
35,8
35.6
34.1
40.7
33,9
36.7
34.4
36.3
34.3
36.0
34.1
34,3
35.2
33.6
35.3
Table 4: Results of visual question answering, visual reason- ing, and image captioning tasks.
13
Perplexity
Model
VQAv2
NLVR2
COCO Caption
20
Expert Number B
test-dev
test-std
dev
test-P
B@4
M
C
S
Expert Number 32
Dense
65.9
66.0
73.8
74.2
35.9
29.3
120.5
19.6
12
X-MoH
68.4
69.7
75.5
76.1
38.1
30.2
122.9
21.3
Head Number it Figure 5: Comparison results for dif- ferent head number h.
MH-MoE
70.1
71.4
77.0
77.8
39.7
33.1
124.1
23.0
(MultiNLI) corpus to languages. We follow the LLM Evaluation Harness pipeline and use the zero- shot setting to evaluate the multi-lingual ability. presents the zero-shot evaluation results on XNLI task. Similarly, X-MoE benefit from the large model capacity and show notable improvement compared with Dense model. Overall, MH-MoE obtains the best performance, surpassing X-MoB by an average performance gain of 0.6 for experts setting and 0.8 for 32 experts setting. Comparing MH-MoE with the X-MoE, it shows that MH-MoE models provide consistent gains on downstream tasks, demonstrating the effectiveness of our proposed MH-MoE on modeling tross-lingua natural language.
Masked Multi-modal Modeling. We evaluate on the widely used vision-language understanding and generation benchmarks, including visual question answering [15], visual reasoning [33 and image captioning 233 We report vqa-score on VQAv2, accuracy for NLVR2. For COCO image captioning, we report BLEU@4 (B@4), METEOR (M), CIDEr (C), and SPICE (S). Table E presents the evaluation results. For VQA task, MH-MoE outperforms both Dense and X-MoE by a large margin, e.g.. 4.24 and 1.69 points gain on test-dev split, respectively. For visual reasoning task, MH-MoE beats all these two baselines on both dev (1.5 points gain than X-MoE) and test-P split (1.7 points gain than X-MoE). For image captioning task, MH-MoE surpasses X-MoE by 4.2%, 10.2%, 9.4% in terms of @4, M and S, respectively. Above results show that X-MoE exhibits enhanced visual information comprehension, which also validates the effectiveness of our proposed MH-MoE in capturing diverse semantic information within text-image pair data.
4.4 Ablation Studies
This section presents experimental analysis to demonstrate the functionality of MH-MoF. In all comparative experiments, ensure parameter equality across models by adjusting the internal dimensions of the experts.
Number of Heads h. We conduct experiments by adjusting the number of heads (h = 2, 4, 6, 8, and 12) MH-MoE. As shown in Figure we find that across all settings h, our model consistently outperforms the X-MoF. Besides, as the value of h increases, we observe an initial improvement followed by a decline in our model's performance. This leads us to hypothesize that as initially increases, the enhancement in performance benefits from MH-MoE by activating greater number of experts, thereby enhancing the model's effectiveness and capturing a wider range of fine-grained token information. However, as h becomes too large, the excessive subdivision of tokens may impair their original semantic content, leading to decreased model performance.
Really Matters? The key motivation of MH-MoE is to split each token into several sub-tokens and then merge the sub-tokens after processing by experts. We use two MLP layers for the splitting and merging processes. We conduct a detailed analysis to determine whether the performance improvement due to the Loken-Split-Merginn (TSM) operation or the additional MLP layers.
Does the Token-Splitting-We
7

---

Table 5: Ablation studies of MH-MoË components: MLP layers and the Token- Splitting-Merginn (TSM, Eq. B and
Lert
Eq. 66 operation.
Model
MLP
TSM
Perplexity
16.23
De
x
16.40
X-MoE X-MoE
14,82
X-MoE
MH-MoE
MH-MoE
x
14.77
Figure 6: Distribution of expert activation in X-MoEE and MH-MoE on both Harness and XNL/ corpus, where
MH-MoRssee TSM
x
14.77
MH-MoE MLP
13:97
MH-MoE
12.72
h denotes the number of heads.
XM6E MoE
Token type
Porpleatty
Accuracy
Count
400
200
XMcE hellaswag
Mi-MMEE @ hellaswag
255
Experts
Number Experts
Assign diversity
(a) Upstream perplexity
(b) Downstream accuracy scores
(c) Sub-tokens assign
Figure 7: (a) Upstream training perplexity (L) when scaling the number of experts for both X-MoE and MH-MoE. (b) Accuracy scores on the hellaswag task when scaling the number of experts for both X-MoE and MH-MoE. (c) Comparison for sub-tokens assign diversity (the number of different experts they are routed to) for P&F and Non P&F tokens.
The results are presented in Table 5 A comparative analysis between Dense v.S. Densew/ MLP. as well as X-MoE v.s. X-MoE,,./ MLP: reveals that introduction of the MLP layer does not enhance the model's performance. Similarly, when comparing MH-MoE with MH-MoE, TSM it becomes evident that the inclusion of only the MLP, in the absence of the TSM, also does not yield any improvement in the model's effectiveness. The parameter quantities of the models being compared pairwise are equal.
An intriguing observation is made when comparing MH-MoE with MH-MoE MLP. Introducing TSM alone, without MLP, results in slight increase in model performance. In contrast, significant enhancement in model performance is only achieved when both MLP and TSM are incorporated simultaneously. We hypothesize that introduction of TSM, without the integration MLP, activates more experts, but the segmentation and merging of the model appears overly straightforward and abrupt in its execution. This limitation hinders the model's 's ability to meaningfully segment tokens into sub-tokens and effectively merge the diverse information gathered from different expert spaces.
We also conducted additional ablation studies examine the impact of varying the number of MLP layers, different splitting-merging methods and different balance loss. The details of these experiments are provided in Appendix
Analysis
Experts Activation. We visualize the activation of each expert varies across parallel expert layers for X-MoE and MH-MoE at Figure 6 It can be observed that: 1) X-MoB demonstrate a more skewed distribution, wherein significant portion of experts remain inactivated all the time. 2) Our MH-MoE achieves a denser expert activation compared to X-MoE, effectively mitigating the issue of low expert utilization. 3) As the number of heads f increases, the expert activation frequency in MH-MoE also rises.
Scalability. We explore the scalability for both X-MoE and MH-MoE by scaling up the number of experts from to 256 (about 7B parameters). For upstream performance, as shown Figure] (a), with the increase of experts, our MH-MoE could bring more gains. It is because MH-MoF. could mitigate the low expert activation problem effectively. With this ability, the superiority of the large-scale SMoE model will be better exerted, thereby achieving the improvement of the upper bound

---

Table 6: Performance across multiple pure vision tasks: classification (CLS) on ImageNet- 1k, object detection (OD) and instance segmentation (IS) on COCO. The number of expert is set to 8.
CLS
OD
2008
250k
Model
Table 7: Assign diversity of sub-tokens split
ACC
AP
APs0
AP
AP mask
from different patches in vision data with respect to training steps (100k + 200k + 250k steps). Brighter regions indicate sub-tokens split from
Dense
70.73
39.81
58,97
44,46
36.42
SMoE
75.66
42.23
60.30
44.58
37.50
MH-MoB
77.34
44.45
63.18
45.85
38.24
this patch are distributed to a greater number of diverse experts.
of SMoE with more experts. Detailed validation perplexity curves for these scaling up experiments can be found in Figure T0 at Appendix [F] For downstream performance shown in Figure 7 (b), for X-MoE, expert number = 64 is the upper bound, meaning that continuing to increase the number of experts will not bring performance gain. Our MH-MoE not only has a performance advantage over the X-MoE with the same number of experts, but also improves the upper bound from 64 to 256, which demonstrates the effectiveness of the scalability of our MH-MoE on downstream tasks.
Experts Assign within Token. We delve into a more granular analysis to validate how the multi- head mechanism aids MH-Mot in capturing diverse and intricate semantic information that is often challenging to comprehend, c.g., polysemous and false cognates words (denoted as PF tokens) in languages, and semantically-rich areas in images.
For languages data, we compute and compare the divergence levels (i.e., the number of different experts these sub-tokens are routed to) of sub-tokens split from PF tokens and Non-PF tokens. The results, presented in Figure (c), clearly demonstrate that the distribution of divergence for PF tokens is significantly skewed towards the right when compared to that of Non-PF tokens. This indicates that during the MH-MoE's inference process, PF tokens route their sub-tokens to a greater number of different experts, thereby capturing diverse semantic information in contrast to Non-PF PF tokens for a better polysemous and false cognates word modeling. Note that we utilized the GPT-4 API [25) to extract polysemous words and false cognates from the XNLI [10] corpus, and the corresponding prompt can be found in Table T2
For image data, we analyzed how the divergence levels of different patches evolve during the training process, as illustrated in Figure [ Notably, as training progresses, divergence levels gradually increase in high-frequency texture regions (or regions with rich semantics), while they decrease in low-frequency texture regions. This suggests that MH-MoE tends to route tokens from complex texture areas to wider variety of experts, thereby enhancing the finer-grained understanding of semantics in those regions. For more visualization examples, please refer to Figure in Appendix G
Experiments on Pure Vision Tasks, We conduct small-scale experiments on pure vision tasks to further validate the effectiveness of MH-MoE. We utilize ViT-B(ase) [1] as the Dense model and AdaMV-MoE [ as the SMoE baseline. AdaMV-MoE is multi-task vision SMoE model demonstrating a superior performance across various vision tasks. We build our MH-MoE upon AdaMV-MoE. The comparison results are shown in Table 6 which demonstrates that MH-MoE exhibits corresponding effectiveness and versatility in pure vision tasks.
Model Collapses. Concerns may arise that MH-MoE could revert to the original SMoE approach, routing all sub-tokens from the same token to a single expert. In MH-MoE's training framework, the load balancing loss Lralance (Eq serves as weak constraint, treating all sub-tokens as independent entities and uniformly distributing them among experts, promoting a balanced allocation. Our observations indicate that most sub-tokens from the same token are distributed among 3-5 different experts (see Figures] 7 and Figures 7 (c)).
Conclusion
In this paper, we study how can achieve denser experts activation without introducing additional cost, while improving the fine-grained understanding ability. With the proposed Multi- Head Mixture-of-Expere the simplicity of MH-MoE allows it to integrate with other SMoE frameworks to enhance performance easily. Extensive empirical results across three tasks demonstrate the effectiveness of MH-MoB.
we can easily implement the aforementioned functionality. Furthermore,

---

References
[1] H. Bao, L. Dong, F. Wei, W. Wang, N. Yang, X. Liu, Y. Wang. J. Gao, S. Piao. M. Zhou, and H. Hon. Uni Pseudo-masked language models for unified language model pre-training. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 642-652, PMLR, 2020.
[2] L. Caccia, E. Ponti, Z. Su, M. Pereira, N. L. Roux, and A. Sordoni. Multi-head adapter routing for cross-task generalization. arXiv preprint arXiv:2211.03833 2022.
[3] S. Changpinyo, P. Sharma, N. Ding, and R. Soricut. Conceptual 12m: Pushing web-scale image-text pre-training recognize long-tail visual concepts. In JEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021, pages 3558-3568. Computer Vision Foundation / IEEE, 2021.
[4] T. Chen, X. Chen, X. Du, A. Rashwan, F. Yang, H. Chen, Z. Wang, and Y. Li. Adamv-moe: Adaptive multi-task vision mixture-of-experts. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 17346-17357, 2023.
[5] Z. Chi, L. Dong, S. Huang, D. Dai, S. Ma, B. Patra, S. Singhal, P. Bajaj, X. Song, X.-L. Mao, et al. On the representation collapse of sparse mixture of experts. Advances in Neural Information Processing Systems, 35:34600-34613, 2022.
[6] H. W. Chung, L. Hou, S. Longpre, B. Zoph, Y. Tay, W. Fedus, Y. Li, X. Wang, M. De- hghani, S. Brahma, et al. Scaling instraction-finetuueed language models. arXiv preprini arXiv:2210. 14416, 2022.
[7] A. Clark, D. I. Casas, A. Guy, A. Mensch, M. Paganini, J. Hoffmann, B. Damoc, B. Hechtman, T. Cai, S. Borgcaud, et al. Unified scaling laws for routed language models. arXiv preprint DXXiv22022.01699 2022.
[8] T. Computer. Redpajama: An open source recipe to reproduce Ilama training dataset, 2023.
[9] A. Conneau, K. Khandelwal, N. Goyal, V. Chaudhary, G. Wenzek, E. Guzmán, E. Grave, M. Ott, L. Zettlemoyer, and V. Stoyanov. Unsupervised cross-lingual representation learning at scale. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, pages 8440-8451. Association for Computational Linguisties, 2020.
[10] A. Conneau, R. Rinott, G. Lample, A. Williams, S. Bowman, H. Schwenk, and V. Stoyanov, XNLI: Evaluating cross-lingual sentence representations. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2475 2485, Brussels, Belgium, Oct.-Nov. 2018. Association for Computational Linguistics.
[11] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[12) N. Du, Y. Huang, A. M. Dai, S. Tong, D. Lepikhin, Y. Xu, M. Krikun, Y. Zhou, A. W. Yu, O. Firat, et al. Glam: Efficient scaling of language models with mixture-of -x-eeerrsss... arXiv preprint arXiv:21 2021.
[13] W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232-5270 2022.
[14] L. Gao, Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, Hsu, A. Le Noac"h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang. L. Reynolds, H. Schoelkopf, A. Skowron, Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou. framework for few-shot language model evaluation, 12 2023.
10

---

[15] Y. Goyal. T. Khot, D. Summers- Stay, D. Batra, and D. Parikh. Making the V in VQA matter: Elevating the role of image understanding in visual question answering. In 2017 IEEE Confer- ence Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017, pages 6325-6334. IEEE Computer Society, 2017.
[16] A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. BBmford, D. Chaplot, D. d. I. Casas, E. B. Hanna, F. Bressand, et al. Mixtral of experts. arXiv preprint F4Xv:2201..04088, 2024.
[17] R. Krishna, Y. Zhu, O. Groth, J. Johnson, K. Hata, Kravitz, S. Chen, Y. Kalantidis, L. Li, D. A. Shamma, M. S. Bernstein, and L. Fei-Fei. Visual genome: Connecting language and vision using crowdsourced dense image annotations. Int. J. Comput. Vis., 123()):32-73, 2017.
[18] T. Kudo and J. Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 66-71, Brussels, Belgium, Nov. 2018. Association for Computational Linguistics.
[19] K. Kumatani, R. Gmyr, F. C. Salinas, L. Liu, W. Zuo, D. Patel, E. Sun, and Y. Shi. Building great multi-lingual teacher with sparsely-gated mixture of experts for speech recognition. arXiv preprint arXiv:21 12 12.58200, 2021.
[20] G. Lample and A. Conneau. Cross-lingual language model pretraining, 2019.
[21] D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16666, 2020.
[22] D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen. (har: Scaling giant models with conditional computation and automatic sharding. In International Conference Learning Representations, 2021.
[23] T. Lin, M. Maire, S. Belongie, J. Hays, P. Perona, D. Ramanan, P. Dollár, and C. L. Zitick.. Microsoft COCO: common objects in context. In D. J. Fleet, T. Pajdla, B. Schiele, and T. Tuytelaars, editors, Computer Vision ECCV 2014 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V, volume 8693 of Lecture Notes in Computer Science, pages 740-755. Springer, 2014.
[24] Y. Liu, M. Out, N. Goyal, Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer. and V. Stoyanov. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/ /ttps/997....999 2019.
[25] OpenAl. Gpt-4 technical report. ArXiv, abs/2303.087744 2023.
[26] V. Ordonez, G. Kulkarni, and T. L. Berg. Im2text: Describing images using million captioned photographs. In J. Seowe-Taylor, R. S. Zemel, P. L. Bartlett, F. C. N. Pereira, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems 24: 25th Annual Conference on Neural Information Processing Systems 2011. Proceedings of a meeting held 12-14 December 2011, Granada, Spain, pages 1143-1151, 2011.
[27] Z. Peng, W. Wang, L. Dong, Y. Hao, S. Huang. S. Ma, and Wei. Kosmos-2: Grounding multimodal large language models to the world. arXiv preprint arXiv:2306. 14824, 2023.
[28] H. Pham, J. Kim, S. Mukherjee, D. P. Woodruff, B. Poczos, and H. H. Awadalla. Task-based moe for multitask multilingual machine translation. arXiv preprint arXiv:2308. /5772, 2023.
[29] A. Radford, K. Narasimhan, T. Salimans, and Sutskever. Improving language understanding by generative pre-training. 2018.
[30] P. Sharma, N. Ding, S. Goodman, and R. Soricut. Conceptual captions: A cleaned, hypernymed, image alt-text dataset for automatic image captioning. In I. Gurevych and Miyao, editors, Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics, ACL 2018, Melbourne, Australia, July 15-20, 2018, Volume I: Long Papers, pages 2556-2565. Association for Computational Linguistics, 2018

---

[31] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gater mixture-of exxeeessss layer. In International Conference on Learning Representations, 2017.
[32] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean. Outra- geously large neural networks: The sparsely-gate mixture-of-experts layer. arXiv preprint arXiv.1701.06538 2017.
[33] A. Suhr, S. Zhou, A. Zhang, I. Zhang, H. Bai, and Y. Artzi. A corpus for reasoning about natural language grounded in photographs. In A. Korhonen, D. R. Traum, and L. Mårquez, editors, Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, pages 6418-6425 Association for Computational Linguistics, 2019.
[34] H. Touvron, Martin, K. Stone, Albert, A. Almahairi. Y. Babaei, N. Bashlykov, Batra, P. Bhargava, S. Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09228, 2023.
[35] H. Trinh and Q Le. A simple method for commonsense reasoning. ArXiv, wwwp//80..02...0 2018.
[36] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and
I. Polosukhin. Attention is all you need. In Advances Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, pages 5998 5998-6008, 2017.
[37] W. Wang, H. Bao, L. Dong, J. Bjorck, Z. Peng, Q. Liu, K. Aggarwal, O. K. Mohammed, Singhal, S. Som, et al. Image as a foreign language: Beit pretraining for all vision and vision-language tasks. arXiv preprint arXiv:2208. 10442, 2022.
[38] A. Yang. J. Lin, R. Men, C. Zhou, L. Jiang, X. Jia, A. Wang, J. Zhang, J. Wang. Y. Li, et al. M6-t: Exploring sparse expert models and beyond. arXiv preprint a0Xiv:2 15082, 2021.
[39] X. Zhao, X. Chen, Y. Cheng, and T. Chen. Sparse moe with language guided routing for multilingual machine translation. In Conference Parsimony and Learning (Recent Spotlight Track), 2023.
[40] Y. Zhu, R. Kiros, R. Zemel, R. Salakhutdinov, R. Urtasun, A. Torralba, and S. Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference on computer vision, pages 19-27, 2015.
12

---

A Pre-training Data of Masked multi-modal modeling task
Table g presents of pre-training data Masked multi-modal modeling task. For multi-modal data, there are about 15M images and 21M image-text pairs collected from five public datasets: Conceptual
12M (CC12M) 3 Conceptual Captions (CC3M) [3), SBU Captions (SBU) 26 COCO 23 and
Visual Genome (VG) [17]. For monomodal data, we use 14M images from ImageNet-21K and
160GB text corpora from English Wikipedia, BookCorpus [40, OpenWebTex CC-News [2
and Stories 35
Table 8: Pretraining data of Masked multi-modal modeling task. All the data are academically accessible.
Data
Source
Size
Image-Text Pair
CC12M, CC3M, SBU, COCO, VG
2 21M pairs
Image
ImageNet-21K
14M images
Text
English Wikipedia, BookCorpus, OpenWebText, CC-News, Stories
160GB documents
B Model Hyperparameters of Language modeling tasks
Table presents the model hyperparameters of X-MoE and MH-MoE for both Aaglish-focusedd lan- guage modeling and bulti-lingual language modeling tasks. The gating temperature is initialized as 0.3 and 0.07 for the softmax gating and sigmoid gating, respectively. We use the same vocabulary as XLM-R 9 with 250K subwords tokenized by SentencePiece [8]
Table 9: Model hyperparameters of Dense, X-MoE and MH-MoE. The SMoE frequency refers to how many experts each token will be assigned to, i.e. the value of k in the Top- expert selection.
Hyperparameters
Dense
X-MoE
MH-MoE
FFNs within layer
2
Expert embedding dimension
16
16/h
Initialized gating temperature TO
0.3
0.07
0.3 / 0.07
Transformer blocks
12
12
12
Hidden size
768
768
768
FFN inner hidden size
3,072
3,072
3,072 x
Attention heads
12
12
12
SMoE frequency
13

---

C Hyperparameters for Pre-training
Table T0 presents the hyperparameters for pre-training across three tasks: Language modeling
tasks (English-focusee language modeling and Multi-lingual language modeling tasks) and Masked multi-modal modeling task.
Table 10: Pre-training hyperparameters for Language modeling tasks (English- language
modeling and Multi-lingual language modeling tasks) and Masked multi-modal modeling task tasks.
Hyperparameters
Language modeling tasks
Multi-modality modeling task
Batch size
256
512 AdamW W
Optimizer
Adam
Batch tokens per task
IM
Adam
le-6
le-6
Adam
(0.9, 0.98)
(0.9, 0.98)
Maximum learning rate
5e-4
2.8c-3
Learning rate schedule
Linear decay
Cosine decay
Warmup steps
10,000
10,000
Weight decay
0.01
0.05
Transformer dropout
0.1
Dropout
0
Load balancing coefficient
le-2
le-2
Table 11: Parameter count setting of X-MoE and MH-MoE in our experiments for English-focused language modeling. Multi-lingual language modeling and Masked multi-modalin modeling tasks. "non-expert param" refers to the parameters that are not part of the expert networks, such as the attention layer, router, etc., while "expert params" represents the total number of parameters in the parallel expert networks. For Dense models, since there are no expert network layers, we only list the total number of parameters. All models under the same task utilize the same architecture and hyperparameters following identical training settings and steps.
Expert Setting
Dense
X-MoF
MH-MoR
Sum
non-expert params
expert params
Sum
non-expert parmms
expert params
Sum
fngliss-focused language modeling
expert
162M
8 experts
134M
227M
361M
141M
213M
354M
experts
454M
588M
141M
S71M
experts
134M
908M
LOI2M
141M
898M
1.039M
64 exports
134M
1.815M
1,949M
141M
1.797M
1,938
128 experts
3,631M
12300
256 experts
134M
7,263M
7,397M
141M
Mudti Jingual language modeling
expert
162M
experts
134M
227M
361M
141M
213M
M5SE
experts
900M
Masked mudii- modeling
expert
277M
experts
191M
323M
514M
195M
310M
505M
32 exports
191M
1,037M
1,228M
195M
1,014M
1,209M
14

---

D PyTorch-style Code
We also provide the PyTorch-style code Algorithm to explain our MH-MoE, which including two main aspects: 1) Stage 1. The creation and initialization of multi-head layer and merge layer. 2) Stage The segmentation of tokens, followed by processing through an expert network, and ultimately merging.
Algorithm The Overall Procedures MH-MoE a PyTorch-lite -ike style.
Input: A MH-MoB model with L parallel SMoE layers M, the number of the experts k.
# Stage 1: Initial parameter of multi-head layer & merge layer
for i in (1nge L):
M[i] www.mmliiiaaa.. www.lmttteaaadaaaaeaeaa__aeee haedii_eaa___ay_eeee layer nn Linear (hidden dim, hidden_dimp
M[i] mmerge_layer nn Ln.ein (hidden_din, hidden_dim)
Initialization
nn. seisstxavier__uiii uniform. [M[i] multii wwwuttueead_____eeeeeeeaaeaaeeee layer. weight, gain=1 math. bath.sgrt(a (2)) nn. www.iiiitttteer.eeemmmmmm
(M[i] lereereellye..... layrge__ayer..eeghh) weight)
nn. init. wonstantt (M[i) merge merge_layer bias, 0.0)
Stage 2: The segmentation and merge of tokens for the i-th MH-MoE layer
MHMoE_Layer "".
Input:
Tensor shape: (batch_size Length, sidden_dimg mask Tensor shape: (batch_ size, Length)
Output:
o Tensor shape: (batch_size Length, hidden_dim)
heads: head number of multi_hear layer ">>
Processed by multi-head layer [[i]] wwwmiihhaaddlaleeeeeaeee (x)ee_xa
Split token & rearrange sub-tokens in parallel
x. reshape (batch_sike size Length heads, hidden dim // heads) (eaad).coottieeeeeeeeeeeeee contiguous
mask = mask. reshape--,, 1) reeeettt(1, heads) reshape (batch_size Length heads)
Standrad SMoE routing block
x, mask router(z, (x, mask)
Merge back the original token form reshape bbatch_size Length, heads, dim // heads) rrehaape (batch_siz) Length, dim) contiguous()
M[i] [[a].merge__ayye meyge__ayerr())
return
15

---

Table 12: Prompt template for identifying polysemous and false cognates in different languages.
Your role identify polysemous and false cognatee in different langueges from the given textusl input (AMA Input), Note that "Polysemous" refers word having two or more completaly different meenings (for exemple, "grouse" meanings related complaining end also refers type bird), while "felse cognatiss different languages" refers words in different languages that have similar forms carry different meanings (for example, English and Italian, "camera" looks similar but represents different semantic concepts).
## Input (Prompt)
provide your identify results following format:
Word: [Make sure that there is only a word here.]
Rationsle: [Rationale why
polysemous cognates, step step)
Output
Word: [Make sure that there is only a word here.]
Rationale: (Rationale why this word polysemous or false cognates, think stop stop)
Table 13: Comparison results for different numbers of MLP layers n. The results are averaged over five runs.
15.5
15.1
15.0
14:86
Perplexity
14.5
Upstream
Downstream
n
14.0
13.92
Perplexity
RTE
PIQA Winogrande
0
13.97
52.9
68.2
51.7
13.5
Mi+-MoE
MH.MoE 5-11)
MH.MoE 5-11
MH-MoEij 3.5
MH Mit-MEE w.s
1
12.72
53.4
69.3
53.8
2
12.66
54.0
68.8
53.3
Comparison results for different splitting method. K denotes the size of kernel while denotes the size of stride in Conv ID.
Figure 8:
3
12.87
53.1
68.8
52.7
E
Ablation Studies
Number of MLP layers. We explore the impact of varying the number of layers (n = 0, 1, 2, 3) MLP on MH-MoE performance. For configurations exceeding single layer, ReLU activation functions were incorporated between MLP layers to ensure the non-linearity of transformations. The parameter quantities of the models being compared are equal. Upon analyzing the results in Table we observe that increasing the number of MLP layers beyond one had negligible impact on the model's 's performance. This indicates that single-layer MLP is sufficient for accomplishing token segmentation and fusion.
What Makes Good Splittigg-Merging Method? We designed ablation experiments investigate whether the tokenization method affects the MH-MoE's s performance. We replace the multi-head layer and merge layer (both are FC layers) with Conv 1d layers (denoted as MH-MoE*)..
We conducted experiments in two settings [English-focusee language modeling and Multi-lingual language modeling tasks) and explored different levels of randomness by using different kernel sizes and strides for the Conyld d layers. We ensure that the parameters are kept as consistent as possible by adjusting the number of Conv layers and the dimensions of experts. Through experimental results presented at Figure 8 we observe that replacing splitting method with Conv 1d (denoted as MH-MoBFF resulted in significantly inferior performance compared the original MH-MoE. This underscores the importance of tokenization methods, indicating that tokenization methods such as Convld, 1d, which disrupt the original input sequence features, are not suitable for this context.
Impact different balance loss We conducted experiments with three groups to assess the impact of different constraints sub-token allocation: (1) MH-MoE MH-MoE trained without the Lbalunce (2) MH-MoBt: MH-MoE trained with the Lislance, and (3) MH-MoE!: MH-MoE trained with both Loslunce and diversity loss Lav. Lai is designed enhance sub-token allocation diversity and defined as: given a set of sub-tokens split from the same token, with P. representing the
gating probability distribution of the 5th sub-token, Laiv = - - E KC((P.)),),, where
16

---

14.4
14,23
8.00
Perplexity
14.2
Perplexity
7.75
7.50
7.41
14.0
7.35
13.92
13.89
7.25
13.8
7.00
MH-MoE1 MH-MoET MH-MoE
MH-MoE" MH-MoEt MH-MoE
Figure 9: Experimental results for varying level constraints sub-token allocation across tasks: focused language modeling (left) and Multi-lingual language modeling (right).
expert
MH+MME with experts
XMoE with 128 experts
XMME with 52 experts
MH MoE with 128 expers
MH+MME
exports
experts
Porplexity
120
2006
Training Step
Figure 10: Validation perplexity reported for both X-MoE and MH-MoE.
K.C(.) indicate Kulbback-Leibler divergence. This constraint increases the inconsistency in gating distributions among sub-tokens in enhancing expert allocation diversity.
The experimental results (perplexity) are summarized in the Figure 9 Key findings include: (I)
Without any constraints (MH-MoE*)) the model's performance degrades, likely due severe imbalance, causing some experts to receive insufficient data, resulting in suboptimal performance. (2) The addition of Lav did not significantly improve model performance compared using only Cbalance. We hypothesize that Lialance alone provides sufficient diversity in sub-token allocation.
F Visualization of training perplexity
We provide the training perplexity curve for model training in the experimental setting of increasing the number of experts (from 8 to 256) in Figure 10
G
Visualization
We provide more visualization of variation assign diversity for sub-tokens split from different patches in vision data at Figure 11
17

---

10k stepž00k stepß50k steps Input
100k stepš00k stepš50k steps
Input
100k stepš00k
stepž50k steps
Input
Figure 11: More visualization examples for assign diversity of sub-tokens split from different patches with respect training steps. We analyze MH-MoE with heads (h=8) the subject. Brighter regions indicate that sub-tokens from this patch are distributed to a greater number of diverse experts, while darker regions indicate that sub-tokens are assigned more of the same experts.
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer:[Yes
Justification: Please refer our abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [NA]
If we discover limitations, we will include them in the latest version of the paper.
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [No]
19

---

Justification:
In our paper, we do not make any assumptions, and consequently, there are no associated proof processes.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide detailed explanations of the model training and data specifics in Sections 3 and Appendix A and C
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
20

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: provide the code for model training and testing the supplementary materials. The training data used are publicly available, and we detail the data specifics in the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nnis..c)
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Please refer to Sections 3 and Appendix A
B and C
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[NA]
Justification: All results reported as the average of five randomly initialized tests.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
21

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Please refer to Sections 3 and Appendix A and
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://ppp...arrs.....///////taa
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
Our model serves the foundational architecture for large models rather than their socio-economic applications. As such, it operates on a different level, and we do not provide an analysis of societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
22

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [No]
Justification:
We will conduct this before our model is publicly released.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Please refer to Section 4
Guidelines:
The answer NA means that the paper does use existing assets. The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, thcoe com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
23

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Please refer the README.md in supplement code. Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moddl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.


Reviews:
[{"review_id": "NTUQUsRHMb", "reviewer": "Reviewer_Rrmq", "summary": "The paper proposed a new architectural extension to large pre-trained models (e.g. LLMs, or Multi-modal models). Specifically, the focus is on Mixture of Experts (MoE) models and the paper proposes a new method by introducing subtokens and routing each subtoken to experts.\n\nThe resulting architectural change performs better on down-stream tasks and demonstrates a higher use of a variety of experts.\n\nThe author's also conduct a variety of ablation studies to test various aspects and changes of their proposed architectural change.", "questions": "Questions:\n1. How does your model compare against some popular & state-of-the-art models (e.g. MoE LLMs like Mistral) in terms of performance?\n2. Referring to \"Weakness 4.b\", what is your explanation of why your method actually works better? \n3. What is the actual cost of running standard MoE vs. yours (in terms of practical training or inference time) - given the same computation budget.", "limitations": "Limitations:\n1. Limitations are not really considered in the paper.\n2. It would be really interesting to understand what does NOT work in the method in practice?\na.  What did not work during developing of the method?\nb. When is the method worse than previous other methods?\nc. Computational limitations?\nd. etc. etc.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Strengths:\n- Achieves strong performance on down-stream tasks and seems to out-perform the baseline methods presented in the paper.\n- Paper claims that they achieve better results with same cost\n- Good reporting on experiments for reproducibility\n- Decent amount of experiments and evaluations", "weaknesses": "Weaknesses:\n1. No actual comparison in terms of \"cost\" is given except for parameter count and complexity calculations in the appendix. It would be interesting to see actual evaluation of cost (e.g. computational cost of training & computational cost of inference vs. MoE alone).\n2. Evaluation and comparison is not done on publicly available and popular (& stat-of-the-art) models of relevant sizes (e.g. Mistral, Aya, Gemma ...). Therefore, it is not clear if this method actually performs better or just on the models and training runs presented in this paper.\n3. Clarity\na. some parts are unclear, e.g. equations 2,3,6,7 speak about the new additions of the method, however, it is not clear at all that there is MLP layer present (only later in the paper it is mentioned), however, it is \"still\" not clear whether and how it applied during equations 2,3,6,7.\nb. some mistakes in the paper e.g. Paper checklist 7, the answer is N/A. However, N/A means no experiments were run, the correct answer should be No.)\n4. The core claim that the problem with MoE is: \"low experts activation issue\" has a few problems:\na. It is never clearly explained why this is an issue (and no evidence is given that this is a real issue).\nb. On one hand the paper presents that the method increases activations, however, the paper also shows that increasing the number of heads (and therefore activations) seems to also decrease performance after a specific number of heads [Figure 5, Section 4.4]. (Great for reporting this, however, this also challenges the introduction / proposition of the method). I.e. it seems that \"low experts activation\" is not always a problem, or only a problem to a certain degree. => Therefore, it would be interesting to revise this hypothesis, and also to provide a counter hypothesis (with evidence).\n5. Reporting on sharding of the models on the DGX-2 system is not quite described. (Which is quite essential to run it sensibly.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful to the reviewer for the extensive review. Due to space constraints, we address your questions in the **Rebuttal** below as well as at the top of this page in the **Author Rebuttal**:\n\n---\n>**Q1**: What is the actual practical training or inference time of running MoE vs. MH-MoE.\n\n**A1**: **For training cost evaluation**: In fact, if traditional static routing [1] (i.e., using zero-padding or dropping tokens to force a uniform distribution of tokens among experts) were applied in MH-MoE, the situation would be different. Since MH-MoE uses sub-tokens for more complex and fine-grained routing decisions, the matrix dimensions required to record routing information expand accordingly during training. This results in more time-consuming matrix operations and higher all-to-all communication costs within nodes, thereby affecting the end-to-end throughput of training MH-MoE. We present a comparison of throughput between XMoE and MH-MoE using static routing (denoted as MH-MoE*) in Table A. Our results show that MH-MoE with static routing is slower than the baseline.\n\nTo address this issue, we have implemented dynamic routing [2] (refer to the implementation: [databricks/megablocks (github.com)](https://github.com/databricks/megablocks)) in our MH-MoE model. Dynamic routing alleviates the constraints imposed by static routing on both software and hardware, enabling a more efficient and unrestricted routing strategy. Table A shows a comparison of the training throughput taken by MH-MoE with dynamic routing (denoted as MH-MoE) and the baseline. We observe that using the more efficient dynamic routing approach results in end-to-end training throughput for MH-MoE that is nearly identical to the baseline.\n\n***Table A. End-to-End Throughput on eight GeForce RTX 3090 GPUs.*** \n\n|Models|End-to-EndThroughput(toks/s)|\n|--|--|\n|SMoEBaseline|32.21|\n|MH-MoE*|25.77|\n|MH-MoE|31.93|\n\n**For inference cost evaluation**, we implemented our MHMoE model on the Mistral 7*8B architecture and used vLLM (https://github.com/vllm-project/vllm) to test the inference speed. vLLM is a widely used LLM inference framework. We tested the inference speed on four A6000 GPUs. All GPUs processed the same prompts and used the same sampling parameters (temperature=0.8, top_p=0.95), with the tensor parallel size set to 4. To ensure the fairness of the inference speed test, given that the MHMoE model outputs garbled text, we set a consistent output length. The results are shown in Table B. We can observe that the inference speed of the MHMoE model is slightly slower than that of the SMoE model. Additionally, as the number of heads in MHMoE increases, the inference speed decreases slightly. However, the difference is not significant, and the MHMoE model still achieves a high inference speed.\n\n***Table B.*** \n\n||InputSpeed(toks/s)|OutputSpeed(toks/s)|\n|------|--------|--------|\n|SMoE|18.47|45.47|\n|MHMoE(head=2)|18.25|44.92|\n|MHMoE(head=4)|18.18|44.74|\n\n**Reference**\n\n[1] Lepikhin, Dmitry, et al. \"Gshard: Scaling giant models with conditional computation and automatic sharding.\" arXiv 2020.\n\n[2] Gale, Trevor, et al. \"Megablocks: Efficient sparse training with mixture-of-experts.\" PMLR 2023. \n\n---\n> **Q2**: Evaluation and comparison is not done on publicly available models of relevant sizes (e.g. Mistral).\n\n**A2**: **The training data and steps for the models you mentioned are either lengthy or not well-defined, making them difficult to compare directly**. To ensure fairness, we have used our own implementation of the **two MoE methods** for comparison:\n\n1. For English-focused language modeling, multilingual language modeling, and masked multimodality modeling tasks, **we implemented our MH-MoE on XMoE [1]** and conducted extensive experiments across these three tasks (ranging from 300M to 7B parameters, with 3 upstream tasks and 26 downstream tasks) to validate the effectiveness of MH-MoE.\n2. For pure vision tasks, as detailed in Section 5 of the main text, **we implemented MH-MoE on AdaMV-MoE [2]**. We validated that MH-MoE can further enhance AdaMV-MoE's performance in vision tasks such as classification, object detection, and instance segmentation. \n\nWe chose XMoE and AdaMV-MoE because they have been proven superior to vanilla SMoE in extensive experiments. **We believe that the extensive experiments on these multiple upstream & downstream tasks with the two MoE structures demonstrate the effectiveness of MH-MoE.** \n\n---\n\n>**Q3**: Clarity a. some parts are unclear.\n\n**A3**: Thank you for pointing this out. We apologize for the unclear or incorrect descriptions in a and b. In fact, in all experimental setups (except for the ablation experiments on MLP layers in Table 5 and the number of MLP layers in Table 13), the MH-MoE model uses both the multi-head layer and the merge layer. These two MLP layers are represented by the $W_\\{head\\}$ and $W_{merge}$ matrices in Equations 2 and 7, respectively. The multi-head layer and the merge layer perform matrix transformations on the last dimension (hidden state dimension) of the input tokens and the merged output tokens for feature modeling. We will revise these sections in the updated version of the paper to ensure a clearer and more accurate presentation.\n\n---\n>**Q4**: The core claim that the problem with MoE is: \"low experts activation issue\" has a few problems: a. It is never clearly explained why this is an issue. b. the paper also shows that increasing the number of heads (and therefore activations) seems to also decrease performance after a specific number of heads [Figure 5, Section 4.4]. (Great for reporting this, however, this also challenges the introduction / proposition of the method). I.e. it seems that \"low experts activation\" is not always a problem, or only a problem to a certain degree?\n\n**A4**: Pleasee see in the **Author Rebuttal section**.\n\n---\n>**Q5**: Reporting on sharding of the models.\n\n**A5**:  Pleasee see in the **Author Rebuttal section**", "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": "xYglPDcu8U", "reviewer": "Reviewer_Y9UT", "summary": "This paper introduces Multi-Head Mixture-of-Experts (MH-MoE), a method for training MoE models with enhanced expert activation. In particular, MH-MoE splits each input token into multiple sub-tokens, which are processed by a set of experts in parallel, and then merges them back into their original token form. The results show improved performance over dense and X-MoE models in parameter-matching settings.", "questions": "Please see the reasons to reject.", "limitations": "No. authors did not disclose any limitations of their method. See \"Weaknesses\" section for more details.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The authors propose a simple yet effective idea of splitting tokens into sub-tokens to boost expert activation. The paper is clear and well-written. The authors verify their approach on a variety of tasks, including English-focused language modeling, multi-lingual language modeling, and masked multi-modality modeling, across different parameter scales (ranging from 300M to 7B by scaling the number of experts).", "weaknesses": "1. Experiments are performed on a limited number of models and baselines, making the results less convincing.\n\n- The experiments are limited by a single model architecture. The authors apply their approach by modifying the previously released X-MoE model.\n\n- The authors compare their performance with the mentioned X-MoE model, whereas there are a variety of MoE approaches that have been proposed and released recently.\n\n- The ablation studies are limited. The performance of SMoE heavily depends on the choice of hyper-parameters, one of the most important being the number of experts to be activated (referred to as top-k). How does the activation map change with an increase in k? Would MH-MoE still perform better, or is there a saturation point? In other words, does MH-MoE seem superior to X-MoE only because a small k=2 was chosen?\n\n2. There is a lack of discussion on limitations. The authors claim that their approach does not have any limitations. I would urge the authors to reflect on the guidelines and reconsider their answer. For example, the weaknesses mentioned above (but not limited to) should be considered.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. We address your feedback point by point below.\n\n---\n\n>**Q1**: The experiments are limited by a single model architecture. The authors apply their approach by modifying the previously released X-MoE model. The authors compare their performance with the mentioned X-MoE model, whereas there are a variety of MoE approaches that have been proposed and released recently.\n\n**A1**: In fact, we conducted experiments on **two MoE structures**:\n\n1. For English-focused language modeling, multilingual language modeling, and masked multimodality modeling tasks, **we implemented our MH-MoE on XMoE [1]** and conducted extensive experiments across these three tasks (ranging **from 300M to 7B** parameters, with **3 upstream tasks and 26 downstream tasks**) to validate the effectiveness of MH-MoE.\n2. For pure vision tasks, as detailed in Section 5 of the main text, **we implemented MH-MoE on AdaMV-MoE [2]**. We validated that MH-MoE can further enhance AdaMV-MoE's performance in vision tasks such as **classification, object detection, and instance segmentation**. \n\nWe chose XMoE and AdaMV-MoE because they have been proven superior to vanilla SMoE in extensive experiments [1, 2]. **We believe that the extensive experiments on these multiple upstream & downstream tasks with the two MoE structures demonstrate the effectiveness of MH-MoE.** If you have any MoE structures of interest, please let us know during the discussion stage, and we would be happy to conduct further experiments :)\n\n**Reference**\n\n[1] Chi, Zewen, et al. \"On the representation collapse of sparse mixture of experts.\" NeurIPS 2022.\n\n[2] Chen, Tianlong, et al. \"Adamv-moe: Adaptive multi-task vision mixture-of-experts.\" ICCV 2023.\n\n---\n>**Q2**: The ablation studies are limited. The performance of SMoE heavily depends on the choice of hyper-parameters, one of the most important being the number of experts to be activated (referred to as top-k). How does the activation map change with an increase in k? Would MH-MoE still perform better, or is there a saturation point? In other words, does MH-MoE seem superior to X-MoE only because a small k=2 was chosen?\n\n**A2**: This is a meaningful suggestion. To explore the issue you mentioned, we conducted additional experiments with k=3 and k=4 to investigate the effects of increasing k on MH-MoE.\n\nWe performed a language model experiment based on the RedPajama dataset, utilizing a 6-layer transformer architecture with each layer having a hidden dimension of 768. For the MoE settings, the model selects the k=3  or k=4 experts for each input. We employ a Multi-Head Mixture of Experts (MH-MoE) configuration with 2 heads. The learning rate is set to 6e-4, token number of each batch is 0.5M. **Due to time and resource constraints, we present results for the model at 25k steps.**\n\nThe experimental results are shown in the Table A below. We find that, **regardless of whether k=2 （see in the maintext）, k=3, or k=4, MH-MoE consistently outperforms the baseline MoE model. Additionally, there is no indication of a saturation point in MH-MoE’s performance as k increases.**\n\nThe activation map's behavior with an increase in k shows effects similar to those observed in Figure 6 of the main text. Specifically, as k increases, the activation map becomes brighter, indicating that experts are activated more frequently. This is because an increase in k means that each expert's probability of being assigned tokens also increases with each token assignment. Consequently, within a fixed number of training steps, the number of times experts are activated naturally increases.\n\n***Table A.*** \n\n| Models        | k=3       | k=4  |\n| ------------- | --------- | ---- |\n| XMoE Baseline | 13.22     |   13.14   |\n| MH-MoE        | **12.68** |   **12.56**   |\n\n---\n\n> **Q3**: There is a lack of discussion on limitations. The authors claim that their approach does not have any limitations. I would urge the authors to reflect on the guidelines and reconsider their answer. For example, the weaknesses mentioned above (but not limited to) should be considered.\n\n**A3**: Thank you for your suggestion. We will include an explicit limitation section in the latest version, summarized as follows:\n\nAn interesting observation is that the performance of our MH-MoE does not consistently improve with an increasing number of heads. We hypothesize that this may be due to each sub-token's dimension becoming too small when the number of heads is too large, leading to information loss within the tokens. This, in turn, makes meaningful and effective routing more challenging, thereby affecting the overall performance of the model. This is an intriguing area for further exploration, and we plan to investigate it in future research.", "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": "gGonVXhhXG", "reviewer": "Reviewer_84mb", "summary": "This paper proposes Multi-Head Mixture-of-Experts (MH-MoE), a simple yet effective routing strategy that splits each input token into multiple sub-tokens for expert routing. This operation significantly enhances the ratio of activated experts for each token, enabling more fine-grained assigning of tokens. Through extensive experimental results across different parameter scales (300M to 7B) and three 15 pre-training tasks, the paper validates the effectiveness of the proposed MH-MoE strategy.", "questions": "Can you provide the balance loss curves for your method and the baselines? I’m curious about the effect of your method on load balancing as it introduces more sub-tokens for routing.", "limitations": "None", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow.\n2. The proposed MH-MoE is simple yet effective. It is easy to be integrated into existing frameworks with MoE optimization. This improves the feasibility of the method.\n3. The experiments cover both CV and NLP modeling scenarios and the results are consistently positive.", "weaknesses": "1. My major concern is about the experimental settings. Implementing MH-MoE based on the X-MoE instead of the vanilla SMoE is strange. The improvements shown by the experiments may not be generalizable to other MoE structures, which makes the proposed method less convincing.\n2. Lack of experiments on scaling computations (activated experts). Though the paper conducts scaling experiments on scaling total parameters (number of experts), the computation remains static. The information conveyed by this kind of scaling is limited as the computation is also important in affecting the performance of MoE. I’m not expecting additional experiments as the time is limited during rebuttal. However, a computational scaling experiment would make the method more convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback. We address your feedback point by point below.\n\n---\n\n>**Q1**: My major concern is about the experimental settings. Implementing MH-MoE based on the X-MoE instead of the vanilla SMoE is strange. The improvements shown by the experiments may not be generalizable to other MoE structures, which makes the proposed method less convincing.\n\n**A1**: In fact, we conducted experiments on **two MoE structures**:\n\n1. For English-focused language modeling, multilingual language modeling, and masked multimodality modeling tasks, **we implemented our MH-MoE on XMoE [1]** and conducted extensive experiments across these three tasks (ranging from 300M to 7B parameters, with 3 upstream tasks and 26 downstream tasks) to validate the effectiveness of MH-MoE.\n2. For pure vision tasks, as detailed in Section 5 of the main text, **we implemented MH-MoE on AdaMV-MoE [2]**. We validated that MH-MoE can further enhance AdaMV-MoE's performance in vision tasks such as classification, object detection, and instance segmentation. \n\n**We chose XMoE and AdaMV-MoE because they have been proven superior to vanilla SMoE in extensive experiments [1, 2]**. **We believe that the extensive experiments on these multiple upstream & downstream tasks with the two MoE structures demonstrate the effectiveness of MH-MoE.** If you have any MoE structures of interest, please let us know during the discussion stage, and we would be happy to conduct further experiments :)\n\n**Reference**\n\n[1] Chi, Zewen, et al. \"On the representation collapse of sparse mixture of experts.\" NeurIPS 2022.\n\n[2] Chen, Tianlong, et al. \"Adamv-moe: Adaptive multi-task vision mixture-of-experts.\" ICCV 2023.\n\n\n\n---\n\n>**Q2**: Lack of experiments on scaling computations (activated experts). Though the paper conducts scaling experiments on scaling total parameters (number of experts), the computation remains static. The information conveyed by this kind of scaling is limited as the computation is also important in affecting the performance of MoE. I’m not expecting additional experiments as the time is limited during rebuttal. However, a computational scaling experiment would make the method more convincing.\n\n**A2**: This is a meaningful suggestion. To explore the issue you mentioned, we conducted additional experiments with top-k=3 and top-k=4 to investigate the effects of increasing activated experts on MH-MoE.\n\nWe performed a language model experiment based on the RedPajama dataset, utilizing a 6-layer transformer architecture with each layer having a hidden dimension of 768. For the MoE settings, the model selects the top-k=3 or top-k=4 experts for each input. We employ a Multi-Head Mixture of Experts (MH-MoE) configuration with 2 heads. The learning rate is set to 6e-4, token number of each batch is 0.5M. **Due to time and resource constraints, we present results for the model at 25k steps.**\n\nThe experimental results are shown in the Table A below. We find that, **regardless of whether k=2 （see in the maintext）, k=3, or k=4, MH-MoE consistently outperforms the baseline MoE model.**\n\n***Table A.*** \n\n| Models        | k=3       | k=4  |\n| ------------- | --------- | ---- |\n| XMoE Baseline | 13.22     |   13.14   |\n| MH-MoE        | **12.68** |   **12.56**   |\n\n---\n>**Q3**:  Can you provide the balance loss curves for your method and the baselines? I’m curious about the effect of your method on load balancing as it introduces more sub-tokens for routing.\n\n**A3**: **In the PDF file located in the Author Rebuttal at the top of the page**, we provide a comparison of the changes in balance loss during training for XMoE Baseline and MH-MoE under two training settings, 32 experts and 8 experts, in the English-focused language modeling experiment. If you have any questions, please feel free to reach out to us during the discussion phase :).", "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": "K5L5RuICbX", "reviewer": "Reviewer_Y3PR", "summary": "The work presents an advanced MoE model architecture to further improve the expert utilization. Authors conduct comprehensive experiments on three pre-training tasks, English-focused language modeling, multi-lingual language modeling and masked multi-modality modeling with models at different scales (300M to 7B).", "questions": "See weakness.", "limitations": "NA", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1) Important research problem. Better expert utilization has great potential to improve the MoE scalability.\n2) Comprehensive experiments. From 300M to 7B. From NLP to CV.\n3) Cool visualization. Figure 1 is very helpful in understanding why and how this model works.", "weaknesses": "I only have one but important concern: the end-to-end training and inference throughput. I understand that the theoretical computation and communication cost is not high, but since we are conducting a more complicated/fine-grained routing decision, the routing process may take longer in practice. Considering that MoE models are sometimes trained and deployed at very large scale, the intra-node all2all is sometimes required. Will the proposed approach slow down this?\n\nI can increase my score to positive if this concern can be solved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback :)\n\n---\n\n>**Q1**: I only have one but important concern: the end-to-end training and inference throughput. I understand that the theoretical computation and communication cost is not high, but since we are conducting a more complicated/fine-grained routing decision, the routing process may take longer in practice. Considering that MoE models are sometimes trained and deployed at very large scale, the intra-node all2all is sometimes required. Will the proposed approach slow down this?\n>\n>I can increase my score to positive if this concern can be solved.\n\n\n**A1:** In fact, if traditional static routing [1] (i.e., using zero-padding or dropping tokens to force a uniform distribution of tokens among experts) were applied in MH-MoE, the situation would be different. Since MH-MoE uses sub-tokens for more complex and fine-grained routing decisions, the matrix dimensions required to record routing information expand accordingly during both training and inference. This results in more time-consuming matrix operations and higher all-to-all communication costs within nodes, thereby affecting the end-to-end throughput of deploying MH-MoE. We present a comparison of throughput between XMoE and MH-MoE using static routing (denoted as MH-MoE*) in Table A. Our results show that MH-MoE with static routing is slower than the baseline.\n\nTo address this issue, we have implemented dynamic routing [2] (refer to the implementation: [databricks/megablocks (github.com)](https://github.com/databricks/megablocks)) in our MH-MoE model. Dynamic routing alleviates the constraints imposed by static routing on both software and hardware, enabling a more efficient and unrestricted routing strategy. Table A shows a comparison of the throughput taken by MH-MoE with dynamic routing (denoted as MH-MoE) and the baseline. We observe that using the more efficient dynamic routing approach results in end-to-end throughput for MH-MoE that is nearly identical to the baseline.\n\nThis is the approach we have taken to address the issue of high end-to-end training and inference throughput in MH-MoE. If you have any further questions, we welcome continued discussion during the discussion phase. :)\n\n***Table A. End-to-End Throughput on eight GeForce RTX 3090 GPUs, each with 24 GB of GPU memory.*** \n\n| Models        | End-to-End Throughput (toks/s) |\n| ------------- | ------------------------------ |\n| SMoE Baseline | 32.21                          |\n| MH-MoE*       | 25.77                          |\n| MH-MoE        | 31.93                          |\n\n**Reference**\n\n[1] Lepikhin, Dmitry, et al. \"Gshard: Scaling giant models with conditional computation and automatic sharding.\" arXiv 2020.\n\n[2] Gale, Trevor, et al. \"Megablocks: Efficient sparse training with mixture-of-experts.\" PMLR 2023.", "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:
## Supplementary rebuttal for Reviewer Rrmq

---

>**Q4**: The core claim that the problem with MoE is: "low experts activation issue" has a few problems: a. It is never clearly explained why this is an issue. b. the paper also shows that increasing the number of heads (and therefore activations) seems to also decrease performance after a specific number of heads [Figure 5, Section 4.4]. (Great for reporting this, however, this also challenges the introduction / proposition of the method). I.e. it seems that "low experts activation" is not always a problem, or only a problem to a certain degree?

**A4**: We address your questions a and b point by point below:

(a). As shown in Figure 7 (a) and (b), we observe that when the number of experts increases significantly (e.g., to 128 or 256), the performance of the SMoE model declines both upstream and downstream. This suggests that with a large number of experts, the SMoE model suffers from an imbalance in expert activation, where most experts have very low activation rates and only a small subset of experts remain active. As a result, the majority of experts do not receive sufficient data for optimization, making it challenging for SMoE to leverage the benefits of a large number of experts. This issue highlights that the "low experts activation issue" is a major factor limiting the further scaling of the SMoE model. **Therefore, we believe that addressing the "low experts activation issue" is crucial, especially as we aim to improve model performance by increasing the number of experts to very large sizes, where this issue may become a significant bottleneck.**

(b). It is indeed possible for the model performance to first increase and then decrease as the number of heads grows. We hypothesize that, although the "low experts activation issue" is mitigated, when the number of heads is excessively large, each sub-token's dimensionality is divided too finely. This can lead to a loss of internal token information and make meaningful and effective routing more challenging, thereby affecting the overall performance of the model. This is a fascinating area for exploration, and we appreciate your observation. We will investigate this further in our future research.



---

>**Q5**: Reporting on sharding of the models.

**A5**:  We apologize for not adequately describing the reporting on sharding of the models in our paper. Specifically, we employed both **data parallelism** and **expert parallelism** strategies. We will include the corresponding details in the revised version of the paper.

Meta-review: This paper presents a mixture-of-experts approach that routes different sub-tokens to different experts. Reviewers found the idea to be simple and effective, the results to be strong compared to existing baselines, the paper clear, and the experiments comprehensive. Concerns regarding missing experiments, unsupported claims and clarity issues have largely been addressed. Some concerns about XXL-scale experiments seem out of scope in my view, as there is plenty of value in models much smaller than 10T parameters. Other requests for additional experiments and baselines seem not essential either.
Decision: Accept (poster)