On the Representation Collapse of Sparse Mixture of Experts
Authors: Zewen Chi
Li Dong
Shaohan Huang
Damai Dai
Shuming Ma
Barun Patra
Saksham Singhal
Payal Bajaj
Xia Song
Xian-Ling Mao
Heyan Huang
Furu Wei
Venue: NeurIPS 2022 Accept (2022)

Abstract: Sparse mixture of experts provides larger model capacity while requiring a constant computational overhead. It employs the routing mechanism to distribute input tokens to the best-matched experts according to their hidden representations. However, learning such a routing mechanism encourages token clustering around expert centroids, implying a trend toward representation collapse. In this work, we propose to estimate the routing scores between tokens and experts on a low-dimensional hypersphere. We conduct extensive experiments on cross-lingual language model pre-training and fine-tuning on downstream tasks. Experimental results across seven multilingual benchmarks show that our method achieves consistent gains. We also present a comprehensive analysis on the representation and routing behaviors of our models. Our method alleviates the representation collapse issue and achieves more consistent routing than the baseline mixture-of-experts methods.

On the Representation Collapse of Sparse Mixture of Experts
Zewen Chit; Li Dong', Shaohan Huang', Damai Daill, Shuming Ma', Barun Patra Saksham Singhal', Payal Bajaj', Xia Song Xian-Ling Mao", Heyan Huang', Furu Wei!
Beijing Institute of Technology
$ Microsoft Corporation
I| Peking University
attps://ake ms mmmrragi
Abstract
Sparse mixture of experts provides larger model capacity while requiring a con- stant computational overhead. It employs the routing mechanism to distribute input tokens to the best-matched experts according to their hidden representations. However, learning such routing mechanism encourages token clustering around expert centroids, implying a trend toward representation collapse. In this work, we propose estimate the routing scores between tokens and experts dimensional hypersphere. We conduct extensive experiments on cross-lingual language model pre-training and fine-tuning on downstream tasks. Experimental results across seven multilingual benchmarks show that our method achieves con- sistent gains. We also present comprehensive analysis on the representation and routing behaviors of our models. Our method alleviates the representation collapse issue and achieves more consistent routing than the baseline mixture-of-experr methods.
low-
1 Introduction
Scaling up model capacities has shown to be a promising way achieve better performance on wide range of problems such as language model pre-training (Radford et al., 2019; Raffel et al., 2020), and visual representation learning (Dosovitskiy et al., 2021; Bao et al., 2022). Despite the effectiveness, increasing the number of parameters leads to larger computational cost, which motivates recent studies to explore Sparse Mixture-of- -xpeeessss (SMoE) models (Shazeer et al., 2017; Fedus et al., 2021; Lepikhin et al., 2021). SMoE increases the model capacity by building several sparsely- activated neural networks. With nearly constant computational overhead, SMoE models achieve better performance than dense models various tasks, including machine translation (Lepikhin et al., 2021), image classification (Riquelme et al., 2021), and speech recognition (Kumatani et al., 202 I).
The routing mechanism plays an important role in SMoE models. Given an input token, the router measures the similarity scores between each token and experts. Then we distribute tokens to the best- matched experts according to the routing scores. Recent studies explored various token assignment algorithms to improve SMoE training. For instance, Lewis et al. (2021) formulate SMoE routing as linear assignment problem that globally maximizes token-expert similarities. Zhou et al. (2022) have experts selecting top tokens rather than assigning tokens to top experts. Roller al. (2021) and Dai et al. (2022) propose to keep routing choices consistent. Many studies in recent years focus on how to design the token-expert assignment algorithm. In this paper, we present that current routing mechanisms tend to push hidden representations clustering around expert centroids, implying trend toward representation collapse, which in turn harms model performance.
'Zewen Chi and Damai Dai contribute during internship at Microsoft Research.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

In order to alleviate the representation collapse issue, we introduce simple yet effective routing algorithm for sparse cextuure-of-xpertt. models. More specifically, rather than directly using the hidden vectors for routing, we project the hidden vectors into a lower-dimensional space. Then, we apply L normalization both token representations and expert embeddings, i.e,, measuring routing scores on low-dimensional hypersphere. Besides, we propose soft expert gate with learnable temperature, which learns to control the activation of experts.
We evaluate the proposed method on cross-lingual language model pre-training and fine-tuning on downstream tasks. Experimental results show that our model consistently outperforms the baseline SMoE models in terms of both language modeling and fine-tuning performance. Moreover, analysis indicates that our method alleviates the representation collapse issue compared with the SMoE baseline. Our method also achieves more consistent routing behaviors during both pre pre-training and fine-tuning, which confirms the effectiveness the proposed routing algorithm.
Our contributions are summarized as follows:
We point out the representation collapse issue in sparse mixture-of-cxpert models, which is under-explored previous work.
We propose estimate routing scores between tokens and experts on a low-dimensional hypersphere in order to alleviate representation collapse.
We conduct extensive experiments on cross- pess-linggul language model pre-training and fine- tuiigg downstream tasks.
We present detailed analysis of routing behaviors and representation properties, which shows that our method improves performance and achieves more consistent routing.
2 Background
2.1 Sparse Mixture of Experts
Sparse Mixture-ote of- of-Experts (SMoE) models take advantage of conditional computation, and have shown to be promising way to scale up the number of parameters. In this work, we consider SMoE for Transformers, where SMoE layers are inserted into neighboring Transformer blocks. Each SMoE layer consists of router and several expert networks. Following most previous work (Fedus et al., 2021), we use feed-forward networks as experts, instead of self-attention modules.
For the input token æ with its hidden representation h E ıR the router computes the routing score between and the i-th expert by dot-product similarity metric s; h e, where e, E RR is learnable expert embedding, and d is the hidden size of the model. Then, the router utilizes a sparse gating function g()) make the expert network conditionally activated.
In this paper, we mainly focus on top-1 routing, i.e., only the expert with the largest routing score is activated. Formally, considering SMoE layer with N experts, the forward function of SMoE can be written as:
arg max s, arg max
SSMoB (h) h (FP
(2)
where SFFN(.) stands for the k-th expert network that is implemented stacked feed-forward networks. Moreover, we explore both softmax gating (Lepikhin et al., 2021; Fedus et al., 2021) and sigmoid gating (Lewis et al., 2021; Dai et al., 2022) for the function g(su):
exp(sx)/ En exp(s;), softmax gating (sk),
g(*x) =
(3)
sigmoid gating
where (()) is the sigmoid function.
2.2 Representation Collapse of Sparse Mixture-of-Experts
We present how representation collapse happens in sparse mixture-of-experts models. For conve- nience, we use h' = JSMoE (h) to denote the output of the SMoE layer as in Equation (2), S₁ g(s))

---

to denote the k-th output of the softmax function, and h = = (h) to denote the output of the k-th expert network. The Jacobian matrix with respect given by:
N
J = (T JE.NE
Σ
Se( S S B3
REN
(4)
where dEj is a Kronecker delta. The equation means that the Jacobian matrix can be decomposed into two terms. The first term J1 represents producing a better token representation given the current activation Sg. The second term J2 means to learn better gating function for appropriate actívation score Sk. After back-propagation, the gradient is received from the above two paths, written as VnL = J1 Vn J2 VnrL L. The second term can be expanded as:
J C Σ Sk(ōks - S,)(H Vn V Σ cjj
(5)
where Cj - Sk(kk - (h h tends to be updated toward a linear combination of the expert embeddings.
The above equation indicates that the token representation
The finding also holds for top-K routing (Lepikhin et al., 2021) where the top K experts (K N) are activated for each token. The forward function of top- K routing is kı, ki,,.... , kx top (SL)
and h = /SMoE(()) = + Ei=1.. ((Sk,) ).E, exp(>x))/ Ej=ı... exp(sk)). Similar to Equation (5), we have
(FFN (h). The gating function is defined as g(8k,) =
K
K
K
J V ΣΣ
Sk. (ōK,
Sk, (1h HNNe, Dle,e,, = Σ CyeNy
(6)
Therefore, the above finding holds for top-K routing.
We consider that such behavior potentially harms the representation capacity of Transformers. Firstly, consider that the N expert vectors span N-dimensional space at most via linear combinations. As N is much smaller than the hidden size d in practice, the spanning subspace does not fully utilize the entire available capacity. Thus, the mechanism renders the Transformer hidden vector h collapsed to an N-dimensional subspace, implying a trend toward representation collapse from R to R where N < d in practice. Secondly, Equation (5) indicates that the hidden vector h tends to be similar to the expert embedding that it is routed to. If the hidden states were routed to the same expert, they are going to be pushed closer. However, we would like to encourage the representations more diverse, so that they can be more expressive and discriminative. The phenomenon possibly restricts the expressibility of hidden states, especially when an expert is inclined to dominate routing.
3 Methods
We introduce the routing algorithm for sparse mixture of experts, which measures the routing scores between tokens and experts on now-dimensional hypersphere. As shown in Figure 1b, we address the representation collapse issue of SMoE by applying dimensionality reduction and Ti normalization for the token representations and expert embeddings. Then, we describe how to incorporate the routing algorithm into an SMoE model under the pre- ree-training-thee thee--fin-tuning paradigm.
3.1 Routing Algorithm
Dimension Reduction In order alleviate the representation collapse issue mentioned in Sec- tion 2.2, we represent the expert embedding e; and the token vector h in a low-dimensional space instead of the original high-dimensional hidden space. Specifically, we first parameterize the experts with lower-dimensional embeddings E dR such that d, is much smaller than the Transformer hidden size Next, we conduct a projection over the hidden states (h), which projects h to the expert embedding space. We use linear projection from Thus, the routing scoring function between the tokens and experts can be written as s; = (Wh) e,, Typically we set d, N/2 (i.e., half of the number of experts) in implementation.
Wh such that W Rd
Inspired by Jing et al. (2022), dimension reduction mitigates the issues described in Section 2.2 from two perspectives. First, linear projection Wh isolates the direct interaction between hidden vector h
3

---

Transfarmer Block
Transformer Block
Router
Networks
Embeediings
Expert Emacdiings
E1
ex
FINZ
by Norm
FEN2
Gor Aricher
N3
Temperature
N3
FRN4
b Norm
FFN4
Transforer Block
Transformer Block
(a) Sparse offture-of-Exertte (SMoE) Layer
(b) X-MOE Layer (Ours)
Figure I: Illustration of a typical SMoE layer and the proposed X-MoE layer. (a) An SMoE layer consists of router and expert networks, where the experts sparsely activated according to dot-product token-expert routing scores. (b) X-MoE improves the routing algorithm via dimension reduction, L2 normalization, and gating temperature.
and expert embedding e;, which tends to relieve cascaded collapse for representations. Second, it is natural apply low-rank projector for hidden vectors, as the number experts is usually much smaller than the hidden size of Transformers. Hence the reduced dimension better fits in with the low-rank nature of routing.
L2 Normalization After dimension reduction, apply L2 normalization to both token represen- tations and expert embeddings. Our routing score is defined as:
(Wh) Whillell
8,
where || || is L2 normalization. Thus, the resulting representations are transformed into a certain scale with stabilized routing scoring.
As described Section 2.2, if an expert dominated a set of hidden states, the representations were pushed toward the expert embedding. In order to fully utilize the space, we favor larger uniformity of representations while avoiding dominated experts. Given a hidden vector h, the dot-product routing 8; (Wh)- is affected by both ||e,|| and cos Wh, e,). So some experts allocated with more tokens because of larger values |e,||- In contrast, L2 normalization projects vectors on the unit hypersphere, which suppresses the undesired effect of |/e,]]. The visualization in Figure 2b also confirms that our method improves the uniformity of learned representations.
Implementation Tips. When scaling the model to more experts, empirically, we observe that the resulting token assignments can be in fluctuation if the expert embedding norm ||e,|| is small. Therefore, we initialize the expert embeddings with L2 norm of 0.1 and keep the norm unchanged during training. Since the expert embeddings are parameterized in the space of IRd,, the change rate of the angle of e; is in inverse proportion to ||e:|l |l As a result, if the norm is small, the angle of e; is updated fast, finally leading to the fluctuation of token assignments, especially when scaling up the model with more experts.
Gating with Learnable Temperature In addition, we add a learnable temperature scalar in the SMoE gating function g( ((sk). Because L2 normalization rescales the routing scores sy the range [],, directly using the scores for SMoE gating tends to make expert activation too conservative. The introduced temperature enables the router adjust the gating (**)) accordingly. To be more specific, our gating function is:
exp(ss/7) EN exp(s)/7) (8k/T),
softmax gating
g($))
(8)
sigmoid gating
where (-)) is the sigmoid function, and the temperature scalar is learnable.

---

3.2 Training Objective
The training objective is jointly minimizing the loss of the target task and auxiliary load balancing loss (Fedus et al., 2021). The load balancing loss is separately computed for each router. For each router, given the frequency t, of how many tokens are routed the i-th expert and the routing score Si, the load balancing loss is computed via:
N
exp(s:/To)
Chalance
B
Σソ= cxp(sy/To)
I
where N is the number of the experts, B is batch of training examples, [B] is the number of tokens, and To stands for a constant temperature. Different from the learnable T in Equation (8), To is kept fixed during training. The overall training objective is to minimize:
C Luskk halance
(10)
where a is a coefficient for load balancing. The term Liss is determined by the specific task that Transformer learns. For example, we employ the masked language modeling loss (Devlin et al., 2019) for pre-training. and the secenceeto-seqenner learning objective for neural machine translation.
3.3 Frozen Routing During Fine-tuning
We evaluate SMoE under the teaining-then--ne-ttuuee paradigm in our work. During fine-tuning, freeze all the parameters of experts, including both the router and expert networks. Because the fine-tuning datasets are usually small compared with pre-training corpora. We find that SMoE models tend to overfit downstream tasks, which often leads to inconsistent routing. Freezing SMoE parameters helps to relieve the above issues. Notice that we still use load balancing loss although the routers are kept fixed, which empirically improves fine-tuning performance in experiments.
4 Experiments
We conduct experiments on cross- loossll language model pre-training (Devlin et al., 2019). We evaluate the performance by fine-tuning the pretrained models various downstream benchmarks. We also compare validation losses of the masked language modeling task. Our method is named as X-MOE in the following sections.
4.1 Experimental Setup
Pre-training Data Following (Chi al., 2021), we use the combination of CCNet (Wenzek et al., 2019) and Wikipedia dump as pre-training corpora. We sample sentences in 94 languages from the corpora, and employ a re-balanced distribution introduced by Conneau and Lample (2019), which increases the probability of low-resource languages.
Model Architecture and Hyperparameters We construct our X-MoE models using the Trans- former (Vaswani et al., 2017) encoder (L. = 12, H = 768, A = 12) with the vocabulary provided by Conneau al. (2020) as the backbone architecture. Following Lewis et (2021), we build 32-expert sparse layer with 3 FFN sub-layers, and insert it after the 6-th Transformer layer. The routing dimension d, is set as 16. The gating temperature To set as 0.3 and 0.07 for the softmax gate and sigmoid gate, respectively. The detailed hyperparameters of X-MOE models can be found in Appendix A. X-MOE models are pretrained with the Adam optimizer (3) = 0.9, 33 0.98) using batch size of 2, 048 for 125K steps. The pre-training procedure takes 2 days on 2 Nvidia DGX-2 Stations. Appendix and Appendix C provide the detailed hyperparameters for X-MoE pre-training and fine-tuning.
Baselines We consider two baselines in our experiments. (1) Dense is a dense Transformer encoder without aparsely-activated modules. (2) SMoE is our implementation of Switch Transformers (Fedus al., 2021). The SMoE baseline built with the same setting with X-MOE. original softmax-gating implementation, we also implement sigmoid-gating (Lewis et al., 2021; Dai al., 2022) variant of Switch Transformers as baseline approach. Notice that the baseline models pretrained with the same training data as X-MOE for a fair comparison.
addition to its
5

---

Table I: Evaluation results the cross-lingual XTREME benchmark. The models fine-tuned on the English training data and directly evaluated in all target languages. SMoE models are grouped according to the choice of gating function. The results averaged over five runs.
Structured Prediction
Question Answering
Classification
Model
Avg
POS
NER
XQuAD
MLQA
TyDiQA
XNLI PAWS-X
Metrics
FI/EM
FI/EM
FI/EM
Acc.
Acc.
Dense (without SMoE)
70.0
61.1
67.3/511
58.7/41.1
42.1 / 28.3
70.1
84.1
61.4
With softmax gating
SMoE Baseline
70.1
60.9
75.3/55.2
62.8/4.3
44.8
50.9/
34.5
71.5
84.6
63.8
X-MoH (Ours)
70.8
63.2
72.4/66.2
64.5 46.3
53.7 38.1
72.0
85.2
65.3
With sigmoid gating
SMOE Baseline
70.6
61.2
/5.7
63.2/45.3
50.1
35.1
71.2
85.1
64.1
X-MoB (Ours) 71.1
62.7
72.3/ 56.3
64.3 46.0
51.5 36.6
72.2
85.2
65.0
Table 2: Results of upstream evaluation. We report the validation perplexities on masked language modeling.
Table 3: Ablation studies of X-MoE components. The models employ various combinations of dimension re- duction, Le normalization, and frozen routing. Average fine-tuning results of five random seeds reported.
Model
Perplexity
Dense (without SMoE)
23.51
Dim. Red.
Lg Norm
Frozen
XNLI
MLQA
With softmax gating
72.2
64.3 46.0
SMoE Baseline
19.02
71.5
63.4 45.2
X-MOE (Ours)
18.72
71.4
63.0 45.2
With sigmoid gating
71.7
63.9 / 45.8
J
71.6
63.6 / 45.5
SMoE Baseline
19.59
X- X-MOE (Ours)
19.12
X
71.2 X
63.2 45..3
4.2 Downstream Evaluation
We conduct downstream evaluation seven widely-used cross-lingual understanding benchmarks from XTREME (Hu et al., 2020). Specifically, we conduct experiments on Universal Dependencies v2.5 part-of-speech tagging (Zeman al., 2019), WikiAnn named entity recognition (Pan et al., 2017; Rahimi et al., 2019), natural language inference (XNLI; Conneau et al. 2018), paraphrase adversaries from word scrambling (PAWS-X; Yang al. 2019), and question answering on MLQA (Lewis et al., 2020), XQuAD (Artetxe et al., 2020), and TyDiQA-GoldP (Clark et al., 2020). Among the benchmarks, we adopt the tross-lingual transfer setting, where the models are fine-tuned with the training data in English and evaluated in all target languages.
Table 1 presents the evaluation results the seven downstream tasks from the XTREME benchmark. For each task, the results are first averaged among the test languages and then averaged over five random seeds. Overall, the softmax-gatine X-MOE model obtains the best performance, achieving an average score of 65.3. Comparing SMoE models with the dense model, SMoE models show notable improvement, indicating that SMoE models benefit from the large model capacity. Comparing X-MOE with the two SMoE baselines, it shows that X-MOE models provide consistent gains on downstream tasks, demonstrating the effectiveness of our proposed routing algorithm. We also validate X- X-MoE under the top-2 2 routing setting. Table 7 presents the evaluation results on XNLI, showing consistent improvements over the baseline for both top-1 and top-2 routing settings.
4.3 Upstream Evaluation
We compare the pretrained models for the upstream performance by the validation perplexity masked language modeling (MLM). We sample multilingual sentences from mC4 (Xue et al., 2020), and construct an MLM validation dataset that contains 65, 536 sequences with lengths around 512.
The results are shown Table 2. Similar the downstream results, we observe that SMoE models perform better than the dense model. In terms of the SMoE models, X-MoE models with both softmax

---

Table 4: BLEU scores multilingual machine translation WMT-10. The models evaluated in the directions of 'x - en'.
Model
ro
CS
et
de
tr
lv
Avg
Dense (without SMoE)
36.9
33.7
29.8
27.8
40.6
25.4
24.6
22.2
20.9
29.1
SMoE Baseline
37.1
33.8
31.0
28.6
42.3
26.0
24.3
23.0
21.2
29.7
X-MoE
37.9
34.3
31.1
29.0
42.2
27.0
24.8
23.5
21.6
30.2
Table 5: Comparison of routing dimensions for dimensionality reduction. 'N' stands for the number of experts.
Table 6: Effects of load balanc- ing during fine-tuning. The mod- els fine-tuned with various weights for the auxiliary load balancing loss.
Table 7:
Evaluation results on XNLI under the top-1 and top-2 routing settings. The models use the softmax gating functions.
Routing Dim.
XNLI
MIQA
Model
top- XNLI
Weight
XNLI
MLQA
N/4
71.4
64.3 46.4
SMOE Baseline
71.5
N/2
72.2
64.3 / 46.0
71.71
64.57
SMOE Baseline
2
73.4
N
71.7
63.8 45.9
10 3
71.65
64.40
2N
71.7
62.7
/ 44.8
10-2
71.93
64.59
X-MOE
1
72.0
4N
71.2
63.1 / /45.0
10-1
71.68
64.50
X X-MOE
2
73.7
and sigmoid gating achieve lower masked language modeling perplexities than their counterparts. Among all the pretrained models, the softmax-gating X-MoE the achieves the lowest validation perplexity. The results show that our method not only works well for learning transferable text representations for downstream tasks, but also brings improvements to the upstream masked language modeling task. Comparing the upstream results with the downstream results, it shows that achieving a lower upstream perplexity does not promise better downstream performance. For instance, the sigmoid-gatinn X- X-MoE model has larger perplexity than the softmax -gaiing SMoE baseline has, but outperforms the fine-tuning performance of the baseline on the downstream tasks.
We also conduct experiments the multilingual machine translation task. As shown in Table present the BI.EU scores on the WMT-10 (Wang et al., 2020) dataset where the models are evaluated the directions of + en', X-MoE consistently outperforms both the dense model and the SMoE baseline in eight translation directions.
4.4 Ablation Studies
Routing Algorithm To better understand routing algorithm, pretrain several variants of sigmoid-gating X-MOE models with various combinations of dimension reduction (Dim. Red.), L2 normalization (L2 Norm), and routing frozen (Frozen). For fair comparison, all the models pretrained and fine-tuned under the same setup, i.e., training data, steps, and the random seeds. We evaluate the models on XNLI and MLQA, and report the results in Table 3. Jointly using the three routing methods achieves the best performance. When ablating one of the three routing methods, the model performs less well, demonstrating that X-MOE benefits from all the three components.
Dimension of Expert Embedding We conduct experiments by adjusting the routing dimension for dimensionality reduction. Specifically, dimensions of N/4, N/2, , N, 2N,, and 4NV, where is the number of the experts. Table shows the downstream performance. It shows that using the routing dimension of N/2 provides the best performance for XNLI and N/4 the best for MLQA. The results also confirm that dimension reduction better fits in with the low-rank nature of SMoE routing.
compare sigmoid-gatin X-MOE models with routing
Load Balaneing During Fine-tuning We explore whether load balancing is beneficial for fine- tuning SMoE models. To this end, we add load balancing loss to the total loss with various weights when fine-tuning X-MOE models on XNLI and MLQA. Table shows the average validation scores where we search the load balancing coefficient ranging from 0 to 10-1 We observe that using balance loss during fine-tuning slightly beneficial for X-MEE, When removing the balance loss, X- K-MoE still remains comparable results on both XNLI and MLQA.

---

Reyessemntio Callogse (RC)
15.0
SM6eE Boriax
12.5
X-Mall
10.0
7.5
0.0
0.5
Pro-mmining Step
le5
(a) SMOE Baseline
(b) X-MOE (Ours)
(c) RC Through Pre-training
Figure 2: Analysis on the representation collapse of the Transformer hidden states. Figure (a) and (b) visualize the spatial structure of the experts. Each data point represents token to be routed, and its color stands for the expert that it assigned to. Figure (c) presents the curves of representation collapse (RC), which measures the within-class variability of hidden states. Larger RC values indicate less collapse.
Analysis
Representation Collapse We qualitatively analyze the representation collapse issue by visualizing experts. Figure 2a and 2b illustrate the spatial structure of the experts of SMoE baseline and X-MoE in hyperbolic space, which is produced by Uniform Manifold Approximation and Projection (UMAP; Melnnes et al. 2018) with n-neighbor of 100 and min-dist of 1. Each data point represents token to be routed, where use the hidden states for SMoE baseline and the projected token representations for X- MOE. Each color stands for an expert that the tokens are assigned to.
Figure 2a shows that most of the points are mixed together with large amount of available room unused, which suggests representation collapse in the expert embedding space. In contrast, X-MOE in Figure 2b shows well-organized feature space with clear distinctions between clusters. It indicates that our routing methods successfully project the tokens to the expert embedding space with routing features preserved.
Additionally, we conduct quantitative analysis on the degree of representation collapse for the learned Transformer hidden states that are fed into SMoE routing. We use the representation collapse metric proposed (Zhu al., 2021). Given the representations be measured, we use Ew and Σg to denote the within-class and between-clase covariance matrices, respectively. The representations collapse (RC) metric is calculated via:
RC = TreE
(11)
where E is the pseudo inverse of Σg. Smaller RC values indicate representation collapse to greater extent. Figure 2c illustrates the metries during pre-training, where the data sampled from the validation set mentioned in Section 4.3. SMoE baseline is unlike unconstrained feature models that can empirically collapse almost zero RC, but still shows consistent descending trend through pre-training, implying a trend toward representation collapse. Differently, X-MoB obtains larger RC scores than SMoE baseline with uptrend through pre-training.
Routing Consistency Through Pre-training We examine whether proposed routing algorithm achieves more consistent routing through training. We measure the routing consistency via the routing fluctuation (RF) ratio metric. Routing fluctuation is defined the change of the target expert of an input token. Correspondingly, the RF ratio measures the ratio of RF between the current and the last checkpoints for the same input. A lower RF ratio indicates better routing consistency. As shown in Figure 3a, we present the RF ratio on the MI. validation set mentioned in Section 4.3. After the 15K step, X- X-MOE shows much lower RF ratio than the SMoE baseline, indicating that model produces more consistent routing behaviors.
Inter-run Consistency Through Fine-tuning In experiments of the downstream evaluation, we find that the routing behaviors of SMoE baseline models can be sensitive to random seeds. As the learned token assignments are various for different training data orders, the final downstream performance can be diverse among runs. Therefore, study the routing behaviors of the SMoE

---

1.0
Routing Pluctuation Ratio
0.6
Inter-run Consistency
0.8
0.5
0.6
SMoE Bascline
0.4
X-MoE
0.4
SMoE Bascline
X-MoE
0.2
0.2
0.25
0.50
0.75
1.00
1.25
Pre-training Step
lc5
Fine-tuming Step
le4
(a) Routing Fluctuation Ratio Through Pre- tre-training
(b) Inter-run Consistency Through Fine-tuning
The routing behaviors of SMoE baseline and X-MOE. (a) Routing fluctuation (RF) ratio measures the ratio the tokens that change their target experts between two checkpoints. Smaller RF values indicate more stable routing. (b) Inter-run consistency measures the correlation among the token assignments of various fine-tuning runs. Larger values indicate more consistent routing.
Figure 3:
baseline and X-MoE models through fine-tuning. To achieve this, we develop a metric, named inter-run consistency, which measures how closely the token assignments converge among the with different seeds. Considering a model with experts, let / = /,.,.... n.y] denote the total load of the experts, where n, stands for the number of the tokens that assigned to the i-th expert. Given two loads I and l2 from two runs with different seeds, the similarity between 1 and l2 is defined as the Pearson correlation coefficient (PCC) between them, which denoted as Pl; J.. Here PCC only serves as a similarity metric rather than measuring linear correlation between variables. By extending it m runs with different seeds for each run, we define the inter-run consistency as the average of correlation matrix IC = _i.fE()...)) Pt.,t/ /m².
We fine-tune X- MoE and SMoE baseline models on XNLI for 12 runs separately. Then we compute the inter-run consistency for every 100 mini-batches, i.e., the expert loads are accumulated for 100 steps. Figure 3b illustrates the inter-run consistency. The SMoE baseline converges toward different routing solutions across multiple runs of fine-tuning, even though the only difference between runs is the random seed. In comparison, X-MOE obtains substantially better inter-run consistency than the SMoE baseline. The curve of X-MOE indicates that the models have various routing behaviors the beginning of the fine-tuning, but finally converge to almost the same routing behaviors.
Related Work
SMoE for Large-Scale Models Sparse Mixture-of-Experts (SMoE) models are introduced by Shazeer et al. (2017), which extends mixture of experts (Jacobs et al., 1991; Jordan and Ja- cobs, 1994) with conditional computation (Bengio et al., 2013; 2015) techniques. Taking advantage of computational computation, SMoE enables massive increase in model capacity while maintaining computational efficiency. To explore the potential of SMoE, recent studies apply SMoE range of machine learning problems such machine translation (Lepikhin et al., 2021), image classification (Riquelme et al., 2021), speech recognition (Kumatani et al., 2021). In addition to the supervised learning scenario, there has been work exploring SMoE under the pre-training- fine-tuning paradigm, and observing discrepancies between strong pre-training quality and poor fine-tuning performance (Fedus et al., 2021; Artetxe al., 2021; Zoph et al., 2022). Besides, the scaling behaviors of SMoE are also studied (Clark et al., 2022; Du et al., 2021).
wide
SMoE Routing Algorithms Many recent studies explore the token assignment algorithms for SMoE routing. BASE layers (Lewis al., 2021) formulate the token routing problem as linear assignment problem. Hash Layers (Roller et al., 2021) employ parameter-free assignment algorithm that routes tokens by hashing. Zhou et al. (2022) let each expert select top-k tokens rather than distribute tokens to experts. Dai et al. (2022) propose to freeze the routing function in order to relieve routing fluctuation. These methods focus on the assignment algorithm in routing, but our routing algorithm focuses improving the underlying routing scoring metric, which is still under-explored.

---

Representation Collapse Representation collapse, also termed neural collapse, is the degeneration of the representations during the training of neural networks. Several studies observe that the within- class variation of the representations in classification networks becomes negligible at the terminal phase of training (Papyan al., 2020; Zhu et al., 2021; Tirer and Bruna, 2022). Besides, this phenomenon has also been observed in language model fine-tuning (Aghajanyan et al., 2021), and visual representation learning (Chen and He, 2021; Ermolov et al., 2021; Jing et al., 2022). These studies focus on centell-activated neural networks. In this work, we point out the representation collapse issue in SMoE models.
Conclusion
In this work, we point out the representation collapse issue in sparse mixture-of-expeet (SMoE) models, and propose a routing algorithm that estimates the routing scores on low-dimensional hypersphere. We conduct extensive experiments on cross-lingual language model pre-training. Experimental results across various benchmarks demonstrate that our method brings consistent improvements over SMoE baselines in terms of both language modeling and fine-tuning performance. Besides, our method alleviates the trend toward representation collapse and achieves more consistent routing. We are going to improve the work from the following perspectives. First, most current X- MoE experiments are conducted on language tasks, such as multilingual language model pre-training, and machine translation. We will also evaluate the proposed method vision pretraining (Bao et al., 2022; Peng et al., 2022) and multimodal pretraining (Wang et al., 2022). Second, we would like to report the results of scaling up model size. The performance gain tends to be greater with larger number of experts.
Ethical Considerations One of the negative societal impacts of training large-scale models is the high computational and environmental cost. Our paper focuses on improving SMoE, which is usually more efficient than dense model training with the same number of parameters. So better SMoE algorithms potentially save required computation and lessen CO2 emissions from computing. Moreover, X-MoE improves multilingual pre-training and fine-tuning, so that cross- lrogs-lingual knowledge from high- to low-resource languages. The bless larger model size brought by SMoE reduces the parameter conflicts of multilinguality, while keeping the computation cost manageable.
better transfer
Acknowledgement
We would like to acknowledge Bo Zheng and Zhiliang Peng for the helpful discussions. Heyan Huang is the corresponding author. Zewen Chi, Xian-Ling Mao, and Heyan Huang are supported by National Key R&D Plan (No. 2020AAA0106600), National Natural Science Foundation of China (No. U219,009, 62172039, and L1924068).
References
Armen Aghajanyan, Akshat Shrivastava, Anchit Gupta, Naman Goyal, Luke Zeutlemoeer and Sonal Gupta. Better fine-tuning by reducing representational collapse. In International Conference on Learning Representationss, 2021. URL /ttpp://openr ewwew..eiww.... wwt.:///ooourr/dddo0//////
Mikel Artetxe, Sebastian Ruder, and Dani Yogatama. On the cross-lingual transferability of monolin- gual representations. In Proceedings of the 58th Annual Meeting of the Association for Computa- tional Linguistics, pages 4623-4637, Online, July 2020. Association for Computational Linguis- tics. doi: 10. acl-main. 421.
URL https:://// httpa://wwwww //www//www///////// wwww.aclwebbbb. aclweb. org/ wwtgg//nthloogy///o /2020.
Mikel Artetxe. Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, et al. Efficient large scale language modeling with mixtures of experts. arXiv preprint arXiv: 12. 10684 2021.
Hangbo Bao, Li Dong. Songhao Piao, and Furu Wei. BETT: BERT pre-training of image transformers. In International Conference on Learning Representations, 2022. URL. ttttpp:///peerevieee het/forummidee -d-n-i-Z2-------
10

