Mixture-of-Transformers: A Sparse and Scalable Architecture for Multi-Modal Foundation Models (2411.04996v2)

Abstract: The development of 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).

• The paper presents a novel MoT architecture that decouples modality-specific parameters while maintaining global self-attention.
• It achieves strong performance across text, image, and speech tasks, using only 55.8% and 37.2% FLOPs in key experimental settings.
• The architecture offers a scalable blueprint for efficient multi-modal learning and inspires future exploration of hybrid sparse models.

Mixture-of-Transformers: A Sparse and Scalable Architecture for Multi-Modal Foundation Models

The presented research introduces the Mixture-of-Transformers (MoT), a novel architecture devised to address computational challenges in large-scale multi-modal model pretraining. MoT focuses on optimizing the processing of various data modalities—such as text, images, and speech—by adopting a sparse architecture that disentangles non-embedding parameters by modality while maintaining global self-attention. This architecture reduces computational costs significantly compared to dense models traditionally used in multi-modal systems.

Decoupling Modality-Specific Parameters

MoT stands apart in its architectural design by decoupling each modality's non-embedding parameters, which include feed-forward networks, attention matrices, and layer normalization components. This design allows for optimal processing of each modality’s unique characteristics while ensuring that cross-modal relations are efficiently captured via a shared attention mechanism. Such an approach not only facilitates efficient computation but also empowers improved model performance across different modalities.

Strong Performance Across Settings

The paper provides an extensive evaluation of MoT architectures across three significant experiments: Chameleon settings for text-and-image objectives, speech integration in the Chameleon framework, and Transfusion settings incorporating diffusion-based objectives. In each of these settings, MoT demonstrates substantial reductions in the computational overhead compared to dense baselines. Notably, in the Chameleon 7B setting, MoT achieves comparable performance to its dense counterpart while utilizing only 55.8% of the FLOPs for text-and-image tasks, while the inclusion of a speech modality achieves baseline performance with only 37.2% of the FLOPs.

Implications and Efficiency Gains

MoT’s efficiency is further highlighted by direct performance metrics gathered using AWS infrastructure with NVIDIA A100 GPUs. The architecture achieves significant practical benefits; for instance, in wall-clock time, MoT matched dense baseline performance in image and text tasks with substantial reductions in required training time. It underscores MoT's ability to deliver high performance while optimizing the use of computational resources.

Exploring Hybrid Models

The paper also considers potential synergies between the proposed MoT framework and other sparse model architectures, such as Mixture-of-Experts (MoE). An exploratory variant that integrates MoE within MoT demonstrates additional gains in efficiency and performance, suggesting a promising pathway for mixed architectures that leverage the strengths of different sparse methodologies.

Critiques and Future Trajectories

While Mixture-of-Transformers provides a robust framework for multi-modal learning, the models were primarily tested in controlled settings constrained to specific FLOP configurations. Future explorations could extend to dynamically adjusting sparsity patterns per modality depending on the real-time computational and application-specific demands.

Additionally, MoT's scalability across even larger datasets and broader modality types remains a promising area for subsequent research. As multi-modal models continue to grow in scale and complexity, frameworks like MoT provide a crucial blueprint for achieving efficient, scalable AI systems capable of seamlessly integrating diverse data modalities. This architecture offers an exciting path forward in the continuous pursuit of advanced AI modeling approaches.