Collaboration of Experts: Achieving 80% Top-1 Accuracy on ImageNet with 100M FLOPs

1 INTRODUCTION . 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 significant 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 refined 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 significant 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 first , since networks learn gradually more complex hypotheses during 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 overfitting 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 ImageNet demonstrate the superiority of CoE . It achieves 78.2/80.7 % top-1 accuracy 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 first 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 % top-1 accuracy with only 100M FLOPs for the first time . 2 RELATED WORK . 2.1 ENSEMBLE LEARNING AND MODEL SELECTION . 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 significant 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 significant runtime cost . Differently , CoE utilizes a delegator to select only one expert for the refined 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 can not 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 finding 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 can not be adopted in CoE . 2.2 DYNAMIC NETWORKS . 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 specific 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 can not 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 . 3 METHOD 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 classification 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.1 INFERENCE PROCEDUREF . CoE firstly 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 final 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 refined prediction . The averaged FLOPs/Instance of CoE ranges from FD to FD + FE by varying τ , FD and FE are FLOPs of delegator and experts . 3.2 LABEL GENERATION MODULE ( 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 ) : TCPj , k = P ( Y = yj |xj , Expertk ) , ( 1 ) where , xj is the j-th sample , Y and yj 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 : Sj , k = TCPj , k −Mean ( TCP : ,k ) Std ( TCP : ,k ) , ( 2 ) where , Mean ( TCP : ,k ) and Std ( TCP : ,k ) are mean value and standard deviation for TCPs of Expertk . Given m samples and n experts , selection labels can be denoted by a binary matrix Lm×n , thus each row of which is a selection label . According to NSA , the sum of each colum vector of Lm×n should be same , namely ∑ j Lj , k = m n for k = 1 , ... , n. Therefore , Lm×n can be obtatined by maximizing the sum of standardized TCP ( i.e . Sj , k in Eq.2 ) : min ∑ j , k −Sj , k ∗ Lj , k s.t . Lj , k ∈ { 0 , 1 } , ∑ k Lj , k = 1 , ∑ j Lj , k = m n ( 3 ) 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 . −Sj , 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 .

Summary: The paper proposes a model for classification (on ImageNet) that involves having a very small network act as a selector and early predictor, and a set of experts that specialise on a subset of the training data. The selector chooses one of them and thus keeps a constant cost irrespective of the number of experts. This is then combined with a number of existing techniques to raise the overall performance. Namely, the authors use a PWRL non-linearity instead of ReLU, they use logit-matching knowledge distillation, and conditional convolution (CondConv). The outcome is a tiny model that yields excellent performance on ImageNet, with performance envelopes superior to existing methods.