Scaling Vision with Sparse Mixture of Experts
Authors: Carlos Riquelme Ruiz
Joan Puigcerver
Basil Mustafa
Maxim Neumann
Rodolphe Jenatton
André Susano Pinto
Daniel Keysers
Neil Houlsby
Venue: NeurIPS 2021 Poster (2021)
Keywords: sparse
model
mixture
experts
sparsity
vision
conditional
computation
routing
router
adaptive
compute

Abstract: Sparsely-gated Mixture of Experts networks (MoEs) have demonstrated excellent scalability in Natural Language Processing. In Computer Vision, however, almost all performant networks are "dense", that is, every input is processed by every parameter. We present a Vision MoE (V-MoE), a sparse version of the Vision Transformer, that is scalable and competitive with the largest dense networks. When applied to image recognition, V-MoE matches the performance of state-of-the-art networks, while requiring as little as half of the compute at inference time. Further, we propose an extension to the routing algorithm that can prioritize subsets of each input across the entire batch, leading to adaptive per-image compute. This allows V-MoE to trade-off performance and compute smoothly at test-time. Finally, we demonstrate the potential of V-MoE to scale vision models, and train a 15B parameter model that attains 90.35% on ImageNet.

Scaling Vision with Sparse Mixture of Experts
Carlos Riquelme Google Brain
Joan Puigcerver Google Brain
Basil Mustafa Google Brain
Maxim Neumann
Google Brain
Rodolphe Jenatton
André Susano Pinto
Daniel Keysers
Neil Houlsby Google Brain
Google Brain
Google Brain
Google Brain
Abstract
Sparsely-gated Mixture of Experts networks (MoEs) have demonstrated excellent scalability in Natural Language Processing. In Computer Vision, however, almost all performant networks are "dense", that is, every input is processed by every parameter. We present a Vision MoE (V-MoE), sparse version of the Vision Transformer, that is scalable and competitive with the largest dense networks. When applied to image recognition, V-MoE matches the performance state-of- the-art networks, while requiring little half of the compute at inference time. Further, we propose an extension to the routing algorithm that can prioritize subsets of each input across the entire batch, leading to adaptive per-image compute. This allows V-MoE to trade-off performance and compute smoothly at test-time. Finally, demonstrate the potential of V-MoE scale vision models, and train 15B parameter model that attains 90.35% on ImageNet.
1 Introduction
Deep learning historically shows that increasing network capacity and dataset size generally improves performance. In computer vision, large models pre-trained large datasets often achieve the state of the art [57]50) (NLP), where large pre-trained models ubiquitous, and perform very well many tasks 48 [8] Text Transformers [6] are the largest models to date, some with over 100B parameters (9]. However, training and serving such models is expensive [56] 46). This is partially because these deep networks are typically "dense"- every example is processed using every parameter -thus, scale comes at high computational cost. In contrast, conditional computation
(50 36 20 3. This approach has had even more success in Natural Language Processing
aims to increase model capacity while keeping the training and inference cost roughly constant by applying only subset of parameters each example. In NLP, sparse Mixture of Experts (MoEs) are gaining popularity [54] 33) 22).
enabling training and inference with fewer resources while unlocking trillion parameter models.
In this work, we explore conditional computation for vision at scale. We introduce the Vision MoE (V-MoF), a sparse variant the recent Vision Transformer (ViT) architecture (20) for image classification. The V-MoE replaces subset of the dense feedforward layers in ViT with sparse MoE layers, where each image patch is "routed" to subset of "experts" (MLPs). Due to unique failure modes and non-differentiabiliy, routing in deep sparse models is challenging. We explore various design choices, and present an effective recipe for the pre- trainini and transfer of V-MoE, notably outperforming their dense counterparts. We further show that V-MoE models are remarkably flexible. The performance vs. inference-con cost trade-off of already trained models can be smoothly adjusted during inference by modulating the sparsity level with respect the imput and/or the model weights. Also, we open-source our implementation and number of V-MoE models trained on ImageNet-21k
'These authors contributed equally. Correspondence to rikel, jpuigcervei basilm | @google.com Mixture of experts code and models available at http:://gtthub..../ com wwwgl//gooolleeeoooooooooooo - research/vmoe
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

