MoExtend: Tuning New Experts for Modality and Task Extension (2024)

Abstract: Large language models (LLMs) excel in various tasks but are primarily trained on text data, limiting their application scope. Expanding LLM capabilities to include visionlanguage understanding is vital, yet training them on multimodal data from scratch is challenging and costly. Existing instruction tuning methods, e.g., LLAVA, often connects a pretrained CLIP vision encoder and LLMs via fully fine-tuning LLMs to bridge the modality gap. However, full fine-tuning is plagued by catastrophic forgetting, i.e., forgetting previous knowledge, and high training costs particularly in the era of increasing tasks and modalities. To solve this issue, we introduce MoExtend, an effective framework designed to streamline the modality adaptation and extension of Mixtureof-Experts (MoE) models. MoExtend seamlessly integrates new experts into pre-trained MoE models, endowing them with novel knowledge without the need to tune pretrained models such as MoE and vision encoders. This approach enables rapid adaptation and extension to new modal data or tasks, effectively addressing the challenge of accommodating new modalities within LLMs. Furthermore, MoExtend avoids tuning pretrained models, thus mitigating the risk of catastrophic forgetting. Experimental results demonstrate the efficacy and efficiency of MoExtend in enhancing the multimodal capabilities of LLMs, contributing to advancements in multimodal AI research. https://github.com/zhongshsh/MoExtend.

heading: 1 Introduction
text: General-purpose large language models (LLMs) have demonstrated their effectiveness across a broad spectrum of application scenarios, such as conversational chatbot (Ouyang et al., 2022), document analysis (Radford et al., 2019), and coding (Chen et al., 2021). While the most powerful LLMs, such as ChatGPT (Radford et al., 2019), Llama (Touvron et al., 2023), and Mixtral (Jiang
et al., 2024), are predominantly trained on textual data, there is a growing interest in extending their capabilities to support a wider array of applications beyond natural language processing, especially with a significant focus on vision-language understanding (Liu et al., 2023a; Zhu et al., 2023; Liu et al., 2023b; Team et al., 2023). While training large models from scratch on multimodal data suffers from insufficient data (Zhu et al., 2023) and significant training costs (Team et al., 2023), most efforts have been focused on enhancing the multimodal capabilities of pretrained LLMs (Zhu et al., 2023; Liu et al., 2023b,a). To accomplish this, LLMs handle new modal data by processing representations extracted by encoders specific to each modality. For instance, the vision transformer pre-trained with CLIP (Radford et al., 2021) is utilized to encode visual images. Then, the model is trained using text-image Q&A pairs to carry out tasks based on these multimodal instructions.
The parameter-efficient approach to bridging the gap between modality-specific encoders and large language models (LLMs) involves the use of a few linear projection layers (Zhu et al., 2023) and Low-Rank Adaptation (LoRA) (Zhang et al., 2023a; Hu et al., 2021). However, this does not entirely mitigate the modality gap, limiting LLMs’ ability to fully understand new modalities. Consequently, State-of-the-art multimodal methods, e.g. LLaVA (Liu et al., 2023b), have sought to further enhance the multimodal capabilities of LLMs by fully fine-tuning these models on multimodal datasets (Lin et al., 2024). Despite these efforts, fully fine-tuning encounters two significant obstacles: 1) Catastrophic Forgetting: LLMs, when fine-tuned to effectively integrate various modalities, tend to lose the knowledge they had acquired previously (Luo et al., 2023). 2) Large fine-tuning cost: With the increasing sizes of LLMs, fully finetuning on larger models is becoming increasingly impractical. As a result, smaller models, like those
80
with 7 billion parameters, are often preferred. However, this preference restricts the exploration and utilization of the capabilities of larger LLMs. How to efficiently extend new modality to large LLM while reduce the side effect of catastrophic forgetting is an urging problem for multimodal LLMs.
Mixture-of-Experts (MoE) architectures enable LLMs to use the gate layer to dynamically select the most relevant experts from a diverse set of specialized experts, e.g. different MLP layers in Transformer, for a given query token. MoE helps to enlarge the model size by increase the number of experts while keeping low inference cost by selecting a sub set of experts for each token. For instance, the Mixtral-8x7B model (Jiang et al., 2024) incorporates 8 MLP experts per block, totaling 46.7 billion parameters, yet it selects only 2 experts, utilizing 12.9 billion parameters per token. Nonetheless, the current MoE models predominantly concentrates on the textual modality.
We introduce an extension strategy for MoE models, named MoExtend, designed to accommodate new modalities. This strategy involves incorporating new modality-specific experts and calibration modules into trained MoE models to enhance their capability to process additional modalities. MoExtend maintains the original MoE model parameters unchanged, while only trains the newly added experts and the corresponding gate layer. By doing so, MoExtend facilitates the efficient adaptation of new modalities into large models while also addressing issues of catastrophic forgetting (Liang et al., 2024, 2022). We observe that the rapid adaptation to new modalities relies on the weight initialization of new experts and gates, and the insertion position of these new experts. Thus, we introduce a simple yet effective scheme for selecting positions and weights of new experts based on evaluating distribution shifts. Utilizing the data from the new modality, we fine-tune the existing gate layers of the MoE model. Then, we infer the new modality data to the models before and after fine-tuning and get the average gate probability distribution for all samples. By comparing the degree of gate probability distributions before and after fine-tuning, we identify the top-k layers for adding experts by examining the magnitude of these shifts. Then, based on the probability distribution after fine-tuning, we determine the expert with the highest probability and replicate the gate and expert weights onto the newly incorporated expert.
Experimental results show that MoExtend
achieves a training speed acceleration ∼6 times faster than full fine-tuning, while also delivering superior performance. The positions selection scheme in MoExtend allows for fewer newly added experts, specifically, half the number of new experts required for the Mixtral model, which reduces training time to ∼30 hours without compromising performance. In addition, MoExtend helps mitigate the risk of catastrophic forgetting when extending MoE LLMs to handle multimodal inputs. Our contributions can be summarized as follows:
• We introduce MoExtend, a strategy designed to augment Mixture-of-Experts LLMs with new modalities by addition of new experts.
• MoExtend offers significant advantages, including substantially reduced fine-tuning costs, no additional costs during inference, and a minimized impact from catastrophic forgetting issue.

heading: 2 Methodology
text: In this section, we introduce MoExtend as an example of extending the visual modality for MoE models, which were originally designed for text modality only. As shown in Fig. 1, MoExtend consists of three stages: alignment, extension with extender, and fine-tuning for the extension part. The purpose of the alignment stage is to initially align the MoE LLM with the newly added visual modality using a pre-trained vision encoder. The extension stage determines which MoE layers should be extended to accommodate the new modality information. The fine-tuning stage is then employed to tune the newly added parameters, achieving the final expansion of multimodal information.

heading: 2.1 Alignment Stage
text: As illustrated in Fig. 1 (a), we train the newly added MLP using image-caption pairs from the LLaVA 1.5-558k dataset. This training aligns the modal information of images through the vision encoder (i.e., CLIP encoder) with textual modalities. Specifically, the caption c from the textual modality is projected via word embedding to T = [ti]Ni=1 ∈ RN×D, where D is the hidden size of LLM. Additionally, the image I is mapped through the vision encoder to V = [vi]Pi=1 ∈ RP×D, where P is the sequence length of visual tokens. Subsequently, the information from both modalities, T and V , is concatenated into the vector x0 ∈ R(N+P )×D. For an L-layer MoE LLM, the forward process can be
formulated as follows:
x′ℓ = MSA(LN (xℓ−1)) + xℓ−1, ℓ = 1 . . . L,
xℓ = MoE ( LN ( x′ℓ ))
+ x′ℓ, ℓ = 1 . . . L, (1)
where MSA represents the multi-head selfattention module and LN represents layer normalization. The final input to the model is LN(xL). During this stage, the structure of the MoE layer with m experts remains unchanged, as depicted in Fig. 2 (Left). The router predicts the probability of each token being assigned to each expert, and each token is computed by the top-k experts with the highest probabilities. The output of the MoE layer is a weighted sum as follows:
MoE(x) = ∑k
j=1 s(x)j · FFN(x)j , (2)
where k ≤ m. Note that the weighted summation in Eq. (2) is related to the outputs of experts with top-k probability. The parameter k has a significant impact on MoE LLMs. However, to consider the trade-off between training efficiency and model performance, it’s common to set k = 2. In this
paper, we also follow this setting. The [FFNi]mi=1 represents m experts, and
s(x)j = e f(x)j/ ∑m h=1 ef(x)h , (3)
where f(x) = Wx and W ∈ RD×m are the parameters of the router.

