Towards Understanding the Mixture-of-Experts Layer in Deep Learning
Authors: Zixiang Chen
Yihe Deng
Yue Wu
Quanquan Gu
Yuanzhi Li
Venue: NeurIPS 2022 Accept (2022)

Abstract: The Mixture-of-Experts (MoE) layer, a sparsely-activated model controlled by a router, has achieved great success in deep learning. However, the understanding of such architecture remains elusive. In this paper, we formally study how the MoE layer improves the performance of neural network learning and why the mixture model will not collapse into a single model. Our empirical results suggest that the cluster structure of the underlying problem and the non-linearity of the expert are pivotal to the success of MoE. This motivates us to consider a challenging classification problem with intrinsic cluster structures. Theoretically, we proved that this problem is hard to solve by a single expert such as a two-layer convolutional neural network (CNN).  Yet with the MoE layer with each expert being a two-layer CNN, the problem can be solved successfully. In particular, our theory shows that the router can learn the cluster-center features, which helps divide the input complex problem into simpler classification sub-problems that individual experts can conquer. To our knowledge, this is the first theoretical result toward formally understanding the mechanism of the MoE layer for deep learning.

Towards Understanding the Mixture-of-Experts Layer in Deep Learning
Zixiang Chen
Yihe Deng
Department of Computer Science University of California, Los Angeles Los Angeles, CA 90095, USA chenzx19@cs ucla.edu
Department of Computer Science University of California, Los Angeles Los Angeles, CA 90095, USA y yihedeng@cs. ucl.. edu
Yue Wu
Quanquan Gu
Department of Computer Science University of California, Los Angeles Los Angeles, CA 90095, USA ywuêcs www.cclccacco edu
Department of Computer Science University of California, Los Angeles Los Angeles, CA 90095, USA qguêcs. ucla. edu
Yuanzhi Li
Machine Learning Department Carnegie Mellon University Pittsburgh, PA 15213, USA yuanzhil@andrez cmu, edu
Abstract
The Mixture-of ofaere-of-xpppeett (MoE) layer, a sparsely- model controlled by a router, has achieved great success in deep learning. However, the understanding of such architecture remains elusive. In this paper, we formally study how the MoE layer improves the performance of neural network learning and why the mixture model will not collapse into a single model. Our empirical results suggest that the cluster structure of the underlying problem and the non-linearity of the expert are pivotal to the success of MoE. This motivates us to consider a challenging classification problem with intrinsic cluster structures. Theoretically, we proved that this problem is hard to solve by single expert such as a two-layer convolutional neural network (CNN). Yet with the MoE layer with each expert being two-layer CNN, the problem can be solved successfully. In particular, our theory shows that the router can learn the cluster-center features, which helps divide the input complex problem into simpler classification sub-problems that individual experts can conquer. To our knowledge, this is the first theoretical result toward formally understanding the mechanism of the MoE layer for deep learning.
Introduction
The Mixture-of of-Expert (MoE) structure (Jacobs et al., 1991 Jordan and Jacobs 1994 is a classic design that substantially scales up the model capacity and only introduces small computation overhead. In recent years, the MoE layer (Eigen et al., 2013 Shazeer et al., 2017). which an extension the MoE model to deep neural networks, has achieved remarkable success in deep learning. Generally speaking, an MoE layer contains many experts that share the same network architecture and are trained by the same algorithm, with gating (or routing) function that routes individual inputs few experts among all the candidates. Through the sparse gating function, the router in the MoE layer
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

can route each imput to the (o 2 2) best experts (Shazeer et al., 2017). or the single (K = 1) best expert (Fedus et al. 202T). This routing scheme only costs the computation K experts for new input, which enjoys fast inference time.
Despite the great empirical success of the MoE layer, the theoretical understanding of such architecture is still elusive. In practice, all experts have the same structure, initialized from the same weight distribution (Fedus et al.) 202I) and are trained with the same optimization configuration. The router is also initialized to dispatch the data uniformly. It is unclear why the experts can diverge to different functions that are specialized make predictions for different inputs, and why the router can automatically learn to dispatch data, especially when they are all trained using simple local search algorithms such as gradient descent. Therefore, we aim to answer the following questions:
Why do the experts in MoE diversify instead of collapsing into single model? And how can the router learn to dispatch the data to the right expert?
In this paper, in order to answer the above question, we consider the natural "mixture of classification" data distribution with cluster structure and theoretically study the behavior and benefit of the MoF layer. We focus on the simplest setting of the mixture of linear classification, where the data distribution has multiple clusters, and each cluster uses separate (linear) feature vectors to represent the labels. In detail, we consider the data generated as a combination of feature patches, cluster patches, and noise patches (See Definition for more details). We study training an MoE layer based on the data generated from the "mixture of classification" distribution using gradient descent, where each expert is chosen to be two-layer CNN. The main contributions of this paper are summarized as follows:
We first prove negative result (Theorem 4.1 that any single expert, such as two-layer CNNs with arbitrary activation function, cannot achieve a test accuracy of more than 87.5% on our data distribution.
Empirically, we found that the mixture of linear experts performs better than the single expert but is still significantly worse than the mixture of non-linear experts. Figure I provides provides such result a special case of our data distribution with four clusters. Although a mixture of linear models represent the labeling function of this data distribution with 100% accuracy, it fails to learn so after training. We can see that the underlying cluster structure cannot be recovered by the mixture of linear experts, and neither the router nor the experts are diversified enough after training. In contrast, the mixture of non-linear experts can correctly recover the cluster structure and diversify. the toy data, we study a sparsely- -agted
Motivated by the negative result and the experiment MoB model with two-layer CNNs trained by gradient descent. We prove that this MoE model can achieve nearly 100% test accuracy efficiently (Theorem 4.22..
Along with the result on the test accuracy, we formally prove that each expert of the sparsely- gated MoE model will be specialized which is determined by the initialization of the weights. In the meantime, the router can learn the cluster-center features and route the input data to the right experts.
specific portion of the data (i.e., least one cluster),
Finally, we also conduct extensive experiments on both synthetic and real datasets to corroborate our theory.
Notation. We use lower case letters, lower case bold face letters, and upper case bold face letters to denote scalars, vectors, and matrices respectively. We denote union of disjoint sets (A₃ : i T)
by U,, For a vector x, we use ||x||2 to denote its Euclidean norm. For a matrix W, we use ||W]] to denote its Frobenius norm. Given two sequences } and [yn we denote xn Ovw) if |rn|| < C1 |n] for some absolute positive constant C1, In M(y.)) if [an] 2 C2ly for some absolute positive constant Cz, and Vr, = (yr,) if Cslynl < [I,,] S Calyaa for some absolute constants C3, C4 > O. We also use ((.) to hide logarithmic factors of in O-). Additionally, we denote In (Jly if x, Ow,)) for some positive constant D, and I, polylog(y.)) if In = poly (log(y,) ().)). We also denote by In, = o(zn.) if lim In/Un yn = 0. Finally we use [N] to denote the index set N).
2 Related Work
Mixture of Experts Model. The mixture of experts model (Jacobs et al. 1991 Jordan and Jacobs 1994) has long been studied in the machine learning community. These MoE models are based