x
MoE Block
Norm
Multi Head Attention
Norm
MLP
Norm
Mult--Head Attention
Norm
Sparse MaE
Images represented different colors. Tokens from the same image distinguished by gradient.
Sparse MoE
Device
Router
Tokens dispatched between pairs of
come different images. Each token routed Independently
Device E
Within device connection
Router
MLP
To All, individual tokens
Figure 1: Overview the architecture. V-MoE is composed of ViT blocks. In some, we replace the MLP with a sparsely activated mixture of MLPs. Each MI. (the expert) is stored on separate device, and processes fixed number of tokens. The communication of these tokens between devices is shown in this example, which depicts the case when k expert is selected per token. Here each expert uses a capacity ratio
4: the sparse MoE layer receives 12 tokens per device, but each
3'
expert has capacity for 16 ( [6611
12 3: see Section 2.4 Non-expert components of V-MoE such as
routers, attention layers and normal MLP blocks replicated identically across devices.
With V-MoE, we scale to model sizes of 15B parameters, the largest vision models to date. We match the performance of state-of-the-arr dense models, while requiring fewer time to train. Alternatively, V-MoE match the cost of ViT while achieving better performance. To help control this tradeoff, we propose Batch Prioritized Routing, a routing algorithm that repurposes model sparsity skip the computation of some patches, reducing compute uninformative image regions.
We summarize our main contributions as follows:
Vision models at scale. We present the Vision Mixture of Experts, distributed aparsely-actvvated Transformer model for vision. We train models with up to 24 MoE layers, 32 experts per layer, and almost 15B parameters. We show that these models can be stably trained, seamlessly used for transfer, and successfully fine-tuned with as few as 000 datapoints. Moreover, our largest model achieves 90.35% test accuracy ImageNet when fine-tuned.
Performance and inference. We show V-MoEs strongly outperform their dense counterparts on upstream, sho- and full fine-tur tuning metrics in absolute terms. Moreover, at inference time, the V-MoE models can be adjusted to either (i) match the largest dense model's performance while using little as half the compute, actual runtime, (ii) significantly outperform it at the same cost. Batch Prioritized Routing. We propose a new priority-based routing algorithm that allows V-MoEs to discard the least useful patches. Thus, devote less compute to each image. In particular, show V-MoEs match the performance of the dense models while saving 20% of the training FLOPs. Analysis. We provide some visualization of the routing decisions, revealing patterns and conclusions which helped motivate design decisions and may further improve understanding in the field.
2 The Vision Mixture of Experts
We first describe MoEs and sparse MoEs. We then present how we apply this methodology to vision, before explaining our design choices for the routing algorithm and the implementation of V-MoEs.
2.1 Conditional Computation with MoEs
Conditional computation aims at activating different subsets of a network for different inputs 5. mixture-of-expertt model is specific instantiation whereby different model "experts" responsible for different regions of the input space B1.
We follow the setting of [54), who present for deep learning mixture of experts layer with E experts MoE(x) Σ g(x)s ex(x) where e DRD is the input the layer, e, EgD + R is

---

the function computed by expert i, and g: IR - ~ R is the "routing" function which prescribes the input-conditioned weight for the experts. Both e, and g are parameterized by neural networks. As defined, this is still a dense network. However, if g is sparse, i.e., restricted to assign only k < E non-zero weights, then unused experts need not be computed. This unlocks super-linear scaling of the number of model parameters with respect to inference and training compute.
2.2 MoEs for Vision
We explore the application of sparsity vision in the context of the Vision Transformer (ViT) 20 ViT has been shown to scale well the transfer learning setting, attaining better accuracies than CNNs with less pre-trate tre-training compute. VIT processes images first divided into a grid of equal-sized patches. These are linearly projected to the Transformer's [6]
sequence of patches. An input image is
hidden size. After adding positional embeddings, the patch embeddings (tokens) are processed by Transformer, which consists predominately of alternating self-attention and MLP layers.
The MIPs have two layers and GeI.U [29] non-linearity: MLP((x= W 2 gelu Wix), For Vision MoE, we replace a subset of these with MoE layers, where each expert is an MLP; see Figure The experts have the same architecture e, (x) MLPP, (x) but with different weights 02 (Wi, W:). This follows a similar design pattern as the M4 machine translation model [3]].
2.3 Routing
For each MoE layer in V-MoE, we use the routing function g(x) = TOP; (softmax (Wx + c)), where TOP; an operation that sets all elements of the vector zero except the elements with the largest k values, and is sampled independently e ~ N (0, E entry-wise. In practice, we use k = I representation of an image token some
or k = 2. In the context of the Vision Transformer, layer of the network. Therefore, V-MoE routes patch representations, not entire images.
The difference between previous formulations [54] is that apply TOP, after the softmax over experts weights [39]. instead of before. This allows us to train with k = 1 (otherwise gradients with respect routings are zero almost everywhere) and also performs better for k 1 (see Appendix [A].
Finally, we add a small amount of noise with standard deviation E to the activations Wx. We empirically found this performed well but that the setup was robust to this parameter. The noise typically altered routing decisions ~ 15% of the time earlier layers, and ~2-3% deeper layers.
2.4 Expert's Buffer Capacity
During training, sparse models may favor only a small set of experts 26 52 This common failure mode can cause two problems. First, statistical inefficiency: in the limit of collapse to single expert, the model is no more powerful than a dense model. Second, computational inefficiency: imbalanced assignment of items to experts may lead to a poor hardware utilization.
To combat imbalance and simplify our implementation, we fix the buffer capacity of each expert (i.e. the number of tokens that each expert processes), and train our model with auxiliary losses that encourage load balancing. This is essentially the same approach as followed by [54] 39 22]. In our case, we use slight variants of two of the auxiliary losses proposed in [54). as described in Appendix A
We define the buffer capacity of an expert (Ba) as a function of the number of images in the batch (N), the number of tokens per image (P), the number of selected experts per token (l)), the total number of experts (F),, and the capacity ratio (C') B round (NNCC)
If the router assigns more than B, tokens to given expert, only Be of them are processed. The remaining tokens are not entirely 'lost' as their information is preserved by residual connections (the top diagram of Figure I). Also, if k > 1, several experts try to process each token. Tokens are never fully discarded. If an expert is assigned fewer than B. tokens, the rest its buffer is zero- zero-pddded.
We use the capacity ratio to adjust the capacity of the experts. With C'> > 1, slack capacity is added to account for a potential routing imbalance. This typically useful for fine-tuning when the new data might come from a very different distribution than during upstream training. With C' < 1, the router forced to ignore some assignments. In Section we propose a new algorithm that takes advantage of setting C << 1 to discard the least useful tokens and save compute during inference.
3

---

3 Transfer Learning
In this section, we first present training different variants of V-MoE on a large dataset (Section 3.2) in order to be used for Transfer Learning afterwards. The ability to easily adapt our massive models to new tasks, using a small amount of data from the new task, is extremely valuable: allows to amortize the cost of pre-training across multiple tasks. We consider two different approaches to Transfer Learning: linear few-shot learning on fixed representations and full fine-tuning of the model.
3.1 Models
We build V-MoE on different variants of ViT 20
ViT-S(mall), ViT-B(ase), ViT-L(arge) and ViT- H(uge), the hyperparameters of which are described in Appendix B.5 There are three additional major design decisions that affect the cost (and potentially the quality) of our model:
Number of MoE layers. Following [9]. we place the MoEs on every other layer (we refer to these V-MoE Every-2). In addition, experimented with using fewer MoE layers, by placing them on the last-n even blocks (thus we dub these V-MoE Last-n). In Appendix [._ we observe that, although using fewer MoE layers decreases the number of parameters of the model, it has typically little impact on quality and can speed-up the models significantly, since less communication overhead is incurred.
Number of selected experts k: The cost of our model does not depend on the total number of experts but the number of selected ones per token. Concurrent works in NLP fix k 22 or = [54] 39] In our case, values of k)), while we found the total number of experts E 32 to be the sweet spot in our setting.
by default k 2 (see Figure 10 in Appendix B for the exploration of different
we use fixed buffer capacity. While this is typically regarded as downside or engineering difficuity to implement these models, we can adjust the capacity ratio to control different trade-offs. We can intentionally set it to a low ratio to save compute, using Batch Prioritized Routing (see Section 4). During upstream training, we set C = 1.05 by default to give small amount of slack without increasing the cost noticeably.
Buffer capacity C: As mentioned in Section 2.4
Note that for given trained model, the latter two-k and C--can be adjusted without further training, whereas the positioning and quantity of expert layers is effectively fixed to match pre- training..
3.2
Data
We pre-train our models on JFT-300M (7
semiiautornatcclly noisy-labeled dataset. It has 305M
training and 50 000 validation images, organised in hierarchy of 18 291 classes (average 1.89 labels per image). We deduplicate it with respect to all our validation/test sets as in previous efforts (36
Our few-shot experiments on ImageNet (i.e. ILSVRC2012) use only 1, 5, or 10 shots per class to adapt the upstream model, evaluating the resulting model on the validation set.
We also fine-tuned the pre-trained models on the full training set (ca. 1M images). We report performance in similar regime for four other datasets in Appendix [B.5] Lastly, we explore the ability to fine-tune our large models in the low-data regime by evaluating them on the Visual Task Adaptation Benchmark (VTAB) [69], a diverse suite of 19 tasks with only 000 data points per task. As well as natural image classification, VTAB includes specialized tasks (e.g. medical or satellite imagery) and structured tasks (e.g. counting or assessing rotation/distance).
3.3 Upstream results
JFT is multilabel dataset, so we measure model performance via precision @1 (see Appendix B.6 for details). Note that as in previous works 220) hyperparameters were tuned for transfer performance, and JFT precision could be improved at the expense of downstream tasks e.g. by reducing weight decay. Figure 2a shows the quality of different V-MoE and ViT variants with respect to total training compute and time. It shows models that select l 2 experts and place MoEs in the last even blocks (n = 5 for V-MoE--H, n = 2 otherwise). but the best results are achieved by V-MoF--H14 Every-2 (see Table 2 14 is the patch size). L/16's s trained for 7 or 14 epochs. See Appendix B.5 for all results.
We also checked the effect of deduplication with respect to the ImageNet training set, showing negligible (within noise) impact fhw-shot results (only 1-shot worsened, see Table 9).

---

F/14
14
H/14
IFT T:300M Precision@1 1%9
55
L16
L16
L16
L16
70
70
50
L/6
50
L/16
65
L/32
65
L/32
L/32
L/32
B/16
8/16
45
B/16
45
8/16
60
60
40
0/32
8/32
Imagele:
55
55
50
30
SC
VIT
V-MoE (Last n]
S/
V-MoE (Last n)
VIT
S VIT
V-MoE (Last n)
45
SC
V-McE (Last n)
VIT
10a
10a
10a
10a
10a
10a
10-
10
Total Training ExaFLOPs
Total Training TPUv3-days
Total Training ExaFLOPs
Total Training TPUv3-days
(a) JFT-300M
(b) ImageNet 5-shot
Figure 2: JFT-300M Precision @1 and ImageNet 5-shot accuracy. Colors represent different ViT variants, markers represent either standard ViT or V-MoFs the last even blocks. The lines represent the Pareto frontier of VIT (dashed) and V-MoR (solid) variants.
Figure 3: ImageNet Fine-Tuning Accuracy. Colors represent different VIT variants, markers represent either standard + ViT + V-MoEs the last n even blocks. Lines show the Pareto frontier of VIT (dashed) and V-MoE (solid).
ImageNet Fine- Tuning Accuracy
L/16
86
85
85
L32
L/32
84
8/16
84
8/16
VIT
V-MoE
Table Scores and 95% confidence intervals for ViT and V-MoE.
VTAB.
83
83
L/16 76.3
77.2.0.4
82
82
H/14 77.6
77.
V-MoE (Last n)
(Last n)
81
B/3
VIT
81
B/S
VIT
10a
10a
10
Total Training ExaFLOPs
Total Training TPUV3-days
Expert models provide notable gains across all model sizes, for only a mild increase in FLOPs, establishing new Pareto frontier (gray lines). Alternatively, of ViT models at lower cost (e.g. V-MoE-L/16 improves upon ViT-H/14), Similar conclusions hold for training time, which includes communication overhead of dispatching data across devices.
match improve performance
3.4 Linear few-shot results
We evaluate the quality of the representations learned using few-shot lincar transfer. Given training examples from the new dataset ((X,Y)) use the pre-trained model M to extract fixed representation M(x;; of each image. We fit a linear regression model mapping M(x;) to the one-hot encoding of the target labels Y, following 20 (see [27] Chapter 5] for background).
Figure 2b shows that the upstream gains are preserved under 5-shot ImageNet evaluation, considering both compute and time; in other words, the quality of the representations learned by V-MoE also outperforms ViT models when looking at new task. further shows the results on {1, 10)100)- -hot for some selected models, and the full detailed results are available in Appendix B.5
3.5
Full fine-tuning results
The typically most performant approach for Transfer Learning 19 consists of replacing the upstream classification head with new task-specific one and fine-tuning the whole model. Though one may expect that massive models like V-MoEs require special handling for fine-tuning, we broadly follow the standard fine-tuning protocol for Vision Transformers. We use the auxiliary loss during fine-tuning as well, although we observe that it is often not needed this step, as the router is already well trained. We explore the two sets of tasks considered therein:
5

---

Table 2: Main V-MoE & VIT models; Table 8
shows results for additional models and datasets.
Model
Parums
JFT prec@1
IN/5shot IN/lshot
IN/1Oshot
IN/Fine-t.
ExaFLOPs
TPUv3-days
VII-H/14
656M
56.68
62.34
76.95
79.02
88.08
4.27k
2.38k
V-MoF-I/16, Every-2
3.4B
57.65
62.41
77.10
79.01
87.41
2.17k
1.20k
V-MoE-H/14. Last-5
2.7B
60.12
62.95
78.08
80.10
88.23
4.75k
2.73k
V-MoE-H1/14, Every-2
7.2B
60.62
80.33
88.36
5.79k
3.47k
V-MoE-15B, Every-2
14.7B
68.66
82.78
84.29
90.35
33.9k
16.8k
NFNet-F4+
89.20
1.86k
MPL
90.20
22.5k
patch discarded.
C 0.9
C= 0.7
C= 0.5
C= 0.3
C= 0.1
Figure White patches are discarded tokens in the first layer of experts, for different capacities, using Batch Prioritized Routing (Section 4.1 with a V-MoE-H/14. See Appendix D for more examples.
Full data. We follow the setup of [20], except that we apply a dropout rate of 0.1 on the expert MLPs (as done in (22)), and we halve the number of fine-tuning steps for all datasets other than ImageNet. Figure the results on ImageNet (averaged over three runs). Here, V-MoE also performs better than dense counterparts, though we suspect the fine-tuning protocol could be further improved and tailored the sparse models. See Table for all details, including results on other datasets.
Low-data regime. On the VTAB benchmark, we use a similar setup and hyperparameter budget as 20 (but fine-tune with half the schedule length).
that, while performance similar for V-MoE-H/14, experts provide significant gains at the ViT-L/16 level, indicating that despite the large size of these models, they can still be fine-tuned with small amounts of data and no further tricks.
3.6 Scaling up V-MoE
Finally, we test how well V-MoE can scale vision models to a very large number of parameters, while continuing to improve performance. For this, increase the size of the model and use a larger pre-training dataset: JFT-3B a larger version of JFT-300M, it contains almost 3B images and is noisily annotated with 30k classes. Inspired by (68). apply the changes detailed in Appendix B.3 and train a 48-block V-MoE model, with every-2 expert placement (32 experts and 2), resulting in model with 14.7B parameters, which denote by V-MoE-15B.
successfully train V-MoB-15B, which is, as far as we are aware, the largest vision model date. has an impressive 82.78% % accuracy on 5-shot ImageNet and 90.35% when fully fine-tuned, as shown in Appendix [B.5] which also includes more details about the model. Training this model required 16.8k -Pryy-coredaay. To contextualize this result, the current state of the art ImageNet is Meta Pseudo-Labelling (MPL) [49]. MPL trains an EfficientNet-bace using ImageNet pseudo-labelling, achieving 90.2% while requiring 22.5k TPUv3-core-days.
model on unlabelled JFT-300M
4 Skipping Tokens with Batch Prioritized Routing
We present a new routing algorithm that allows the model to prioritize important tokens (corresp. patches). By simultaneously reducing the capacity of each expert, tokens. Intuitively, not every patch is equally important to classify a given image, e.g., most background patches can be dropped to let the model only focus the ones with the relevant entities.
discard the least useful
4.1 From Vanilla Routing to Batch Prioritized Routing
With the notation from Section the routing function g is applied row-wise to a batch of inputs X NPDx A batch contains N images composed of P tokens each; each row of X corresponds

---

V.MoE
V-MoE (BPR)
H-M4 (BPR)
V-MoE-H/14 (Vanilla)
80
77-300M Precision@) (%)
ImageNet 5-shot Accuracy (%)
78.0
(FT-300M Precision@: (%)
60
VIT-H/14
maggNNet 5-shot Accuracy (%)
VIT-H/14
59
77.5
50
70
77.0
V/T.H/14
50-
40
76.5
VIT-H/14
76.0
30-
56-
75.5
20
50
55
75.0
10
40
74.5
VIT-L/26
200
300
400
100
300
400
00
00
Inference GigaFLOPs//mage
0
Capacity Ratio
Figure 5: Reducing compute with priority routing. Performance vs. inference FLOPS for large models. V-MoEs with the original vanilla routing
Figure 6: Priority routing works where vanilla fails. Performance vs. inference capacity ratio for V-MoE-HH14 model with k 2. Even for large O's BPR outperforms vanilla; at low C' the difference stark. BPR is competitive with dense by processing only 15-30% of the tokens.
represented by while shows V-MoEs where BPR and mix of C'e {0.6, 0.7, 0.8) and & {1,2] are used re- duce compute. ViT models shown as x.
the D-dimensional representation of particular token of an image. Accordingly, 9(X),, denotes the routing weight for the t-th token and the i-th expert.
In all routing algorithms considered, for < j. every TOP-i assignment has priority over any TOP-- assignment. The router first tries to dispatch all Th expert choices before assigning any jth choice"
Given the TOP-i position, the default-on or vanilla-routing as used in [54 39 22], assigns tokens to experts as follows. It sequentially goes over the rows of g(X) and assigns each token to its TOP-i expert when the expert's 's buffer is not full. As result, priority is given to tokens depending on the rank of their corresponding row. While images in a batch are randomly ordered, tokens within an image follow a pre-defined fixed order. The algorithm is detailed in Algorithm[] I of Appendix C
Batch Prioritized Routing (BPR). To favour the "most important" tokens, we propose to compute priority score s(x) on each token, and sort g(X) accordingly before proceeding with the allocation. We sort tokens based on their maximum routing weight, formally s(X); max; g(X),,, The sum of TOP-k weights, i.e. s s(X): = Σ, g( (X,,,, worked equally well. These two simple approaches outperformed other options we explored, e.g., directly parameterising and learning the function s.
We reuse the router outputs as proxy for the priority of allocation. Our experiments show this preserves the performant predictive behaviour of the model, even though the router outputs primarily encode how well tokens and experts can be paired, not the token's 's "importance" for the final classifica- tion task. Figure visualizes token prioritisation with Batch Prioritized Routing for increasingly small capacities. Since all tokens across all images in the batch X compete with each other, different images may receive different amounts of compute. We summarize BPR in Algorithm in Appendix C
4.2 Skip tokens with low capacity
Batch Prioritized Routing opens the door redueing the buffer size by smartly selecting which tokens to favor. This can have dramatic impact in the computational cost of the overall sparse model. We discuss now inference and training results with C' defined in Section 2.4 in the regime C << 1.
At inference time. Prioritized routing is agnostic to how the model was originally trained. Figure 6 shows the effect of reducing compute at inference time by using BPR versus vanilla routing, on a V-MoE-H//1e model trained using vanilla routing. The difference in performance between both methods is remarkable especially for C 0.5, where the model truly starts fully dropping tokens, as
A token may however successfully assign all its TOP-K choices while another may not allocate single one. This can happen for instance if the latter selects very popular experts that out of capacity.
7

---

MoE Layer 1
Layer 9
MoE Layer
MoE Layer
1
image class
600
800
1000
20
expert
expert
expert
expert
Figure 7: Deeper routing decisions correlate with image classes. We show 4 MoE layers of V-MoE-HH14. The r-axis corresponds to the 32 experts in a layer. The y-axis are the 1 000 ImageNet classes; orderings for both axes are different across plots. For each pair (expert e, class c) we show the average routing weight for the tokens corresponding to all images with class for that particular expert e. Figure 29 includes all the remaining layers; see Appendix E.2 for details.
k: = 2. Also, BPR allows the model to be competitive with the dense one even at quite low capacities. As shown in Figure 5 for V-MoE-L/16 and V-MoE-H/14, Batch Prioritized Routing and low allow V-MoE to smoothly trade-off performance and FLOPS at inference time, quite unique model feature. More concretely, Table Thoshows V-MoE models can beat the dense VIT-H performance by using less than half the FLOPs and less than 60% of the runtime. Conversely, FLOPs cost and preserve a one- moi-point accuracy gain in ImageNet/5shot and almost three- point JFT precision one (Table TTp. Dense models generally require less runtime for the same amount of FLOPPs due to the data transfer involved in the V-MoB implementation.
match the inference
At training time. Batch Prioritized Routing can also be leveraged during training. In Appendix C we show how expert models with max-weight routing can match the dense performance while saving around 20% of the total training FLOPs, and strongly outperform vanilla with similar FLOP budget.
5 Model Analysis
Although large-scale sparse MoEs have led to strong performance 22 39 54J, little is known and understood about how the internals of those complex models work. We argue that such exploratory experiments can inform the design of new algorithms. In this section, we provide the first such analysis at this scale, which guided the development of the algorithms presented in the paper.
Specialized experts. Intuitively, routers should learn to distribute images across experts based on their similarity. For instance, if the model had three experts, and the task mainly involved three categories- -ay animals, cars, and buildings-one would expect an expert to specialize in each of those. We test this intuition, with some obvious caveats: (a) experts are placed at several network depths, (b) & experts are combined, and (c) routing happens at the token rather than the image level.
Figure 7 illustrates how many images of given ImageNet class use each expert. The plots were produced by running fine-tuned V-MoE-H Every-2 model. Interestingly, we saw similar patterns with the upstream model without finettunig. Experts specialize in discriminating between small sets of classes (those primarily routed through the expert). In earlier MoE layers we do not observe this. Experts may instead focus on aspects common all classes (background, basic shapes, colours) for example, Figure 30 (Appendix
F) shows correlations with patch location in earlier layers.
The value of routers. After training sparse MoF, is natural to study the usefulness of the learned routers, in the light of several pitfalls. For example, the routers may just act as a load balancer experts end up learning very similar functions, or the routers may simply choose poor assignments. In Appendix E.I we replace, after training, one router time with a uniformly random router. The models are robust to early routing changes while more sensitive to the decisions in the last layers.
Routing weights distributions. We analyse the router outputs in Appendix E and observe the distribution selected weights varies wildly across different mixture of experts layers.
Changing l: at inference time. We have observed expert models are remarkably flexible. Somewhat surprisingly, sparse models are fairly robust to mismatches between their training and inference configurations. In Appendix E.4 explore the effect of training with some original value of & while

---

applying the model at inference time with different k' k. This can be handy to control (decrease or increase) the amount of FLOPs per imput particular production system.
Related work
Conditional Computation. To grow the number of model parameters without proportionally increas-
ing the computational cost, conditional computation [5] 15 12 only activates some relevant parts of
the model in input-dependent fashion, like in decision trees 7 In deep learning, the activation of portions of the model can use stochastic ncurons
66] or reinforcement learning
17 53].
Mixture of Experts. MoEs 31 via router an input-dependent way. MoEs have successfully used this form of conditional
10 66 combine the outputs of sub-models known as experts
computation in a range of applications 23 30
58
55 67 An input can select either all experts [2]
or only a sparse mixture thereof as in recent massive language models [54]] 22
MoEs for Language. MoEs have recently scaled language models up to trillions of parameters. Our approach is inspired by [54] who proposed a top-k gating in LSTMs, with auxiliary losses ensuring
the expert balance [26]. (39 further scaled up this approach for transformers, showing strong gains for neural machine translation. With over one trillion parameters and one expert per input, [22 sped up pre-training compared to dense baseline (50] while showing gains thanks to transfer and distillation. [40] alternatively enforced a balanced routing by solving a linear assignment problem.
MoEs for Vision. For computer vision, previous work on MoEs 21 25
64 focused on
architectures whose scale is considerably smaller than that of both language models and our model. In DeepMoB [63] the "experts" are the channels of convolutional layers that are adaptively selected by multi-headed sparse gate. This is similar to (64 where the kernels of convolutional layers activated on a per-example basis. Other approaches use shallow MoEs, learning a single router, either disjointly [25) or jointly (2]. together with CNNs playing the role of experts. [[] further have cost-aware procedure to bias the assignments of inputs across the experts. Unlike shallow MoEs, we operate with up to several tens of routing decisions per token along the depth of the model. Scaling up routing depth was marked as major challenge in 511 which we successfully tackle in our work.
7 Conclusions
We have employed sparse conditional computation to train some of the largest vision models to date, showing significant improvements in representation learning and transfer learning. Alongside V-MoE, we have proposed Batch Prioritized Routing, which allows successful repurposing of model sparsity to introduce sparsity with respect the inputs. This can be done without further adapting the model, allowing the rs-use of trained models with sparse conditional computation.
This has interesting connotations for recent work in NLP using sparse models; recent analysis shows model sparsity is the most promising way reduce model CO2 emissions (46) and that 90% of the footprint stems from inference costs we present an algorithm which takes the most efficient models and makes them even more efficient without any further model adaptation.
This is just the beginning of conditional computation at scale for vision; extensions include scaling up the expert count, reducing dependency on data and improving transfer of the representations produced by sparse models. Directions relating to heterogeneous expert architectures and conditional variable-length routes should also be fruitful. We expect increasing importance of sparse model scaling, especially in data rich domains such as large scale multimodal or video modeling.

---

Acknowledgments and Disclosure of Funding
We thank Alex Kolesnikov, Lucas Beyer and Xiaohua Zhai for providing continuous help and details about scaling VIT models; Alexey Dosovitskiy, who provided some of the pre-trained VIT models; Ilya Tolstikhin, who suggested placing experts only in the last layers; Josip Djolonga for his early review of the manuscript; Dmitry Lepikhin for providing details about the original GShard implementation; Barret Zoph and Liam Fedus for insightful comments and feedback; James Bradbury, Blake Hechtman and the rest of JAX and TPU team who helped us running our models efficiently, and many others from Google Brain for their support.
References
[1] A. Abbas and Y. Andreopoulos. Biased mixtures of experts: Enabling computer vision inference under data transfer limitations. IEEE Trans. Image Processing, 29:7656-7667, 2020.
[2] K. Ahmed, M. H. Baig, and Torresani. Network of experts for large-scale image categorization.
In ECCV, 2016.
[3] A. Arnab, M. Dehghani, G. Heigold, C. Sun, M. Lučié, and C. Schmid. ViVIT: A video vision transformer. arXiv preprint arXiv:2103. 15691, 2021.
[4] E. Bengio, P.-L. Bacon, Pineau, and Precup. Conditional computation in neural networks for faster models. arXiv preprint arXiv: 15 1.06297, 2015.
[5] Y. Bengio. Deep learning of representations: Looking forward. In International Conference on Statistical Language and Speech Processing, pages 1-37, 2013.
[6] Y. Bengio, N. Léonard, and A. Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint 3xXXv:330..332 2013.
[7] L. Breiman, J. Friedman, C. J. Stone, and R. A. Olshen. Classification and regression trees.
CRC press, 1984.
[8] A. Brock, S. De, S. L. Smith, and K. Simonyan. High-performance large-scale image recognition without normalization. arXiv preprint wxXi:22102.06171, 2021.
[9] B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, Neelakantan,
P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. arXiv preprint a0Xiv:2005. 14165, 2020.
[10] K. Chen, L. Xu, and H. Chi. Improved learning algorithms for mixture of experts in multiclass 1290:122291122, 1999.
classification. Neural networks,
[11] G. Cheng, J. Han, and X. Lu. Remote sensing image scene classification: Benchmark and state of the art. Proceedings of the IEEE, 105(10):185--188, Oct 2017.
[12] K. Cho and Y. Bengio. Exponentially increasing the capacity-to-computae tional computation in deep learning. arXiv preprint arXiv: 1406.7362, 2014.
ratio for condi-
[13] M. Cimpoi, S. Maji. I. Kokkinos, Mohamed, and A. Vedaldi. Describing textures in the wild.
In Computer Vision and Pattern Recognition (CVPR), 2014.
[14] E. D. Cubuk, B. Zoph, J. Shlens, and Q. Le. Randaugment: Practical automated data augmenta- tion with reduced search space. H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, 2020.
[15] A. Davis and I. Arel. Low-rank approximations for conditional feedforward computation in deep neural networks. arXiv preprint arXiv: 1312,4461, 2013.
[16] Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. ImageNet: A large-scale hierarchical image database. In CVPR, 2009.
[17] I.. Denoyer and P. Gallinari. Deep sequential neural network. arXiv preprint 100 2014.
[18] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In NAACL, 2019.
[19] G. S. Dhillon, P. Chaudhari, A. Ravichandran, and S. Soatto. A baseline for few-shot image classification. In ICLR, 2020.
10

---

[20] A. Dosovitskly, L. Beyer, A. Kolesmikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, Uszkoreit, and N. Houlsby. An image worth 16x16 Transformers for image recognition at scale. In ICI,R, 2021.
words:
[21] D. Eigen, M. Ranzato, and I. Sutskever. Learning factored representations in a deep mixture of experts. arXiv preprint arXiv:1312.4314, 2013.
[22] W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint AXii:2101.03%00 2021.
[23] D. M. Gavrila and S. Munder. Multi-cue pedestrian detection and tracking from a moving vehicle. International journal of computer vision, 73(1):41-59, 2007.
[24] A. Geiger, P. Lenz, and R. Urtasun. Are we ready for autonomous driving? The KITTI vision benchmark suite. In Computer Vision and Pattern Recognition (CVPR), 2012.
[25] Gross, M. Ranzato, and A. Szlam. Hard mixtures of experts for large scale weakly supervised vision. In CVPR, 2017.
[26] J. V. Hansen. Combining predictors: comparison of five meta machine learning methods. Information Sciences, 119(1-2)91--100 1999.
[27] Hastie, Tibshirani, and Friedman. The elements of statistical learning: data mining, inference, and prediction. Springer, 2017.
(28) P. Helber, B. Bischke, A. Dengel, and D. Borth. EuroSAT: A novel dataset and deep learning benchmark for land use and land cover classification. IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, 12(7):2217-222 2019.
[29] D. Hendrycks and K. Gimpel. Gaussian error linear units (gelus). arXiv preprint 1606v:606.08440 08410, 2016.
[30] Y. H. Hu, S. Paireddy, and W. J. Tompkins. A patient-adaptable ECG beat classifier using mixture of experts approach. IEEE Trans. Biomedical Engineering, (499):991900) 1997.
[31] R. A. Jacobs, M. Jordan, S. Nowlan, and G. E. Hinton. Adaptive mixtures of local experts. Neural Computation, 3(1):79-88 1991.
[32] Z. Jiang, Q. Hou, L. Yuan, D. Zhou, X. Jin, A. Wang, and J. Feng. Token labeling: Training a 85.5% t accuracy vision transformer with 56m parameters on imagenet. arXiv preprint arXiv:2104. 10858, 2021.
[33] J. Johnson, B. Hariharan, L. van der Maaten, F.-F. Li, C. Lawrence Zitnick, and R. Girshick. CLEVR: A diagnostic dataset for compositional language and elementary visual reasoning. In Computer Vision and Pattern Recognition (CVPR), 2017.
[34] M. I. Jordan and R. A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural computation, 6(2):181-214, 1994.
[35] Kaggle and EyePacs. Kaggle diabetic retinopathy detection, 2015.
[36] A. Kolesnikov, L. Beyer, X. Zhai, J. Puigcerver, J. Yung. S. Gelly, and N. Houlsby, Big transfer (BiT): General visual representation learning. ECCV, 2020.
[37] A. Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[38] Y. LeCun, F. J. Huang, and L. Bottou. Learning methods for generic object recognition with invariance to pose and lighting. In Computer Vision and Pattern Recognition (CVPR), 2004.
[39] D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen. GShard: Scaling giant models with conditional computation and automatic sharding. In ICLR, 2021.
[40] M. Lewis, S. Bhosale, T. Dettmers, N. Goyal, and L. Zettlemoyer. Base layers: Simplifying training of large, sparse models. arXiv preprint arXiv:2 16716, 2021.
[41] F.F.. Li, R. Fergus, and P. Perona. Learning generative visual models from few training examples: An incremental Bayesian approach tested on 101 object categories. Computer Vision and Pattern Recognition Workshop, 2004.
[42] I. Matthey, Higgins, D. Hassabis, and A. Lerchner. dSprites: Disentanglement testing sprites dataset, 2017.
11

---

[43] Y. Netzer, T. Wang. A. Coates, A. Bissacco, and A. Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning 201 1, 2011.
[44] M.-E. Nilsback and A. Zisserman. Automated flower classification over large number of classes. In Sixth Indian Conf. on Computer Vision, Graphics & Image Processing, 2008.
[45] O. M. Parkhi, A. Vedaldi, A. Zisserman, and C. Jawahar. Cats and dogs. In CVPR, 2012.
[46] D. Patterson, J. Gonzalez, Q. Le, C. Liang, L.-M. Munguia, D. Rothchild, D. So, M. Texier, and
J. Dean. Carbon emissions and large neural network training. arXiv preprint arXiv:2104. 10350, 2021.
[47] S. Pavlitskaya, C. Hubschneider, M. Weber, R. Moritz, F. Huger, P. Schlicht, and M. Zollner. Us- ing mixture of expert models to gain insights into semantic segmentation. In CVPR Workshops, 2020.
[48] M. E. Peters, M. Neumann, M. lyyer, M. Gardner, C. Clark, K. Lee, and L. Zettlemoyer. Deep contextualized word representations. In NAACL, 2018.
[49] H. Pham, Z. Dai, Q. Xie, M.-T. Luong, and Q. Le. Meta pseudo labels. arXiv preprint arXiv:2003. 10580, 2020.
[50] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint
arXiv: 1910. 10683, 2019.
[51] P. Ramachandran and Q. Le. Diversity and depth in per-example routing models. In ICLR, 2018.
[52] C. Rosenbaum, I. Cases, M. Riemer, and T. Klinger. Routing networks and the challenges of modular and compositional computation. arXiv preprint 1rX4 12774, 2019.
[53] C. Rosenbaum, T. Klinger, and M. Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. arXiv preprint arXiv: 2017.
[54] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Le, G. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In ICLR, 2017.
[55] C. Sminchisescu, A. Kanaujia, Z. Li, and D. Metaxas. Learning to reconstruct 3D human motion from Bayesian mixtures of experts. A probabilistic discriminative approach. Dept. Comput. Sci., Univ. Toronto, Tech. Rep. CSRG-502, 2004.
[56] E. Strubell, A. Ganesh, and A. McCallum. Energy and policy considerations for deep learning in NLP. arXiv preprint 190iv
2019.
[57] C. Sun, A. Shrivastava, S. Singh, and A. Gupta. Revisiting unreasonable effectiveness of data in deep learning era. In ICCV, 2017.
[58] J. Tani and S. Nolfi. Learning perceive the world as articulated: an approach for hierarchical learning in sensory motor systems. Neural Networks, 2(7-8):1131-144 1999.
[59] H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, and H. Jégou. Training data-efficient image transformers & distillation through attention. arXiv preprint arXiv:2012. 12877, 2020.
[60] H. Touvron, M. Cord, A. Sablayrolles, G. Synnaeve, and H. Jégou. Going deeper with image transformers. arXiv preprini arXiv:2 100.003..3333 /7239, 2021.
[61] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need, In NeurIPS, 2017.
[62] B. S. Veeling, J. Linmans, J. Winkens, T. Cohen, and M. Welling, Rotation equivariant CNNs for digital pathology. In Medical Image Computing and Computer Assisted Intervention (MICCAI), 2018.
[63] X. Wang, F. Yu, L. Dunlap, Y.-A. Ma, R. Wang, A. Mirhoseini, T. Darrell, and J. E. Gonzalez. Deep mixture of experts via shallow embedding. In Uncertainty in Artificial Intelligence, 2020.
[64] B. Yang. G. Bender, Q. V. Le, and J. Ngiam. Condconv: Conditionally parameterized convolu- tions for efficient inference. arXiv preprint urXiv:1904.044 04971, 2019."
[65] I.. Yuan, Y. Chen, Wang, W. Yu, Y. Shi, F. E. Tay, Feng, and Yan. Tokenn-to-token vit: Training vision transformers from scratch imagenet. arXiv preprint arXiv:2101.11986, 2021.
12

---

[66] S. E. Yuksel, J. N. Wilson, and P. D. Gader. Twenty years of mixture of experts. IEEE transactions on neural networks and learning systems, 2378)1117-1199, 2012.
[67] A. Zeevi, R. Meir, and R. J. Adler. Time series prediction using mixtures of experts. In NeurIPS, 1997.
[68] X. Zhai, A. Kolesnikov, N. Houlsby, and L. Beyer. Scaling vision transformers, 2021.
[69] X. Zhai, J. Puigcerver, A. Kolesnikov, P. Ruyssen, C. Riquelme, M. Lucic. Djolonga, A. Pinto, M. Neumann, A. Dosovitskiy, L. Beyer, O. Bachem, M. Tschannen, M. Michalski, O. Bousquet, S. Gelly, and N. Houlsby. A large-scale study of representation learning with the visual task adaptation benchmark. arXiv preprint arXiv: 0e11199004000 2019.
[70] X. Zhai, J. Puigcerver, A. Kolesnikov, P. Ruyssen, C. Riquelme, M. Lucic, J. Djolonga, A. S. Pinto, M. Neumann, A. Dosovitskiy, al. A large-scale study of representation learning with the visual task adaptation benchmark. arXiv preprint 1xX::1910.048800 2019.
13


Reviews:
[{"review_id": "uWjdXhrHm9a", "reviewer": "Reviewer_wGNF", "summary": "This paper aims at scaling up Vision Transformer by Sparsely-gated Mixture of Experts networks (MoEs). Finally, they developed a 15B parameter model, and by using extra 3B data, they achieve 90.24% top-1 accuracy on ImageNet.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Pros:\n1. The 15B model get very high results were obtained using MoE, proving the potential of MOE in training very large models in computer vision.\n2. Well written and easy to follow.\n3. The proposed Batch Prioritized Routing although sounds engineering but useful.\n\nCons:\n1. The results are not reproducible since the data is not public and computer cost is too high. It will be more convincing and appreciated if the author can do more experiments on some public dataset, like ImageNet-1k/22k.\n2. Some work has found many of the labels of ImageNet validation set are incorrect (“Pervasive Label Errors in Test Sets Destabilize Machine Learning Benchmarks”). I'm curious about the reliability of the 90.24% accuracy and if it is just hacking the validation set.\n3. The few-shot setting makes me confused since the pretraining stage has already seen a large quantity of images with labels of the same classes. Although the main body is fixed and the linear head is re-initialized, but it is not exactly matched the \"few-shot\" setting.\n4. Data leaking of the imagenet-1K validation data is also a potential issue since the de-duplication method is not so reliable when the data number is so large(3B).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Please refer to Cons in Main Review.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "dUR1WgM5YNe", "reviewer": "Reviewer_5sB5", "summary": "The paper proposes a sparse Mixture of Experts transformer model for vision tasks (mainly image classification). The authors show that the model has comparable performance to state-of-the-art models and potential advantages in computational requirements due to its sparsity. The model is mostly based in previously proposed sparse MoE models for language with a few tweaks.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "**Originality**\n\nThe paper adapts the idea of sparse MoE from previous work in NLP and adapts it to vision transformers, with small modifications, mainly to make the amount of compute adaptative to each processed example. Therefore, it is a relatively straightforward application of a previously proposed idea to a different modailty.\n\n**Clarity**\n\nThe paper is well-written. The model formulation is clear, the experiments are focused on evaluating important aspects of the model and overall it is easy to follow the main ideas in this submission.\n\n**Quality**\n\nThe paper claims are sound, with a clear adaptation of the sparse MoE idea to vision transformers and modifications that are well motivated and empirically justified by the experiments. \n\n**Significance**\n\nThe paper is well-written and the experimental section adequately assesses the model performance. However, the relatively low novelty of the model, the fact that the code is proprietary and the fact that the model requires extensive computational resources to be trained, and just comparable performance to other current approaches will reduce its impact.\n\nOverall, as a successful implementation of the sparse MoE idea for vision, I believe it is of interest to the community and therefore it should be accepted, although I believe it will have limited practical relevance to the community.\n\n**Review updates**\n\nAfter reading the rest of the reviews and the authors' response, I stand by my previous rating. The technical novelty of the paper remains very limited, and the approach is very expensive to train while comparable to other alternatives, which will limit its adoption by the community. However, adapting sparse MoE for vision has some value, and therefore I stand by my rating of 6.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper includes a limited assessment of the negative societal impacts, centered on the energy impact of training and running these kind of models. Additional discussions on potential malicious uses of these models would improve this section.", "review_text": "", "extra_scores": {"ethical_concerns": "There are no ethical concerns with the paper.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "CeqwML-YiW", "reviewer": "Reviewer_QFfS", "summary": "This paper extends Mixture of Experts transformers to Computer Vision by training various configurations of MoE transformers on JFT dataset and finetuning on Imagenet and VTAB. Along with this, they show very interesting analysis and modification to the router that leads to a nice trade off between inference performance and runtime. \n\nOverall, I think the paper is well written and the experiments are well executed, which will benefit the community. There are some suggestions in the Main Review, but apart from that, I think the paper has enough merits to be accepted.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "**Contributions of the paper**:\n\n1) The paper extends mixture of experts (MoE) transformers introduced in gshard (and switch) paper to computer vision. \n2) They train various configurations of MoE vision transformers on JFT300M dataset and fine-tune it on various downstream datasets in multiple settings of few shot and full shot.\n3) They also scale up the model to 15B parameters trained on 3B images from JFT dataset which achieves good performance across a wide variety of downstream tasks.\n4) They did lot of really insightful analysis around (a) routing decisions, (b) varying values of k (#experts) during inference and finetuning etc\n5) They also introduced a simple extension of the gating router to prioritize tokens based on their importance\n\n\n**Originality**:\n1) Overall their architecture itself is mostly similar to the one introduced in gshard. There are some issues about the novelty specially on the architecture side but I think this is still a good contribution because MoE transformers can be unstable in training and finetuning the large models is hard. \n2) Overall their findings that MoE transformers work in Vision also successfully will open up some more work in this direction to push more compute efficient (or more specifically conditional compute) architectures. \n\n\n**Quality**\n1) The analysis in the paper is quite clear and very useful. The experiment setup is very detailed and most of the analysis is followed by a clear takeaway.\n2) Apart from the few suggestions made in the suggestion section down, I think most of the claims are backed by thorough experiments.\n\n**Clarity**\n1) The paper is well organized and well written.\n2) I'd suggest authors to focus a bit more around why experts beyond 32 didn't work or share/ discuss failure cases but apart from that most sections are well motivated. \n\n**Significance**\n1) Overall I think that this work will be useful for the CV community and will provide a good direction for further work in MoE transformers in computer vision.\n2) I would find some issue with reproducibility of the work. I can understand the data being proprietary but not releasing code can limit the impact on the community overall.\n\n\n**Questions / comments**:\n1) Authors say that they found sweet spot at 32 experts, but I couldn’t find further discussion around it. Did you see divergence issues with models with larger experts or were they not getting good results? In NLP, at least on the upstream task, experts have been scaled to 2048 without much issue, do you have any thoughts on why it seems different for computer vision? \n2) The paper defines and uses slightly different auxiliary loss than gshard, curious if there are any results with gshard auxiliary loss? \n3) Figure 16’s two entries of L/16 are confusing, I could find that they are trained for 7 and 14 epochs in the appendix, but would be ideal its more explicit somewhere either in the main text or caption of those images. Also L/16 7 epochs is mid checkpoint of 14 epoch model or actually trained from scratch such that LR schedule is fixed for 7 epochs?\n4) Table 6 and section B.3 says 0.03 weight decay for the “kernel”, what is it referring to exactly.\n5) Any reasons to use inverse sqrt schedule for the biggest model? Was it just to avoid mentioning total updates upfront?\n6) Section B.3 says V-MoE 15B, does those changes apply to all MoE models or just the 15B? I.e did you train other MoE models with adafactor also apart from 15B\n\n\n**Suggestions for possible additional analysis**:\n1) Any small scale experiments for comparing gradient moving average in float32 vs bf16\n2) Any small scale experiments to compare adafactor vs adam on MoE vision transformer? \n3) Add how many epochs / training images during training on JFT3B did V-MoE-15B go over?\n\n**Typos / broken citation**:\n1) broken citation [54] in table 3 appendix?\n2) Typo in Table 8 (B.5) 15B model for imagenet shows 94.24%, it is 90.24%, right? \n3) Conclusions, Line 311: “Improving transfer of the representations sparse models” -> “Improving transfer of the representations of the sparse models”?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, I don't see much negative societal impact of this work apart from general scaling model sizes which is address by the authors.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "~15-20 hours"}, "extra_text": {}}, {"review_id": "9v26uvTEOOP", "reviewer": "Reviewer_kbeB", "summary": "This paper proposed a Vision MoE motivated by MoE used in NLP, which helps to scale vision models. The novel routing algorithm leads to dynamic per-image computation. Extensive experiments show its powerful performance while requiring half of the inference time comparing with SOTA. Amazingly, the largest model attains 90.35% top-1 accuracy on ImageNet.\n\n", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Strength\n\n- This paper firstly presents MoE for vision transformer, which helps to scale models. It outperforms other dense models while reducing most inference time. And the largest model even achieves 90.35% top-1 accuracy. \n- Extensive experiments and figures clearly show how it works and demonstrates its powerful performance.\n- The effective and novel routing algorithm further helps to allocate per-image computation dynamically.\n\nWeakness\n\n- Although experiments show the effectiveness of V-MoE, the idea is not surprising enough as MoE has been popular use in NLP. Can MoE be generalized to CNN-based models?\n- The authors demonstrate that the performance drops quickly when using the longer random router. However, the experiments are conducted for testing, how about using a random router in training?\n- It seems that the parameters are not efficiently used in MoE. For example, the **ViT-L/16 (323.1M, top-1 87.12%)** is much better than **V-MoE-B/16, Last 2 (393.3M, top-1 85.39%)** in Table8. Should parameters efficiency also be considered when scaling models besides computation efficiency?\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "No", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper extends transformer Mixture of Experts to vision and achieves very competitive performance on ImageNet.

Most of the experiments use a dataset that is proprietary which prevents the work from being reproducible. Additionally, the code won’t be release. This potentially limits the impact of this work.

However, this paper demonstrates that MoE can successfully be applied to vision which is of interested to the community.  I therefore support acceptance.

Decision: Accept (Poster)