Mixture-of-Transformers: A Sparse and Scalable Architecture for Multi-Modal Foundation Models
Weixin Liang, Lili Yu, Liang Luo, Srinivasan Iyer, Ning Dong, Chunting Zhou, Gargi Ghosh, Mike Lewis, Wen-tau Yih, Luke Zettlemoyer, Xi Victoria Lin
2024-11-08
Summary
This paper presents Mixture-of-Transformers (MoT), a new architecture designed to efficiently handle multiple types of data, like text, images, and speech, while reducing the computational resources needed for training.
What's the problem?
As large language models (LLMs) evolve to process different types of information (multi-modal systems), they require much larger datasets and more computing power. This makes it challenging to train these models effectively without using excessive resources.
What's the solution?
The authors introduce MoT, which separates the model's components based on the type of data being processed. This allows the model to focus on specific tasks while still using a global approach to understand all the input data. They tested MoT across various scenarios and found that it can achieve similar performance to traditional models while using significantly less computational power. For example, in some cases, MoT used only about 55% of the resources compared to other models but still produced high-quality results.
Why it matters?
This research is important because it provides a more efficient way to train models that can handle multiple types of data. By reducing the amount of computing power needed, MoT makes it easier for researchers and developers to create advanced AI systems that can perform well in real-world applications without requiring massive resources.
Abstract
The development of large language models (LLMs) has expanded to multi-modal systems capable of processing text, images, and speech within a unified framework. Training these models demands significantly larger datasets and computational resources compared to text-only LLMs. To address the scaling challenges, we introduce Mixture-of-Transformers (MoT), a sparse multi-modal transformer architecture that significantly reduces pretraining computational costs. MoT decouples non-embedding parameters of the model by modality -- including feed-forward networks, attention matrices, and layer normalization -- enabling modality-specific processing with global self-attention over the full input sequence. We evaluate MoT across multiple settings and model scales. In the Chameleon 7B setting (autoregressive text-and-image generation), MoT matches the dense baseline's performance using only 55.8\% of the FLOPs. When extended to include speech, MoT reaches speech performance comparable to the dense baseline with only 37.2\% of the FLOPs. In the Transfusion setting, where text and image are trained with different objectives, a 7B MoT model matches the image modality performance of the dense baseline with one third of the FLOPs, and a 760M MoT model outperforms a 1.4B dense baseline across key image generation metrics. System profiling further highlights MoT's practical benefits, achieving dense baseline image quality in 47.2\% of the wall-clock time and text quality in 75.6\% of the wall-clock time (measured on AWS p4de.24xlarge instances with NVIDIA A100 GPUs).