---

Mdxtur of norilnear experts
Initialization
Training
Finished
Miturre of Unear experts
Figure 1: Visualization of the training of MoE with nonlinear expert and linear expert. Different colors denote router's dispatch different experts. The lines denote the decision boundary of the MoE model. The data points are visualized on space via t-SNE (Van der Maaten and Hinton 2008). The MoE architecture follows section where nonlinear experts use activation function a(z) = z" For this visualization, we let the expert number M = 4 and cluster number K = 4. Wc generate n = 1, 600 data points from the distribution illustrated in Section 3 with (0.5, 2), 8 (1,2), E (1,2), and op 1. More details of the visualization are discussed Appendix
on various base expert models such support vector machine (Collobert et al. 2002 Gaussian processes (Tresp) 2001), or hidden Markov models (Jordan et al.) 1997). In order to increase the model capacity deal with the complex vision and speech data, Eigen et (2013) extended the MoE structure to the deep neural networks, and proposed a deep MoE model composed of multiple layers of routers and experts. [hazzeer al. (207)) simplified the MoE layer by making the output of the gating function sparse for each example, which greatly improves the training stability and reduces the computational cost. Since then, the MoE layer with different base neural network structures (Shazeer et al., (017: Dauphin et al.) 2017] Vaswani et al. (2017) has been proposed and achieved tremendous successes a variety language tasks. Very recently, Fedus et al. (202T) improved the performance of the MoE layer by routing one example to only single expert instead of K experts, which further reduces the routing computation while preserving the model quality.
In this paper, we consider a "mixture of clas- sification" model. This type of models can be dated back to (De Veaux 1989; Jordan and Ja- cobs 1994 Faria and Soromenho 2010) and has been applied to many tasks including object recognition Quattoni et 2004) human action recognition (Wang and Mori 2009), and machine translation (Liang et al.) 2006). In order to learn the unknown parameters for mixture of linear regressionn/Classifcttiio model. (Anandkumar et al. 2012 Hsu al., 2012 Chaganty and Liang 2013 Anandkumar et al. 2014: Li and Liang 2018) studies the method of moments and tensor factorization. Another line work studies specific algorithms such as Expectation-Maxmizzaiin
Mixture of Linear Ceggression/CCassiicatttio
(EM) algorithm (Khalili and Chen 2007 Yi al. 2014 Balakrishnan et al. 2017 Wang et al. 2015).
Theoretical Understanding of Deep Learning. In recent years, great efforts have been made to establish the theoretical foundation of deep learning. A series of studies have proved the convergence (Jacot et al. 2018 Li and Liang 2018 Du et al. 2019) Allen-Zhu et al. 2019b; Zou et al. (2188) and generalization (Allen-Zhu et al. 2019a| Arora et al. 2019a|b b] and Gu (019) guarantees in the so-called "neural tangent kernel" (NTK) regime, where the parameters stay close the initialization,
and the neural network function is approximately linear in its parameters. A recent line of works
(Allen-Zhu and Li 2019 Bai and Lee 2019 Allen-Zhu and Li 2020a blc, Li et al.) 2020 Cao et al.
2022 et al. 2021 Wen and Li 2021) studied the learning dynamic of neural networks beyond NTK regime. is worthwhile mention that our analysis of the MoE model is also beyond the
NTK regime.
3

---

3 Problem Setting and Preliminaries
We consider MoE. layer with each expert being a two-layer CNN trained by gradient descent (GD) over n independent training examples [x,, [i= generated from a data distribution 'D. In this section, we will first introduce our data model D, and then explain neural network model and the details of the training algorithm.
3.1 Data distribution
We consider binary classification problem over P-patch inputs, where each patch has d dimensions.
In particular, each labeled data is represented by (x, y), where input x = (x(*), x(() xx)))).... , x(") e (Rd)) (ed))in a collection of P patches and E {=1} is the data label. We consider data generated from K clusters. Each cluster & E [K] has label signal vector V3 and cluster-center signal vector ea with | = | = 1. For simplicity, we assume that all the signals {vx { U (C& [ are orthogonal with each other.
Definition 3.1. A data pair (x,8) E (RR)P= x {+1} is generated from the distribution D as follows.
Uniformly draw pair (k, k") with l + k'' from (1....,k) K}.
Generate the label y E {+1] uniformly, generate a Rademacher random variable e E {+1}. Independently generate random variables o., B, y from distribution Da. D3. Dy. In this paper, we assume there exists absolute constants Cı, C2 such that almost surely 0 < C1 S a, 8,,877 S Cz.
Generate x as collection of /' patches: x (x(*), x(2) ,...,,",,,,,,,,,,,,,..,,,,,,,.,. ((?)) E (IRd),, where
Feature signal. One and only one patch is given by yovk.
Cluster-center signal. One and only one patch is given by Bcg.
Feature noise. One and only one patch is given by eyva'.
Random noise. The rest of the p 3 patches are Gaussian noises that are independently drawn from N (O, /d) Ia) where op is an absolute constant.
How to learn this type of data? Since the positions of signals and noises not specified in Definition it is natural to use the CNNs structure that applies the same function to each patch. We point out that the strength of the feature noises can be large as the strength of the feature signals o. As we will see later in Theorem [.1 this classification problem is hard to learn with a single expert, such any two-layer CNNs (any activation function with any number of neurons). However, such a classification problem has an intrinsic clustering structure that may be utilized to achieve better performance. Examples can be divided into K clusters UKE]K] based on the cluster-center signals: an example (x, y) E S2 if and only if at least one patch of x aligns with Cy. It is not difficult to show that the binary classification sub-probe problem over Q4 can be easily solved by an individual expert. We expect the MoE can learn this data cluster structure from the cluster-center signals.
Significance of our result. Although this data can be learned by existing works on a mixture of linear classifiers with sophisticated algorithms (Anandkumar et al. 2012, Hsu al., 2012 Chaganty and Liang 2013), ) the focus of our paper is training a mixture of nonlinear neural networks, a more practical model used in real applications. When an MoE is trained by variants of gradient descent, we show that the experts automatically learn to specialize on each cluster, while the router automatically learns dispatch the data to the experts according to their specialty. Although from a representation point of view, it is not hard to see that the concept class can be represented by MoEs, our result is very significant we prove that gradient descent from random initialization can find a good MoE with non-linear experts efficiently. To make our results even more compelling, we empirically show that MoE with linear experts, despite also being able to represent the concept class, cannol be trained to find a good classifier efficiently.
3.2 Structure of the MoE layer
An MoE layer consists of a set of M "expert networks" f1..... fM. and gating network which is generally set be linear (Shazeer et al. 2017: Fedus al. 2021). Denote by fo (x; W) the output of the m-th expert network with input x and parameter W. Define an M-dimensional vector h((; Q)
EpErri p θ (P) as the output of the gating network parameterized by e (0,... 0m) E Rd
The output F of the MoE layer can be written follows:
Fxx: θ, W) = =EmET Tm Tm(x: 0) (x; W),