---

Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. arXiv preprint arXiv:1511.06299 2015.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:/388.344 340.3433, 2013.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 15750-15758. 2021.
Zewen Chi, Li Dong, Bo Zheng, Shaohan Huang, Xian-Ling Mao, Heyan Huang, and Furu Wei. Improving pretrained cross-lingual language models via self-labeled word alignment. In Pro- ceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume I: Long Pa- pers), pages 3418-3441 141--433, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021 long.lonn..... URL https: /aw.aaaaaanttooyy....ooomooo.oae org/ 2021 acl- -ong..og.... 265.
Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann. Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. arXiv preprint wwXivv2202.01166 2022.
Jonathan H. Clark, Eunsol Choi, Michael Collins, Dan Garrette, Tom Kwiatkowski, Vitaly Nikolaev, and Jennimaria Palomaki. TyDi QA: A benchmark for information--sekking question answering in typologically diverse languages. Transactions of the Association for Computational Linguistics, 8:454-470, 2020. doi: 10.1162/tacl a 00317. URL https: ///////www// /www. aclweb. org/ anthology/ 2020. tac1-1.30.
Alexis Conneau and Guillaume Lample. Cross-lingual language model pretraining. In Advances in Neural Information Processing Systems 32: Anmual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, pages 7057-7067, 2019. URL https://procceedang htwpppoooooonnnnoooooor..eee.pp. neurips. hhtttracctttll....... html.
Alexis Conneau, Ruty Rinott, Guillaume Lample, Adina Williams, Samuel Bowman, Holger Schwenk, and Veselin 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.
2018. Association for Computational Linguistics. doi: 10.18653/w1/DI8. 1269. URL https: / /www//wwwwoooo..o aclweb. org/ aathology/D18- 1269.
Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, Edouard Grave, Myle Ott, Luke Zettlemoyer, and Veselin Stoyanov. Un- supervised cross-lingual representation learning at scale. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, pages 8440-8451, As- sociation for Computational Linguistics. 2020. doi: https::/ddl: org/10 10. 18653 ht///v1/220222 acl mam.mman777.
747. URL
Damai Dai, Li Dong, Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. StableMoE: Stable routing strategy for mixture of experts. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (olumee 1: Long Papers), pages 7085-7095, Dublin, Ireland, May 2022. Association for Computational Linguistics. URL https: : h/tp:/aalllaann hacathhllonh org/2022.coo aon-oog..........
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HILT 2019, Volume 1 (Long and Short Papers), pages 4171- 4186. Association for Computational Linguistics, 2019. doi: 10.18653/v1/n19-11 1423. URL https: org/10 10. 13655/vv//n99
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner. Mostafa Dehghani, Mathhias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL. https: //tp///epe//..eeee///// net/ set//forumiidd/iccaaaaie
11

