Towards Understanding the Mixture-of-Experts Layer in Deep Learning
Anonymous Author(s)
AfﬁliationAddress email
Abstract
The Mixture-of-Experts (MoE) layer, a sparsely-activated model controlled by arouter, has achieved great success in deep learning.
However, the understanding ofsuch architecture remains elusive.
In this paper, we formally study how the MoElayer improves the performance of neural network learning and why the mixturemodel will not collapse into a single model.
Our empirical results suggest thatthe cluster structure of the underlying problem and the non-linearity of the expertare pivotal to the success of MoE. This motivates us to consider a challengingclassiﬁcation problem with intrinsic cluster structures.
Theoretically, we provedthat this problem is hard to solve by a single expert such as a two-layer convolutionalneural 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 showsthat the router can learn the cluster-center features, which helps divide the inputcomplex problem into simpler classiﬁcation sub-problems that individual expertscan conquer.
To our knowledge, this is the ﬁrst theoretical result toward formallyunderstanding the mechanism of the MoE layer for deep learning.
Introduction
The Mixture-of-Expert (MoE) structure (Jacobs et al., 1991; Jordan and Jacobs, 1994) is a classicdesign 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 is an extension of the
MoE model to deep neural networks, has achieved remarkable success in deep learning.
Generallyspeaking, an MoE layer contains many experts that share the same network architecture and aretrained by the same algorithm, with a gating (or routing) function that routes individual inputs to afew experts among all the candidates.
Through the sparse gating function, the router in the MoE layercan route each input to the top- K ( K ≥ 2) best experts (Shazeer et al., 2017), or the single ( K = 1 )best expert (Fedus et al., 2021).
This routing scheme only costs the computation of K experts for anew input, which enjoys fast inference time.
Despite the great empirical success of the MoE layer, the theoretical understanding of such architectureis still elusive.
In practice, all experts have the same structure, initialized from the same weightdistribution (Fedus et al., 2021) and are trained with the same optimization conﬁguration.
Therouter is also initialized to dispatch the data uniformly.
It is unclear why the experts can diverge todifferent functions that are specialized to make predictions for different inputs, and why the router canautomatically learn to dispatch data, especially when they are all trained using simple local searchalgorithms such as gradient descent.
Therefore, we aim to answer the following questions:
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?
In this paper, in order to answer the above question, we consider the natural “mixture of classiﬁcation”data distribution with cluster structure and theoretically study the behavior and beneﬁt of the MoElayer.
We focus on the simplest setting of the mixture of linear classiﬁcation, where the datadistribution has multiple clusters, and each cluster uses separate (linear) feature vectors to representthe labels.
In detail, we consider the data generated as a combination of feature patches, clusterpatches, and noise patches (See Deﬁnition 3.1 for more details).
We study training an MoE layer basedon the data generated from the “mixture of classiﬁcation” distribution using gradient descent, whereeach expert is chosen to be a two-layer CNN.
The main contributions of this paper are summarized asfollows:
• We ﬁrst prove a negative result (Theorem 4.1) that any single expert, such as two-layer CNNswith arbitrary activation function, cannot achieve a test accuracy of more than 87 .
5% on our datadistribution.•
Empirically, we found that the mixture of linear experts performs better than the single expert butis still signiﬁcantly worse than the mixture of non-linear experts.
Figure 1 provides such a resultin a special case of our data distribution with four clusters.
Although a mixture of linear modelscan represent the labeling function of this data distribution with 100% accuracy, it fails to learn soafter training .
We can see that the underlying cluster structure cannot be recovered by the mixtureof linear experts
and neither the router nor the experts are diversiﬁed enough after training.
In
contrast, the mixture of non-linear experts can correctly recover the cluster structure and diversify.
• Motivated by the negative result and the experiment on the toy data
we study a sparsely-gated
MoE model with two-layer CNNs trained by gradient descent.
We prove that this MoE model can
achieve nearly 100% test accuracy efﬁciently (Theorem 4.2).
• Along with the result on the test accuracy, we formally prove that each expert of the sparselygated MoE model will be specialized to a speciﬁc portion of the data (i
e., at least one cluster),which is determined by the initialization of the weights.
In the meantime, the router can learn thecluster-center features and route the input data to the right experts.• Finally, we also conduct extensive experiments on both synthetic and real datasets to corroborateour theory.
Notation.
We use lower case letters, lower case bold face letters, and upper case bold face letters todenote scalars, vectors, and matrices respectively.
We denote a union of disjoint sets ( A i : i ∈ I )by (cid:116)
i ∈
I
A i .
For a vector x , we use (cid:107) x
(cid:107) 2 to denote its Euclidean norm.
For a matrix W , we use(cid:107) W (cid:107) F to denote its Frobenius norm.
Given two sequences { x n } and { y n } , we denote x n = O ( y n )if | x n | ≤ C 1 | y n | for some absolute positive constant C 1 , x n = Ω( y n ) if | x n | ≥ C 2 | y n | for someabsolute positive constant C 2 , and x n = Θ( y n ) if C 3 | y n | ≤ | x n | ≤ C 4 | y n | for some absoluteconstants C 3 , C 4 > 0 .
We also use (cid:101) O ( · ) to hide logarithmic factors of d in O ( · ) .
Additionally, wedenote x n = poly( y n ) if x n = O ( y Dn ) for some positive constant D , and x n = polylog( y n ) ifx n = poly(log( y n )) .
We also denote by x n = o ( y n ) if lim n →∞ x n /y n = 0 .
Finally we use [ N ] todenote the index set { 1 , . .
. , N } .
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 basedon various base expert models such as support vector machine (Collobert et al., 2002) ,
Gaussianprocesses (Tresp, 2001), or hidden Markov models (Jordan et al., 1997).
In order to increase themodel capacity to deal with the complex vision and speech data, Eigen et al.
(2013) extended the
MoE structure to the deep neural networks, and proposed a deep MoE model composed of multiplelayers of routers and experts.
Shazeer et al.
(2017) simpliﬁed the MoE layer by making the output ofthe gating function sparse for each example, which greatly improves the training stability and reducesthe computational cost.
Since then, the MoE layer with different base neural network structures(Shazeer et al., 2017; Dauphin et al., 2017; Vaswani et al., 2017) has been proposed and achievedtremendous successes in a variety of language tasks.
Very recently, Fedus et al.
(2021) improved theperformance of the MoE layer by routing one example to only a single expert instead of K experts,which further reduces the routing computation while preserving the model quality.
Mixture of Linear Regressions/Classiﬁcations.
In this paper, we consider a “mixture of classiﬁcation” model.
This type of models can be dated back to (De Veaux, 1989; Jordan and Jacobs, 1994; Faria and Soromenho, 2010) and has been applied to many tasks including object
[Figure]
recognition (Quattoni et al., 2004) human action recognition (Wang and Mori, 2009), and machinetranslation (Liang et al., 2006).
In order to learn the unknown parameters for mixture of linearregressions/classiﬁcation model, (Anandkumar et al., 2012; Hsu et al., 2012; Chaganty and Liang,2013; Anandkumar et al., 2014; Li and Liang, 2018) studies the method of moments and tensorfactorization.
Another line of work studies speciﬁc algorithms such as Expectation-Maximization(EM) algorithm (Khalili and Chen, 2007;
Yi et al., 2014; Balakrishnan et al., 2017; Wang et al., 2015).
Theoretical Understanding of Deep Learning.
In recent years, great efforts have been made toestablish 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., 2018) andgeneralization (Allen-Zhu et al., 2019a; Arora et al., 2019a,b; Cao and Gu, 2019) guarantees in theso-called “neural tangent kernel” (NTK) regime, where the parameters stay close to 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,b,c; Li et al., 2020; Cao et al.
,2022; Zou et al., 2021; Wen and Li, 2021)
studied the learning dynamic of neural networks beyondthe NTK regime.
It is worthwhile to mention that our analysis of the MoE model is also beyond the
NTK regime.
Problem Setting and Preliminaries
We consider an MoE layer with each expert being a two-layer CNN trained by gradient descent (GD)over n independent training examples { ( x i , y i ) } ni =1 generated from a data distribution D .
In thissection, we will ﬁrst introduce our data model D , and then explain our neural network model and thedetails of the training algorithm.
3.1 Data distribution
We consider a binary classiﬁcation 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 (1) , x (2) , . . .
, x ( P ) ) ∈( R d ) P is a collection of P patches and y ∈ {± 1 } is the data label.
We consider data generated from
K clusters.
Each cluster k ∈ [ K ] has a label signal vector v k and a cluster-center signal vector c k with (cid:107) v k (cid:107) 2 = (cid:107) c k (cid:107) 2 = 1 .
For simplicity, we assume that all the signals { v k } k ∈ [ K ] ∪ { c k } k ∈ [ K ] areorthogonal with each other.
Deﬁnition 3.1.
A data pair ( x , y ) ∈ ( R d ) P × {± 1 } is generated from the distribution D as follows.
• Uniformly draw a pair ( k, k (cid:48) ) with k (cid:54) =
k (cid:48) from { 1 , . .
. , K } .•
Generate the label y ∈ {± 1 } uniformly, generate a Rademacher random variable (cid:15) ∈ {± 1 } .•
Independently generate random variables α, β, γ from distribution D
α , D β , D γ .
In this
paper, weassume there exists absolute constants C 1 , C 2 such that almost surely 0 < C 1 ≤ α, β, γ ≤ C 2 .
• Generate x as a collection of P patches: x = ( x (1) , x (2) , . . .
, x ( P ) ) ∈ ( R d ) P , where
– Feature signal.
One and only one patch is given by yα v k .
– Cluster-center signal.
One and only one patch is given by β c k .
– Feature noise.
One and only one patch is given by (cid:15)γ v k (cid:48) .
– Random noise.
The rest of the P − 3 patches are Gaussian noises that are independently drawnfrom N (0 , ( σ 2 p /d )
· I d ) where σ p is an absolute constant.
How to learn this type of data?
Since the positions of signals and noises are not speciﬁed in
Deﬁnition 3.1, it is natural to use the CNNs structure that applies the same function to each patch.
Wepoint out that the strength of the feature noises γ could be as large as the strength of the feature signalsα .
As we will see later in Theorem 4.1, this classiﬁcation problem is hard to learn with a single expert,such as any two-layer CNNs (any activation function with any number of neurons).
However, sucha classiﬁcation problem has an intrinsic clustering structure that may be utilized to achieve betterperformance.
Examples can be divided into K clusters ∪ k ∈ [ K ] Ω k based on the cluster-center signals:an example ( x , y ) ∈
Ω k if and only if at least one patch of x aligns with c k .
It is not difﬁcult to showthat the binary classiﬁcation sub-problem over Ω k can be easily solved by an individual expert.
Weexpect the MoE can learn this data cluster structure from the cluster-center signals.
Signiﬁcance of our result.
Although this data can be learned by existing works on a mixture oflinear classiﬁers with sophisticated algorithms (Anandkumar et al., 2012; Hsu et al., 2012; Chagantyand Liang, 2013), the focus of our paper is training a mixture of nonlinear neural networks, a morepractical model used in real applications.
When an MoE is trained by variants of gradient descent, weshow that the experts automatically learn to specialize on each cluster , while the router automaticallylearns to dispatch the data to the experts according to their specialty .
Although from a representationpoint of view, it is not hard to see that the concept class can be represented by MoEs, our result isvery signiﬁcant as we prove that gradient descent from random initialization can ﬁnd a good MoEwith non-linear experts efﬁciently.
To make our results even more compelling, we empirically showthat MoE with linear experts, despite also being able to represent the concept class, cannot be trainedto ﬁnd a good classiﬁer efﬁciently.
3.2 Structure of the MoE layer
An MoE layer consists of a set of M “expert networks” f 1 , . . .
, f M , and a gating network which isgenerally set to be linear (Shazeer et al., 2017; Fedus et al., 2021).
Denote by f m ( x ; W ) the output ofthe m -th expert network with input x and parameter W .
Deﬁne an M -dimensional vector h
( x ; Θ ) = 154 (cid:80)
p ∈ [ P ] Θ (cid:62) x ( p ) as the output of the gating network parameterized by Θ =
[ θ 1 , . .
, θ M ] ∈ R d × M .
The output F of the MoE layer can be written as follows:
[Equation]
where T x ⊆ [ M ] is a set of selected indices and π m ( x ; Θ ) ’s are route gate values given by
[Equation]
Expert Model.
In practice, one often uses nonlinear neural networks as experts in the MoE layer.
Infact, 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, we consider a convolution neural network as follows:
[Equation]
where w m,j ∈ R d is the weight vector of the j -th ﬁlter (i.e., neuron) in the m -th expert, J is thenumber of ﬁlters (i.e., neurons).
We denote W m =
[ w m, 1 , . .
. , w m,J ] ∈ R d × J as the weight matrix of the m -th expert and further let W = { W m } m ∈
[ M ] as the collection of expert weight matrices.
Fornonlinear CNN, we consider the cubic activation function σ ( z ) =
z 3 , which is one of the simplestnonlinear activation functions (Vecci et al., 1998).
We also include the experiment for other activationfunctions such as RELU in Appendix Table 6.
Top-1 Routing Model.
A simple choice of the selection set T x would be the whole experts set
T x =
[ M ] (Jordan and Jacobs, 1994), which is the case for the so-called soft-routing model.
However,it would be time consuming to use soft-routing in deep learning.
In this paper, we consider “switchrouting”, which is introduced by Fedus et al.
(2021) to make the gating network sparse and save thecomputation time.
For each input x , instead of using all the experts, we only pick one expert from[ M ] , i.e., |T x | = 1 .
In particular, we choose T x = argmax m { h m ( x ; Θ ) } .
Require: Number of iterations T , expert learning rate η , router learning rate η r , initialization scale σ 0 , training set S
= { ( x i , y i ) } ni
=1 .
[Table] [Table]
3.3 Training Algorithm
Given the training data S =
{ ( x i , y i ) } ni =1 , we train F with gradient descent to minimize the followingempirical loss function:
[Equation]
where (cid:96) is the logistic loss deﬁned as (cid:96) ( z )
=
log(1 + exp( − z )) .
We initialize Θ (0) to be zero andinitialize each entry of W (0) by i.i.d N (0 , σ 20 ) .
Zero initialization of the gating network is widelyused in MoE training.
As discussed in Shazeer et al.
, it can help avoid out-of-memory errorsand initialize the network in a state of approximately equal expert load (see (5.1) for the deﬁnition ofexpert load).
Instead of directly using the gradient of empirical loss (3.2) to update weights, we add perturbation to the router and use the gradient of the perturbed empirical loss to update the weights.
Inparticular, the training example x i will be distributed to argmax m { h m ( x i ; Θ ( t ) )
+ r ( t ) m,i } instead,where { r ( t ) m,i } m ∈
[ M ] ,i ∈ [ n ] are random noises.
Adding noise term is a widely used training strategyfor sparsely-gated MoE layer (Shazeer et al., 2017; Fedus et al., 2021), which can encourage exploration across the experts and stabilize the MoE training.
In this paper, we draw { r ( t ) m,i } m ∈
[ M ] ,i ∈ [ n ]independently from the uniform distribution Unif [0 , 1] and denotes its collection as r ( t ) .
Therefore,the perturbed empirical loss at iteration t can be written as
[Equation]
where m i,t = argmax m { h m ( x i ; Θ ( t ) )
+ r ( t ) m,i } .
Starting from the initialization W (0) , the gradientdescent update rule for the experts is
[Equation]
where η > 0 is the expert learning rate.
Starting from the initialization Θ (0) , the gradient update rulefor the gating network is
[Equation]
where η r > 0 is the router learning rate.
In practice, the experts are trained by Adam to make surethey have similar learning speeds.
Here we use a normalized gradient which can be viewed as asimpler alternative to Adam (Jelassi et al., 2021).
Main Results
In this section, we will present our main results.
We ﬁrst provide a negative result for learning with asingle expert.
Theorem 4.1 (Single expert performs poorly) .
Suppose D α = D γ in Deﬁnition 3.1, then any functionwith the form F ( x ) =
(cid:80) Pp =1 f ( x ( p ) ) will get large test error P ( x ,y )
∼D (cid:0) yF ( x ) ≤ 0 (cid:1) ≥ 1 / 8 .
Theorem 4.1 indicates that if the feature noise has the same strength as the feature signal i.e.,
D α = D γ , any two-layer CNNs with the form F ( x ) =
(cid:80)
[-Jj =
-]j ∈
[ J ] a j (cid:80) p ∈ [ P ] σ ( w (cid:62) j x ( p ) + b j ) can’tperform well on the classiﬁcation problem deﬁned in Deﬁnition 3.1 where σ can be any activationfunction.
Theorem 4.1 also shows that a simple ensemble of the experts may not improve theperformance because the ensemble of the two-layer CNNs is still in the form of the function deﬁnedin Theorem 4.1.
As a comparison, the following theorem gives the learning guarantees for training an MoE layer thatfollows the structure deﬁned in Section 3.2 with cubic activation function.
Theorem 4.2 (Nonlinear MoE performs well) .
Suppose the training data size n = Ω( d ) .
Chooseexperts number M = Θ( K log K log log d ) , ﬁlter size J =
Θ(log
M log log d ) , initialization scaleσ 0 ∈
[ d − 1 / 3 , d − 0 . 01 ] , learning rate η = (cid:101) O ( σ 0 ) , η r = Θ( KM 3 ) η .
Then with probability at least1 − o (1) , Algorithm 1 is able to output ( Θ ( T ) , W ( T ) ) within T = (cid:101)
O ( η − 1 ) iterations such that thenon-linear MoE deﬁned in Section 3.2 satisﬁes
• Training error is zero, i e., y i F ( x i ; Θ ( T ) , W ( T ) ) > 0 , ∀ i ∈ [ n ] .•
Test error is nearly zero, i
e., P ( x ,y ) ∼D (cid:0) yF ( x ; Θ ( T ) , W ( T ) ) ≤ 0 (cid:1)
=
o (1) .
More importantly, the experts can be divided into a disjoint union of K non-empty sets [ M ] =(cid:116)
k ∈
[ K ] M k and
• (Each expert is good on one cluster)
Each expert m ∈ M k performs good on the cluster Ω k ,
P ( x ,y ) ∼D ( yf m ( x ; W ( T ) ) ≤ 0 | ( x , y ) ∈
Ω k )
=
o (1) .
• (Router only distributes example to good expert)
With probability at least 1 − o (1) , an examplex ∈ Ω k will be routed to one of the experts in M k .
Theorem 4.2 shows that an MoE performs well on the classiﬁcation problem in Deﬁnition 3.1.
In addition, the router will learn the cluster structure and divide the problem into K simpler subproblems, each of which is associated with one cluster.
In particular, each cluster will be classiﬁedaccurately by a subset of experts.
On the other hand, each expert will perform well on at least onecluster.
Furthermore, together with Theorem 4.1, Theorem 4.2 suggests that there exist problem instances in
Deﬁnition 3.1 (i.e., D α
= D γ ) such that an MoE provably outperforms a single expert.
Overview of Key Techniques
A successful MoE layer needs to ensure that the router can learn the cluster-center features anddivide the complex problem in Deﬁnition 3.1 into simpler linear classiﬁcation sub-problems thatindividual experts can conquer.
Finding such a gating network is difﬁcult because this problem ishighly non-convex.
In the following, we will introduce the main difﬁculties in analyzing the MoElayer and the corresponding key techniques to overcome those barriers.
Main Difﬁculty 1: Discontinuities in 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 smallperturbation of the gating network outputs h ( x ; Θ ) + δ may change the router behavior drastically ifthe second largest gating network output is close to the largest gating network output.
Key Technique 1: Stability by Smoothing.
We point out that the noise term added to the gatingnetwork output ensures a smooth transition between different routing behavior, which makes therouter more stable.
This is proved in the following lemma.
Lemma 5.1.
Let h , (cid:98) h ∈ R M to be the output of the gating network and { r m } Mm =1 to be the noiseindependently drawn from Unif[0,1].
Denote p , (cid:98) p ∈ R M to be the probability that experts get routed,i.e., p m =
P (argmax m (cid:48) ∈
[ M ] { h m (cid:48)
+ r m (cid:48) }
= m ) , (cid:98) p m = P (argmax m (cid:48) ∈
[ M ] { (cid:98) h m (cid:48)
+ r m (cid:48) } = m ) .
Thenwe have that (cid:107) p − (cid:98) p (cid:107) ∞ ≤ M 2 (cid:107) h − (cid:98) h (cid:107) ∞ .
Lemma 5.1 implies that when the change of the gating network outputs at iteration t and t (cid:48) is small, i.e.,(cid:107) h ( x ; Θ ( t ) )
− h
( x ; Θ ( t (cid:48) ) )
(cid:107) ∞ , the router behavior will be similar.
So adding noise provides a smoothtransition from time t to t (cid:48) .
It is also worth noting that Θ is zero initialized.
So h ( x ; Θ (0) )
=
andthus each expert gets routed with the same probability p m = 1 /M by symmetric property.
Therefore,at the early of the training when (cid:107) h ( x ; Θ ( t ) ) − h ( x ; Θ (0) )
(cid:107) ∞ is small, router will almost uniformlypick one expert from [ M ] , which helps exploration across experts.
Main Difﬁculty 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 featuresbecause all the experts look the same: they share the same network architecture and are trained bythe same algorithm.
The only difference would be the initialization.
Moreover, if the router makes amistake at the beginning of the training, the experts may amplify the mistake because the experts willbe trained based on mistakenly dispatched data.
Key Technique 2: Experts from Exploration.
Motivated by the key technique 1, we introduce anexploration stage to the analysis of MoE layer during which the router almost uniformly picks one expert from [ M ] .
This stage starts at t = 0 and ends at T 1 = (cid:98) η − 1 σ 0 .
5 0 (cid:99) (cid:28) T = (cid:101)
O ( η − 1 ) and the gatingnetwork remains nearly unchanged (cid:107) h
( x ; Θ ( t ) )
− h
( x ; Θ (0) ) (cid:107) ∞
=
O ( σ 1 . 5 0 ) .
Because the experts aretreated almost equally during exploration stage, we can show that the experts become specialized tosome speciﬁc task only based on the initialization.
In particular, the experts set [ M ] can be divided into
K nonempty disjoint sets [ M ] = (cid:116) k M k , where M k := { m | argmax k (cid:48) ∈
[ K ] ,j ∈ [ J ] (cid:104) v k (cid:48) , w (0) m,j (cid:105) = k } .
For nonlinear MoE with cubic activation function, the following lemma further shows that experts indifferent set M k 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) , thefollowing equations hold for all expert m ∈ M k ,
[Equation]
Lemma 5.2 implies that, at the end of the exploration stage, the expert m ∈ M k can achieve nearlyzero test error on the cluster Ω k but high test error on the other clusters
Ω k (cid:48) , k (cid:48) (cid:54) =
k .
Main Difﬁculty 3: Expert Load Imbalance.
Given the training data set S
= { ( x i , y i ) } ni =1 , theload of expert m at iterate t is deﬁned as
[Equation]
where P ( m i,t = m ) is probability that the input x i being routed to expert m at iteration t .
Eigenet al.
ﬁrst described the load imbalance issues in the training of the MoE layer.
The gatingnetwork may converge to a state where it always produces large Load ( t ) m for the same few experts.
This imbalance in expert load is self-reinforcing, as the favored experts are trained more rapidlyand 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 alsoimpedes the theoretical analysis of the expert training.
Key Technique 3:
Normalized Gradient Descent.
Lemma 5.2 shows that the experts will divergeinto (cid:116) k ∈
[ K ] M k .
Normalized gradient descent can help different experts in the same M k being trained at the same speed regardless the imbalance load caused by the router.
Because the selfreinforcing circle no longer exists, we can prove that the router will treat different experts in thesame M k almost equally and dispatch almost the same amount of data to them (See Section D.2 in
Appendix for detail).
This Load imbalance issue can be further avoided by adding load balancingloss (Eigen et al., 2013;
Shazeer et al., 2017; Fedus et al., 2021), or advanced MoE layer structuresuch as BASE Layers (Lewis et al., 2021; Dua et al., 2021) and Hash Layers (Roller et al., 2021).
Experiments
[Table]
[Table]
6.1 Synthetic-data Experiments
Datasets.
We generate 16 , 000 training examples and 16 , 000 test examples from the data distributiondeﬁned in Deﬁnition 3.1 with cluster number K = 4 , patch number P = 4 and dimension d = 50 .
We randomly shufﬂe the order of the patches of x after we generate data ( x , y ) .
We consider twoparameter settings: 1.
α ∼ Uniform (0 . 5 , 2) , β ∼ Uniform (1 , 2) , γ ∼ Uniform (0 . 5 , 3) and σ p = 1 ; 2.α ∼ Uniform (0 . 5 , 2) , β ∼ Uniform (1 , 2) , γ ∼ Uniform (0 . 5 , 3) and σ p = 2 .
Note that Theorem 4. shows that when α and γ follow the same distribution, neither single linear expert or single nonlinearexpert can give good performance.
Here we consider a more general and difﬁcult setting when α andγ are from different distributions.
Models.
We consider the performances of single linear CNN, single nonlinear CNN, linear MoE, andnonlinear MoE.
The single nonlinear CNN architecture follows (3.1) with cubic activation function,while single linear CNN follows (3.1) with identity activation function.
For both linear and nonlinear
MoEs, we consider a mixture of 8 experts with each expert being a single linear CNN or a singlenonlinear CNN.
Finally, we train single models with gradient descent and train the MoEs with
Algorithm 1.
We run 10 random experiments and report the average accuracy with standard deviation.
Evaluation.
To evaluate how well the router learned the underlying cluster structure of the data, wedeﬁne the entropy of the router’s dispatch as follows.
Denote by n k,m the number of data in cluster Kthat are dispatched to expert m .
The total number of data dispatched to expert m is n m =
(cid:80)
Kk =1 n k,mand the total number of data is n = (cid:80)
Kk =1 (cid:80)
Mm =1 n k,m .
The dispatch entropy is then deﬁned as
[Equation]
When each expert receives the data from at most one cluster, the dispatch entropy will be zero.
And auniform dispatch will result in the maximum dispatch entropy.
As shown in Table 1, 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 ), [Table] the difference between the nonlinear MoE and linear MoE becomes even more signiﬁcant.
We also observe that the ﬁnal dispatch entropy of nonlinear MoE is nearly zero while that of the linear MoE islarge.
In Figure 3, we further demonstrate the change of dispatch entropy during the training process.
The dispatch entropy of nonlinear MoE signiﬁcantly decreases, while that of linear MoE remainslarge.
Such a phenomenon indicates that the nonlinear MoE can successfully learn the underlyingcluster 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 theclustering data structure to the MoE layer in deep neural networks.
Datasets.
We consider the CIFAR-10 dataset (Krizhevsky, 2009) and the 10-class classiﬁcation task.
Furthermore, we create a CIFAR-10-Rotate dataset that has a strong underlying cluster structurethat is independent of its labeling function.
Speciﬁcally, we rotate the images by 30 degrees andmerge the rotated dataset with the original one.
The task is to predict if the image is rotated, whichis a binary classiﬁcation problem.
We deem that some of the classes in CIFAR-10 form underlyingclusters in CIFAR-10-Rotate.
In Appendix F, we explain in detail how we generate CIFAR-10-Rotateand present some speciﬁc examples.
Models.
For the MoE, we consider a mixture of 4 experts with a linear gating network.
For theexpert/single model architectures, we consider a CNN with 2 convolutional layers (architecture detailsare illustrated in Appendix F.)
For a more thorough evaluation, we also consider expert/single modelswith architecture including MobileNetV2 (Sandler et al., 2018) and ResNet
(He et al., 2016).
Thetraining process of MoE also follows Algorithm 1.
The experiment results are shown in Table 2, where we compare single and mixture models of differentarchitectures over CIFAR-10 and CIFAR-10-Rotate datasets.
We observe that the improvement of
MoEs over single models differs largely on the different datasets.
On CIFAR-10, the performanceof MoEs is very close to the single models.
However, on the CIFAR-10-Rotate dataset, we canobserve a signiﬁcant performance improvement from single models to MoEs.
Such results indicatethe advantage of MoE over single models depends on the task and the cluster structure of the data.
Conclusion and Future Work
In this work, we formally study the mechanism of the Mixture of Experts (MoE) layer for deeplearning.
To our knowledge, we provide the ﬁrst theoretical result toward understanding how the MoElayer works in deep learning.
Our empirical evidence reveals that the cluster structure of the data playsan important role in the success of the MoE layer.
Motivated by these empirical observations, westudy a data distribution with cluster structure and show that Mixture-of-Experts provably improvesthe test accuracy of a single expert of two-layer CNNs.
There are several important future directions.
First, our current results are for CNNs.
It is interestingto extend our results to other neural network architectures, such as transformers.
Second, our datadistribution is motivated by the classiﬁcation problem of image data.
We plan to extend our analysisto other types of data (e.g., natural language data).
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section.
• Did you include the license to the code and datasets?
[No] The code and the data are
proprietary.• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes] This paper gives the ﬁrst 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 byadding 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 to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] See our Appendix.
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 a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e 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., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If 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
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
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 to participants and the total amount
spent on participant compensation?
[N/A]