PaperlayerView exact revision on arXiv

Source-linked AI summary

Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts

Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, Neil Houlsby

arXiv:2206.02770v1cs.CV

TL;DR

Existing large sparsely activated models are usually trained one modality at a time, motivating multimodal architectures that can process images and text together. The paper introduces LIMoE, a contrastively trained sparse MoE with entropy-based regularization, and finds improved performance over compute-matched dense models across scales, reaching 84.1% zero-shot ImageNet accuracy at H/14. The authors also identify multimodal routing failure modes and naturally emerging modality-specific experts.

  • Problem

    Large sparsely activated models are typically trained on a single modality, leaving multimodal sparse modeling as an open area of study.

  • Method

    LIMoE uses one sparse mixture-of-experts architecture for images and text, trained with contrastive learning and entropy-based regularization.

  • Results

    LIMoE significantly outperforms compute-matched dense baselines across scales and reaches 84.1% zero-shot ImageNet accuracy with LIMoE-H/14.

  • Takeaways & Limitations

    Sparse multimodal MoEs can achieve strong performance across model scales while exhibiting modality-specific expert specialization.

  • Takeaways & Limitations

    Routing interference between multiple modalities remains not fully understood, and LIMoE does not explicitly address potential harms associated with large-scale models, contrastive models, and web-scale multimodal data.

Abstract

from arXiv · show

Large sparsely-activated models have obtained excellent performance in multiple domains. However, such models are typically trained on a single modality at a time. We present the Language-Image MoE, LIMoE, a sparse mixture of experts model capable of multimodal learning. LIMoE accepts both images and text simultaneously, while being trained using a contrastive loss. MoEs are a natural fit for a multimodal backbone, since expert layers can learn an appropriate partitioning of modalities. However, new challenges arise; in particular, training stability and balanced expert utilization, for which we propose an entropy-based regularization scheme. Across multiple scales, we demonstrate remarkable performance improvement over dense models of equivalent computational cost. LIMoE-L/16 trained comparably to CLIP-L/14 achieves 78.6% zero-shot ImageNet accuracy (vs. 76.2%), and when further scaled to H/14 (with additional data) it achieves 84.1%, comparable to state-of-the-art methods which use larger custom per-modality backbones and pre-training schemes. We analyse the quantitative and qualitative behavior of LIMoE, and demonstrate phenomena such as differing treatment of the modalities and the organic emergence of modality-specific experts.

1 Introduction

LIMoE applies sparse mixture-of-experts modeling to a single architecture that jointly processes images and text through contrastive learning. Entropy-based regularization addresses multimodal routing challenges, while LIMoE outperforms compute-matched dense baselines across scales.

  • LIMoE is presented as the first large-scale multimodal mixture-of-experts model for jointly processing images and text.The architecture aligns image and text representations using contrastive learning.
  • Entropy-based regularization stabilizes training and improves expert utilization for multimodal learning.The approach responds to failure modes that arise when multiple modalities are fed into one MoE architecture.
  • LIMoE models significantly outperform compute-matched dense baselines across a range of model scales.Relative improvements in zero-shot ImageNet accuracy range from 7% to 13% over equivalent dense models.
  • 84.1% zero-shot ImageNet accuracy is achieved by LIMoE-H/14, comparable to state-of-the-art contrastive models using per-modality backbones and pre-training.The comparison uses models that apply 3–4× more parameters per token.
  • The paper analyzes LIMoE’s quantitative and qualitative behavior and its design decisions through ablations.

2 Multimodal Mixture of Experts