heading: 2.2 Extension Stage
text: To address the incorporation of additional modality information via extending the MoE layer, the most straightforward approach is to add a new expert to each MoE layer. However, this approach not only increases the parameter count significantly, leading to greater computational costs during training but also poses a potential risk of overfitting due to blindly adding a large number of parameters.
Therefore, in the extension stage, inspired by the concept of neural network pruning (Li et al., 2016; Gao et al., 2020), we construct an Extender to adaptively determine whether each MoE layer needs extension. Specifically, we randomly sample 10,000 instruction data related to the vision modality from the LLaVA 1.5-mix-665k dataset (Liu et al., 2023b)
as the validation set Se, with the remaining data forming the sub-training set St.
Next, for the model κ obtained from the alignment stage training, we make all routers of the MoE layers trainable while freezing all other parameters. Utilizing St, we tune κ for 1,000 steps to obtain κ′. Furthermore, we input Se into both κ and κ′, and count the occurrences of each expert being selected in every MoE layer, resulting in
Rκ = {rκij}m×L, Rκ′ = {rκ ′
ij }m×L. (4) After normalization as follows, we can estimate the probability distributions of each expert being selected in every MoE layer:
R̄κ = Rκ/(r κ 11 + r κ 21 + ...+ r κ m1),
R̄κ′ = Rκ′/(r κ′ 11 + r κ′ 21 + ...+ r κ′ m1).
(5)
It is worth noting that for 1 ≤ i ≤ L, ∑mi=1 rκi1 =∑m i=1 r κ ij and ∑m i=1 r κ′ i1 = ∑m i=1 r κ′ ij . Then, we can estimate the distribution differences of expert selections in each MoE layer between the two models κ and κ′ by calculating dj as follows:
dj = Std m i=1(r̄ κ ij − r̄κ ′ ij ), 1 ≤ j ≤ L, (6)
where Std denotes standard deviation. If dj is small, it implies that the MoE layer j exhibits minimal response variation to the current data of the image-text modality, hence, there’s no necessity to add new experts to this layer. Conversely, for MoE layers with larger dj , adding new experts can effectively address the learning of new modality information. We rank the MoE layers based on dj and introduce a new expert FFNm+1 to the top ⌊pL⌋ layers for original MoE LLM κ, with p set to 0.5 in this paper. In fact, the adaptive extension stage proposed in this section not only reduces computational costs during training and mitigates the
risk of overfitting but also accelerates the training of MoE LLM. For detailed analysis, please refer to Section 4.

heading: 2.3 Fine-tuning Stage
text: In addition to introducing an additional expert in certain MoE layers for the original κ, as mentioned in Section 2.2, and illustrated in Fig. 2, we also need to augment the parameters of the corresponding routers for these experts, i.e.,
Wnew = [W;vnew] ∈ RD×(m+1), (7)
where vnew ∈ RD×1, Furthermore, we add some Calibration modules to all experts in the MoE layers to mitigate changes in gate weights due to the addition of modalities. These newly introduced trainable parameters constitute the extension part. In this section, we fine-tune the extension part using the LLaVA 1.5-mix-665k dataset to enhance the final performance of LLM.
Specifically, we first consider the initialization of the newly added m+ 1-th expert and its corresponding router parameters vnew. In this work, for the j-th MoE layer, we consider directly copying the expert and router parameters corresponding to
max(rκ1j , r κ 2j , · · · , rκmj), (8)
as initialization for the new parameters. This is because intuitively, the newly added expert is primarily intended to address the new modalities, and it is appropriate to initialize it with the existing expert that has the highest response to the new modalities. In Section 4, we will demonstrate that the initialization of the new parameters significantly affects the probability of an expert being selected by the MoE mechanism, thereby affecting the final performance of the MoE LLM.
Furthermore, since some MoE layers have added experts, s(x)j will change according to Eq. (3). For example, for a fixed input x, the new probability s(x)′j satisfies
s(x)′j = e f(x)j/( ∑m h=1 ef(x)h + ef(x)m+1)
≤ ef(x)j/ ∑m
h=1 ef(x)h = s(x)j ,
(9)
This causes the feature distribution of the original MoE κ regarding previously learned knowledge to change during forward propagation, resulting in some degree of forgetting of existing knowledge by the model, thereby affecting performance. To address this issue, we add a Calibration module sc(·) for each expert such that
MoE(x) = ∑k
j=1 s(x)j · [1 + sc(x)] · FFN(x)j ,
(10) and sc(·) is a two-layer GELU neural network W1(GELU(W2(·))). Here, the weights of W1 are initialized to 0, and W2 uses normal initialization. This initialization ensures that the calibration term sc(x) = 0, maintaining consistency with the model’s output features when sc(·) is not added, thus preventing significant interference with model output features due to the addition of sc(·), which could lead to abnormal loss and affect model training. For a fair comparison, all training hyperparameters, training methodologies, and loss functions with LLaVA 1.5-558k and LLaVA 1.5-mix-665k in all stages remain consistent with LLAVA.

heading: 3 Experiments

heading: 3.1 Experimental Setup
text: Model Settings. To ensure fairness in experimental comparisons, we follow the settings outlined in LLaVA 1.5. We utilize CLIP (Radford et al., 2021) as the vision encoder, two linear layers with GELU (Hendrycks and Gimpel, 2016) as the vision projection, and other training hyperparameters are shown in Appendix Table 6.
Dataset. We utilize the same dataset as LLaVa 1.5 to train the model, consisting of LLaVA 1.5-558k for pretraining stage and LLaVA 1.5-mix-665k for instruction tuning stage (Liu et al., 2023b). The computational cost of MoExtend is ∼15 hours of pretraining and ∼30 hours of visual instruction tuning, while MoExtend-Full, the model trained like LLaVA, need ∼200 hours of instruction tuning.

heading: 3.2 Image Understanding Evaluation
text: Image Question Answering. As shown in Table 1, we assess MoExtend performance across four widely-used image question answering benchmarks. Compared to the state-of-the-art method LLaVA-1.5 (Liu et al., 2023b), MoExtend exhibits robust image understanding capabilities and achieves performance very close to that of LLaVA1.5. Specifically, MoExtend, which is trained with only 3B LLM parameters, surpasses LLaVA-1.5 13B, trained with 13B LLM parameters, by 3.1%, and outperforms the recent vision-language model HyperLLaVA (Anonymous, 2024) by over 4.8% on SQA. Remarkably, MoExtend achieves comprehensive superiority over IDEFICS-80B (Laurençon et al., 2024) with only 13B activated parameters, underscoring the strong comprehension abilities of MoE-LLaVA in vision features.
Performance on Multimodal Benchmarks. To comprehensively evaluate multimodal comprehension capabilities of MoExtend, we evaluate its performance across five widely-used benchmark toolkits, as shown in Table 1. Experimental results indicate that, under the same dataset and training settings, MoExtend, fine-tuned with only 3B LLM parameters, achieves performance on par with the state-of-the-art model on most benchmark toolkits. Particularly, MoExtend has significantly superior performance on MME, surpassing the existing leading model LLaVA 1.5-13B by 178.8 points, indicating that MoExtend facilitates a efficient expansion of modalities.
Comparison with Forgetting. To mitigate catastrophic forgetting in LVLMs, MoExtend fine-tunes LLM through calibration and the addition of new experts, thereby preserving the performance of LLM’s original modalities. To evaluate the superiority of our fine-tuning strategy in preserving the understanding capabilities of LLM’s original modalities, we evaluate the performance of LVLMs using different fine-tuning methods on pure text metrics as shown in Table 2. Specifically, we compare the performance of LLaVA-1.5, MoExtendFull, MoE-LLaVA, and MoExtend with original LLMs in Table A. Across all metrics, MoExtend exhibits performance similar to the original LLM. Additionally, we observe only slight decreases for LLaVA-1.5, while MoE-LLaVA and MoExtendFull show significant declines relative to the original LLM model in pure text evaluation metrics,
suggesting that full-parameter fine-tuning may lead to catastrophic forgetting for MoE-type LLMs, whereas non-MoE-type LLMs are less affected.

heading: 4 Ablation Study and Analysis
text: Effect of Model Architectures. We investigate the impact of different architectures on the performance of MoExtend. While the intuitive approach of adding new experts to all layers might seem optimal, our experiments, detailed in Table 3, reveal comparable performance between models with ex-
perts added to every layer (All layer), the first half (First-half), or every alternate layer (Interval). Additionally, results from models with experts added only to the first quarter (First-quarter) or every alternate layer starting from the first layer (Firstinterval) indicate performance degradation when too few layers receive additional experts. This finding informs our extension stage design, where experts are appropriately added to half of the layers.
As depicted in Fig. 3 (Left), our extension stage identifies layers requiring new experts. MoExtend based on our proposed strategy, as demonstrated in Table 3, performs on par with the current optimal insertion strategy (First-half, Interval). Furthermore, Fig. 3 (Right) shows that our extension strategy converges at a rate comparable to the op-
timal insertion strategy during training, validating its effectiveness on accurately determining the appropriate layers for adding new experts without extensive experimentation.
Effect of Initialization. As depicted in Table 4, we analyze the impact of expert and router initialization on the performance of MoExtend. If the parameters of the new experts and router dimensions are directly copied from fixed positions i of experts and corresponding dimensions of routers at each layer (Copy(i)), the performance of copying experts from different positions is relatively close and lower than that of MoExtend.
Additionally, we explore the performance when the router parameters are not directly copied from the corresponding router parameters of the i-th expert, but initialize directly with zeros or with the mean of the initial parameters of the eight experts (Mean). Experimental results indicate that initializing the router with zeros generally results in poorer performance compared to direct copying (Ours). Mean initialization implies that the new experts are a few selected in the initial state, and later in the instruction tuning stage the new experts are selected through gradient updates. In fact, this performance difference is mainly due to the fact that such an ini-
tialisation will lead to the newly added experts not being easily selected during the training process, so that the newly added experts are not fully trained or not used for new modality. Specifically, take the "Mean" initialisation as an example. Since the MoE layer generally selects the top-2 probability of experts for feature integration, the initialisation of "Mean" makes it difficult for the new experts to be selected with a large probability. Since the new router parameters and experts are rarely updated, it is difficult to improve this situation during the training process.
However, experimental results show that this initialization method leads to inferior performance. Furthermore, to investigate the impact of initialization methods on performance, we calculate the ratio of expert selection for different initializations as shown in Fig. 4, and find that models initialized with Zero and Mean are both unbalanced in expert selection, while MoExtend is more balanced. This finding indicates that the balance of expert selection is closely related to model performance.
The Design of Calibration Modules. As shown
in Fig. 5, we design two concise calibration modules (Type1, Type2) to investigate the impact of these modules on MoExtend performance under two integration modes (Liang et al., 2020; Huang et al., 2020; Zhong et al., 2023d,c): addition (a) and multiplication (b). Type1 consists of a simple learnable parameter 1×m, while Type2 consists of two simple linear layers connected by the GELU activation function. To minimize the disruption of router performance by calibration modules in the initial state, we mitigate the initial impact of calibration modules on routers through special initialization as shown in Table 5. In the additive mode of Type1, we use Zero initialization for calibration modules, while in the multiplicative mode, we use One initialization.
In the additive mode of Type2, we initialize the first linear layer normally and zero-initialize the second linear layer. In the multiplicative mode, it is hard to reduce the impact of calibration modules through appropriate initialization, so we opt for simple normal initialization for both linear layers. Type2 (b) does not exhibit any evaluation result in Table 5 because of gradient explosion, and the
experimental results indicate that Type2 (a) calibration module structure performs better than others.

