Task-customized Masked Autoencoder via Mixture of Cluster-conditional Experts
Authors: Zhili LIU
Kai Chen
Jianhua Han
Lanqing HONG
Hang Xu
Zhenguo Li
James Kwok
Venue: ICLR 2023 notable top 25% (2023)

Abstract: Masked Autoencoder (MAE) is a prevailing self-supervised learning method that achieves promising results in model pre-training. However, when the various downstream tasks have data distributions different from the pre-training data, the semantically irrelevant pre-training information might result in negative transfer, impeding MAE’s scalability. To address this issue, we propose a novel MAE-based pre-training paradigm, Mixture of Cluster-conditional Experts (MoCE), which can be trained once but provides customized pre-training models for diverse downstream tasks. Different from the mixture of experts (MoE), our MoCE trains each expert only with semantically relevant images by using cluster-conditional gates. Thus, each downstream task can be allocated to its customized model pre-trained with data most similar to the downstream data. Experiments on a collection of 11 downstream tasks show that MoCE outperforms the vanilla MAE by 2.45\% on average. It also obtains new state-of-the-art self-supervised learning results on detection and segmentation.

Published as a conference paper at ICLR 2023
TASK-CUSTOMIZED MASKED AUTOENCODER VIA MIXTURE OF CLUSTER-CONDITIONAL EXPERTS
Zhili Liu 12, Kai Chen', Jianhua Han", Lanqing Hong", Hang Xu²,
Zhenguo Li James T. Kwok
Department of Computer Science and Engineering, Hong Kong University of Science and Technology
Huawei Noah's Ark Lab
(zhili 1i liu, kai. chen ( @connect.......c..c............. ust.hk
(hanjianhua4, hongla ag, xu.hang, li. zhenguo)@huawea com jamesk@cse sk@cse. www.cseesst.....eeeeeeeeeeee
ABSTRACT
Masked Autoencoder (MAE) is prevailing self-supervised learning method that achieves promising results in model pre-training. However, when the various downstream tasks have data distributions different from the pre-training data, the semantically irrelevant pre- vre-training information might result in negative transfer. impeding MAE's scalability. To address this issue, we propose novel MAE- based pre-training paradigm, Mixture of Cluster-conditiona Experts (MoCE), which can be trained once but provides customized pre-training models for diverse downstream tasks. Different from the mixture of experts (MoE), our MoCE trains each expert only with semantically relevant images by using cluster-conditiona gates. Thus, each downstream task can be allocated to its customized model pre- trained with data most similar to the downstream data. Experiments on a collection of 11 downstream tasks show that MoCE outperforms the vanilla MAE by 2.45% on average. It also obtains new state-of-the-ar self-supervised learning results on detection and segmentation.
INTRODUCTION
Self-supervised learning (SSL), which learns effective transferable representations without human annotations. has become a prevailing model pre-training paradigm (He et al. 2020; CCen et al.) 2021a Bao al.
2022). Currently, the most prevalent SSL method is the Masked Autoencoder
(MAE) (He et al. 2022), which constructs supervision signals from raw image data by masking
random input patches and then reconstructing the missing pixels. This simple strategy has proved
efficient in the training of large-scale models. For example, ViT (Dosovitskiy et al. 2021 shows
impressive performance on popular benchmarks such as the ImageNet (Dengg et al.
2009) How-
ever, does MAE really scale well for various downstream tasks (Deng et al. 2009 Lin et al. 2014
Zhou et al. 2019: Han et al. 2021 Li et al. ?022a)
Preliminary studies (in Section 3.1 show that the MAE indeed suffers from negative transfer (Liu
et al., 2022) when transferring to downstream tasks with very different semantics. [(]] that on 9 of 11 downstream tasks, MAE pre- trained on the full ImageNet data is outperformed by the one that pre-trained on only the semantically relevant data subsets. Hence, using pre-training data that are semantically irrelevant can hurt transfer performance.
The above observation motivates the need for task-customized pre-training. A promising model for this is the Mixture of Experts (MoE) (Shazeer et al. 2017| Riquuemee al. 2021). which uses multi-expert architecture to provide customized models for different input tokens. However, unlike supervised pre-training, self-supervised learning lacks semantic labels, and thus the experts differ more on low-level information than semantics, referring to Figure [bb] Experiments in Section 4.2 show that a naive adoption MoE to the MAE has inferior performance. Since various downstream tasks contain different semantics, semantic-related experts may be preferred.
We refer ImageNet-1K ImageNet if specified in this paper.

---

Published as a conference paper at ICLR 2023
01-
Transfee accuracy (%)
#1,
90-
M2. oxpert
85-
80-
Eame semonics Diflerent soparts 41 M2
73-
85. oxgror
Diflurenc semar riccs Same experts ay
65
Ahondd
Care
SUN
DTD
Flawos
pood
reo
Clu
Caltedh
voc
C)
(a) Negative transfer phenomenon on MAE.
(b) Problem with TokenMoE.
Figure 1: (a) Transfer performance of MAES pre-trained on Split-A (blue), Split-B (red) and full ImageNet data (white). Only two of the eleven downstream tasks benefit from using the full Ima- geNet data for pre-training (more details Section 3.1p. (b) TokenMoF uses pixel RGB values as reconstruction targets. Thus, tokens with similar pixel values tend to be routed to the same expert, leading to two types of mistakes: (i) same semantics but routed to different experts, (ii) different semantics but routed to the same expert.
In this paper, we propose the Mixture of Cluster-conditional Expert (MoCE), a novel paradigm to achieve task-customized self-supervised pre-training by data clustering and explicitly training each expert with images of similar semantics. The MoCE procedure has three stages. First, we cluster the whole dataset by using a pre-trained, dense MAE model. We then construct the MoCE with multi-expert structure. Each expert is trained using clusters selected by routing tokens based on cluster embedding (instead of token embedding). To stabilize training and enhance confidence of the gate results, a regularization loss is proposed. Finally, with the arrival of a downstream task, we propose a search procedure to select the closest cluster. Empirically, the proposed MoCE shows superior performance over MAE collection of 11 downstream tasks. Besides, one only a MoCE sub-model on deployment, thus saving inference time and model capacity.
To summarize, main contributions are:
1. We systematically analyze the negative transfer phenomenon of MAE, and show that naively adopting the MoE to MAE cannot improve transfer performance of downstream tasks.
2. We propose the MoCE, which trains each expert with semantics-awand aware clusters so that similar clusters can be routed to the same expert.
3. We demonstrate effectiveness of the proposed MoCE collection of 11 downstream tasks, and achieve up to 2.45% performance improvement in Top-1 accuracy. State-of- the-art self-supervised results also achieved on the detection and segmentation tasks. To the best of our knowledge, this the first work that achieves state-of-the-ar transfer performance by training vision MoE models with ImageNet under the SSL setting.
2 RELATED WORK
Self-supervised Learning. Previous works mainly focus on the design of pretext tasks with im- age transformations (Doersch al., 2015 Gidaris et al. 2018), inpainting (Pathak al. 2016). colorization (Zhang et al. 2016), contrastive learning (Chen et al. 2020 He et al. 2020 Grill et al.)
2020 Caron et al.] 2020 Radford et al.
102
Yao et al. 2022b), and for specific downstream
la Chen et al. 202 la Yao et al.) 2022a). Motivated by the design of BERT (Devlin et al. | 2018). masked image modeling (MIM) is recently proposed to
tasks (Wang et al.) 2020 Xie et al. 2020
2021a
learn by reconstructing masked images. BEiT (Bao et al. 2022) is the pioneering work that predicts
visual tokens generated by pre-trained tokenizor (Radford et al. 2021a). SimMIM (Xic et al.)
2021ch simplifies the framework by directly utilizing the pixel RGB values reconstruction targets.
2

