PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2411.04996v2cs.CL

TL;DR

Multi-modal foundation models require substantial computation because text, images, and speech create distinct optimization challenges within unified systems. MoT decouples non-embedding transformer parameters by modality while preserving global self-attention, and it matches or exceeds dense baselines across settings at lower computational cost. The gains extend to practical training time, although text improvements are less pronounced in Transfusion and warrant further investigation.

  • Problem

    Multi-modal foundation models process text, images, and speech together but require substantially more data and computation than text-only models while facing modality-specific optimization challenges.

  • Method

    MoT decouples non-embedding parameters by modality, including feed-forward networks, attention matrices, and layer normalization, while retaining global self-attention over the full sequence.

  • Results

    MoT matches or exceeds dense baselines across Chameleon, Chameleon+Speech, and Transfusion while using substantially fewer FLOPs, including 55.8% in 7B Chameleon and less than one third for 7B Transfusion image performance.

  • Takeaways & Limitations

    MoT provides a scalable approach to multi-modal pretraining with lower FLOP and wall-clock costs while preserving cross-modal interactions and supporting large-scale distributed training.

  • Takeaways & Limitations

    Text gains are marginal to none in Transfusion, and the authors identify further investigation of text performance and hybrid strategies as future work.

Abstract

from arXiv · show

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).

1 Introduction

Multi-modal foundation models face high computational demands because modalities introduce distinct optimization challenges within unified transformers. MoT addresses this by decoupling non-embedding parameters by modality while preserving global self-attention, matching dense-model performance with substantially fewer FLOPs across settings.

  • Architecture: MoT applies modality-specific feed-forward networks, attention projections, and layer normalization to each token while retaining global self-attention across the full multi-modal sequence.This extends modality-aware sparsity across the entire transformer and preserves the dense counterpart’s computational structure and FLOP count.
  • Chameleon: 7B MoT matches the 7B dense baseline in Chameleon using 55.8% of the FLOPs across multiple data distributions.Results remain consistent at 37M, 94M, 443M, and 1.5B scales.
  • Chameleon with Speech: 37.2% of training FLOPs suffices for MoT to achieve performance similar to the dense baseline when speech is added as a third Chameleon modality.The result is consistent across multiple model scales.
  • Transfusion: A 760M MoT outperforms a 1.4B dense Transfusion baseline across image metrics, while 7B MoT matches dense image performance with less than one third of the FLOPs.The comparisons include CLIP, FID, CIDEr, image training loss, diffusion validation loss, and image captioning.
  • Efficiency: System profiling shows 7B MoT reaches dense-model image performance in 47.2% of the time and text performance in 75.6% of the time.Measurements used AWS p4de.24xlarge instances with NVIDIA A100 Tensor Core GPUs.
  • Sparse Baselines: MoT consistently outperforms a 4-expert MoE baseline, particularly on image and speech modalities, with an even larger advantage in wall-clock time.MoE-4x contains more sparse parameters than MoT, but those parameters do not increase training or inference FLOPs.

2 Method: Mixture-of-Transformers Architecture

Mixture-of-Transformers (MoT) decouples transformer parameters by modality while retaining global self-attention across interleaved multimodal tokens. This design supports modality-specific processing and diverse token representations or training objectives within one architecture.

  • Motivation: Multimodal inputs exhibit distinct feature-space clustering by modality despite uniform token processing, motivating modality-aware parameterization.The clustering is observed across multiple layers in PCA analyses of a Chameleon+Speech dense model.
  • Mixture-of-Transformers Architecture: MoT applies modality-specific weights to non-embedding transformer components while preserving global self-attention across the full multimodal sequence.The decoupled components include feed-forward networks, attention projections, and layer normalization, while global attention captures cross-modal relationships.
  • Flexible Modalities and Objectives: MoT supports both discrete autoregressive image tokens and continuous diffusion-trained image tokens, enabling different multimodal training objectives.The architecture is illustrated for Chameleon-style discrete tokens and Transfusion-style continuous tokens.
  • Mixture-of-Transformers Architecture: Tokens are grouped by modality, processed with modality-specific projections, globally attended, and passed through modality-specific output projections, normalization, and feed-forward networks.The computation concludes by combining outputs with residual connections and layer normalization.
  • Experimental Setting: In the Chameleon setting, text and images use a unified autoregressive objective with images represented as 1,024 discrete tokens.The image tokens are produced by a pretrained VQ-VAE model.

3 Experiments