heading: 5 Conclusion
text: In this work, we introduce MoExtend, an effective framework tailored to streamline the modality adaptation and extension of Mixture-of-Experts (MoE) models. MoExtend introduces new experts into MoE models by putting them at the parallel positions of the experts in MoE. Then MoExtend designs a method to select previous experts in MoE for initilizing the new experts. Finally, it only tunes the new experts on the corresponding modal data and tasks. This endows MoE with novel knowledge without necessitating the tuning of pretrained models such as MoE and vision encoders, thus avoiding the catastrophic forgetting issue. Furthermore, MoExtend facilitates rapid adaptation and extension to new modal data or tasks, thereby effectively addressing the challenge of accommodating new modalities within LLMs. Empirical results show the efficacy and efficiency of MoExtend in augmenting the multimodal capabilities of LLMs.

heading: 6 Limitation
text: In this work, due to limited GPU resource, we take the visual task as one example to validate the effectiveness our proposed MoExtend. So one limitation of MoExtend is that its performance is not investigated on the other modal data, such as speech, and other tasks, e.g., continue learning and streaming tasks. However, as aforementioned, MoExtend is a general approach to extend the MoE model to other modal data or tasks, because our design principle is to endows MoE with novel knowledge via tuning the new integrated experts, and does not involve any specific tasks or modality. Accordingly, we believe that by replacing the vision encoder in MoExtend with other modal encoder and inserting new experts like MoExtend, one can easily extend MoExtend to other modal data and tasks, which is also left as our future work to thoroughly test.

heading: 7 Related Work

heading: 7.1 Mixture of Experts
text: Mixture of Experts (MoE) (Masoudnia and Ebrahimpour, 2014; Riquelme et al., 2021; Zhou et al., 2022; Lin et al., 2024; Jiang et al., 2024) is a technique that leverages multiple sub-networks, also referred to as experts, to integrate features generated by different experts through adaptive strategies, thereby enhancing the overall performance of neural networks. The MoE layer, when processing each token, employs a router module to assign tokens to different experts, thereby reducing interference between different types of samples and keep low inference cost. In specific computational frameworks, MoE can achieve performance comparable to LLMs with a large amount of computational cost (Masoudnia and Ebrahimpour, 2014). Consequently, with the rapid advancement and application of LLMs, MoE is emerging as a promising and noteworthy paradigm for further enhancing LLM performance (Masoudnia and Ebrahimpour, 2014; Team et al., 2023).

heading: 7.2 Multimodal Model
text: Multimodal Learning involves leveraging various types of data, such as text, images, speech, and video, to train machine learning models for a more comprehensive understanding and inference capability (Bayoudh et al., 2022; Xu et al., 2023; Zhong et al., 2023b,a). By integrating and jointly modeling different modalities of data, multimodal learning enhances machines’ ability to comprehend and
express rich real-world information, thereby improving performance in tasks like image description, sentiment analysis, speech recognition, and video understanding.
Recently, with the advancement of LLM technologies, multimodal learning methods have been rapidly integrated into LLM to expand its understanding and analysis of different modalities, especially visual modality (Liu et al., 2023b; Bai et al., 2023). Recent efforts have focused on enhancing performance through methods such as adjusting datasets (Liu et al., 2023b), optimizing training strategies (Zhang et al., 2023b; Zhong et al., 2022), improving image resolution (Bai et al., 2023), enhancing image encoders (Fan et al., 2024; Gao et al., 2024), aligning inputs (Radford et al., 2021), and projecting layers (Wu et al., 2023; Liu et al., 2023b). These approaches, by fine-tuning datasets and model scales through expanded visual instructions, have endowed LLM with robust visual comprehension capabilities. However, most current methods for expanding modalities generally involve fine-tuning a significant portion of or all parameters on multimodal data, leading to substantial computational costs and risking performance degradation due to forgetting. Facing this dilemma, in this paper, we consider leveraging the strong base performance of MoE LLM to explore cost-effective methods for expanding LLM modalities by introducing new experts.

heading: 8 Hyperparameters

heading: 9 Acknowledgments
text: This work was supported by National Natural Science Foundation of China (No.61876045, 623B2099, U1711264). Pan Zhou acknowledges support from the Singapore Ministry of Education (MOE) Academic Research Fund (AcRF) Tier 1 grant.

heading: Future Work
text: Accordingly, we believe that by replacing the vision encoder in MoExtend with other modal encoder and inserting new experts like MoExtend, one can easily extend MoExtend to other modal data and tasks, which is also left as our future work to thoroughly test