---

Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-expene experts.. arXiv preprint wrxiv:2112.06900 2021.
Aleksandr Ermolov, Aliaksandr Siarohin, Enver Sangineto, and Nicu Sebe. Whitening for self- supervised representation learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 3015-3024. PMLR, 18 24 Jul 2021. URL htttpp:::ppoceeeeeeeeeeo mlr. wwa...ress/eee vvtps//vv/////// http//1199///momooe/ ermolov21a. htmm..lll.....
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961 2021.
Junjie Hu, Sebastian Ruder, Aditya Siddhant, Graham Neubig, Orhan Firat, and Melvin Johnson. XTREME: A massively multilingual multi-task benchmark for evaluating poos-llnggal generaliza- tion. arXiv preprint 1xXiv:2033.11080 2020.
Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. In International Conference Learning Represeniations, 2022. URL.
Michael I Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2):181-214, 1994.
Kenichi Kumatani, Robert Gmyr, Felipe Cruz Salinas, Linquan Liu, Wei Zuo, Devang Patel, Eric Sun, and Yu Shi. Building a great multi-lingual teacher with sparsely-gated mixture of experts for speech recognition. arXiv preprint 1:Xiv:2//2.05822 2021.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. (GS /hard: Scaling giant models with condi- tional computation and automatic sharding. In International Conference Learning Representa- tions, 2021. URL https: pestps:/oopene /ttp///pppeeeeeeee////ee/eeeeeee net. met/forum?ide qrwprruunddduueeeaeeaee 7 XHTMYb.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning, pages 0265-6274 PMLR, 2021.
Patrick Lewis, Barlas Oguz, Ruty Rinott, Sebastian Riedel, and Holger Schwenk. MLQA: Evaluating
cross-lingual extractive question answering. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 7315-7330. Online, July 2020. Association for Computational Linguistics. URL https:/ /wwww aclweb. org/ anthology/2000 acl-main
653.
Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv: 034822000266 2018
Xiaoman Pan, Boliang Zhang. Jonathan May, Joel Nothman, Kevin Knight, and Heng Ji. Cross-lingual name tagging and linking for 282 languages. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 1946-1958, Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 10. 10118663/1//117-1173 URL https: / aclweb. org/ aatpaantthooyy/geeeaeeeeetteaatt PPwp/ 1178.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences. 117(40): 24652-24663, 2020.
Zhiliang Peng, Li Dong, Hangho Bao, Qixiang Ye, and Furn Wei. BEiT v2: Masked image modeling with vector-quantized visual tokenizers. ArXiv, abs/2208.06346 2022.
12

