Towards Understanding the Mixture-of-Experts Layer in Deep Learning
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 classiﬁcation 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 classiﬁcation sub-problems that individual experts can conquer.
To our knowledge, this is the ﬁrst theoretical result toward formally understanding 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 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 is an extension of 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 a gating (or routing) function that routes individual inputs to a few experts among all the candidates.
Through the sparse gating function, the router in the MoE layer can 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 a 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., 2021) and are trained with the same optimization conﬁguration.
The router is also initialized to dispatch the data uniformly.
It is unclear why the experts can diverge to different functions that are specialized to 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 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 MoE layer.
We focus on the simplest setting of the mixture of linear classiﬁcation, 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 Deﬁnition 3.1 for more details).
We study training an MoE layer based on the data generated from the “mixture of classiﬁcation” distribution using gradient descent, where each expert is chosen to be a two-layer CNN.
The main contributions of this paper are summarized as follows:
• We ﬁrst prove a 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 signiﬁcantly worse than the mixture of non-linear experts.
Figure 1 provides such a result in a special case of our data distribution with four clusters.
Although a mixture of linear models can 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 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 the cluster-center features and route the input data to the right experts. 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 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 some absolute positive constant C 2 , and x n = Θ( y n ) if C 3 | y n | ≤ | x n | ≤ C 4 | y n | for some absolute constants C 3 , C 4 > 0 .
We also use (cid:101) O ( · ) to hide logarithmic factors of d in O ( · ) .
Additionally, we denote x n = poly( y n ) if x n = O ( y Dn ) for some positive constant D , and x n = polylog( y n ) if x 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 ] to denote 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 based [Figure] on various base expert models such as 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 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 multiple layers of routers and experts.
Shazeer et al.
(2017) simpliﬁed 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., 2017; Dauphin et al., 2017; Vaswani et al., 2017) has been proposed and achieved tremendous successes in a variety of language tasks.
Very recently, Fedus et al.
(2021) improved the performance 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 recognition (Quattoni et al., 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 regressions/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 tensor factorization.
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 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., 2018) and generalization (Allen-Zhu et al., 2019a; Arora et al., 2019a,b; Cao and Gu, 2019) guarantees in the so-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 beyond the 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 this section, we will ﬁrst introduce our data model D , and then explain our neural network model and the details 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 ] are orthogonal 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, we assume 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 drawn from 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.
We point out that the strength of the feature noises γ can 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, such a classiﬁcation problem has an intrinsic clustering structure that may be utilized to achieve better performance.
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 show that the binary classiﬁcation sub-problem over Ω k can be easily solved by an individual expert.
We expect 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 of linear classiﬁers with sophisticated algorithms (Anandkumar et al., 2012; Hsu et 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 to 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 signiﬁcant as we prove that gradient descent from random initialization can ﬁnd a good MoE with non-linear experts efﬁciently.
To make our results even more compelling, we empirically show that MoE with linear experts, despite also being able to represent the concept class, cannot be trained to ﬁ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 is generally set to be linear (Shazeer et al., 2017; Fedus et al., 2021).
Denote by f m ( x ; W ) the output of the m -th expert network with input x and parameter W .
Deﬁne an M -dimensional vector h
( x ; Θ ) = (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.
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, 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 the number 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.
For nonlinear CNN, we consider the cubic activation function σ ( z ) =
z 3 , which is one of the simplest nonlinear activation functions (Vecci et al., 1998).
We also include the experiment for other activation functions such as RELU in Appendix Table 7.
Top-1 Routing Model.
A simple choice of the selection set T x is the whole experts set T x =
[ 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 in 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., |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]
Generate each entry of W
[Table]
0) independently from
[Table]
[Table]
[Table]
do
[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 following empirical loss function: [Equation] [Equation]
[Equation]
) where (cid:96) is the logistic loss deﬁned as (cid:96) ( z )
=
log(1 + exp( − z )) .
We initialize Θ (0) to be zero and initialize each entry of W (0) by i.i.d N (0 , σ 20 ) .
Zero initialization of the gating network is widely used in MoE training.
As discussed in Shazeer et al.
, it can help avoid out-of-memory errors and initialize the network in a state of approximately equal expert load (see (5.1) for the deﬁnition of expert 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.
In particular, 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 strategy for 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 gradient descent update rule for the experts is
[Equation]
where η > 0 is the expert learning rate.
Starting from the initialization Θ (0) , the gradient update rule for the gating network is
[Equation]
where η r > 0 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 can be viewed as a simpler 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 a single expert.
Theorem 4.1 (Single expert performs poorly) .
Suppose D α = D γ in Deﬁnition 3.1, then any function with 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)
j ∈
[ J ] a j (cid:80) p ∈
[ P ] σ ( w (cid:62) j x ( p ) + b j ) can’t perform well on the classiﬁcation problem deﬁned in Deﬁnition 3.1 where σ can be any activation function.
Theorem 4.1 also shows that a 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 deﬁned in Theorem 4.1.
As a comparison, the following theorem gives the learning guarantees for training an MoE layer that follows 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 ) .
Choose experts 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 = Θ
( M 2 ) η .
Then with probability at least 1 − o (1) , Algorithm 1 is able to output ( Θ ( T ) , W ( T ) ) within T = (cid:101)
O ( η − 1 ) iterations such that the non-linear MoE deﬁned in Section 3.2 satisﬁes that
• 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 example x ∈ Ω k will be routed to one of the experts in M k .
Theorem 4.2 shows that a non-linear 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ﬁed accurately by a subset of experts.
On the other hand, each expert will perform well on at least one cluster.
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 and divide the complex problem in Deﬁnition 3.1 into simpler linear classiﬁcation sub-problems that individual experts can conquer.
Finding such a gating network is difﬁcult because this problem is highly non-convex.
In the following, we will introduce the main difﬁculties in analyzing the MoE layer 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 small perturbation of the gating network outputs h ( x ; Θ ) + δ may change the router behavior drastically if the 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 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 , (cid:98) h ∈ R M to be the output of the gating network and { r m } Mm =1 to be the noise independently 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 ) .
Then we 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 smooth transition from time t to t (cid:48) .
It is also worth noting that Θ is zero initialized.
So h ( x ; Θ (0) )
=
and thus 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 uniformly pick 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 features because all the experts look the same: they share the same network architecture and are 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 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 gating network remains nearly unchanged (cid:107) h
( x ; Θ ( t ) )
− h
( x ; Θ (0) ) (cid:107) ∞
=
O ( σ 1 . 5 0 ) .
Because the experts are treated almost equally during exploration stage, we can show that the experts become specialized to some 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 in different 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) , the following 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 nearly zero 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 ) } n i =1 , the load 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 .
Eigen et al.
ﬁrst described the load imbalance issues in the training of the MoE layer.
The gating network 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 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 (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 of the imbalance load caused by the router.
Because the selfreinforcing circle no longer exists, the load imbalance issue will get mitigated. In particular, the router will treat different experts in the same M k 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 the reason
Theorem 4.2 requires early stopping.
This load imbalance issue can be further avoided by adding load balancing loss (Eigen et al., 2013; Shazeer et al., 2017; Fedus et al., 2021), or using advanced
MoE layer structure such as BASE Layers (Lewis et al., 2021; 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 is presented in Appendix E.
The training process can be decomposed into several stages.
The ﬁrst stage is called Exploration stage . During this stage, the experts will diverge into K professional groups (cid:116) Kk =1 M k =
[ M ] .
In particular, we will show that M k is not empty for all k ∈
[ K ] .
Besides, for all m ∈ M k , f m is a good classiﬁer over Ω k .
The second stage is called router learning stage .
During this stage, the router will learn to dispatch x ∈ Ω k to one of the experts in M k .
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 1 .
Setting 1: α ∈ (0 . 5 , 2) , β ∈ (1 , 2) , γ ∈ (0 . 5 , 3) , σ p =
Test accuracy ( % )
Dispatch Entropy
Single (linear) 68 .
NA
Single (nonlinear) 79 .
NA
MoE (linear) 92 .
99 ± 2 .1 .
300 ± 0 .
MoE (nonlinear) 99 . 46 ± 0 .0 .
098 ± 0 .
Setting 2: α ∈ (0 . 5 , 2) , β ∈ (1 , 2) , γ ∈ (0 . 5 , 3) , σ p =
Test accuracy ( % )
Dispatch Entropy
Single (linear) 60 .
NA
Single (nonlinear) 72 .
NA
MoE (linear) 88 .
48 ± 1 .1 .
294 ± 0 .
MoE (nonlinear) 98 .
09 ±
[Table]
[Table]
171 ±
[Table]
[Figure]
     
[Figure]
6.1 Synthetic-data Experiments
Datasets.
We generate 16 , 000 training examples and 16 , 000 test examples from the data distribution deﬁ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 two [Table] parameter 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 nonlinear expert 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, and nonlinear 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 single nonlinear 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, we deﬁne the entropy of the router’s dispatch as follows.
Denote by n k,m the number of data in cluster K that are dispatched to expert m .
The total number of data dispatched to expert m is n m =
(cid:80)
Kk =1 n k,m and 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 a uniform 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 ), 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 is large.
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 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 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 structure that is independent of its labeling function.
Speciﬁcally, 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 is a binary classiﬁcation problem.
We deem that some of the classes in CIFAR-10 form underlying clusters in CIFAR-10-Rotate.
In Appendix A, we explain in detail how we generate CIFAR-10-Rotate and present some speciﬁc examples.
Models.
For the MoE, we consider a mixture of 4 experts with a 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 a more thorough evaluation, we also consider expert/single models with architecture including MobileNetV2 (Sandler et al., 2018) and ResNet
(He et al., 2016).
The training process of MoE also follows Algorithm 1.
The experiment results are shown in Table 2, where we compare single and mixture models of different architectures 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 performance of MoEs is very close to the single models.
However, on the CIFAR-10-Rotate dataset, we can observe a signiﬁcant 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.
In Figure 4, we visualize the latent embedding learned by MoEs (ResNet18) for the 10-class classiﬁcation task in CIFAR-10 as well as the binary classiﬁcation task in CIFAR-10-Rotate.
We visualize the data with the same label y to see if cluster structures exist within each class.
For CIFAR-10, we choose y = 1 ("car"), and plot the latent embedding of the data using t-SNE on the left sub-ﬁgure, which does not show an salient cluster structure. For CIFAR-10-Rotate, we choose y = 1 ("rotated") and visualize the data using t-SNE in the middle sub-ﬁgure. Here, we can observe a clear clustering structure even though the class signal is not provided during training.
We take a step further to investigate what is in each cluster in the right sub-ﬁgure.
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.
[Equation]
Conclusion and Future Work
In this work, we formally study the mechanism of the Mixture of Experts (MoE) layer for deep learning.
To our knowledge, we provide the ﬁrst 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-Experts provably improves the 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 interesting to extend our results to other neural network architectures, such as transformers.
Second, our data distribution is motivated by the classiﬁcation problem of image data.
We plan to 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
[-[Yes] We note in Section 
-]
QG are supported in part by the National Science Foundation CAREER Award 1906169, BIGDATA
IIS-1855099, and the Sloan Research Fellowship.
TZ is supported in part by the GRF 16201320.
The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies.