On the Benefits of Learning to Route in Mixture-of-Experts Models (2023)

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.

## introduction with the advent of large transformer based models @xcite @xcite , model size has increasingly become an issue with neural networks. larger models continue 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 @xcite which uses a mixture-of-experts (moe) architecture. mixture-of-expert models (figure 1 ) in the simplest and most-often used setting are described by a set of experts e 1 , . . . , e n and a routing function (aka gating mechanism). each expert is itself a small neural network (a popular choice for the expert architecture is 2 fully connected layers with a non-linearity in between). 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. the core principle behind using moe models is that only one (or a few) of the e i 's is invoked per example. thus, when normalizing for compute time, the learned model can draw from the benefits of having many more parameters than a dense model (which can be viewed as having a single expert). in this work, we look closely at a core component of moe models: the router, which routes different inputs to different experts. by design, moe models confer computational advantages, and several works @xcite @xcite have demonstrated that this speed-up is achieved without any loss in accuracy. we outline two major reasons for this advantage: we seek to theoretically and empirically understand what role the router plays in delivering the advantages that moe models enjoy over dense models. specifically, we study the following questions about the routing mechanisms in moe models: ## preliminaries we use [n] to denote the list of integers from [1, 2, . . . , n]. we denote the probability distribution of a d-dimensional gaussian with mean vector µ and covariance σ by n (µ, σ). in our work, we will deal with spherical gaussians where the covariance matrix σ = σi d . given two d-dimensional spherical gaussians n (µ 1 , σ 1 i d ), n (µ 2 , σ 2 i d ), we say they are c- @xcite ). 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 moe models by specifying the router and the experts. for example, a linear router with constant experts refers to a trained model which 1. on input x, computes w x where w denotes the router weights. the input is assigned the expert j where @xmath0.e. j th entry of w x is the largest entry. we note that this description applies to trained models. while training, we do a weighted combination of experts instead of routing to the top expert. all our results are for a trained network for which we do top-1 routing. see appendix c for more details. ## experimental results we present our experiments on synthetic (sections 3.1-3.6) and real datasets (section 3.7) in this section. full details of our experimental setup are in appendix c. ## 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 relevant directions and remove spurious directions. we prove this result for any moe architecture which passes the input through a softmax router first with a regularization term for the router weights. ## 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 capable of learning latent cluster structure in some scenarios and less capable in more complex scenarios 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 to learn a sparse yet interpretable solution to a complex problem.