Across Chameleon, speech-augmented Chameleon, and Transfusion settings, MoT accelerates multimodal pretraining while preserving or improving modality performance. Gains are strongest for image and speech, remain consistent across scales, and can exceed dense baselines at lower compute.

  • Chameleon: MoT matched the 7B dense Chameleon model’s performance using 55.8% of training FLOPs, with comparable or lower validation losses across text and image.At 55.8% of training steps, MoT achieved validation losses comparable to or below the dense model’s final losses.
  • Speech integration: In speech-augmented Chameleon, MoT matched dense speech pretraining loss using 22.9% of training steps and retained comparable or lower validation losses for text and image.MoT also outperformed baselines on speech validation datasets LL60K and PPL30K.
  • Speech integration: MoT’s modality-aware parameter partitioning avoids MoE load-balancing issues and supports stable speech performance across training and validation metrics.The architecture decouples non-embedding parameters by modality, including FFNs and projection matrices.
  • Transfusion: MoT’s text advantage was marginal to nonexistent in Transfusion, where separate image and text objectives already yielded near-optimal text performance.The authors identify this setting-specific result as requiring further investigation and possible hybrid strategies.
  • Transfusion: A 760M MoT model using half the FLOPs of a 1.4B dense baseline improved CLIP, FID, and CIDEr image-generation metrics.Scores were CLIP 0.214 vs 0.206, FID 21.145 vs 24.688, and CIDEr 0.320 vs 0.286.

4 Modality Separation in MoT: Leave-One-Out Analysis

The Leave-One-Modality-Out analysis shows that sharing transformer towers across modalities consistently worsens performance, supporting modality-specific parameter allocation in MoT.

  • Setup: The analysis compares three LOO variants against baseline MoT and a dense single-tower model while maintaining equivalent FLOPs.
  • Results: Sharing transformer towers across modalities consistently increases training and validation losses, supporting separate modality-specific towers in MoT.The LOO variants combine two modalities while isolating one, and the dense baseline combines all modalities in one tower under equivalent FLOPs.
  • Results: LOO-text yields the lowest text loss, while LOO-image and LOO-speech yield the lowest losses for their isolated modalities.
  • Results: Separating speech preserves speech gains but removes image and text improvements, whereas isolating image preserves image gains while merged text and speech degrade both.The effects are non-reciprocal across modality combinations.

5 Combining the Best of Both Worlds – Mixing Heterogeneous Transformers

The hybrid MoT + Text MoE-4x architecture improves text training and validation performance while preserving MoT’s image efficiency in both Chameleon and Transfusion settings.

  • Method: The approach replaces MoT’s text FFN with a four-expert MoE layer while leaving the image transformer unchanged.
  • Chameleon setting: In Chameleon, adding MoE-4x to MoT’s text tower accelerates text loss reduction while retaining MoT’s image benefits.
  • Transfusion setting: In Transfusion, the hybrid retains MoT’s image training-loss and speed advantages while accelerating text training-loss reduction.
  • Scope: The authors characterize these results as a preliminary proof of concept requiring further study of scalability and generalizability.
  • Transfusion setting: MoT + Text MoE-4x achieves the best text performance while maintaining comparable or slightly improved image performance relative to MoT in Transfusion.It significantly outperforms both the dense model and MoE-4x in the reported validation comparisons.

6 ML Systems Aspects of Mixture-of-Transformers

MoT’s modality-based sparsity can reduce parameter-to-FLOPs overhead and improve practical training efficiency, but its benefits depend on hardware, modality distribution, and implementation.

  • Communication volume: MoT adds 7(K−1)D^2 parameters versus more than 3(E−1)D^2 for MoE, so its typically smaller modality count K can yield a lower parameter-to-FLOPs ratio.
  • Compute efficiency: MoT’s overheads arise from CPU-GPU synchronization and sequential modality processing, which can underutilize GPUs when modality tokens are unevenly distributed.
  • Horizontal scaling: MoT’s image matching point improves from 42.1% to 21.6% of training steps and text matching from 75.7% to 50.9% as GPUs increase from 16 to 256.The horizontal-scaling experiment increases global batch size and total training tokens while holding training steps constant.
  • Scope: MoT’s reported wall-clock advantages may vary across clusters because the profiling used a specific AWS infrastructure configuration.
  • Wall-clock performance: MoT matches dense-model image quality in 47.2% of GPU training time and text quality in 75.6% under the reported 256-GPU setup.The comparison is against the dense Transformer baseline in wall-clock GPU training time.

7 Related Work

Related work advances multimodal generation through tokenized non-text sequences and sparse architectures, while MoT targets modality-aware sparsity across multimodal transformers.

  • Multimodal generation: Early multimodal systems primarily used late fusion for understanding, but lacked multimodal generation capabilities.
  • Multimodal generation: Tokenizing images and speech into discrete sequences enables unified autoregressive modeling for multimodal generation.Chameleon uses 1,024 discrete image tokens from a pretrained tokenizer for combined text-image sequences.
  • Sparse architectures: MoE methods dynamically select parameters for inputs, and recent work adapts them to multimodal tasks to address modality feature-space gaps.
  • MoT distinction: MoT differs from prior approaches by supporting both image and text generation while applying modality-aware processing across its transformer architecture.

