Sparse MoEs meet Efficient Ensembles

1 INTRODUCTION . Neural networks typically use all their parameters to process an input . Sustaining the growth of such models—reaching today up to 100B parameters ( Brown et al. , 2020 ) —is challenging , e.g. , due to their high computational and environmental costs ( Strubell et al. , 2019 ; Patterson et al. , 2021 ) . In this context , sparse mixtures of experts ( sparse MoEs ) employ conditional computation ( Bengio et al. , 2013 ) to combine multiple submodels ( experts ) and route examples to certain experts ( Shazeer et al. , 2017 ; Lepikhin et al. , 2021 ; Fedus et al. , 2021 ; Riquelme et al. , 2021 ; Yang et al. , 2021 ) . Conditional computation can decouple the growth of the number of parameters from the training and inference costs , by only activating a subset of the overall model in an input-dependent fashion . Paralleling this trend , the deployment of ML systems in safety-critical fields , e.g. , medical diagnosis ( Dusenberry et al. , 2020b ) and self-driving cars ( Levinson et al. , 2011 ) , has motivated the development of reliable deep learning , e.g. , for calibrated and robust predictions ( Ovadia et al. , 2019 ) . Among the approaches , ensembles of neural networks have remarkable performance for calibration and accuracy under dataset shifts ( Ovadia et al. , 2019 ) . These methods improve reliability by aggregating the predictions of individual submodels ( ensemble members ) . While sharing conceptual similarities , these two classes of models—MoEs and ensembles—have different properties . Sparse MoEs adaptively combine their experts depending on the inputs , and the combination generally happens at internal activation levels . Ensembles typically combine several models in a static way and at the prediction level . Moreover , these two classes of models tend to be benchmarked on different tasks : few-shot classification for MoEs ( Riquelme et al. , 2021 ) and uncertainty-related evaluation for ensembles ( Ovadia et al. , 2019 ; Gustafsson et al. , 2020 ) . CONTRIBUTIONS : In this paper , we study the interplay between sparse MoEs and ensembles . This results in two sets of contributions . Contribution 1 : Complementarity of MoEs and ensembles . We show that sparse MoEs and ensembles have complementary features and benefit from each other . Specifically : • The adaptive computation in sparse MoEs and the static combination in ensembles are orthogonal , with additive benefits when associated together . Their association results in insightful performance versus FLOPs trade-offs while varying the ensemble size and sparsity . • In sparse MoEs , combining models at the prediction level leads to improved uncertainty estimates . • Over tasks where either sparse MoEs or ensembles are known to perform well , naive—and computationally expensive—ensembles of MoEs provide the best predictive performance . Our benchmarking effort includes the first evaluation of sparse MoEs on uncertainty-related vision tasks , which builds upon the empirical work of Riquelme et al . ( 2021 ) . 1 Under review as a conference paper at ICLR 2022 Contribution 2 : Partitioned batch ensembles . We propose partitioned batch ensembles ( pBE ) , see Figure 1 , an efficient ensemble approach tailored to sparse MoEs . Specifically : • pBE improves over sparse MoEs across metrics including few-shot performance , likelihood and calibration error . pBE matches the performance of deep ensembles for 30 % -43 % fewer FLOPs . • pBE gracefully scales up to vision Transformers with up to 2.7B parameters . • pBE is both simple ( requiring only minor implementation changes ) and convenient because standard sparse-MoE checkpoints can be used directly to initialize pBEs for fine-tuning . 2 PRELIMINARIES . We focus on classification tasks where we learn classifiers of the form f ( x ; θ ) based on some training data D = { ( xn , yn ) } Nn=1 . A pair ( xn , yn ) corresponds to an input xn ∈ RP together with its label yn ∈ { 1 , . . . , C } belonging to one of the C classes . The model f ( · ; θ ) is parametrized by θ and outputs a C-dimensional probability vector . We use ◦ to refer to matrix element-wise product . 2.1 VISION TRANSFORMERS AND SPARSE MOES . Vision Transformers . Throughout the paper , we choose the model f to be a vision Transformer ( ViT ) ( Dosovitskiy et al. , 2021 ) . ViT is growing in popularity for vision , especially in transferlearning settings where it was shown to outperform convolutional networks while requiring fewer pre-training resources . ViT operates at the level of patches . An input image is split into equal-sized patches ( e.g. , 32× 32 , 16× 16 , or 14× 14 pixels ) whose resulting sequence is ( linearly ) embedded and processed by a Transformer ( Vaswani et al. , 2017 ) . The operations in the Transformer then mostly consist of a succession of multiheaded self-attention ( MSA ) and MLP layers . ViT is defined at different scales ( Dosovitskiy et al. , 2021 ) : S ( mall ) , B ( ase ) , L ( arge ) and H ( uge ) ; see specifications in Appendix A . For example , ViT-L/16 stands for a large ViT with patch size 16× 16 . Sparse MoEs and V-MoEs . The main feature of sparsely-gated mixture-of-experts models ( sparse MoEs ) lies in the joint use of sparsity and conditional computation ( Bengio et al. , 2013 ) . In those models , we only activate a small subset of the network parameters for a given input , which allows the total number of parameters θ to grow while keeping the overall computational cost constant . The subparts of the network that are activated on a per-input fashion are known as experts . Central to our study , Riquelme et al . ( 2021 ) recently extended ViT to sparse MoEs . Their extension , referred to as V-MoE , follows the successful applications of sparse models in NLP ( Shazeer et al. , 2 Under review as a conference paper at ICLR 2022 2017 ) . Riquelme et al . ( 2021 ) show that V-MoEs dominate their “ dense ” ViT counterparts on a variety of tasks for the same computational cost . In the specific case of V-MoEs , the experts are placed in the MLP layers of the Transformer , a design choice reminiscent of Lepikhin et al . ( 2021 ) in NLP . Given the input h ∈ RD of such a layer , the output of a single MLP ( h ) is replaced by MoE ( h ) = E∑ e=1 ge ( h ) · MLPe ( h ) with { ge ( h ) } Ee=1 = topK ( softmax ( Wh ) ) , ( 1 ) where the routing weights { ge ( h ) } Ee=1 combine the outputs of the E different experts { MLPe } Ee=1 . To sparsely select the experts , topK sets all but the K largest weights to zero . The router parameters W ∈ RE×D are trained together with the rest of the network parameters . We call the layer defined by ( 1 ) an MoE layer . In practice , the weights { ge ( h ) } Ee=1 are obtained by a noisy version of the routing function topK ( softmax ( Wh+ σε ) ) with ε ∼ N ( 0 , I ) , which mitigates the nondifferentiability of topK when combined with auxiliary losses ( Shazeer et al. , 2017 ) . We use the shorthand gateK ( z ) = topK ( softmax ( z + σε ) ) and take σ = 1/E ( Riquelme et al. , 2021 ) . In this paper , we take the “ last-n ” setting of Riquelme et al . ( 2021 ) wherein only a few MoE layers are placed at the end of the Transformer ( n = 2 for the { S , B , L } scale and n = 5 for H ) . This setting retains most of the performance gains of V-MoEs while greatly reducing the training cost . 2.2 ENSEMBLES OF NEURAL NETWORKS . Ensembles . We build on the idea of ensembles , which is a known scheme to improve the performance of individual models ( Hansen & Salamon , 1990 ; Geman et al. , 1992 ; Krogh & Vedelsby , 1995 ; Opitz & Maclin , 1999 ; Dietterich , 2000 ; Lakshminarayanan et al. , 2017 ) . Formally , we assume a set of M model parameters Θ = { θm } Mm=1 . We refer to M as the ensemble size . Prediction proceeds by computing 1M ∑ θ∈Θ f ( x ; θ ) , i.e. , the average probability vector over the M models . To assess the diversity of the predictions in the ensemble , we will use the KL divergence DKL ( f ( xt ; θm ) ‖f ( xt ; θm′ ) ) between the predictive distributions f ( xt ; θm ) and f ( xt ; θm′ ) , averaged over the test input xt and all pairs ( m , m′ ) of ensemble members . Batch ensembles . Ensembles differ in the way Θ is defined . Central to our study , batch ensembles ( BE ) ( Wen et al. , 2019 ) build the ensemble as a collection of submodels , with the parameters θm ∈ Θ sharing components . This mitigates the computational and memory cost of ensembling , enabling one to improve the performance of the original model at little extra cost . We focus on the example of a single dense layer in f with parameters U ∈ RD×L , assuming no bias . BE defines M copies of parameters { Um } Mm=1 so that Um = U ◦ ( rms > m ) , where U are parameters shared across ensemble members , and rm and sm are separate D- and L-dimensional vectors for ensemble member m. Given an input , the BE produces M outputs , and the M outputs are averaged after applying all layers . Despite the simple rank-1 parametrization , BE leads to remarkable predictive performance and robustness ( Wen et al. , 2019 ) . Notably , the efficiency of BE relies on tiling the inputs to simultaneously predict with the M ensemble members , an insight that we also exploit . 2.3 UPSTREAM PRE-TRAINING AND DOWNSTREAM FINE-TUNING . Large-scale Transformers pre-trained on upstream tasks were shown to have strong performance when fine-tuned on smaller downstream tasks , across a variety of domains ( Devlin et al. , 2018 ; Dosovitskiy et al. , 2021 ; Radford et al. , 2021 ) . We follow this paradigm and focus on the finetuning of models pre-trained on JFT-300M ( Sun et al. , 2017 ) , similar to Riquelme et al . ( 2021 ) . We will thus assume the availability of already pre-trained ViT and V-MoE model checkpoints . Our assumption relies on the growing popularity of transfer learning , e.g . Kolesnikov et al . ( 2020 ) , and the increasing accessibility of pre-trained models in repositories such as www.tensorflow . org/hub or www.pytorch.org/hub . The fine-tuning of all the approaches we study here , including extensions of ViT and V-MoE , will be either directly compatible with those checkpoints or require only mild adjustments , e.g. , reshaping or introducing new downstream-specific parameters ( see Appendix C ) . Also , unless otherwise mentioned , the performance we report will always be downstream , e.g. , for ImageNet ( Deng et al. , 2009 ) or Cifar10/100 ( Krizhevsky , 2009 ) . In all our comparisons , we will use the downstream training floating point operations per second ( FLOPs ) , or GFLOPs ( i.e. , 109×FLOPs ) , to quantify the computational cost of the different methods . 3 Under review as a conference paper at ICLR 2022

Summary: The authors show empirically that Sparse MOEs and Ensembles have complementary features, and suggest that combining the two should lead to improved performance. Authors build on the Vision Transformer (ViT) for their experiments. To efficiently combine Sparse MOEs and Ensembles, the paper presents Partitioned Batch Ensembles (PBE), where the parameters of the self-attention layers are shared, and an ensemble of Sparse MOEs are used for the MLP layers of the Transformer blocks.