Uni-MoE: Scaling Unified Multimodal LLMs with Mixture of Experts

Yunxin Li; Shenyuan Jiang; Baotian Hu; Longyue Wang; Wanqi Zhong; Wenhan Luo; Lin Ma; Min Zhang

Source-linked AI summary

Uni-MoE: Scaling Unified Multimodal LLMs with Mixture of Experts

Yunxin Li, Shenyuan Jiang, Baotian Hu, Longyue Wang, Wanqi Zhong, Wenhan Luo, Lin Ma, Min Zhang

arXiv:2405.11273v1May 18, 2024cs.AIcs.CLcs.CVcs.MM

TL;DR

Scaling multimodal language models and broadening their data can improve capability but creates substantial computational costs, while prior MoE applications used fewer experts or limited modalities. Uni-MoE introduces a unified sparse-MoE MLLM with modality-specific encoders and connectors, parallel training, and progressive expert training. Across multimodal evaluations, it reports stronger performance and improved collaboration and generalization, with a stated video representation limitation.

Problem
Scaling MLLMs and expanding multimodal instruction data can impose substantial computational costs, while prior MoE models generally covered fewer experts or limited modalities.
Method
Uni-MoE combines modality-specific encoders and connectors with a sparse MoE LLM, expert-level model parallelism, modality-level data parallelism, and progressive training with LoRA.
Results
Uni-MoE consistently outperforms dense MLLMs on almost all evaluation benchmarks and improves multi-expert collaboration and generalization through modality-specific expert pre-training.
Takeaways & Limitations
The results support sparse MoE as a promising framework for unified MLLMs spanning multiple modalities and mixed multimodal tasks.
Takeaways & Limitations
Video inputs use average-pooled representations from randomly extracted frames, which can lose visual information relative to fine-grained temporal video representations.

Abstract

from arXiv · show

Recent advancements in Multimodal Large Language Models (MLLMs) underscore the significance of scalable models and data to boost performance, yet this often incurs substantial computational costs. Although the Mixture of Experts (MoE) architecture has been employed to efficiently scale large language and image-text models, these efforts typically involve fewer experts and limited modalities. To address this, our work presents the pioneering attempt to develop a unified MLLM with the MoE architecture, named Uni-MoE that can handle a wide array of modalities. Specifically, it features modality-specific encoders with connectors for a unified multimodal representation. We also implement a sparse MoE architecture within the LLMs to enable efficient training and inference through modality-level data parallelism and expert-level model parallelism. To enhance the multi-expert collaboration and generalization, we present a progressive training strategy: 1) Cross-modality alignment using various connectors with different cross-modality data, 2) Training modality-specific experts with cross-modality instruction data to activate experts' preferences, and 3) Tuning the Uni-MoE framework utilizing Low-Rank Adaptation (LoRA) on mixed multimodal instruction data. We evaluate the instruction-tuned Uni-MoE on a comprehensive set of multimodal datasets. The extensive experimental results demonstrate Uni-MoE's principal advantage of significantly reducing performance bias in handling mixed multimodal datasets, alongside improved multi-expert collaboration and generalization. Our findings highlight the substantial potential of MoE frameworks in advancing MLLMs and the code is available at https://github.com/HITsz-TMG/UMOE-Scaling-Unified-Multimodal-LLMs.

1 INTRODUCTION

Uni-MoE addresses the computational cost of scaling multimodal models by applying sparse MoE architecture to a unified MLLM spanning multiple modalities. Its progressive training strategy improves expert collaboration and generalization, while evaluations report strong performance across multimodal benchmarks.

Motivation: Scaling MLLMs improves multimodal capability but usually incurs substantial training and inference overhead.The motivation includes larger models and more diverse instruction data, alongside the computational costs of increasing model size.
Motivation: The framework addresses a gap in prior MoE MLLMs, which largely focused on text-only or image-text models with fewer experts and fewer modalities.The paper identifies scaling beyond four experts and beyond image-text inputs as largely unexplored.
Training Strategy: Uni-MoE uses progressive training: cross-modality alignment, modality-specific expert training, and unified MoE tuning with LoRA on mixed multimodal data.The stages map modalities into a shared language space, refine expert domain proficiency, and integrate the experts into the unified model.
Framework: Uni-MoE is a sparse MoE-based unified MLLM integrating video, images, text, audio, and speech through modality-specific encoders, connectors, and an MoE-equipped LLM.It uses expert-level model parallelism and modality-level data parallelism to support increased experts and diverse multimodal data.
Training Strategy: Pre-training experts on individual modalities enhances multi-expert collaboration and generalization compared with standard MoE tuning using identical initial expert parameters.The contribution statement attributes this finding to detailed experiments in Table 8.
Results: Uni-MoE consistently outperforms dense MLLMs on almost all evaluation benchmarks, including complex out-of-domain tasks.The paper also reports that auxiliary balancing loss is unnecessary for superior collaboration and generalization in some settings, though its benefits grow with more experts and routing options.