8 Conclusion

MoT reduces computational costs for multimodal pretraining by decoupling non-embedding parameters by modality while preserving global self-attention. Across Chameleon, speech, and Transfusion settings, it maintains or improves performance with substantially lower computational and wall-clock costs.

  • MoT decouples non-embedding transformer parameters by modality while retaining global self-attention across multimodal sequences.This enables modality-specific processing without removing cross-modal interactions.
  • MoT matches or exceeds dense baselines across Chameleon, Chameleon+Speech, and Transfusion settings while using substantially fewer FLOPs.The reported gains extend to settings with distinct training objectives for different modalities.
  • MoT reduces wall-clock time for both text and image tasks, with further improvements observed when scaled across GPUs.The conclusion also reports preliminary potential for combining MoT with MoE-4x without increasing computational costs.

A.1 Diffusion for Image Generation

Diffusion models generate images by progressively adding noise and learning to reverse that process. The described implementation uses cosine noise scheduling, classifier-free guidance, and latent diffusion to reduce computational requirements.

  • Diffusion models generate data by adding Gaussian noise through a forward Markov process and learning iterative denoising in reverse.As the timestep approaches T, the data distribution approaches isotropic Gaussian noise.
  • The reverse process uses a neural network to predict added noise at each timestep while conditioning on extra context such as a text prompt.Sampling starts from Gaussian noise and applies the learned denoising steps to produce a data point.
  • The implementation uses a cosine scheduler and classifier-free guidance to improve qualitative image generation.Classifier-free guidance contrasts context-conditioned and unconditioned predictions.
  • Latent diffusion performs forward and reverse diffusion in a lower-dimensional latent space to reduce computational requirements.The original data is first encoded into a latent representation using a variational autoencoder.

A.2 Transfusion Model Architecture

The Transfusion architecture processes text and images with a shared transformer while using modality-specific input and output modules. It trains with a combined language-modeling and diffusion objective and alternates between text generation and image denoising during inference.

  • A single transformer processes the combined sequence across modalities, with lightweight modality-specific modules mapping inputs into a shared vector space.Text uses embedding and output matrices, while images use a U-Net to compress local patch windows.
  • The model minimizes a combined objective consisting of per-token language-modeling loss and per-image diffusion loss.A balancing coefficient λ weights the diffusion component of the objective.
  • The Transfusion objective fixes λ at 5 following preliminary experiments, while further tuning is left for future work.
  • During inference, the model alternates between autoregressive language modeling and diffusion sampling after generating a beginning-of-image token.It appends noise as image patches, denoises them iteratively, and appends an end-of-image token afterward.

B MoT Transfusion Fine-tuning Results

Fine-tuned Transfusion MoT and dense models are compared on text-prompted image generation. MoT produces more visually appealing or prompt-faithful images on some prompts, while both models struggle on harder prompts.

  • MoT produces images that are more visually appealing or faithful to prompts than the dense model for examples shown in Figure 20.The comparison concerns qualitative image-generation behavior under text prompts.
  • Both MoT and dense models successfully follow the prompts in the examples shown in Figure 21.
  • Both models struggle to perfectly follow text prompts and capture all details in the hard examples shown in Figure 22.The passage suggests that text faithfulness may improve with extended training.

C Supplementary Figures

The supplementary figures examine modality structure and MoT’s scaling behavior across Chameleon, Chameleon+Speech, and Transfusion settings. Across these evaluations, MoT shows modality-specific feature organization and favorable loss or generalization behavior relative to dense and MoE-4x models.

  • Latent feature space: PCA visualizations show text, speech, and image activations occupying distinct feature-space regions across training checkpoints and layers.The visualization covers 4%, 24%, 50%, and 100% checkpoints and Layers 1, 5, 17, and 32.
  • Chameleon: MoT delivers a substantial image-modality speedup over dense and MoE-4x models in Chameleon, with the advantage increasing across model scales.MoE-4x’s gains diminish with scale and disappear at 7B for image training.
  • Chameleon+Speech: MoT consistently improves image and text validation loss over model scales in the Chameleon+Speech setting under FLOPs-controlled training.The experiments use matching model scales and pretraining from scratch.
  • Transfusion: MoT matches dense-model training and validation loss for text in Transfusion and improves captioning generalization as measured by CIDEr.MoE-4x attains lower training losses but shows poorer text generalization than the dense model.
Loading 2411.04996v2…