Towards Being Parameter-Efficient: A Stratified Sparsely Activated Transformer with Dynamic Capacity (2023)

Abstract: Mixture-of-experts (MoE) models that employ sparse activation have demonstrated effectiveness in significantly increasing the number of parameters while maintaining low computational requirements per token. However, recent studies have established that MoE models are inherently parameter-inefficient as the improvement in performance diminishes with an increasing number of experts. We hypothesize this parameter inefficiency is a result of all experts having equal capacity, which may not adequately meet the varying complexity requirements of different tokens or tasks. In light of this, we propose Stratified Mixture of Experts (SMoE) models, which feature a stratified structure and can assign dynamic capacity to different tokens. We demonstrate the effectiveness of SMoE on three multilingual machine translation benchmarks, containing 4, 15, and 94 language pairs, respectively. We show that SMoE outperforms multiple state-of-the-art MoE models with the same or fewer parameters.

## introduction scaling up the model and data size has shown tremendous success in enhancing model performance across a large number of nlp tasks @xcite @xcite . sparsely gated mixture of experts (moe) @xcite provides an effective way to greatly scale the model size under the same computational cost and achieves state-ofthe-art performances on various tasks including natural language understanding @xcite , machine translation (nllb @xcite , language modeling @xcite , etc. the efficiency comes from sparsely activating a subset of the neural network weights for each incoming sample. however, moe is reported to be parameter-inefficient @xcite @xcite i.e., there are diminishing improvement returns from adding more experts. for example, switch transformer @xcite only outperforms t5 @xcite by an average of 0.7 on the glue benchmark @xcite despite being 35× larger. similarly, in the translation task, a moe model with 20 times more parameters only offers an average improvement of 0.3 bleu on its ablation dataset (moe-64 vs. 1.3b dense) @xcite . we hypothesize that this parameter inefficiency stems from the equal capacity assignment, where we particularly define 'capacity' as the number of parameters used for the incoming token. for the current moe models, the capacity of experts are the same used for serving all tokens. however, different tokens may demand varying capacities. for instance, in the context of multilingual machine translation, certain translation directions may necessitate a greater capacity to prevent overfitting, while others only require a smaller capacity. to address this limitation, our hypothesis posits that the dynamic allocation of capacity to tokens results in more efficient utilization of parameters. thus, we propose stratified mixture of experts (smoe) models, characterized by a stratified structure, which allows for the dynamic assignment of capacity to incoming tokens. a high-level comparison of vanilla moe and smoe is presented in figures 1a and 1b . in vanilla moe, a single routing gate connects to all e experts and sends tokens to the top-k experts. here, we take @xmath0. in smoe, the experts are divided into two strata. each stratum has its own routing gate that connects to all experts in the current stratum as well as all experts in each stratum has a gate that is connected to all subsequent experts. tokens can be directly sent to the last stratum to only experience one expert, or be sent to both strata and have more capacity. hence, the dynamic capacity of a token depends on how many experts it needs to pass. c) a detailed architectural design, where a comprehensive explanation of the design components will be presented in section 3. the subsequent strata. if gate 1 assigns tokens to expert 4 or 5, the tokens will only need to pass through a single expert (an ffn layer). however, if tokens are sent to experts in the first stratum (experts 1 to 3), they will need to go through the next stratum as well, meaning that another expert will be assigned by gate 2 before exiting the smoe block. this allows smoe to dynamically assign capacity to different tokens. in addition, a comprehensive illustration of the architectural design of the smoe model is provided in figure 1c . a thorough explanation of the design elements will be provided in section 3. our main contributions are summarized as follows: ## background and related work massively multilingual machine translation models have been developed to handle several translation directions simultaneously in a single model @xcite . however, the use of shared parameters for different languages often leads to negative transfer and decreased performance @xcite . in contrast to dense mmt models, sparsely gated mixtureof-experts (moe) models, which activate a subset of parameters for each input, have been shown to significantly improve translation performance @xcite . @xcite first demonstrated the benefit of adding moe layers to scale rnn models for improved translation performance, and @xcite extended this work to transformer architectures @xcite . moe layers in the transformer model replace a single feedforward network (ffn) layer with e ffn layers, denoted with {ffn 1 , . . . , ffn e }. each ffn layer is an expert. given an input token x, we have ∀@xmath1, . . . , e}, ffn e (x) = w (e) 1 relu(w (e) 2 • x), (1) where w (e) 1 and w (e) 2 are the weights of ffn e . a trainable routing gate with weights w g predicts scores for these experts to use for the input x, in the form of a routing vector @xmath2: we select the set of top-k experts, denoted with e ⊂ {1, • • • , e}, and compute the output of the moe layer as follows: moe models suffer from the notorious load imbalance issue, where the gate weights could collapse and send most tokens to the same expert. as a result, recent research has focused on designing better auxiliary load balancing loss functions to encourage tokens to be evenly distributed across experts, e.g., @xcite formulated token-to-expert allocation as a linear assignment problem, @xcite modified the feedforward layer to hash to different sets of weights depending on the current token, @xcite proposed a router z-loss that resolves instability issues, and @xcite reversely design an expert-to-token allocation algorithm. other lines of investigation in moe include regularization techniques such as gating dropout @xcite and output masking (eom and fom) @xcite , as well as novel moe architectures, such as conditional moe routing (cmr) that add an extra branch beside moe layer @xcite , or pyramid-residual moe @xcite , a hybrid dense and moe model with more experts in the last layers. however, all previous work default to equal capacity for all tokens regardless of their language, frequency, or any other property. in the subsequent sections, we present a stratified mixture of experts (smoe) model that automatically assigns dynamic capacities to different types of tokens. ## experiments we evaluate the proposed smoe on a many-tomany multilingual neural machine translation task. ## analysis the advantage of smoe is that it can assign dynamic capacity to different tokens, with some tokens passing through only one expert and others passing through multiple experts. here, we define the requested capacity (rc) as the average number of experts that a token need to pass in one smoe block. rc of a token is dictated by how the smoe gates route it throughout the different strata. to understand what may affect the rc of tokens, we examine three potential influencing factors: the language of the input token, the frequency of the token, and the depth of the smoe block. all analysis is conducted using the smoe-4-4-4-4 model trained on the m15 dataset. ## conclusion this work presents stratified mixture of experts (smoe) models, a novel design for moes that is capable of dynamically assigning capacity to input tokens. through experimental evaluation at three scales (m4, m15, and opus-100), we have demonstrated that the proposed smoe model surpasses the performance of many current state-of-the-art moe methods. this proves that dynamically assigning capacity to tokens in moe models is a viable solution to address the moe's parameter inefficiency. additionally, we conduct a thorough analysis to investigate the factors that 12865 influence dynamic capacity assignment, including the language of the tokens and the location of the smoe block within the model architecture.