From Sparse to Soft Mixtures of Experts
Authors: Joan Puigcerver
Carlos Riquelme Ruiz
Basil Mustafa
Neil Houlsby
Venue: ICLR 2024 spotlight (2024)
Keywords: transformers
mixtures of experts
computer vision
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: Soft MoEs outperform Transformers and Sparse MoEs in computer vision.

Abstract: Sparse mixture of expert architectures (MoEs) scale model capacity without significant increases in training or inference costs.
Despite their success, MoEs suffer from a number of issues: training instability, token dropping, inability to scale the number of experts, or ineffective finetuning.
In this work, we propose Soft MoE, a fully-differentiable sparse Transformer that addresses these challenges, while maintaining the benefits of MoEs.
Soft MoE performs an implicit soft assignment by passing different weighted combinations of all input tokens to each expert.
As in other MoEs, experts in Soft MoE only process a subset of the (combined) tokens, enabling larger model capacity (and performance) at lower inference cost.
In the context of visual recognition, Soft MoE greatly outperforms dense Transformers (ViTs) and popular MoEs (Tokens Choice and Experts Choice).
Soft MoE scales well: Soft MoE Huge/14 with 128 experts in 16 MoE layers has over 40x more parameters than ViT Huge/14, with only 2% increased inference time, and substantially better quality.

Published as a conference paper at ICLR 2024
FROM SPARSE TO SOFT MIXTURES OF EXPERTS
Joan Puigcerver' Google DeepMind
Carlos Riquelme" Google DeepMind
Basil Mustafa Google DeepMind
Neil Houlsby
Google DeepMind
ABSTRACT
Sparse mixture of expert architectures (MoEs) scale model capacity without signif- icant increases in training or inference costs. Despite their success, MoEs suffer from a number of issues: training instability, token dropping, inability to scale the number of experts, or ineffective finetuning. In this work, we propose Soft MoE, a fully-differentiable sparse Transformer that addresses these challenges, while maintaining the benefits of MoEs. Soft MoE performs by passing different weighted combinations of all input tokens to each expert. As in other MoEs, experts in Soft MoE only process a subset of the (combined) tokens, enabling larger model capacity (and performance) at lower inference cost. In the context of visual recognition, Soft MoE greatly outperforms dense Transformers (ViTs) and popular MoEs (Tokens Choice and Experts Choice). Furthermore, Soft MoE scales well: Soft MoE Huge/14 14 with 128 experts over 40x more parameters than VIT Huge/14, with only 2% increased inference time, and substantially better quality.
implicit soft assignment
MoE layers has
INTRODUCTION
Larger Transformers improve performance at increased computational cost. Recent studies suggest that model size and training data must be scaled together to optimally use any given training compute budget (Kaplan et al., 2020; Hoffmann et al., 2022; Zhai et al., 2022a). A promising alternative that allows to scale models in size without paying their full computational cost is sparse mixtures of experts (MoEs). Recently. number of successful approaches have proposed ways to sparsely activate token paths across the network in language (Lepikhin al., 2020; Fedus et al., 2022), vision (Riquelme et al., 2021), and multimodal models (Mustafa et al., 2022).
Sparse MoE Transformers involve a discrete optimization problem to decide which modules should be applied to each token. These modules commonly referred to as experts and usually ML.ss Many techniques have been devised to find good token-to-expert matches: linear programs (Lewis al., 2021), reinforcement learning (Bengio et al., 2015), deterministic fixed rules (Roller al., 2021), optimal transport (Liu et al., 2022), greedy top-k experts per token (Shazeer et al., 2017), or greedy top-k tokens per expert (Zhou et al., 2022). Often, heuristic auxiliary losses are required to balance utilization of experts and minimize unassigned tokens. These challenges can be greater in out-of-distribution settings: small inference batch sizes, novel inputs, or in transfer learning.
We introduce Soft MoE, that overcomes many these challenges. Rather than employing a sparse and discrete router that tries to find a good hard assignment between tokens and experts, Soft MoEs instead perform a soft assignment by mixing tokens. In particular, we compute several weighted with weights depending on both tokens and experts and then we process each weighted average by its corresponding expert.
averages of all tokens witttss--hhhhhhhhh
Soft MoE 1/16 outperforms VIT H/14 on upstream, few-shot and finetuning while requiring almost half the training time, and being 2 f'sterr at inference. Moreover, Soft MoE B/16 matches ViT H/14 on few-shot and finetuning and outperforms on upstream metrics after comparable amount of training. Remarkably, Soft MoE B/16 is 5.7x faster at inference despite having 5.5x the number of parameters of ViT H/14 (see Table and Figure 5 for details). Section demonstrates Soft MoE's potential to extend to other tasks: we train a contrastive model text tower against the frozen vision
"Equal contribution. The order was decided by a coin toss.

---

