C OLLABORATION OF E XPERTS : A CHIEVING 80% T OP -1 A CCURACY ON I MAGE N ET WITH 100M FLOP S
Anonymous authors
Paper under double-blind review
A BSTRACT
In this paper, we propose a Collaboration of Experts (CoE) framework to assemble the expertise of multiple networks towards a common goal.
Each expert is an individual network with expertise on a unique portion of the dataset, contributing to the collective capacity.
Given a sample, delegator selects an expert and simultaneously outputs a rough prediction to trigger potential early termination.
For each model in CoE, we propose a novel training algorithm with three components: weight generation module (WGM), label generation module (LGM) and selection reweighting module (SRM). WGM adapts the losses of experts, enabling them to focus on different portions of the dataset. LGM generates the label to constitute the loss of delegator for expert selection. SRM aims to promote delegator to select experts better.
CoE achieves the state-of-the-art performance on ImageNet, 80.7% top-1 accuracy with 194M FLOPs.
Combined with PWLU activation function and CondConv, CoE further boosts the accuracy to 80.0% with only 100M FLOPs for the ﬁrst time.
Furthermore, experimental results on the translation task also demonstrate the strong generalizability of CoE. CoE is hardware-friendly, yielding a 3 ∼ 6x acceleration compared with existing conditional computation approaches.
I NTRODUCTION
There are many approaches for model collaboration, among which ensemble learning (Hansen &
Salamon, 1990; Wen et al., 2020; Wenzel et al., 2020) is a popular one.
Ensemble learning uses a consensus scheme to decide the collective result by vote.
However, it requires multiple forward passes, leading to a signiﬁcant runtime cost.
MIMO (Havasi et al., 2021) draws inspiration from model sparsity (Frankle & Carbin, 2019) and tries to ensemble several subnetworks within one regular network.
It only needs one single forward pass of the regular network but is incompatible with compact models.
Conditional computation methods (Cheng et al., 2020; Yan et al., 2015; Shazeer et al., 2017) alleviates this issue via delegation scheme, i.e. assigning one or several, rather than all models, conditionally to make the prediction.
Some recently proposed conditional computation methods (Zhang et al., 2020b; Yang et al., 2019; Zhang et al., 2021) have achieved remarkable performance based on dynamic convolution.
Nonetheless, they usually have high memory access cost (MAC) and a low degree of parallelism, which increases the real latency (Ma et al., 2018).
Motivated by this, we propose the Collaboration of Experts (CoE) framework to both eliminate the need for multiple forward passes and keep hardware-friendly.
CoE consists of one delegator and multiple experts. Firstly, delegator gives a rough prediction and makes the expert selection. If the rough prediction is unreliable, the selected expert will make the reﬁned prediction.
Otherwise, the procedure will be early terminated to save FLOPs. Moreover, we only need to load the selected expert into memory, thus keep the ratio of MAC to FLOPs as a constant. By contrast, dynamic convolution methods (Zhang et al., 2020b; 2021) need to load a large number of parameters, namely basis models or experts, to synthesize the input-dependent ones. It enlarges MAC and reduces the degree of parallelism, resulting in a signiﬁcant deceleration.
To make each model in CoE play its role, we propose a novel training algorithm (as shown in Fig.1) which consists of three components: weight generation module (WGM), label generation module (LGM) and selection reweighting module (SRM). LGM generates the label (selection label) to constitute loss of delegator for expert selection (selection loss).
Selection label is a one-hot vector, indicating the suitable expert for each given input. Due to the random initialization of experts, selection labels are irregular in the early training stage. Nonetheless, delegator tends to learn generalizable patterns ﬁrst, since networks learn gradually more complex hypotheses during 𝓛𝓛 S
Delegator 0 1 0
PlanePanda Iron Jean Task Label:
Expert Selection
Rough Prediction
Expert 1 Expert 2 Expert
Experts
RefinedPredictions
Label Generation   𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸 1 𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸 2 𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸𝐸
Selection Label:
Weight Generation : Data flow w/o gradient : Data flow with gradient
TCP Values
Selection Reweighting 𝑣𝑣
CE CE : Cross-entropy loss function
TCP : TCP function 𝓛𝓛 𝑪𝑪𝑪𝑪 𝑳𝑳𝑳𝑳𝑳𝑳𝑳𝑳 𝑷𝑷 𝓛𝓛 𝑷𝑷 𝓛𝓛 T𝟏𝟏 𝓛𝓛 T 𝓛𝓛 T
W 1 𝑊𝑊𝑊𝑊𝓛𝓛 T :
Weighted sum
Training
[Figure]
CE CE
TCP
CE training (Arpit et al., 2017).
With delegator as the bridge, WGM can partition the training data into portions based on generalizable patterns, and reweight expert losses to impel each expert to focus on one portion.
This procedure makes selection labels more regular in return.
Thus, delegator avoids overﬁtting to the irregular labels. SRM aims to promote delegator to select experts better.
It is achieved by enforcing delegator to focus on samples whose recognition results are sensitive to expert selection.
Experimental results on ImageNe demonstrate the superiority of CoE. It achieves 78.2/80.7% topaccuracy with only 100/194M FLOPs, while the accuracy for ensembled models (Hansen & Salamon, 1990) is 79.6% with 920M FLOPs.
Compared with dynamic network approaches, CoE is more hardware-friendly.
It not only outperforms the SOTA dynamic method BasisNet which achieves 80.0% accuracy with 198M FLOPs (Zhang et al., 2021), but also accomplishes a 3.1x speedup on hardware.
Besides, CoE can be equipped with CondConv and further improve the accuracy to 79.2/81.5% with 102/214M FLOPs.
More surprisingly, we further boost the accuracy to 80.0% with only 100M FLOPs for the ﬁrst time by using PWLU activation function (Zhou et al., 2021).
Experimental results on the translation task also demonstrate the strong generalizability of CoE.
The contributions of this paper can be summarized as follows:
• We propose a collaboration framework named Collaboration of Experts (CoE) and demonstrate that it can lead to outstanding performance with little computation cost.
Moreover, it is hardware-friendly and able to achieve real speedup.
• We present a novel optimization strategy for CoE.
The core insight is to promote diversity within experts by distributing their expertise over different portions of the dataset.
• We update the state-of-the-art on ImageNet for mobile setting, achieving 80.0% topaccuracy with only 100M FLOPs for the ﬁrst time.
R ELATED W ORK 2. E NSEMBLE L EARNING AND M ODEL S ELECTION
Ensemble learning (Hansen & Salamon, 1990) aims at combining the predictions from several models to get a more robust one.
Some recently proposed literatures (Wen et al., 2020; Wenzel et al., 2020) demonstrate that signiﬁcant gains can be achieved with negligible additional parameters compared to the original model.
However, these methods still require multiple (typically, 4-10) forward passes for prediction, leading to a signiﬁcant runtime cost.
Differently, CoE utilizes a delegator to select only one expert for the reﬁned prediction, thus at most two forward passes are needed.
MIMO (Havasi et al., 2021) draws inspiration from model sparsity (Frankle & Carbin, 2019) and holds the view that multiple independent subnetworks can be concurrently trained within one regular network because of the heavy parameter redundancy.
Therefore, those subnetworks can be ensembled with a single forward pass of the regular model.
However, MIMO cannot be applied to compact models which have already been pruned or the ones constructed by AutoML methods (Zhong et al., 2018; Zhang et al., 2020a;
Cai et al., 2020).
It is because these models are compact enough and have few redundant parameters.
By contrast, CoE is free from the compactness of experts and compatible with various models.
Recently, some works about model selection are proposed (Li et al., 2021b; You et al., 2021).
These methods are concerned with ranking a number of pre-trained models and ﬁnding the one transfers best to a downstream task of interest.
Therefore, they select models task-wisely.
By contrast, CoE aims at improving the task performance via selecting the most suitable expert for each sample instance-wisely.
Moreover, these methods conduct model selection based on a set of samples (training set) which cannot be adopted in CoE.
2. D YNAMIC N ETWORKS
Dynamic networks achieve high performance with low computation cost by conditionally varying the network parameters (Zhang et al., 2020b; Yang et al., 2019) or network architectures (Yuan et al., 2020).
HD-CNN (Yan et al., 2015) and HydraNet (Mullapudi et al., 2018)
select branches based on the category, they cluster all categories into n groups, where n is the number of branches.
While CoE learns the model selection pattern automatically, it can be based on any property, rather than limited to the category.
MoE (Shazeer et al., 2017) and Switch Transformer (Fedus et al., 2021) enable the direct training of Router by scaling the output feature of experts with the predicted gate-values of Router.
These methods aim at conditionally selecting a speciﬁc layer or block.
Differently, CoE can take more advantage of conditional computation as the selection of whole network makes every parameter input-dependent.
The recently proposed Dynamic Convolution methods (Zhang et al., 2020b; Yang et al., 2019; Chen et al., 2020) share the same idea and achieve remarkable performance with low
FLOPs but high latency.
It is because these methods need to load many basis models or experts to synthesize the dynamic parameters, causing high MAC and low degree of parallelism (Ma et al., 2018).
By contrast, CoE only needs to load the selected expert into memory, avoiding these problems.
Finally yet importantly, batch processing is an important method to enhance the degree of parallelism.
Because of the input-dependent parameters (Zhang et al., 2021) or architectures (Yuan et al., 2020), these methods cannot process samples in batch.
Differently, CoE supports batch processing because the number of experts is limited and each one of them corresponds to many test samples.
M ETHOD
[Figure]
As shown in Fig.1, CoE consists of a delegator and n experts, a total of n + 1 individual neural networks.
To make the best of limited capacity, each expert is encouraged to focus on just one unique portion of the dataset.
This is achieved by WGM which reweights training losses of experts.
Given a sample, delegator selects an expert and simultaneously outputs a rough prediction to trigger potential early termination.
It is trained with cross-entropy loss and labels generated from LGM that indicate the suitable expert.
Moreover, SRM aims to boost delegator by enforcing it to focus on samples whose recognition results are sensitive to expert selection.
Since the inference of delegator is conducted all the time, we prefer to make delegator more lightweight than expert.
Delegator consists of three modules: feature extractor, task predictor and expert selector as shown in Fig.2.
Based on the feature derived from feature extractor, task predictor and expert selector output the probabilities for classiﬁcation and expert selection respectively.
We describe the inference procedure and training strategy of CoE comprehensively in the following subsections.
The number of samples and experts are denoted as m and n respectively.
3. I NFERENCE P ROCEDUREF
CoE ﬁrstly uses delegator to obtain the rough prediction and determine the selected expert for each sample.
Afterward, Maximum Class Probability (MCP, Corbière et al. 2019) of rough prediction is calculated.
It is the probability of the predicted class.
For samples with MCP larger than a given threshold τ , the ﬁnal recognition result is derived from the rough prediction (early termination).
Other samples are partitioned into n groups based on which expert is selected.
Then batch processing can be conducted within each group to obtain the reﬁned prediction.
The averaged FLOPs/Instance of
CoE ranges from F D to F D + F E by varying τ , F D and F E are FLOPs of delegator and experts.
3. L ABEL G ENERATION M ODULE (LGM)
LGM generates the label (selection label) to constitute loss of delegator for expert selection (selection loss). Selection label is a one-hot vector, indicating the suitable expert for a given input.
Model accuracy can be measured by True Class Probability (TCP, Corbière et al. 2019):
[Equation]
where, x j is the j-th sample, Y and y j are the predicted and true class.
But accuracy is not the only factor for suitability. For example, when models are of different sizes, the larger one usually has a higher TCP.
But it may not be more suitable, due to the large inference cost. Given that our concern is not the optimization of network architecture, we can suppose no expert is superior to others ( No Superiority Assumption , NSA).
Motivated by NSA, we leverage the standardized TCP as the metric for sutability:
[Equation]
where, Mean ( TCP : ,k ) and Std ( TCP : ,k ) are mean value and standard deviation for TCPs of Expert k .
Given m samples and n experts, selection labels can be denoted by a binary matrix L m × n , thus each row of which is a selection label.
According to NSA, the sum of each colum vector of L m × n should be same, namely (cid:80) j L j,k
= mn for k = 1 , ..., n . Therefore, L m × n can be obtatined by maximizing the sum of standardized TCP (i.e. S j,k in Eq.2):
[Equation]
This problem can be modeled as the balanced transportation problem (BTP, Shore 1970), where each sample is a supply source with a supply of one, each expert is a demand source with a demand of m/n . − S j,k is the per-unit transportation cost from the j-th supply source to the k-th demand source. We solve this problem via Vogel approximation method (VAM, Shore 1970) as introduced in appendix A.1, which is a short-cut approach to invariably obtain a good solution.
3. W EIGHT G ENERATION M ODULE (WGM)
To maximize the collective capacity of CoE, the dataset needs to be partitioned into portions then
WGM encourages each expert to focus on one portion by reweighting losses of experts.
A naive partition can be based on expert suitability, namely, partitioning the dataset with selection labels generated from LGM.
However, it results in a poor generalization to delegator. Assuming Expert k is suitable on a sample x j , this partition manner will make Expert k more suitable in return, owing to the reweighting in WGM. Thus selection labels cannot be updated. Moreover, selection labels are irregular in the early training stage because of the random initialization, thus selection labels will remain irregular consistently.
With the training going on, delegator will overﬁt to those irregular labels, yielding poor generalization as shown in Fig.3a. As a result, CoE performs poorly because delegator can hardly select a suitable expert during validation. This is also veriﬁed in Appendix B.4.3. Since networks learn gradually more complex hypotheses during training (Arpit et al., 2017), delegator tends to learn generalizable patterns ﬁrst.
Therefore, the partition can be based on generalizable patterns with delegator as the bridge.
In this way, selection labels get more regular in return thanks to the reweighting of expert losses in WGM. As shown in Fig.3b, delegator avoids overﬁtting to the irregular labels. Given m samples and n experts, delegator outputs a probability matrix P m × n ∈ R m × n , whose element P j,k ∈ [0 , 1] represents the probability of assigning the j-th sample to the k-th expert. The partition is indicated by an assignment matrix
A m × n , with one-hot row vectors.
A j,k = 1 means the j-th sample is assigned to the k-th expert.
As analyzed above, it is better to partition the training data based on P m × n , thus A m × n can be obtained by maximizing (cid:80) j,k P j,k ∗ A j,k .
Moreover, according to NSA, the number of samples assigned to each expert should be same, i.e. A m × n needs to satisfy (cid:80)
j A j,k
= m/n .
Thus, A m × n is optimized by:
[Equation]
[Figure]
Hypothes
[Figure]
sofdelegator
[Figure]
Hypothes i sofdelegator :
[Figure]
This problem can also be modeled as BTP, and solved via VAM as described in section 3.2.
In the early training stage, experts are underﬁtted.
Thus we cannot trust A m × n and need to make the gap between loss weights for experts smaller.
We achieve this by smoothing A m × n to A m × n with Eq.5, where α grows linearly from 0.2 to 0.8 with the training going on, [Equation] Finally, the output of WGM (i.e. W m × n ) is obtained by normalizing A m × n with the coefﬁcient
Z = (cid:80)
[Equation] 3. S ELECTION R EWEIGHTING M ODULE (SRM) Expert suitablity can be measured with Standardized TCP (Eq.2).
If experts have similar suitablities for a given sample, the expert selection will have little inﬂuence on the ﬁnal performance of CoE. As the capacity of delegator is limited, it should pay less attention to those samples. This is achieved by SRM, which reweights losses of delegator based on statistic that reﬂects suitablity similarity. The output of SRM is denoted as { v j | j = 1 , ..., m } , which are weights for selection losses over m samples.
The suitabilities for experts over the j-th sample, i.e. { S j,k | k = 1 , ..., n } are denoted as S j, : .
A small value of the standard deviation Std ( S j, : ) indicates experts have similar suitabilities on the j-th sample, thus v j gets smaller.
Therefore SRM determines the value of v j by:
[Equation]
3. T RAINING METHOD
The training framework of CoE is shown in Fig 1, which consists of two stages.
During the ﬁrst stage, only feature extractor and task predictor of delegator are trained to minimize L P , cross-entropy loss for the rough prediction of delegator.
During the second stage, we ﬁx these two modules, jointly optimize the expert selector and experts with L Total :
[Equation]
here, the hyperparameter η is set as 0.8.
L S is used to optimize the expert selector.
Based on the selection label L j, : , we can get the crossentropy loss L jCE for the j-th sample.
Then L S is determined by the weighted sum of {L jCE | j = 1 , . . .
, m } with weights { v j | j = 1 , . . .
, m } output by SRM:
[Equation]
L T is used to optimize the experts.
Based on the class labels of m samples, we can get m × n cross-entropy losses {L j,kT | j = 1 , . . .
, m ; k = 1 , . . . , n } , where L j,kT is the cross-entropy loss for the k -th expert on the j -th sample.
Then L T is obtained by the weighted sum of L j,kT with weights W j,k output by WGM:
[Equation]
Each of our experiments includes either four or sixteen experts in this paper.
When using four experts, the training method is same as described.
However, we will meet slow-convergence problem if the number of experts is sixteen.
To alleviate this problem, we propose a strategy which is described in Appendix A.2.
E XPERIMENTS
We conduct the main experiments on ImageNet classiﬁcation task (Deng et al., 2009).
After comparing with some popular methods in terms of computation and memory cost, we verify the superiority of
CoE over some other existing model collaboration methods.
Moreover, we try to generalize CoE to the translation task and re-evaluate CoE using Reassessed Labels (ReaL) (Beyer et al., 2020).
Finally, we try to analyze the reasonability of learned expert selection patterns.
Due to the page limit, elaborated ablation studies are illustrated in Appendix B.4.
4. I MPLEMENTATION D ETAILS
We conduct experiments with two settings: CoE-Small and CoE-Large .
For CoE-Small, we take
TinyNet-E (Han et al., 2020b) with 24M FLOPs as the feature extractor of delegator by removing the last fully connected layer.
We use OFA-110 (Cai et al., 2020) with 110M FLOPs as the expert.
For the CoE-Large, MobileNetV3-Small (Howard et al., 2019) with 56M FLOPs is adopted to construct the delegator by analogy.
We use OFA-230 as the experts.
We have also tried to introduce
CondConv (Zhang et al., 2020b) and PWLU activation fuction (Zhou et al., 2021) to achieve the extreme performance.
To combine with CondConv, we replace the convolutions within each inverted residual block of the experts with CondConv ( expert _ num = 4 ).
To take advantage of PWLU, we replace all activation layers except those that have tiny input feature maps as illustrated in Zhou et al.
Models are trained using SGD optimizer with 0.9 momentum.
We use a mini-batch size of 4096, and a weight decay of 0.00002.
Cosine learning rate decay is adopted and the number of training iterations is 313000.
We use ﬁxed auto-augment (Cubuk et al., 2019) as well.
Inspired by
BasisNet, we use knowledge distillation with EfﬁcientNet-B2 (Tan & Le, 2019; Xie et al., 2020) as the teacher.
The learning rate is 0.8/1.6 for CoE-Small/Large and dropout rate is 0.2.
The stochastic depth (Huang et al., 2016) is used except for TinyNet-E with a survival probability of 0.8.
4. R ESULTS AND A NALYSIS 4.2. A CCURACY AND C OMPUTATION C OST
Accuracy curves for CoE are drawn in Fig.4a by varying the threshold of early termination (section 3.1).
We pick out a point from each curve to compare with some efﬁcient networks in Table 1.
Our method achieves 78.2% and 80.7% top-1 accuracy with 16 experts and averaged FLOPs/Instance as 100M and 194M respectively.
Compared with OFA, CoE reduces the FLOPs from 230M to 100M and from 595M to 194M, with better top-1 accuracy.
Compared with EfﬁcientNet-B1 with noisy student training, CoE also reduces the FLOPs by 3.6x while improving the accuracy by 0.5%.
Though dynamic networks like GFNet, CondConv and BasisNet are more efﬁcient than traditional networks,
CoE still has signiﬁcantly higher accuracy with smaller FLOPs.
Compared with these approaches, CoE improves the accuracy by 2.2/2.4/0.7% respectively.
When combined with CondConv, we can achieve 79.2% and 81.5% top-1 accuracy with only 102M and 214M FLOPs respectively, which indicates CoE is complementary to dynamic networks like CondConv.
On the contrary, as CondConv and BasisNet share similar essence, namely using a group of basis to dynamically synthesize the input-dependent convolution kernel, the combination of them only arouses little collaborative beneﬁt with the top-1 accuracy of only 80.5%.
More surprisingly, we achieve the accuracy of 80.0% with only 100M FLOPs for the ﬁrst time by further making use of PWLU.
4.2. I NFERENCE S PEED AND M EMORY C OST
Compared with conditional computation methods (Yang et al., 2019; Zhang et al., 2021), CoE is more hardware friendly.
To verify the advantage, we also analyze the inference latency on hardware.
[Table] The experiments are conducted on CPU platform (Intel(R) Xeon(R) CPU E5-2699 v4 @
2.20GHz) with PyTorch version as 1.8.0.
We report the averaged latency on the ImageNet validation set in
Table 2. We notice the discrepancy between FLOPs and real speed.
For example, OFA-230 has 1.6x FLOPs compared with GhostNet 1.0x, but the speed is 1.2x faster.
Moreover, this discrepancy can be enlarged by CondConv.
CondConv-EfﬁcientNet-B0 has similar FLOPs with the original EfﬁcientNet-B0, but the speed is 1.7x slower.
BasisNet synthesizes the dynamic parameters all at once, rather than the “layer by layer” fashion like CondConv, thus is more efﬁcient.
However, it still needs to load a large number of parameters for this synthesis, which brings a large MAC.
This is why CoE (16 experts) can reduce 14.09% latency than BasisNet when the mini-batch size is one.
Finally yet importantly, BasisNet and CondConv do not support batch processing, while CoE ( experts) can take advantage of it (seciton 3.1) to further achieve a 3.1/6.1x speedup compared with them. We analyze the memory cost from two perspectives: the number of parameters and MAC. As can be seen from Table 2, the accuracy of CoE-Large (4 experts) is no worse than BasisNet and
CondConv-EfﬁcientNet-B0 when using similar parameters. Besides, the averaged MAC/Instance of CoE is much smaller than theirs. Compared with GhostNet 1.3x, the accuracy for Coe-Large ( experts) is 5.0% higher with a smaller MAC.
[Table]
M 67.4% GhostNet 1.0x 39. 141M 5.2M 5.2M 73.9% TinyNet-B 34.
202M 3.7M 3.7M 75.0% MobileNetV3-Large 31.
219M 5.4M 5.4M 75.2% GhostNet 1.3x 43. 226M 7.3M 7.3M 75.7% OFA33. 230M 5.8M 5.8M 76.9% EfﬁcientNet-B 49. 391M 5.3M 5.3M 77.2% TinyNet-A 45.
339M 5.1M 5.1M 77.7% CondConv-EfﬁcientNet-B 81.
- 413M 24.0M 24.0M 78.3% BasisNet 40. - 198M 24.9M 24.9M 80.0% CoE-Large (4 experts) 38.
220M 6.6M 25.7M 79.9% CoE-Large (16 experts) 34. 194M 6.0M 95.3M 80.7% [Figure] [Figure]
4.2. A NALYSIS OF THE TRAINING COST
To achieve superior performance with less inference FLOPs and latency [Figure] CoE may consume more training time.
For example, the applying of CoE (4 experts) on OFA-230 improves the accuracy from 78.0% to 79.9%, but at the expense of a 2.2x training cost.
To verify whether the improvement still exists with similar training cost, we get a series of accuracies by varying the number of training epochs as shown in Fig.4b, where “xx ep.
” means the number of training epochs is “xx” and
Tesla-V100-PCIe-16GB GPUs are used for training.
It is seen that CoE can boost the performance (from 78.3% to 79.9%) even if the training cost is controlled to be similiar (127 vs 141 GPU days).
4. C OMPARISON WITH O THER M ODEL C OLLABORATION A PPROACHES 4.3. C OMPARISON WITH R EGULAR M ODEL E NSEMBLE
We train four OFA-230 models with different initialization seeds.
The results are shown in Table 4.
The random seed only causes minor variety in accuracy.
However, an improvement of 1.6% is still achieved via ensemble.
It is because these models fall into different local minima, yielding the diversity of output.
Compared with the ensemble, CoE achieves 0.3/1.3% higher accuracies.
It indicates CoE can realize more potential of each individual model.
Besides, CoE keeps the computation cost constant, while model ensemble increases the computation cost by four times.
4.3. C OMPARISON WITH M ODEL S ELECTION M ETHOD
HD-CNN (Yan et al., 2015) and HydraNets (Mullapudi et al., 2018)
select branches based on the category.
MoE (Shazeer et al., 2017) and Switch Transformer (Fedus et al., 2021) enable the direct training of Router by scaling the output feature of experts with gate-values predicted from Router.
Despite these methods are originally designed to conditionally select a speciﬁc layer or block, we apply them to the expert selection.
To select the expert based on category, the categories should be partitioned into n groups, where n is the number of experts.
We try two schemes: random partition (RP) or clustering-based partition (CBP).
Then, an expert can be selected according to the rough prediction of delegator.
During the training procedure, we also reweight losses of each expert based on the assignment matrix A m × n with Eq.5 and 6.
A m × n is obtained directly based on the rough prediction.
We can enable the direct training of expert selector as well.
Directly trained selector (DTS) is obtained by scaling the loss of each expert, rather than the output feature of each branch as the original paper (Shazeer et al., 2017;
Fedus et al., 2021).
While CoE optimizes the expert selector via weighted cross-entropy loss L s (Eq.9).
The results with 4 experts are shown in Table 3.
It can be seen that CoE outperforms the compared methods, demonstrating a better collaboration pattern is learned.
4. A NALYSIS OF THE G ENERALIZATION
To verify the generalizability, we conduct two extra experiments: generalizing CoE to translation task and using Reassessed Labels (ReaL) (Beyer et al., 2020) to re-evaluate CoE .
We mainly introduce the ﬁrst one here, another one are discussed in Appendix B.1. methods.
“RP”,“CBP” and “DTS” means “Random Partition”,“Clustering-Based Partition” and “Directly Trained Selector”
[Table] [Table] To generalize CoE to translation task, we build a CoE-Transformer model based on Transformer (base model) (Vaswani et al., 2017).
Considering the decoding procedure is much more time-consuming than encoding because target tokens are generated one by one, only the decoder of CoE-Transformer is conditionally selected.
CoE-Transformer has four decoders, given a sentence, one decoder will be selected to decode the features output by encoder.
To select the decoder, an extra constant token is added at the beginning of each sentence, the feature of this token output by encoder is input to the Expert Selector (Fig.2) for expert (decoder) selection.
During training, the TCP of a sentence is obtained by averaging the TCPs of each token.
Following (Vaswani et al., 2017; Ott et al., 2019), CoE-Transformer is trained on the standard WMT 2014 English-German dataset.
As mentioned, an extra token will be added into this vocabulary.
We adopt the same training and evaluating setting as (Ott et al., 2019), more details are shown in Appendix B.2.
Although CoE-Transformer has a larger number of parameters compared with Transformer (base model), the MAC is nearly identical.
We calculate MAC based on the number of parameters loaded given a sentence and report it on Table 5 as well.
From Table 5 we can see, CoE-Transformer outperforms the regular Transformer (base model) by a large margin and achieves similar performance with Transformer (big). [Table]
[Table]
A NALYSIS OF THE L EARNED E XPERT S ELECTION P ATTERNS
We also conduct experiments to analyze the expert selection patterns of CoE and ﬁnd them quite reasonable.
When experts have different architectures, the delegator learns to assign easy samples to smaller experts and complex samples to heavier experts.
When experts share the same architecture, delegator learns the expert selection patterns automatically, it can be based on any property (e.g. whether humans are contained), rather than limited to the category.
The details are shown in Appendix B.3.
C ONCLUSION
We propose a CoE framework to pool together the expertise of multiple networks towards a common aim.
Experiments in this paper demonstrate the superiority of CoE on both accuracy and real speed.
We also analyze the collaboration patterns and ﬁnd them has interpretability.
In the future, CoE will be extended to the trillion parameters level.
Meanwhile, we will try to implement CoE to more tasks and verify its compatibility with quantiﬁcation and other technologies.
Besides, CoE can be conducted to solve the problems of lifelong learning by updating experts.
A E XTRA D ETAILS FOR M ETHOD A. I NTRODUCTION OF V OGEL A PPROXIMATION M ETHOD (VAM)
In weight generation module (WGM) and label generation module (LGM), we need to solve the balanced transportation problem (BTP, (Shore, 1970)) via Vogel approximation method (VAM, (Shore, 1970)).
We will introduce it in this section with the number of samples and experts as m and n respectively.
The BTP involved in WGM and LGM has m supply sources, each of which is denoted as Silo j with a supply of one, as well as n demand sources, each of which is denoted as Mill k with a demand of mn .
C j,k is the per-unit transportation cost from Silo j to Mill k .
Speciﬁcally, C j,k = − P j,k in
WGM and C j,k = − S j,k in LGM.
To make it clear, we illustrate this algorithm with a toy example, where the problem is simpliﬁed as Fig.5 (a) with m = 4 , n = 2 .
In the ﬁrst step, we calculate the penalty cost pc row j for each row and pc col k for each column of the tableau in Fig.5 (a).
Penalty cost is determined by subtracting the lowest unit cost in the row (column) from the next lowest unit cost.
The penalty costs of the respective rows and columns have been marked in red color for clarity in
Fig.5 (b).
Since the third row has the largest penalty cost ( pc row 3 =11) and C 3 , 1 is the lowest unit cost of that row, Silo 3 is allocated to Mill 1 , i.e. A 3 , : =
[1 , 0] in WGM or L 3 , : =
[1 , 0] in LGM.
Then the corresponding row should be crossed out and the demand of Mill 1 should minus one, if this results in a zero demand, the ﬁrst column will be crossed out as well.
After adjusting penalty cost for each row and column, the tableau becomes Fig.5 (c), where the changed values are marked in orange.
The described procedure will be looped until no rows remained.
Considering the calculation of pc col k is much more time-consuming compared with pc row j because m (cid:29) n in WGM and LGM, we modify VAM by only seeking lowest penalty cost among { pc row 1 , ..., pc row m } .
We ﬁnd this modiﬁcation makes VAM more efﬁcient while keeps the superiority of the solution.
It is because the mechanics of VAM makes it meaningful to take pc col k into account only when the demand of Mill k is one, which rarely happens.
Thus, we adopt this modiﬁcation to promote efﬁciency in this paper.
C 1,1 =
C 2,
[Table]
=
C 3,
[Table]
= C 4,
[Table]
= C [Table]
[Table]
=
C [Table] [Table] =
C
[Table]
[Table]
= C [Table] [Table]
Mill
[Table]
𝑝𝑝𝑐𝑐
[Table]
𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆 1 :𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠=
[Table]
𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆
[Table]
:𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠
[Table]
𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆
[Table]
:𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠=
[Table]
𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆
[Table]
:𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠=
[Table] Mill 1 : [Table] =
[Table] Mill [Table] [Table]
C 1,1 =
[Table]
C 2,1 =
C 3,1 = C 4,1 = C 1,2 = C 2,2 =
C 3,2 =
C 4,2 =𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆 1 :𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠= 1 𝑝𝑝𝑐𝑐 𝑟𝑟𝑐𝑐𝑟𝑟1 =
𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆 2 :𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠=1𝑝𝑝𝑐𝑐 𝑟𝑟𝑐𝑐𝑟𝑟2 =𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆 3 :𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠=1𝑝𝑝𝑐𝑐 𝑟𝑟𝑐𝑐𝑟𝑟3 =𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆 4 :𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠 =1 𝑝𝑝𝑐𝑐 𝑟𝑟𝑐𝑐𝑟𝑟4 =
Mill 1 : demand= [Table] 𝑝𝑝𝑐𝑐 𝑐𝑐𝑐𝑐𝑐𝑐1 =
Mill 2 : demand= 𝑝𝑝𝑐𝑐 𝑐𝑐𝑐𝑐𝑐𝑐2 = 𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆 1 :𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠 =1 𝑝𝑝𝑐𝑐 𝑟𝑟𝑐𝑐𝑟𝑟1 =𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆 2 :𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠=1𝑝𝑝𝑐𝑐 𝑟𝑟𝑐𝑐𝑟𝑟2 =𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆 4 :𝑠𝑠𝑠𝑠𝑝𝑝𝑝𝑝𝑆𝑆𝑠𝑠 =1 𝑝𝑝𝑐𝑐 𝑟𝑟𝑐𝑐𝑟𝑟4 =
Mill 2 : demand= 𝑝𝑝𝑐𝑐 𝑐𝑐𝑐𝑐𝑐𝑐2 =
C 1,1 =
C 2,1 =
C 3,1 =
C 4,1 =
C 1,2 = C 2,2 = C
[Table] = C
[Table] [Table] =
[Figure]
A. A S TRATEGY TO F ACE THE L ARGE N UMBER OF E XPERTS
Each of our experiments involves either four or sixteen experts.
For the four-experts setting, the training method is the same as described.
However, we will meet the slow-convergence problem when the number of experts is sixteen.
It motivates us to decompose the task into four subtasks, each of which involves four experts and can be trained with the proposed training method.
With this strategy, the number of samples assigned to each expert increases from m 16 to m 4 .
Because these samples contribute most to the optimization of the expert, the rate of convergence becomes nearly four times faster.
To fulﬁll task decomposition, we introduce a new module to delegator, named subtask selector as shown in Fig.6.
The subtask selector is used to allocate the input samples into different subtasks.
The expert selector outputs sixteen probabilities, which are partitioned into four groups as well.
For each subtask, only one group of probabilities is visible.
The experts within each subtask and the corresponding weights of the expert selector are jointly optimized.
As for the feature extractor, task predictor, and subtask selector, their weights directly derive from the delegator trained with the setting of four experts and then ﬁxed.
During this procedure, the weights of subtask selector derive from the expert selector.
[Figure]
Selector
[Figure]
Subtask…
[Figure]
Expert 13−
B E XTRA DETAILS FOR E XPERIMENTS B. R E - EVALUATION WITH R EASSESSED L ABELS
As described in paper (Beyer et al., 2020), the validation set labels have a set of deﬁciencies that makes the recent progress on ImageNet classiﬁcation benchmark suffer from overﬁtting to the artifacts.
To verify the generalization, we use the Reassessed Labels (ReaL) (Beyer et al., 2020) to re-evaluate our method.
The results are shown in Table 6.
It can be seen that our method still has a remarkable performance, achieving higher accuracy than the compared methods with signiﬁcantly smaller FLOPs.
[Table]
B. E XPERIMENT DETAILS FOR THE TRANSLATION TASK
Following (Vaswani et al., 2017; Ott et al., 2019), CoE-Transformer is trained on the standard WMT 2014 English-German dataset, which has a shared source-target vocabulary of about 37000 tokens.
As mentioned, an extra token will be added into this vocabulary.
The default training setting is identical with the one described in (Vaswani et al., 2017), except for the batch size and learning rate become larger following (Ott et al., 2019).
Moreover, the parameter α in Eq.5 grows linearly from 0.1 to 0. with the training going on.
We report BLEU on news2014 with a beam width of 4 and length penalty of 0.6 based on a single model obtained by averaging the last 5 checkpoints following (Vaswani et al., 2017; Ott et al., 2019).
B. A NALYSIS OF THE L EARNED E XPERT S ELECTION P ATTERNS B.3. E XPERT S ELECTION P ATTERNS WHEN EXPERTS HAVE VARIOUS ARCHITECTURES
Considering TCP can measure the complexity of a given sample if the model for inference is ﬁxed, we experiment to analyze the relationship between sample complexity and expert selection.
If the experts have different computation cost, it is reasonable to assign easy samples (the ones with large TCP) to smaller experts and complex samples to heavier experts.
To verify this, we take four architectures searched via OFA (Cai et al., 2020) as the experts, i.e. OFA-110, OFA-163, OFA-230 and OFA-595.
The delegator is also MobileNetV3-small as described in section 4.1.
We obtain the TCP value for each sample based on the delegator.
We count the selection probability for each expert at different TCP values on the validation set.
As shown in Fig.7, it meets our expectation that the selection probability for smaller models increases the input sample getting simpler (with the increase of TCP).
[Figure]
S e l e c t i o n
P r o b a b ili t y
True Class Probability (TCP)
OFAOFA- OFAOFA-
E XPERT S ELECTION P ATTERNS WHEN EXPERTS SHARE THE SAME ARCHITECTURE
We have analyzed the selection patterns when the experts have different architectures, here we focus on the case that all experts share the same architecture.
We adopt the CoE-Large setting with four experts.
Considering many works (Yan et al., 2015; Mullapudi et al., 2018) select branches based on the category, we ﬁrstly experiment to observe the relationship between selection patterns and rough prediction of the delegator on ImageNet validation set.
Based on the predicted class of rough prediction, the validation set can be partitioned into 1000 subsets. Then we calculate the probabilities to select each expert within each subset and get 1000 probability vectors.
After clustering, we plot the probability vectors on Fig.8, each column of which is a probability vector.
It can be seen that samples with the same rough prediction class are assigned to different experts.
Therefore, we can conclude the expert is not always selected based on category.
Besides, we further make qualitative analysis on the ImageNet validation set and ﬁnd some interesting patterns.
For example, we ﬁnd that images predicted as “meat market” are most likely to be assigned to the fourth expert if humans are contained.
We show those images in Fig.9.
It can be seen, images are assigned to the fourth expert, among which 22 images contain humans with a ratio of 81.5%.
By contrast, among the 32 images that are assigned to the other experts, only 7 images contain humans with a ratio of 21.9%.
This indicates CoE learns the expert selection patterns automatically, it can be based on properties other than the category.
[Figure]
n
P
[Figure]
ili t y [Figure] [Figure]
We analyze the number of experts in this section, including 1, 4, and 16 experts. The results are shown in Table 7.
Using one expert brings little improvement compared with the original model. When increasing the number of experts, the accuracy becomes 1.9% better with four experts and 2.9% better with sixteen experts. It demonstrates CoE can make full use of multiple experts, leading to a large collaborative beneﬁt. What‘s more, the accuracy also reaches 79.9% by combining CondConv with
OFA-230.
In this manner, CoE can further enhance the accuracy to 80.8/81.5% with 4/16 experts.
[Table]
E FFECT OF E ARLY T ERMINATION
The original OFA-230 has 78.0% top-1 accuracy with 230M FLOPs.
We can introduce a
MobileNetV3-Small to conduct early termination.
By varying the threshold, we get a series of accuracies and FLOPs as shown in Fig.10.
It can seen that the accuracy becomes 78.0% with 220M
FLOPs.
This indicates the computation cost brought by MobileNetV3-Small is eliminated via early termination strategy.
Inspired by this, we expect to eliminate the computation cost brought by delegator as well.
It does reduce the computation cost by 60/66M FLOPs, demonstrating the effectiveness of early termination
[Table]
[Figure]
[Table] [Figure] [Figure] [Figure]
B.4. A BLATION STUDY FOR EACH COMPONENT OF C O E
CoE consists of 3 major components: WGM, LGM and SRM.
Besides directly removing one component, we also try to alter some elements inside them.
We propose several modiﬁed versions of
CoE for ablation as below:
• CoE WGM :
Remove WGM from CoE.
Thus, losses of experts have identical weights for each sample.
• CoE WGM (cid:63) :
WGM partitions the training data based on expert suitablity, i.e. the assignment matrix
A m × n in WGM equals to the output matrix L m × n of LGM.
• CoE WGM ◦ : Remove the “ (cid:80)
j A j,k = m/n ” constraint in Eq.4, so that A m × n neglects the No Superiority Assumption (NSA).
Then, take A m × n as the output of WGM without the smoothing (Eq.5) and normalizing(Eq.6).
• CoE WGM • : Remove the progressive sharpening of assignment in WGM.
Speciﬁcally, using a constant 0.8 for α in Eq.5, instead of linearly increasing it from 0.2 to 0.8.
• CoE LGM : Remove LGM from CoE. Thus, CoE collapse to a single expert with delegator to trigger the early termination.
• CoE LGM (cid:63) : Abandon the reﬁning of suitability criterion (Eq.2) and remove the “ (cid:80) j L j,k = m/n ”constraint in Eq.3.
So that L m × n neglects the No Superiority Assumption (NSA).
• CoE SRM : Remove SRM from CoE.
We conduct experiments for those CoE versions with the CoE-Large setting and 4 experts.
Results are shown in Table 8, and the conclusions are listed below:
WGM and LGM are the most important components in CoE. The removal of WGM and LGM reduce the accuracy from 79.9% to 78.1% and 78.0%, respectively.
In WGM, the training data should be partitioned based on delegator.
Otherwise, delegator will overﬁt to irregular selection labels as illustrated in section 3.3.
That is why CoeLarge WGM (cid:63) only achieves an accuracy of 78.4%.
The No Superiority Assumption (NSA) is importart for CoE. Without this assumption, CoE-Large WGM ◦ and CoE-Large LGM (cid:63) only reaches an accuracy of 79.2% and 79.4%.
The progressive sharpening of assignment in WGM can also boost the performance, thus the accuracy for CoE-Large WGM • is 0.5% lower than CoE-Large. 5.
The component SRM is also useful.
It promotes delegator to select experts better, yielding a 0.4% improvement for accuracy.
[Table]
B.4. A BLATION STUDY FOR ELEMENTS OF THE TRAINING STRATEGY
Knowledge distillation (KD), auto-augment (AA) and stochastic depth (SD) are widely-used strategies to overcome the overﬁtting problem.
We think only when the overﬁtting problem is solved can task accuracy reﬂect model capacity exactly.
Because this paper is concerned with improving model capacity with limited computation cost, we use these strategies.
Nonetheless, we conduct ablations for them in this section.
We adopt the CoE-Large setting and use 4 experts.
Results are shown in
Table 9.
We ﬁnd KD extremely important for CoE, it may indicate CoE is easy to be overﬁtted.
In addition, SD decreases the accuracy of CoE. By removing SD, CoE-Large (4 experts) boosts the accuracy from 79.9% to 80.2%.
Perhaps, it is because SD makes the capacity of delegator and each expert too tiny (Gontijo-Lopes et al., 2021).
[Table]