2 RELATED WORK

Related work established multimodal representation and MoE-based efficiency as complementary directions. Uni-MoE builds on these developments to explore a unified MoE model across a broader set of modalities.

Unified Multimodal Models: Transformer-based models enabled the integration of image, language, speech, audio, and video into unified representational spaces.The cited background connects this integration to generative and interpretive capabilities across data types.
Unified Multimodal Models: MLLMs commonly combine a pretrained visual backbone, an LLM, and a vision-language connector for image and video interpretation.The connector harmonizes visual encoder outputs with the language model representation.
Mixture of Experts: MoE architectures improve large-model efficiency by selectively activating network components instead of processing every input with all parameters.Related examples include language and image-text models that achieve competitive performance with fewer activated parameters.
Research Gap: Prior multimodal MoE work primarily addressed text-only or image-text models, leaving unified models with more than four experts and broader modalities largely unexplored.This gap motivates Uni-MoE’s focus on a unified multimodal MoE architecture.

3 UNI-MOE

Uni-MoE combines modality-specific encoders and connectors with a sparse MoE language model to process diverse modalities in a unified representation. Its progressive training strategy specializes experts before joint mixed-modality tuning, improving convergence and stability.

Architecture: Uni-MoE uses audio, speech, and visual encoders with modality connectors that project inputs into the LLM’s language space.Audio and speech connectors use Q-Former modules and linear projection layers to produce language-compatible representations.
Architecture: The model concatenates modality tokens into one sequence and applies a sparse router that selects Top-k experts for each token.Selected expert outputs are combined according to routing probabilities.
Training Strategy: LoRA activates each expert while the pretrained LLM and added expert parameters remain frozen, reducing the parameters that must be updated.Each expert receives learnable low-rank parameters added to its pretrained linear weights.
Training Strategy: Uni-MoE training proceeds through cross-modality alignment, modality-specific expert training, and joint tuning on mixed multimodal instructional data.The optimization framework initializes modality encoders and connectors, performs alignment, then updates linear and LoRA parameters during instruction tuning.
Training Strategy: Experts refined during the modality-specific stage converge faster and more stably on mixed-modality data than less specialized configurations.The reported comparison concerns complex mixtures including video, audio, and images; the figure caption also identifies MoE-Task3(4-2) as having more stable coverage than smaller or identical expert settings.

4 EXPERIMENTS

Uni-MoE is evaluated across speech-image, audio-text, image-text, and video understanding benchmarks using specialized datasets, metrics, and comparisons with dense or single-expert models. Results report strong performance across modalities, including gains on audio and video tasks, while video evaluation uses a simplified frame representation that can lose visual information.

Evaluation Setup: Speech-image and long-speech experiments test zero-shot reasoning on speech-converted visual questions, held-out long-speech tasks, and the English High School Listening Test.A-OKVQA, OK-VQA, and VQAv2 use speech-synthesized questions; MMBench-Audio and RACE-Audio evaluate long-speech understanding.
Evaluation Setup: The evaluation spans speech-image, long-speech, audio-text, image-text, and video-understanding benchmarks, including a newly collected English High School Listening Test.The benchmark suite includes five speech-related, three audio-understanding, and two video-understanding tasks.
Speech-Image and Speech-Text Understanding: Uni-MoE improves unseen audio-image reasoning and challenging long-speech understanding compared with dense and single-modality expert variants.The reported comparisons include A-OKVQA, OK-VQA, VQAv2, MMBench-Audio, and the English High School Listening Test.
Audio-Text Understanding: Uni-MoE surpasses baseline performance by 8.4% on ClothoV1 and 10% on ClothoAQA for audio-text understanding.The audio benchmarks focus on environmental audio and use audio captioning and question-answering tasks.
Video QA: Video evaluation uses average-pooled representations from eight randomly extracted frames, which can lose visual information relative to fine-grained spatiotemporal video processing.This simplified representation is used to speed training and inference.
Video QA: Uni-MoE achieves the best ActivityNet-QA performance among the compared video baselines, exceeding Video-ChatGPT by 7.5%, while remaining competitive on MSVD-QA.The results attribute this to integrating cross-modality reasoning with experts focused on distinct audio and image modalities.