---

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 2019. URL http: http::/wwwwwwwwwagaeeee persagen. com/fflessmiiscrrddee or1201 9language. pdf.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. URL http://jmlr: wwg/papeeeeeee....o...o.....eeee wwwppss/we///2/ 074. html.
Afshin Rahimi, Yuan Li, and Trevor Cohn. Massively multilingual transfer for NER. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 151-164, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 11865//11//19-0000 URL https: / /www. aclweb. org/ anthology/P19, 1015.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34, 2021.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, al. Hash layers for large sparse models.
Advances in Neural Information Processing Systems, 34, 2021.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gates mixture-of-expert layer. In International Conference on Learning Representations, 2017. URL https / http///ppe//e/eeeeeeeee net/ forum?id BiimmmiddBBikkiii ckkMDqlg..
Tom Tirer and Joan Bruna. Extended unconstrained features model for exploring deep neural collapse. arXiv preprint arXiv:2202.000 08087, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural In- formation Processing Systems 30: Annual Conference on Neural Information Processing Sys- tems 2017, pages 5998-6008, 2017. URL httpps//ppocceeddiii
wwwereiiinnnnnnnnnnneeeeeee cc/pappr/2011/ Seeeh33f5ee22344deeee htt.://deeeeaa00occeeeeeoo c4a44ea4444aaaaaaa Abstract html.
Wenhui Wang, Hangbo Bao, Li Dong, Johan Bjorek, Zhiliang Peng, Qiang Liu, Kriti Aggarwal, Owais Mohammed, Saksham Singhal, Subhojit Som, and Furu Wei. Image as a foreign language: BEiT pretraining for all vision and vision-language tasks. ArXiv, abs/2208. 10442, 2022.
Yiren Wang, ChengXiang Zhai, and Hany Hassan. Multi-task learning for multilingual neural machine translation. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages Linguistics. doi: 00.18653/v1/2020. 10118663111/002..mmllppp. m5.n... URL https: pettps:/accaattoo/// /aclaatooaoalooyyooaa org 2020. mmmlp-main.F.
Online, November 2020. Association for Computational
Guillaume Wenzek, Marie-Anne Lachaux, Alexis Conneau, Vishrav Chaudhary, Francisco Guzman, Armand Joulin, and Edouard Grave. CCNet: Extracting high quality monolingual datasets from web crawl data. arXiv preprint arXiv:1911.00353 2019.
Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami AI-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mT5: massively multilingual pre-trained text-to-text transformer. arXiv preprint arXiv:2010.119 11944, 2020.
Yinfei Yang, Yuan Zhang, Chris Tar, and Jason Baldridge. PAWS-X: A cross-lingual adversarial dataset for paraphrase identification. In Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNIP-IJCNILP), pages 3687-3692, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10. 10.1865//1/DD191133 URL https: //www//ww/.// aclweb. org/
anthology/D19- 1382.
Daniel Zeman, Joakim Nivre, Mitchell Abrams, and et al. Universal dependencies 2.5, 2019. URL http://hdl http::/hdd..hand... hwwdle..eee..cc nwt/11234/1-311000 LINDAT/CLARIAII-CZ digital library at the Institute of Formal and Applied Linguistics (ÚFAL.), Faculty of Mathematics and Physics, Charles
University.
13

