CR-MoE: Consistent Routed Mixture-of-Experts for Scaling Contrastive Learning
Authors: Ziyu Jiang
Guoqing Zheng
Yu Cheng
Ahmed Hassan Awadallah
Zhangyang Wang
Venue: Accepted by TMLR ()

Abstract: While Contrastive Learning (CL) achieves great success in many downstream tasks, its good performance heavily relies on a large model capacity. As previous methods focus on scaling dense models, training and inference costs increase rapidly with model sizes, leading to large resource consumption. In this paper, we explore CL with an efficient scaling method, Mixture of Experts (MoE), to obtain a large but sparse model.  We start by plugging in the state-of-the-art CL method to MoE. However, this naive combination fails to visibly improve performance despite a much larger capacity. A closer look reveals that the naive MoE+CL model has a strong tendency to route two augmented views of the same image token to different subsets of experts: such ``cross-view instability" breaks the weight-sharing nature in CL and misleads the invariant feature learning. To address this issue, we introduce a new regularization mechanism, by enforcing expert-routing similarity between different views of the same image (or its overlapped patch tokens), while promoting expert-routing diversity of patches from different images. The resultant method, called CR-MoE, improves by 1.7 points in terms of 1\% semi-supervised learning accuracy on ImageNet, compared to the naive combination baseline. It further surpasses the state-of-the-art CL methods on ImageNet pre-training of Vision Transformer (ViT) by 2.8 points, at the same computational cost. Our findings validate CR-MoE as an effective and efficient image representation learner. Code is available at https://github.com/VITA-Group/CRMoE.

Published in Transactions on Machine Learning Research (01/2024)
CR-MoE: Consistent Routed Mixture-of-Experts for Scaling Contrastive Learning
Ziyu Jiang Texas A&M University
Bangsign@tamu.edd edu
Guoqing Zheng
zhengS@microsoft..omm
Microsoft Research
Yu Cheng
chenggn@ccs..uhhleeee
The Chinese University of Hong Kong
Ahmed Hassan Awadallah
Jassemamm@mmierrooff..comm
Microsoft Research
Zhangyang Wang University of Teros Austin
allaswnn@uutezus edu
Reviewed on OpenReview: https: //
//tp://ppppe openreview. net/ forum? id= qKIvn9nL1R
Abstract
While Contrastive Learning (CL) achieves great success in many downstream tasks, its good performance heavily relies on a large model capacity. As previous methods focus on scaling dense models, training and inference costs increase rapidly with model sizes, leading to large resource consumption. In this paper, we explore CL with efficient scaling method, Mixture of Experts (MoE), to obtain a large but sparse model. We start by plugging in the state-of-the-art CL method to MoE. However, this naive combination fails to visibly improve performance despite a much larger capacity. A closer look reveals that the naive MoE+CL model has a strong tendency to route two augmented views of the same image token to different subsets of experts: such "cross-view instability" breaks the weight-sharing nature in CL and misleads the invariant feature learning. To address this issue, introduce regularization mechanism, by enforcing expert-routing similarity between different views of the same image (or its overlapped patch tokens), while promoting expert-routing diversity of patches from different images. The resultant method, called CR-MoE, improves by 1.7 points in terms of 1% semi-supervised learning accuracy on ImageNet, compared to the naive combination bascline. It further surpasses the stato-of-tho-art CL methods on ImageNet pre- training of Vision Transformer (ViT) by 2.8 points, the same computational cost. Our findings validate CR-MoF. is available at https://githin wwt.///IA-GGrruu//CHHt....
effective and efficient image representation learner. Code
1 Introduction
U Unsupervised contrastive Learning (CL) has been popularly explored as it demonstrate strong performance on many downstream tasks, which could even beat its supervised counterpart (Chen et al. 2020c Grill et al. 2020 Caron al.] 2020 Chen et al. 2021bb Caron et al. 2021p. However, the performance of CL heavily relays on the large capacity of the employed model. For instance, in semi-supervised learning with few labels. important application of self-supervised learning Tiin al. 2020b), SimCI.R-v2 (Chen al. 2020b)

---

Published in Transactions on Machine Learning Research (01/2024)
Contrastive Learning with dense network
Contrastive Learning with Mixture of Experts (MoE)
Expert
Expert 2
Expert 3
Contrastive Learning
Contrastive Learning
Expert
Expert 2
Expert
Share weight
Different routing breaks the weight sharing mechanism
Figure 1: Routing comparison between the traditional dense network (left) and naive MoE+CL (right). In contrastive learning of the dense network, two branches always share the same weights. However, naively adopting MoE to CL can lead to different routing predictions for the same patch and break the weight- sharing mechanism.
demonstrates that scaling model parameters from 24M to 795M brings performance improvement by 17%. However, scaling dense models significantly increases the training and inference cost. For instance, The 795M model would increase the training time by 41 times, and training to full performance (1000 epochs) on TetgeNet-1K (Deng et al. 2009h requires 7000 GPU (V100) days.
In this paper, we study employing an efficient scaling method, a sparse Mixture Experts (MoE) Shazeer et al., 2017), for CL, without sacrificing training and inference efficiency. In contrast dense models that process each sample with all parameters, MoF leverages a dynamic sparse model: each sample is routed to small subset of experts. And cach expert is small Muti-Layer Perceptron (MLP) network. In this way, large candidate pools of experts can be built while only activaling a small part for each sample, making it possible to leverage large model capacity while maintaining small computational costs for training and inference. MoE has been applicd successfully in NLP applications (Lepikhin et al., 2020 Fedus et al. 2021 and was recently introduced to vision tasks but only for supervised settings (Riquelme et al. 2021
We start with directly applying CL on vision MoE models (e.g. Riquelme al. (2021)) ). However, we find this naive combination only yields marginal performance improvement compared to its dense counterpart despite much larger capacity. Looking closer, we observe that different augmented views of the same image tokens are mostly routed to different subsets of experts (as illustrated in Figure 1), This essentially breaks the conventional design of contrasting shared weight branches (Chen et al. 2020a He et al. 2020 Grill et al.] 2020) and turns to contrasting independent branches, which we show hurts performance with further empirical evaluations.
To enforce consistency in expert selections for augmented image views, naive way is to always assign them the same set of experts. However, this leaks the learning target of CL: the instance identity, causing the model to overfit such trivial nuisance without learning meaningful image representations (Chen et al. 2021a). Instead, as shown in Figure we proposc simple yet effective regularization mechanism to enforce the consistency of expert. selection based on visual overlapping. Specifically, first. we pair all image tokens based on the overlapping between patches. Then we pull the selection of experts of paired tokens to similar while differentiating that for tokens from different images through the proposed Overlapping-basse Gate Aligning Regularization (OGAR). The resulting method, termed CR-MoE, significantly improves the consistency of the experts selection for different augments of the same image and the 1% semi-supervised
2

---

Published in Transactions on Machine Learning Research (01/2024)
performance by 1.7 points compared to the naive plugin, which is also 2.8 points higher than competing state-of-the-art CL methods on VIT.
Our contributions are summarized as follows:
We propose CR-MoE, which efficiently scales Contrastive Learning (CL) with the sparse Mixture of Experts, pushing the limit of CL towards large model capacity while maint similar computation cost.
We identify the problem of naively combining MoF and CL, which essentially routes semantically similar images to different sets of experts thus hurting performance, and address it by proposing a novel regularization loss.
Extensive experiments verifies the effectiveness of the proposed regularization term. Compared to competitive state-of-the-art CT. methods on ViT, the proposed CR-MoR achieves an improvement of 2.8 points at the same computational cost.
2 Related works
2.1 Self-supervised training
Inspired by the observation that conducting instance recognition could yield good representation that naturally clusters the same class images (Alexey et al. 2016 Wu et al. 2018), various works devote to designing self-supervised learning through pulling the reprosentations of the same images together while pushing those of different images apart (Chen et al. 2020c & He et al. 2020 Tian et al. 2020a), also known as contrastive learning. Some works also recognize that negative samples are not necessary (Grill et al. 2020; Misra & Maaten 2020; Chen He. 2021: Zbontar et al. 2021). A trend was observed and verified by (Chen et al.) [2020a]p b) that contrastive learning yields better performance with a longer training schedule and a large backbone model. However, training large models with CL for long schedule imposes significantly high training costs. In this work, to scale CL we investigate an efficient scaling option based on Mixture- of- df-ExprrtsMMo)... While recent work (Meng et al., 2022) also starts to explore sparsifying the contrastive learning with dynamic pruning strategies, MoE has its unique strength on memory elficiency and combining it with contrastive learning is still not explored.
Other works on self-supervised learning focus on the handcrafted pretext tasks (Trinh et al., 2019 like rotation prediction (Gidaris et al. 2018), jigsaw (Noroozi & Favaro 2016 Carlucci et al.) 2019) and col- orization (Gidaris et al., 2018). Recent advances in transformer highlight the possibility of a new class of self-supervised learning methods through masked image modeling (Bao et al., 2021 He et al. 2021b Xie ct al. 2021). These conceptually different directions can also be combined with contrastive learning to fur- ther boosting the performance (Dangovski et al. 2021 Zhou et al. 2021). In this work, we focus on studying contrastive lcarning while lcaving other directions as potential future work.
2.2 Sparse Mixture of Experts
The traditional Mixture of Experts Network is composed of multiple sub-models and conduct input con- ditional computation (Jacobs et al.) 1991 Jordan & Jacobs 1994 Chen et al.) 1999 Yuksel et al. 2012 Roller et al. 202T). While contrastive learning can also be improved with the traditional MoF. (Tsai et al. 2020), suffers from intensive computation since the model are dense and all experts are activated. Recent work (hazeer et al., (017)) proposes the Sparse Mixture of Experts Layer and demonstrates better results on language modeling with lower computational cost. Following works devise methods to further address the communication cost (Fedus et al. 2021 Lewis et al. 2021) and stability (Zoph et al.] 2022 issues. GLaM (Du
et al. performance.
2021 studies the MoE for language self-supervised task and achieve significant downstream few-shot
MoE recently applied for computer vision tasks (Riquelme al. 2021
Gross al. 2017
Xue et al. 2021
Wang et al. 2020 Tsai et al. [2018 Ahmed et al. 2016 Yang et al.
2019
Pavlitskaya et 2020 However,
most of these works focus only on supervised weakly supervised learning. Recently, L.MM Mustaffa et al.
3

---

Published in Transactions on Machine Learning Research (01/2024)
starts explore applying MoE on self-supervised language-image pairing tasks, where they propose local and global entropy design to balance different modalities. In this work, we reveal and address the challenge from inconsistent expert routing when applying MoE to self-supervised vision tasks.
2022
3 Method
3.1 Preliminaries
Contrastive learning Contrastive learning is a self-supervised method via maximizing instance discrimi- nativeness. For example, it enforces the similarity of positive pairs while enlarging the distance of negative pairs (Wu et al. 2018):
N log 8+ (Ui 17 Dujev ev 87(11; vi)
87 (Vi, v25))
M(G, ot V . = N
where otth is considered a positive sample of sample o while the set V- consists of negative samples. 8, (1t, wrf) exp (1. uz /r) measures the similarity of positive pair (v,, ozt) while 8, (v,, v, measure the similarity negative pair (1), T is the temperature controlling the magnitude of all terms.
MoCo-v3 CCen al. 2021b is one of the stateoof-th-art self-supervised methods devised for VIT (oso vitskly al. 2020). It encodes two crops C1 and C2 for each image under random data augmentation. The images are Ihen encoded with network and its Exponential Moving Average (EMA). MoCo-v3 also introduce random token projection to stabilize the learning process. The loss of MoCo-v3 is defined as
8 (21,
LeL =
log ST f2
(2)
N
E Ei- 8ff1,f-)
where the features (21 encoded from (C1, C2), respectively, are employed as positive samples while negative set {f}- is composed by the fcatures of views from other images.
Sparse Mixture of Experts MoE reduces the computational cost via activating a small subset compu- tational graph for each sample. The basic building block of MoE the sparse MoE layers, which consists of ne expert networks (E,,Eg, En.)) Formally, MoE layer is defined as
y Σ G(z)LE()))
(3)
where and are the input and output, respectively. G is the gating function that outputs a vector containing scores for each expert network Ez),, typically instantiated with Softmax. By picking the top-k scored experts (k <<< ne), the model only activates a small subset of expert networks for each sample. For G, we employ the noisy top-k gating design introduced in Riquelme et al. 2021 as
G(z) = Top <), k)
(4)
where W is learnable weight and denotes Gaussian noise sampled from N (0, +/2)) Wr controls the clean score of the gating function while noise e benefits the load balancing between experts. The sum of the score are then normalized with Softmax function and sparsified with TopK defined as
if is in the top ls elements of 0 otherwise.
TopK(u, kt;;
In this work, we focus on studying applying MoE for the VIT Dosovitskiy al.) 2020) backbone. We follow the strategy of Riquelme et al. (2021) to replace every other multi-layer perceptron (MLP) layers with sparse
4

---

Published in Transactions on Machine Learning Research (01/2024)
Load balanceness loss c*
Overlapping based Gate Aligning Regularization ç
Gate
MLP
Multi-Head Attention
Norm
Norm
MLP
MLP
Overiap based matching
Contrastive
Moving-average
Loss
Stop Gradient
Figure 2: Pipeline of the proposed CR-MoE. It replaces every other block of ViT to sparse MoE layer. Overlapping based gate aligning regularization is applied for training the proposed network.
MoE layers. Each expert network is of the same architecture: Mz) = W20 gelu (Wiz), where Wi E Rd xdj and W22 xdm are learanble weights while gelu is the non-ling -ineiinerr activation layer (Hendrycks & Gimpel 2016). It is worth noting that MoE is applied to multiple visual tokens, where each token could have different expert choice.
We also employ an auxiliary loss encourage the load balancedness following Shazeer al. (017 termed Llb to prevent the over-solection of few experts.
3.2 Sparse Mixture of Experts for Contrastive Learning
To enforce the consistency of expert selection while not leaking image identity, we introduce new regulariza- tion term called Overlapping based Gate Aligning Regularization (OGAR). In ViT, MoF layer would choose experts for each token. The token sequence includes one classification token and multiple patch tokens. We then introduce how OGAR is applied for classification and patch tokens.
OGAR for classification tokens As the classification token is at the image level, enforcing consistency can be easily realized by applying the similarity constraint among classification tokens for augments of the same image. Formally, it is defined
CECLI M(G GiCLS) (G(cLs) .T)
sy sr G[LLS) (iLLs)
(6)
log 8, (CCLss
N
Eoers 8y (G|Css) G GiLL))
where G CCLS and [[CS) denote the gating function output (G(x) in Equation 4 of classification tokens from pair positive samples. [G[cLs]] denotes those from negative samples. 7 is the temperature, where we use the same value as LCL. We employ the form of Moco V3 loss to enforce the consistency for preventing all the gate functions collapse to always outputting the same prediction.
5

---

Published in Transactions on Machine Learning Research (01/2024)
OGAR for patch tokens Unlike classification tokens, different patches lack one-to-one correspondence the patches are randomly sampled from different regions of the original image. Hence matching the patches is required before conducting the regularization. Previous studies reveal that the transformer can automatically learn object segmentation that aligus well with input in terms of the spatial location (Caron ct al., 2021 which indicales the strong spatial correlation belween input and features learned by CL. Inspired by this observation, we design matching method based on the spatial location of the patches. As shown Figure 2 each patch of one view is paired with the most overlapping patch from the other view. For one patch that is not overlapped enough with any other patches (below certain overlapping threshold A), leave it umpaired. Only those paired patches are utilized for calculating the loss. Formally, the proposed loss on patch Pm is defined as
1 N
e.(Gm.Gn) E log s, (0mman) Lo-etor- Es0r-- 8, (-m
if ToU mm > A
pm
n = arg max loUnni
(7)
0
otherwise.
where Pr denotes the patch that has the largest Intersection over Union (IoU) with pm- IoU mn represents the ToU between patch m and n. Gm and Gn are gating function outputs for Pm and Pu; respectively. {G} denotes those from negative patch samples. When the loUnn is less than threshold X, the loss is 0. Otherwise, the consistency loss between Gm and Gn would be employed. The overall gating loss is averaged over all patches as
N
C
LG
Np
Pm
where No is the number of patches.
Some previous works study similar problem: enforcing the regional regularization of CL (Li al., 2021
Wang et al. 2021), which also requires matching the local features. They match the features across two views based on the feature distance (e.g. cosine similarity). However, we empirically find this approach yield less significant improvement in our case. The intuition behind this is that the paired features in inter-mediate layers may lack strong feature similarity. The proposed matching method allows the existence of non-paired patches while the design of (Wang et al. 2021) assumees all local features can be paired, which is prone to noise in learned features and also in general does not hold in practice.
We balance the two regularization terms with a convex combination controlled with a weight a (0 < a < 1). Formally, the resulting OGAR is
_ = - a) + ac
(9)
The overall optimization target for CR-MoE To sum up, the overall loss is
L= Lct. + wibC wa
(10)
where wpb and we are the scaling factor of the loading balancedness losses and OGAR, respectively. By employing OGAR on naive MoE+CL, the resultant CR-MoE framework can efficiently scaling contrastive learning with MoE.
4 Experiment
4.1 Settings
Pre-training Our pre-training experiments are conducted on ImageNel-1KK (Deng et al. 2009) follow- ing common practice (Chen et al. 2020aa He et al. 2020). For pre-training framework, we employ Moco (Chen et al. (0221b), and we follow the same settings as Moco v3 on data augmentations and learning specification: 3-layer MLP projection head, temperature T 0.2, momentum m 0.99, random patch projection, cosine decay schedule (Loshchilov & Hutter 2016), and 40-epoch warmup. For optimization, we

---

Published in Transactions on Machine Learning Research (01/2024)
employ AdamW [oshchilov & Hutter 2017 optimizer and weight decay of 0.1. We employ linear scaling rule (Goyal et al. 2017) and search for the best base learning rate (Ir) on 100-epoch results with grid of {1.5e ,330e 1,00= 5.0e 1.0e.-- The best searched Ir is 5.0e-4 x BatchSize 2256. For model ablations, we employ a shorter schodule of 100 epochs with a relatively small batch size of 1024. When comparing with stale-of-the-ar methods, we scale up and employ 300 epochs with a batch size of 3072.
Linear probing Linear probing measures the quality of learned representations from pre-training. After self-supervised pre-training, we remove the MLP heads and train classifier with the frozen backbone. Following Moco V3, employ the SGD optimizer with batch size of 4096 and weight decay of 0 for 90 epochs, with only random resized cropping and flipping augmentation. The Ir is swept following common practice (Chen et al. 2021b Zhou et al., 2021).
Semi-supervised and transfer few-shot learning Learning with few labels is an important application for contrastive learning, which pertains to both semi-supervised and transfer few-shot learning (Chen et al. 2020b: Tian al.] b020bbb Islam el. al.) 2021). Specifically, for semi-supervised learning, we consider 1% or 10% available labels (following the sampling in [Cen et al. (2020b)) of ImageNet. For transfer few-shot
learning, consider 4-shot and 10-shot settings for three datasets: CIFAR10 (Krizhevsky et al. 2009
Pet37 (Parkhi et al. 2012) and Food101 (Bossard et al. 2014).
For these two applications, we consider a two steps paradigm: The model is first pre-trained on the pre-train and then it is supervised fine-tune on the seed or few-shot dataset. For the supervised fine-tune step, we employ different settings for different tasks. As suggested in Tian al. (2020b); Zhou et al.] (2021), we train linear classifier frozen features for ImageNet 1% semi-supervised task and all transfer few-shot tasks. We optimize for 800 epochs with batchsize of 256 while other scttings keeps the same lincar probing. For ImageNet 10% semi-supervised task, we follow Chen et al. (2020b); [Zhou et al. (2021) fine-tuning from the first layer of the MLP head. The epochs mumber is set as 200 while the Ir are searched with grid of (1e-5, ne- 5,,3-5, 1e-,,3e-1
Hyper-pane Syper-parameters for Mixture-of-EEperts Model and loss For MoE network, we by default employ 16 expert candi- dates (n. 16) and always activate 2 of them (k- = 2). For the employed loss terms, employ À 0.2, 0.3, wib 0.01 and we = 0.001, which are searched on 100-epoch training.
Model
Parameters
FLOPs
ResNet50
25M
4.1G
VIT-S/16
22M
4.6G
VMoF-S/16 ViT-B/16
72M
4.6G
For each expert network, we choose dj = 2dm instead of dy- = 4dm in Chen et al.] (2021b) to keep the computational cost of activating 2 experts the same as that in VIT. The em- ployed model is VMoEE-/166, as shown in table [ its FLOPs are comparable to ViT-S/16. Moreover, we further compare the training and inference computation costs in terms of GPU time cost. For inference of a single image one A6000 GPU, the time costs are 1.25ms and 1.07ms for VMoE-S/16 and ViT/S-16, respectively. For training a batch of 1024 images on A6000 GPUs, the time costs are 1.579s and 1.425s for VMoE-S/16 and ViT/S-16, respectively. VMoE-S/16 is only marginally slower than ViT-S/16 in both cases.
87M
17.6G
Table 1: Network architecture comparison for four different architectures. CR-MoE uses VMoE-S/16 as the backbone.
Computation Framework Our implementation is based on Pytorch (Pasuke et al. 2019) and Fast-MoE He et al. 2021a) library. Models are pre-trained 32 Nvidia. V100 GPUs.
4.2 Naive Combination of MoE and CL Does Not Work
In this section, we look into the "cross-view instability" issuc of directly plugging MoE to CL and show how the proposed regularization address this problem.
The routing is inconsistent To check the consistency of the expert decision, as shown in Figure Ba we exclude random cropping and flipping from data augmentations to ensure we can locate the different views of the same patches: they are always in the same position in this way. Further, we define these patches with the same content as corresponding tokens while defining the tokens from other images the non-corresponding
7

---

Published in Transactions on Machine Learning Research (01/2024)
Corresponding patch
Non-oomesponding potch
corresponding and non-corresponding patches
corresponding tokens
corresponding tokens
of Shared Experts
non-corresponding tokens
1.4
non-correspondin tokens
1.5
1.0
0.8
1.0
of
0.4
0.2
layer
layer
(b) [CLS] tokens w/o regu
(c) Patch tokens w/o regu
corresponding
corresponding
Number Shared Experts
2.0
non- non-cornespending tokens
Shared Experta
1.4
non-corresponding tokens
1.2
1.5
0.8
0.6
0.2
Mot layer Index
Mot layar Index
(d) (CLS) tokens w/ regu
(e) Patch tokens w/ regu
Figure 3: (a) Illustration for the definition of corresponding and non-corresponding patches. The rest. four figures compare the average number of shared experts for G(x) between corresponding and non-corresponding tokens. (b) shows the number of shared experts for classification tokens while (c)(e) represents the number of shared experts for patch tokens. All of them measured across different layers. The "w/o regu" in (b) denotes they are from the naive combination CL and MoF, model. In contrast, (d) is from the proposed CR-MoE. The x-axis of the last four figures is the index of the MoE layer in VMoE.
tokens. Then, we calculate the average number of shared experts (the number of experts selected by both tokens in the pair) for corvesponding tokens and non-correspen non-corresponding tokens and make comparison.
As shown Figure 3b for classification tokens the naive combination, the gating function always selects the similar number of shared experts between corresponding and non-corresponding patches, which means the difference between corresponding patches and non-corresponding patches can hardly be distinguished. For the patch tokens, presented in Figure 3e the boundary between corresponding and non-correspoodinn patches get blurred in the deep layers. This would change the standard contrasting shared weight backbone fashion of CL to contrasting (partially) non-shared weight contrastive learning
Inconsistent routing leads performance dropping Unfortunately, the proof-of-concept. experiments verify that performance of (partially) non-shared weight contrastive learning can drop. Specifically, we
'Contrasting with moving average network can regarded sharing weight the moving average would converge the online value when training stabilizes.
8

---

Published in Transactions on Machine Learning Research (01/2024)
Table 2: Linear probing (denote as linear) and 1% imagenet semi-superrvsed (denote as 1%) performance comparison for pre-training and evaluation on ImageNet. All the reported accuracy is top accuracy (%). expert 0/1 for sep-Vis ViT-V denote the two different paths of sep- ViT-S.
Method
Model
Linear
1%
Moco v3
ViT-S/16
69.7
53.5
Moco v3
sep- SIp-ViT-S/10 (expert 0)
68.4
48.5
Moco v3
sep-ViT-S/16 (expert 1)
68.5
48.7
Moco v3
V-MoE-S/16 V-MoE-S/16
69.9
54.1
CR-MoE (Ours)
70.7
55.8
Table 3: Comparison with State-of-The-Ard methods in terms of linear probing (denote Linear), 1% and 10% semi-supervised performance (denote 1% and 10%, respectively). All the reported accuracy is top 1 accuracy (%). The SD denotes self-distillation.
Match method
Model
Linear
1%
10%
SimCLR v2 Chen et al. 2020b)
Resnet50
71.7
57.9
68.1
SimCLR v2 + SD (Chen et al. 2020b)
Resnet50
71.7
60.0
70.5
Moco v3 CCen al. 2021b
ViT-S/16
73.4
59.4
72.2
CR-MoE (ours)
V-MoE-S/16
74.1
62.2
73.0
designed special network called sep- which has the same backbone architecture as MoE with 2 expert candidates. For routing, would activate different experts for different branches. In this way, these two branches would not share weight in the expert network. The result is illustrated in Table 2 the ViT-V (expert 0) decrease the performance by 0.8% and 5% for linear probing and 1% semi-supervised performance compared to the baseline, respectively, indicating that (partially) non-shared weight can hurt the performance for CT. (especially for the semi-supervised performance).
The proposed CR-MoE improves both consis- tency and performance After employing the pro- posed classification alignment and OGAR, shown in Figure 3d and 3e the proposed CR-MoE suc- cessfully increase the number of shared experts for corresponding tokens while reducing or keeping the number of the shared experts for non-corresponding tokens. Also, shown in Table 2 in contrast the naive combination of CT. and MoF that only im- proves the baseline Moco V3 by small margin of 0.2% and 0.6% in terms of linear probing and 1% semi-supervised performance, the proposed CR-MoF. increase this margin to 1.0% and 2.3%, demonstrat- ing the effectiveness of the proposed method.
Table 4: Transfer few-shot performance comparison across different datasets between MocoV3 and the proposed CR-MoE with ViT-S/16 and V-MoE-S/16, respectively. 4-shot and 10-shot denote and 10 sam- ples available for each class for downstream tasks, re- spectively. All the reported accuracy top accu- racy (%).
Dataset
Method
4-shot
10-shot
Moco V3
72.9
80.1
CIFAR10
CR-MoE
74.4
80.7
Moco V3
71.8
81.4
Pet37
CR-MoE
74.4
84.3
4.3 Comparison with State-of of- The-Art Methods
Moco V3 CR-MoE
35.2
48.8
Food101
37.4
50.1
In this section, compare the proposed CR-MoE with state-of-the-art methods. For fair comparison, we employ a longer training schedule of 300 epochs
following Chen et al. (0021b); Caron et. al.
2021 Zhou et al. 20221)
CR-MoE yield better in-domain performance As shown in Table the proposed CR-MoE achieves
highest performance in terms of Linear probing, 1% and 10% semi-supervised learning. Remarkably, com- pared to Moco v3 on VIT-S/16, the proposed CR-MoF, significantly improves the 1% semi-supervised per-

---

Published in Transactions on Machine Learning Research (01/2024)
DE
MOJU
3
D0G
60
(a) Expert 1 (Characters)
(b) Expert 2 (Faces)
(c) Expert 4 (Pool/Sea)
(d) Expert (Forest/
Visualization of the patch tokens routed to different experts in the 7th layer of CR-MoE on ImageNet. The patches with different patterns are routed to different experts.
Figure 4:
formance by 2.8%. Mcanwhile, there also a non-trivial improvement on linear evaluation and 10% semi- supervised performance by 0.6% and 0.8%, respectively. Since the Moco V3 and CR-MoE share the same CL framework, this demonstrate the effectiveness of MoE framework and the proposed regularization. The large improvement semi-supervised performance also matches the observation at Chen et al. (022b that large capacity helps more for few-shot learning.
CR-MoE yields better transfer few-shot performance We then study if the strong in-domain few- shot performance can transfer to downstream datasets. As demonstrated in Table ! the proposed CR-MoE also yields consistent improvement of (1.5%,0.6%). [2.6%,2.9%] and 12.2%,1.3% for CIFAR10, Pet37 and Food101, respectively, in terms of [4-shot, 10-shot] performance, demonstrating the proposed CR-MoE can also significantly improve the downstream few-shot performance.
4.4 Ablation Studies
visualize the routing distri- bution of CR-MoE in Figure g Even though no semantic caption or labels are involved in the training, we find that the patches routed to different tokens show distinct semantic patterns. For example, the patches of (Characters, Faces, Pool/Sea, Forest/Trree] are routed to Expert [1, 2, 4, 5]. respectively.
Visualization of routing choices Following LIMoE (Mustafa et al. 2022
OGAR loss for patch tokens matters As shown in Table 5 when removing OGAR loss for patch tokens by setting
the linear evaluation and 1% semi-supervised performance would drop by [0.1%, 0.4%). which demonstrate the effectiveness of the proposed OGAR loss for patch tokens. Other hyper- paramtt changes like wa 0.01 and À 0.1 only marginally change the performance.
Table 5: Comparison between different hyper- parameters settings of the proposed CR-MoE. Lin- ear probing (denote as linear) and 1% imagenet semi-supervised (denote as 1%) performance are re- ported. All the reported accuracy is top accuracy (%). The first row denotes the employed hyper- parameter setting. Error bar calculated by run- ning 3 times with different random seeds.
Table 6: OGAR loss ablation regarding different. matching methods and whether to employ nega- tive samples. FSM denotes the Feature Similarity- based Matching method employed in Li et. al.
2021
Wang et al. (021
WG
Linear
1%
Matching Method
Negative samples
Linear
1%
0.001
0.3
0.2
70.7+-0.07
55.8++0.25
0.001
0.0
0.2
70.6::0.1
55.4+10.14
FSM
70.5
55.6+10.37
0.01
0.3
0.2
70.5
55.8
Overlap
70.2
54.2
0.001
0.3
0.1
70.6
55.9
Overlap
/
70.7+0.07
55.8+0.25
10

---

Published in Transactions on Machine Learning Research (01/2024)
In T we ablation study the proposed OGAR loss. When discarding the negative samples for OGAR loss and only enforcing consistency as in Grill et al. 2020), we observe the gating function tent to choose the same experts for all samples even though we have employed the loading balance loss. Meanwhile, the performance would largely decrease, showing that negative samples are necessary for OGAR. When switching from the overlap-based matching method to the Feature Similarity-based Matching method (FSM), the linear probing and 1% semi-supervised performance would both incur a drop of 0.2%. Moreover, we further compare with FSM in terms of transfer few-shot learning and we find that FSM matching method achieves 68.5% and 80.2% in terms of 4-shot and 10-shot transfer few-shot accuracy Pet37, respectively. In contrast, the proposed overlapping-
matching method significantly improves 4-shot and 10-shot transfer few-shot accuracy by 1.5% and 0.9%, respectively, demonstrating the effectiveness of the proposed matching method.
The number of experts matters We conducted an ablation study concerning the number of experts Re: as detailed Table 7 Our findings suggest that increasing the number of experts can lead to an in- crease in both linear evaluation performance and 1% few-shot performance. For instance, by increasing the experts' number from 2 to 16, the linear evaluation and 1% few-shot performance significantly increase by 1.3% and 5.6%, respectively.
Table 7: Comparison between different numbers of
experts ne for the proposed CR-MoE. Linear probing (denoted as linear) and 1% imagenet semi-supervised (denoted as 1%) performance are reported. All the reported accuracy is top 1 accuracy (%).
Linear
1%
69.4
50.2
4
70.2
53.1
5 Conclusion
8
70.6
55.3
16
70.7
55.8
In this work, we study an efficient way of scaling con- trastive learning with sparse Mixture of Experts. We start from naively plugging in the MoE to CL and observe that the naive combination lends route different views of the same image to different subsets of experts, thus breaking invariant feature learning and hurting the performance of downstream tasks. To tackle this problem, we propose novel regulariza- tion framework to promote consistency of experts selection on the same (or overlapped) image tokens while encouraging diversity of the experts selection for different images. Extensive evaluations on multiple down- stream tasks demonstrate the proposed framework, CR-MoE, effectively improves the routing consistency and the overall performance of downstream tasks without incrcasing the computation cost.
Broader Impact and Limitation The proposed CR-MoF show the possibility of scaling Contrastive Learning with large sparse neural network, which greatly reduces the training and inference time and energy consumption while achieving stato-of-the-art performance. They can serve the goal of GreenAI for self-supervised learning. On the other hand, in this work, we mostly focus on academic datasets. However, in practice, unlabeled datasets in the wild may come with imbalances and adversarial samples, which could lead to performance or fairness issues. One of the future directions is to extend CR-MoE to such imbalance or adversarial settings.
References
Karim Ahmed, Mohammad Haris Baig, and Lorenzo Torresani. Network of experts for large-scale image categorization. In European Conference on Computer Vision, pp. 316-532. Springer, 2016.
Dosovitskiy Alexey, Philipp Fischer, Jost Tobias, Martin Riedmiller Springenberg, and Thomas Brox. Dis- criminative, unsupervised feature learning with exemplar convolutional, neural networks. IEEE TPAMI, 38(9))1734 1747, 2016.
Hangbo Bao, Li Dong, and Furu Wei. Beit: Bert pre-training of image transformers. ar preprint arXiv:2106.0880 2021.
Lukas Bossard, Matthieu Guillaumin, and Lue Van Gool. Food-101 mining discriminative components with random forests. In European Conference on Computer Vision, 2014.
11

---

Published in Transactions on Machine Learning Research (01/2024)
Fabio M Carlucci, Antonio D' D'Innocente, Silvia Bucci, Barbara Caputo, and Tatiana Tommasi. Domain generalization by solving jigsaw puzzles. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 2229 2238, 2019.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsu- pervised learning of visual features by contrasting cluster assignments. Advances in Neural Information 2020.
Processing Systems,
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 9650-9660, 2021.
Ke Chen, Lei Xu, and Huisheng Chi. Improved learning algorithms for mixture of experts in multiclass classification. Newral networks, (299):122911582 1999.
Ting Chen, Simon Korublith, Mohammad Norouzi, and Geoffrey Hinton. simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020a.
Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey E Hinton. Big self- supervised models are strong semi-supervised learners. Advances in neural information processing systems, 33:22243 23:22243-22225,, 2020b.
Ting Chen, Calvin Luo, and Lala Li. Intriguing properties of contrastive losses. Advances in Neural Infor- mation Processing Systems, 34, 2021a.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pottern Recognition, pp. 15750 15758, 2021.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming Hc. Improved basclines with momentum contrastive
learning, Xiv preprint arXiv:2003.044 04297, 2020c.
Xinlci Chen, Saining Xie, and Kaiming He. An empirical study training self-supervised vision transformers. In Proceedings of the IREE/CVE International Conference on Compuler Vision, pp. 9640 9649, 2021b.
Rumen Dangovski, Li Jing, Charlotte Loh, Seungwook Han, Akash Srivastava, Brian Cheung, Pulkit Agrawal, and Marin Soljačić. Equivariant contrastive learning. arXiv preprint arXix:211..00 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. 2009 IEEE conference on computer vision and pottern recognition, pp. 248-255. Icee, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa. Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words:
Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1 11929, 2020.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yangi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture- of-experts. Xiv preprint. ar Xiv:2: 12. 06905, 2021.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXim:2101.0320 001611 2021.
Spyros Gidaris, Praveer Singh, and Nikos Komodakis. Unsupervised representation learning by predicting image rotations. arXiv preprint arXiu:1808.077e5 2018.
Priya Goyal, Piotr Dollár, Ross Girshick, Picter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint or Xie: 1706. 02677, 2017.
12

---

Published in Transactions on Machine Learning Research (01/2024)
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in Neural Information Processing Systems, 23:21271-21284, 2020.
Sam Gross, Marc' Aurelio Ranzato, and Arthur Szlam. Hard mixtures of experts for large scale weakly supervised vision. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition. pp. (885-68777 2017.
Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang, Jidong Zhai, and Jie Tang. Fastmoe: A fast mixture-of expert training system. arXiv preprint arXiv:2103. 13262, 2021a.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF c on computer vision and pattern recognition, pp. 1729-97338 2020.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. arXiv preprint orXiv2211..0671 2021b.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXio:1606.08/21 2016.
Ashraful Islam, Chun-Fu Richard Chen, Rameswar Panda, Leonid Karlinsky, Richard Radke, and Roge- rio Feris. A broad study the transferability of visual represenlations with contrastive learning. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 8845 8855, 2021.
Robert A Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79 87, 1991.
Michael I Jordan and Robert A Jacobs. Hicrarchical mixtures of experts and the algorithm. Neural compulation, 6(2):181 214, 1994.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of fcatures from tiny images. 2009.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint ftt:2006.16666 2020.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning, pp. 0265-6274. PMLR, 2021.
Chunyuan Li, Jianwei Yang, Pengchuan Zhang, Mei Gao, Bin Xiao, Xiyang Dai, Lu Yuan, and Jianfeng Gao. Efficient self-supervised vision transformers for representation learning. or preprint arXiv:2106.000 06709785 2021.
Tya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint
ar xrxiv:6608.038 2016.
Iya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXie preprint orXiv:1711.05101. 2017.
Jian Meng, Li Yang, Jinwoo Shin, Deliang Fan, and Jac-sun Seo. Contrastive dual gating: Learning sparse features with contrastive learning. In Proceedings of the IEEF/CVE Conference on Computer Vision and Pattern Recognition, pp. 12257-12265, 2022.
Ishan Misra and Laurens van der Maaten. Self-supervised learning of pretext-invariant representations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pottern Recognition, pp. 6707-6717, 2020.
13

---

Published in Transactions on Machine Learning Research (01/2024)
Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Multimodal con- trastive learning with limoe: the language-inage mixture of experts. or ar preprinl or orXim:2206.0270 2022.
Mehdi Noroozi and Paolo Favaro. Unsupervised learning of visual representations by solving jigsaw puzzles. European conference on computer vision, pp. 69 84. Springer, 2016.
Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and CV Jawahar. Cats and dogs. In 2012 IEEE conference computer vision and pattern recognition, pp. $498-3508 IEEE, 2012.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Svetlana Pavlitskaya, Christian Hubschneider, Michael Weber, Ruby Moritz, Fabian Huger, Peter Schlicht, and Marius Zollner. Using mixture of expert models to gain insights into semantic segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pp.
342 343, 2020.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Su- sano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture experts. Adoances in Newral Information Processing Systems, 34, 2021.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34, 2021.
Noam Shazeer, Azalia Mirhoseini, Kraysziof Maziarz, Andy Davis, Quoc Le, Geolfrey Hinton, and Jell Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint 1oXi::701.0663 2017.
Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola. What makes for good views for contrastive learning? Advances in Neural Information Processing Systems, 33:6827-6839, 2020a.
Yonglong Tian, Yue Wang, Dilip Krishnan, Joshua B Tenenbaum, and Phillip Isola. Rethinking few-shot image classification: a good embedding is all you need? In European Conference on Computer Vision, pp. 266 282. Springer, 2020b.
Trieu Trinh, Minh-Thang Luong, and Quoe V Le. Selfie: Self-supervised pretraining for image embedding. aixiv preprint 1rXi::1906.022 02940, 2019.
Tsung Wei Tsai, Chongsuan Li, and Jun Zhu. Mice: Mixture of contrastive experts for unsupervised image clustering, In International Conference on Learning Representations, 2020.
Yao-Hung Hubert Tsai, Paul Pu Liang, Amir Zadeh, Louis- Philippe Morency, and Ruslan Salakhutdinov. Learning factorized multimodal representations. arXiv preprint 1Xii::806.060 2018.
Xin Wang, Fisher Yu, Lisa Dunlap, Yi-An Ma, Ruth Wang, Azalia Mirhoseini, Trevor Darrell, and Joseph F, Gonzalez. Deep mixture of experts via shallow embedding. In Uncertainty in artificial intelligence, pp. 552-562. PMLR, 2020.
Xinlong Wang, Rufeng Zhang, Chunhua Shen, Tao Kong, and Lei Li. Dense contrastive learning for self- supervised visual pre-training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3024-3033, 2021.
Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via non-parametric instance discrimination. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3733 3742, 2018.
14

---

Published in Transactions on Machine Learning Research (01/2024)
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. arxiv preprint arXiv:21 000-2111.09880 2021.
Fushao Xue, Ziji Shi, Futao Wei, Yuxuan Lou, Yong Liu, and Yang You. Go wider instead of deeper. arXiv preprint ar (rXf:21007118117 2021.
Brandon Yang, Gabriel Bender, Quoe V Le, and Jiquan Ngiam. Condconv: Conditionally parameterized convolutions for efficient inference. Advances in Neural Information Processing Systems, 32, 2019.
Seniha Esen Yuksel, Joseph N Wilson, and Paul Gader. Twenty years of mixture of experts. IEEE transactions neural networks and learning systems, 2378):1177-1193, 2012.
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphanc Deny. Barlow twins: Self-supervised learning via redundancy reduction. International Conference on Machine Learning, pp. 12310-1221 12320. PMLR, 2021.
Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. ibot: Image bert pre-training with online tokenizer. arXiv preprint ar 11x::2211.07832, 2021.
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazcer, and William Fedus. Designing effective sparse expert models. arXiv preprint arXiv:2202.085 08906, 2022.
A Appendix
A.1 With Other Backbones
Table 8: Illustration of the performance in terms of linear probing (denote Linear), 1% semi-supervised performance (denote as 1%). Pet37@4-shot and Pet37@10-shot represent the transfer few-shot learning performance on Pet37 when 4 and 10 samples are available for each class, respectively. All the reported accuracy is top 1 accuracy (%).
Match method
Model
Linear
1%
Pet37@4-shot
Pet37@10-shot
Moco v3 Chen et al. 2021bb
ViT-B/16
76.7
63.9
74.2
84.5
CR-MoE (ours)
V-MoE-B/16
76.3
64.9
76.7
85.0
different backbone. As shown in Table S CR-MoE with V-MoE-B/16 surpasses the Moco V3 with VrT-B/16 by 1% in terms of the 1% few-shot performance while leading to small drop of 0.4% on linear evaluation performance. Moreover, it improves the transfer few-shot learning performance on Pet37 by 2.5% and 0.5% for 4-shot and 10-shot performance, respectively.
In this section, we explore the performance of the proposed method
(a) Expert 0
(b) Expert
(c) Expert
(d) Expert
Figure 5: Visualization of the patch tokens routed to different experts in the 1st layer of CR-MoF. ImageNet. The patches with different patterns are routed to different experts.
15

---

Published in Transactions on Machine Learning Research (01/2024)
nsti
MMM
P9
FO
E
(a) Expert
(b) Expert
(c) Expert (d) Expert 13
Figure Visualization of the patch tokens routed to different experts the 11th layer of CR-MoE on ImageNet. The patches with different patterns are routed different experts.
A.2 Visualizing experts routing of more layers
We further analyze the routing pattern of 1st and 11th layers for CR-MoE, which correspond to the first and last MoE layers in the network, respectively. As shown in Figure 5 and Figure [6] similar patches would also be routed to the same experts on these layers. Moreover, we find that the patches with the same low-level pattern (c.g. edges) are often routed to the same expert in the shallow layer (c.g. the 1st layor). Mcanwhile, the patches with similar semantic information are often routed the same expert in the deep layer (e.g. the 7th and 11th layers).
16


Reviews:
[{"review_id": "3xzRgFqZwq", "reviewer": "Reviewer_Fz5M", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The method is simple and effective. It may be an interesting paper for other researchers doing MoE.\n2. The paper is well-written and easy to follow.\n3. The experiments are clear and well-presented.\n\nWeaknesses:\n1. The citation is really bad. I think the authors should read the style template to distinguish between \\citet{} and \\citep{}.\n2. I think the authors should write the loss function completely (Eq. 6 and 7) instead of just referencing it. \n3. In Sec. 3.2, the authors reference two papers and say that those papers also proposed to enforce regional regularization. However, in the experiments, there are no comparisons against such methods. \n4. As other vision MoE (e.g., vMoE) paper shows, after training, similar patches will automatically be routed to the same expert. Why such a phenomenon is different in CL, i.e., you need extra regularization to enforce similar patches being routed to the same expert? \n5. How generalizable is this method? E.g., how about applying it to the supervised training?\n6. I think the training and inference computation costs should also be compared. MoE approach should be slower in practice.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed Consistent Routed Mixture-of-Experts (CR-MoE) for self-supervised contrastive learning (CL). Based on the native combination of CL and MoE, this paper introduced consistent routing regularization, which encourages the corresponding tokens to be routed to the same set of experts.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Modify the citation format.\n2. Write the loss function.\n3. Comparison with similar methods.\n4. Comparison of the computation costs.", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "yBt9BPVkzj", "reviewer": "Reviewer_cRMt", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\nIt is often shown that self-supervised methods benefit from scaling more than supervised counterparts.\nProposing ways to train self-supervised methos in more efficient/effective ways is a valuable contribution to the field.\nThe proposed method with roughly 3x more parameters than the baseline (but with similar FLOPs) improves the baseline on downstream tasks.\n\n### Weaknesses\n\nAccording to the motivation of the paper, i.e. maintaining the consistency of gating tokens that belong to the same patch, I am not sure how Equations-6/7 help in that regard.\nFirst of all, notation is not clear.\nThese equations follow the notation from Riquelme et al. 2021, but there is no need to specify tokens for different images.\nCurrently, Equation-6 is no different than Equation-2 (the original InfoNCE loss of MoCo-v3), yet Equations-2/6 both appear in the ultimate loss (Equation-10) used to train the model.\nSimilarly, Equation-7 just \"filters\" tokens which do not overlap with any other token according to some threshold $\\lambda$.\nWhich means that, Equation-7 is actually a mechanism to do a better \"pair selection\".\nTherefore, I'd strongly encourage the authors to compare to works like [A] or to baselines where MoCo-v3 is trained by positive pairs which overlap with at least some threshold.\n\nBased on my previous comment, the illustrations shown in Figure-1/3 might be misleading.\nAs Equation-7 filters some patch tokens**$\\dagger$, the number of expert selection is reduced in consequence.\n\nExperiments do not extensively validate the main claim of the paper: scaling the number of parameters.\nThe proposed method is only trained with the ViT-S architecture containing 72M parameters (for 16 experts) and compared against ViT-S, with 22M parameters.\nTo validate the claims, it'd be nice to see:\n- How the number of experts impact the results\n- The performance of a network with the same number of parameters (72M)\n\nThe weight of the proposed loss $w_G$ is so small, i.e. 1e-3.\nI wonder whether this is because the range of this loss term is high.\nOtherwise it is surprising that the proposed loss actually improves performance.\nThis makes me question even more whether the improvements come from the careful selection of tokens for positive pairs (my comment above)\n\nVisualization of the patch tokens routed to different experts in Figure-4 is obtained by the assignments of the CR-MoE layer in the 7th transformer block.\nI wonder if such behaviour is generalized across the network, or Figure-4 is just cherry-picked.\n\n[A] What Makes for Good Views for Contrastive Learning? Tian et al. 2020, NeurIPS 2020\nFinds best positive pairs for self-supervised contrastive learning.\n\n$\\dagger$: It would also be good to show how many tokens are discarded based on different thresholds.\n\n\n### Some minor points\n\n- Figure-1 is misplaced, or its caption should be much more informative.\nWhile reading the introduction first time, I had to check the method and experiments sections to actually understand this figure.\nFor instance, it wasn't clear what \"expert\" was, what \"corresponding tokens\" meant, or out of how many experts the y-axis ranged\nMaybe Figure-3a can be moved next to Figure-1.\n- What is FFN?\n- applying similarity constrain --> applying THE similarity constrain", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Scaling up the number of parameters for neural networks often brings perfomance gains, at the cost of increased compute budget.\nThe idea of \"Mixture of Experts\" (MoE) from Riquelme et al. 2021 addresses this issue, where there are multiple MLP layers in each transformer block  and only a few of them are selected to process each token by a gating mechanism.\nThis way, the number of parameters is increased (hence the complexity of the original network) while adding a marginal computational complexity overhead.\n\nThis paper proposes a method that allows training contrastive self-supervised learning (SSL) methods with more parameters while maintaining their compute requirements.\nIt is claimed that MoE cannot be applied to MoCo-v3 naively in this setting, because patches that are shared between the two images of a positive pair can be assigned to different \"experts\".\nTherefore, the weight-sharing mechanism of contrastive SSL is not valid anymore.\nTo circument this problem, a new loss term is proposed that reqularizes the assignment of patch tokens across experts.\nThe experiments show that the proposed model trained on ImageNet-1K learn better representations when evaluated on ImageNet-1K or 3 transfer datasets (CIFAR10, Pet37, Food101).", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I expect the authors to respond to my concerns and improve the manuscript accordingly.", "broader_impact_concerns": "I do not see a need for a broader impact statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ratzNHtxVD", "reviewer": "Reviewer_N628", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- The methodology of encouraging consistency between two gating functions on the global and local levels is well motivated.\n- The proposed method shows a wide margin on transfer learning.\n\n**Weaknesses**\n- The MoE model is not really widely adopted, hence affecting the impact of the paper.\n- The difference between regular ViT and regular MoE using MoCo is small. Table 2 last two rows of results show that the added loss is not really helping too much, so the only noticeable benefit is on transfer learning. \n- The difference between the proposed method and the cited FSM method is non-existent (Table 6).\n- The motivation of the paper is to scale models more efficiently using sparsely connected experts. The experiments do not showcase that. We don’t know if MoE with the proposed changes is really going to scale vision models as the results focus on relatively light weight models.\n- There is no analysis on the alpha parameter. It would be good to have a sweep of hyperparameters to show the sensitivity and the necessity of each loss component.\n- The paper writing has a large room to improve. In the introduction, it assumes knowledge of the MoE model with a lot of jargon, which is hard to parse before reading the background section. Figure 1 is not understandable until introducing the model section. E.g. what is corresponding token, what is shared expert, what is shared weight branch and independent branch. I would suggest moving it to experiments and properly defining the terminologies.\n- All of the citations need to be fixed. Choose properly between \\citep and \\citet.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes to add a consistency loss in MoE-based visual encoder models for contrastive visual representation learning. MoE is a type of visual transformer model where each block has multiple parallel expert functions and a gating function that selects a distribution of experts. The problem with MoE is that the gating function can select different experts for the same image, and similar experts for different images, which makes the contrastive learning objective less effective. To mitigate the problem, the paper proposes to add a contrastive loss on the gating distribution. For classification tokens, it treats the gating function as an additional “feature,” and encourages the gating to be similar on positive pairs and dissimilar on negative pairs. For patch tokens, the local overlap IoU is taken into consideration when computing the loss on the gating functions. The resulting method show case that it can improve slightly upon regular vision transformers, especially on lower data regime and transfer learning settings. However, the differences do not seem too significant. It is not much better than standard contrastive learning, except on transfer learning tasks. It is not better than FSM. In conclusion, the paper may be technically correct but the results may not be significant.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Add large scale experiments\n- Add hyperparameter studies\n- Improve paper writing\n- Fix citations", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies applying contrastive learning (CL) to Mixture of Expert (MoE) model, in order to address the scaling issue brought by training with large, dense models. The paper identifies that a naive MoE+CL approach tends to route two augmented views of the same image token to different subsets of experts, which contrasts with the CL working mechanism and affects invariant feature learning; the paper then proposes a regularization by enforcing expert-routing similarity between different views of the same image and promoting expert-routing diversity of patches from different images. Results show the efficacy of the proposed regularization.

Reviewer cRMt has doubts on the validity of Eqns (6) and (7), i.e., the proposed regularization,  and thinks that its efficacy for performance improvement is not fully tested; cRMt thinks that Eqns (6) and (7) function as "pair selection", rather than as regularization of the routing in MoE, and consequently the improvement may come from pair selection. AE reads the paper and thinks that cRMt may have misunderstanding of the paper - although  Eqns (6) and (7) look similar to Eqn. (2) in terms of the function form, their effect on learning is totally different; in fact they impose regularization on routing of experts, while Eqn. (2) of CL is for feature learning itself. Authors have made such clarifications in the response and revised paper. Another important suggestion from cRMt is to study the routing effects from the proposed regularization for transformer blocks other than the 7th one, for which the authors conducted further experiments and show similar routing effects across the ViT network.

Both reviewers cRMt and N628 point out that applying the proposed method on ViT-S cannot fully justify the proposed regularization; the authors conducted experiments on the larger ViT-B and got consistent results. Reviewer N628 worries about impact of the work since MoE models are not widely used, which however, does not affect decision making for TMLR. In the revised paper, the authors also added additional analysis on the alpha parameter, as required by N628.

Reviewer Fz5M suggests to compare the training and inference computation costs for the MoE version of ViT; the authors conducted the comparison in the revision and the results confirm the claims made in the paper. Reviewers also point out writing issues of the paper; most of these have been fixed in the revised paper, but there exist some remaining ones. For example, there is a typo in the denominator of Eqn. (6); in P5, spatially correlation -> spatial correlation. 

Overall, the paper is technically correct but the results are not very significant. Given that two reviewers support acceptance of the paper, while the other reviewer has a misunderstanding of the paper, we recommend acceptance. The authors please thoroughly check the paper again and make sure there are no more typos. Please also make sure all the revisions made in the response phase are included in the final version.
Decision: Accept (poster)