---

Published as a conference paper at ICLR 2023
MAE (He et al., 2022) proposes asymmetric encoder-decoder architecture for better training ef- ficiency. MixedAE (Chen et al. 2023) further explores image mixing for object-aware pre- -reitiaining... In this paper, we will focus the MAE due to its effectiveness and efficiency.
While self-supervised learning methods have achieved improved transfer performance, most of them only provide unified representation to various downstream tasks. This may suffer from negative transfer as demonstrated in Section 3.1 The work most relevant to ours is SDR (Liu et al.] 2022). which trains 256 subnets with 256 disjoint ImageNet subsets simultaneously. However, this paper differs from SDR in three ways: (i) the mapping from subsets to subnets in SDR is randomly selected and fixed during pre- tre-training while MoCE achieves self-adaptive mapping with ctuster-condtional
gates; (ii) Progressive training is required in SDR, while MoCE enjoys one-time time end-to-end training; (iii) During the transfer process, SDR uses brute force to select the best sub-model, while MoCE reuses the clustering module to achieve more efficient selection.
Mixture of Experts. The mixture of experts (MoE) has a long history (Jacobs et al. 1991] [ Jordan
& Jacobs
1994 Shazeer et al. (017). Recently, is considered as an effective tool for model scale-
up in natural language processing (Lepikhin et al. 2020 Fedus et al. 2021 Yang et al. 2021 Lewis
al.
202T). With the growing interest of the Vision Transformer (Dosovitskiy et al. 202] Liu
et al.
2021 Wang et al., [221] Xie et al.) (02Tb), MoE for vision (Riquelme et al. [021] Wu et al.,
2022
is also explored recently. However, there is still no self-supervised MoE model that can be
trained on medium-sized datasets such as the ImageNet-lk.
Kudugunta et al. (202T); ;: Ma et al. (2018) regard the MoE as multi-task learning model, and use it for muuli-languaaee translation and recommendation systems, respectively. In this paper, we show that for self-supervised learning on images, an additional clustering component is crucial in the learning of highly performant MoE model. Moreover, while the downstream tasks should follow the pre-training task in (Kudugunta et al. [022]] Ma et al.) 2018), the MoCE can be used with any downstream task due to its unsupervised pre-training, Puugcerver et al. (2020) shares a similar setting with us, but their model is pre-trained in a supervised learning manner. Moreover, their mapping between experts and data is pre-defined and fixed during training, while that for the MoCE is learned dynamically and achieves better performance.
Multi- Taski Learning aims to learn a model that is appropriate for multiple tasks. sharing, which uses shared backbone with multi-heads for the different tasks, has been shown to be effective on time series, language and graph data (Liu et al., [2199] Hu et al., 2019) McDermott al., (021]. Gao et al. (202I) claims that the network design may further benefit from the use of task relationships, and trains masks for different tasks. However, they require the task information
be available during model training, which is not possible for downstream tasks in SSL pre-training.
3 PROPOSED METHOD
In this section, we first empirically demonstrate the negative transfer phenomenon in MAE (Sec- tion We then discuss the limitations of adopting TokenMoB (Riquelme et al. 202T) with MAE (Section 3.2). and propose the Mixture of Cluster- conditioal Experts (MoCE), novel paradigm achieving customized pre- gre-training for various downstream tasks (Section 3.3
3.1 NEGATIVE TRANSFER IN MASKED AUTOENCODER
In this section, we evaluate the transfer performance of MAE models pre-trained with data of dif- ferent semantics on various downstream tasks. As in (Huh et al. 2016 Liu et al., 2022). we first split the ImageNet data into two disjoint subsets, Split-A and Split-B, based on the labels' semantic dissimilarities in the WordNet tree (Miller 1998). Split-A mainly contains inanimate objects (such as cars and airplanes), while Split-B primarily involves organisms (such as plants and animals). We then pre-train MAES on Split-A, Split-B and the full ImageNet without data annotation, and evaluate the three resulting models on 11 downstream tasks. See more implementation details in Section 4.1
As shown in Figure (a)) the MAE tre-trained with Split-A performs best on Aircraft, Cars, SUN397 and DTD, while the MAE pre-trained with Split-B performs best on Flowers, Food, Pets, Cifar- 10 and Cifar-100. Only two of the eleven tasks (Caltech and VOC) benefit from using the full
3

---

Published as a conference paper at ICLR 2023
Clooe Inforation
Expart
Export
heework
Galing
Expert B
Clustering
Cusiig Network
Iepai
Expart
Inpul
Expert
TakenMoE
Clutees
(o)
Figure 2: Model design comparison between (a) TokenMoE (Riquelme et al., 2021) and (b) MoCE. Both methods utilize the multi-expert architecture with the main difference about the input of the gating network. MoCE adopts the corresponding cluster embedding of the current token as Eqn. instead of the token embedding in Eqn. 3.2 Therefore, each expert can be trained by semantically similar images to alleviate the negative transfer phenomenon.
ImageNet data. This suggests that for tasks whose semantics are close to inanimate objects, adding pre-training data from Split-B is not useful, and vice versa for tasks whose semantics are close to organisms. To conclude. the introduction of semantically irrelevant pre-training data may impede transfer performance for downstream tasks. This negative transfer phenomenon motivates us to develop an efficient and automatic paradigm for task-customized pre-training.
3.2 EXPLORING TOKENMOE TO MASKED AUTOENCODER
Overview of TokenMoE. TokenMoB (Riquelme et 2021 is successful customized super- vised pre-training model built upon the VIT (Dosovitskiy et al. 2021). which mainly consists of transformer blocks with alternating multi-head self-attention (MSA) and multi-layer perceptron (MLP). Specifically, the TokenMoE converts several transformer blocks to Mixture of Expert (MoE) blocks by expanding the MI.P layer N times, each of them is considered as an expert (denoted as F.(.), 1= 1,2,... N). Conditional computation on the N experts is controlled by a gafe, which is a linear layer whose input is the token embedding æ, and the output is the top-K probabilities on the experts: G(æ)) Topk (ap (W x E)), where K is the number of experts to be activated, W, is the gate parameter, is the softmax function, and N (0, N TopK() (-) returns the K largest entries of o (Wyx y + unchanged but set the others to zero. Thus, each token is routed to its corresponding experts. The final output is represented
N E()
As in (Riquelme al. 2021), importance loss and load loss are also used to enforce balanced use of the experts. Unless otherwise specified, we set K = and N = 8 in all our experiments.
Limitation of TokenMoE. As will be shown in the experimental results (Table 3). naively adopt- ing TokenMoE to the MAE cannot improve performance, even with intense hyper-parameter tuning and data augmentations (e.g.. Repeat Augment (Hoffer et al. 2020) and RandAugment (Cubuk al. 2020) with larger magnitude). Figure ((a)ree((aauue ((a))((ow shows the routing heatmaps of the pre-trained TokenMoE model. As can be seen, the routing process has little correlation with the ImageNet labels. More- over, expert 3 is selected most of the time (91.8% of the classes). This degenerates the multi-expert network into single-expert network. As demonstrated in Figure [(b) we speculate that this is due to the use of low-level pixel values (instead of semantic class labels in the original TokenMoE) as reconstruction targets. This also observed in et al. (2022b).
3.3 MIXTURE CLUSTER-CONDIITOONE EXPERTS
To address the limitations of TokenMoE, propose the Mixture of Cluster-conditional Experts (MoCE), which trains each expert in a semantic-aware manner. The procedure consists data clustering, architecture and gate design, and deployment.

---

Published as a conference paper at ICLR 2023
Data Clustering. To train each expert semantically, clustering procedure is first performed to simulate the label partitioning in Section 3.1 With a pre-trained MAE model, we collect all the image features f. 's (normalized to unit length || 1), and represent the feature matrix as F = J1,J2:.... fm] E RRdx,, where n is the number of images and d is the dimension the feature. The learnable cluster centroids represented C C2.... Cm. Rdxmm (with ||c 1), where m is the desired number of clusters. The assignment of feature to clusters is computed as A F Following Asano et al. (2019), let Q gmxx be the posterior distribution of clustering, whose objective is
- 72
max Tr(Q" A) + (H)Q) s.t. QI = m 1m Q"1m Q
(2)
where 1m is the m-dimensional vector of all ones, H is the entropy function, and the constraints force the clustering results to be balanced. Q and C are optimized iteratively. For given C, Q is solved by the Sinkhorn-Knopp algorithm (Cuturi 2013); while for a given Q, C is obtained by minimizing the cross entropy between and A with SGD. We take the final C' and as the cluster centroids and clustering assignments, respectively. The implementation details are in Appendix A.
The whole network is trained on the full ImageNet data, with each expert trained by images from selected clusters decided by the MoCE gates' routing results. As on average each data cluster has only a fraction of 1/ K of the original sample size, the training time of each expert is also K times shorter than the other parameters with dense modeling (e.g.. MSA parameters (Riquelme et al., 202I)), we further adopt distillation loss Ldistali. which is defined as the lz distance between features generated by the whole network and each expert. This loss function can be formulated as
Architecture.
min Σ (M; Cdistill,
where D; is the ith cluster, 0; is the parameter used for training Di, and AMAR(D;; reconstruction loss for masked image modeling. Og consists of several experts in different layers, as explained in the following.
(;; 0z) is the
Gate Design. As in the TokenMoE, we replace several MLP layers in the ViT with layers equipped with MoCE gates. TokenMoE, routings of the tokens experts are considered separately. In MoCE, we route tokens from images of the same cluster to the same expert. The MoCE gate output can thus be written as
G(x) TopK C2) +c)),
where W, is the gate parameter, and Cix) is the embedding of the cluster that x belongs to. Em- pirically, we find that the confidence of G'(x) (the max entry) is low and consequently, the mapping between clusters and experts varies a lot during pre-training. Inspired by the importance and load losses (Riquelme et al. 202T), we add the following loss Limbalance to enhance the confidence of the gates. Since makes G(xx shaper, call it imbalance loss.
std(G(x):: 2
Limbalance -
(5)
mean(G(xx).
For practical implementation, the loss is calculated over the samples in batch. The imbalance loss penalizes on the negative variance of the gate confidence.
Deployment. On deployment, customized experts are selected from MoCE, and fine-tuned for each downstream task. As shown in Section [.I] we prefer to use the experts that is pre-trained from data whose semantics is closest to that of the downstream task. This can be obtained by re- using the data clustering module. Specifically, we feed images for the downstream task through the pre-trained MAE model and collect all the image features as Frask- The assignment of downstream images to the clusters is then computed as Ataek = F C. We select the largest cluster with assigned downstream images, and use the corresponding experts (a sub-model of the whole MoCE model) for deployment. In the case when only one expert is activated at each MoCE layer (K = 1), regular ViT model is needed for downstream fine-tuning, which is much more efficient than MoF.
5

---

Published as a conference paper at ICLR 2023
Table 1: Transfer accuracy (%) of self-supervised learning models on 11 downstream tasks.
Aireraft Caltech
Cars
C10 C100
DTD
Flowers
Food
Pets
SUN
VOC
Avg.
ResNer-50
BYOL
82.39
90.12
87.33
96.28
82.15
74.57
95.96
82.13
88.52
64.41
83.97
84.35
DeepCluster-v2
78.75
90.51
86,33
96.48
82.28
75.43
96.16
83.68
90.33
66.68
81.37
84.36
Vision Transformer
Supervised
76.55
89.98
86.
96.79
83.96
75.09
93.94
85.17
92.54
64.54
87.22
84.72
DINO
66.50
91.65
76.37
86.69
75.73
96.40
93.77
93.97
59.33
86.62
84.10
MoCo v3
76.29
91.64
85.18
97.99
86.98
72.64
95.33
83,94
92.35
65.54
84.21
84.74
BEiT
53.16
79.02
68.11
94.34
73.54
68.04
91.33
79.59
84.02
56.13
65.65
73.90
MAE
72.38
90.47
83.51
95.69
68.40
75.48
96.10
79.98
92.35
62.43
84.79
81.96
MAE
72.71
91.24
84.47
77.33
75.05
96.25
80.49
92.78
62.46
85.02
83.09
MoCE (Ours)
78.73
90.61
88.56
97.79
84.68
74.04
96.94
86.24
93.07
65.05
85.26
85.54+2-45
Table 2: Transfer accuracy (%) on detection and segmentation.
ADE20K
COCO
Method
mInU
Ap
AP
AP
Ap
AP
AP
Supervised
46.9
48.8
68.7
52.7
42.5
65.9
45.5
DINO
46.9
49.5
69.1
53.6
42.9
66.0
46.3
MoCo v3
46.8
47.2
66.9
50.8
41.1
63.6
44.1
BRIT
45.6
40.8
59.4
44.1
36.0
56.8
38.2
MAE
48.1
50.6
69.4
55.0
43.8
66.6
47.5
MoCE
48.3
51.1
69.8
55.4
44.2
67.0
48.1
EXPERIMENTS
In this section, we first introduce the setup of pre-training and fine-tuning stage of MoCE in Sec. Then we demonstrate the effectiveness of MoCE by evaluating pre-trained models on collection of 11 downstream tasks with detailed analysis of our MoCE superior to vanilla MAE and TokenMoE Finally we take ablation studies on the key components of MoCE in Sec. 4.3
in Sec. 4.2
4.1
SETUP
For all experiments, replace MLP layers with MoCE layers in the original VTT-B (Doso- vitskiy et al. (202I). Following Wu et al. (2022). layers with the greatest gradient magnitude are selected (which are the last two MLP layers in our experiments). Unless otherwise specified, the number of experts is 8 and the number of clusters is 256. Our model utilizes the officially released 1600-epoch pre-trained MAE mode?" and continues train for an extra 200 epochs. Each expert is initialized by the corresponding dense model with small weight perturbation. The training pro- cedure mainly follows that of MAE, except that we multiply the base learning rate by 0.1. All regularization loss weight is set to 0.01 by default.
To ensure fair comparison with the vision transformer on downstream classification tasks, we mainly follow the hyper-parameter settings in (Dosovitskiy al.) 2021: Riquelme al. 2021) and the benchmark settings (Ericsson et al. 2021). The proposed model is compared with var- ious self-supervised models, including DINO (Caron al. 2021). MoCo v3 (Chen al. 202 ResNet-50 models of BYOL (Grill et al. 2020)
BEiT (Bao et al.) (2022), and the highly-performe and DeepCluster-v2 (Caron al. 2018). We also compare with the supervised pre-trained model DeiT' (Touvron et al. 2021). To make a fair comparison of training time, we continue to train a 1600-epoch pre-trained MAE for 200 epochs with total ImageNet
baseline, and is denoted as
MAE* in Table [ For detection and segmentation tasks, following Bao et al. (2022), we perform
experiments ADE20K (Zhou et al., 2019) and COCO (Lin et al. 2014). We utilize the officially released checkpoints for all baselines. Details are in Appendix A.5

---

Published as a conference paper at ICLR 2023
Table 3: Transfer accuracy of MAE, TokenMoF, SDR and MoCE. SDR(VIT) is our implementation of SDR under ViT. We observe that TokenMoE cannot outperform vanilla MAE, while SDR(ViT) achieves better performance, which is further outperformed by MoCE.
Aireraft
Caltech
Cars
C10
C100
DTD
Flowers Food
Pets
SUN
VOC
Avg.
MAE
72.71
91.24
84.47
96.15
77.33
75.05
96.25
80.49
92.78
62.46
85.02
83.09 81.04
TokenMoE
70.51
89.70
81.40
95.18
76.44
73.67
95.09 77.45
90.71
61.12 80.15
SDR
75.77
89.73
86.65
95.31
83.60
73.62
95.53
84.77
91.25
64.64
83.51
84.03
SDR(ViT)
76.57
90.04
86.95
96.92
81.42
73.09
96.14
82,90
92.65
64.40
85.37
84.22
MoCE
78.73
90.61
88.56
97.79
84.68
74.04
96.94
86.24
93.07
65.05
85.26
85.54
4.2
RESULTS
Transfer Results, The classification transfer performance of various self-supervised models are shown in Table [T] As can be seen, MoCE achieves a 2.45% improvement over MAE" and reaches the state-of-the-art averaged accuracy, demonstrating the effectiveness of the task-customized pre- training paradigm. On fine-grained datasets such Aircraft, Cars and Food, MoCE outperforms the baseline model by large margin. This is because these fine-grained tasks are similar to only a subset of the pre-training dataset. Hence, McCE can alleviate negative transfer by using the model that is trained by the cluster most similar to the particular downstream task. On the other hand, MoCE shows only limited improvement on tasks such as Caltech, Cifar- 10 and VOC. These tasks are more general and contain images covering the various semantics in the pre-training dataset, and thus negative transfer does not exist.
Table 2 shows the transfer performance on the detection and segmentation tasks. As can be seen, MoCE outperforms MAE and the other baselines (including the supervised one), and achieves state- of of-the-art art results.
Comparison between MoCF, TokenMoF, MAE and SDR. In this experiment, we compare MoCE with the following models: (i) MAE, (ii) TokenMoE, (iii) SDR (Liu et al. 2022). a task- customized model that aims at alleviating negative transfer, and (iv) SDR(ViT), which re-implements SDR with the VIT architecture. T the transfer accuracy on 11 downstream tasks. As can be seen, TokenMoE performs is not desirable. Both MoCE and SDR(ViT) outperform MAE, demonstrating the effectiveness of task-customized methods for alleviating negative transfer. MoCE further outperforms SDR(ViT), indicating the importance of self-adaptive routing.
worse than MAE, suggesting that naively adopting MoE to MAE
Figure B(d) shows the peak signal-to-noise ratio (PSNR) (Sara al. 2019), which reflects the gener- ation quality of these autoencoder models. MoCE exhibits improvement over TokenMoE and MAE on most datasets. We also provide the comparisons in the case of a fair parameter count, large architectures, and training from scratch in the Appendix A.2
A.3 and A.4, respectively.
Analysis on experts. Figure 3(a) and Figure ] show the routing heatmaps for TokenMoE and MoCE, respectively. As can be seen, routing of the TokenMoE experts has little correlation with semantics. On the other hand, each MoCE expert is trained by several classes, showing a more balanced assignment of images more effective learning of the experts. Moreover, notice that the importance loss and load balance loss (Riquelme et al. 202T) are applied and indeed work as "expected" because they are only applied with respect to patch tokens instead of semantic classes. On the other hand, McCE can balance the experts both at the token level and semantic level.
experts. This verifies that the improvement of MoCE is due to
Figure 3(b) shows example pre-training samples for 3 random MoCE experts. Note that expert is mostly trained by images containing clothes, experts is pre-trained mostly by bird images, while expert 3 is pre-trained mostly by dog images.
Next, we show that each expert is trained by samples with similar semantics. Following MMikolov et al. 2013), select the label set used by each expert, and then compute the & distances between CLIP embeddings (Radford et al. 2021a) of labels used by the same expert and by different experts. TTe average distance between labels used by the same expert is 0.84, while that between labels
7

---

Published as a conference paper at ICLR 2023
200
400
800
1
2
5
(a) Routing heatmap for TokenMoR: experts.
(h) Examples pre-training samples for expert I (top), ex-
pert 2 (middle), and expert (bottom).
MAE CA
200
McCt
400
Relative PSWR
-1
800
1000
5
caars
cont
Plowert
food
SUM2097
(c) Routing heatmap for MoCE experts.
(d) Relative PSNR improvement over MAE.
Figure 3: (c), Routing heatmaps for experts in TokenMoE and MoCE. The x-axis is the expert ID, and the y-axis is the ImageNet semantic label ID. Darker green means a higher proportion of tokens belonging to the corresponding class are allocated to the expert. The label is sorted di differently in each figure to make readable. (b): Example samples from the pre-training dataset of 3 MoCE experts. (d): Relative PSNR improvement of TokenMoE and MoCE over MAE for each downstream task.
used by different experts is 0.92, indicating that the MoCE gate automatically aggregates labels with similar semantics to each expert, thus benefiting downstream transfer.
Training and testing efficiency. Table 4 compares the efficiencies of MAE, TokenMoE and MoCE during training and testing. As can be seen, all of them have similar FLOPs. However, TokenMoE needs to use the whole model during both training and testing, while McCE only needs to use a single expert and thus halves the required number of parameters when testing. In addition, the training and testing speeds are improved by respectively 18% and 37%, which is attributed to the reduction token shuffle operations as tokens in the same image do not need to be split and are dispatched to the same expert, significantly reducing the communication overhead.
4.3 ABLATION
Search method. When a downstream task arrives, it is expensive to fine-tune all experts to choose the best one. To find the task-customized expert (K 1), we compare the method proposed in Section 3.3 with (i) early stop, (ii) KNN (Liu et al. 2022). (iii) LogME (You et al. 202T). The experiment is performed on the task with the most samples (Food), the task with the least samples (Flowers), and the one with a medium number of samples (Aireraft). For comparison, we addition- ally show the performance of the best and worst experts based on an exhaustive search. As can be seen in Table 5, McCE performs stably among different sizes of the dataset, and the search cost is negligible as we only need to infer the downstream task once and feed it to the clustering module. This illustrates another advantage of combining clustering and pre-training in single paradigm.
MoCE Architectures. In this experiment, we study the different architecture preer-parameters of MoCE in three aspects. First, vary the number of experts in each MoCE layer. As be seen from Table using more experts leads to consistent improvement on the accuracy.
Next, vary the location of the MoCE layers. As mentioned in Section we select the MoCE layers based on the gradient magnitudes. In the experiments, MoCE selects the 11th and 12th MLP layers. On the other hand, TokenMoE chooses the last 2 even-numbered (i.e., 10th and 1 12th) ML.
8

---

Published as a conference paper at ICLR 2023
Table 4:
Efficiency during training (top)
Table 5: The search cost (in GPU hours) for different expert search algorithms.
and testing (bottom).
MAE
TokenMoE
MoCE
Aircraft Flowers Food GPU hours
Params (M)
111.91
178.03
178.03
Best Worst
79.92
97.96
86.24
288
FLOPs (G)
9.80
9.81
9.81
69.84
94.97
81.51
288
Speedt
1,41x
1x
1.18x
Early stop
77.00
96.83
85,33
144
# Params (M)
85.88
152.00
85,88
KNN
71.40
95.10
83.32
5
FLOPs (G)
16.88
16.88
16.88
LogME
73.84
96.54
85.11
5
Speedt 1.37x
Ix
1.37x
MoCE
78.73
96.94
86.24
1
Table 6: Accuracies with dif- ferent numbers of experts in MoCE layer. (default setting used is in bold).
Accuracies with dif- ferent numbers and locations of the MoCE layers (default setting used is in bold).
Table 7:
Table 8: Accuracies with
different numbers of clusters (default setting used is bold).
experts
Acc (%)
# MoCE layers
Acc (%)
# clusters
Acc (%)
83.09
83.09
16
82.00
2
83.01
(10th & 12th)
83.08
84.02
4
84.22
2 (11th & 12th) 4
85.54
256
85.54
8
85.54
85.59
512
85.33
layers. Furthermore, exhibit the performance with only 1 and 4 MoCE layers, which are also selected based on the gradient magnitudes. As shown in Table 7 we notice that essential to choose the right layer to be MoCE layer. Adding more MoCE layers shows little improvement.
We also train MoCE with different numbers clusters. As shown in Table the accuracy increases up to 256 clusters, and then begins to drop. We hypothesize that with moderate number of clusters, MoCE can produce a variety cask-custoiized models. With even more clusters, the number of experts become the bottleneck and performance starts saturate.
CONCLUSION
In this work, first show that the negative transfer phenomenon exists in the prevailing self- supervised learning method MAE through extensive experiments. It will impede the scalability of MAE as more pre-training data may instead degenerate the downstream performance. In order to lackle the problem, we introduce Mixture of Expert to MAE as the multi-experts design can equip MAE with different ability that aids transfer. However, different from supervised pre-training, To- kenMoE suffers from the fact that the gate shows no correlation to the semantics and the transfer ability is not improved. Based on this, we propose MoCE to explicitly train each expert with different clusters through the MoCE gate design and several losses to stabilize the training process. A search algorithm for selecting the best model for transfer is also proposed based on the clustering priors. Extensive experiments show that MoCE trains each expert with meaningful semantics and achieves state-of-the- thettrrrt transfer performance on a collection 11 downstream tasks and both detection and segmentation tasks. It is the first work that successfully trains self-supervised learning MoE model on ImageNet only. We hope such a design will motivate more research on the self-supervised MoE models.
ACKNOWLEDGMENTS
We gratefully acknowledge the support of MindSpore, CANN (Compute Architecture for Neural Networks) and Ascend Al Processor used for this research.

---

Published as a conference paper at ICLR 2023
REFERENCES
Yuki Markus Asano, Christian Rupprecht, and Andrea Vedaldi. Self-labelling via simultancous clustering and representation learning. Preprint 1rXiv:1911.05371, 2019.
Yutong Bai, Zeyu Wang, Junfei Xiao, Chen Wei, Huiyu Wang, Alan Yuille. Yuyin Zhou, and Cihang Xie. Masked autocncoders enable efficient knowledge distillers. Preprint uuxiv2200.1125555 2022.
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. Beit: Bert pre-training of image transformers. In ICLR, 2022.
Zhaowei Cai and Nuno Vasconcelos. Cascade R-CNN: delving into high quality object detection. In CVPR, 2018.
Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze. Deep clustering for unsu- pervised learning of visual features. In ECCV, pp.
2018.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. NeurIPS, 33:9912- 9924, 2020.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In ICCV, pp. 9650- 9660, 2021.
Kai Chen, Lanqing Hong, Hang Xu, Zhenguo Li, and Din Yeung. Multisiam: Self-supervised multi-instance siamese representation learning for autonomous driving. In ICCV, pp. 7546-7554, 2021a.
Kai Chen, Zhili Liu, Lanqing Hong, Hang Xu, Zhenguo Li, and Dit-Yan Yeung. Mixed autoencoder for self-supervised visual representation learning. In CVPR, 2023.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. ICML, 2020.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training suff-supereieed vision transformers. In ICCV, pp. 9640-9649, 2021b.
Kevin Clark, Minh-Thang Luong, Quoe V Le, and Christopher D Manning. Blectra: Pre-training text encoders as discriminators rather than generators. In ICLR, 2020.
Ekin Dogus Cubuk, Barret Zoph, Jon Shlens, and Quoe Le. Randaugment: Practical automated data augmentation with a reduced search space. NeurlPS, volume 33, pp. 18613-18624. Curran Associates, Inc., 2020.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. NeurIPS, 26, 2013.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009.
Jacob Devlin, Ming-Wes Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep 2018
bidirectional transformers for language understanding. Preprint
Carl Doersch, Abhinav Gupta, and Alexei Efros. Unsupervised visual representation lcarning by context prediction. In ICCV, pp. 1422-1430, 2015.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x 16 words: Transformers for image recognition at scale. In ICLR, 2021.
Linus Ericsson, Henry Gouk, and Timothy M Hospedales. How well do self-supervised models transfer? In CVPR, pp. 5414-5423, 2021.
10

---

Published as a conference paper at ICLR 2023
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.
Dehong Gao, Wenjing Yang, Huiling Zhou, Yi Wei, Yi Hu, and Hao Wang. Network clustering for multi-task learning. Preprint arXiv:2101. 01009018, 2021.
Spyros Gidaris, Praveer Singh, and Nikos Komodakis. Unsupervised representation learning by predicting image rotations. In ICLR, 2018.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, al. Bootstrap your own latent: A new approach to self-supervised learning. In NeurIPS, 2020.
Jianhua Han, Xiwen Liang, Hang Xu, Kai Chen, Lanqing Hong, Jiageng Mao, Chaoqiang Ye, Wei Zhang. Zhenguo Li, Xiaodan Liang, and Chunjing Xu. SodalOm: A large-scale 2d self/semi- supervised object detection dataset for autonomous driving. Preprint arXiv:2106. 1118, 2021.
Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In ICCV, 2017.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In CVPR, pp. 9729-9738, 2020.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In CVPR, pp. 16000-16009, 2022.
Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, and Daniel Soudry. Augment your batch: Improving generalization through instance repetition. CVPR, June 2020.
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang. Vijay Pande, and Jure
Leskovec. Strategies for pre-training graph neural networks. Preprint 1tXiv:1905.12265 2019.
Minyoung Huh, Pulkit Agrawal, and Alexei Efros. What makes imagenet good for transfer learn- ing? Preprint wrXiv:1608.08614 2016.
Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Michael I Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2):181-214, 1994.
Sneha Kudugunta, Yanping Huang, Ankur Bapna, Maxim Krikun, Dmitry Lepikhin, ThangThang Luong, and Orhan Firat. Beyond distillation: Task-level mixture-of-experts for efficient inference. Preprint arXiv:21 10.03742, 2021.
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. Preprint arXiv:2006. 2020.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In ICML, pp. 6265-6274. PMLR, 2021.
Kaican Li, Kai Chen, Haoyu Wang, Lanqing Hong, Chaoqiang Ye, Jianhua Han, Yukuai Chen, Wei Zhang, Chunjing Xu, Dit- Yan Yeung, al. Coda: A real-world road corner case dataset for object detection in autonomous driving. Preprint arXiv:2203. 007724, 2022a.
Xiaotong Li, Yixiao Ge, Kun Yi, Zixuan Hu, Ying Shan, and Ling- Yu Duan. me-beit: Multi-choice discretization for image bert pre-training. Preprint AXXi:2203.15377 2022b.
Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. ECCV, 2014.
Xiaodong Liu, Pengcheng He, Weizhu Chen, and Jianfeng Gao. Multi-task deep neural networks for natural language understanding. Preprint arXiv:1901.1155 2019.

---

Published as a conference paper at ICLR 2023
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. ICCV, pp. 10012- 10022, 2021.
Zhili Liu, Jianhua Han, Kai Chen, Lanqing Hong, Hang Xu, Chunjing Xu, and Zhenguo Li. Task- customized self-supervised pre-training with scalable dynamic routing. AAAI, 55:65, 2022.
lya Loshchilov and Frank Hutter.
Decoupled
weight decay regularization.
Preprint
arXiv:1711.001 2017.
Jiaqi Ma, Zhe Zhao, Xinyang Yi, Jilin Chen, Lichan Hong, and Ed H Chi. Modeling task relation- ships in multi-task learning with multi-gate mixture-of -exture-ff-experts... In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 1930-1939, 2018.
Matthew McDermott, Bret Nestor, Evan Kim, Wancong Zhang, Anna Goldenberg, Peter Szolovits, and Marzyeh Ghassemi. A comprehensive ehr timeseries pre-training benchmark. In Proceedings of the Conference on Health, Inference, and Learning. pp. 257-278, 2021.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word represen- tations in vector space. Preprint urXiv:1301.311 2013.
George A Miller. WordNet: An electronic lexical database. MIT press, 1998.
Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A Efros. Context encoders: Feature learning by inpainting. In CVPR, 2016.
Joan Puigcerver, Carlos Riquelme Ruiz, Basil Mustafa, Cedric Renggli, André Susano Pinto, Syl- vain Gelly, Daniel Keysers, and Neil Houlsby. Scalable transfer learning with expert models. In ICLR, 2020.
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, and J. Clark. Learning transferable visual models from natural language supervision. In ICMI,, 2021a.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In ICML, pp. 9748-8763 PMLR, 2021b.
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. NeurlPS, 34:8583-8595, 2021.
Umme Sara, Morium Akter, and Mohammad Shorif Uddin. Image quality assessment through fsim, ssim, mse and psnr-4 a comparative study. Journal of Computer and Communications, 7(3):8-18, 2019.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-expprrs layer. In ICLR, 2017.
Hugo Touvron, Matthicu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-cfficient image transformers & distillation through attention. In ICML, volume 139, pp. 10347-10357, July 2021.
Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. In ICCV, pp. 568-578, 2021.
Xinlong Wang, Rufeng Zhang, Chunhua Shen, Tao Kong, and Lei Li. Dense contrastive learning for self-supervised visual pre-training. Preprint arXiv:2011.09 1.09157, 2020.
Lemeng Wu, Mengchen Liu, Yinpeng Chen, Dongdong Chen, Xiyang Dai, and Lu Yuan. Residual mixture of experts. Preprint SSYv:2204.096366 2022.
12

---

Published as a conference paper at ICLR 2023
Enze Xie, Jian Ding, Wenhai Wang, Xiaohang Zhan, Hang Xu, Zhenguo Li, and Ping Luo. Detco: Unsupervised contrastive learning for object detection. Preprint Wtt.:::212.0488000 2021a.
Enze Xie, Wenhai Wang, Zhiding Yu, Anima Anandkumar, Jose M Alvarez, and Ping Luo. Seg- former: Simple and efficient design for semantic segmentation with transformers. In NeurIPS, 2021b.
Zhenda Xie, Yutong Lin, Zheng Zhang, Yue Cao, Stephen Lin, and Han Hu. Propagate your- self: Exploring pixel-level consistency for unsupervised visual representation learning. Preprint DrXiv:2011.100 100443, 2020.
Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: A simple framework for masked image modeling. Preprint arXiv:2111.09886 2021c
An Yang, Junyang Lin, Rui Men, Chang Zhou, Le Jiang, Xianyan Jia, Ang Wang, Jie Zhang, Jiamang Wang. Yong Li, al. M6-t: Exploring sparse expert models and beyond. Preprint 1.iv:2105.15082, 2021.
Lewei Yao, Jianhua Han, Youpeng Wen, Xiaodan Liang, Dan Xu, Wei Zhang. Zhenguo Li, Chunjing Xu, and Hang Xu. Detclip: Dictionary-enriched visual-concept paralleled pre-training for open- world detection. In NeurIPS, 2022a.
Lewei Yao, Runhui Huang, Lu Hou, Guansong Lu, Minzhe Niu, Hang Xu, Xiaodan Liang, Zhenguo Li, Xin Jiang, and Chunjing Xu. Filip: Fine-grained interactive language-image pre-training. In ICLR, 2022b.
Kaichao You, Yong Liu, Jianmin Wang, and Mingsheng Long. Logme:
Practical assessment of pre-trained models for transfer learning. In ICML, pp. 22133-12144 PMLR, 2021.
Richard Zhang, Phillip Isola, and Alexei A Efros. Colorful image colorization. In ECCV, 2016.
Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ade20k dataset. In IJCV, 2019.
Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. ibot: Image bert pre-training with online tokenizer. Preprint arXiv:2111.07832, 2021.
13

---

Published as a conference paper at ICLR 2023
A APPENDIX
A.1 DETAILS FOR CLUSTERING,
For data clustering, the features are computed by inferring the pre-train MAE, and the matrix Q and C are solved by the Sinkhorn- Knopp algorithm, we set the iteration number as 3. The learning rate SGD is set to 0.1, the momentum is 0.9 and weight decay is set to 0.9 for the sparse assignment of cluster results. We train 10 epochs in total and costs 3 minutes and 20 seconds on average for a single GPU.
algorithm and SGD optimizer iteratively. For the Sinkhorn-
A.2 COMPARISON UNDER FAIR PARAMETER COUNTS.
The setting used in our work focuses on a fair comparison of FLOPs, referring to Table 4 in the main paper. Since TokenMoE and MoCE always activate only one expert throughout the whole pre-training and fine-tuning procedure, the FL.OPs value is maintained close to MAE. Apart from this criterion, further provide the comparison equal parameter counts. As shown in Table 9 we train MAE under the same parameter count as the whole model of MoCE, and MoCE still outperforms MAE consistently.
Table 9: Comparison of MAE and MoCE under equal parameter counts. We train MAE with larger model that shares the same parameter count as the whole model of MoCE.
# Params Aircraft Caltech
Cars
C10
C100
DTD Flowers Food
Pets
SUN
VOC
Avg.
MAE
178.03
74.43
90.30
85.50 96.90 83.80
74.84
81.86 92.97 62.98 85.51 84.13 96.94
96.30
MoCE
178.03
78.73
90.61
88.56 97.79 84.68
74.04
86.24
93.07 65.05
85.26 85,54
A.3 MOCE FOR LARGER ARCHITECTURE.
Here we provide the analysis on the larger archttettureVViT-L. 2.57x larger than ViT-B) to explore the scalability of MoCE. We first demonstrate that negative transfer still exists for larger architecture by training MAE with total ImageNet, Split-A A and Split-B following the same setting in Sec. 3.1 As shown in the first three rows of Table. 10 similar phenomenon is observed that MAE-L traied by Split-A performs better in Aireraft, Cars, DTD and SUN while Split-B in Flowers, Food, and Pets. On the other hand, MoCE-L can still alleviate the problem and therefore transfers better. We believe that the negative transfer phenomenon mainly exists when common pre-trained model is used for various downstream tasks, due to the inevitable semantic gaps between the pre-trate pretrranning and downstream datasets, rather than the architecture.
Table 10: Comparison of MAE and MoCE on ViT-L. We also train MAE with 2 subsets of ImageNet, namely Split-A and Split-B, following the same setting mentioned Sec. This table shows that negative transfer still exists on larger architectures, while MoCE can alleviate this problem and
achieve better transfer results.
Aireraft Caltech Cars C10
C100 DTD Flowers
Food
Pets
SUN
VOC
Avg.
MAE-L (full set)
74.30
93.97
88.60
97.85 82.47 77.61
96.67
81.22 93.97 67.99 88.30
85.72
MAE-L (Split-A)
79.70
91.59
89.33
96.97
80.38 78.67
95.44
82.97
92.49 68.73 82.41
85,33
MAE-L (Split-B)
73.42
90.80
86.00
96.18
78.73 77.34
96.75
83.63
94.92 66.06 85.85
84.52
MoCE-L
87.04
94.86
90,72
98.29 87.49
76.65
97.38
88.21
95.89
69.49
89.13
88.65+2.0
A.4 PERFORMANCE OF MOCE WITHOUT PRE-TRAINING.
We provide results of MoCE trained from scratch for 200 epochs and 1600 epochs Table II In this experiment, for clustering, we first pre-train MAE for 50 epochs and perform clustering. We then train MoCE from scratch for 200 epochs and 1600 epochs based on the clustering results. Although it is a common practice to utilize pre-trained dense models initialization accelerate pre- pre-trrining (Wu et al., 2022) Bai et al.) 2022), MoCE still outperforms MAE consistently in various downstream tasks when trained from scratch.
14

---

Published as a conference paper at ICLR 2023
Table 11: Comparison of MAE and MoCE both training from seratch for 200 epochs (first two rows) and 1600 epochs (last two rows).
Aircraft Caltech
Cars
C10
C100
DTD
Flowers
Food
Pets
SUN
VOC
Avg.
MAE
64.73
85,91
77.10
92.92
72.50
73.30
93.11
73.14
88.70 57.84 73.27
77.50 84.25
MoCE
71.16
90.55
82.46
96.06
76.56
74.57
95.70
79.67
92.58 62.20
82.34
MAE MoCE
72.38
90.47
83.51
95.69
68.40
75.48
96.10
79.98
92.35
62.43
84.79
81.96
78.75
91.64
87.04
97.15 83.12 73.62
96.08
83.84
93.06
65.49
85.81
85.05
A.5 EVALUATION DETAILS FOR DOWNSTREAM TASKS.
Classification. We mainly follow the settings Ericsson et al. 2021 to make a fair comparison. Specifically, all models are trained by SGD with a momentum of 0.9. Weight decay is set to be and the learning rate is searched among [le-4, 3e-4, le-3, 3e-3, le-2, 3e-2, le-1, 3e-1]. Each model is fine-tuned for 2500 steps with cosine learning rate decay, batch size of 64, and x24x224 resolution. We fine-tune each model 3 times and report the average performance. We find such a setting generates stable result.
Semantic segmentation. We evaluate MoCE on the semantic segmentation task that aims to pre- dict the class for each pixel in the image. We report the metric of mean Intersection of Union (mloU) averaged over all semantic categories in ADE20K (Zhou et al. 2019). We choose the best expert by applying ADE20K images to our clustering module and selecting the cluster that contains the most images. We use Adam (Loshchilov & Hutter (017)) as the optimizer. The learning rate is set to le-3 with layer-wise learning rate decay (Clark al. 2020) be 0.65. We conduct fine-tuning for 160K steps. The batch size is 16. The detailed hyper- parameess can refer to Bao et al. (202)).
Detection and Instance segmentation are also evaluated on COCO (Lin et al. 2014). We follow the same deployment method as the one used in the semantic segmentation task choose the best expert. Following iBOT (Zhou et al. (202)) we adopt the Cascade Mask R- R-CNN (Cai & Vasconcelos/ 2018 He et al. 2017) and the multi-scale training. The shorter side is randomly resized between 480 and 800 while the longer one is no longer than 1333. The batch size is 16, and the initial learning rate is le-4. The layer-wise learning rate decay ratio (Clark al. 2020) is set to 0.75, We train the model for 12 epochs and decrease the learning rate by 10x at epoch 9 and 11.
15


Reviews:
[{"review_id": "hf1llfdmaK", "reviewer": "Reviewer_ycEm", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "", "summary_of_the_paper": "The paper studies the negative transfer phenomenon of MAE pre-trained models. Authors first experimentally show the existence of negative transfer and then show that natively using MoE can not solve the problem. Then, a clustering-based method is proposed to solve the problem.", "summary_of_the_review": "I like the idea but believe the current experiments are not sufficient. If my concerns are solved, I will be willing to raise the score.\n\n\n------ After rebuttal --------\nIn general, some of my concerns have been solved, e.g., training from scratch, and transferring to det and seg tasks. Thus, I would like to raise my score to borderline accept considering the interesting design of MOE.\n\nHowever, I still have some concerns and hope the authors can address them in the final version. First, the added supervised pertaining results remind me that comparing MoCE with [1] would be necessary, which shows that simple MLPs can significantly improve the transferring performance of supervised pretraining.\nSecond, as you claim/argue the method is pretraining not transferring, it would be better to update all the results of MoCE with training from scratch instead of based on MAE to avoid confusion. This issue has also been pointed out by other reviewers (I believe it's feasible for the authors to train MoCE 1600 epochs from scratch since they already show some results about 200 epochs).\n\n[1] Revisiting the Transferability of Supervised Pretraining: an MLP Perspective, CVPR 2022", "strength_and_weaknesses": "Generally, the paper is well-organized and easy to follow. The research topic is relatively new and interesting. The proposed MoCE seems correct and is well supported by the experimental results. \nHowever, I also have the following concerns:\n1. In the abstract and introduction, the authors position this work as a pretraining method. However, MoCE is based on a pre-trained MAE. Considering the extremely heavy cost of training foundation models, it would not be feasible. \n2. Based on 1, the cost of training MoCE would be twice of MAE. To make a fair comparison, what if we train MAE with a dataset of twice scale or train MAE with more epochs?\n3. If we focus on the negative transfer phenomenon, not combining MoE, there would be other transfer learning and model adaptation methods (e.g., prompt tuning) that could ease the negative transfer phenomenon. Could you compare it with some of them?\n4. Since the input of Eq.5 is cluster centroid and there are only m clusters, is it enough to learn W_g well?\n5. For Eq.6, is the loss implemented over all samples in the dataset or a batch? \nIn the experiment, only 11 small-scale classification tasks are evaluated, how about other more important tasks, e.g., object detection, semantic and instance segmentation?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Tz3lJ2LOTeC", "reviewer": "Reviewer_uKwM", "summary": "", "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": "", "summary_of_the_paper": "The paper initiates an interesting exploration of MAE with mixture of experts (MoE). The method is quite well-motivated, with interesting and fairness-in-mind designs, and a good amount of experiments devoted to it. The final results are reported using a suite of 11 downstream classification tasks, typically used to evaluate self-supervised learning methods. The proposed method is shown to significantly improve the overall accuracy, while maintaining efficiency.", "summary_of_the_review": "I would be on the acceptance side for the paper given the pros and cons listed above. I hope the authors can address the clarification questions I have, and try to tackle the empirical study (e.g. with larger models) as best as they can.", "strength_and_weaknesses": "Strengths:\n- This is one of the first works that I am aware to study MAE and MoE jointly. The exploration on this direction is interesting and of significance.\n- The paper motivates the approach by pointing to an empirical result of MAE suffering from the negative transfer phenomenon, and the final approach (guided by the motivation) is able to significantly outperform the MAE baseline. Overall it looks like a healthy research cycle to me.\n\nWeaknesses:\n- One major concern I am having with MoE-kind approach is about larger models. While used in a sparse way, MoE is still having a lot of parameters, and if directly comparing against MAE of the same backbone, it does not sound too fair to me. So I would like to see two things: 1) with a larger MAE model that roughly has the same number of parameters as a model with MoE, what's the comparison? and 2) Whether the improvements of the current model can transfer to larger models. Related: in table 3, I am not able to find the column for MAE. All I find are columns for TokenMoE and MoCE, while the caption says MAE is under comparison.\n- Related to the above, I am not sure the \"negative transfer\" phenomenon still exists with larger and larger MAE models? The hypothesis here is that maybe with larger models, it can capture both the man-made objects and the natural ones from ImageNet?\n- Maybe it is demanding too much computation overhead, but I noticed that the method in the paper is built from pre-trained MAE, and not training from scratch.  So I am wondering what would MoCE be like when pre-trained from scratch?\n- For the distillation loss, I am not sure why it still maintains an efficiency advantage over TokenMoE because the \"whole\" network is used as a \"teacher\", and presumably it can be costly to forward through all the parameters.\n- (minor as I don't know where to put it in the review) Why does MoCo v3 achieve so good results on Table 1? Given that MoCo v3 is better than MAE, isn't it more reasonable to start MoE explorations from MoCo v3 given the results listed in the table? In the same table, I am also curious to see what the supervised ViT will achieve on these downstream datasets. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "A_3vlRIYP9L", "reviewer": "Reviewer_Emir", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "", "summary_of_the_paper": "\nThe paper has a good motivation of transferring a pretrained model into diverse tasks by tackling the problem of negative transfer. The solution is having an abstraction layer of semantic, where input is grouped into different clusters based on their semantics before routing to experts. Experiments show a very good improvement of 2.54% on ImageNet. \n", "summary_of_the_review": "\nThe paper motivates a good problem in transferring a vision task and proposes a good solution in which semantics is used to train the mixture-of-expert model, and perform well with diverse tasks. Although there are some parts that need more clarification, generally it’s a good and well written paper. \n", "strength_and_weaknesses": "Strength: 1/ Having a good motivation on an important application, where we should take advantage of pretrained models to save cost.  2/ The testing efficiency is really good. \n\nSome Weaknesses\n\n1/ About cost saving: the proposed solution has a non-negligible overhead of data clustering, which is itself an optimization problem. My understanding is this overhead is not included in the training analysis (please correct if I am wrong). \n\n2/ It’s not entirely clear about the location and number of MoCE layers chosen even with Fig 4b and Table 6,  e.g. the intuition as to why placing them on, say, 10th and 12th layers are good but not before that? This is also along with the claim of hierarchical learning by the model concerning Fig 4b. Why don’t we try a couple of other patterns such as interleaving odd and even layers, or early in the layers while the input is already arranged in clusters, which is somewhat more intuitive? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9nNtVFkmTYu", "reviewer": "Reviewer_Wv7b", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "", "summary_of_the_paper": "This paper investigates practical techniques to apply the Mixture of Experts (MoE) model to the MAE-based pre-training paradigm. \n\nTaking a supervised MoE model for ViT as a baseline, the authors introduced data clustering as a strong baseline for experts' partition and assignment. \n\nSeveral additional terms, such as imbalance regularization and distill loss, are also proposed in the self-supervised MoE design. \n\nOn 11 downstream tasks, the proposed method outperforms MAE, supervised TokenMoE and several recent self-supervised single model methods (non-MoE). \n", "summary_of_the_review": "Although this paper obtains improved results over MAE and other self-supervised methods, the technical contribution is kind of incremental w.r.t (Liu et al. 2022) and TokenMoE.   \nAnd there is a lack of discussion and comparison w.r.t (Liu et al. 2022)   \n\nOverall, this paper is below the acceptance threshold. \n\n------------------------------\nAfter Rebuttal:   \nThe authors addressed most of my concerns and made significant efforts. I would like to see this paper accepted and increased the overall score. \n", "strength_and_weaknesses": "# Strength\n- This paper proposed practical techniques to address the expert collapse issue in TokenMoE. \n- The proposed MoCE improves the baseline performance on 11 downstream tasks, showing the effectiveness of the proposed method. \n- Reasonable visualization of the expert routing and expert visualization is shown in the experiment. \n- Several ablation studies are conducted in Table 4-7. \n\n\n# Weaknesses\n- Although better performance is achieved on the benchmark dataset, the technical contribution is kind of incremental. The main algorithm re-uses multiple design components of (Liu et al., 2022) and applies them to a slightly different scenario, i.e. MAE self-supervised training with a Mixture of Experts. Compared with (Liu et al., 2022), the improvements are (1) exploring the new self-supervised paradigm MAE, (2) applying to the Transformer backbone, and (3) different routing mechanisms with MoE. Both (1) and (2) are straightforward extensions to new technologies. As to (3), it offers some insights and benefits on self-supervised pre-training. \n- As mentioned above, the proposed method heavily relies on design modules from (Liu et al., 2022), e.g. Data Clustering (Eq. (3)), Distillation term. However, most of the methodology and experimental comparison are conducted with respect to TokenMoE, a specific supervised MoE model. As also verified by the authors in Figure 3(a), TokenMoE is not specifically designed for the setting of self-supervised pre-training. Therefore, more comparisons and discussions should be made w.r.t. (Liu et al., 2022) rather than TokenMoE. \n- TokenMoE used a load loss to balance the expert assignment, so it looks a little strange in Figure 3(a) the expert collapsing. \n- In Table 3, the results for MAE are missing. \n- It seems the results of DeepClusterV2 and BYOL in Table 1 are different from the results from (Liu et al., 2022), both on Resnet 50. \n- Other minors:   \n> \"decades(Jacobs et al.,\" and \"processing(Lepikhin et al., 2020\", missing space before left parentheses.   \n> Equ. 5 is not so common, usually Eqn. 5 or Eq. 5.   \n> In Eq. (3), $Tr(Q^TA), Q\\in \\mathbb{R}^{n\\times m}, A\\in \\mathbb{R}^{n\\times m}$, please check if this is correct. A more common notation for Sinkhorn is using matrix product $<Q, A>$.   \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%