On the Benefits of Learning to Route in Mixture-of-Experts Models
Authors: Nishanth Dikkala
Nikhil Ghosh
Raghu Meka
Rina Panigrahy
Nikhil Vyas
Xin Wang
Venue: EMNLP 2023 Main (2023)
Keywords: mixture-of-experts
transformer
router
efficiency
conditional compute
sparsely activated models
theory
TL;DR: We study the role of the router in Mixture-of-Expert models and show empirical and theoretical evidence that a learnable router is better than a non-trainable one, (ii) the router can learn to discover latent cluster structure.

Abstract: Mixture-of-Expert (MoE) Transformer models, such as the Switch Transformer, allow us to successfully scale up model sizes while keeping the amount of compute time fixed. Prior work has established the computational efficiency benefits of using these models. A core component of these models is a router that routes input tokens to different experts in a layer. We show theoretical and empirical evidence that the router's ability to route tokens intelligently confers a significant advantage to MoE models. We study synthetic settings where the input data is distributed in clusters and show theoretically and empirically that the router learns to route the inputs according to these clusters. Then we perform experiments on real data using the T5X library, where we observe that a trainable router confers a non-trivial benefit instead of a non-trainable router.

On the Benefits of Learning to Route in Mixture-of-Experts Models
Nishanth Dikkala
Nikhil Ghosh
Raghu Meka
Google Research
UC Berkeley
UCLA
Rina Panigrahy
Nikhil Vyas Harvard University *
Xin Wang Google Research
Google Research
Abstract
Expert
Mixxture-of-Exper (MoE) models, such as the Switch Transformer, allow scale model sizes while keeping the amount of compute time fixed. Prior work established the com- putational benefits of MoE models. We inves- tigate whether they offer benefits other than scaling up. core component of these models is router that routes input tokens to differ- ent experts a layer. We show theoretical and empirical evidence that the router's ability to route intelligently confers significant ad- vantage to MoE models. We study synthetic settings where the input data is distributed in clusters and show theoretically and empirically that the router learns the cluster structure, Then we perform experiments real data using the TSX library, where we observe that a trainable router confers non- trivial benefit instead of non-trainable router.
Input
Router
Expert 2
Output
Expert 3
Figure 1: schematic diagram showing the abstract
Mixture-ff Experts architecture for Neural Networks.
small neural network (a popular choice for the ex- pert architecture is 2 fully connected layers with non-linearity in between). The routing mechanism is described by a trainable weight matrix W and routes each input x an expert (or a few experts) based on the soft-max function applied to W æ. The core principle behind using MoE models is that only one (or a few) of the F2 's is invoked per ex- ample. Thus, when normalizing for compute time, the learned model can draw from the benefits of having many more parameters than dense model (which can be viewed having single expert). In this work, we look closely at core com- ponent of MoE models: the router, which routes different inputs to different experts. By design, MoE models confer computational advantages, and several works (Fedus et al., 2022b; Riquelme et al., 2021; Fedus et al., 2022a) have demonstrated that this speed-up is achieved without any loss in accu- racy. We outline two major reasons for this advan- tage:
1 Introduction
With the advent of large Transformer based models (Vaswani et al., 2017; Kaplan et al., 2020; Brown et al., 2020), model size has increasingly become an issue with neural networks. Larger models con- tinue to improve quality; however, this has started to become very expensive. A promising idea to scale up model size without increasing compute time is that of conditional compute architectures which activate only a part of the network on any example. Conditional compute architectures hold significant potential for saving compute as not all model parameters are used in every example. One of the most successful realizations of this idea is the Switch Transformer (Fedus et al., 2022b) which uses ofxters (MoE) architec- ture. Mixture-of -aree-of-Exeeeee Expert models (Figure in the simplest and most-often used setting are described by a set of experts E,...., Ex and a routing func- tion (aka gating mechanism). Each expert is itself
First, an MoE model using the same compute as dense model has a larger number of effec- tive parameters, and hence can use different sets parameters for different regions in the
.....,,,,
Work done while at Google Reseurch

---

input space. Such an advantage would also hold, for example, a neural network that uses a fixed hash function to assign the input to an expert. Note here that the assignment of inputs to experts need not be learned.
1.1 Our Empirical Results
Several works (see Section 1.3 and references therein) have empirically demonstrated that MoE models exhibit significant advantages over dense models for very large models. We aim to under- stand this phenomenon further, identify the benefits (if any) of learned routing vs. fixed routing, and determine whether routers trained with gradient descent identify latent structures in data, An im- portant design choice for MoE models is the router architecture. We use a Softmax-layer router which is the most popular choice for MoE models (Fedus et al., 2021). The routing mechanism is described by a trainable weight matrix W and routes each input x to an expert (or a few experts) based on the soft-max function applied to W x. There are more recent variants such as Zoph et al. (2022b); Zhou et al. (2022); Puigcerver et al. (2023) which perform better than Switch Transformer but all of these works use the same Softmax-layer router backbone. We do note that other routing mecha- nisms which do not use a Softmax-layer have been studied (See Section 1.3) but we do not study those in our work. For interpretability and to identify the more fine points we seek, we introduce the study of the performance of MoE models under various syn- thetic problems. Most of these problems quite well-studied in clustering and unsupervised learn- ing and are independently interesting. We describe these settings first:
The second and more interesting reason is that the router is learning an intelligent rout- ing function that cleverly partitions the input space amongst the experts. An example would be router that routes to experts based on the clusters in the input i.e., all examples in cluster routed to single expert.
We seek to theoretically and empirically under- stand what role the router plays in delivering the advantages that MoE models enjoy over dense mod- els. Specificallyy, we study the following questions about the routing mechanisms in MoE models:
Can the router learn intelligent and inter- pretable routing functions? For instance, can the router identify any latent clusters in the data?
How important is the learnability of the rout- ing mechanism as opposed to routing based on fixed data-independent router (for instance, using a random locally-sensiivee hash func- tion)?
Can we prove that MoE models achieve good accuracy in well-studied clusterable settings such as mixtures of Gaussians?
Mixtures of Gaussians: One of the most stud- ied models for clusterable data is mixtures of Gaussians. We study the performance of MoE models when the input examples come from a mixture of Gaussians and the labels are spe- cific to each cluster (either constant or linear functions).
Zoph et al. (2022a) and others have tried to iden- tify patterns in what the router is learning and have uncovered some insights. However, they are hard to interpret due to the complexity of real data. Here, we study MoE models on various synthetic tasks, which sheds clearer light on the role of the router. To begin with, look at tasks where the input has a clusterable structure, such as a mixture distribu- tion over separated Gaussians, and study whether the router can learn to route inputs according to the cluster identity. We show theoretical and empiri- cal results in the affirmative for the router learning the underlying cluster structure of the data. We also perform experiments comparing MoE models with trainable vs. non-trainable router the Common Crawl (mC4) dataset and observe that a learned routing function performs significantly better than a routing function decided in a data- independent manner.
2. Mixtures of Subspaces: We move from Mix- tures of Gaussians to Mixtures of subspaces, this is a setup in which the input clusters are not necessarily separable in Euclidean dis- tance.
3. Tokenized mixtures of Gaussians: The Transformer architecture which MoE models try to improve on is sequence prediction model. Consequently, an important design feature of many MoE architectures! is that one routes each token in the sequence indepen- dently of the other tokens.
To capture this sequence setting, we consider a setup where we have a sequence of tokens

---

with each token being sampled from a mixture of Gaussians.
curves (Figures 2(a), 3(a) 4 for Mixtures of Gaus- sians and Figure for Mixtures of Subspaces) as well as pointwise comparisons for the CIFAR-100 based task and the T5X task.
4. Dictionary Learning: This is a more compli- cated set up where each input is created using combined information from or more clus- ters. This helps model practical data more accurately as there is some amount of struc- tured overlap between concepts here. Here, although we perform better than dense models in certain settings, we do not observe intelli- gent routing as frozen router performs just as well as Iearnable router.
1.2 Our Theoretical Results
We seek to show that MoE models provably achieve good performance (and intelligent routing) in classically well-studied instances of clusterable distributions. To our knowledge, such results are not known'. As first step. in this direction, we study simple models with a single router that routes inputs to experts, and each expert either applies a linear function on the input routed to it or outputs fixed vector. We also analyze convergence un- der gradient descent with population gradients to minimize technical machinery and highlight the important points. The arguments would extend naturally to sample gradients long as we have enough samples by standard concentration inequal- ities.
5. CIFAR- R-100 class clusters: In this setting we consider clusters defined by CIFAR-100 classes. The output is a linear function per cluster.
Finally, we perform some experiments real data. We pre-train T5 transformer architectures (Raffel al., 2020) and measure the pre-training accuracy. We compare three different models: (i) a dense 12-layer T5 model, (ii) a 12-layer MoE model with three sparse layers and top-2 routing and trainable router, (iii) a 12-layer MoE model with three sparse layers and top-2 routing and a fixed non-trainable router.
In our first theoretical result, Theorem 1, we show that the router can identify relevant directions and remove spurious directions. Concretely, we show that if we have d-dimensional inputs that are clusterable in subspace of dimension d' < d (un- known to us) and has independent spherical Gaus- sian noise along the other router aligns itself with the hidden d' -ddimensiioal subspace. This tells us that the router can perform intelligent projection operations. We prove this result more generally for any neural network archi- tectures with regularization term for the bottom- layer weights, which could be of interest indepen- dently. The proof exploits rotational symmetry and deeoppecccoosaaility of the Gaussian distribution. Next, we study the setting of a mixture of & Gaus- sians. We assume that the k Gaussians are spherical and live in well-separated high-dimensional space where the notion of well-separability is analogous to the classical notion in Dasgupta and Schulman (2007). The second assumption we make in the initialization of the weights of the router matrix. Rather than randomly initializing them in a data- independent manner, we set each row of the router as a randomly sampled training example. This is similar in spirit to Dasgupta and Schulman (2007). Under these assumptions, we show that each
In the above settings, have two main results corresponding to our questions on understanding the routing behavior. For details on the results and the experimental methodology in Section 3.
d' directions, the
Routers learn to route intelligently based on input structure: In our empirical experiments, we find that routers of MoE models can adapt to the input structure. Specifically.
In Mixtures of Gaussians experiments, with constant or linear outputs per Gaussian clus- ter, we find that routers learn to map clusters to experts. That is, nearly all examples from cluster are routed to the same expert (Fig- ures 2(c), 3())).
When trained on Mixtures of Gaussians with spurious dimensions, i.e., dimensions along which the marginal of all Gaussian clusters is the same, we find that the router learns to ignore these directions (See the last paragraph of Section 1.
Learning to route offers an advantage MoE models: In all of our settings described above, we universally find that learned routers significantly outperform MoE models with random routers. We study these effects by plotting the data scaling
"A described in Section 1.3, Chen et al. (2022) recently showed an elegant provable guarantees for MoE models. How- ever, their result applies to a data distribution they define with a specific purpose and perhaps does not correspond to an independently interesting clusterable instance

---

expert specializes specific cluster center. By using ( log k) experts, we can ensure that every cluster has at least one expert specializing in that cluster. By choosing the router learning rate appro- priately, we can also argue that while training the router, weights do not change too much, while the experts specialize in their clusters. Our result is given in Theorem 2.
parameterized model, the benefit of MoE models having more parameters goes away. Still, one could hope that MoE models outperform dense models, particularly when inputs are clustered, as they might have a better implicit bias for such in- puts. We don't observe this in our experiments in Figure 2(a) where the largest dense model out- performs the MoE model even though both are over-parameterized.
We also show a result similar to the above for mixture of two Gaussians with no assumption on the initialization in Appendix F.
2 Preliminaries
We use [n] to denote the list of integers from [1. .n]. We denote the probability distri- bution of d-dimensional Gaussian with mean vector μ and covariance Σ by ((, Σ). In our work, we will deal with spherical Gaus- sians where the covariance matrix Σ = ola- Given two dddimensional spherical Gaussians N (P1, a1la), σ2la). we say they are c- separated if || 101--2210 la 2 cmax(o1. a2 v (from Dasgupta and Schulman (2007)).
1.3 Prior work
Mixture-of-Exper models (Jacobs et al., 1991; Shazeer et al., 2017: Fedus et al.. 2022b; Riquelme et al., 2021) have been used in various forms in machine learning over the last three decades. We refer to the excellent survey of Fedus al. (2022a) for their history. In our work, we focus specifically on their use in neural networks. We describe the most relevant works to ours below.
Roller et al. (2021) studied routing based on Hashing the input tokens. This is an example of a routing scheme where the assignment of inputs to experts is not learned but is not necessarily the strongest such routing. Clark et al. (2022) showed that such a routing scheme does not scale as well as architectures with trainable routers. We consider a more general setup where the routing is not fixed at initialization, but the ability to learn to route is limited. We show that even in this setting, learning to route is essential.
The neural networks we use in our synthetic experiments have the following form: Every input is first processed by the router and assigned to a single expert, and the output of this specific expert is the neural network's output. We describe MoL models by specifying the router and the experts. For example, a linear router with constant experts refers to trained model which
On input x, computes Wz where W denotes the router weights. The input is assigned the expert j where j = arg max; x);; i.e. jth entry of W x is the largest entry.
Recent works of Chen et al. (2022) and Baykal al. (2022) study the theoretical benefits of MoE models over dense models. Similar to us, Chen et al. (2022) focus on problems where the input structure has clusters. We differ from this work in our focus on more elaborate synthetic and real- world examples. We give a more in depth compari- son of our results with those of Chen et al. (2022) along with discussion of additional related work in Appendix B.
2.
The expert j outputs learned vector wj and hence wy is the output of the neural network.
We note that this description applies to trained models. While training, we do a weighted com- bination of experts instead of routing to the top expert. All our results are for trained network for which we do top-1 routing. See Appendix C for more details.
Remark 1. In trying to compare the performance of MoE models with dense models, it is impor- tant to understand the difference between over- parameterized and under-parameterized settings². In mnder-parameterizer settings, irrespective of the router learning intelligent routing, MoE mod- els have a clear advantage over dense models as they simply have more parameters. In an over-
2.1 Metrics
We will use various metrics in our synthetic experiments to understand the behaviour of MoE models. Here we list these metrics:
'The output the expert in this case does not depend on the input as the expert is "constant" expert. By "constant" we mean independent of input, not that it is fixed throughout training, i.c. wy is still learned. we had a linear expert with weight wy, the output would be uz
WWo define these settings based on whether train loss can reach near 0 or not respectively

---

Test Loss: The loss of the MoË model on the underlying distribution. If test loss tends we are able to argue that gradient descent can find a network which can represent the true labelling function.
3.1 Mixture-of-Gausiann inputs and constant vector output per Gaussian
This is the simplest setting we consider. The input is a mixture of well-separated Gaussians, and the output is a random constant vector per Gaussian. Note that an MoE model can represent this label function with a linear router and "constant" experts, i.e., each expert outputs fixed albeit learned vector independent of the input.
Average Sparsity per Cluster: MoE routers typically output a probability distribution over the experts indicating which experts' to rely on for the output. Our input data will usually be uniform mixture of k underlying (clus- tered) distributions C1, C2, ........ Ck. We want to understand how well the router is able to "discover" these clusters. To do this, we de- fine a notion of 'sparsity' which is a contin- uous quantity which captures how sparse the probability vector output by the router is. We consider the following quantity:
Can gradient descent discover underlying clusters? In our first experiment (Figure 2), we set the dimension of the input to be 24 and take the input distribution to be a uniform mixture of 64 well-separated spherical Gaussians. The output is a randomly sampled 10-dimensional vector per Gaussian cluster. We use an MoE model with lin- ear router and "constant" experts. See Appendix C for more details. In Figure 2(a), we see that the number of samples increases, the loss tends 0. A similar effect is seen in Figure 2(b) where we run the same experiment but vary the number of experts (the number of samples set very large value). When the number of experts is less than the number of clusters (64) we see that the performance cannot achieve 0 test loss. This is expected as the architecture is unable to represent the actual function. We see that when the number of experts equals the number of clusters (64), the loss is near 0. This is surprising in comparison our theoretical result (Theorem 2), where we need an additional log(#clusters) factor to prove con- vergence to a good solution. A similar logarithmic factor is incurred in the analyses of Dasgupta and Schulman (2007), Chen et al. (2022).
Exck H (B,,E,
where H is the entropy function and R(x) is the vector of assigned probabilities to each expert by the router R on input x. Intuitively, we take the average (over a cluster) of the as- signment to experts and calculate the entropy and then we also average over all the clusters. A smaller average sparsity metric across in- puts from cluster indicates that all the inputs from a cluster are being routed to a small set of experts. similar metric used in Chen et al. (2022) as well.
Average Sparsity per Cluster (Shuffled Router): Suppose R is a linear router i.e. R(x) = (ettmxx(x)) where W E IR fexperts input dimension Let W' be a random
Next, we explore the behavior of the router. In Figure 2(c) we see how sparsity (average sparsity per cluster, see Section 2.1) changes with increase in training data size. We find that sparsity tends 1 with increasing dataset size. Hence each cluster gets mapped to a single expert. Comparison with the shuffled router, which has much higher sparsity, shows this is not due to spurious factors such as large norm router weights.
shuffling of the columns of W. Let R' de- note the router obtained by using W" in- stead of W. This metric is the quantity FieR (Bzec, R'(x)). This is the exact same quantity as the previous metric except that we use the shuffled router R' instead of R. We use this metric as a baseline for the previ- ous metric. This serves as a good baseline as various statistics such as average Lp norm of entries are shared between W and W"
Benefits of learning to route Our next question is whether trainable router helps the performance of MoE models. To do this, we can compare the performance of an MoB model with frozen router. But with frozen router, an MoE model with a linear router and "constant" experts will not be able to reach a train loss of 0. To rectify this, we move to an MoE model with depth-3 experts and compare
3 Experimental Results
We present our experiments on synthetic (Sec- tions 3-3.3.6) and real datasets (Section 3.7) in this section. Full details of our experimental setup are in Appendix C.

---

MCE
MDE
6.33
1005
Sparsity
Stuflled
96
96
192
Number of samples cluster
Numbar experts
Namber of samples per cluster
(a)
(c)
Figure 2: Mixture-of- Sinture-o-GGausiins inputs and constant vector output per Gaussian. We use neural network with a linear router, and experts that output constant (learned) vector. In Figure 2(a), we take the number of experts equal the number of clusters and that the number of samples increases, the loss tends we set the number of samples to be very large and see that soon the number of experts matches the number of clusters (64), the loss drops to near O. Finally, in Figure 2(c), we examine the learned router and find that the router learns map each cluster unique expert while shuffled router does not.
In Figure 2(b),
Average sparsity per cluster
Loss
.1
Router
Shuffled Router
.05
MOE
MOE (frozen router)
.025
(small width) (large width)
6
12
24
3
6
12
24
Number samples per cluster
Number samples per cluster
Figure 3: Bixture-of-Gaussiinn inputs and constant vector output per Gaussian We use neural network with a single layer router i.c. a linear router and depth-3 MLPs as experts. In Figure 3(a) we take number of experts to equal the number of Gaussian cluster and while the performance of the model with frozen router is much worse. In Figure 3(b) we examine the learned router and find that the router learns map Gaussian clusters to unique experts while shuffled router not able do so.
that the the number of samples increases the loss tends 0
with the same MoE model with depth-3 experts but with frozen router. Figure 3(a) shows that the MoE model with learned router outperforms the MoE model with frozen router. Note that the training loss of all models is a gap in generalization than in training error. In Figure 3(b), we again see similar behavior to what we saw in Figure 2(c). See Appendix C.1 for a related experiment.
MLP models to have one more depth than the MoE models. We find that the MoE model outperforms the smaller MLP but slightly underperforms the larger MLP model.
zero; thus, this
Identifying non-spurious dimensions: We peat the same experiment with 64 clusters, but this time we embed these 240 dimensions, out which 240 - 24 = 216 are spurious dimensions, i.e., they are i.i.d. Gaussians for all clusters. We again find that training with this setup leads to a small test loss (0.01, being trivial, and being perfect). This indicates that gradient descent con- verges near-perfect solution. Let w; denote the weights of the router corresponding to expert and wi denote the projection of wi on the 24 non-spurious dimensions. We find that the quan-
Comparison to Dense models: Figure 3(a) also compares MoE. models with two dense MLP networks. One matches the number of parameters of single expert (so compute time is similar to MoE), and another matches the total number of parameters of the entire MoE model (much more expensive computationally). As MoE models have one routing layer, for a fair comparison, we allow

---

Average sparsity per cluster
0.4
12
0.2
Loss
0.1
0.05
MOE
Router
0.025
MOE (frozen router?
Shuffled Router
Number of samples per cluster per dimension
Number of samples per cluster per dimension
(
(b)
Figure 4: Mixture-of -aausssann inputs and linear function per Gaussian We use neural network with single layer router, i.c. a linear router, and depth-3 MLPs as experts. Figure 4(a), we see that as the number of samples increases, the loss tends to 0, and the model's performance with frozen router is much worse. In Figure 4(b), we examine the learned router and find that the router learns to map Gaussian clusters to a small set of experts while the shuffled router does not.
Ed tity E,1,, which intuitively denotes how much the router focuses on the non-spurious dimensions, goes from .1 24/240 at initialization to .84 after training without weight decay. Adding weight de- cay increases it to .9 without loss in performance. Hence the router learns focus mostly on the 24 non-spurious dimensions out of the 240 dimensions. This connects with our Theorem 1, which predicts this behavior.
router.
3.3 Non-linear clusters: Mixture-of-subsaace In this setting, the input is a mixture of & ran- domly chosen rank-t subspaces with added noise to samples. We defer more details on the setting to Appendix C. In Figure 6, we compare the MoE model with a learned router to one with frozen router. We find that learning to route is important for the performance for large range of dataset sizes.
3.2 Mixture-of-Gaussiee inputs and linear function per cluster
3.4 CIFAR-100 class clusters:
We next generalize the previous setting. The in- put distribution remains the same but has separate linear function for each cluster. For each Gaussian N (H, Ia), we now sample a random vector w linear function w, and on input x from this Gaussian, output is wn (x - (.) 4. We train an MoE model with a linear router and depth-3 MLP experts on this task. In Figure 4(a), we observe that the loss goes down as we increase the number of samples per cluster that the network can learn this func- tion. In Figure 4(b), analyze what the router learns. To do this, plot the average sparsity per cluster as we increase the number of samples. We see that the sparsity decreases as the number of samples increases. As a baseline, we plot the same quantity after shuffling the router weights and observe no improvement in sparsity in this setting. This demonstrates that the router routes according the latent clusters much better than the shuffled
We consider the following regression task. We construct a dataset where each input is a pair the form (CIFAR-100 100 image, random vector v). The ground truth target output is defined as ui!' where i is the class of the image, and w; is class- specific random vector. The model
of the following form: have a CNN router and depth- MLP experts. We set the number of experts to 200 (Theorem suggests that for k clusters we need O log k) experts). The CNN is only given the image as input, and experts are given the vector and the penultimate layer of the CNN. The loss when we allow the CNN router to learn is .5, while is 80 when we freeze it.
3.5 Sequence Data:
One difference between real transformer models and our experiments above is that we were not processing sequence data where the input is split into tokens. We also perform experiments which
WW center the linear function as clse, the values hecome close to constant per cluster again.
*This helps the model in the case we freeze the router by making sure that experts also get information about the imuge.

---

demonstrate that our insights can be extended to when the input is given as a sequence of tokens. See Appendix C.2 for more details in this setting.
model with frozen router is significantly outper- formed by an MoE model with trainable router. The performance of these three models is com- pared in Figure 5. We note that our artificial con- straint on the model dimension makes it hard to directly draw conclusions for the setting of real- world Transformer models. However, it is oth- erwise challenging to disentangle implicit router learning by the pre-router layers in large model from the actual role of allowing a router to train. See Appendix C.4 for experiments at additional scales with T5 models.
3.6 A Dictionary Learning Inspired Setting
We now explored whether MoE can route in a more complex setting. Taking inspiration from dic- tionary learning, we create an experiment where there are & atoms {4)] C IRd. Each of these atoms corresponds to a cluster. An input x is gen- erated by selecting a few of these atoms (2 or 3 in our experiments), sampling Gaussian random variables centered at each of the selected atoms, and finally adding up the sampled variables. The target output y scalar obtained by multiplying bi with an atom-specific vector wi; for the selected atoms and adding up the results. Depending on the problem parameters, this problem can be easy or hard. Although we observed the MoE model achieve a better test loss than dense models (0. 19 0.03) in some settings, we did not observe spar- sity in routing similar to what we saw in previous sections. Moreover, the test loss of 0.03 was also achieved when the router was frozen at initializa- tion. In addition, the test loss was not very close to 0 indicating that this sctting is more challenging for MoE routers to learn the latent structure. See Appendix C.3 for more details.
Pretnaining T5 mC4
Accuraccyn (44)
MaE
Figure 5: Pretraining Accuracy on mC4 Each run is 12 layer model with model dimension reduced to 32 (from 768). This has been done to better study the role of learning to route. See Section 3.7.
3.7 Language Modeling mC4
One of our aims the study the benefit of learn- ing to route in MoE models on real data. We study language modeling objective on the Com- mon Crawl (mC4) dataset with T5 transformers. A first approach one could consider would be to com- pare neural networks where the router is learned versus networks where the router is frozen. Unfor- tunately, this cannot be directly achieved in stan- dard Transformer-basee MoE models as the router is also used between layers. Hence even if we were to freeze the router itself, the earlier layers could learn to do the routing implicitly.
4 Theoretical Results
We state and describe our theoretical results in this section deferring full proofs to the Appendix. Our first result proves that MoE routers can identify rele- vant directions and remove spurious directions. We prove this result for any MoE architecture which passes the input through a softmax router first with regularization term for the router weights.
4.1 Router Learns to Ignore Spurious Directions
We consider setting where are given inputs from mixture of i spherical Gaussians in Rd The Gaussians are well separated in a dimension of rank 2, but there is no clustering when projected to other a dimensions. The ground truth is fixed vector Ve corresponding to the Gaussian x is drawn from.
To get around this, we a) reduce the model di- mension to 32 and b) freeze the router. We do (a) to reduce previous layers' capacity to learn to route. Note that there could still be some learning to route happening implicitly in the earlier layers. Hence, the performance of this model will be an upper bound on the performance of a model which does not learn to route. With these changes, we first find that an MoE model with a frozen router significantly outperforms a dense model with sim- ilar FLOPs. More relevant to our study, the MoE
Theorem 1. (Informal) Given target function v(x) which depends on only projection of x E B to d' < d dimensional subspace V and a prediction function f which is of the form gu (Wa x + b) for some function (parameterized by U), any local

---

minimum of the regularized loss Eyy f) must have each row of W lie within V.
for Mixture of 2 Gaussians which do not need initialize the router intelligently in Appendix F. In all our analyses we assume balanced mixture of Gaussians as that simplifies the proof considerably. However, we believe our results of Theorem and 4 can be generalized to hold even when the mixture weights are not uniform.
This shows us that routers indeed have the ca- pacity to learn to filter out signal directions from noise directions even when the final function being computed is a non-lincar function of the input. Next, we show that routers can learn latent clus- ters when inputs come from mixture of Gaussians. We choose mixture of Gaussians as it is known to be a highly powerful modeling framework in prac- tice (Khansari-Zaade jan et al., 2013) while also theoretically capable of representing urbitrary distributions using Gaussians as basis functions. They are also well-studied and are known to be amenable for theoretical analysis.
Conclusion
We studied the role of learning a router in MoE models and identified a clear advantange in many synthetic and real settings for learning to route. In addition, we also saw that the MoE router is ca- pable of learning latent cluster structure in some scenarios and less capable in more complex scenar- ios like the dictionary learning setting. We believe that future designs of MoE routers could be guided by our experiments to bolster their capabilities for certain types of routing tasks. In addition, the fact that MoE routers trained with gradient descent are able to learn meaningful and interpretable clusters is a promising sign for the field of Conditional Compute models as it shows that not only do we save compute but we can also hope learn sparse yet interpretable solution to a complex problem.
and Billard, 2011; I; Varadara-
4.2 Router Learning Latent Clusters for a Mixture of k well-separated Gaussians We consider setting with a mixture of k Gaus- sians. Let {ge- -V oola)) be the k spheri- cal component Gaussians of mixture distribution defined by Σe= Wofe. We assume that all the mix- ture weights we are uniformly = /k: for simplicity.
MoE architecture: The neural network is of the form: f(x)= = E P:uty where P: = softmax (wf x) and uy wyys 's are parameters. The wy's represent the router and the u, 's represent the experts.
5.1 Open Questions
An important open direction from our work is designing new neural architectures given the in- sights we uncover about the role of routers in MoF. models. We have seen that the standard MoE router struggles in certain settings. This necessitates the design of more complex and capable router mod- ules. Another direction which we believe can help improve our understanding of MoE models further is to extend the theoretical results to more prac- tical settings. For instance, the following would be an interesting question to analyze theoretically. In practice, we don't 't have a clean separation of concepts or tasks as we did in our theory. Rather, concepts might be organized in hierarchy or more generally graph with varying amounts of overlap between nodes. Given such a graph, can we design router and prove guarantees on it being able to uncover the latent graph structure?
First, we note strong similarities this setting has with the k-means problem in Lemma 1. Then show that if we initialize the router weights in a specific way using the training data, each expert specializes to a specific cluster center. By using @kk log k) experts, we can ensure that every cluster has at least one expert specializing in that cluster. Theorem 2. (Informal) Given n = O(kk log k) samples drawn from a mixture of k spherical Gaus- sians in d-dimensions which are e-separated for some constant as described above, consider an instantiation of the above MoE architecture with (k log k) experts. If we initialize the router weights wi to randomly drawn examples from the train set, the router will learn to route examples according to the cluster they belong to.
In Theorem 2, we had to initialize our router weights in a smart manner using the inputs to en- sure that they identify the latent clusters. Such utilization of stratified subsampling of dataset is not uncommon in Machine Learning. is used in K-means algorithms, for instance. It is also a key component of Curriculum Learning (Bengio et al., 2009). We also have stronger theoretical results
Limitations
We outline the key limitations of our work here. While we see empirical and theoretical evidence of an MoE router being able to learn intelligent and interpretable routing in many settings, they are brittle in slightly more complex settings such as the dictionary learning setting. We acknowledge that

---

today's MoE routers perhaps do not learn latent structures always. But we believe there scope for them to develop the ability to do so in the future. We performed a number of our experiments on carefully crafted synthetic data to ablate clearly the effect of learning to routc. However a downside of this is that we lose some amount of generalization of insights to real data settings where the picture might be different. We would like to mention, how- ever, that there is significant evidence in the deep learning field for analyzes such as ours to yield use- ful insights which generalize to real data and offer actionable ways to improve models. Lastly, some of our experiments were done on reasonably large T5 models which require significant compute.
Nan Du Yanping Huang, Andrew M Dai, Simon Tong. Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. 2022. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference on Machine Learning, pages 5547-5569. PMLR.
Angela Fan, Shruti Bhosale, Holger Schwenk, Zhiyi Ma, Ahmed El-Kishky, Siddharth Goyal, Mandeep Baines, Onur Celebi, Guillaume Wenzek, Vishrav Chaudhary, et al. 2021. Beyond english-centric mul- tilingual machine translation. J. Mach. Learn. Res.. 22(107):1 48.
William Fedus, Jeff Dean, and Barret Zoph. 2022a. A review of sparse expert models in deep learning.
William Fedus, Barret Zoph, and Noam Shazeer. 2021. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.
Ethics Statement
William Fedus, Barret Zoph, and Noam Shazeer: 2022b. Switch transformers: Scaling trillion parame- models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232- 5270.
Our work helps understand one of the key aspects of Mixture of Expert style deep learning models better. There are direct negative societal or dis- criminatory impacts of our work. Rather, we hope that the broader impact our work has positive ef- fects for society at large.
Hussein Hazimeh, Zhe Zhao, Aakanksha Chowdh- ery, Maheswaran Sathiamoorthy, Yihua Chen, Rahul Mazumder, Lichan Hong. and Ed Chi. 2021. Dselect- k: Differentiable selection in the mixture of experts with applications to multi-task learning. Advances in Neural Information Processing Systems, 34:29335- 29347.
References
Cenk Baykal, Nishanth Dikkala, Rina Panigrahy, Cyrus Rashtchian, and Xin Wang. 2022. A theoretical view on sparsely activated networks. arXiv preprint arXiv:2208.0444 04461.
Hidden. 2022. Hidden for double blind.
Robert A Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey E Hinton. 1991. Adaptive mixtures of local experts. Neural computation, 3(1):79-87.
Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston, 2009. Curriculum learning. In Proceedings of the 26th annual international confer- ence on machine learning, pages 41-48.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scout Gray, Alec Radford. Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836)
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et 2020. Language models are few-shot learners. Advances in neural information processing systems, 13:1877-1901.
S Mohammad Khansari- Zadeh and Aude Billard. 2011. Learning stable nonlinear dynamical systems with gaussian mixture models. IEEE Transactions on Robotics, 27(5):9439957.
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. 2022. Towards understanding mixture of experts in deep learning.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2020. Gshard: Scaling giant models with conditional com- putation and automatic sharding. arXiv preprint arXiv:2006. 16668.
Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini. Jordan Hoffmann, Bog- dan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, al. 2022. Unified scaling laws for routed language models. In International Conference on Machine Learning, pages 4057-4086. PMLR.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. 2021. Base layers: Simplifying training of large, sparse models. In In- ternational Conference on Machine Learning, pages 62744 PMLR.
Sanjoy Dasgupta and Leonard J Schulman. 2007. probabilistic analysis of em for mixtures of separated, spherical gaussians. Journal of Machine Learning Research, 8:203-226.
Ankur Moitra. 2018 Algorithmic aspects of machine learning. Cambridge University Press.

---

Rina Panigrahy, Xin Wang. and Manzil Zaheer. 2021. Sketch based memory for neural networks. In In- ternational Conference Artificial Intelligence and Statistics, pages 3169-3177. PMLR.
Barret Zoph, Irwan Bello, ameer Kumar, Nan Du, Yan- ping Huang. Jeff Dean, Noam Shazeer, and William Fedus. 2022a. Designing effective sparse expert mod- els. arXiv preprint arXiv:2202. 0900009906.
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. 2023. From sparse to soft mixtures of experts. arXiv preprini
Barret Zoph. Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. 2022b. St-moe: Designing stable and transferable sparse expert models, arXiv preprint arXiv:220..88006.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter Liu, et al. 2020. Exploring the limits of transfer learning with unified text-to-text trans- former. J. Mach. Learn. Res., 21(140):1-67.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Su- sano Pinto, Daniel Keysers, and Neil Houlsby. 2021. Scaling vision with sparse mixture of experts. Ad- vances in Neural Information Processing Systems, 8598583-8595.
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. 2021. Hash layers for large sparse models. In Advances Neural Information Pro- cessing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 14, 2021, virtual, pages 17555-17566
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff' Dean. 2017. Outrageously large neural networks: The sparsely-gates mixture-of-expers layer. arXiv preprinl 1Xxx:1701.065::
Tay, Mostafa Dehghani. Samira Abnar, Hyung Won Chung. William Fedus, Jinfeng Rao, Sharan Narang, Vinh Q Tran, Dani Yogatama, and Donald Metzler. 2022. Scaling laws vs model architectures: How does inductive bias influence scaling? arXiv preprint arXiv:2207.100 1055I.
Sriram Varadarajan, Paul Miller, and Huiyu Zhou. 2013. Spatial mixture gaussians for dynamic background modelling. In 2013 10th IEEE international confer- ence advanced video and signal based surveil- lance, pages 63-68, IEEE.
Ashish Vaswami. Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural information processinnnpooessinng systems, 30.
Lemeng Wu, Mengehen Liu, Yinpeng Chen, Dongdong Chen, Xiyang Dai, and Lu Yuan. 2022. Residual mixture experts. arXiv preprint arXiv:2204.09636
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, al. 2022. Mixture-of-experts with expert choice routing. Advances Neural Informa- tion Processing Systems, 35:7103-7114,

---

A Additional Background
B Additional Related Work
We first discuss comparison of our work with (Chen et al., 2022) in more detail. The focus in Chen al. (2022) is on designing a data distribution where MoE models have provable advantage over a non-expert model. In contrast, we are interested in the benefit of MoE models in well-studied clustering instances (such as mixtures of Gaussians) and their variants. such settings, we prove the first theoretical results showing that MoE models can get very good accuracy. In addition, Chen et al. (2022) focus on convolutional layers, whereas we focus on fully connected layers, which is closer to the standard Transformer architecture. Further, we seek to compare MoE models with dense models and aim to understand the advantage of learned routing vs. non-learned routing. Regarding real-world experiments, our focus is on language data and understanding the benefit of learned routing when used within a transformer architecture (T5X). In contrast, the experiments in Chen et al. (2022) language data study the benefits of routing when using the representations generated by large-lnnguage model (such as BERT) black-box manner.
We discuss some additional related works now. Clark et al. (2022); Tay et al. (2022) established that the benefits of MoE models scale to larger sizes (although the scaling law seems to slow down at extremely large scales). There have been lot of efforts to study the best way route tokens which includes ideas from load balancing (Zoph et al., 2022a; Lewis et al., 2021), reinforcement learning (Clark al., 2022), and smooth tok-k expert selection (Hazimeh et al., 2021). Other choices for the routing function include non-learable ones such as Locality Sensitive Hashing (Panigrahy et al., 2021), which generally maps similar inputs the same expert, Hash Layers (Roller al., 2021) which uses ooken-based random hashing, and language-specific deterministic routing (Fan et al., 2021). Another approach Residual Mixture of Experts (Wu et al., 2022) separates the expert weights into input-independent and input-dependent components. MoE ideas have made their impact on prominent large scale systems such GLAM (Du et al., 2022), GShard (Lepikhin et al., 2020).
Dictionary Learning and Mixture of Gaussians learning problems are well-studied by long line of works. For references on these problems we refer the reader to the textbook Moitra (2018) and the references within.
Additional Experimental Details
C.1
Mixture of Gaussians/ubspaaee experiments
Details on the Mixture of Subspaces Setting In this setting, the input is mixture of k randomly chosen rank-t subspaces with added noise samples. For each subspace, let be spherical Gaussian restricted to the subspace, and is spherical Gaussian in all dimensions (with a smaller norm than r). The examples are generated as x + 6, and the label wT where is random vector chosen for the subspace. Note that this is no longer possible with a linear router we want to route based on the subspace, as the subspaces are not linearly separable. Due to this, we shift to using a depth-3 MLP as the router. There are k experts, and each expert is also a depth-3 MLP. We set k = 64 and t = 8. The results of our experiments are shown in Figure 6.
Is the router learning bipartite matching between clusters and experts? In Figure 3(b) that each cluster nearly gets routed single expert. As each expert is overparameterized it is possible that the router learns to focus on only few experts for all clusters. To verify that this is not the case we compute how uniformly distributed the norms of weights feeding into different experts are (note that the router is linear). Let M, denote the weights feeding into expert i. We compute
[ M 12] [[ M, 1
This quantity would be 1 if all ||M|]| were equal and it would be #eaperts if only one ||Ma]] was nonzero. the right most point for the MOE with learned router in Figure 3(a) this quantity is 1.025 while #eaperts is 64. Hence it is not the case that the model learns to focus on only a few experts.

---

Other details for Figures
In Figure 2 we trained networks with Adam and a learning rate of 3.2c-4 for the experts and 3.2c-3 for the router.
In Figure 3 and we trained networks with SGD (with .9 momentum) and a learning rate le-2 for the experts and le-1 for the router. We choose higher learning rate for the router as otherwise the overparameterized experts may learn to overfit before the expert learns a good routing. The hidden layer width is 1 21 = 96.
MOE MOE (frazen router)
0.8
0.6
0.4
32
64
128
256
Number samples per cluster per dimension
Figure 6: Mixturee-f-subspaes inputs and linear function per subspace. We that MoE model with learned router outperforms an MoE model with a frozen router across various dataset sizes.
In Figure we trained networks with Adam and a learning rate of 1e-3 for the experts and 3.2e-3 for the router. We choose a higher learning rate for the router as otherwise the overparameterized experts may learn to overfit before the expert learns good routing. The hidden layer width is 4 24 = 96.
C.2 More Details on Tokenized Variant Experiments
We first describe our setup for the experiments on the tokenized variant of mixture of Gaussians in more detail. We have an equally weighted mixture distribution over 32 Gaussians whose centers randomly chosen from a 32 dimensional space. The inputs generate consist of sequence of tokens where each token is a randomly drawn sample from the mixture. We vary the length of the sequence s. The output is one of two functions:
Linear Output: Each Gaussian component in the mixture has an associated vector wy which yields the final output as follows
y
Wgg Te.
where is the identity of the cluster of the tth token in the sequence. For most of our experiments in this section, we use this output function.
Quadratic Output: As before, we have was corresponding to each Gaussian. The final output is
(w xz,)²
where gi is the identity of the cluster of the 7th token in the sequence. This is a harder setting for model to solve without figuring out good routing function.
We experiment with learning the above function using soft variant of an MoE model which takes in token at a time and routes different experts. The expert outputs are weighted using the router sofumax probabilities to give the final output. We run experiments with experts which are layer linear functions and compare the performance with a dense MLP model with 2 fully-connected layers. This additional layer allows for a fair comparison as the router can be thought of as a layer in itself.

---

Results. On the linear output function, at small sequence lengths we observe stronger performance by MoE in comparison to the dense model. As the sequence length increases, we start to see more of the clusters the mixture present in the input. Since the final output is a linear function of the tokens, a dense model can solve this approximately by learning a single linear function which represents the expected value of the output. We start to see this behavior come out when sequence length increases and the dense model starts to perform just as well as the MoE. We observe that the MoE model attains a non-trivial sparsity and the sparsity is lower than what we get under shuffled router. On the quadratic output function, we observe that MoE with 32 experts strongly outperform dense model with the MoE model achieving MSE loss of 0.0075 compared to MSE of 0.018 by the dense model. The sparsity shuffled router sparsity of I 11.5 showing that the model has learned cluster directions to a non-trivial degree and is using this learnt routing to solve the problem more effectively than a dense model with an additional layer.
measure of the MoE model is 8.8 compared
Model 'Sequence Length
4
8
16
Dense model with 2 layers
0.0044
0.0021
0.0015
MoE with 32 linear experts
0.0016
0.0017
0.0015
MoE with 32 linear experts and frozen router
0.0046
0.0020
0.0027
Table 1: MSE for Tokenized Variant of Mixture of Gaussians experiments.
Model \Sequence Length
4
8
16
Sparsity
13
14.14
14.29
Sparsity with shuffled router (trained)
14.29
15.41
15.6
Sparsity with frozen router
28.52
28.51
28.53
Table 2: Sparsity for Tokenized Variant of Mixture of Gaussians experiments.
C.3 More Details on Dictionary Learning Experiments
We first describe the setup of our experiments in Section 3.6 in more detail. We sample e random directions in a d dimensional space each corresponding to the center of spherical Gaussian. We ensure the Gaussians are well-separated from each other. An input is generated by first sampling from of the Gaussians (without replacement) and adding them. Here denotes the rank of the dictionary learning problem and we use r 2 in all our experiments. Additionally, we have vector w, associated with each Gaussian which are randomly sampled. For an input x which is generated from Gaussians and j, the output is given by
(w; + w;) To.
Results. We first run experiments with 32, 32. In this setting, our MoE model uses 32 linear experts. We compare it with dense model with 2 fully connected layers and ReLU activation between them. Firstly, we observe that both the dense and the MoE model perform similarly and achieve MSE scores of 0.043 and 0.041 respectively. When we fix the dense model also be of depth 1, we observe a the MSE increase to 0.19 indicating a larger gap from MoE's performance. However, we don't 't observe the router learning the cluster center directions. We see a sparsity of 4.5 but the same sparsity is observed even under shuffled router. This indicates that the MoE model is achieving sparsity in this setting by blowing up the router norm rather than learning useful directions. We do, however, see some evidence that the model is learning something cluster specific. When we measure the average sparsity across samples from different pairs of Gaussians observe value of 8.1. This is mcasure of how many total experts the model is learning use. When we fix one of the Gaussian and vary the second one, we observe the average sparsity drop to 6.4. This indicates that on average, when one Gaussian is fixed, the router uses 6.4 experts versus when 8.1 experts used when both Gaussians are varied. To investigate whether this is happening because the problem setting was too easy to solve, we increase c to 128 and keep the number of experts as 32. In this setting we observe the MoE model achieve an MSE of 0.23 as opposed to 0.26

---

achieved by the dense model. Here, the problem is hard enough that none of the models are able to achieve close to 0 error. Moreover, the sparsity and shuffled sparsity of the MoE model are both 3.01 making it unclear if the MoE model learning a useful routing. To understand this, we train an MoE model with 32 experts and a frozen router. This achieves a much worst MSE of 0.38 indicating that there is a clear benefit to learning to route, although there is still a lot of room for the MoF. model to improve as it is not learning the optimal routing which can fully solve the problem). Overall, we conclude that this setting is a challenging one for the simple form of MoE routing to handle. We leave it open for future work to come up with routing schemes which work well for the dictionary learning setting we propose here.
C.4 More T5 Experiments with varying number of Experts
We describe the setup of our experiments in Section 3.7 in more detail. We adapt the T5-Base model which an Transformer mucoder-Transfore
decoder model with 12 layers each. uses model dimension of 768 and hidden dimension of 2048 along with 12 attention heads each of dimension 64. We run pre-training for this model on a language modeling task on the mC4 dataset and measure the next-token prediction accuracy achieved. We use an implementation of the MoE idea in this model which uses 3 sparse layers equally spaced among the 12 layers. It routes to the top-2 experts (instead of top-1) and performs load balancing using expert capacity factor of 1 1.25 during training. During evaluation the expert capacity factor is set 2. We train these models for 500k steps using Adam optimizer.
Results. On this model, we wish to study how important it is to learn a good routing function. Even if one freezes the router to its randomly initialized values, due to the earlier layers being allowed to train, the model can implicitly learn clever routing functions due to the large model dimension is has at its disposal. Indeed, we see evidence to support this when we try to learn a Base-sized model with frozen router, the pre-training accuracy is not affected significantly (70.34% vs 70.27%). To prevent this implicit learning of routing functions from happening severely restrict the model dimension and hidden dimension of the model. We reduce the model dimension to 32 and hidden dimension 256 and keep everything else the same. Now we compare training with a frozen router to training with a trainable router and observe a consistent strong improvement in performance when the router is allowed to train (Figure 7) across different total number of experts used.
To add to the experiments in Section 3.7, we perform similar comparisons between the pre-training accuracy of an MoE model trained with frozen routers vs a model with trainable routers at differing total number of experts.
D Proof of Theorem
We first re-state formal variant of the theorem here.
Theorem 3.
Let Jo,wu Rd -> R be a prediction function parametrized by weights b E R.A, W E B and U such that f,,wy (x) = gu (W for some function g. Let y Rd R$ be target function that depends only on some d' < d coordinates; iee., y is a function of a projection of x onto d' d'-dimensional subspace V. Then any local minimum of the regularized expected loss
L(b, - (.,.,((,.... + A1llW]| 2211U112
A >
must have each row of W within the relevant subspace V. That is, for any i e [M], Projv-(W.) = O.
Proof.
We prove the theorem for the where V is axis-aligned and is in particular the span of the first d' coordinates; the general case follows by a similar argument once you use rotational symmetry.
We write x = (L) x (2)) where (1(1) is the first d' < d coordinates of the input æ. Note that 1) (1)) and x (2) are independent, 2) (2) is Gaussian.
Consider local minimum (b, W. U) of L. Let (W1, W2) denote the column- split of W such that Wa (1 + (2) For the sake of contradiction assume that Wa not zero.
Take some e E (0, 1) and define a corresponding W" (W1, W3)= E glodd such that W = (1 - W2. Let & E T be a Gaussian random vector such that 2 is independent of x and Wzx((2) + 2 equals Wz (2)

---

Trainable Routers
52.0
Frozen Routers
51.5
Accuracy (%)
51.0
50.5
50.0
20
60
80
100
120
Total Number Experts
Figure 7: Pretraining Accuracy on mC4 using different T5 models. Each run is 12 layer model with model dimension reduced to 32 (from 768). This has been done to better capture the role of learning to route. Each model is allowed train for 500k steps.
in distribution. Since (2) is Gaussian is suffices that
B[z] aE| (22 Cov(z) / (1 --))))2 Cor
Since (11) and .2))) are independent, we have in fact (x(1), W1x(1) W2x(() 2) and (z(1), Wix(1) (2x are equal in distribution. We can therefore write
E, (x. xx()))
200) 90 b))
= E,E, E,E,o ()) ) (Wix(1) 2zz((2) + b)).
Define
2, arg min 200 ae
+ W3r(() b)).
and D' = =4 b. Then it is clear that
Exl(y fo ,2(2)))
E, E,E,, 200 (lv(x())), gu (W]x + Wzzx(2) + + b))
> Egeu m00 ((v(x()), su (W + Wy + 2, + b))
E (2) ((y(x()), ), (W'x +b))
and since [W] |W"] and d1 O,
L(b, W,)) L(b' W'.
Since & was arbitrary it follows that (b, W, U) cannot be a local minimum, which is a contradiction.

---

E Router Learns Latent Cluster Structure in Mixture of Gaussians
We begin by describing the problem setting we consider. Let {ge = N Mue, de) be the & spherical component Gaussians of a mixture distribution defined by Σe= Wege. In our work, we assume that all the mixture weights We are uniformly = 1/k for simplicity. Let (vole= be & cluster specific vectors. Given x. ~ Ei=1 9:/k; define the ground truth function f*(x) Vo(x) where d(x) represents the mixture component from which z was sampled.
First, we observe that if l/c- ee for all the components, then the problem has close resemblance to the problem of k-means clustering.
k-Means Clustering: Given input data {z}== E R the k-means objective tries to find & cluster centers [uy}z=1 E Rd such that the average squared distance between each point and its closest cluster center
7
h
1 = arg min I0,
w 22
[E(k]
is minimized.
We present a simplified MoE architecture with weight sharing to formalize this connection with k-means.
MoE Architecture with Weight Sharing: The neural network of the form: f(x) Ei= P:us where p: = > i = arg min; ||u; - x||2 and u 's are its parameters (experts).
We first note that if we set the ground truth output to be the input x; in our setting, then the best solution would be uj - My which is what would optimize the k-means objective as well. Thus k-means can be viewed supervised problem in this sense.
Lemma 1. Gradient Descent square loss in the above setting with the weight sharing MoE architecture is equivalent to optimizing the k-means objective using step size.
Proof. This loss functions in the two settings match and hence gradient descent will make identical steps.
Next, we generalize our architecture slightly by dropping the weight sharing assumption to get some- thing closer to an MoE model.
Ei= pru; where Pt = softmax(uu) wi 2) and ug wy's 's are parameters. The wi 's's represent the router and the u4 's represent the experts. Recall that in this setting we have {g: = (pe, == as the k spherical component Gaussians of mixture distribution defined by Σ Wigt- We assume that all the mixture weights wy are uniformly 11/k for simplicity. We also assume that the centers 14 are all randomly chosen in d-dimensional space such that || l2 < av d. {uitez be & cluster specific vectors. Given z ~ Σ gi/k: the ground truth function (x) = Ve(x) where c(x) represents the mixture component from which x was sampled.
A more general MoE architecture: The neural network is of the form: f(x)
MoE architecture: The neural network is of the form: f(z) = Ei=1 Piuj where p; = softmax w x) and uy; whess 's are trainable parameters. The wy's 's represent the router and the u;'s represent the experts. Given the above, we restate Theorem 2.
@dkk log samples drawn from a mixture of k spherical Gaussians in d- dimensions which are c-separated for some constant described above, consider an instantiation of the above MoE architecture with O log k) experts. If we initialize the router weights w, randomly drawn examples from the train set, the router will learn to route examples according to the cluster they belong to.
Theorem. Given
Proof. Recall that [g: = NN os = are the & spherical component Gaussians. Let X = {x;) be the train samples and let Wo = {wyl [ We initialize the router weights to Wo. A simple coupon collector argument would give us that with
O(klog())) be O(k log k) randomly selected samples from X.

---

probability 1/poly(k)) there exists least one sample from each of the l component Gaussians in Wo. Moreover, due to the c-separated nature of the Gaussians, for any two components i j
Ing 3e 9g || -yyl S ccd]
1 S d
(1)
for a small enough constant C. Let E denote the set of all the experts. Note that there is a one-to-one association between an expert ej and router weight vector wj. Partition the experts into & sets where Eg is the set of experts whose corresponding router weights were initialized using samples from the h component 9k- We call the set of experts Eg as cluster k's specialists.
When an input 24 Ok needs be routed, the total probability weight the router assigns to experts not Eg is I exp(-O(d)) Hence, assuming 2 k, the output is going to have contributions largely from the experts Eig- This implies that the contribution of the gradient from the loss on input æ; is minimal on the experts not in Eg. Among the experts in Eg, there will cxist at least one expert which each gradient will push in the direction of wj. This can happen with multiple experts in Eg at the same time as well. Moreover, using the same argument as above, the experts in Σ do not get significantly affected by the gradients coming from examples Ij which are sampled from Gaussians other than Dk- By picking small enough learning rate for the router, this leads to the following configuration after a logarithmic number of gradient descent steps. For every there will exist expert in j E3 such that luy w 12 S o(1). In addition, the train loss will also be close to Once this configuration has been reached, the subsequent steps of gradient descent only serve to boost the norms of the router weights without changing their direction by much which leads to a further cementing of the cluster based specialization of the experts.
F
Mixture of two truncated Gaussians
Task: We are given inputs drawn from a balanced mixture of two spherical truncated Gaussians in IR each with a Ia covariance and means t R' where || | 1. The Gaussians are truncated at radius r where < 1. The ground truth output is fixed to for the two clusters respectively.
Network architecture and optimization setup: The neural network is of the form: pu+ (1- (2p 1) where p(x) = Sigmoid w z) and w, u, b are parameters. Here p corresponds to the router and =u to the 2 experts. We will train with population gradient descent with MSE loss. We initialize the w and u as random vectors with norms tending to O. We will use learning rates Tu and qu for parameters u and w such that 1/u >> Qw-
Theorem Gradient descent on population loss in the above setting will converge to a solution in which =2, ||l| + oo and u = tv.
Proof. Let Wh denote the initial value of w, we assume wlog that wi > 0. Let G and G' denote the distributions of the two truncated Gaussians. The loss is:
Eseg [(2p(x) 1)u w)| Ezec (2p(2) (2ec=)l(2p())
1)u + +)]l
We note that the probability of sampling under G is the same probability of sampling - under G'. This allows to rewrite the above loss
Execl (2p(z) 1)u --vr 11(2p11())))))) +)
Using p(-x)= =1 p(x) we get that the loss functions equivalent to
L(u, w) = 2Ercoll 1
where q(x) = -2p(x) 1. As nu > The we can assume assume that throughout the training trajectory aL 0. Hence Bu

---

Ex
solving for gives
u² (w) (x)
E (x)
Now let us consider the partial derivative of loss wrt w
aL (Exco(((r) - 1) (ug' (x)())) (x))x)
EEreod (x) vf (=))))(( -F)g
By a symmetry argument the second term is of the form ew where is scalar. Let us focus on the coefficient of in first term (omitting constant factors). As Th The we can assume that u"ww.. hence we start by substituting that
v²u)= - 1110-1110 11 (x) ((xx) Ereag(x) 2
Exegq(2) (xx) Execq (ar)
Exegd (=l[)))))((
BzEgq
Brecola)) )² |lul Exxe''z ((xx)
Exccg(x)
Exccqp(x))
(xcoqq(x))
Execq(2)
Ex (x) ) -1)
EvECI(()))
||u||?
(x)) (((r))
(x)
EVEGq (x)
||u||l
(x)))llg()))()))BB<<c))) (Enegq' (x)) (x) (=)))
(E (x))² (z))
Claim 1. If wf 0 then
EveGq(Z) >
Proof. Let w = az + wi where and wF= O. For any point x Bz + a1 where xf== is easy to see that the point x' = B3 #1 occurs with the same probability as x in G. Hence
(xeGg(x) (xxec((()) + q(x'))()
Let a = w' x and b wTx x' then a + 202"2 = 2a > 0. As h(t) = 2Sigmoid(t) is anti-symmetric and monotonically inereasing function, a > 0 implies that h(a) + h(b) > 0. This in turn implies that q(x) + q(x')= > 0 as q(z) h(a) and q(x') h(b). This is sufficient to prove that
EmEGq(z) Exeo(g(r) + q(x'))22
Claim 2.
(x)))(4)))))) (Esecg (x)) (x))
Proof. We will prove the following stronger statement:
E.cold (=)))) (x)) < E2ccll2 (xccco (x)
is easy to show that o' (ι) = (1 q² (x))/2 Let D be the distribution of [g(x)) where x G. We can equivalently show that

---

where f(y) = (1 /2) This statement directly follows by applying Lemma 2.
This shows us that our gradient update which is proportional to aL is of the form ow + Bz where ow 3 O. This implies that GD will converge to solution which is parallel to Hence can restrict ourselves doing GD with being parallel to Let us rewrite the loss with the constraint that w az for some scalar a.
Luu w) = ee (a(x) u)
Denote by the component of x along where is a scalar. Let G11 be the distribution of x1. As ④ =z we know that (2x) is determined by Hence we can rewrite that loss
L(u, a) EG, EG, ((q1 v)12
where qı (x1) - (((12) = 1 - 2 u² 2) = 1 2Sigmoid(amr)) as |>| 1.
Al 13 EErc (x)))(((1)) T)Tu
1Erec da
- (a)
Beccol(z) (Beecy((1)
)
Encco(() (x1) Execgy2 (x)
-||lo| (1 (((1)(1)))(0))))
-Ilul| (I (Ba,, (aq] (x) (Eregg (x)) (Evegag (En (
Assume w.l.o.g. that a > O, assuming this we will prove that AL < 0 which will be sufficient to prove that as we keep running gradient descent we will get á ||
8a
co.
As G was truncated Gaussian we have that that > implies that (x1) > 0 for all x G1. Hence only need to prove that
EG, (x] < BeC (x))
Claim 3. a > 0 then
ExieG, (z)g(x)) (x)C()))))))(( <
Proof.
We will prove the following stronger statement:
Emicc, (x) (x)
Execly² (a)
It is easy to show that o' (x: (1 ( (x)) and In [+g(2) [q(x) where G. We can equivalently show that
/a. Let D be the distribution of (-
(RyED[((y)v] By < By
which is decreasing function of y for < < hence f is concave on (0, 1). Using f(0) = 0 and applying concavity on the three points 0, a, where 0 < < b <<1 gives us that
where f(y) = In 1+0 (1 - )2)/) Note that (v) = (1- - y In
1+y )/=))>/ )/a
1-
f(a)
f(b) b
With this property of f the required statement directly follows by applying Lemma with T

---

Hence u will converge satisfying i and || co. As started with truncated Gaussians, we will have g(x) 1 for all G1 and hence u² (w) +
We will use the following lemma from (Hidden, 2022), which we reproduce for completeness.
Lemma 2. (Hidden, 2022)] Let f be function which satisfies that for all 0 < b < T we have
f(a) f(b)
a
Then for any distribution DD supported only on (0, T) we have
E [[=)))(())) E[
Proof. Let y be an independent copy of æ. We note that
-E = fww) -
-Elyff(x)) - Elyxf(x)
= ((te-))(=f())))00 -f(=))
Because of our assumption the random variable (o (x=))(=f(1)) y. (x) is non-negative since (w)/ fxx / has the same sign as y. The desired incquality follows.


Reviews:
[{"review_id": "7MgT1wdMB5", "reviewer": "Reviewer_6A1b", "summary": "", "questions": "See weaknesses", "limitations": "", "rating": 3, "confidence": 2, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and feedback. We are glad you found our theoretical insights valuable. We address your concerns below\n1. **Limited motivation provided**: We would like to highlight that we devoted almost the entire first page to motivating the specific directions we study in our work. The page limits force us to be selective about the length of the motivation section and we could not delve deeper into how MoE models came to be and what some of the major open questions are in this direction. However, we understand that to a reader new to this field, the motivation provided might be a bit terse. We will add additional background into the problem and prior work in the Appendix in the future versions of our paper.\n1. **Presentation Issues**: We thank you for highlighting the length of Section 1.2. We will polish this in future versions of the paper. You mentioned that the paper was very difficult to understand to readers new to this field. Could you please highlight which particular parts of the paper were difficult to understand?\n\nApart the presentation issues, there do not seem to be any technical concerns raised. If you feel our research is technically sound, we urge you to consider raising your soundness score.", "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": "1. The paper offers theoretical insights into the benefits of Mixture-of-Expert (MoE) models, specifically those like the Switch Transformer. It establishes the advantages of these models beyond mere scalability, which can contribute to the theoretical foundation of the field.\n\n2. The paper is well-structured, presenting its findings and reasoning clearly. ", "reasons_to_reject": "1. The paper does not talks much about the problem setup. The problem is not well motivated.\n2. The paper is very difficult to understand to readers new to this field.\n3. Paper has some presentation issues as section 1.2 could be made small by just putting key contributions instead of a large section.\n4. There is a very small section of problem setup and most of the paper is covered by experiments. However, it would be better if the authors make the presentation better and place things under the correct sections.", "questions_to_authors": "See weaknesses", "paper_topic_and_main_contributions": "In summary, the paper explores the broader benefits of MoE models, such as the Switch Transformer, beyond just scaling. It demonstrates that the intelligent routing capabilities of the router lead to advantages in effectively utilizing parameters and partitioning input data. This understanding provides valuable insights into the potential of MoE models in various applications, shedding light on their advantages in both theoretical and practical contexts.\nThe paper first examines synthetic scenarios where input data is distributed in clusters. Theoretical and empirical analyses demonstrate that the router learns to capture the underlying cluster structure, enhancing its effectiveness. The research then extends to real-world data using the T5X library. It's observed that a trainable router offers a substantial benefit over a non-trainable one, indicating the practical advantages of the learned routing function.", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "No", "Typos_Grammar_Style_and_Presentation_Improvements": "No"}, "extra_text": {}}, {"review_id": "yfYJ7LnAqD", "reviewer": "Reviewer_iWp2", "summary": "", "questions": "Please comment on the second reason to reject- the question of routing utility for large transformers seems to remain unanswered. Can you suggest future work that may be able to address it without artificially constraining the model width?\n\nRegarding Theorem 2, could the theorem hold if the clusters are non-uniform in size?  In a practical use case such as the classification setting is there any utility in stratified sampling or other intelligent sampling of initialization examples to ensure coverage across all the classes?\n\n", "limitations": "", "rating": 3, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your careful reading of our paper and for the constructive comments. We are glad you found our theoretical and empirical insights valuable. We address your concerns and questions below.\n1. **Contributions Limited to Simple Model Architectures**: We want to highlight that our paper focuses on the Transformer architecture and the MoE variant of the Transformer architecture. The Transformer architecture is well established to be an incredibly powerful architecture which can scale well to sizes ranging across many orders of magnitudes. Conceptually, apart from the attention component, the Transformer is indeed a simple repetition of a 2-layer MLP architecture which allows us, for the purpose of theoretical study or synthetic data study, to focus on MLPs as an example. Proving theoretical results or developing interpretable results for the Transformer in its entirety turns out to be extremely challenging. Hence, we focused on simpler yet key pieces of the architecture to develop insights.\n1. **Artificial Constraint on the Width in Experiments**: Thank you for highlighting this point. To clearly ablate out the benefit of having a trainable router, we chose to make the width small. We originally tried freezing the router in a normal width transformer (the T5-Base model). In this setting, we observed that the difference in pre-training accuracy between the model with a frozen vs trainable router was not very large (although there was still an advantage to having a trainable router). As noted in the paper, we hypothesize that this happens due to the frozen router learning implicit routing using the earlier layers in the model.\nDue to this reason it is tricky to ablate out the precise effect of the trainability of the router. \nWe would like to point out that recent works ([1,2,3]) in the space of MoE models all use trainable routers even when the routers are very different from each other.\nWe believe developing an experiment that could demonstrate the gains of a trainable router in large transformer models is an important yet challenging future direction.\n1. **Theorem 2 with non-uniform clusters**: Thank you for raising this interesting question. We believe that the results of Theorem 2 can indeed be generalized to hold when the mixture weights are non-uniform as well. This will complicate the proof significantly but there is nothing crucial holding us back from generalizing the theorem to this setting.\n1. **Utility of Stratified Sampling or Other Intelligent Samplings in practical use cases**: This is an interesting research direction. The use of intelligent samplings in dataset selection is a very actively studied area where we do have some insights such as the following:\n. Selecting a good balanced dataset is important for learning good classifiers\n. Significantly decreasing the dataset size by selecting examples of the right difficulty.\n. Curriculum learning where we present examples in order of increasing difficulty to the model.\nA tangential direction which our work touches upon is that of using data to initialize the parameters of a model. This is known to give benefits in simpler models such as K-means clustering algorithms. For the complex setting of deep neural networks, it is unclear how to initialize a deeper layer using the data embeddings we have at the input layer. However, there is work [4] which shows that attaching an input token based lookup at the middle of a transformer layer has benefits in scaling up the model capacity without adding too much compute. In addition, there has been work which uses the gradients obtained on the data at initialization to sparsify the network ([5,6]).\nAs far as we are aware, there isn’t work which uses input samples to smartly initialize the deeper layer routers of an MoE model. We think this is an interesting direction with potential for efficiency savings.\n\n- [1] - https://arxiv.org/abs/2209.15466\n- [2] - https://arxiv.org/abs/2308.00951\n- [3] - https://arxiv.org/abs/2103.16716\n- [4] - https://arxiv.org/abs/2302.00003\n- [5] - https://arxiv.org/abs/2202.12002\n- [6] - https://arxiv.org/abs/1911.13299", "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": "\nThe paper provides some clear evidence of when a trainable router will outperform a fixed router, and when the benefits of MoEs might be diminished (ie when competing with large dense models).\n\nThe work also shows that trainable routers can learn to ignore spurious dimensions in the input. \n\nThe theoretical results hint at a good initialization strategy for linear routers.\n\n", "reasons_to_reject": "The contributions are relatively modest and largely limited to simple model architectures.\n\nBy artificially constraining the width of the T5 model, the observation that a trainable router outperforms a frozen router seems to be purely academic- the important practical question is whether trainable routing can improve performance for a large (wide) transformer.\n", "questions_to_authors": "Please comment on the second reason to reject- the question of routing utility for large transformers seems to remain unanswered. Can you suggest future work that may be able to address it without artificially constraining the model width?\n\nRegarding Theorem 2, could the theorem hold if the clusters are non-uniform in size?  In a practical use case such as the classification setting is there any utility in stratified sampling or other intelligent sampling of initialization examples to ensure coverage across all the classes?\n\n", "paper_topic_and_main_contributions": "The paper examines MoE models from an empirical and theoretical perspective, aimed at understanding the value of trainable vs fixed routers and providing some theoretical results for relatively simple mixture of gaussian settings.\n", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xkPCboN8mG", "reviewer": "Reviewer_qHJT", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your careful reading of our paper and for the constructive comments. We are glad you found our theoretical and empirical insights valuable. We address your concerns and questions below.\n1. **Most Experiments are on Synthetic Data**: One of our main goals was to study whether MoE routers learn interpretable functions. This is hard to achieve with real data. This was one of the main reasons we focused so much on synthetic data in our work. We would like to point out that we do study the importance of having a trainable router on real data (the mC4 dataset).\n1. **Theoretical results tackle simple cases like mixtures of Gaussians**: Firstly, we would like to highlight that mixtures of Gaussians is a powerful framework which, if unrestricted, can model any arbitrary distribution in practice by using the Gaussians as ‘basis’ functions. While we agree it would be nice to have theory for more complex distributions like the ones encountered in practice we want to highlight that it is incredibly difficult to prove theoretical results about more complex distributions in general. Moreover, many theoretical insights in Machine Learning such as the interplay of overparameterization and robustness were obtained by studying Gaussian distributions.\n1. **Latest SoTA Models**: We would like to highlight that our main motivation in this work was to understand the importance of having a trainable router and to study whether the router learns interpretable functions. Hence, we focused on a clean version of MoE routers which is the most popular variant used today. The SOTA MoE models which exist today (e.g. [1,2,3]), while having some variations over the original MoE paper [4], all use the same fundamental router as a backbone to create different routing decisions.\n1. **Clarity in Theoretical Results**: We apologize for the lack of clarity in our description of theoretical results. We will undertake the following actions to polish our writing for future versions of the paper.\n\t. We will clean up Section 1.2 focusing more on the implications of our theoretical results rathethan the techniques used to prove them and punt a discussion of the techniques to later in the paper.\n. We will present clearer connections to practice for each of our theoretical results.\n. We will have a discussion of some extensions of our theoretical results that would be interesting and have stronger implications.\n1. **Implications of our Results for Architecture Design**: This is indeed an intriguing research direction. However, we feel that it is worth a full set of papers on its own as there are many things one could try and it is a challenging task to figure out how to use the insights we gained through our experiments to help design better MoE routers.\n\n- [1] - https://arxiv.org/abs/2202.09368\n- [2] - https://arxiv.org/abs/2202.08906\n- [3] - https://arxiv.org/abs/2308.00951\n- [4] - https://arxiv.org/abs/2101.03961", "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": "1. The paper provides both theoretical analysis and empirical evidence on the advantages of learning to route in Mixture-of-Experts (MoE) models. This sheds light on an important aspect of MoE models that is not yet well understood.\n\n2. Through experiments on synthetic and real data, the paper gives evidence that MoE routers can learn to route inputs intelligently based on latent structure in the data. This is an interesting finding that helps us better understand the behaviors and capabilities of MoE routing mechanisms.\n\n3. The theoretical results prove that MoE models with learned routing can achieve good accuracy on classically studied clustering tasks like mixtures of Gaussians. This provides a theoretical grounding for the empirical advantages of learned routing. \n\n4. The paper introduces several carefully constructed synthetic tasks like mixtures of subspaces, tokenized mixtures of Gaussians etc. that help clearly demonstrate the effects of learned routing. These could be useful testbeds for future work on routing mechanisms.\n\n5. The experiments on real NLP data (language modeling using T5) provide evidence that learned routing gives a significant boost over fixed routing schemes. This validates the relevance of insights from synthetic tasks.", "reasons_to_reject": "1. The paper lacks experimental results on real-world data. Most of the experiments are on synthetic data specifically designed to demonstrate the benefits of mixture-of-experts models. It would be stronger to show these benefits conclusively on complex, real-world datasets.\n\n2. The theoretical results are limited to simple cases like mixtures of Gaussians and may not extend to more complex data distributions encountered in practice. The assumptions made in the theoretical analysis may not hold in real settings.\n\n3. The paper does not compare to the latest state-of-the-art models. Many advances have been made in sparse expert models since this paper was written. The paper would be stronger if it benchmarked against the current best methods.\n\n4. The writing could be improved for clarity in some areas, especially in describing the theoretical results.\n\n5. The implications of the results are not fully explored. For example, the insights into what the router learns could guide architecture designs, but this direction is not pursued.", "questions_to_authors": "", "paper_topic_and_main_contributions": "This paper studies the benefits of learning to route in Mixture-of-Experts (MoE) models. The key contributions are:\n\n1. Empirical studies on synthetic and real data showing that learning to route intelligently provides significant benefits over routing randomly or using a fixed routing function. The experiments show the router can identify and exploit latent cluster structure in data. \n\n2. Theoretical analysis proving that with simple MoE architectures and gradient descent, the router provably learns to route inputs from a mixture of Gaussians to the correct expert. This holds under standard assumptions on the mixture components being well-separated.\n\n3. Identification that the benefits of MoE models arise from two factors - having more parameters due to sparsity and the router learning an intelligent routing function. The work provides evidence that the second factor of learned routing is very important.\n\n4. Comparisons between MoE models and dense models. While MoE models outperform smaller dense models, the gap reduces with larger dense models in the overparameterized regime.\n\n5. Insights into what type of cluster structures are easy or hard for simple MoE routers to identify. The router excels at separable mixture components but struggles on more complex latent structures.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper examines MoE models from an empirical and theoretical perspective, aimed at understanding the value of trainable vs fixed routers and providing some theoretical results for relatively simple mixture of gaussian settings.

Pros

It provides theoretical analysis and empirical evidence on the advantages of learning to route in Mixture-of-Experts (MoE) models

 It establishes the advantages of these models beyond mere scalability, which can contribute to the theoretical foundation of the field.

Through experiments on synthetic and real data, the paper gives evidence that MoE routers can learn to route inputs intelligently based on latent structure in the data and that learning to route yields a significant boost over fixed routing schemes.

Cons

Most of the experiments are on synthetic data. It would be more impactful to show its practical usefulness in real-world complex datasets 

The authors have done empirical analysis using quite small models (T5-base). In the current dynamics of this space, in order to make this paper really relevant it would be good to understand whether their findings hold true for large enough models.  While for the theoretical part i can understand simplifying the model architecture, but for the empirical analysis its important to show its generalization to larger LLMs. 

Presentation of the paper and writing style needs to be improved so that the motivation and the exact problem setup and discussion of the results is clear to readers, especially those who may not be very familiar with the topic
Decision: Accept-Main