A uto M o E : Heterogeneous Mixture-of-Experts with Adaptive Computation for Efficient Neural Machine Translation (2023)

Abstract: Mixture-of-Expert (MoE) models have obtained state-of-the-art performance in Neural Machine Translation (NMT) tasks. Existing works in MoE mostly consider a homogeneous design where the same number of experts of the same size are placed uniformly throughout the network. Furthermore, existing MoE works do not consider computational constraints (e.g., FLOPs, latency) to guide their design. To this end, we develop AutoMoE – a framework for designing heterogeneous MoE’s under computational constraints. AutoMoE leverages Neural Architecture Search (NAS) to obtain efficient sparse MoE sub-transformers with 4x inference speedup (CPU) and FLOPs reduction over manually designed Transformers, with parity in BLEU score over dense Transformer and within 1 BLEU point of MoE SwitchTransformer, on aggregate over benchmark datasets for NMT.Heterogeneous search space with dense and sparsely activated Transformer modules (e.g., how many experts? where to place them? what should be their sizes?) allows for adaptive compute – where different amounts of computations are used for different tokens in the input. Adaptivity comes naturally from routing decisions which send tokens to experts of different sizes. AutoMoE code, data, and trained models are available at https://aka.ms/AutoMoE .

## introduction sparsely activated models like the mixture-of-experts (moe) @xcite perform conditional computation in which only a subset of the weights of the network are activated per input. selective compute allows us to design neural networks with a large number of model parameters, without significant increase in the computational cost. with increased capacity, these sparse models have demonstrated state-of-the-art performance in natural language tasks such as neural machine * correspondence to {ganeshjwhr@gmail.com, subhabrata.mukherjee@microsoft.com}. translation (nmt) @xcite @xcite . moe architectures require several design choices: (a) expert placement: identifying transformer layers for introducing expert sub-networks. (b) number of experts: how many experts to place in different layers? (c) expert ffn size: what should be the feedforward network (ffn) size for each expert? given the large search space of potential architectures and the exorbitant computational cost of training and evaluating them, existing approaches manually design moe architectures from a highly-restricted homogeneous space. for instance, they use the same number of experts of the same capacity in different layers and make ad-hoc decisions like introducing experts in every other layer @xcite @xcite @xcite or every four layers @xcite . while these moe's support conditional computation, homogeneity (specifically, fixed-size experts) results in the same amount (albeit different subsets) of weights to be applied to each input. we hypothesize that this is not an optimal solution and that we can reduce the number of experts (in some layers) to reduce communication cost, and the size (of some experts) to reduce computation cost resulting in reduction in model size, flops and latency without much quality degradation. this naturally extends moes to be adaptive compute models (similar to work on early exit @xcite ) where different amounts of computations are used for different inputs. the adaptivity comes naturally from the routing decisions which would send tokens to experts of different sizes. the above observations are depicted in ## designing heterogeneous mixture-of-experts we now present the components of automoe framework (illustrated in figure 1 ) for designing efficient moe's under computational constraints. ## conclusion automoe is the first framework to design heterogeneous moe's under computational constraints. it supports adaptive computation i.e. variable compute for different inputs with variable-size experts. it leverages nas to explore a heterogeneous search space with variable number of experts, sizes, and placement choices; alongside other standard transformer architectural hyper-parameters. automoe generated moe subnetworks reduce flops and latency over both manually designed and nassearched architectures on benchmark mt tasks.