---

Yanqi Zhou, Tao Lei, Han-Chu Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M. Dai, Zhifeng Chen, Quoc Le, and James Laudon. -xxpers-of-experts with expert choice routing. arXiv preprint arXiv:2202.093668, 2022.
Zhihui Zhu, Tianyu DING, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geometric analysis neural collapse with unconstrained features. In A. Beygekzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https::/t// //tp://////i//e//////e///e
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. Designing effective sparse expert models. arXiv preprini arXiv:2002008066, 2022.
14


Reviews:
[{"review_id": "wjtdW0Gsv1q", "reviewer": "Reviewer_gm6K", "summary": "This paper makes the following contributions:\n1. Identify that common MoE routing mechanisms that rely on matching token representations to expert representations suffer from representation collapse - learnt token representations become too similar to expert representations collapsing to the subspace of expert representations.\n2. Propose a simple fix to the representation collapse problem - utilizing a low rank projection and cosine similarity metric for routing.\n3. Evaluate their approach and baseline routing algorithms on multilingual text MLM pre-training and downstream evaluation on the Xtreme benchmark.\n4. Qualitative and quantitative analysis demonstrating that the approach presented in the paper improves the representation collapse in MoEs.", "questions": "Please see weaknesses above.", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Identify a previously unknown representation collapse problem in Sparse Mixture-of-Experts models. Propose and evaluate a simple modification to the routing architecture that addresses the problem.\n2. Careful experimentation evaluating the approach on multilingual text MLM pre-training and on downstream tasks from the Xtreme benchmark.\n3. Intuitive qualitative and quantitative analysis of the learnt token representations that highlight the collapse and that the proposed modification addresses the problem.\n4. The paper is well motivated, clearly written and makes it easy to understand the details of the experiments.\n\nWeaknesses:\n1. Limited downstream evaluations: The downstream evaluations on the Xtreme benchmark only evaluate on English-only fine-tuning and zero-shot evaluations. The zero-shot evaluation setting is a little brittle and can benefit from regularization effects and minor fine-tuning differences. It would strengthen the results to: (i) include language-level evaluation results (perhaps in the appendix). (ii) Conduct evaluations in non zero-shot settings perhaps using translated data to simulate multilingual fine-tuning (Translate-Train).\n2. Non-standard upstream evaluation: The upstream evals evaluate MLM perplexity on a custom validation split. This is not a standard upstream evaluation setting and would be hard to replicate if there are minor differences in experimental setups or vocabularies or test split. I would advise the authors to repeat these experiments on a more standard benchmark; perhaps a causal LM with perplexity evaluations on standard test sets.\n3. Limited experimental setting: The paper analyzes and evaluates MoEs trained with Top-1 routing. While theoretically you would expect generalized MoEs trained with Top-K routing to suffer from similar representational collapse, it would strengthen the paper to validate that their approach helps when MoEs are trained with Top-2 routing as well.\n\n\nOverall I recommend weak accept. This paper identifies and addresses a problem with the sparse-MoE architecture which could be valuable to large model training for a variety of different tasks and applications. However, there are several limitations with their evaluation and the scope of experiments presented in the paper.\n\nEDIT: Updated recommendation to accept following author rebuttal. The proposed algorithm seems to result in small but consistent improvements for MoE training on multilingual text pre-training and Machine Translation.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "HH2WkBS2QuG", "reviewer": "Reviewer_5h1o", "summary": "This paper looks at the problem of representational collapse in Mixture Of Experts (MoE) based transformer layers. Specifically, they show that for a softmax based gating function (for combining the MoE layer output with the hidden representation from the previous layer) , the gradient updates to token representations consists of a term that is simply a linear combination of the $N$ expert vectors, and since $N$ is much less than the dimensionality of the token vectors themselves, this portion of the gradient update is “low-dimensional” (specifically, $<= N$). \n\nThe paper argues that this could lead to representational collapse where token vectors receive updates that do not fully utilize the entire dimensionality of the vectors, leading to poorly learnt representations. In response, the paper proposes a simple (and somewhat effective) fix for this: Simply project the representations into a lower dimensional space before feeding it into the MoE layer. This ensures that the token vectors “interact” with the MoE layers (and hence with the corresponding gradients) via the low dimensional mapping. From experiments, they show that this simple fix results in representations with much lesser representational collapse (figure 2c) and in-turn modest gains on down-stream performance (Table-1)", "questions": "* Method: Could the authors clarify the role of other technical ``bells and whistles’’ for preventing representational collapse: using cosine similarity instead of dot-products (L2 normalization), using a temperature parameter in the softmax? I can see these being applied to the baselines as well - was this also done?\n* Experiments: Could the authors comment on why the MoE parameters were kept frozen while finetuning? I only somewhat buy the overfitting argument because pre-trained transformers have been finetuned on extremely small data without much “overfitting” and several tricks can be employed to prevent it (such as early stopping, using low learning rates etc).\n\nPlease also see the weaknesses section from above for a more comprehensive list of questions on experiments. ", "limitations": "Limitations adequately addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n* Paper makes a clever observation regarding representational collapse in mixture of experts models, which to the best of my knowledge, has not been studied or pointed out previously. \n* Conditional computation is becoming increasingly important for scaling to ever larger model sizes and thus the overall theme of the paper is extremely relevant.\n* The experiments on cross-lingual NLP show some modest imrprovements.\n\nWeaknesses:\n* Method: This work only shows representational collapse to be a problem for softmax gating, but this doesn’t seem to be the case for sigmoid? In particular, the contribution from the $J_2$ term only ends up pushing the token embedding towards $e_k$ (the chosen expert). So while one can say that all tokens that are assigned expert $k$ will receive a gradient which has a direction along $e_k$ (hence some amount of clustering), the $J_1$ term will ensure that they also get a gradient pointing in a different direction. (This second comment also applies to the case when you use a softmax gating).\n* Experiments: \n    * Maybe I misunderstood the model, but it seems like only a single MoE layer is added after layer 6, instead of having several MoE layers (perhaps every other layer is an MoE layer). The reason why i ask this is because if there are more MoE layers, that increases the chances of representational collapse, and could probably show bigger gains for your approach. If there is only one MoE layer, then the representational collapse might not be that beneficial. \n    * On the flipside, it would be great to have a toy experiment on simple toy data (even linear regression) where you have only a single MoE layer with a softmax gating, where you can clearly demonstrate representational collapse.\n    * It’s unclear whether the difference b/w Dense (no MoE) and MoE models is because of the MoE layers or just because of more parameters i.e. did you ensure that the Dense models have the same number of parameters? \n    * The UMAP plots don’t seem to suggest representational collapse: It’s really hard to conclude that 2a and 2b have different amounts of spread in the representations. It appears that for both of these, tokens that are mapped to the same expert are very close together. The plot in fig2c is nice but I would like really like to see more analysis / quantitative results on representational collapse. Some suggestions:\n        * you could sample a large enough number of tokens and see if the subspace spanned by them is larger for your approach vs the baseline (presumably, for the baseline you’d expect the subspace to be much lower dimensional). \n        * you could also look at spread i.e. given a cluster of all token vectors that are assigned the same expert, what is thethe average L2 distance of each token vector from the mean?  \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "6fvwyEPtAu-", "reviewer": "Reviewer_JmFj", "summary": "Paper proposes a remedy to the representation collapse issue of sparsely gated (activated) mixture of expert modes (sMoE) in Transformers. The hypothesis, and later supported with observations is that, the top-1 routing mechanism pushes the hidden representations clustering around expert centroids, introducing a tendency toward representation collapse - which might be hurting the model accuracy. \n\nTo mitigate the representation collapse, deemed to be harmful, paper proposes adding two more operations to the routing procedure: (1) dimensionality reduction of the hidden representations that the routing decisions used to base on, and (2) L2 normalization of the activations (logits) of pre-softmax/sigmoid at the top-1 expert selection routine. \n\nTwo enhancements combined, paper demonstrates improved upstream (pretraining perplexity) and downstream (accuracy after fine-tuning)  quality gains over the dense and regular top-1 sMoE models on various benchmarks. ", "questions": "Q1 - What does \"conservative\" mean in line 128?\n\nQ2 - What are the initial and final values of \\tau (in line 131) that are learned by the model - does the model learn higher or lower values than the one initially set?\n\nQ3 - In line 142, you mention \"we employ ... seq2seq for nmt\" but there is no nmt experiments in the paper, is this intended?\n\nQ4 - What is the alpha / temperature for the language mixing (section 4.1 - Pre-training Data)?\n\nQ5 - You mention in 4.3 (last sentence) that the sigmoid gating has different behavior in upstream evaluation, but do not provide any data points / tables / figures - have I missed those?\n\nQ6 - Representation collapse (and others in Figure 2 & 3) are using which set to collect data? Dev / test / train ? Please specify. \n\nQ7 - Why do you think routing consistency should compare the first and one of the last checkpoints to measure the stats? What is the hypothesis there? The model may change the chosen of expert throughout the course of training, which should be okay (please correct me if I'm wrong), rather than committing to an expert at the beginning of the training and sticking with it. Wouldn't it make more sense to take a look at the last n-checkpoints and computing the statistics over that? Please word your hypothesis why comparing first with the last checkpoint supports the argument for routing consistency. \n\nQ8 - All models seem to have trained for the same number of tokens, but some of them use extra parameters and extra priors (dim-reductions + L2 norm). Given these, wouldn't it make more sense to compare two models with the same upstream (pretraining) perplexity on their downstream quality? How do we make sure the #tokens that we set in the pretraining not biasing the proposed method over the others ?\n\nQ9 - Follow up on Q8, have we adjusted the learning rate schedules for all the different approaches (a la chinchilla) ?\n", "limitations": "adequate ", "rating": 3, "confidence": 5, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- Identification of \"representation collapse\" of MoE models as one of the potential issue of these family of models sheds light on the failure modes of them. \n\n- New metrics, like routing fluctuation ratio, and UMAP of token-expert assignment, are novel and great contributions to the community moving forward.\n\nWeakness:\n\n- Proposed method is focusing only on top-1 expert choice sMoE models which are not widely used, and might be attempting to solve a problem that only exist in that space. Which is okay, but the paper is not scoped, structured and narrated that way. Do we know, the widely used top-2 expert choice, where the first expert is chosen according to eq. (3) and the second is sampled randomly from the same distribution, has the same problem? \n\n- The assumption that the number of experts are going to be much smaller than the hidden dimension does not always hold (e.g., Lepikhin et al. 2020 use 2048 experts w/ 1024 hidden dimensions). Similar to the first weakness, scoping of the problem is one of the weaknesses.  \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept