Routers in Vision Mixture of Experts: An Empirical Study
Authors: Tianlin Liu
Mathieu Blondel
Carlos Riquelme Ruiz
Joan Puigcerver
Venue: Accepted by TMLR ()

Abstract: Mixture-of-Experts (MoE) models are a promising way to scale up model capacity without significantly increasing computational cost. A key component of MoEs is the router, which decides which subset of parameters (experts) process which feature embeddings (tokens). In this paper, we present a comprehensive study of routers in MoEs for computer vision tasks. We introduce a unified MoE formulation that subsumes different MoEs with two parametric routing tensors. This formulation covers both sparse MoE, which uses a binary or hard assignment between experts and tokens, and soft MoE, which uses a soft assignment between experts and weighted combinations of tokens. Routers for sparse MoEs can be further grouped into two variants: Token Choice, which matches experts to each token, and Expert Choice, which matches tokens to each expert. We conduct head-to-head experiments with 6 different routers, including existing routers from prior work and new ones we introduce. We show that (i) many routers originally developed for language modeling can be adapted to perform strongly in vision tasks, (ii) in sparse MoE, Expert Choice routers generally outperform Token Choice routers, and (iii) soft MoEs generally outperform sparse MoEs with a fixed compute budget. These results provide new insights regarding the crucial role of routers in vision MoE models.

Published in Transactions on Machine Learning Research (04/2024)
Routers in Vision Mixture of Experts: An Empirical Study
Tianlin Liu' University Basel
Mathieu Blondel Google Deep
Carlos Riquelme! Stability Al
Joan Puigcerver Google DeepMind
Work done as an intern at Google DeepMind. Work done while Google DeepMind.
Reviewed OpenReview: https: // openreview. net/ forum? id= deccebeccchfffttoeeeess-777-
Abstract
Mixture-of-Experts (MoE) models are a promising way to scale model capacity without significantly increasing computational cost. A key component of MoEs is the router, which decides which subsct of parameters (experts) process which feature embeddings (tokens). In this paper, we present a comprehensive study of routers in MoEs for computer vision tasks. We introduce unified MoE formulation that subsumes different MoFs with two parametric routing tensors. This formulation covers both sparse MoE, which uses binary or hard assignment between experts and tokens, and soft MoE, which uses soft assignment, belween experts and weighted combinations of tokens. Routers for sparse MoEs can be further grouped into two variants: Token Choice, which matches experts to each token, and Expert Choice, which matches tokens to each expert. We conduct head-to-head experiments with 6 different routers, including existing routers from prior work and new ones introduce. We show that (i) many routers originally developed for language modeling can be adapted to perform strongly in vision tasks, (ii) in sparse MoE, Expert Choice routers generally outperform Token Choice routers, and (iii) soft MoEs generally outperform sparse MoEs with a fixed compute budget. These results provide new insights regarding the crucial role of routers in vision MoE models.
Deep learning is using ever larger models. However, larger models require extensive compulational resources for training and deployment. To further scale up model capacity without linearly increasing the compu- tational cost,
promising direction is to use mixturo-of-oxperts (MoE) layers in neural networks. The general idea of MoE layer is to assign different feature embeddings, also known as tokens, to subsets of neural network parameters, known experts. Compared to traditional dense models, MoFs yield better performance compute trade-off (Shazeer et al., 2017; Riquelme et al., 2021; You et al., 2021; Mustafa et al., 2022; Puigcerver al., 2023).
At the heart of the MoF model is its router, responsible for matching tokens with experts. Routers used before the deep learning era were reviewed in Yuksel al. (2012), where they termed router as a gate. In the deep learning era, Bengio et al. (2016) cast the input-to-expert routing problem Markov decision process, and used reinforcement learning to train the router. Shazeer et al. (2017); Lepikhin et al. (2021); Fedus et al. (2022b) introduce diffcrentiable training approaches of routers without the complication of reinforcement learning. Roller et. al. (2021) use non-learnable router that hashes inpul samples in a deterministic way. Clark et al. (2022); Kool et al. (2021); Liu al. (2023) parametrize routers based on optimal transport formulations. Sander et al. (2023) proposod differentiable top-k operators that help the lcarning of token- expert matching. All aforementioned examples are sparse MoEs that soek hard, binary token-expert match. In contrast., the recently proposed soft MoE (Puigcerver et al., 2023) router matches linearly

---

Published in Transactions on Machine Learning Research (04/2024)
combined imputs with experts. This soft MoE approach avoids the discrete, non-diffcrentiable transforms often used in spanse MoFs while having similar computational cost.
MoEs as layers in deep neural networks were first explored in language modeling (Shazeer et al., 2017) before being applied to computer vision tasks (Riquelme et al., 2021). It is natural to hypothesize that many language MoF. routers can be adapted to work well in vision MoE. models, given that both rely on transformer architectures (Vaswani et al., 2017; Dosovitskiy et al., 2021). However, we currently lack a comprchensive comparison between routers specifically designed for vision MoEs and those repurposed from language MoEs. For instance, the Sinkhorn router (Clark et al., 2022), initially introduced for language MoE, can be potentially used for vision tasks as well; but to the best of our knowledge, their performances have not been reported in prior vision MoE work. In this work, we address this gap by conducting an empirical study of routers in vision MoEs. We compare different vision MoE models head-to-head by varying their underlying routers. We then evaluate these models through few-shot transfer learning and fine-tuning on ImageNet, providing a comprehensive analysis of the performance of different routers image recognition lasks. The main contributions of this paper are summarized below.
We propose a unified formulation of MoE layers by way of parametrizing its dispatch tensor and combine tensor, which we collectively refer to as routing tensors. This generalizes the dispatch weights and combine weights of Puigcerver et al. (2023), which focuses on Soft MoE. This unified approach allows for a comparison of existing MoE routers and motivates to introduce new roulers.
We study sparse MoEs by decoupling (i) how they parameterize token-expert affinity matrices (e.g., based on softmax and Sinkhorn transform) and (ii) how they determine which experts use: Expert Choice (where each expert selects tokens) and Token Choice (where each token selects experts). We demonstrate that factor (ii) outweighs factor (i) in sparse MoE routers.
We show that the soft MoE router generally outperforms the sparse variants both previously re- ported and our newly introduced one in Section 2.4 under a fixed compute budget.
Notations. Throughout this paper, we use bold letters for vectors, matrices, and tensors. For a vector a, we let ali] be its i-th entry. For a matrix A E mmxx,, we let Aii be its (i.j)-th entry, let A[::j] E Rm be its j-th column, and :] R= be the i-th row of A, transposed into a column vector. We write (1......... We denote all-one vector in R. by 1p.
The paper is structured as follows: Section details a unified formulation of Mixture of Experts (MoF) layers, defined through their routers. Section 2 explores various instantiations of MoE layers within this formulation. Section 3 presents the numerical experiments conducted on MoE layers.
1 A Unified Formulation of MoE Layers
In this section, we first introduce the most commonly used router, which is based on the composition of the softmax function and the top-k function. We then provide a unified formulation of MoE layers that encompasscs many types of MoE layers as special cases. Throughout our work, without loss of generality, we let. each expert be Multi-Layer Perceptron (MLP).
1.1 A motivating example of MoE layer
In the most widely used MoE formulation, each token æ E F chooses & experts through softmax function (Shazeer et al., 2017; Fedus et al., 2022b; Riquelme et al., 2021; A Allingham al., 2022):
E
MoE(x)= = Σ (x) (x) with (x) = top: (softmax W x + + (r) R. Vr [ (1)
where the gating weights (Gate, (x)) } lincarly combine the outputs of the E experts (M)P,())]EL The function top, E - E sets all but the l & F. largest numbers to zero, thereby selecting few experts for the
2

---