Published as a conference paper at ICLR 2024
Sparse MoË
Soft MoË
Assign
Weighted average
Expert
Expert 1
Expert
Expert
Export
Expert
Slots
Figure 1: Sparse and Soft MoE layers. While the router in Sparse MoE layers (left) learns to assign individual input tokens to each of the available slots, in Soft MoE layers (right) each slot is the result of a (different) weighted average of the input tokens. Learning to make discrete assignments introduces several optimization and implementation issues that Soft MoE sidesteps. Appendix G visualizes learned distributions of soft-assignments by Soft MoE.
tower, showing that representations learned via soft routing preserve their benefits for image-text alignment.
2 SOFT MIXTURE OF EXPERTS
2.1 ALGORITHM DESCRIPTION
The Soft MoE routing algorithm is depicted in Figure We denote the inputs tokens for one sequence by X E Rmndd where m is the number of tokens and is their dimension. Each MoE layer uses set of expert functions' applied on individual tokens, namely {.fs Rd F ) Each expert processes p slots, and each slot has corresponding d-dimensional vector of parameters, & E (x.pp)))
In particular, the input slots E (n-p)) xd are the result of convex combinations of all the m input tokens, X:
exp((X9)
Dij ΣX exp (xp)())))))
D
(1)
Notice that D, which we call the dispatch weights, simply the result of applying a softmax over columns of X Then, mentioned above, the corresponding expert function is applied on each slot (i.c. of X) to obtain the output slots: Y; s/rl (X,).
(X,).
Finally, the output tokens are computed as a convex combination all (n p) output slots, Ý, whose weights are computed similarly as before:
exp( (
Cjj Σラe
Y
We refer to C as the combine weights, and it is the result of applying a softmax over the rows of X
Following the usual design for Sparse MoEs, we replace a subset the Transformer's MLP blocks with Soft MoE blocks. We typically replace the second half of MLP blocks. The total number of slots a key hyperparameter of slots rather than on the number of experts. One set the number of slots equal to the input sequence length to match the FL.PPs of the equivalent dense Transformer.
Soft MoE layers because the time complexity depends on the number
2.2 PROPERTIES OF SOFT MoE AND CONNECTIONS WITH SPARSE MoEs
Fully differentiable Sparse MoE algorithms involve an assignment problem between tokens and experts, which subject to capacity and load-balancing constraints. Different algorithms approximate the solution in different ways: for example, the top-k or "Token Choice" router (Shazeer et al., 2017;
In practice, all experts apply the function with different parameters, usually ML.P.
2

---

Published as a conference paper at ICLR 2024
Figure 2: Soft MoE rout- ing details. Soft MoE com- putes scores or logits for every pair of input token and slot. From this it com- putes slotsxtokens matrix of logits, that are normal- ized appropriately com- pute both the dispatch and combine weights. The slots themselves are allocated to experts round-robin.
Soft MoE Layer
likie
Wrin
Bilk
Talen
cupu lissa
Lepikhin et al., 2020; Riquelme et al., 2021) selects the top-k-scored experts for each token, while there slots available in such expert (i.e. the expert has not filled its capacity). The "Expert Choice" router (Zhou et al., 2022) selects the crapacittyscore tokens for each expert. Other works suggest more advanced (and often costly) algorithms to compute the assignments, such approaches based on Linear Programming algorithms (Lewis et al., 2021), Optimal Transport (Liu et al., 2022; Clark al., 2022) or Reinforcement Learning (Clark et al., 2022). Nevertheless virtually all of these approaches are discrete in nature, and thus non diffferefniiable In contrast, all operations in Soft MoE layers are continuous and fully differentiable. We can interpret the weighted averages with softmax scores as soft assignments, rather than the hard assignments used in Sparse MoE.
_oft_mme____ def soft moe_ Lore__aerreee (X, Phi, experts):
2
# Compute the dispatch and combine weights.
logits = jnp. cinsum (' md, dnp->mmp' X, Phi)
jax. nn. softmax (logits, axis=(0,))
jax.nn.
(logits, axis= 2))
# The input slots are a weighted average of all the input tokens, # given by the dispatch weights.
Xs = jnp. emp.uinsum(' md, mnp->npd', , X, D)
# Apply the corresponding expert function to each input slot.
10
Ys stp.stackil
11
f (Xsli, :, :1) for i, f in enmerate (experts)
12
axis=0)
The output takens are weighted average of all the output slots, I given by the combine weights.
14
15
Y jnp.einsum ('ndu, mnp->nd Ys, C)
16
rsturn Y
Algorithm I: Simple JAX (Bradbury et al., 2018) implementation a Soft MoE layer. Full code is available at htttps://gitee
com/ google/ reeeeereeeeeeeeeeeeeeeeeeeeeeeee vmoe.
No token dropping and expert unbalance The classical routing mechanisms tend to suffer from issues such as "token dropping" (i.e. some tokens are not assigned any expert), or "expert unbalance" (i.e. some experts receive far more tokens than others). Unfortunately, performance can be severely impacted as a consequence. For instance, the popular top-k or "Token Choice" router (Shazeer et al., 2017) suffers from both, while the "Expert Choice" router (Zhou et al., 2022) only suffers from the former (see Appendix B for some experiments regarding dropping). Soft MoEs are immune to token dropping and expert unbalance since every slot is filled with weighted average of all tokens.
Fast The total number of slots determines the cost of Soft MoE layer. Every input applies such number of MLPs. The total number of experts is irrelevant in this calculation: few experts with many slots per expert or many experts with few slots per expert will have matching costs if the total number of slots is identical. The only constraint must meet is that the number of slots has to be greater or equal to the number of experts (as each expert must process at least one slot). The main advantage of Soft MoE is completely avoiding sort or top-k operations which are slow and typically not well suited for hardware accelerators. As a result, Soft MoE significantly faster than most sparse MoEs (Figure 6). See Section 2.3 for time complexity details.
3

---

Published as a conference paper at ICLR 2024
Features of both sparse and dense The sparsity in Sparse MoEs comes from the fact that expert parameters are only applied to a subset of the input tokens. However, Soft MoEs are not technically sparse, since every slot is weighted average of all the imput tokens. Every input token fractionally activates all the model parameters. Likewise, all output tokens are fractionally dependent on all slots (and experts). Finally, notice also that Soft MoEs are not Dense MoEs, where every expert processes all input tokens, since every expert only processes a subset of the slots.
Per-sequence determinism Under capacity constraints, all Sparse MoE approaches route tokens in groups of fixed size and enforce (or encourage) balance within the group. When groups contain tokens from different sequences inputs, these tokens compete for available spots in expert buffers. Therefore, the model is no longer deterministic at the sequence-level, but only the batch-level. Models using larger groups tend to provide more freedom to the routing algorithm and usually perform better, but their computational cost also higher.
2.3 IMPLEMENTATION
Time complexity Assume the per-token cost of a single expert function is O(k). The time com- plexity of a Soft MoE layer then (mmpd + npk).. By choosing p O(m/n) slots per expert, Le. the number of tokens over the number of experts, the cost reduces to O(m² 2 + mk).. Given that each expert function has its own set of parameters, increasing the number of experts n and scaling p accordingly. allows us to increase the total number of parameters without any impact on the time complexity. Moreover, when the cost of applying an expert is large, the mà term dominates over mdd, and the overall cost of Soft MoE layer becomes comparable to that of applying single expert on all the input tokens. Finally, even when m is not dominated, this is the same as the (single-headed) self-atte self-attention cost, thus it does not become a bottleneck in Transformer models. This be seen in the bottom plot Figure 6 where the throughput Soft MoE barely changes when the number of experts increases from 8 to 4 096 experts, while Sparse MoEs take significant hit.
Normalization In Transformers, MoE layers are typically used to replace the feedforward layer in each encoder block. Thus, when using pre-normalization as most modern Transformer architectures (Domhan, 2018; Xiong et al., 2020; Riquelme et al., 2021; Fedus et al., 2022), the inputs the MoE layer are "layer normalized". This causes stability issues when scaling the model dimension d, since the softmax approaches a one-hot vector as d + oc (see Appendix E). Thus, Line 3 of algorithm we replace X and Phi with 12_normalize (X, axis=1) and scale 12 m__ormmaliz (Phi, axis-0), respectively; where scale is a trainable scalar, and 12_normaliza normalizes the corresponding axis to have unit (L2) norm, as Algorithm 2 shows.
def 12_normal iz (x, axis, eps-le-6)
norm inp.sqrt (iup.suuarr square (x) sum (axis=axis, keepd (eepdims=True) return x jnp. jnp. rapprrecipocal (norm eps)
Algorithm 2: JAX implementation of the L2 normalization used in Soft MoE layers.
For relatively small values of d. the normalization has little impact on the model's quality. However, with the proposed normalization in the Soft MoE layer, we can make the model dimension bigger and/or increase the learning rate (see Appendix E).
Distributed model When the number of experts increases significantly, it not possible to fit the single device, especially during training or when using MoEs top of
entire model in memory large model backbones. In these cases, we employ the standard techniques to distribute the model across many devices, as in (Lepikhin et al., 2020; Riquelme et al., 2021; Fedus et al., 2022) and other works training large MoE models. Distributing the model typically adds an overhead in the cost of the model, which is not captured by the time complexity analysis based on FLOPs that we derived above. In order to account for this difference, all of our experiments we measure not only the FLOPs, but also the wall-clock time in TPUvv-chip-hours..
4

---

Published as a conference paper at ICLR 2024
3 IMAGE CLASSIFICATION EXPERIMENTS
Training Pareto frontiers. In Section 3.3 we compare dense ViT models at the Small, Base, Large and Huge sizes with their dense and sparse counterparts based on both Tokens Choice and Experts Choice sparse routing. We study performance at different training budgets and show that Soft MoE dominates other models in terms of performance at given training cost or time.
Inference-time optimized models. In Section 3.4, we present longer training runs ("overtrain- ing"). Relative to VIT, Soft MoE brings large improvements in terms of inference speed for fixed performance level (smaller models: S, B) and absolute performance (larger models: I., H).
Model ablations. In Sections 3.5 and 3.6 investigate the effect of changing slot and expert counts, and perform ablations on the Soft MoE routing algorithm.
3.1 TRAINING AND EVALUATION DATA
We pretrain our models on JFT-4B (Zhai et al., 2022a), a proprictary dataset that contains more than 4B images, covering 29k classes. During pretraining, evaluate the models on two metrics: upstream validation precision-at- on JFT-4B, and ImageNet 10-shot accuracy. The latter is computed by freezing the model weights and replacing the head with new that only trained on dataset containing 10 images per class from ImageNet-Ik (Deng et al., 2009). Finally, we provide the accuracy the validation set of ImageNet- after finetuning on the training set of ImageNet- Ik (1.3 million images) at 384 resolution.
3.2 SPARSE ROUTING ALGORITHMS
Tokens Choice. Every token selects the top-K K experts with the highest routing score for the to- ken (Shazeer et al., 2017). Increasing K typically leads better performance at increased computa- tional cost. Batch Priority Routing (BPR) (Riquelme et al., 2021) significantly improves the model performance, especially in the case of K (Appendix F, Table 7). Accordingly we use Top- K routing with BPR and K E 21, We also optimize the number of experts (Appendix F, Figure 11).
Experts Choice. Alternatively, experts can select the top-C' tokens in terms of routing scores (Zhou al., 2022). C is the buffer size, and we set E C c T where E is the number of experts, T is the total number of tokens in the group, and is the capacity multiplier. When = 1, all tokens can processed via the union of experts. With Experts Choice routing, it is common that some tokens are simultaneously selected by several experts whereas some other tokens are not selected at all. Figure 10, Appendix B illustrates this phenomenon. We experiment with 0.5, 1,2.
3.3 TRAINING PARETO-OPTIMAL MODELS
We trained VIT- (S/8, S/16, S/32, B/16, B/32, L/16, L/32, H/14) models and their sparse counterparts. We trained several variants (varying K, C' and expert number), totalling 106 models. We trained for 300k steps with batch size 4096, resolution 224, using a reciprocal square root learning rate schedule.
Figures 3a and 3b show the results for models in each class that lie on their respective training cost/performance Pareto frontiers. On both metrics, Soft MoE strongly outperforms dense and other sparse approaches for any given FLOOPs or time budget. Table 9, Appendix J, lists all the models, with their parameters, performance and costs, which are all displayed in Figure 19.
3.4 LONG TRAINING DURATIONS
We trained number of models for much longer durations, up 4M steps. We trained number of Soft MoEs on JFT, following a similar setting to Zhai et al. (2022a). We replace the last half of the blocks in ViT S/16, B/16. L/16, and H/14 with Soft MoE layers with 128 experts, using slot per expert. We train models ranging from 1B to 54B parameters. All models were trained for 4M steps, except for H/14, which trained for 2M steps for cost reasons.
5

---

Published as a conference paper at ICLR 2024
0.18
0.56
0.56
JPT-48 brcctontoos
0.52
0.70
0.50
0.65
0.48
Sed Mal!
0.50
Seet Mol!
Choice
Exgers
0.44
Trkerm Choice
0.55
Toker Choice
Dome
Deme:
0.50
10'
102
10
10
Total Training Edgss
Thaa Trinnng ELaPLOP
Toul Traning Triss
Tixal Thiining ExaFL.OP
(a) JFT-4B Precision-at-1
(b) ImageNet 10-shot Accuracy
Figure 3: Train Pareto frontiers. Soft MoE dominates both ViTs (dense) and popular MoEs (Experts and Tokens Choice) the training cost performance Pareto frontier. Larger marker sizes indicate larger models, ranging from S/32 to H/14. Cost is reported in terms of FLOPS and TPU-v3 training time. Only models on their Pareto frontier displayed, Appendix shows all models trained.
84%
H/14
H/14
L/16
HI14
L/16
InageNer Accurscy
B/16
B/16
B/16
JFT-48 Precision ak
88%
81%
SS116
ImageNet Finctanc
5/16
87%
77%
73%
54%
85%
Soft MoE
Danse
678
51%
84%
10a
102
10
Total Training ExaFLOPs
Total Training EsaFLOPs
Total Training EsaFLOPs
Figure 4: Models with long training durations. Models trained for 4M steps (H/14 trained only for 2M steps). Equivalent model classes (S/16, B/16, etc.) have similar training costs, but Soft MoE outperforms VIT on all metrics at a fixed training budget.
Figure 4 shows the JFT-4B precision, ImageNet 10-shot accuracy, and the ImageNet finetuning accuracy for Soft MoE and VIT versus training cost. Appendix F, Table 8 contains numerical results, and Figure 16 shows performance versus core-hours, from which the same conclusions can be drawn. Soft MoE substantially outperforms dense ViT models for given compute budget. For example, the Soft MoE S/16 performs better than ViT B/16 on JFT and 10-shot ImageNet, and it also improves finetuning scores the full ImageNet data, even though its training (and inference) cost significantly smaller. Similarly, Soft MoE B/16 outperforms ViT L/16 upstream, and only lags 0.5 behind after finetuning while being 3x faster and requiring almost 4x fewer FLOPs. Finally, the Soft MoE 1/16 model outperforms the dense H/14 one while again being around 3x faster in terms of training and inference step time.
We continue training the small backbones up 9M steps obtain models of high quality with low inference cost. Even after additional (over) training, the overall training time with respect to larger ViT models is similar smaller. For these runs, longer cooldowns (linear learning rate decay) works well for Soft MoF. Therefore, we increase the cooldown from 50k steps to 500k steps.
Figure 5 and Table present the results. Soft MoE B/16 trained for 1K TPUv3 days matches or outperforms VIT H/14 trained on similar budget, and is 10x cheaper at inference in FLOPs (32 vs. 334 GFLOPS/img) and > 5x cheaper in wall-clock time (1.5 vs. 8.6 ms/img). Soft MoE B/16 matches the ViT H/14 model's performance when double ViT-H/14's training budget (to 2k TPU- days). Soft MoE L/16 outperforms all VIT models while being almost 2x faster at inference than VIT H/14 (4.8 vs. 8.6 ms/img).

---

Published as a conference paper at ICLR 2024
VIT
Sall MoB
Sell M&E (amg?
1844
1216
1016
806
1/16
1514
1.16
625
S115
1216
1014
Bris
675
755
6:16
Br.b
766
2716
2016
216
5/16
30
10"
10'
10!
101
101
Brduniicecco cent (OPLCC
Brdiaicc (OPLC()kng)
Transsg con (sxaFLOH)
Tratsin con (staFLO0)
37x
5.7x
85%
Beugyta (b Mamm
B216
625
40%
Lits
7014
Imspehinn (i-con
60%
6000
1934
13%
frt.6
57%
Br1s
550
hrt6
70%
70%
10
10
10'
10"
Extturtre tim (TPIN mwtrigg
Bodturios time (TPMNS musrng)
Twainggite (TPUNS
Tokking (TO0
Figure 5: Models optimized for inference speed. Performance models trained for more steps, thereby optimized for performance given inference cost (TPUv3 time FLOPs).
Table 1: Models trained for longer durations (cooldown steps in parentheses).
Model
Params
Train
Train TPU-days exaFLOP ms/img GFLOP/mg
Train
Eval
JFT P@
INct
INet
steps (cd)
10shot finetune
VIT S/16
33M
4M (50k)
153.5
227.1
0.5
9.2
51.3
67.6
84,0
VIT B/16
108M
4M (50k)
410.1
864.1
1.3
35.1
56.2
76.8
86.6
VIT L/16
333M
4M (50k)
1290.1
3025.4
4.9
122.9
59.8
81.5
VIT H/14
669M
IM (50k)
1019.9
2060.2
8.6
334.2
58.8
82.7
88.6
VIT H/14
669M
2M (50k)
2039.8
4120.3
8.6
334.2
59.7
83.3
88.9
Soft MoF. S/14 256E
1.8B 10M (50k)
494.7
814.2
0.9
13.2
60.1
80.6
87.5
Soft MoE B/16 128E
3.7B
(500k)
1011.4
1769.5
32.0
62.4
82.9
88.5
Soft MoE L/16 128E
13.1B 4M (500k)
1355.4
2734.1
4.8
111.1
63.0
84.3
89.2
3.5 NUMBER OF SLOTS AND EXPERTS
We study the effect of changing the number of slots and experts the Sparse and Soft MoEs. Figure shows the quality and speed of MoEs with different numbers of experts, and numbers of slots per token; the latter equivalent to the average number experts assigned per token for Sparse MoEs. When varying the number of experts, the models' backbone FLOPs remain constant, changes in speed are due to routing costs. When varying the slots exots-ter-expert, the number tokens processed in the expert layers increases, so the throughput decreases. First, observe that for Soft MoF, the best performing model at each number of slots-per-token is the model with the most experts (i.e. one slot per expert). For the two Sparse MoEs, there is a point at which training difficulties outweigh the benefits of additional capacity, resulting in the a modest optimum number of experts. Second, Soft MoE's throughput is approximately constant when adding more experts. However, the Sparse MoEs' throughputs reduce dramatically from 1 experts, see discussion in Section 2.2.
3.6 ABLATIONS
We study the impact of the components of the Soft MoE routing layer by running the following ablations: Identity routing: Tokens are not mixed: the first token goes to first expert, the second token goes to second expert, etc. Uniform Mixing: Every slot mixes all input tokens the same way: by averaging them, both for dispatching and combining. Expert diversity arises from different initializations of their weights. Soft / Uniform: We learn token mixing on input to the experts to create the slots (dispatch weights), but average the expert outputs. This implies every input token is
7

---

Published as a conference paper at ICLR 2024
Soft
Exports Choice
Tekens
Nieher Stoa per Thara
-06.5 66.7 705
71.6 720 723
65.4
66.6
67.1 68.2 582
664
64.3
66.2
66.3
(4.3
55.0
67.4 09.5 706
72.0 72, 15.6 74.0 74.4
67.1
68.6
69.8
704 706 702 7611 60.7
454
65.5
67.5
67.
678
57.3
54.
71.2 12.2 7311 74.8
67.3
AM-3
10A
11.4 718 73.4 T0LR
08.5
69.6
703 142 75L
725 732
71.
-558
66.6
66.3
62.1
12.1
Numbbroo Expert
Numberr of Experti
Nambered Expera
Experts Choice
Tekens
Nieeera Staqpe Totr
379
576
315
571
367
345
-314
373
359
361
350
544
366
3.0
347
324
355
255
285
212
178
282
267
255
274
275
274
271
26S
267
155
362
222
2722
266
259
245
251
243
216
200
212
196
2B1
303
194
172
188
1M1
134
153
1::
1:9
129
107
65
100
(022
3H
1004
3HI
Eapara
orEapetss
Figure 6: Top: Performance (ImageNet) for MoEs with different number of experts (columns) and poots-per-token / assignments-per peeermeeeet---tttoeen (rows). Bottom: Training throughput of the same models. Across the columns, the number of parameters increases, however, the theoretical cost (FLOPS) for the model (not including routing cost) remains constant. Descending the rows, the expert layers become more compute intensive as more tokens/slots are processed in the MoE layers.
Table 2: Ablations using Soft MoE-S/14 with 256 experts trained for 300k steps.
Method
Experts
Mixing
Leamndd Dispatch
Learned Combine
JFT p@1
IN/1Oshot
Soft MoE Soft / Uniform Uniform / Soft Uniform Identity
54.3%
74.8%
53.6%
72.0%
V
52.6%
71.8%
51.8%
70.0%
51.5%
69. %
VIT
48.3%
62.3%
identically updated before the residual connection. Uniform Soft. All slots filled with uniform average the input tokens. We learn slot mixing of the expert output tokens depending on the input tokens. Table 2 shows that having slots is important; Identity and Uniform routing substantially underperform Soft MoE, although they do outperform VIT. Dispatch mixing appears slightly more important than the combine mixing. See Appendix A for additional details.
4 CONTRASTIVE LEARNING
We test whether the Soft MoB's representations are better for other tasks. For this, try image- text contrastive learning. Following Zhai et al. (2022b), the image tower is pre-trained on image classification, and then frozen while training the text encoder a dataset of image-tex pairs. We re-use the models trained on JFT in the previous section and compare their performance zero-shot on downstream datasets. For contrastive learning we train WebLI (Chen et al., 2022), proprietary dataset consisting of 10B images and alt-texts. The image encoder is frozen, while the text encoder is trained from scratch.
Table shows the results. Overall, the benefits observed image classification also in this setting. For instance, Soft MoE-L/16 outperforms ViT-L/16 by more than 1% and 2% on ImageNet and Cifar-100 100 zero-shot, respectively. However, the improvement COCO retrieval modest, and likely reflects the poor alignment between features learned on vlosbl-vocabulary JFT and this open-vocabulary task.
Finally, in Appendix F.1 we show that Soft MoEs also surpass vanilla ViT and the Experts Choice router when trained from scratch on the publicly available LAION-400M (Schuhmann et al., 2021). With this pretraining, Soft MoEs also benefit from data augmentation, but neither ViT nor Experts Choice seem to benefit from it, which is consistent with our observation Section 3.5, that Soft MoEs make better use of additional expert parameters.
8

---

Published as a conference paper at ICLR 2024
Table LIT-style evaluation with ViT-g text tower trained for 18B input images (~ 5 epochs).
Model
Experts IN/Oshot Cifar 100/0shot Pet/Oshot Coco Img2Text Coco Tex12Img
ViT-S/16
74.2%
56.6%
94.8%
53.6%
37.0%
Soft MoE-S/16
128
81.2%
67.2%
96.6%
56.0%
39.0%
Soft MoE-S/14
256
82.0%
75.1%
97.1%
56.5%
39.4%
ViT-B/16
79.6%
71.0%
96.4%
58.2%
41.5%
Soft MoE-B/16
128
82.5%
74.4%
97.6%
58.3%
41.6%
ViT-L/16
82.7%
77.5%
97.1%
60.7%
43.3%
Soft MoE-L/16
128
83.8%
79.9%
97.3%
60.9%
43.4%
Souped Soft MoE-L/16
128
84.3%
81.3%
97.2%
61.1%
44.5%
ViT-H/14
83.8%
84.7%
97.5%
62.7%
45.2%
Soft MoE-H/14
256
84.6%
86.3%
97.4%
61.0%
44.8%
RELATED WORK
Many existing works merge, mix or fuse input tokens to reduce the input sequence length (Jaegle al., 2021; Ryoo al., 2021: Renggli et al., 2022; Wang et al., 2022), typically using attention- weighted averages with fixed keys, to try to alleviate the quadratic cost of self-attention with respect the sequence length. Although our dispatch and combine weights are computed to these approaches, our goal is not to reduce the sequence length (while it is possible), and we actually recover the original sequence length after weighting the experts' outputs with the combine weights, at the end of each Soft MoE layer.
similar fashion
Multi-headed attention also shows some similarities with Soft MoE, beyond the use of softmax in weighted averages: the h different heads can be interpreted as different (linear) experts. The distinction is that, if m is the sequence length and each input token has dimensionality d, each of h heads processes m vectors size d/h. The m resulting vectors are combined using different weights for each of the m' output tokens (i.e. the attention weights), on each head independently, and then the resulting (d/h) d. Our experts (d/h))dimeensee -dimensional vectors from each head are concatenated into one of dimension non-linear and combine vectors of size d, at the input and output of such experts.
Other MoE works use weighted combination of the experts parameters, rather than doing sparse routing of the examples (Yang al., 2019; Tian et al., 2020; Muqeeth et al., 2023). These approaches are also fully differentiable, but they can have higher cost, since 1) they must average the parameters of the experts, which can become time and/or memory bottleneck when experts with many parameters are used; and 2) they cannot take advantage of vectorized operations as broadly as Soft (and Sparse) MoEs, since every input uses a different weighted combination of the parameters. We recommend the "computational cost" discussion in Muqeeth et al. (2023).
CURRENT LIMITATIONS
Auto-regressive decoding One of the key aspects of Soft MoE consists in learning the merging of all tokens in the input. This makes the use of Soft MoEs in auto-regressive decoders difficult, since causality between past and future tokens has to be preserved during training. Although causal masks used in attention layers could used, one must be careful not introduce any correlation between token and slot indices, since this may bias which token indices each expert is trained on. The use of Soft MoE in auto-regressive decoders a promising research avenue that we leave for future work.
Lazy experts & memory consumption We show in Section that one slot per expert tends to be the optimal choice. In other words, rather than feeding expert with two slots, it is more effective to use two experts with one slot each. We hypothesize slots that use the same expert tend to align and provide small informational gains, and expert may lack the flexibility to accommodate very different slot projections. We show this in Appendix I. Consequently, Soft MoE can leverage a large number of experts and-while its cost is still similar to the dense backbone- the memory requirements of the model can grow large.

---

Published as a conference paper at ICLR 2024
REFERENCES
Emmanuel Bengio, Pierre-Lue Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. arXiv preprint /00i::/5/110029999 2015.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, al. JAX: compos- able transformations of Python+ NumPy programs, 2018.
Xi Chen, Xiao Wang. Soravit Changpinyo, AJ Piergiovanni, Piotr Padlewski, Daniel Salz, Sebastian Goodman, Adam Grycner, Basil Mustafa, Lucas Beyer, et al. Pali: A jointly-scaled multilingual language-image model. arXiv preprint arXiv:2209.088 06794, 2022.
Aidan Clark, Diego De Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann. Bogdan Damoc, Blake Hechtman, Trevor Cai, Sehastian Borgeaud, et al. Unified scaling laws for routed language models. International Conference on Machine Learning, pages 4057 4086. PMLR, 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-24882 255. leee, 2009.
Tobias Domhan. How much attention do you need? a granular analysis of neural machine translation architectures. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Pupers), pages 1799-1808, 2018
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1): 5232-5270. 2022.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249-222922222 256. JMLR Workshop and Conference Proceedings, 2010.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-1034, 2015.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In International conference on machine learning, pages 4661-4664. PMLR, 2021.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08360 2020.
Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self- normalooma neural networks. Advances in neural information processing systems, 30, 2017.
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. arXiv preprint arXiv:2006. 16668, 2020.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning. pages 6265-6274. PMLR, 2021.
Tianlin Liu, Joan Puigcerver, and Mathieu Blondel. Sparsity-constraine optimal transport. arXiv preprint arXiv:2209. /5466, 2022.
10

---

Published as a conference paper at ICLR 2024
Mohammed Muqeeth, Haokun Liu, and Colin Raffel. Soft merging of experts with adaptive routing, 2023.
Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Multi- modal contrastive learning with limoe: the language-image mixture of experts. arXiv preprint arXiv:2206.02770 2022.
Cedric Renggli, André Susano Pinto, Neil Houlsby, Basil Mustafa, Joan Puigcerver, and Carlos Riquelme. Learning to merge tokens vision transformers. arXiv preprint arXiv:2202.12015, 2022.
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. Advances in Neural Information Processing Systems, 34:8583-8599 2021.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34:17555-17566, 2021.
Michael S Ryoo, AJ Piergiovanni, Anurag Arnab, Mostafa Dehghani, and Anelia Angelova. Token- learner: What can learned tokens do for images and videos? arXiv preprint arXiv:2 12977 2021.
Christoph Schuhmann, Richard Vencu, Romain Beaumont, Robert Kaczmarczyk, Clayton Mullis, Aarush Katta, Theo Coombes, Jenia Jitsev, and Aran Komatsuzaki. LAION-400m: Open dataset of CLIP-filtered 400 million image-text pairs. arXiv preprint arXiv:21 1.02/ 14, 2021.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gate mixture-of-exxeree layer. arXiv preprint arXiv: .06538,, 2017.
Zhi Tian, Chunhua Shen, and Hao Chen. Conditional convolutions for instance segmentation. In Computer -CCV 2020: 16th European Conference, Glasgow, UK, August 23-28 2020, Proceedings, Part I 16, pages 282-298. Springer, 2020.
Yikai Wang, Xinghao Chen, Lele Cao, Wenbing Huang, Fuchun Sun, and Yunhe Wang. Multimodal token fusion for vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 12186-12195, June 2022.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization the transformer architecture. In International Conference on Machine Learning, pages 10524-10533, PMLR, 2020.
Brandon Yang, Gabriel Bender, Quoe V Le, and Jiquan Ngiam. Condconv: Conditionally parameter- ized convolutions for efficient inference. Advances in Neural Information Processing Systems, 32, 2019.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 121031 2022a.
Xiaohua Zhai, Xiao Wang, Basil Mustafa, Andreas Steiner, Daniel Keysers, Alexander Kolesnikov, and Lucas Beyer. Lit: Zero-shot transfer with locked-image text tuning. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 18123-18133, 2022b.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mixture-of-experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114, 2022.
11


Reviews:
[{"review_id": "JSv3xrDx7J", "reviewer": "Reviewer_HCxh", "summary": "The following paper proposes Soft Mixture-of-Experts (MoE), a fully-differentiable sparse transformers capable of passing different weighted combinations of all input tokens to each expert. Compared to other MoEs, experts in Soft MoE only process a subset of (combined) tokens, resulting in larger model capacity (and performance) with lower inference cost. Thanks to this, Soft MoE are able to alleviate training instability, token dropping, scalability issues, and ineffective finetuning process that is apparent in previous type of MoE. Experiments on visual recognition tasks in Section 3 and 4 show the superiority of Soft MoE over dense Vision Transformers (ViT) and popular MoEs not limited to Tokens Choice and Experts Choice. In addition to that, Soft MoE provably scales well; since while the Huge/14 version of it has $40\\times$ more parameters than ViT Huge/14, it bring substantial improvement in terms of classification accuracy with only 2% gain in inference.", "questions": "- Is it possible to see the result of MoE in GLUE or SuperGLUE benchmark as seen on [1]?\n\n[1] Zhou et al. Mixture-of-experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103–7114, 2022. https://arxiv.org/abs/2202.09368", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Paper is well-written, easy to understand, and provided with source codes that model how Soft MoE is being implemented.\n- Experiments are comprehensive and detailed with respect to its domain, with emphasis not only on performance, but also in terms of inference and wall-clock time.\n- High potential to bring improvement when it is deployed towards various modalities such as Large-Language Models (LLMs).", "weaknesses": "Aside from the weaknesses mentioned in the paper, I would like to address concerns that is apparent in the paper:\n- The experiments performed in Sections 3 and 4 seem to focus only on vision-related tasks. It would be great to be able to observe results on different modalities such as NLP-related tasks based on GLUE or SuperGLUE benchmark that is performed in [1]. \n- Unfortunately, the dataset that is being used for training is not publicly available; making it hard to be used for benchmarking with other papers.\n- As mentioned in Section 6, the memory footprint will be huge when we leverage a large number of experts, and usage in auto-regressive decoders is problematic due to the preservation of past and future tokens during training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VXRBriqdAh", "reviewer": "Reviewer_fyPo", "summary": "This paper introduces a novel soft MOE method that keeps the benefits of sparse MOE (processing a subset of tokens for lower cost), and also has full differentiability and better performance. The proposed method let the MOE experts to process a number of slots, where each slot is a combination of all tokens. The idea of processing a number of slots is how the inference cost is controlled. It outperforms dense (non MOE) networks and also typical sparse MOE networks: with similar inference cost, the accuracy metric is much better.", "questions": "I hope the authors can release as much details as possible of the implementation details. Looking forward to the adaptation into different fields.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Originality\nThe idea of this soft MOE is new. I have not encountered previous work that uses this idea. It combines the idea of sparse MOE with an attention-like mechanism to get the benefits of both. The major part that I like is that this combination is very practical in large scale model that requires model parallelism. The author also clearly contrasted its difference with other existing works and multi-head attention in Section 5. \n\nQuality\nThe quality of the paper is high. \n-\tThe author provides very good details of how the method is implemented and some details of implementation.\n-\tThe experiments settings are clearly stated and the hyper-parameter searching is transparent and properly done for contrast models (dense models and other MOE models) as well.\n-\tThe design choices (slots per token and expert numbers) are well studied and concluded.\n-\tThe experiments conducted are at large scale datasets and on various tasks/settings, and the improvement of the proposed methods are convincingly large.\n\nClarity\nThe clarity of the paper is good. The demonstration of the method and experiment conclusions are clear and well supported by experiments.\n\nSignificance\nThis is very exciting work that can impact the community greatly. The proposed method shows convincing performance dominancy over other MOE methods or dense models on image classification task and contrastive learning tasks. It shows that with similar inference cost, it can achieve much better accuracy metric, even against strong baselines. It is also very practical that it can be used in modern large scale parallel learning platforms without too much extra optimization needed.", "weaknesses": "as mentioned already by the author, the main weakness of this method from my perspective is: each expert do not handle multiple tokens well (i.e. one expert one token is better). In practice, this may cause inefficient increase of number of parameters (memory). But the authors have already recognized it, and I think it doesn't hurt the significance of the existing contribution of this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tf93Hr9BOA", "reviewer": "Reviewer_wHYN", "summary": "Instead of sending each expert a bag of vectors from the input the vectors entering each expert are composed of a weighted sum of all input vectors. My understanding of the implementation was assisted by [Phil Wang's replication][lucidrains] in PyTorch, which I have shortened below:\n\n```\ndef forward(self, x):\n        \"\"\"\n        einstein notation\n        b - batch\n        n - sequence length\n        e - number of experts\n        s - number of slots per expert\n        d - feature dimension\n        \"\"\"\n        slot_embeds = self.slot_embeds\n\n        logits = einsum('b n d, e s d -> b n e s', x, slot_embeds)\n\n        # get dispatch and combine weights (softmax across right dimensions)\n\n        dispatch_weights = logits.softmax(dim = 1)\n\n        combine_weights = rearrange(logits, 'b n e s -> b n (e s)')\n        combine_weights = combine_weights.softmax(dim = -1)\n\n        # derive slots by weighted average of input tokens using the dispatch weights from above\n\n        slots = einsum('b n d, b n e s -> b e s d', x, dispatch_weights)\n\n        # route the slots per expert to each expert\n\n        out = self.experts(slots) # vectors are mapped across dim e to each expert\n\n        # combine back out\n\n        out = rearrange(out, ' b e s d -> b (e s) d')\n        out = einsum('b s d, b n s -> b n d', out, combine_weights)\n```\n\nThe contributions of this paper are mainly the experimental results, the experiments aim to demonstrate that a soft-MoE layer allows faster inference and shorter training time on some large scale training benchmarks. Specifically:\n\n- Soft MoE L/16 vs ViT H/14:\n    - Half the training time\n    - 2x faster inference\n- Soft MoE B/16 vs ViT/14:\n    - 5.7x faster at inference time\n\nThese results are demonstrated comprehensively with ablation experiments investigating the different assumptions, such as whether the weighted sum is necessary, whether the weighted sum can be uniform and whether the weighted sum at the output is necessary; finding that the decisions all improve performance.\n\nThe results are demonstrated on expensive large scale benchmarks, including training many ViT style models on JFT-4B and further experiments testing these models on a contrastive benchmark using WebLI. The results presented support the contributions of the paper, demonstrating improved performance across the experiments compared to other MoE methods.\n\nThe authors also note that the proposed method would violate the causal assumption if used in an autoregressive model, meaning it can't be used as defined above for autoregressive language modelling. \n\n[lucidrains]: https://github.com/lucidrains/soft-moe-pytorch/blob/main/soft_moe_pytorch/soft_moe.py", "questions": "The authors mention that autoregessive modelling is an area for future work, but it's not totally clear to me how to achieve that with this type of model. The Token mixing weights would need to have a causal mask similar to the attention causal mask, but that seems like it might affect performance or throughput. Do they have any ideas how to approach this issue?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper makes a clear contribution of an architectural improvement that requires extensive experimental justification to prove. It then provides the experimental results to back up this assertion on expensive benchmarks. The results are presented clearly and it is easy for the reader to find relevant information.", "weaknesses": "Minor presentation issue:\n\n- Figure 2 is confusing, the relationship between dispatch and combine weights to the tokens is illustrated with two downward arrows, but they don't really mean anything so it doesn't help the reader\n- Better signposting about all the different results that may be found in the experiments section as it is the most valuable part of the paper. For example, in the introduction some of the inference time and training time benefits are mentioned but not where these results are found.\n\nIt would be useful to see the complexity results described in Section 2.3 verified by experiment, perhaps looking at throughput in practice. I think Figure 6 implies something to do with this but the relationship is not entirely obvious.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DvZZkrfhzg", "reviewer": "Reviewer_pKE3", "summary": "This work focuses on a number of issues in vanilla Sparse MoEs, and proposes Soft MoE. Instead of hard assignment of tokens, Soft MoE “softs” the tokens by weighted averaging them before assigned to experts. Experiments with ViT models and CV tasks demonstrate that Soft MoE outperforms both Sparse MoEs and dense models.", "questions": "When assessing the training and inference speed, the authors chose dense ViT as the baseline. I would be better to include the results for Sparse MoEs additionally.\n\nThe experiments focus on CV tasks (in particular, ViT-based MoEs). It makes me wonder how Soft MoE would be on NLP tasks. Although in Section 6 the authors have discussed that Soft MoE is difficult to use in auto-regressive models, attempts on BERT- or T5-based MoEs should be suggested.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This work is novel and sound. The soft assignment by learnable weighted combinations helps tackle many limitations of vanilla Sparse MoEs. Extensive empirical results show that Soft MoE outperforms vanilla Sparse MoEs and dense models in terms of model ability (precisions/accuracies/...) and improves the training and inference speed/throughput compared with dense models.", "weaknesses": "The proposed technique cannot be integrated with main-stream foundation models (e.g., auto-regressive LLMs), which may limit the applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)