MoE-LLaVA

Analysis

Why This Paper Matters

Large vision-language models (LVLMs) like LLaVA have achieved impressive results but come with high computational costs. MoE-LLaVA addresses this by introducing a sparse mixture-of-experts (MoE) architecture that activates only a subset of parameters per token, dramatically reducing inference cost while maintaining—or even exceeding—the performance of larger dense models. This is a practical step toward deploying capable multimodal AI on resource-constrained devices.

How It Works

MoE-LLaVA builds on the LLaVA framework but replaces the standard feed-forward network (FFN) in each transformer layer with an MoE layer. The MoE layer contains multiple expert FFNs and a learnable router. For each input token, the router computes a probability distribution over experts and activates only the top-k experts. The outputs of these experts are weighted by the router's softmax scores and summed.

The architecture processes an RGB image through a vision encoder to produce visual tokens, which are projected to the LLM's hidden dimension. Text tokens are embedded similarly. Both token types are concatenated and fed into the LLM with MoE layers.

The router is a simple linear layer: it computes logits ( f(x) = W \cdot x ), normalized by softmax. Each token is processed by the top-k experts with highest probabilities.

A key innovation is the three-stage MoE-Tuning strategy. Stage 1 freezes the LLM and trains only the visual projection MLP on image captioning to align visual tokens. Stage 2 trains the entire model on complex multimodal instructions (reasoning, text recognition). Stage 3 replicates the trained FFN to initialize experts, then fine-tunes the MoE layers with a loss combining auto-regressive language modeling and a differentiable load-balancing auxiliary loss to prevent expert collapse.

The load-balancing loss encourages uniform token assignment across experts. It is defined as the product of the fraction of tokens assigned to each expert and the average routing probability for that expert.

Training data is reorganized from existing datasets like MIMIC-IT and SViT for the three stages.

Results

MoE-LLaVA achieves strong results across multiple benchmarks. MoE-LLaVA-Phi-2.7B×4 (2.7B total parameters, 4 experts, top-2 active) surpasses LLaVA1.5-7B by 2.7% on ScienceQA with only 3.6B activated parameters. MoE-LLaVA-StableLM-1.6B×4 outperforms the 80B-parameter IDEFICS-80B with just 2.0B activated parameters. On MMBench, MoE-LLaVA-Qwen-1.8B×4 beats Qwen-VL-7B by 21.5% despite lower image resolution.

In object hallucination evaluation, MoE-LLaVA-1.8B×4 surpasses LLaVA-1.5-13B across all sampling methods (adversarial, popular, random), indicating better consistency in generating relevant objects.

Quantitative analysis of expert routing reveals that experts develop distinct roles: expert 3 dominates deeper layers, while experts 2 and 4 collaborate in shallow layers. Importantly, routing distributions for text and image tokens are highly similar, showing no modality bias and strong multimodal interaction.

Significance

MoE-LLaVA shows that sparse MoE architectures can match or exceed dense models in vision-language tasks while using far fewer activated parameters. This is a practical advance for deploying capable multimodal AI on edge devices. The MoE-Tuning strategy provides a recipe for adapting MoE to LVLMs without performance degradation. The analysis of expert specialization offers insights into how sparse models learn to divide multimodal tasks, which could guide future designs of efficient, scalable multimodal systems.