T ASK - CUSTOMIZED M ASKED A UTO E NCODER VIA M IXTURE OF C LUSTER - CONDITIONAL E XPERTS
Anonymous authors
Paper under double-blind review
A BSTRACT
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, named Mixture of Cluster-conditional Experts (MoCE), which can be trained once but provide 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 clusterconditional gates, which can automatically format the hierarchical cluster architecture of the pre-training data.
Thus, each downstream task can be allocated to its customized model pre-trained with data most similar with the downstream data.
Experimental results show that our MoCE achieves 2.45% average performance improvement compared to the vanilla MAE on a collection of 11 different downstream tasks.
I NTRODUCTION
Self-supervised learning (SSL) has become a prevailing model pre-training paradigm, which learns effective transferable representations without human annotations (He et al., 2020;
Chen et al., 2021a;
Bao et al., 2021a).
Among different SSL methods, Masked Autoencoder (MAE)
(He et al., 2022) might be the most prevalent one currently, which proposes to construct supervision signals with raw image data by masking random patches of the input image and reconstructing missing pixels.
This simple but effective self-supervisory strategy has been proved efficient to train large-scale models, such as ViT (Dosovitskiy et al., 2020) with impressive performance on widely adopted benchmarks, such as ImageNet-1K 1 (Deng et al., 2009).
However, can MAE, the so-called “scalable vision learner”, really be efficient for various downstream tasks?
Unfortunately, we observe that MAE still suffers from the negative transfer phenomenon (Liu et al., 2022) when transferring towards downstream tasks with distinct semantics.
Specifically, for a given downstream task, involving more data in pre-training might instead degrade transfer performance, due to the negative effects of semantic irrelevant pre-training data.
A preliminary is conducted in
Sec. 3.1, where MAE pre-trained on the full ImageNet is outperformed by MAE pre-trained on the semantically relevant subsets on 9 of 11 downstream tasks, as shown in Fig. 1(a).
To achieve task-customized pre-training, Mixture of Expert (MoE) (Riquelme et al., 2021) seems promising, which proposes a multi-expert architecture to provide customized models for different input tokens.
In the scenarios of self-supervised learning, however, it is difficult to train experts with diversity due to the missing semantic labels used in supervised pre-training, as demonstrated in Fig. 1(b).
Our experiments in Sec.
4.2 show that naively adopting MoE to MAE, referring as
TokenMoE, is not feasible as the experts tend to degenerate to one or a few experts without the data labels during pre-training.
In this paper, we propose the Mixture of Cluster-conditional Expert (MoCE), a novel paradigm to achieve task-customized self-supervised pre-training with data clustering and explicitly training each expert with images of similar semantics.
Specifically, our MoCE consists of three stages.
We first a i
[Figure]
c r a f t c a r s s u n 3 9 7 d t d f l o w e r s f o o d p e t s c
i f a r 1 0 c i f a r 1 0 0 c a l t e c h v o c
T r a n s f e r p e r f o r m a n c e ( a cc u r a c y , % )
[Figure]
cluster the pre-trained dataset based on a pre-trained dense MAE model.
Our MoCE is then constructed under the multi-expert structure, and we train each expert with different clusters through routing tokens in the image based on the cluster embedding instead of token embedding .
Additional regularization loss is proposed to enhance the confidence of gate results, stabilizing the training procedure.
Finally, when the downstream tasks come, we propose a searching method reusing the cluster module by feeding downstream images to it and the closet cluster is selected by the majority vote.
We show that our MoCE achieves remarkable strong performance on a collection of 11 downstream tasks over MAE and achieves state-of-the-art averaged accuracy.
Besides, the MoCE also enjoys the benefit of sub-model deployment that saves inference time and model capacity.
To summarize, our main contributions are:
We systematically analyze the negative transfer phenomenon with MAE, and show that naively adopting MoE to MAE cannot improve the transfer performance in downstream tasks.
We propose MoCE that train each expert explicitly with semantic-aware clusters by automatically routing similar clusters with the same expert.
We demonstrate the effectiveness of our model on a collection of 11 downstream tasks and achieve up to 2.45% performance improvement in Top-1 accuracy.
To the best of our knowledge, this is the first work that successfully trains Vision MoE models under SSL settings with ImageNet only and achieves state-of-the-art transfer performances.
R ELATED WORK
Self-supervised Learning Previous works mainly focus on the design of pretext tasks with image transformations (Doersch et al., 2015; Gidaris et al., 2018), inpainting (Pathak et al., 2016), colorization (Zhang et al., 2016), contrastive learning (Chen et al., 2020;
He et al., 2020; Grill et al., 2020b;
Caron et al., 2020), and for specific downstream tasks (Wang et al., 2020; Xie et al., 2020; 2021a;
Chen et al., 2021a).
Motivated by BERT Devlin et al.
(2018), masked image modeling (MIM) is recently proposed to learn by reconstructing masked images.
BEiT (Bao et al., 2021a) is the pioneering work to predict visual tokens generated by a pre-trained tokenizor (Radford et al., 2021).
SimMIM (Xie et al., 2021b) further simplifies the framework by directly utilizing pixel RGB values as the reconstruction target, and MAE (He et al., 2022) proposes an asymmetric encoder-decoder architecture for better training efficiency.
Due to its effectiveness and efficiency, MAE is selected as the primary baseline.
Although achieving significant transfer performance, most previous methods can only provide a unified representation for various downstream tasks, suffering from the risk of negative transfer, as demonstrated in Sec. 3.1.
The most relevant work to ours is SDR (Liu et al., 2022), which can train 256 subnets with 256 disjoint ImageNet subsets simultaneously.
Instead, our work differs from three perspectives, 1) the mapping between subnets and subsets in SDR is randomly selected and fixed during pre-training, while MoCE achieves self-adaptive mapping with the cluster-conditional gates; 2) progressive training is required in SDR, while MoCE shares the same dataloader with MAE; 3)
SDR is built upon SimSiam (Chen & He, 2020), a representative contrastive method, while MoCE focuses on the recently proposed MAE.
Mixture of Experts Mixture of experts(MoE) has a long history through decades (Jacobs et al., 1991; Jordan & Jacobs, 1994; Shazeer et al., 2017).
Recently, MoE is regarded as an effective way to scale up the model in natural language processing (Lepikhin et al., 2020; Fedus et al., 2021; Yang et al., 2021; Lewis et al., 2021).
With the growth of Vision Transformer, MoE for vision (Riquelme et al., 2021; Wu et al., 2022) is also explored recently.
However, we still lack a self-supervised learning MoE model that can be trained on the moderate size of the dataset, like ImageNet-1k.
We notice (Kudugunta et al., 2021; Ma et al., 2018) most resemble our work.
They regard MoE as a multi-task learning method and focus on the multi-language translation and recommendation systems respectively.
We show that under the self-supervised learning setting for image modality, an additional clustering component is crucial for learning a meaningful MoE model using ImageNet-1K only.
Moreover, benefiting from the unsupervised pre-training, MoCE is able to be deployed to any image downstream task, which differs from these two works in that downstream tasks should follow the pre-training task.
(Puigcerver et al., 2020) shares a similar setting with us except their model is pre-trained supervisedly.
However, their mapping between the expert and the data is pre-defined and fixed during training, while MoCE is learned dynamically to achieve a better result, referring to
Table 2.
Their idea is implemented as MoCE(fix) in the table.
Multi-Task Learning Multi-task Learning targets to learn a model that is appropriate for several tasks.
Some works propose to adopt a shared backbone together with multi-heads for different specific tasks, namely hard-parameter sharing, showing the effectiveness on time-series data, language and graph data separately (Liu et al., 2019;
Hu et al., 2019; McDermott et al., 2021).
Some works (Gao et al., 2021) claim that the network design may benefit from the relationship of tasks, and trains masks for different tasks.
These works assume that the information of tasks are available during model training.
In SSL pre-training, on the other hand, information of the downstream tasks are unknown.
M ETHOD
In this section, we first demonstrate the negative transfer phenomenon in MAE with a preliminary experiment in Sec. 3.1.
After that, we discuss the limitations of adopting TokenMoE (Riquelme et al., 2021) with MAE in Sec. 3.2, and propose our Mixture of Cluster-conditional Experts (MoCE), a novel paradigm for customized pre-training for various downstream tasks in Sec. 3.3.
3. N EGATIVE T RANSFER IN M ASKED A UTO E NCODER
To demonstrate the negative transfer phenomenon on MAE, in this section, we evaluate the transfer performance of MAE models pre-trained with data of distinct semantics on various downstream tasks.
First, we separate ImageNet into two disjoint splits, called Split-A and Split-B, based on their label semantic dissimilarity in WordNet tree (Miller, 1998) following (Huh et al., 2016; Liu et al., 2022).
Split-A mainly contains inanimate objects, such as cars and airplanes, while Split-B primarily involves organisms, such as plants and animals.
Then, we pre-train MAE on Split-A, Split-B and the full ImageNet respectively without data annotation, and evaluate the resulting three models on 11 downstream tasks.
Check Sec.
4.1 for more implementation details.
As demonstrated in Fig.
1, MAE pre-trained with Split-A performs the best on Aircraft, Cars,
SUN397 and DTD, while achieving the highest accuracy on Flowers, Food, Pets, Cifar-10 and Cifar100 when pre-trained with Split-B. Only two of the eleven tasks benefit from the usage of the full
ImageNet, suggesting that for datasets whose semantics are close to inanimate objects, like Aircraft,
[Figure]
Cars, SUN397 and DTD, adding pre-training data from Split-B would in contrast degrade the downstream model performance, and vice versa for datasets whose semantics are close to organisms, such as Flowers, Food, and Pets.
To conclude, the introduction of semantically irrelevant pre-training data might impede the transfer performance for specific downstream tasks.
This negative transfer phenomenon motivates us to develop an efficient and automatic paradigm for task-customized pre-training.
3. E XPLORING T OKEN M O E TO M ASKED A UTO E NCODER
Overview of TokenMoE
Here, we introduce TokenMoE (Riquelme et al., 2021), a successful customized supervised pre-training paradigm built upon ViT (Dosovitskiy et al., 2020), which mainly consists of Transformer blocks with alternating multi-head self-attention (MSA) and multi-layer perceptron (MLP).
Specifically, TokenMoE converts several Transformer blocks to Mixture of Expert (MoE) ones by expanding the MLP layer N times, with each considered as an expert , denoted as
E i ( · ) , i = 0 , 1 , 2 , ..., N .
Conditional computation is performed between N experts controlled by gate , denoted as G ( · ) , which is a linear layer whose input is the token embedding x , and the output is the top- K probability among experts as,
[Equation]
where W g is the parameters of the gate and σ is the softmax function.
TopK ( · ) will remain the top- K largest values unchanged, but set others to be 0, and K is the number of activated experts.
Thus, each token will be routed to its corresponding experts.
The final output can be represented as,
[Equation]
Importance loss and load loss are also defined to enforce the balance usage of experts as in Riquelme et al.
If not specified, we set K to be 1 and N to be 8 by default in all our experiments.
Limitation of TokenMoE. As shown in Table 2, naively adopting TokenMoE to MAE cannot achieve explicit improvement even if with intense hyper-parameter tuning and complicated data augmentations ( e.g. , Repeat Augment and RandAugment with larger magnitude).
We visualize the routing heatmaps of the pre-trained TokenMoE model in Fig. 3(a), and observe that the routing process demonstrates little correlation with ImageNet labels.
Moreover, Expert-3 nearly dominates all the semantic classes, selected by 91.8% classes as the most frequently chosen expert, which degenerates the multi-expert design by utilizing single expert frequently.
Check detailed analysis in
Sec. 4.2.
As demonstrated in Fig. 1(b), we hypothesize that the degeneration is due to the usage of low-level pixel RGB values as the reconstruction target in MAE instead of the semantic class labels used in the original TokenMoE, which is also observed in Li et al.
3. M IXTURE OF C LUSTER - CONDITIONAL E XPERTS
To address the degeneration of TokenMoE, we propose the Mixture of Cluster-conditional Experts (MoCE), which trains each expert in an explicit semantic-aware manner.
The general procedure mainly consists of data clustering, architecture and gate design, and deployment.
Data Clustering.
To train each expert semantically, a clustering procedure is first performed to simulate the label partitioning in Sec. 3.1.
With a pre-trained MAE model, we collect all image features normalized to be unit vector, so the feature matrix can be represented as F = [ f 1 , f 2 , ..., f n ] , ∀∥ f i ∥ 2 = 1 .
We denote n as the number of images and m as the desired number of clusters.
The learnable centroids of the clusters are represented as C =
[ c 1 , c 2 , ..., c m ] , ∀∥ c i ∥ 2 = 1 .
The assignment of each feature to the clusters can be calculated as A = F T C .
Following Asano et al.
(2019), we set Q ∈ R m ∗ n as the posterior distribution of clustering, and the objective can be described as,
[Equation]
where 1 m is the vector of all ones in dimension m , and H is the entropy function.
The constraints force the clustering results to be balanced.
During optimization, Q is solved by Sinkhorn-Knopp algorithm (Cuturi, 2013), while C is trained by SGD to minimize the cross entropy between Q and
A .
We optimize Q and C iteratively, and take the final C and Q as the centroid of each cluster and the clustering assignments respectively.
Architecture.
The whole network is trained on the full ImageNet simultaneously, with each expert trained by images of certain clusters depending on the routing results of the MoCE gates.
Considering the size of each cluster could be relatively small and the effective training time of each expert is K times shorter by expectation than other parameters with dense modeling ( e.g. , MSA parameters (Riquelme et al., 2021)), we further adopt a distillation loss L distill defined as the L2 distance between the feature generated by the whole network and each expert.
In general, the loss function can be formulated as,
[Equation]
where D i means the i -th cluster and θ i denotes the parameters trained by D i , which consists of several experts in different layers, as explained in the following.
Gate Design.
Following the design of TokenMoE, we also replace several MLP layers in ViT with our MoCE layers equipped with the MoCE gates.
Different from TokenMoE where each token is routed to different experts separately, we route the tokens from the images in the same cluster to the same expert.
Therefore, the output of our MoCE gates can be represented as,
[Equation]
where Cluster Emb ( · ) returns the embedding of the cluster that x belongs to.
In practice, we fix it to be centroid matrix C .
Furthermore, we empirically find that the confidence of G ( x ) is low, and consequently the mapping between clusters and experts varies a lot during pre-training.
Inspired by the design of importance and load losses, we adopt another regularization loss L imbalance to enhance the confidence of the gate result as,
[Equation]
n is the number of samples here.
For practical implementation, the loss is calculated over the samples in a batch.
The imbalance loss penalizes on the negative variance of the gate confidence.
We find
MoCE benefits from such design with more stable training.
Deployment.
During the deployment stage, we try to customize the model for fine-tuning with best transfer results.
As shown in the Sec. 3.1, the model sharing the closest semantics to the downstream task is preferred.
We achieve this by re-using the data clustering.
Specifically, we
[Table]
feed the downstream images through the pre-trained MAE model and collect all features F task .
Then, the assignment of downstream images to the clusters can be computed by A task =
F Ttask C .
We select the cluster with the most downstream images, and extract out the corresponding experts for deployment, which we call a sub-model of the total MoCE model.
Therefore, only a regular
ViT model is used for downstream fine-tuning, which is much more efficient than TokenMoE, as demonstrated in Table 3.
E XPERIMENT
In this section, we first introduce the implementation details of pre-training and fine-tuning stage of
MoCE in Sec. 4.1.
Then we demonstrate the effectiveness of MoCE by evaluating the pre-trained models on a collection of 11 downstream tasks with detailed analysis of our MoCE superior to vanilla MAE and TokenMoE in Sec.
Finally we take ablation studies on the key components of
MoCE in Sec. 4.3.
4. S ETUP
Implementation details.
For all experiments, we replace two MLP layers with MoCE layers in the original
ViT-B. Following (Wu et al., 2022), layers with the greatest gradient magnitude are selected ( i.e. the last two MLP layers in ViT-B according to our experiments).
The number of experts and the clusters are set to be 8 and 256 by default if not specified.
Our model utilizes the officially released 1600-epoch pre-trained MAE model and performs continuous training with extra 200 epochs.
Each expert is initialized by the corresponding dense model in addition with a small magnitude of noise.
The training procedure mainly follows MAE except we multiply the base learning rate by 0.1.
All regularization loss weight is set to 0.01 by default.
For distillation loss, we store the features of the whole network at the end of each epoch to avoid heavy computation for inferring the whole network at each iteration.
Downstream fine-tuning.
For downstream classification tasks, we try to set up a fair benchmark for Vision Transformers.
We mainly follow the hyper-parameter settings of existing works (Dosovitskiy et al., 2020;
Riquelme et al., 2021) and benchmark settings of (Ericsson et al., 2021).
Specifically, all models are trained by SGD with a momentum of 0.9.
Weight decay is set to be 0 and the learning rate is searched among [1e-4, 3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1].
Each model is fine-tuned for 2500 steps with cosine learning rate decay, a batch size of 64, and 224 × 224 resolu-
[Table]
[Figure]
[Table]
[Figure]
We find such a setting generates a stable result.
We compare our method with current state-of-the-art ViT self-supervised models, DINO (Caron et al., 2021), MoCoV3 (Chen et al., 2021b), BEiT (Bao et al., 2021b), in addition to the several strong-performing ResNet-50 models, e.g. BYOL (Grill et al., 2020a) and
DeepClusterV2 (Caron et al., 2018).
4. R ESULTS AND A NALYSIS
Transfer Results.
The transfer performance of different self-supervised models is summarized in
Table 1.
It can be seen that MoCE achieves 2.45% improvement over the MAE baseline and reaches the state-of-the-art averaged accuracy, demonstrating the effectiveness of the task-customized pretraining paradigm.
In addition, MoCE exhibits distinct superiority on different downstream tasks.
For the fine-grained datasets like Aircraft, Cars and Food, MoCE outperforms the baseline model by a large margin, while for tasks like Caltech, Cifar-10 and VOC, MoCE shows limited improvement, which is reasonable, since for downstream fine-grained tasks that might share similar distributions with only a subset of the pre-training dataset, MoCE is able to find the model that is trained by the cluster the most similar with the downstream task, and thus alleviate negative transfer.
However, when the downstream tasks contain images covering the whole semantic of the pre-train dataset, the negative transfer does not exist.
Therefore, MoCE only gets limited improvement.
Comparison between MoCE, TokenMoE and MAE.
The comparison is provided in Table 2, with all models trained with the same epochs for a fair comparison.
We also provide the results of SDR (Liu et al., 2022), a task-customized model aiming at alleviating negative transfer, and further re-implement their idea under ViT architecture notated as SDR(ViT).
It can be seen that
TokenMoE performs even worse than MAE, suggesting that naively adopting MoE to MAE is not feasible.
Both MoCE and SDR models exceed vanilla MAE, demonstrating the effectiveness of task-customized method for alleviating negative transfer.
MoCE further outperforms SDR(ViT), indicating the importance of self-adaptive routing.
Besides transfer accuracy, we also provide the mean of peak signal-to-noise ratio (PSNR)
(Sara et al., 2019) comparison to show the generation quality in Fig. 4(a) and MoCE still exhibits consistent improvement over TokenMoE and MAE on most datasets.
A i r c r a f t C a l t e c h C a r s C i f a r 1 0 C i f a r 1 0
D T D F l o w e r s F o o d P e t s S U N 3 9 7 v o c 2 0
R e l a t i v e
P S N R
i m p
r
o v e
m e n t
[Figure]
Analysis on experts.
We contribute the improvement of MoCE to the effective learning of experts.
As shown in the routing heatmap in Fig. 3(a), each expert of TokenMoE shows little correlation with the semantics, while each expert of MoCE in Fig.
3(b) is trained by several classes, showing a more balanced assignment of images to experts.
We also visualize the samples for 3 random experts in
Fig. 3(c).
It is interesting to find that expert 1 is mostly trained by images containing clothes, and experts 2 and 3 are trained mostly by birds and dogs respectively.
In addition to the visualization, we quantitatively show that each expert is trained by close semantics.
Following (Mikolov et al., 2013), we select the corresponding label set that is trained by each expert and compute the L2 distance between the labels inside and between the sets with label embedding produced by CLIP Radford et al.
The L2 distance inside the set averaged on 8 experts is 0.84, smaller than that between the sets 0.92, indicating that MoCE gates automatically cluster similar semantics to each expert, and consequently benefit the downstream transfer.
We further show that by stacking multiple MoCE layers, a hierarchy cluster structure is automatically formatted.
As shown in Fig. 4(b), the former MoCE layer dispatches the clusters into several experts, and the latter MoCE layer further splits the clusters into its experts.
The latter MoCE layer handles the clusters in a finer way.
Training and testing efficiency.
Besides the transfer performance of MoCE, we provide the training and testing efficiency comparison with TokenMoE. As shown in Table 3, all three models share similar FLOPs.
However, TokenMoE needs to be deployed the full model during both the training and testing phase, while MoCE only needs to deploy a single expert and therefore halves the required number of parameters when testing.
In addition, training and testing speed are improved by 10% and 30% respectively, which is attributed to the reduction of token shuffle operation as tokens in the same image do not split and will be dispatched to the same expert, significantly reducing the communication overhead.
4. A BLATION
Searching method.
When the downstream task comes, it is expensive to fine-tune each expert to choose the best one, requiring an efficient searching method to provide the task-customized expert.
Here we ablate several methods including early stop, KNN, LogME You et al.
(2021), and our method based on the clustering results, as discussed in Sec.
We also provide the best and worst sub-model for reference.
We choose Aircraft, Flowers, Food as representatives since the size of these datasets covers the smallest and largest ones.
It can be seen in Table 4 that MoCE performs stably among different sizes of the dataset, and the searching time is negligible as we only need to infer the downstream task once.
We attribute this to the nature design of MoCE that combines clustering and pre-training into a single paradigm.
[Table]
[Table] [Table]
[Table] [Table]
Scalability.
We ablate on the scalability of MoCE in three aspects.
All default settings are marked in bold .
Firstly, we vary the number of experts in each MoCE layer.
As shown in Table 5, with more experts comes explicitly consistent improvement to the averaged accuracy.
We further ablate the number and location of MoCE layers.
As mentioned in Sec.
4.1, we select MoCE layers based on the gradient magnitude, which is the 11th and 12th MLP layers (3rd row in Table 6) according to our experiments, while the last 2 even MLP layers ( i.e. , 10th and 12th layers, 2nd row in Table 6) are chosen in TokenMoE by default.
As shown in Table 6, we notice that it is essential to replace the right layer to be MoCE layer.
Adding more MoCE layers shows little improvement.
We also train MoCE with different number of clusters.
As shown in Table 7, the accuracy continuously increases when the number of clusters goes up to 256, but start to drop later.
We hypothesize that with moderate number of clusters, MoCE can produce a variety of task-customized models.
However, with even more clusters, the number of experts become the bottleneck and performance starts to saturate.
C ONCLUSION
In this work, we first show the negative transfer phenomenon that exists in the prevailing selfsupervised 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 tackle 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, TokenMoE suffers from the fact that the gate shows no correlation to the semantic 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 stablize the training process.
A search algorithm for selecting the best model for transfer is also proposed based on the clustering priors.
Extensive experiment shows that MoCE trains each expert with meaningful semantics and achieves state-of-art transfer performance on a collection of 11 downstream tasks.
In addition to the transfer accuracy, MoCE can automatically format the hierarchy structure of pretrain data and is also deployed more lightly than TokenMoE. It is the first work that successfully trains a self-supervised learning MoE model on ImageNet only.
We hope such a design will motivate more research on the self-supervised MoE models.
D ETAILS 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 and SGD optimizer iteratively.
For the SinkhornKnopp algorithm, we set the iteration number as 3.
The learning rate of 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 it costs 3 minutes and 20 seconds on average for a single GPU.
A. C OMPARISON UNDER FAIR PARAMETER COUNTS .
The setting used in our work focuses on a fair comparison of FLOPs , referring to Table 3 in the main paper.
Since TokenMoE and MoCE always activate only one expert throughout the whole pre-training and fine-tuning procedure, the FLOPs value is maintained close to MAE.
Apart from this criterion, we further provide the comparison on equal parameter counts.
As shown in Table 8, we train MAE under the same parameter count as the whole model of MoCE, and MoCE still outperforms MAE consistently.
[Table]
A.
M
O CE FOR LARGER ARCHITECTURE .
Here we provide the analysis on the larger architecture(ViT-L, 2.57× 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 and Split-B following the same setting in Sec. 3.1.
As shown in the first three rows of Table.
9, a similar phenomenon is observed that MAE trained by Split-A performs better in Aircraft, 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 a common pre-trained model is used for various downstream tasks, due to the inevitable semantic gaps between the pre-training and downstream datasets, rather than the architecture.
[Table]
A. P ERFORMANCE OF M O CE WITHOUT PRE - TRAINING .
We provide results of MoCE trained from scratch for 200 epochs in Table 10.
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 based on the clustering results.
Although it is a common practice to utilize pre-trained dense models as initialization to accelerate pre-training (Wu
[Table]
R ESULTS ON DETECTION AND SEGMENTATION .
We evaluate MoCE on the semantic segmentation task that aims to predict the class for each pixel in the image.
We report the metric of mean Intersection of Union (mIoU) 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, 2017) as the optimizer.
The learning rate is set to 1e-3 with layerwise decay to be 0.65.
We conduct fine-tuning for 160K steps.
The batch size is 16.
The detailed hyperparameters can refer to (Bao et al., 2021a).
Detection and instance segmentation task are also evaluated on COCO (Lin et al., 2014). We follow the same deployment method as instance segmentation tasks to choose the best expert. Following iBOT (Zhou et al., 2021)
we adopt the Cascade Mask 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 1e-4.
The layer-wise learning rate decay ratio is set to 0.75. We train the model for 12 epochs and decrease the learning rate by 10x at epoch 9 and 11. It can be seen in Table 11 that MoCE still outperforms MAE and achieves state-of-the-art results in both detection and semantic segmentation tasks. The results demonstrate that our task-customized procedure is also beneficial to these downstream tasks
[Table]
[Table]