Published in Transactions on Machine Learning Research (04/2024)
combination. The gate parameters W R trained in conjunction with other network parameters. The vector 6 ~ (0, I) is a noise injected to the logits XW, where E R controls the strength of the noise; through auxiliary losses, the noise helps to improve the numerical stability of router (Appendix A). We refer to the formulation (1) as Softmax Token Choice MoE (not be confused with Soft-MoE, described in Section 2.6). This name emphasizes that each token selects lt experts based on softmax scores.
Compared to a traditional multi-layer perceptron (MLP) layer, the MoE layer (1) offers more flexible balance between the model's ability to fit the data and computational cost. Evaluating the MoE layer is equivalent to evaluating fixed number of MLP layers in parallel, where the number of layers is determined by the value of k. By keeping k fixed, as the number of experts E increases, the MoE layer's ability to fit the data increases, but the computational cost remains roughly the same, to the computation of the gating weights (z)e, In this way, the MoE layer provides a way to control the model- model-fitting capacity and computational cost trade-off by adjusting the number of experts.
1.2 Unifying MoE layers through routers
We now introduce a unified formulation of MoE layers, which encompasses the Softmax Token Choice layer and other MoE layers as special cases. This formulation is for a minibatch setting, where a MoE layer processes a group of T tokens {æ1,...,, Ør) c TE in parallel. An additional hyperparameter used in this minibatch setting is the buffer capacity of experts. It specifies the maximum number of tokens that each expert can handle in minibatch. For efficient hardware usage, it is ideal for each expert to have small buller capacily C < T, such as C= [T/E] or C= [2T/R) E] (Riquelme el al., 2021). A minibatched MoE layer with buffer capacity C' is defined as follows.
A Unified Formulation of MoE Layers
Let (x1,..., ØT} C R be a minibatch of T tokens and let X E RTxD be matrix of row-wise concatenation of tokens with X[t,:] 2,. Let Dx E TxxExC and Cx E RTxEXO be certain dispatch and combine tensors that are functions of tokens X. A MoE layer with fixed buffer capacity C' E N, dispatch tensor Dx and combine tensor Cx is defined as
E c
MoE [):+; = ΣΣ Cxltn MLP, ( (x (xDx(()))) ERD.
(2)
r=
Here, Dx is termed dispatch tensor, since it is responsible for sending different tokens to different experts; Cx called the combine tensor, as it used to linearly combine the expert outputs.
Recovering the Softmax Token Choice layer as a special case. On a single token level, the Softmax Token Choice layer in (1), can be recovered using the more general formulation (2) by the following combine tensor Cx and the dispatch tensor Dx:
Cx(t,r,c]== = 0,
(at), t otherwise,
and Dx = 1 [(x 0).
(3)
With these choices, the MoE output in (2) is then
E
C
MoE x [b: Σ
Σ
Cx(br, MMPP (XTTxxxxxxix
E
Σ
Cx(4n: 1011:,.0MMMP,P..
(X (x Dx Dx((rr1))
(4)
E
(a:)) MLP, (2:)
3

---

Published in Transactions on Machine Learning Research (04/2024)
which recovers the Softmax Token Choice layer (1) single-token level. Note that in (3) C= T. In practical implementation, however, usually choose C' < T to reduce the computational cost. This is detailed in Section 2.1.
MoE routers. The Softmax Token Choice router uses a specific combine tensor Cx and the dispatch tensor Dx. But they can be paramctrized in more flexible ways as generic transforms of X. We refer this mapping X + (Cx; Dx) as MoE router, which is central to our study. It is defined follows.
MoE Router
A MoE router takes minibatch input X E RTxD and yields the dispatch tensor Dx E RTXEXC and the combine tensor Cx E TTXXXC that are used by the MoE layer (2):
Router: X - (Dx; Cx).
(5)
The dispatch tensor Dx and the combine tensor Cx are jointly referred to the routing tensors.
2
MoE layers instantiated by different routers
In this section, we present family of MoE layers with different underlying routers.
2.1 Softmax Token Choice router
We first revisit the Softmax Token Choice in Section 1.1, now in a minibatch setting. To begin, we build softmax affinity matrix that represents the similarity between each token-exper pair:
II softmax = softmax (XW + oE) E ETTXE
(6)
where the softmax is applied to each row (normalized across experts). We next use the affinity matrix Iosttmax to allocate the routing tensors Dx and Cx. Since experts have a fixed buffer capacity C, the dispatch tensor Dx and the combine tensor Cx must allocated in a way such that each expert receives at most C tokens. We achieve this by sequentially going through the rows of affinity matrix Igffmax and assigning each (oken to its top-scored expert as long as the chosen expert's buffer is not full (i.e., below the buffer capacity C*). This procedure is called Token Choice allocation and is described in Algorithm 1.
The forward pass of the Softmax Token Choice router is summarized below.
Softmax Token Choice router
Compute the loken expert affinity matrix Igsfmmax in (6) on a batch of input tokens X.
2. Use II [oftmax to allocate the routing tensors Dx and Cx through Algorithm 1
Balancing expert usage. In practice, it is beneficial regularize the parameters W for balanced usage of experts. Without any regularization on W, the top scores at each row of the affinity matrix Hoffumx in (6) may concentrate on a few column indices. These indices correspond to certain popular experts that many tokens prefer to choose. But since each expert has a fixed buffer capacity, popular experts may drop most of the tokens in the allocation procedure prevent this from happening, previous work considered different auxiliary losses on W. Details of these losses are provided in Appendix A for reference.
Algorithm 1, thus deteriorating the training performance. To
4

---

Published in Transactions on Machine Learning Research (04/2024)
Algorithm 1: Token Choice allocation
Data:
tokens X E RTxxD
token-expert affinity matrix II E RTXE
buffer capacity C' T N
(iv) number of selected experts E N per token (typically or 2)
Result:
The routing tensors Dx and Cx E RTXEXC for the MoE layer (2)
Initialize Dx E RXX and Cx E TX zero tensors.
for the top i-th choice with = l,...,k do for token index =1 1I.... w,r max,th [t.
T do
The value and index of the i-th selected score. // The number tokens already dispatched to the r-th expert
if c < C then Dx(t,r,e=1 =1 Cx(hr,
2.2 Sinkhorn Token Choice router
To perform well, the Softmax Token Choice router requires auxiliary losses that balance the expert usage. To eliminate the need for auxiliary losses, recent work (Kool et al., 2021; Clark al., 2022) proposed the Sinkhorn Token Choice router.
Specifically, a softmax affinily matrix Hoftmax = softmax XW) E TxE in (6) can be seen as the solution to an uutropy-regularin optimization problem:
Hoftmax = arg max II
(I, (II, log log
II 0, III = 1rs
subject to
where 17 is a column vector of size with values 1. This entropy-regularizee characterization decouples the underlying optimization problem and the constraints imposed on the solution. The constraints in Equa- tion (7) are that (i) all token-expert affinity scores have positive values and (ii) the coken-expert affinity matrix has a unit row sum, meaning that the alfinity scores sum to along the expert axis.
The idea of Sinkhorn Token Choice is to add more constraints to (7). As discussed in Section 2.1, unbalanced expert usage of Softmax Token Choice occurs when top token-expert affinity scores concentrate on few columns of the affinity matrix II Msotmaax- To promote a balanced expert usage, the Sinkhorn Token Choice router thus requires that all columns are normalized, leading to the following entropy-regularived optimal transport formulation
nent = arg max
(II - (II, log II)
II 0,
(8)
subject to
IIl g 1r, II 1r (T/E)1g.
The optimization problem (8) can be solved by Sinkhorn's algorithm (Sinkhorn, 1964; Cuturi, 2013; Peyré & Cuturi, 2019), which iterates between row-wise and column-wise normalization. We then allocate the routing tensors using entropy-regularized transportation plan I. In principle, the same Hat can be used
5

---

Published in Transactions on Machine Learning Research (04/2024)
for dispatching Dx and combining tensors Cx. But it is empirically more beneficial' to only allocate the dispatch tensor Dx using the transportation plan Hent while still allocating the combine tensor Cx using the softmax matrix ooftmax = softmax (tttmax(XW)) W) This approach has the benefit that the backward pass of mrndient-basee training does not go through the Sinkhorn algorithm, which results in faster and more stable training. We provide a case study of this Figure 3 of the results section; while Clark et al. (2022) recommend using 1, we have also experimented with k = and observe consistent results.
The forward pass of Sinkhorn Token Choice router is summarized follows.
Sinkhorn Token Choice router
1. Compute the token token-expert affinity matrix eent in (8) on batch of input tokens X.
2. Use eent to allocate the dispatch tensors Dx through Algorithm
Use softmax = MaltmaxxXXV to allocate the combine tensor Cx through Algorithm
2.3 Softmax Expert Choice router
Both Softmax Token Choice and Sinkhorn Token Choice share the router allocation algorithm in Algorithm 1. This allocation approach lets each token choose its top-scored expert. However, a limitation of this approach is that it can result underused experts, which are experts that use fewer tokens than their allowed capacity. This can be seen the inner for-loop of Algorithm 1. There, the number of tokens taken by each expert restricted to not exceed the buffer capacity C, that is, D S C for each expert When the inequality strict, || Dx(s,,, < C, the expert r is underused as could have been possible to allocate more tokens to that expert.
To address the underusage of experts in token-choice allocation, an alternative approach proposed in Zhou et al. (2022) is the Expert Choice allocation. In this approach, each expert chooses fixed amount of tokens. Specifically, given token expert affinity matrix II, sequentially go through the columns of II and assign each expert to its top-scored tokens. This procedure is outlined in Algorithm 2.
Al Algorithm 2: Expert Choice allocation
Data:
tokens X E RTxD
token expert affinity matrix II E RTxE
(iii buffer capacity C N
Result:
The routing tensors Dx and Cx RTXExC for the MoE layer (2)
Initialize Dx E TXCEX and Cx E XXXXXC zero tensors.
for expert index I,..., do
for capacily index = I,....C C do w max,th (aite,(x)), Gate, Dx[t,r,(==) =1 Cx(Lr, w.
The value and index of the i-th selected gates.
The Expert Choice allocation in Algorithm 2 offers several benefits compared to the Token Choice allocation in Algorithm 1. Firstly, it ensures all experts receive the same amount of tokens, preventing any from being over or under-used. Secondly, since experts independently select their tokens, many experts can process a single token. This contrasts with Token Choice allocation, where each token is typically assigned to only
Personal communications with the authors of Clark al. (2022).

---

Published in Transactions on Machine Learning Research (04/2024)
k or k experts. Such diverse expert participation enhances MoE layer performance (Zhou et al., 2022). Moreover, in the Expert Choice allocation, certain tokens might not be chosen by any experts.
In Zhou et al. (2022), a softmax token expert affinity matrix aofomax = softmax (X) E R is used in the Expert Choice allocation of Algorithm 2. We refer it to as the Softmax Expert Choice router, whose procedural steps are summarized as follows.
Softmax Export Choice router
Compute the token-exppe affinity matrix gottmax in (6) on a batch of input tokens X.
2. Use goftmax to allocate the routing tensors Dx and Cx through Algorithm 2
2.4 Sinkhorn Expert Choice router
We now introduce the Sinkhorn Expert Choice router, which fits naturally with the unified MoE layer formulation Section 1.2. To our knowledge, has not boen proviously explored literature.
Recall that the Softmax Expert Choice router offers versatile expert usage, but can skip numerous tokens. On the one hand, skipping tokens can be an advantage of Softmax Expert Choice, the buffer capacity saved by skipped tokens can be allocated to remaining tokens, allowing them to be processed by many experts. On the other hand, skipping too many tokens may lead to inferior results.
To strike a balance between the heterogencous usage of experts and the goal of minimizing token dropping, we can use the entropy-regularried optimal transport plan een to allocate the routing tensors through Algorithm As introduced in Section 2.2, the optimal transport plan eent has normalized column sums. Intuitively, while een reduces token dropping, still allows a variable number of experts applied to tokens. The procedure of this Sinkhorn Expert Choice router is summarized follows.
Sinkhorn Expert Choice router
Compute the token expert affinity matrix Hent (8) on batch of input tokens X.
2. Use Ient to allocate the dispatch tensor Dx through Algorithm 2
Use Moofmmax softmax XW) to allocate the combine lensor Cx through Algorithm
Note that, as in the Sinkhorn Token Choice, we only allocate the dispatch tensor Dx using Im while calculating the combine tensor using Cx using softmax
2.5 cparsity Expert Choice router
All MoË layers introduced so far use specific sorting heuristics to allocate routing tensors. The sorting is applied either to the expert axis, as in Algorithm 1, or the token axis, as in Algorithm 2, to create sparse tensors D. This is necessary because the token-en token-expert expert affinity matrix II is neither sparse nor constrained by the buffer capacity.
An alternative approach is to promote sparsity in II directly meet the bufler capacity constraint (Liu et al., 2023). This is more principled way create sparse allocation of experts. To this end, consider token-expert affinity matrix II Isparse that solves sparsity-constrained optimal transport problem:
II aaarre
arg
III softmax softmax(XW))
n
II 2 0
(9)
II = 1r.
subject to
II 1r = (T/E)1g
| S C for all r r [E].

---

Published in Transactions on Machine Learning Research (04/2024)
There are several key differences between the sparsity-constrainee optimal transport in (9) and the entropy- regularized optimal transport in (8). First, constrty- optimal transport in (9) uses quadratic regularization || instead of the entropy regularization (II, II) on the transportation plan II. This quadratic regularization, as shown in Blondel et al. (2018), preserves the sparsity in the transportation plan and allows for an elficient numerical solver. Second, the formulation (9) introduces an additional constraint to upper bound the number of nonzeros of each column of II by the buffer constraint C. This ensures that all experts use up to C' tokens. Third, in the first term of the sparsity constrained optimal transport in (9), evaluate the imner product of II with positive utility matrix softmax (X W) opposed to the raw matrix product XW in (8). The positive utility matrix encourages each column of II to have exactly C nonzeros: even though we only upper bound the number of nonzeros in each column of II by C in (9), a learned optimal transport plan Msparse will usually lake exactly C' nonzeros per column since it needs to maximize its inner product.
Finding an exact minimizer Ippare (9) can be intractable in general due the non-convexity of this constrained optimization problem. Nonetheless, an approximate II sparse can be obtained using semi-dual dual formulation; for more details refer Liu et al. (2023). The overall procedure of the sparsity-constraneet MoE layer is summarized as follows.
Sparsity-constrained router
1. Approximate the token expert affinity matrix Il sppaae in (9) on batch of input tokens X.
Use IIsarre to allocate the dispatch tensor Dx through Algorithm 2
3. Use Softmma softmax X W) to allocate the combine tensor Cx through Algorithm 2
2.6
Soft MoE
All MoE formulations we have introduced so far aim to find hard assignments between tokens and experts; these hard assignments are represented by the binary, {0, 1Jvalued dispatch tensors Dx in token choice allocation (Algorithm 1) and expert choice allocation (Algorithm 2). In these formulations, each expert can process either an entire token or none of it. In a different approach, Soft MoEs (Puigcerver el al., 2023) instead allow experts to process weighted combinations of tokens, offering more flexibility.
We now present Soft MoE using our unified formulation of MoE layers (2). One key feature of SoftMoE is that we define T as the token count in a single image, since tokens are linearly combined exclusively from the same input. This contrasts with other routers discussed, where T represented the number of tokens in batch of multiple inputs.
We let E DxCC be a learnable tensor, where each vector F. E TRP represents the features of the e-th buffer slot of the r-th expert. For input X E RTDD we compute the affinity scores between all tokens and the (r, c)-th slot with:
Zl.,r,c Xol.dER..
(10)
The Soft MoE approach then computes the routing tensors Dx and Cx by normalizing the affinity scores in Z E TxEc.
exp Z rcc Dx(,nd= En exp exp Z rr Cx(br, En.e exp Z(t,r',e)
exp c[
(11)
Here, the dispatch tensor Dx is normalized along the token axis: we dispatch fractions of every token to each slot, with the fractions summing up to Similarly, the combine tensor Cx is normalized along the slot dimensions c), making each MoE output token a convex combination of all slot outputs with combining weights summing up to 1. We summarize the forward pass of Soft. MoF follows.
8

---

Published in Transactions on Machine Learning Research (04/2024)
Soft MoE router
Compute the affinity tensor Z E RRTxExC for tokens and expert slots via (10).
2. Compute the routing tensors Dx and Cx in (11).
Note that the routing tensors Dx and Cx differ from previous approaches by not being binary or sparse. Nonetheless, Soft MoEs are computationally efficient. This because each expert only processes a fixed number of slots per input that significantly smaller than the sequence length. In addition, compared other methods, Soft MoFs have speed advantage they avoid sorting and optimal transport algorithms.
3 Experiments
We report comprchensive comparisons of vision MoE models with the routers we introduccd carlier: (1) Softmax Token Choice, (2) Sinkhorn Token Choice, (3) Softmax Expert Choice, (4) Sinkhorn Expert Choice, (5) Sparsity-constrained Expert Choice, and (6) Soft MoE router. We evaluate MoE models using both large- scale pre-training and few-shot adaptation experiments.
3.1 Models
As in earlier vision MoE work (Riquelme et al., 2021), we replace subset of the dense feedforward layers in vision transformer (Dosovitskiy et al., 2021) with MoE layers. Specifically, use the Every-2 variant in Riquelme et al. (2021) that places the MoEs every other layer. For the comparison to be comprehensive, we use models of different sizes with naming convention consistent to vision transformer: B(asc)32, B(ase) 16, and L(arge)16 16. The numbers 32 and 16 here refer to 32 x 32 and 16 x 16 patch sizes. Models that use a smaller patch size result more patches, and thus model B16 has a greater capacity and computational cost than model B32. We fix the total number of experts to be 32; that is, E = 32.
For each architeclure (B32, B16, and L16) and each router, we experiment with two configurations to experiment with different capacity C:
For Softmax Token Choice and Sinkhorn Token Choice routers that process each token with k experts, we experiment with & and In this way, the buffer capacity (the number of tokens an expert can process most in batch) of these variants is C round(k- T/E).
For Softmax Expert Choice, Sinkhorn Expert Choice, and prarsity-constranee variants, we control the buffer capacity C through a capacity factor which plays a role similar to k. The buffer capacity C' is defined through round(c- T/E). We experiment with or 2, which match the buffer capacity of
and & in the Token Choice cases.
For Soft MoE routers, set the number of slots per expert to be C' = round(c- T/E) with = or 2, just like in the Expert Choice case.
3.2 Upstream task results: Pre-training on JFT-300M
For the pre-training experiments, all models were trained the JFT-300M (Sun et al., 2017), which contains about 305 million training images and 50,000 validation images, organized in hierarchy of 18,291 different classes. To avoid overlap with the validation and test sets of JFT-300M, the images in the dataset were deduplicated, as done in Kolesnikov et al. (2020). Our main metric for JFTT-300 is the top-1 classification accuracy (Prec@1).
Accuracy comparison. Expert Choice variants Choice models, those using a Sinkhorn-basee token-expert affinity matrix (++ perform better than those with Softmax-based matrix (x), aligning with previous findings in language MoE. research (Clark al., 2022).
Our pretraining results, shown in Figure In terms accuracy (y-axis), the generally surpass Token Choice variants (x, 4) accuracy. Among Token
9

---

Published in Transactions on Machine Learning Research (04/2024)
Softmax Token Choice
Softmax Expert Choice Sinkhorn Expert Choice
aparsity-constraane Expert Choice SoftMoE
Sinkhorn Token Choice
B32
B16
L16
0.51
0.590
0.45
JFT-300M Precision-1
0.50
0.585
0.44
0.580
0.49
0.575
0.43
0.48
0.570
0.42
0.47
0.565
0.46
0.560
40
50
60
150
200
250
800
1000
1200
1400
Total Training TPUv3-days
Total Training TPUv3-days
Total Training TPUv3-days
Figure Comparison of training time and performance the JFT300M dataset for image classification. The marker size represents the router's capacity, with smaller and larger sizes indicating lower and higher capacities.
However, in Expert Choice models, the choices of token-expert affinity matrix Softmax Sinkhorn or sparsity-constrained (4) less significant. impact on performance. We hypothesize that Sinkhorn transform effectively balances the expert usage in the Token Choices case, thus improving the performance relative to the Softmax Token Choice; but such balance is longer needed for the Expert Choice cases, and therefore its impact is small. Compared to Token Choice and Expert Choice routers, SoftMoF. routers (*) outperform both by noticcable margin.
There one instance of SoftMoE for L-16 at c = 2 not included in Figure 1 (one * missing from the right-most panel of Figure 1). This is due to training instability that occurred at that specific instance. This might relate to using only experts (E 32) for consistent comparisons across all routers. Notably, SoftMoE performs best with larger number of experts (e-g., E 2048) and few slots per expert, as shown in Puigcerver et al. (2023, Figure 6). Despite not being the most optimal setup, SoftMoE surpasses other configurations in Figure 1, where instability was not an issue.
Training cost comparison. Considering the computational cost, the sparsity-connstraind router (A) is more expensive compared to others, primarily because of its sorting and sparse projection steps (Liu et al., 2023). In contrast, SoftMoE routers have the best compute-accuracy tradeoff.
For comprehensive numerical comparison of routers using the JFT300M dataset, refer to Tables 1, 2, and in Appendix B.
3.3 Downstream task results: Few-shot transfer on ImageNet-1k
To assess how well pre-trained MoE models adapt to new tasks, we conducted few-shot adaptation experi- ments using the ImageNet-1k dataset (Deng et. al., 2009). Tn these experiments, we used 10 image samples per class from ImageNet-Ik. The pre-trained model extracts fixed feature embedding for cach image, which is then used train linear regression model. This linear model maps the extracted features to the one-hot encoded target labels. This procedure is in line with the 10-shot evaluation procedure described by Dosovitskiy et al. (2021); Riquelme et al. (2021).
Our few-shot results, as shown in Figure 2, confirm findings from the pretraining phase. Overall, the MoE. router (*) consistently outperformed the Expert Choice routers (,,,). A), which themselves surpassed the performance of the Token Choice routers (x, +). In most scenarios, token choice routers marked by + that use a Sinkhorn-derived token-expert affinity matrix tend to outperform those using a Softmax Doftmax-based affinity
matrix (x). However, the performance differences between the Expert Choice routers
), are relatively
10

---

Published in Transactions on Machine Learning Research (04/2024)
minor.
This emphasizes that the choice of routing tensor allocation algorithm is more crucial than the specific method of parameterixing the token-expert affinity matrix.
Softmax Token Choice
Softmax Expert Choice Sinkhom Expert Choice
Sparsity-constrained Expert Choice SoftMoE
Sinkhorn Token Choice
B32
B16
L16
0.68
0.7875
0.73
0.67
0.7850
ThageNet-10shol
0.7825
0.66
0.72
0.7800
0.65
0.7775
0.64
0.7750
0.70
0.63
0.7725
40
250
Total Training TPUv3-days
Total Training TPUv3-days
Total Training TPUv3-days
Figure 2: Comparison of training time and performance in 10-shot Transfer Task on the ImageNet-lk Dataset. The marker size represents the router's capacity, with smaller and larger sizes indicating lower and higher capacities.
For a comprehensive numerical comparison of routers few-shot adapted on the ImageNet- dataset, refer to Tables 1, 2, and 3 in Appendix B.
3.4 Analysis: the usage of Softmax combine tensors in Sinkhorn routers.
In Section 2.2, we mentioned useful trick the Sinkhorn Token Choice router: using the Sinkhorn-baas matrix Nent to allocate the dispatch tensor Dx and the Softmax-based matrix Tostuma for the combine tensor Cx. This approach, also applied in Sinkhorn Expert Choice and Sparsity-constrammed Expert Choice routers, primarily aims to boost processing speed by bypassing the optimal transport algorithun in the backward pass. Perhaps surprisingly, it also enhances performance. Figure 3 showcases a Sinkhorn Token Choice case study. We contrast two methods: one assigns Dx with Ient and Cx with Loftmax (labeled "with softmax combine"), while the other uses Ient for both tensors (labeled "without softmax combine").. The "with softmax combine" variant achieves approximately 1% higher accuracy on the JFT300M dataset.
4 Related work
MoE models have been increasingly used language and vision domains. We briefly review the application of them in this section. We refer the interested readers Fedus al. (2022a) for a comprehensive review,
MoEs for language. MoFs have been successfully applied to language modeling and machine translation. One of the earliest successes of sparse MoEs language modeling and machine translation demonstrated by Shazeer et al. (2017). They insert MoE layers between LSTM layers (Hochreiter & Schmidhuber, 1997) to increase the model capacity while maintaining high computational efficiency. This approach achieved state-of-the-art at that time, with lower computational cost than bascline models. Sparse MoEs have further advanced language modeling when combined with Transformers (Vaswani et al., 2017). The GShard (Lepikhin et al., 2021) and Switch Transformers (Fedus et al., 2022b) are among the earliest works that replace food-forward layers Transformers with sparse MoE layers. addition, rescarch efforts have been made analyze and simplify the routers language MoEs, such deterministic routers (Lewis et al., 2021) and routers based reinforcement. learning and optimal transport (Kool et al., 2021; Clark et al., 2022).
11

---

Published in Transactions on Machine Learning Research (04/2024)
k
0.44
with softmax combine
without softmax combine
0.42
Validation accuracy
0.40
0.38
0.36
0.34
0.32
350000
400000 Training iters
450000
500000
350000
400000
450000
500000
Training iters
Figure Assessing the impact of using Softmax-based combine tensors in Sinkhorn Token Choice routers. The number of selected expets are (left panel) and (right panel). Both routers are used in B32 architecture. The axis shows the training iteration number, while the axis shows the validation accuracy the JFT300M dataset.
MoEs for vision. the realm of vision, early work on MoFs (Figen et al., 2013; Ahmed al., 2016; Gross et al., 2017; Abbas & Andreopoulos, 2020; Wang al., 2020; Pavlitskaya al., 2020; Yang al., 2019) were mostly based convolutional neural networks and were applied specific tasks. The tremendous success of Vision Transformers (ViTs) (Dosovitskiy et al., 2021) has motivated research in creating sparse MoEs based ViTs. Riquelme et al. (2021) introduced V-MoE, vision MoE architecture based on ViTs, and showed that it outperforms dense VIT counterparts at the same computational cost. Beyond accuracy, vision MoEs have been shown to offer better robustness against adversarial atlacks (Puigcerver et al., 2022). Furthermore, while (Riquelme et al., 2021; Puigcerver et al., 2022) focus on image classification problems, Wu et al. (2022) show the capability of vision MoEs to solve high-rcsolution vision tasks such as segmentation and detection. Li et al. (2022) proposed variant of vision MoE with enhanced domaaingennrallizttion abilities. Puigcerver et al. (2023) formulated soft MoEs that fully differentiable while being efficient as sparse MoEs.
MoEs for multimodal and multitask learning. Mustafa et al. (2022) presented the first multimodal sparse MoF. model, called Language-Image MoE (LIMoF). LIMoE processes both images and text in modality-agnostic fashion, to align image and text embeddings via contrastive learning. LIMoE matches the performance of state-of-the-art dense models and surpasses dense basclines at cquivalent computational costs. MoEs have also been used in multitask settings (Ma et al., 2018; Chen et al., 2023).
Other aspocts of MoEs. There are a few research studies MoEs from perspectives that are not tied to specific data modalities. For example, Hwang et al. (2023) proposed an efficient pipeline that optimizes the implementation
MoF layers on GPUs. From a theoretical perspective, Chen et al. (2022) studies why simple sparse MoEs do not collapse into a single model. Komatsuzaki et al. (2023) create sparse MoEs from pro-oxisting dense models as way to reuse the sunk cost for training dense modcls. This simple yet powerful approach can be generally applied to tasks in different modalities.
5
Conclusion
MoEs offer a promising solution to largo-scale machine learning applications. Our paper presents the first comprehensive study transformer-base sparse and soft MoEs in computer vision tasks, achieved through unified MoF layer formulation with routers. We show that the strong performance of many language
12

---

Published in Transactions on Machine Learning Research (04/2024)
MoEs carries over to vision. We found routing tensor allocation crucial for sparse MoEs, more than other factors. Notably, soft MoEs outperfor all sparse alternatives tested.
Research on efficient sparse MoEs for vision problems at scale is still in its early stages. As MoEs can handle large amounts of data while keeping computational costs low, is expected that they will become increasingly important for data-rich tasks in the future. Understanding how different MoE models perform in these tasks crucial. Our paper takes a step in this direction and opens new opportunities for further study of MoEs at scale.
References
Alhabib Abbas and Yiannis Andreopoulos. Biased mixtures of experts: Enabling computer vision inference under data transfer limitations. IEEE Transactions on Image Processing, 29:7656 7666,,7677-7 2020. (Cited on page 12.)
Karim Ahmed, Mohammad Haris Baig, and Lorenzo Torresani. Network of experts for large-scale image categorization. In Proceedings of the European Conference on Compuler Vision (ECCV), pp. 516 532. Springer, 2016. (Cited on page 12.)
James Urquhart Allingham, Florian Wenzel, Zelda. E. Mariet, Basil Mustafa, Joan Puigcerver, Neil Houlsby, Ghassen Jerfel, Vincent Fortuin, Balaji Lakshninarayanan, Jasper Snoek, et al. Sparse moes meet efficient ensembles. Transactions on Machine Learning Research, 2022. (Cited page 2.)
Emmanuel Bengio, Pierre-Lue Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. Tn Proceedings of the International Conference on Learning Representations (ICLR), 2016. (Cited on page 1.)
Mathieu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. In Proceedings of the International Conference on Artificial Intelligence and Statistics (AISTATS), pp. 880 889. PMLR, 2018. (Cited on page 8.)
Zitian Chen, Yikang Shen, Mingyu Ding, Zhenfang Chen, Hengshuang Zhao, Erik Learned-Mille and Chuang Gan. Mod-sqqad: Designing mixtures of experts modular multi-task learners. In Proceedings of the IEEE Conference Computer Vision and Pottern Recognition (CVPR), 2023. (Cited page 12.)
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards undorstanding the mixturc-of- experts layer in deep learning. In Proceedings of the Advances in Neurul Information Processing Systems (NeurIPS), 2022. (Cited on page 12.)
Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In Proceedings of the International Conference on Machine Learning (ICML), 2022. (Cited on pages 1, 2, 5, 6, 9, and 11.)
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Proceedings of the Advances in Newral Information Processing Systems (NeurIPS), volume 26, 2013. (Cited on page 5.)
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognation (CVPR), pp. 248 255, 2009. (Cited on page 10.)
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In Proceedings of the International Conference on Learning Representations (ICLR), 2021. (Cited on pages 2, 9, 10, and 12.)
David Eigen, Mare' Aurelio Ranzato, and Ilya Sutskever. Learning factored representations in a deep mixture of experts. Xiu preprinl ar 1812 2013. (Cited on page 12.)
13

---

Published in Transactions on Machine Learning Research (04/2024)
William Fedus, Jeff Dean, and Barret Zoph. A review of sparse expert models in deep learning. preprint or Xrv:22209.0 01667, 2022a. (Cited page 11.)
arXiv
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1 -9, 2022b. (Cited pages 2, and 11.)
Sam Gross, Marc' Aurelio Ranzato, and Arthur Sxlam. Hard mixtures of experts for large scale weakly supervised vision. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognation (CVPR), 2017. (Cited on page 12.)
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Newral Computation, pp. 1735-1780, 1997. (Cited on page 11.)
Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, et al. Tutel: Adaptive mixture-of-experts at scale. In Proceedings of Conference on Machine Learning and Systems (MLSys), 2023. (Cited on page 12.)
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (BiT): general visual representation learning. In Proceedings of the European Con- ference on Computer Vision (ECCV), pp. 201-507. Springer, 2020. (Cited on page 9.)
Aran Komatsuzaki, Joan Puigcerver, James Lee- Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training of from dense chockpoints. In Proceedings of the International Conference on Learning Representations (ICLR), 2023.
(Cited on page 12.)
Wouter Kool, Chris J. Maddison, and Andriy Mnih. Unbiased gradient estimation with balanced assignments for mixtures of experts. NeurIPS Can't Believe It's Not Better (ICBINB) Workshop, 2021. (Cited on pages 1, 5, and 11.)
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations, 2021. (Cited on pages
and 11.)
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Basc layers: Simplifying training of large, sparse models. In Proceedings of the 38th International Conference on Machine Learning (ICML), pp. 6265-6220 2021. (Cited on page 11.)
Bo Li, Yifei Shen, Jingkang Yang, Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of-experts are domain generalizable learners. In NewrIPS 2022 Workshop on Distribution Shifts: Connecting Methods and Applications, 2022. (Cited page 12.)
Tianlin Lüu, Joan Puigcerver, and Mathieu Blondel. Sparsity-constrainee optimal transport. In Proceedings of International Conference on Learning Representation (ICLR), 2023. (Cited on pages 1, 7, 8, and 10.)
Jiaqi Ma, Zhe Zhao, Xinyang Yi, Jilin Chen, Lichan Hong, and Ed H. Chi. Modeling task relationships in multi-task learning with multi-gate mixture--f-exxperss In Proceedings of the International Conference on Knowledge Discovery and Dala Mining (KDD). Association for Computing Machinery, 2018. (Cited on
page 12.)
Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Multinodal con- trastive learning with LIMoE: the language-image mixture of experts. In Proceedings of the 36th Annual Conference on Neural Information Processing Systems (NeurIPS), 2022. (Cited on pages and 12.)
Svetlana Pavlitskaya, Christian Hubschneider, Michael Weber, Ruby Moritz, Fabian Huger, Peter Schlicht, and Marius Zollner. Using mixture of export models to gain insights into semantic sogmentation. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition Workshops, pp. 342 343, 2020. (Cited on page 12.)
14

---

Published in Transactions on Machine Learning Research (04/2024)
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport: With Applications to Data Science. Foundations and Trends in Machine Learning, (1(-6)):355-000 2019. (Cited on page 5.)
Joan Puigcerver, Rodolphe Jenatton, Carlos Riquelme, Pranjal Awasthi, and Srinadh Bhojanapalli. On the adversarial robustness of mixture of experts. In Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), 2022. (Cited on page 12.)
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. arXiv preprint ar Xio:2308 00951, 2023. (Cited on pages 1, 8, 10, and 12.)
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. In Proceedings of Advances in Newral Information Processing Systems (NeurIPS), 2021. (Cited on pages 1, 2, 3, 9, 10, and 12.)
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason E Weston. Hash layers for large sparse models. In Proceedings of Advances Neural Information Processing Systems (NeurIPS), 2021. (Cited on page 1.)
Michael Eli Sander, Joan Puigcerver, Josip Djolonga, Gabriel Peyré, and Mathieu Blondel. Fast, differentiable and sparse lop-k: convex analysis perspective. In Proceedings of the International Conference on Machine Learning (ICML), 2023. (Cited on page 1.)
Noam Shazcer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparscly-gated mixture-of-cxperrs. layer. In Proceedings of the International Conference on Leurning Representations (ICI.R, 2017. (Cited on pages 1, 2, and 11.)
Richard Sinkhorn. A Relationship Between Arbitrary Positive Matrices and Doubly Stochastic Matrices. The Annals of Mathematical Statistics, 35(2):876 879, 1964. (Cited on page 5.)
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable effectiveness of data in deep learning era. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), pp. 843-852, 2017. (Cited on page 9.)
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of Advances in Neural Information Pro- cessing Systems (NeurIPS), volume 30, 2017. (Cited on pages and 11.)
Xin Wang, Fisher Yu, Lisa Dunlap, Yi-An Ma, Ruth Wang, Azalia Mirhoseini, Trevor Darrell, and Joseph E Gonzalez. Deep mixture of experts via shallow embedding. In Proceedings of the Conference on Uncertainty in Artificial Intelligence (UAI), 2020. (Cited on page 12.)
Lemeng WU, Mengehen Liu, Yinpeng Chen, Dongdong Chen, Xiyang Dai, and Lu Yuan. Residual mixture of experts. arXiv preprial XrXiv:2204.0000 096666, 2022. (Cited on page 12.)
Brandon Yang, Gabriel Bender, Quoe V Le, and Jiquan Ngiam. CondConv: Conditionally parameterized convolutions for efficient inference. Proceedings of the Advances in Neural Information Processing Systems (NeurIPS). 2019. (Cited page 12.)
Zhao You, Shulin Feng, Dan Su, and Dong Yu. SpcechMoE: Scaling to large acoustic models with dynamic routing mixture of experts. arXiv preprint arXiv:2105. 03000000 2021. (Cited page 1.)
Seniha Esen Yuksel, Joseph N Wilson, and Paul Gader. Twenty years of mixture of experts. IEEE Transactions on Neural Networks and Learning Systems, 23(8):1177 1193, 2012. (Cited on page 1.)
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Y Zhao, Andrew M. Dai, Zhifeng Chen, Quoc V Le, and James Laudon. Mixture-of-experus with expert choice routing. In Proceedings of the Advonces in Neural Information Processing Systems (NeurIPS), 2022. (Cited pages 6 and 7.)
15

---

Published in Transactions on Machine Learning Research (04/2024)
A Auxiliary losses for routers
Importance loss. The importance loss penalizes the variability of cach expert's overall gating weights on a minibatch of tokens. Let X E R be a minibatch of T tokens, where each token is row of X. We let
Nw x softmax (eoftmax(XW) RxxE
be a matrix token-to-cxperts routing probabilities, where the softmax operator applies to each row of XW. Summing over rows, the vector
Tw,x- = W lr E R
describe each export's token-roceiving probabilitics summed on that minibatch of tokens. Since we want cach expert to receive roughly the same amount of tokens, the token-receiving probabilities of experts should not vary considerably, To minimize the variation, the importance loss is defined
E 1=11
Immp X, W) = CV
where CV denotes the empirical coefficient variation.
Load loss. The load loss of vision MoE has the form
where 8 is valued random variable with Gaussian distribution N/0, }.)
load,(æ, W. = Ps5 (War) [i] + > max (W x + e)) k-th
(12)
Ps max k-th
WW c) (Wz)
(13)
= + (w
max (Wx+c))
(14)
k-th
where the cumulative distribution function of the Gaussian distribution A(0, +
Summing the loads (12) over all tokens yields the expected numbers of tokens received for each expert. We define the load loss as the coefficient of variation of such expected numbers of received tokens across experts:
Liond(X; (X; W) = CV (V(Iooad((x; )² with load,(X, W) = Σ (x W W,e ()). t=1
(15)
B Detailed experimental results
16

---

Published in Transactions on Machine Learning Research (04/2024)
Accuracy
Training TPUv3-days
Softmax (k 1)
41.23 %
34.31
Sinkhorn (k 1)
41.63 %
34.73
Softmax Expert Choice (Clactor =
43.83 %
33.73
Sinkhorn Expert Choice (Cfactor =
43.65 %
34.40
Sparsity constrained (Ciactor = 1)
43.43 %
56.88
SoftMoE
45.23 %
36.95
Softmax Token Choice (k 2)
43.47 %
41.46
Sinkhorn Token Choice (A: 2)
44.26 %
41.88
Softmax Expert Choice (Clactor =
44.40 %
40.60
Sinkhorn Expert Choice (Cfactor =
44.30
41.21
Sparsity constrained (Ciactor = 2)
44.30
63.13
SoftMoE
45.28 %
43.59
Table 1: Comparing routers in B32 architecture using the JFT dataset.
Accuracy
Training TPUv3-days
Softmax Token Choice (A: 1)
46.35 %
122.12
Sinkhorn Token Choice (A: 1)
47.37 %
123.04
Softmax Expert Choice (Ctactor = 1)
48.64
118.31
Sinkhorn Expert Choice (Cfactor = 1)
48.64
119.80
Sparsity constrained (Cractor = 1)
48.67
231.66
SoftMoE
50.00 %
116.31
Softmax Token Choice (k: = 2)
48.86 %
153.96
Sinkhorn Token Choice (k: 2)
49.80
157.93
Sofumax Expert Choice (Ctactor = 2)
50.23
153.42
Sinkhorn Expert Choice (Ciactor = 2)
50.52 %
155.03
Sparsity constrained (Ofactor = 2)
50.29 %
279.89
SoftMoE
50.53 %
148.47
Table 2: Comparing routers in B16 architecture using the JFT dataset.
Accuracy
Training TPUv3-days
Softmax Token Choice (k: = 1)
56.07
818.80
Sinkhorn Token Choice (A: 1)
56.19
827.70
Softmax Expert Choice (Ciactor =
57.51 %
805.60
Sinkhorn Expert Choice (Cfactor =
57.84 %
818.95
Sparsity constrained (Cractor 1)
SoftMoE
58.51%
790.57
Softmax Token Choice (k 2)
58.05
1023.92
Sinkhorn Token Choice (A: 2)
57.69 %
1033.43
Softmax Expert Choice (Ctactor = 2)
58.43 %
1014.71
Sinkhorn Expert Choice (Cfactor = 2)
58.88 %
1028.06
Sparsity constrained (Ctactor 2)
58.26 %
1537.41
Soft
Table Comparing routers in L16 architecture using the JFT dataset.
17

---

Published in Transactions on Machine Learning Research (04/2024)
Accuracy
Softmax Token Choice (k: 1)
62.77 %
Sinkhorn Token Choice (k 1)
63.42 %
Softmax Expert Choice (Ciactor = 1)
64.99 %
Sinkhorn Expert Choice (Ciactor = 1)
65.34 %
Sparsity constrained (Cractor =
64.50 %
SoftMoE
67.37 %
moftmax Token Choice (A = 2)
65.91 %
Sinkhorn Token Choice (k = 2)
65.87 %
Softmax Expert Choice (Ciactor = 2)
66.38 %
Sinkhorn Expert Choice (Clactor = 2)
66.52 %
Sparsity constrained (Cfactor = 2)
65.75 %
SoftMoE
67.85 %
Table 4: Comparing routers in B32 architecture on the ImageNet 10-shot task.
Table 5: B16 on ImageNet 10-shot
Accuracy
Softmax Token Choice (k = 1)
69.62 %
Sinkhorn Token Choice (k: = 1)
70.04 %
Softmax Expert Choice (Cractor = 1)
71.74 %
Sinkhorn Expert Choice (Cfactor = 1)
71.95 %
Sparsity constrained (Cractor = 1)
71.28 %
SoftMoE
72.84 %
Softmax Token Choice (k = 2)
72.12 %
Sinkhorn Token Choice (k: = 2)
72.26 %
Softmax Expert Choice (Cinetor = 2)
72.68 %
Sinkhorn Expert Choice (Cractor =
72.74 %
Sparsity constrained (Cfactor = -)) SoftMoE
72.76 %
73.26 %
Table 6: Comparing routers in B16 architecture on the ImageNet 10-shot task.
Table 7: L16 on ImageNet 10-shot
Accuracy
Softmax Token Choice (A: = 1)
77.42 %
Sinkhorn Token Choice (k: = 1)
77.18 %
Softmax Expert Choice (Cractor = 1)
77.91 %
Sinkhorn Expert Choice (Cractor 1)
78.08 %
Sparsity constrained (Ctactor = 1)
SoftMoE
78.72 %
Softmax Token Choice (A: = 2)
78.00 %
Sinkhorn Token Choice (k: = 2)
78.10 %
Softmax Expert Choice (Cractor =)
78.38 %
Sinkhorn Expert Choice (Cractor
78.19 %
Sparsity constrained (Cractor = 2)
78.24 %
SoftMoL
Table 8: Comparing routers in L16 architecture on the ImageNet 10-shot task.
18


Reviews:
[{"review_id": "W14pPTVUiO", "reviewer": "Reviewer_rgad", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n* The problem studied is of significant importance to the research community and practitioners, as MoE formulation allows one to have a model that is both expressive and efficient to inference. This work investigates various design options of MoE and compares them in a unified setup, reliably showing superiority of one approach over another. \n* Dispatch-combine tensor formulation is novel and incorporates the known variations of MoE layers as well as allows for a design of new one.\n* The experimental study is quite extensive, involving models of different size and hyperparameter choices. \n\nWeaknesses\n* The comparison between token, expert and soft MoE is not novel. In the referenced work [1] Figure 3 contains detailed comparison between different expert families at different training budgets. The contribution of this work is more fine-grained-comparison involving the algorithm of router assignment, that is important, but somehow incremental. The choice of router softmax/sinkhorn seems to matter much less compared to the choice of Token/Expert/Soft MoE.\n* The experimental evaluation involves only JFT upstream accuracy and ImageNet-1k few-shots. I would suggest adding full-model finetuning on ImageNet-1k and evaluation on a couple more datasets from VTAB benchmark.\n\n---\n[1] Puigcerver, Joan, et al. \"From sparse to soft mixtures of experts.\" arXiv preprint arXiv:2308.00951 (2023).\n\n[2] https://github.com/google-research/task_adaptation", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes a unified formulation of Mixture-of-Experts (MoE), that incorporates several known options from the literature - Token Router, Expert Router, Soft MoE, as well as proposes a new option of router. All mentioned options are compared with each other on large-scale computer task via measuring the performance on upstream and downstream tasks. It is shown that Soft MoE consistently outperforms alternatives.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Proposed changes and additions\n\n* If I understood correctly, all models are evaluated with a fixed number of training steps and there are 2 options of capacity for each MoE formulation. I would suggest comparing the performance of the models at different training budgets, as it could be the case that one option may converge faster than another and the ranking may be dependent on the amount of training. In addition, would be interesting to verify, how does the performance of each approach depend on the number of experts, specifically, whether the gap between MoE formulations increases or decreases with the number of experts. \n* Would be also nice to show comparison between the formulations of MoE in terms of **inference** latency and the overheads of routers. \n\n*Minor*\n\nI would suggest coloring each of the formulations - Token, Expert, SoftMoE with a fixed color and vary the marker shapes for Softmax, Sinkhorn on Figure [1] for the figure to be more illustrative.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XQHVoowwCX", "reviewer": "Reviewer_fYVm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n1. Unified Formulation of MoE: The introduction of a unified formulation for Mixture of Experts (MoE) models is a significant contribution to the field. This approach not only simplifies MoEs by providing a common framework to understand different architectures but also facilitates a more systematic study and comparison of various routing strategies. By harmonizing the underlying principles of MoE models, the authors enable researchers and practitioners to more easily identify the core components that drive performance in these complex systems, leading to more intuitive developments and enhancements.\n\n2. Comprehensive Experiments: The exhaustive experiments conducted using the JFT300M dataset serve as a robust foundation for evaluating the performance of the proposed unified MoE model. These experiments are particularly valuable due to the scale and diversity of the JFT300M dataset, which ensures that the findings are broadly applicable. Furthermore, the transfer results on ImageNet provide critical insights into the generalizability of the model across different domains and datasets. This extensive empirical evaluation underscores the effectiveness of the proposed model and enhances the credibility of the research.\n\n# Weaknesses\nReproducibility and Comparison with Original MoE Models: A notable concern with the proposed unified MoE model is the need for a direct comparison with the results reported in original MoE papers (where routing strategies unified in this framework that were originally proposed). It is crucial to establish that the unified model can reproduce or at least approximate the performances of the original models it seeks to encompass. Such a comparison would serve as a validation of the unified model's effectiveness and ensure that the implementation accurately reflects the advancements it claims to integrate. This step is essential for the research community to assess the model's reliability and for ensuring that the results are genuinely falsifiable, which is a cornerstone of scientific progress.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents an empirical study of routers in Mixture of Experts (MoE) models for computer vision tasks, introducing a unified MoE formulation that integrates various MoE layers through parametric routing tensors. This formulation enables systematic comparison and analysis of sparse and soft MoE models. The study includes experiments with six different routers on the JFT-300M dataset for pre-training and ImageNet-1k for downstream tasks, finding that soft MoEs generally outperform sparse ones within a fixed compute budget. Expert Choice routers typically surpass Token Choice routers in performance. The research highlights the significance of routers in vision MoE models and suggests that many routers developed for language modeling are also effective in vision tasks. The broader impact of this study is aimed at advancing MoE research by proposing a standardized testbed for future explorations, urging the public release of the code and benchmark results for community use.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Why are some fonts in orange (e.g. on top of page 2)?", "broader_impact_concerns": "Accessibility of Code and Benchmark Results: The broader impact of this research could be significantly enhanced by making the code and benchmark results publicly available. Providing open access to these resources would democratize the advancements made by this work, allowing a wider range of researchers to engage in MoE research. A public, well-documented codebase would serve as a fair and standardized testbed, facilitating reproducibility and encouraging innovation. This transparency is vital for fostering a collaborative environment where the community can collectively push the boundaries of what is possible with MoE models.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ELiwcgAvei", "reviewer": "Reviewer_2yy7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**strengths:**\n\nThe paper is well-written and I am glad to read it. The mathematical derivations of this paper are clear and a unanimous mathematical framework that encompassed different routing of MoEs were proposed. The authors offered addressed weaknesses of sparse MoEs and put forward novel routing schemes to compensate for the uneven distributions of experts during training while maintaining efficiency.\n\n**Weaknesses:**\n\nThe paper mainly evaluates different MoE variants on the few-shot transfer on ImageNet-1k, which I find to be insufficient as the paper claimed to empirically study MoEs’ performance in vision. Further experiments on other vision tasks (e.g., dense prediciton or visual multi-task benchmark) should be considered to amplify the results.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed a unified Mixture-of-Experts (MoE) formulation, where they examined 6 different MoE routers under an unanimous mathematical framework. This paper is the first to comprehensive study the transformer-based sparse and soft MoEs in computer vision, and provides clear derivations on the routing of MoEs. Through their experiments, the authors empirically showed that the success of language MoEs can be transferred to vision and that soft MoEs generally perform better than sparse MoEs, which offers practical insights as to how MoEs should be routed in vision tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It is better to provide the experiments on more vision tasks  (e.g., dense prediciton or visual multi-task benchmark) to provide more strong evidence to support the calims that  the success of language MoEs can be transferred to vision.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "tG8VdCHWfm", "reviewer": "Reviewer_VVoS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, these contributions seem small, but may still be of interest.  Measuring on another dataset or task substantially different from JFT would greatly strengthen this paper, as most of the previous works on which this paper is based focus on JFT.\n\nTo me, phrasing all of these different experts in the unified framework is useful, but the framework isn't too far away from most MoE descriptions in the first place.  Still, it's nice to see how they all relate surveyed in one place, in the context of ViTs.  I'd note, though, that the dispatch D, combine C is the same as in Puigcerver et al 2023, and that this work rephrases the other MoEs in that framework.  The framework itself seems to be drawn from this prior work.\n\nMost of the main evaluation results are also already examined in Puigcerver et al 2023 \"From Sparse to Soft Mixtures of Experts\" --- Figs 3 and 22 in Puigcerver show exactly the same measurements comparing Token choice, Expert choice and Soft MoE that this paper describes, but on JFT-4B instead of 300M.  However, this paper fills in additional points of comparison, also looking at Sinkhorn routing in both Token and Expert choice contexts (with this paper being the first comparison using Expert choice context, according to the authors).\n\nAdditional questions and typos:\n\n* p.7: is the sign wrong for the $||\\Pi||_F^2$ ?  Combined with a constraint that the rows (or cols) sum to 1, a sparser vector maximizes this norm, not minimizes it, as is now mentioned in the text.  Or, it could be that the final constraint $||\\Pi[:,r]||_0 \\le C$ is what imposes sparsity (use of <= C entries), and the sign on the 2 norm should minimize, not to encourage sparsity, but the opposite, to encourage full, balanced use of all C entries?  Could you please check this description?\n\n* p.9 second bullet: should this read \"expert choice\"?  both bullet points now say \"token choice\".\n\n* p.10: text says black x for sparsity-constrained router, should be yellow triangle", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper surveys and measures several current MoE variations in the context of image classification using MoE MLPs insert into ViT.  The factorial combination of (Softmax, Sinkhorn) x (Token choice, Expert choice) is examined, as well as the newer soft MoE.  Each is summarized using the same formal framework, and measured on both JFT-300M classification and downstream ImageNet 10-shot classification.  The results corroborate recent findings that soft MoE outperforms Expert choice, and Expert choice outperforms Token choice, for these settings.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* more explicitly describe the unified formulation (esp D, C tensors) as drawn from Puigcerver et al, including in descriptions and the first contribution summary points\n\n* edit descriptions for the second contribution point to be more concrete and suitable to the contributions; at the moment it is a little vague.  to me these seem to be the measurement of Sinkhorn to fill out the factorial combination of (token,expert) x (softmax,sinkhorn), and a corroboration of prior results using JFT-300M, but there may be others\n\n* fix smaller issues mentioned above in review above\n\n* Additional measurements on another dataset or task substantially different from JFT or ImageNet would strengthen this paper considerably, particularly since Puigcerver focuses on JFT.", "broader_impact_concerns": "none", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a systematic framework for MOE, which includes various strategies to construct a powerful MOE model. All the reviewers approve of the manuscript's contribution and its clear presentation. During the rebuttal phase, the authors also addressed the reviewers' queries regarding experimental details and other aspects, supporting the acceptance of the paper.

Overall, this manuscript fits the TMLR acceptance criteria and I recommend accepting this paper.
Decision: Accept (poster)