LIMoE uses a shared one-tower Transformer with sparse, input-dependent experts to process images and text under contrastive learning. Multimodal routing creates modality-imbalance and expert-collapse failure modes, addressed with per-modality entropy regularization and routing stabilization.

  • Architecture: LIMoE aligns image and text representations using a single shared architecture trained with a contrastive objective.Modality-specific input projections feed tokens into a shared Transformer, whose final representations are average-pooled.
  • Architecture: Sparse MoE layers route each token to a small number of input-dependent MLP experts using learned gating weights.Each token is routed among E experts, with the activated experts’ outputs combined according to their gate weights.
  • Routing challenges: Fixed expert capacity causes token dropping when routing collapses, reducing success rates and often harming training stability.Top-K routing can collapse onto one expert, leaving many tokens unprocessed.
  • Routing challenges: Multimodal routing introduces failure modes because image tokens outnumber text tokens by 3–17×, allowing minority-modality tokens to overload one expert.Global routing statistics can still appear balanced when majority-modality tokens are distributed across experts.
  • Entropy regularization: LIMoE combines classic auxiliary losses with local and global per-modality entropy losses to stabilize routing and prevent modality-specific collapse.The local loss acts per token, while the global loss acts after marginalizing over tokens; together they correspond to negative modality-conditioned mutual information.
  • Routing stabilization: Batch Priority Routing prioritizes tokens with larger routing weights when expert buffers fill, while avoiding systematic preference for either modality.This addresses inevitable token dropping under Top-K routing and is designed to favor informative tokens.

3 Experiments

Experiments compare LIMoE with compute-matched dense one-tower baselines across scales and then evaluate a larger model. LIMoE consistently improves performance, reaching 84.1% zero-shot ImageNet accuracy with a shared backbone and sparse per-token computation.

  • Controlled study across scales: LIMoE’s compute-performance frontier dominates dense models across scales from S/32 through L/16.The largest average gains are 10.1% for zero-shot ImageNet classification and 12.2% for 10-shot ImageNet classification.
  • Controlled study across scales: 10.1% and 12.2% average absolute improvements are reported for zero-shot and 10-shot ImageNet classification, respectively.COCO text-to-image retrieval gains are strongest at small scales and remain significant but more modest at larger scales.
  • Scaling up LIMoE: 84.1% zero-shot ImageNet classification accuracy is achieved by LIMoE-H/14, a 5.6B-parameter model applying 675M parameters per token.The large experiment increases architecture size, training duration, and data size, including JFT-4B images in 25% of each batch.
  • Scaling up LIMoE: LIMoE-H/14 is fully trained from scratch without pretrained components and is presented as the first competitive one-tower model with a shared backbone.The table compares it with state-of-the-art zero-shot classification models, including two-tower systems.
  • Scaling up LIMoE: LIMoE-H/14 is comparable with the best two-tower models despite its relatively modest architecture scale.The paper attributes the result to a modality-agnostic shared-backbone approach with specialist components emerging within the model.

4 Ablations

The ablations examine LIMoE’s baselines, entropy-based routing losses, modality balancing, stabilization, and scaling choices. Across these studies, entropy routing improves robustness and performance, while BPR and other architectural choices further affect model quality.

  • Baselines: LIMoE greatly outperforms dense baselines on ImageNet 0-shot and 10-shot evaluation, while retrieval confidence intervals overlap with two-tower models.The two-tower baseline is twice as large and expensive but still falls behind the sparse model.
  • Auxiliary losses: Both entropy losses are important for text, whereas the global image loss is not impactful and the local image loss is harmful.The final loss combination also considered validation accuracy, training stability, and routing success rate.
  • Entropy thresholds: The text threshold typically behaves as a soft minimum using approximately S = e^τtext text experts, while performance is robust provided τtext is not too low.The image threshold is usually inactive and can likely be removed in future work.
  • Alternative losses: The mutual-information auxiliary loss performs slightly worse than the entropy-loss setup, with 1.5% lower zero-shot and 0.1% lower 10-shot performance.The mutual-information formulation combines local and global entropy losses without threshold parameters.
  • Modality balancing: Entropy routing remains superior to the classical setup even with balanced 1:1 image:text sequences, while longer image sequences improve entropy-routed performance.The result indicates robustness across different modality balances rather than only correction of imbalance.
  • Routing stabilization: BPR improves training stability and mitigates token dropping, whereas FIFO routing performs extremely poorly regardless of whether images or text are routed first.Random token shuffling only partially ameliorates the problem.
  • Router structure: Separate modality-specific expert pools are more stable and need no auxiliary-loss regularization, but they are harder to scale to many modalities and tasks.Per-modality routers and experts otherwise perform comparably to the generic modality-agnostic setup.
  • Scaling ablations: Increasing selected experts per token raises zero-shot accuracy from 55.5% with K = 1 to 61.0% with K = 5.The study modifies BPR and the local auxiliary loss to support K > 1.