5 EXPERIMENTAL ANALYSIS

Uni-MoE’s routing and ablation analyses show that specialized experts collaborate across modalities, improving mixed-modality robustness and generalization compared with dense or identical-expert configurations. The results also identify design choices that strengthen performance, including two active experts, four total experts, auxiliary balancing loss, and interval MoE placement.

Routing and pathways: Experts develop modality-sensitive task divisions, with routing assigning tokens to specialized experts according to fine-tuning knowledge.Experts 2 and 4 dominate different workloads, while text-token routing shifts with audio or image context and supports cross-modality reasoning.
Routing and pathways: Uni-MoE’s token pathways show modality preferences across experts and layers, supporting strong multimodal interaction and learning.The pathway analysis uses PCA to identify the top-10 activated pathways, whose expert indexes do not strictly correspond to earlier routing labels.
Comparative performance: Uni-MoE consistently outperforms dense models on almost all evaluation benchmarks and shows less performance bias and better out-of-domain generalization on mixed multimodal data.It improves long-speech benchmarks even when long-speech training data is a small proportion of the mixture.
Ablations: Two active experts improve performance over one, and four total sparse experts provide better comprehensive cross-modal performance than dense expert configurations.The gains are attributed to collaboration among modality-specific experts, particularly those trained on visual and speech tasks.
Ablations: The tri-phase training strategy improves multimodal benchmarks, while identical experts yield negligible gains and can produce unstable performance as their number increases.Auxiliary balancing loss further enhances expert synergy and overall performance across modalities.
Comparative performance: Larger dense LVLMs outperform Uni-MoE on some image-text benchmarks, but Uni-MoE excels over similar-sized MLLMs with matched instruction data and surpasses unified multimodal baselines in other modalities.Uni-MoE also shows improved video understanding when image-text data is added.
Qualitative analysis: Uni-MoE handles cross-modality inputs and recognizes long human speech and speech in video outside its training data.These capabilities are reported in generated cases alongside quantitative evaluations.

6 CONCLUSION

The paper develops Uni-MoE by integrating sparse Mixture-of-Experts architecture into unified multimodal modeling and introducing a three-phase training strategy. Experiments and ablations report improved cross-modal performance, reduced bias, robustness, and adaptability across diverse multimodal datasets.

Conclusion: Uni-MoE integrates MoE architecture into large multimodal models and uses a three-phase strategy to improve stability and generalization.The strategy is evaluated across cross-modality comprehension, long-form speech, and video reasoning benchmarks.
Conclusion: Uni-MoE surpasses existing benchmarks and conventional MoE models with identical experts across cross-modality and mixed-modality frameworks.Ablation studies support the importance of the tailored training strategy for robustness and adaptability.

Comparative Visualization Analysis of Uni-MoE (4 experts) trained with MoE-Task2 and Pure-Task1

The visualization analysis contrasts expert routing under MoE-Task2 and Pure-MoE-Task1 across multimodal combinations. MoE-Task2 develops clearer expert preferences, whereas Pure-MoE-Task1 shows more balanced and disorderly assignments.

Experimental setup: The analysis covers five multimodal combinations using 200 data-pair samples per figure, with routing measured after one training epoch.The combinations include image-text, audio-text, video-text, image-audio, and video-audio-text.
Routing distributions: Pure-MoE-Task1 produces relatively balanced expert loads and modality preferences, unlike the more differentiated routing observed in MoE-Task2.The Pure-MoE-Task1 model also performs below other Uni-MoE variants in the ablation analysis.
Expert specialization: Finetuning individual experts transfers single-modality capabilities, while identical experts do not distinctly separate their functions.The analysis nevertheless finds that identical experts can reveal useful expert-specific patterns in some cases.
Token pathways: Pure-MoE-Task1 has more disorderly and diverse token paths, whereas MoE-Task2 exhibits a distinctive preference for particular experts.The pathway trends align with the corresponding routing distributions.
Visualization overview: Figures 8–13 visualize expert loadings, modality distributions, and activated token pathways for Pure-MoE-Task1 and MoE-Task2.The figure captions organize these views by task and analysis type.