Parameter-Efficient Mixture-of-Experts Architecture for Pre-trained Language Models (2022)

Abstract: Recently, Mixture-of-Experts (short as MoE) architecture has achieved remarkable success in increasing the model capacity of large-scale language models. However, MoE requires incorporating significantly more parameters than the base model being extended. In this paper, we propose building a parameter-efficient MoE architecture by sharing information across experts. We adopt matrix product operator (MPO, a tensor decomposition from quantum many-body physics) to reconstruct the parameter matrix in the expert layer and increase model capacity for pre-trained language models by sharing parameters of the central tensor (containing the core information) among different experts while enabling the specificity through the auxiliary tensors (complementing the central tensor) of different experts. To address the unbalanced optimization issue, we further design the gradient mask strategy for the MPO-based MoE architecture. Extensive experiments based on T5 and GPT-2 show improved performance and efficiency of the pre-trained language model (27.2x reduction in total parameters for the superior model performance, compared with the Switch Transformers). Our code is publicly available at https://github.com/RUCAIBox/MPO/MPOE .

## introduction large-scale pre-trained language models (plms), such as bert @xcite and t5 @xcite , have become the de facto standard in natural language processing (nlp). by involving a huge number of parameters pre-trained on the general-purpose corpus, plms can achieve excellent performance in many nlp tasks. in order to increase the model capacity, a promising direction is to explore the scaling properties with the mixture-of-experts (moe) paradigm @xcite for developing more powerful plms. by incorporating multiple expert networks, moe schedules the learning of data samples through a routing component that is usually implemented by some gating function, which increases model capacity without a proportional increase in computation costs. despite the effectiveness, it has been shown that the moe architecture is parameter inefficient @xcite , considering the yielded improvement w.r.t. the involved costs. most of the existing studies @xcite @xcite attribute this issue to the unbalanced load of experts, focusing on improving the routing strategies. however, an important question about the moe architecture has been neglected in previous studies: whether the increased parameters from the experts are all necessary to increase the model capacity. as different experts from an moe network are often trained with correlated data samples (e.g., sample correlation from training data), it is likely to lead to parameter redundancy across experts. indeed, expert redundancy has been identified in existing studies, where @xcite distills sparse moe models into dense models and @xcite prunes experts to compress moe models. this finding motivates us to develop a parameter-efficient moe architecture by reducing its parameter redundancy. intuitively, a straightforward approach is to share a certain proportion of parameters among experts. however, it is difficult to identify and optimize the key parameters that encode the shared information across experts, since expert networks typically consist of dense matrices. to address this issue, we propose a novel parameter sharing approach inspired by the matrix product operators (mpo) decomposition from quantum many-body physics @xcite , which decomposes a matrix into a sequential product of local tensors (either central or auxiliary tensors). unlike other matrix decomposition methods, mpo can effectively reorganize and aggregate important information of the original matrix into the central tensor. the auxiliary tensors, on the other hand, serve to complement the central tensor for recovering the original matrix @xcite . in the setting of moe, considering the small parameter variations among experts, we speculate that the central tensors of different experts (with mpo decomposition for each expert) are likely to be very similar. if the central tensors could be shared for all expert networks, we would significantly reduce the parameters of the moe architecture. to this end, we propose a novel mpobased parameter-efficient moe architecture, called mpoe. based on classic moe architecture @xcite , our approach introduces a major extension allowing experts to share a global central tensor while keeping expert-specific auxiliary tensors. in our setting, the parameter matrix of in a single expert is formed by the product of the globally shared central tensor and the corresponding auxiliary tensors. since the central tensor contains most of the parameters from an mpo decomposition, our mpoe approach can significantly reduce the parameters of the original moe architecture. another major merit of mpo is that auxiliary tensors are closely entangled with the central tensor @xcite , and it is theoretically guaranteed that any change from auxiliary tensors can be propagated to the central tensor. that is to say, though a large proportion of parameters are shared, local auxiliary tensors still enable the experts to capture specific variations or differences according to routing data samples. however, directly optimizing the mpoe architecture is likely to lead to an unbalanced optimization issue, where the central tensors are updated more frequently than auxiliary tensors during fine-tuning. therefore, we further propose a gradient mask strategy that masks the central tensor gradient to effectively alleviate the unbalanced optimization issue. to the best of our knowledge, this is the first attempt to reduce the parameter redundancy of the moe architecture with structural matrix decomposition. we conduct extensive experiments to evaluate the effectiveness of the mpoe architecture on two representatives plms, t5 and gpt. experiments have demonstrated the effectiveness of our approach in increasing model capacity (27.2x fewer parameters for the superior model performance, compared with several competitive moe-enhanced plms. ## approach to reduce the information redundancy across different experts, we design an mpo-based moe architecture for increasing the model capacity in a parameter-efficient way. we firstly describe the mpo-based moe architecture and then introduce an improved optimization algorithm for learning the parameters in this architecture. ## experiments in this section, we first set up the experiments and then report the results and analysis. then, we conduct a detailed analysis under different experimental settings. here, we use t5 @xcite and gpt-2 @xcite models as the base model in our experiments. ## related work we will review the related works in four aspects. ## conclusion in this paper, we proposed a parameter-efficient moe architecture for increasing model capacity based on the mpo decomposition. first, we shared the central tensors among different experts based on mpo decomposition, which largely reduced the model parameters of moe architecture. then, we designed the gradient mask strategy to alleviate the unbalanced optimization issues and ensured that different tensors capture different types of information efficiently. extensive experiments have shown that our approach outperforms several competitive plm scaling strategies, especially in terms of improving the parameter efficiency of the moe architecture. in the future, we will enhance the proposed mpoe approach with recently proposed routing methods, such as baselayer @xcite , hashlayer @xcite and gshard @xcite . we will also consider exploring additional decomposition methods for developing parameter-efficient moe architecture.