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

TL;DR

The paper tackles parameter inefficiency in Mixture-of-Experts (MoE) by introducing MPOE, which uses Matrix Product Operator decomposition to share a global central tensor across all experts while keeping expert-specific auxiliary tensors. A gradient-masking technique is proposed to stabilize fine-tuning and address unbalanced optimization. Empirical results on T5 and GPT-2 show substantial parameter savings (up to 27.2x) with competitive or improved performance, including strong gains in multi-task GLUE settings. This approach offers a principled path to scalable, parameter-efficient MoE architectures without sacrificing expressivity, and it is complemented by extensive ablations and analyses.

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 among experts. We adopt the 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/MPOE.

Figures (2)

• Figure 1: MPO decomposition for matrix $\mathbf{W}_{I\times J}$ with five local tensors. Auxiliary tensors ($\{\mathcal{A}_{i}\}_{i=1}^{4}$) and central tensor ($\mathcal{C}$) are marked in orange and blue, respectively.
• Figure 2: Illustration the proposed MPOE architecture and gradient mask strategy. We decompose the weight matrix of each expert in the MoE architecture into five local tensors using MPO, containing four auxiliary tensors and one central tensor, which are marked in orange and blue, respectively. In our approach, the central tensor of the $n$ experts is shared in the MPOE architecture. During optimization, each backward propagation process updates a set of auxiliary tensors while updating the central tensor with a probability of $p_b$ (the mask probability of the central tensor), which can effectively avoid the unbalanced optimization of the central tensor.