5 Model Analysis

The analysis finds that LIMoE’s experts specialize without explicit modality assignments. Both modality-specific and multimodal experts emerge, alongside semantically meaningful visual specializations and modality-dependent routing priorities.

  • Multimodal experts: Modality-specific and multimodal experts emerge organically despite no explicit specialization constraint.The observation holds across the analyzed LIMoE configurations.
  • Qualitative analysis: Visual experts specialize in features ranging from colors and lines to faces and text.These specializations indicate semantically meaningful internal organization.
  • Routing behavior: BPR does not always prioritize text: in later layers, it often ranks important image patches ahead of text.The local loss encourages high max-routing weights for text, which determines BPR’s ranking signal.

6 Related work

Related work spans unimodal Transformer models, multimodal architectures that fuse or share representations, sparse mixture-of-experts models, and multimodal contrastive learning. LIMoE occupies the one-tower, shared-backbone setting while building on sparse routing and image-text contrastive training.

  • Multimodal modeling: Multimodal models process multiple data types using a single neural network, while unimodal NLP and vision research increasingly converges on Transformer architectures.The paper distinguishes multimodal modeling from long-standing unimodal modeling.
  • Multimodal modeling: Multimodal approaches either fuse modalities, share weights across distinct tasks without alignment, or combine unimodal and fused components.LIMoE is closer to approaches that retain unimodal feature-extractor behavior.
  • Sparse mixture of experts: Sparse MoE models use learned gating to activate K of E ≫ K experts for each input, scaling parameters while limiting computation.Prior work studied these models independently in computer vision and NLP.
  • Contrastive learning: Multimodal contrastive learning trains on aligned data and supports image classification and cross-modal retrieval without downstream fine-tuning.The approach was scaled from medical image-report data to noisy web data and later combined with larger models or multitask training.

7 Conclusions and Future Work

The paper concludes that LIMoE is a scalable multimodal sparse MoE with strong performance and expert specialization. It also identifies unresolved routing interactions, societal and environmental concerns, and open questions for broader multimodal scaling.

  • Conclusions: LIMoE improves FLOP-matched dense baselines by an average of +10.2% zero-shot accuracy across model scales.The large H/14 model reaches 84.1% accuracy and is competitive with current state-of-the-art approaches.
  • Conclusions: LIMoE’s entropy-based auxiliary losses stabilize training and produce performant multimodal sparse MoE models.The work presents LIMoE as the first multimodal sparse mixture-of-experts model.
  • Limitations: LIMoE does not explicitly address potential harms associated with large-scale models, contrastive models, and web-scale multimodal data.The authors also note that pruning can cause low-resource groups to be forgotten, while specialization might improve underrepresented-group performance by analogy.
  • Limitations: Training large models is environmentally costly, although MoEs are reported to be significantly more efficient during inference.The authors identify LIMoE as a candidate for efficient, large-scale multimodal foundation models.
  • Future work: Routing interference across modalities remains incompletely understood, and conclusions from NLP MoEs do not transfer perfectly to vision or multimodal settings.The paper identifies differing behavior between image and text routing as a future-work issue.
  • Future work: Extending sparse models to more modalities is an open direction because interactions between data types and routing algorithms become more difficult with additional modalities.The authors frame sparse models as a natural way to scale while handling different tasks and data.

A Training details

LIMoE training combines multimodal routing with auxiliary losses designed to balance expert use, control router behavior, and stabilize optimization. The largest model also uses checkpoint souping to select a final model from multiple cooldown trajectories.

  • Training configuration: Training uses Adafactor, cosine decay with linear warmup, and TPU-v2, v3, and v4 hardware.The default learning rate is 1 × 10^-3 with decoupled weight decay of 1 × 10^-5.
  • Training configuration: Models use 32 experts placed every 2 layers, with architecture parameters otherwise following ViT.The final representation is cast to bfloat16 for reduced all-to-all costs and improved memory efficiency.
  • Checkpoint souping: The largest model uses checkpoint souping, greedily selecting checkpoints from multiple cooldowns to maximize zero-shot ImageNet validation accuracy.The final soup contains 8 checkpoints.
  • Auxiliary losses: Auxiliary losses include importance and load losses for balancing gating weights and assignments, plus z-loss for controlling router activation magnitudes.The load loss estimates assignment probabilities under noisy routing because expert assignments are discrete.
  • Auxiliary losses: The mutual-information loss directly regularizes dependence between experts and modalities.It is presented as a variant of local and global entropy losses.
  • Auxiliary losses: The final auxiliary loss averages the component losses using a shared regularization weight set to 0.04 in all experiments.The authors report this value was robust after preliminary tuning.

B.2 In-depth analysis of global entropy threshold

The global entropy threshold acts as a soft minimum on modality-specific expert use. Text thresholds affect stability and specialization, whereas image thresholds generally have little systematic effect because image routing entropy is already high.

  • Threshold interpretation: A threshold τ corresponds to a soft minimum S_T = e^τ on the number of experts a modality is encouraged to use.The loss is satisfied when the modality uses at least this soft minimum of experts.
  • Performance: Low text thresholds produce unstable models with poor performance, while performance becomes fairly consistent beyond a threshold limit.The image threshold shows no systematic performance effect inside or outside the unstable region.
  • Routing entropy: Image routing entropy remains high, so the global entropy loss is usually inactive for images.Text entropy instead closely tracks its target threshold.
  • Expert specialization: Increasing the text threshold yields more text experts and fewer image experts.The number of experts needed to reach a routing success rate serves as a proxy for modality-specific expert count.
  • Open questions: The differing text and image behavior remains unresolved, with imbalance during training and fundamental routing differences identified as possible explanations.The passage frames these as open questions rather than established causes.

C.1 Scaling comparison, and architecture definitions.

The scaling study evaluates LIMoE across model sizes and routing configurations while examining expert capacity and modality imbalance. Multimodal routing can create modality-specific experts, causing token dropping when only one modality is processed.

  • Scaling comparison: LIMoE scaling results and the LIMoE-H/14 result are reported together from the Figure 4 study.Table 4 contains the corresponding results.
  • Ablations: Scaling comparisons include dense baselines, routing-order alternatives, modality balancing, and different expert configurations.The supplied figure and table references cover these ablation dimensions.
  • Capacity: Expert capacity assumes tokens are balanced across experts, with slack factor C_R ≥ 1 providing additional capacity.Tokens assigned beyond capacity are not processed.
  • Modality imbalance: When only one modality is present, modality-specific expert use can leave other experts unused and increase token dropping.The capacity was set assuming all experts would be used during multimodal training.
  • Evaluation: Evaluation uses C_R = 16 to reduce token dropping when image and text inputs are processed separately.Processing both modalities simultaneously or pruning modality-specific experts are described as alternative solutions.

D.2 Understanding the compute costs of LIMoE

LIMoE’s compute profile is favorable under total-FLOPs comparisons, but communication and downstream batching assumptions matter. Routing ablations show that auxiliary-loss and dispatch changes are needed for stable multi-expert operation, while modality-aware designs and pruning offer additional options.

  • Evaluation assumptions: Compute comparisons assume full batches of images and texts, which may disadvantage or advantage models differently when downstream modality volumes are highly uneven.The authors note that text-heavy or image-heavy workloads can change the relative cost profile.
  • Compute costs: LIMoE is Pareto optimal with respect to total FLOPs, although this proxy omits some MoE communication costs.The profiling uses TPUv3 measurements and compares multiple compute-cost proxies.
  • Selected experts: Using more than one selected expert per token requires modifying local entropy targets and the BPR score to preserve stability.The modified BPR prioritizes tokens by the sum of top-K probabilities rather than the maximum probability.
  • Expert pool size: Larger expert pools consistently improve LIMoE performance.The study varies the total number of experts while keeping k = 1.
  • Router design: Disjoint modality-specific experts achieve comparable performance, are generally more stable, and do not require entropy regularization, but may be less general for many modalities or tasks.Per-modality routers and joint routers are the comparison alternatives.
  • Pruning: Router-pred pruning performs significantly better than router-drop when many experts are removed, while both methods perform similarly after limited pruning.The comparison uses ImageNet-10shot linear few-shot evaluation.
  • Pruning: Text evaluation can prune to 4 experts without significant performance loss, whereas retaining more image experts improves performance.Image and text experts can be pruned independently for separate-modality evaluation.
  • Grouped routing: Grouped routing remains effective through 4 groups, but more groups destabilize training; 4 groups reduce step time by 30%.The setup processes 128 examples per device and splits 27136 tokens when group size is 1.

E.6 Experiments on public data

On public LAION-400M data, LIMoE B/16 models outperform a dense baseline, with the largest gains reported for ImageNet zero-shot and 10-shot evaluation. Routing analyses show balanced training-distribution behavior alongside modality-dependent trajectories and priority patterns under distribution shift.

  • E.6 Experiments on public data: +5.0% absolute ImageNet zero-shot performance was observed for LIMoE B/16 models relative to the dense baseline.The corresponding relative improvement was +8.9%.
  • E.6 Experiments on public data: +6.6% absolute ImageNet 10-shot performance was observed for LIMoE B/16 models relative to the dense baseline.The corresponding relative improvement was +13.8%.
  • Routing Distributions: Training-distribution routing works very well in both B/32 and B/16 models, with most experts handling both image and text tokens.Figures 17 and 18 visualize these routing distributions for models with eight experts.
  • Routing Distributions: Training examples show very little token dropping because auxiliary losses lead to balanced routing.This observation is reported for both B/32 and B/16 training-data distributions.
  • Token Trajectories: Image-token trajectories are fairly interconnected across the first two or three MoE layers in both B/32 and B/16 models.The reported trajectories compare image tokens in red with text tokens in blue.
  • Priority Distributions: The local entropy loss concentrates text routing predictions, while BPR prioritizes tokens with high pmax; the resulting buffer sorting is examined on training and shifted data.The supplied passage frames these plots as evidence about how tokens from both modalities are sorted into buffers.

G LIMoE-H/14 Analysis

LIMoE-H/14 develops increasingly confident image routing and shows qualitative specialization among image and text experts. Because image tokens dominate, only a small subset of experts handles text, with some evidence of semantic specialization and padding-focused experts.

  • LIMoE-H/14 Analysis: Image experts naturally specialize in semantic concepts including body parts, textures, fauna, food, and doors.These specializations are described qualitatively for the largest model.
  • LIMoE-H/14 Analysis: Image-token routing becomes increasingly confident across LIMoE-H/14 MoE layers, while most text predictions are close to one-hot under the local entropy loss.Figure 38 summarizes the per-token maximum routing probability for each modality.
  • LIMoE-H/14 Analysis: Text-routing analysis is constrained by SentencePiece subword tokenization, since subwords from a word are not necessarily handled by the same expert.This prevents perfect parsing of every token processed by each expert.
  • LIMoE-H/14 Analysis: Some LIMoE-H/14 experts process mixed token types, while others show semantic specialization or mostly handle padding tokens.One cited example is an expert in Layer 14 that focuses on nouns and adjectives.
  • LIMoE-H/14 Analysis: With a text:image token ratio near 1:27, only 3–4 experts handle text in the LIMoE-H/14 setting.Figure 39 reports token counts per modality and expert across MoE layers.
Loading 2206.02770v1…