---

where Tx S [M] is set of selected indices and Tm (x; Os's are route gate values given by exp(Imm +))
7ym X -
Vm E [M].
exp Am (x; )))
Expert Model. In practice, one often uses nonlinear neural networks experts in the MoE layer. In fact, we found that the non-linearity of the expert is essential for the success of the MoE layer (see Section 6.. For m-th expert, consider convolution neural network follows:
fm (x; W) = E Ep=r (( (wm., x())))
(3.1)
where Wm.j I& the weight vector of the j-th filter (i.e., neuron) in the m-th expert, J is the number of filters (i.e., neurons). We denote W,,, m [Wm.1.... Wm.J] Rdx as the weight matrix of the m-th expert and further let W = (W m JmE[M] as the collection of expert weight matrices. For nonlinear CNN, we consider the cubic activation function ((2) = 2², which is one of the simplest nonlinear activation functions (Vecci al. 1998). We also include the experiment for other activation functions such as RELU in Appendix Table
Top-1 Routing Model. A simple choice of the selection set Tx is the whole experts set Tx = [M] (Jordan and Jacobs 1994). which is the case for the so-called soft-routing model. However, it will be time consuming to use soft-routing deep learning. In this paper, we consider "switch routing". which is introduced by Fedus et al. (2021) to make the gating network sparse and save the computation time. For each input x, instead of using all the experts, we only pick one expert from [M], i.e..
|Tx = 1. In particular, we choose Tx argmax,,, (h. (x; e)}.
Algorithm 1 Gradient descent with ran- dom initialization
Require: Number of iterations T. expert learning rate n, router learning rate Yj initialization scale da, training set S =
Expert
Exper:
Expert
Expert
{(x:, )==)
Router
1: Generate each entry of ((()) indepen-
Boftmax
Gatting
dently from N [0, o
Networs
2: Initialize each entry of θ (0) as zero.
3: for t 0, ,, 1 do
Input
Generate each entry of (()) indepen- dently from Unif[0,1].
Figure Illustration of MoE layer. For each input x, the router will only select one expert to perform computations. The choice is based on the output of the gating network (dotted line). The expert layer returns the output of the selected expert (gray hox) multiplied by the route gate value (softmax of the gating function output).
Update (t-1) as in 3.4
6: Update O (c+1) as in 3.5
end for
return (Q)T) W())
3.3 Training Algorithm
Given the training data S = {xx, Mi) we train F with gradient descent to minimize the following empirical loss function:
4 W) 1 En F(x;; W)),
(3.2)
where é is the logistic loss defined {(z) log(1 exp(-2)). We initialize ()) to be zero and initialize each entry of (0)) by i.i.d A' ).) Zero initialization of the gating network is widely used MoE. training. As discussed in Shazeer et (2017). it can help avoid out-of-memory errors and initialize the network in state of approximately equal expert load (see (5.I) for the definition of expert load).
Instead of directly using the gradient of empirical loss 3.2 to update weights, we add perturba- tion to the router and use the gradient of the perturbed empirical loss to update the weights. In particular, the training example x; will be distributed to argmax,,, (hm (x:; (()) + ( instead, where (1 (E( random noises. Adding noise term is widely used training strategy
5

---

for sparsely-gater MoE layer (Shazeer et al., 2017 Fedus et al. 2021 which can encourage explo- ration across the experts and stabilize the MoE training. In this paper, we draw mm. independently from the uniform distribution Unif[0, 1] and denotes its collection as r()) Therefore, the perturbed empirical loss at iteration can be written as
C ( W Er (yyTT, (
fm..));; W())),
(3.3)
where mi, = argmax,, ( x (()) + m. Starting from the initialization W("), the gradient descent update rule for the experts is
W
m 7 Vw., L ( / (O(), W()) Vm E [M), (3.4)
where 1j > 0 is the expert learning rate. Starting from the initialization (0), the gradient update rule for the gating network is
0
= - mJ Vo, ( (t), W(O), Am [M]
(3.5)
where > O is the router learning rate. In practice, the experts are trained by Adam to make sure they have similar learning speeds. Here we use a normalized gradient which be viewed simpler alternative to Adam (Jelassi et al. 202T
Main Results
In this section, we will present our main results. We first provide negative result for learning with single expert.
Theorem 4.1 (Single expert performs poorly). Suppose Da = Dr in Definition 3.1 then any function with the form F(x) Exn (x(!)) will get large test error P ((x.))>D((// (yF(x) S 0) 2 1/8.
Theorem 4.1 indicates that if the feature noise has the same strength as the feature signal i.e., ((w]xx()) x(0) + 6,) can't perform well on the classification problem defined in Definition 3.1 where a can be any activation function. Theorem 4.1 also shows that simple ensemble of the experts may not improve the performance because the ensemble of the two-layer CNNs is still in the form of the function defined in Theorem 4.1
Da = Dy,, any two-layer CNNs with the form F(x) = Σ aj EpE[P] (w
As comparison, the following theorem gives the learning guarantees for training an MoE layer that follows the structure defined in Section 3.2 with cubic activation function.
Theorem 4.2 (Nonlinear MoE performs well). Suppose the training data size 2(d). Choose experts number M = O(K log K log log d), filter size J = O(log M log log d), initialization scale do [d-1/3 d=0.01] learning rate 1) O(ao), nr @(M²)n. Then with probability at least 1 -1), Algorithm is able to output (T) W(T) within T Ö'n iterations such that the non-linear MoE defined in Section 3.2 satisfies that
Training error is zero, i.e., y, F(x;; (T), W(T) > 0, Vi [m]
Test error nearly zero, i.e., (x)))) (x) (T)) w < 0) o()).
More importantly, the experts can be divided into a disjoint union of K non-empty sets [M] U and
(Fach expert is good on one cluster) Each expert m E MA performs good on the cluster Ωg.
PD w(T) < (x y) E Q4) = o(1.. (Router only distributes example to good expert) With probability at least - o(1), an example x E Ω4 will be routed to one of the experts in Mr.
Theorem 4.2 shows that non-linear MoE performs well on the classification problem in Definition 3.1 In addition, the router will leam the cluster structure and divide the problem into K simpler sub- problems, each of which is associated with cluster. In particular, each cluster will be classified accurately by subset experts. On the other hand, each expert will perform well on at least one cluster.
suggests that there exist problem instances in Definition 3.1 (i.c., D. D.)) such that an MoE provably outperforms single expert.
Furthermore, together with Theorem 4.1 Theorem 4.2

---

Overview of Key Techniques
successful MoE. layer needs ensure that the router learn the cluster-center features and divide the complex problem in Definition [3.1 into simpler linear classification sub-problems that individual experts can conquer. Finding such gating network is difficult because this problem is highly non-convex. In the following, we will introduce the main difficulties in analyzing the MoE layer and the corresponding key techniques overcome those barriers.
Main Difficulty 1: Discontinuities Routing. Compared with the traditional soft-routing model, the sparse routing model saves computation and greatly reduces the inference time. However, this form of sparsity also causes discontinuities in routing (Shazeer et al.) 2017). In fact, even a small perturbation of the gating network outputs h(x; θ) + ô may change the router behavior drastically if the second largest gating network output close to the largest gating network output.
Key Technique 1: Stability by Smoothing. We point out that the noise term added to the gating network output ensures a smooth transition between different routing behavior, which makes the router more stable. This is proved in the following lemma.
Lemma 5.1. Let h, h E RM to be the output of the gating network and {Tm Ji to be the noise independently drawn from 1n. Denote p. E R.M to be the probability that experts get routed, i.e, Pm = P(argmax,, m² [MM]
( h + T'm m). Pm = P(argmax E|MM (Im + Tme = m). Then
we have that ||p - 00 S 2 I bll
Lemma 5. implies that when the change of the gating network outputs at iteration I and L' is small, i.e., ||b(x; ( h(x; e (f)))] oo. the router behavior will be similar. So adding noise provides smooth transition from time to f'. It is also worth noting that O is zero initialized. So h(x; ((")) = 0 and thus each expert gets routed with the same probability Pm = 1/M by symmetric property. Therefore, the early of the training when []h(x; (()))- h(x; O (0)) )loc is small, router will almost uniformly
pick one expert from [M], which helps exploration across experts.
Main Difficulty 2: No "Real" Expert. At the beginning of the training, the gating network is zero, and the experts are randomly initialized. Thus it is hard for the router to learn the right features because all the experts look the same: they share the same network architecture and trained by the same algorithm. The only difference is the initialization. Moreover, if the router makes a mistake at the beginning of the training, the experts may amplify the mistake because the experts will be trained based on mistakenly dispatched data.
Key Technique 2: Experts from Exploration. Motivated by the key technique 1, we introduce an exploration stage to the analysis of MoE layer during which the router almost uniformly picks one ex- pert from [M]. This stage starts at & = and ends at Ti 1 og 0.5
< T = O17 and the gating ]| x = (.)..)). Because the experts are
network remains nearly unchanged || hh(x: (()) - hxx; θ
treated almost equally during exploration stage, can show that the experts become specialized to some specific task only based on the initialization. In particular, the experts set [M] can be divided into
K nonempty disjoint sets [M] LleMa, where M == (ml argmaxx'< JE[J] (V ; w ())) k)- For nonlinear MoE with cubic activation function, the following lemma further shows that experts in different set Mr will diverge at the end of the exploration stage.
Lemma 5.2. Under the same condition as in Theorem 4.2 with probability at least 1 - o(1), the following equations hold for expert m E MMs
P P(x.))>D (yfm (x; W(T) < (x(xv) v) E 24) = o()), ((xx))= (vfm. (T)) - 0 (x, v) E Qx) ((1/K) Vk' + k.
Lemma 5.2 implics that, at the end of the exploration stage, the expert m E Ma can achieve nearly zero test error on the cluster Or but high test error the other clusters Qu. K k.
Main Difficulty 3: Expert Load Imbalance. Given the training data set = {(x,, 9.) the load of expert m at iterate t is defined as
Load ESiennp((mm = m).
(5.1)
where Pmm,, m) is probability that the input x; being routed to expert m at iteration Figen et al. 2013) first described the load imbalance issues in the training of the MoE layer. The gating
7

---

network may converge to state where it always produces large Load (+) m. for the same few experts.
This imbalance in expert load is reiffreinf,,,,ng, as the favored experts are trained more rapidly and thus are selected even more frequently by the router (Shazeer et al. 2017 Fedus et al. 2021] Expert load imbalance issue not only causes memory and performance problems in practice, but also impedes the theoretical analysis of the expert training.
Key Technique 3: Normalized Gradient Descent. Lemma 5.2 shows that the experts will diverge into M Normalized gradient descent can help different experts in the same Mr being trained the same speed regardless of the imbalance load caused by the router. Because the self- reinforcing circle no longer exists, the load imbalance issue will get mitigated. In particular, the router will treat different experts in the same M₁ almost equally and dispatch almost the same amount of data to them during the early stage of training (See Section E.2 in Appendix for detail), which is enough for the router to learn the cluster-center features. However, we can't guarantee load balance for an arbitrary long training period if we only use normalized gradient descent. That's is the reason Theorem 4.2 requires early stopping. This load imbalance issue can be further avoided by adding load balancing loss (Eigen et al. 20133 Shazeer et al.) 2017 Fedus et al. (0211) or using advanced MoE layer structure such as BASE Layers (Lewis et al. 202 Dua et al. 2021 and Hash Layers
(Roller et al. 2021).
Road Map: Here we provide the road map of the proof of Theorem 4.2 and the full proof presented in Appendix [ The training process can be decomposed into several stages. The first stage is called Exploration stage. During this stage, the experts will diverge into K professional groups Lik I= = [M. In particular, we will show that Mg is not empty for all k [K]. Besides, for all m E M&, J' is good classifier over S2g- The second stage is called router learning stage. During this stage, the router will learn to dispatch
Qg to one of the experts in Mr. Finally, we will give the generalization analysis for the MoEs from the previous two stages.
Experiments
In this section, we conduct experiments to validate our theory. The code and data for our experiments can be found on Github
Setting 1:a E (0.5, 2), 3 E (1,2), 7 (0.5,3). op =1
Test accuracy (%)
Dispatch Entropy
1.4
Single (linear)
68.71
NA NA 1.300 + 0.044 0.098 + 0.087
1.2
Single (nonlinear)
79.48
1.0
MoE (linear)
92.99 + 2.11
Lincar-1
MoE (nonlinear)
99.46 = 0.55
0.8
Noninaar-1
Linzar-2
Noriinear-2
Setting 2: (0.5,2), 8 (1,2), E (0.5, 3), o, = 2
0.4
Test accuracy (%)
Dispatch Entropy
0.2
Single (linear)
60.59
NA NA 1.294 + 0.036
100
200
300
400
500
Single (nonlinear)
72.29
Traning Epochs
MoE (lincar)
88.48 + 1.96
Figure 3: Illustration of router dis- patch entropy. We demonstrate the change of entropy of MoE during train- ing the synthetic data. MoE (linear)- and MoF. (nonlinear)-1 refer to Set- ting Table MoE (linear)-2 and MoE (nonlinear)-2 refer to Setting 2 in Table
MoE (nonlinear)
98.09 + 1.27
0.171 + 0.103
Table 1: Comparison between MoE. (linear) and MoF. (nonlinear) in setting. We report results of top-1 gating with noise for both linear and nonlinear models. Over ten random experiments, we report the average value + standard deviation for both test accuracy and dispatch entropy.
6.1 Synthetie-data Experiments
Datasets. We generate 16, 000 training examples and 16, 000 test examples from the data distribution
defined in Definition 3.1 with cluster number = 4, patch number P and dimension d = 50. We randomly shuffle the order of the patches of x after we generate data (x, y). We consider two
butts://gttibb.om/ccaa/////
8

---

Table 2: Comparison between MoE and single model on CIFAR-I0 and CIFAR- RO-Rotate datasets. We report the average lest accuracy over random experiments + the standard deviation.
CIFAR-10 (%)
CIFAR-1O-Rotate (%)
Single
80.68 = 0.45
76.78 + 1.79
CNN
MoE
80.31 + 0.62
79.60 1.25
Single
92.45 + 0.25
85.76 + 2.91
MobileNetV2
MoE
92.23 0.72
89.85 2.54
Single
95.51 0.31
88.23 0.96
ResNet18
MoE
95.32 + 0.68
92.60 + 2.01
parameter settings: 1. a ~ Uniform(0.5, (0.5, 3) (iniform(1,2).1 Uniform(O.5,33 and op = 1: Uniform(0.5, 2), B ~ Uniform(1,2)) 2). Uniform(0.5, 3) and og = 2. Note that Theorem 4.1 shows that when a and y follow the same distribution, neither single linear expert or single nonlinear expert can give good performance. Here we consider a more general and difficult setting when a and are from different distributions.
Models. We consider the performances of single linear CNN, single nonlinear CNN, linear MoE, and nonlinear MoE. The single nonlinear CNN architecture follows (3.I) with cubic activation function, while single linear CNN follows (3.1) with identity activation function. For both linear and nonlinear MoEs, we consider mixture of experts with each expert being a single linear CNN or a single nonlinear CNN. Finally, we train single models with gradient descent and train the MoEs with Algorithm We run 10 random experiments and report the average accuracy with standard deviation. I
Evaluation. To evaluate how well the router learned the underlying cluster structure of the data, we define the entropy of the router's dispatch as follows. Denote by fik,m the number of data cluster K that are dispatched to expert m. The total number of data dispatched to expert m is Tim = Σk= 7k.m and the total number of data Σk
Σ
Em=1 Tk.m- The dispatch entropy then defined as
M
entropy - M
2k=1
llog
(6.1)
When each expert receives the data from at most one cluster, the dispatch entropy will be zero. And a uniform dispatch will result in the maximum dispatch entropy.
As shown in Table T the linear MoE does not perform as well as the nonlinear MoE in Setting 1, with around 6% less test accuracy and much higher variance. With stronger random noise (Setting 2). the difference between the nonlinear MoE and linear MoE becomes even more significant. We also observe that the final dispatch entropy of nonlinear MoË is nearly zero while that of the linear MoE is large. In Figure 3 we further demonstrate the change of dispatch entropy during the training process. The dispatch entropy of nonlinear MoE significantly decreases, while that of linear MoE remains large. Such a phenomenon indicates that the nonlinear MoE can successfully learn the underlying cluster structure of the data while the linear MoE fails to do so.
6.2 Real-data Experiments
We further conduct experiments on real image datasets and demonstrate the importance of the clustering data structure the MoE layer in deep neural networks.
Datasets. We consider the CIFAR-10 dataset (Krizhevsky 2009) and the 10-class classification task. Furthermore, create CIFAR-10-Rotate dataset that has strong underlying cluster structure that is independent of its labeling function. Specifically, we rotate the images by 30 degrees and merge the rotated dataset with the original one. The task is to predict if the image is rotated, which a binary classification problem. We deem that some of the classes in CIFAR- 10 form underlying clusters in CIFAR-10- 1 In Appendix Al explain in detail how generate CIFAR-10 10--100-otttt and present some specific examples.
Models. For the MoE, we consider mixture of 4 experts with linear gating network. For the expert/single model architectures, we consider a CNN with 2 convolutional layers (architecture details are illustrated in Appendix A For more thorough evaluation, also consider expert/single models

---

with architecture including MobileNetV2 (Sandler et al. 2018 and ResNet18 (He et al. 2016 The training process of MoE also follows Algorithm
The experiment results are shown in T where we compare single and mixture models of different architectures over CIFAR-10 and CIFAR-10-Rotie datasets. We observe that the improvement of MoEs over single models differs largely on the different datasets. On CIFAR-IO, IO, the performance of MoEs is very close to the single models. However, on the CIFAR- ME-R--0--Roate dataset, we can observe a significant performance improvement from single models to MoEs. Such results indicate the advantage of MoE over single models depends on the task and the cluster structure of the data.
Visualization. Figure we visualize the latent embedding learned by MoEs (ResNet18) for the 10-class classification task in CIFAR-10 well as the binary classification task in CIFAR- 10- 10-Rotate We visualize the data with the same label to see cluster structures exist within each class. For CIFAR-10, choose y 1 ("car"), and plot the latent embedding of the data using t-SNE the left sub-figure, which does not show an salient cluster structure. For SEFAR-10-Rotan
we choose y ("rotated") and visualize the data using t-SNE the middle sub-figure. Here, we can observe a clear clustering structure even though the class signal is not provided during training, We take step further to investigate what in each cluster in the right sub-figure. We can observe that most of the examples in the "frog" class fall into one cluster, while examples of "ship" class mostly fall into the other cluster.
0.4
0.50
64-
0.-
-0.25
0.2
-050
0.75
0.5
-44
6.2
5.0
02
04
-05
0.0
0.5
1.0
-03
y=1 (car)
y=1 (rotated)
(frog, ship)
Figure 4: Visualization of the latent emhedding CIFAR-10 10 and CIFAR- 10-Rotate with chosen label y. The left sub-figure denotes the visualization of CIFAR-I0 when label is chosen to be 1 (car). The central sub-figure represents the visualization of -O-AR-----otate when label y is chosen red denotes that the data is from the ship class, and blue denotes that the data is from the frog class.
(rotated). On the right sub-figure,
7 Conclusion and Future Work
In this work, formally study the mechanism of the Mixture of Experts (MoF) layer for deep learning. To our knowledge, we provide the first theoretical result toward understanding how the MoE layer works in deep learning. Our empirical evidence reveals that the cluster structure of the data plays an important role in the success of the MoE layer. Motivated by these empirical observations, we study a data distribution with cluster structure and show that Mixture-of- -f-pert provably improves the test accuracy of a single expert of two-layer CNNs.
There several important future directions. First, current results for CNNs. It is interesting to extend our results to other neural network architectures, such as transformers. Second, our data distribution is motivated by the classification problem of image data. We plan extend our analysis to other types of data (e.g., natural language data).
Acknowledgments and Disclosure of Funding
We thank the anonymous reviewers and area chair for their helpful comments. ZC, YD. YW and QG are supported in part by the National Science Foundation CAREER Award 1906166 169, BIGDATA IIS- 1855099, and the Sloan Research Fellowship. YL. is supported in part by the NSF R12007517. The views and conclusions contained this paper are those of the authors and should not be interpreted representing any funding agencies.
10

---

References
ALLEN-ZHU, Z. and LI. Y. (2019). What can ResNet learn efficiently, going beyond kernels? In Advances in Neural Information Processing Systems.
ALLEN-ZHU, Z. and Lı, Y. (2020a). Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.044:
ALLEN-ZHU, Z. and LI, Y. (2020b). Feature purification: How adversarial training performs robust deep learning. arXiv preprint arXiv:2005.
ALLEN-ZHU, Z. and LI, Y. (2020c). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:2012.09816
ALLEN-ZHU, Z., LI, Y. and LIANG, Y. (2019a). Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in Neural Information Processing Systems.
ALLEN-ZHU, Z., LI, Y. and SONG, Z. (2019h). convergence theory for deep learning via over-paramettrizzaion.. In International Conference Machine Learning.
ANANDKUMAR, A., GE, R., HsU, D., KAKADE, M. and TELGARSKY, M. (2014). Tensor decompositions for learning latent variable models. Journal of machine learning research 15 2773-2832.
ANANDKUMAR, A., HSU, D. and KAKADE, M. (2012). method of moments for mixture models and hidden markov models. In Conference Learning Theory. JMLR Workshop and Conference Proceedings.
ARORA, S., DU, S., HU, W., LI, Z. and WANG, R. (2019a). Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning.
ARORA, S., DU, S. S., HU, W., LI, Z., SALAKHUTDINOY, R. and WANG, R. (2019b). On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems.
BAI, Y. and LEE, J. D. (2019). Beyond linearization: On quadratic and higher-order approximation of wide neural networks. arXiv preprint arXiv:1910.06699
BALAKRISHNAN, S., WAINWRIGHT, M. J. and Yu, B. (2017). Statistical guarantees for the em algorithm: From population to sample- based analysis. The Annals of Statistics 45 77-120.
BLARD, T. (2020). French sentiment analysis with bert. httpss//ggaabb.co//mpoppllldlell analysis- with- bert
CAO, Y., CHEN, Z., BELKIN, M. and GU, Q. (2022). Benign overfitting in two-layer convolutional neural networks. arXiv preprint arXiv:2202.06526.
CAO, Y. and Gu, Q. (2019). Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems.
CHAGANTY, A. T. and LIANG, P. (2013). Spectral experts for estimating mixtures of linear regressions. In International Conference on Machine Learning. PMLR.
COLLOBERT, R., BENGIO, S. and BENGIO, Y. (2002). A parallel mixture of svms for very large scale problems. Neural computation 14 1105-1114.
DAUPHIN, Y. N., FAN, A., AULI, M. and GRANGIER, D. (2017). Language modeling with gated convolutional networks, In International conference on machine learning. PMLR.
DE VEAUX, R. D. (1989), Mixtures of linear regressions. Computational Statistics & Data Analysis 227-245.
DEVLIN, J., CHANG, M., LEE, K. and TOUTANOVA, K. (2018). BERT: pre-training of deep bidirectional transformers for language understanding. CoRR abs/1810.04805.
11

---

DU, S. S., ZHAI, X., PocZOs, B. and SINGH, A. (2019). Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations.
DUA, D., BHOSALE, S., GOSWAMI, V., CROSS, J., LEWIS, M. and FAN, A. (2021). Tricks for training sparse translation models. arXiv preprint 1rXiv:2110.08246.
EIGEN, D., RANZATO, M. and SUTSKEVER, (2013). Learning factored representations in a deep mixture of experts. arXiv preprini arXiv:1312.4311
FARIA, S. and SOROMENHO, G. (2010). Fitting mixtures of linear regressions. Journal of Statistical Computation and Simulation 80 201-225.
FEDUS, W., ZOPH, B. and SHAZEER, N. (2021). Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, arXiv preprint arXiv:2101.33611.
GO, A., BHAYANI, R. and HUANG, L. (2009). Twitter sentiment classification using distant supervision. CS224N project report, Stanford 1 2009.
HE, K., ZHANG, X., REN, S. and SUN, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition.
HSU, D. J., KAKADE, S. M. and LIANG, P. (2012). Identifiability and unmixing of latent parse trees. Advances in neural information processing systems 25.
JACOBS, R. A., JORDAN, M. I., NOWLAN, S. J. and HINTON, G. E. (1991). Adaptive mixtures of local experts. Neural computation 79-87.
JACOT, A., GABRIEL, F. and HONGLER, C. (2018). Neural tangent kerne:: Convergence and generalization in neural networks. In Advances neural information processing systems.
JELASSI, S., MENSCH, A., GIDEL, G. and LI, Y. (2021). Adam is no better than normalized sgd: Dissecting how adaptivity improves gan performance.
JORDAN, M. I., GHAHRAMANI, Z. and SAUL,, L. K. (1997). Hidden markov decision trees. Advances in neural information processing systems 501-507.
JORDAN, M. I. and JACOBS, R. A. (1994). Hierarchical mixtures of experts and the algorithm.
Neural computation 6 181-214.
KHALILI, A. and CHEN, J. (2007). Variable selection in finite mixture of regression models. Journal of the american Statistical association 102 1025-1038.
KRIZHEVSKY, A. (2009). Learning multiple layers of features from tiny images. Tech. rep.
LEWIS, M., BHOSALE, S., DETTMERS, T., GOYAL, N. and ZETTLEMOYER, L. (2021). Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning. PMLR.
LI, Y. and LIANG, (2018). Learning overparameterized neural networks via stochastic gradient descent on structured data. In Advances in Neural Information Processing Systems.
LI, Y.. MA, T. and ZHANG, H. R. (2020). Learning over-parametrized two-layer neural networks beyond ntk. In Conference on learning theory. PMLR.
LIANG, P., BOUCHARD-CÔTTÉ A.. KLEIN, D. and TASKAR, B. (2006). An end-to-end discrim- inative approach to machine translation. In Proceedings of the 21st International Conference on Computational Linguistics and 44th Annual Meeting of the Association for Computational Linguistics.
QUATTONI, A., COLLINS, M. and DARRELI,, T. (2004). Conditional random fields for object recognition. Advances in neural information processing systems 17.
ROLLER, S., SUKHBANTAR, S., WESTON, J. FT AL.. (2021). Hash layers for large sparse models. Advances in Neural Information Processing Systems 34 17555-17566.
12

---

SANDLER, M., HOWARD, A., ZHU, M., ZHMOGINOV, A. and CHEN, L.-C. (2018). Mobilenetv2: Inverted residuals and linear bottlenecks. Proceedings of the IEEE conference on computer vision and pattern recognition.
SHAAZER,, N., MIRIIOSEINI, A., MAZIARZ, K., DAVIS, A., LE, Q.. HINTON, G. and DEAN, J. (2017). Outrageously large neural networks: The sparsely-gaue mixture-oo saaeee--f--xppeee offeeeertts layer. arXiv preprint arXiv:1701.0653:
SMETANIN, S. and KOMAROV, M. (2019). Sentiment analysis of product reviews in russian using convolutional neural networks. In 20/9 IEEE 21 conference on business informatics (CBI), vol. 1. IEEE.
TRESP, V. (2001). Mixtures of gaussian processes. Advances in neural information processing systems 654-660.
VAN DER MAATEN, L. and HINTON, G. (2008). Visualizing data using t-sne. Journal of machine learning research 9.
VASWANI, A., SHAZEER, N., PARMAR, N., USZKOREIT, J., JONES, L., GOMEZ, A. N., KAISER, E. and POLOSUKHIN, I. (2017). Attention is all you need. In Advances in neural information processing systems.
VECCI, L., PIAZZA, F. and UNCINI, A. (1998). Learning and approximation capabilities of adaptive spline activation function neural networks. Neural Networks 11 259-270.
WANG, Y. and MORI, G. (2009). Max-margin hidden conditional random fields for human action recognition. In 2009 IEEE Conference on Computer Vision and Pattern Recognition. IEEE.
WANG, Z., GU, Q., NING, Y. and LIU, H. (2015). High dimensional em algorithm: Statistical optimization and asymptotic normality. Advances in neural information processing systems 28.
WEN, Z. and LI, Y. (2021). Toward understanding the feature learning process self-supervised contrastive learning. In International Conference Machine Learning. PMLR.
YI, X., CARAMANIS, C. and SANGHAVI, S. (2014). Alternating minimization for mixed linear regression. In International Conference on Machine Learning. PMLR.
Zou, D., CAO, Y., LI, Y. and GU, Q. (2021). Understanding the generalization of adam in learning neural networks with proper regularization. arXiv preprint arXiv:2108.11371
Zou, D., CAO, Y., ZHOU, D. and GU, Q. (2018). Stochastic gradient descent optimizes over- parameterized deep relu networks. arXiv preprint arXiv:1811.08888.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] This paper gives the first theoretical result toward formally understanding the mechanism of the MoE layer for deep learning.
(b) Did you describe the limitations of your work? [Yes] We note in Section that our analysis of the MoE layer need early stopping which we believe can be waived by adding some well-signed regularization. We will explore this in future work.
(c) Did you discuss any potential negative societal impacts of your work? [No] We seek to mathematically understand the MoE layer in Deep Learning, it is not clear what potential negative impacts a deeper theoretical understanding of this algorithm would bring.
(d) Have you read the ethics review guidelines and ensured that your paper conforms them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] See our Appendix.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you're 'r using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "zyGL88n0xSb", "reviewer": "Reviewer_TAGw", "summary": "The authors provide theoretical and empirical study on the mixture of expert layer, which has recently become increasingly used and for which there is limited theoretical understanding. They show that cluster structure in the data is important for success of MoE, that under specific data distribution with cluster structure a single expert 2 layer CNN cannot generalize while a linear MoE can. They also show the advantages of non-linear MoE.  ", "questions": "none", "limitations": "None", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\n*Strengths*\n\n- Analysis on an important class of model/network components that are not well understood\n- Theoretical analysis appears sound\n- The writing is clear\n\n*Weakness*\n- The distribution studied is limited, i would like to see a more in depth discussion regarding why the distribution (with K centroids) is representative of real data. Indeed the experiments hint naturally many data does not have this underlying structure (e.g. CIFAR-10 performance is not improved by using MoE)\n- (minor) There is some slight disconnect between the original question posed \"Why do the experts in MoE diversify instead of collapsing into a single model? And how can the router learn to dispatch the data to the right expert?\" and the work done, would be goo to make this flow smoother.\n\n\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": ["I don’t know"]}, "extra_text": {}}, {"review_id": "lyDSHGn0Go9", "reviewer": "Reviewer_nUCK", "summary": "The presented paper aims to crack open the blackbox of sparsely activated mixture-of-expert models and tries to understand what leads to the improved performance on many neural network tasks. Through experiments on a toy task, CIFAR-10, as well as a new CIFAR-10-Rotate dataset, the key-findings include that performance improvements are highly connected to the intrinsic cluster structure of the data  and dependent on the non-linear activation function in the experts.  ", "questions": "see above", "limitations": "Limitations are appropriately addressed in the conclusion and future extensions are given.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n- Intuitive visualisations on a toy task to understand the behaviour when training sparsely activated models and additional experiments on a created dataset named CIFAR-10-Rotate with strong underlying cluster structure.\n- Empirical evidence that MoE's are heavily dependent on the non-linear activation function to learn the target distribution backed by visualisations for the router dispatch entropy in Figure 3.\n-  Great effort of putting theory behind some of the empirical observations found in the field. Theory papers like this are essential for effectively moving forward in the empirical-driven field of deep learning.\n\n---\n\n### Weaknesses: \n\n### [Update]: Weaknesses have been addressed in the rebuttal and I've increased the score from 6 to 7. \n\n- **(medium)**: Table 1 + 2 only compare Single with linear + nonlinear activation function to the MoE models and shows that the Single model isn't able to learn the problem at hand to the same extend as the MoE variants. Since effectively MoE models have more parameters to play with (albeit them being in parallel branches), I believe adding total + effective parameter counts and adding one more line where we extend the Single model to the total number of parameters in the MoE model would be beneficial. This would verify that MoE's aren't just finding winning lottery tickets in their parameter initialization which allows them to perform better.\n- **(small)**: The paper does consider other non-linear activation functions in the Appendix but here the setting is different than in the main paper. I would like to see Table 6 for the settings in Table 1, i.e. with $\\alpha \\in(0.5,2), \\beta \\in(1,2), \\gamma \\in(0.5,3), \\sigma_{n}=1$ and  $\\alpha \\in(0.5,2), \\beta \\in(1,2), \\gamma \\in(0.5,3), \\sigma_{n}=2$.\n- **(large)**: Comparing the results of CIFAR-10 and CIFAR-10-Rotate, I find it surprising that MoE's don't seem to benefit CIFAR-10 although it should also have intrinsic clusters (similar to CIFAR-10-Rotate) which different experts should be able to utilize. In a paper that tries to understand MoE's and states that the cluster structure is important, I was hoping for a better explanation (or at least a hypothesis) in l.340 -- l.345 that would help to assess when MoE's are beneficial and what properties to look for in the dataset. In its current state, I do not know how I would transfer some of the findings in the paper to other real-world datasets e.g. multilingual machine translation. Maybe adding additional datasets from other domains or other computer vision datasets would help to solve this problem. \n- **(large)**: Continuing from the previous point, a nice analysis would be to visualise the latent embeddings on CIFAR-10 & CIFAR-10-Rotate similar to Figure 1 using t-SNE and see if we can distill any inherent characteristics about the clustering behaviour. My assumption would be that in CIFAR-10-Rotate the inter-cluster separation/distance is a lot higher than in regular CIFAR-10 which enables better expert learning. Thinking this further, maybe we could add additional regularization during the the training of MOE's on CIFAR-10 to enforce better separation and simplify the training for the different experts.\n\n\n---\n\n### Minor Improvement Suggestions:\n\n- Related literature that seems relevant enough to cite: \n    - \"*BASE Layers: Simplifying Training of Large, Sparse Models*\" (Lewis et al., ICML 2021)\n    - \"*Hash Layers For Large Sparse Models*\" (Roller et al., NeurIPS 2021)\n    - \"*Tricks for Training Sparse Translation Models*\" (Dua et al., arXiv 2021)\n- Typos & Formatting:\n    - Table 1 can be made more easy to read and parse by using the $\\LaTeX$ `booktabs` package with proper spacing. Similarly, Table 2--6 would benefit from it as well.\n    - The definition the empirical loss after l. 175 & dispatch Entropy after l. 316 should have an Equation number\n    -  l. 121 Radamacher -> Rademacher\n- l. 165 refer to the experiments with other activation functions in Table 6 in the appendix where you consider more commonly used choices such as $\\mathrm{ReLU}$ / $\\mathrm{GeLU}$.\n- l. 193 cite Adam paper\n- Generally refer more to the Appendix in the main paper. It is really hard to follow the reading flow without it.\n\n---\n\nApart from the mentioned points, the paper is well written and presentation, structure, and general layout of thoughts are up to NeurIPS standard. Related work is mentioned and cited accordingly apart from the previously mentioned papers.\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": "K-h9OiPLGVa", "reviewer": "Reviewer_RKSY", "summary": "This paper engages in a formal study of a mixture-of-experts model (MOE) by considering a specific synthetic dataset with multiple clusters and a particular choice of the MOE model architecture. It outlines several theoretical and empirical results that show the importance of having multiple experts to properly fit the training data and the importance of nonlinearity in the architecture of an individual expert. It also mentions several typical problems that occur when training MOE models and discusses a set of techniques that address these issues. Finally, it extends empirical results to a more realistic rotated-CIFAR-10 dataset.\n", "questions": "1. While Theorem 4.1 requires the $\\alpha$ and the $\\gamma$ distributions to match, empirical results discussed in Section 6.1 were obtained in the $\\mathcal{D}_\\alpha \\ne \\mathcal{D}_\\gamma$ case. Are there reasons to think that it would be interesting to confirm expected results for a single-expert model and more importantly verify that mixture-of-experts models perform sufficiently well in this case (where the distributions of the signal and the noise multipliers are effectively indistinguishable)?\n2. It is not clear if the result of Theorem 4.2 actually requires the experts to be nonlinear functions, can this question be clarified? If Theorem 4.2 also holds for linear experts, then perhaps it should not be summarized as \"Nonlinear MOE performs well\"? Currently it is a little confusing to first see authors mention that both linear and nonlinear MOEs are expected to fit the training data, yet only see the result seemingly only holding for nonlinear functions.\n", "limitations": "I did not identify any potential negative societal impacts of this work. While also being true for a vast majority of other theoretical studies, the proposed framework is somewhat limited by a number of simplifying choices that make it difficult to extend most results beyond the studied scenario.\n", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n1. This paper outlines a simple synthetic data distribution and a MOE model architecture that make it possible to conduct theoretical analysis of this problem and show a conceptual difference between a single- and multiple-expert scenarios. The proposed system is a simple model showing the essence of the advantage of the mixture-of-experts networks versus their simpler single-expert counterparts.\n2. The empirical study goes beyond the simple synthetic setup and is extended to a rotated CIFAR-10 dataset showing the improvement provided by the mixture-of-experts model in this simple yet realistic case.\n3. The paper outlines several practical techniques that altogether allow one to train mixture-of-experts models that are otherwise prone to various training issues.\n\n### Weaknesses\n\n1. The paper has multiple minor misprints and inaccuracies and may need another careful pass (numerous minor misprints, missing or unnecessary articles, missing parentheses in equations, for example in Lemma 5.2, etc.).\n2. Several aspects of the theoretical analysis including the importance of expert nonlinearity for specific theoretical results may currently confuse the reader.\n3. The problem setup is somewhat simplistic and cannot be easily extended to more realistic setups. But this limitation can of course be difficult to avoid if one wants to obtain meaningful theoretical predictions.\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": "4f4S5adHKS", "reviewer": "Reviewer_By4k", "summary": "The paper analyzes several features of mixture of experts (MoE) and derive a few conclusions such as the superiority of MoE in terms of the classification accuracy and it good loading balance.", "questions": "The key limitation is the analysis is not deep enough. If the analysis is empirical, more experiments are needed. If the analysis is based on math, then concrete prove should be added. Most of the math shown in this paper is used to describe the design of the experiments and does not help me understand the working mechanisms of MoE. Then more analysis on larger model and datasets are also necessary.", "limitations": "Yes", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength: The paper is well written and clear motivation. The results that show MoE is more effective on clustered data is interesting.\n\nWeakness\n\nThe analysis about MoE is too superficial. In the experimental, the author shows MoE can beat single expert, which is actually a common sense. It is also not surprising that a non-linear MoE is better than a plain MoE. For most of MoE papers, they use the MLP as the expert, which is a non-linear module.\n\nThe experiments are too toy to reflect the real working mechanisms behind MoE. Nowadays, MoE is often deployed on ViT architectures that are trained with JFT 300M or ImageNet 22k data. It is really hard to say the analysis built on synthetic data and cifar 10 is still applicable to large models. For instance, in L34-35, the author posted a question about why experts are not collapsed into a single expert. From my view, it is due to the loading balance loss that is used during training, otherwise it will indeed convert to one expert. I hope the author should at least further clarify his/her point about this.", "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": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept