Solving Token Gradient Conflict in Mixture-of-Experts for Large Vision-Language Model¶

🕒 Published (v1): 2025-01-01 · Source: ICLR · Venue: ICLR 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

Existing Mixture-of-Experts (MoE) routers for Large Vision-Language Models (LVLMs) ignore gradient-level interference between tokens co-assigned to the same expert, causing sub-optimal training. STGC identifies "conflicting tokens" — those whose per-token gradients have negative cosine similarity to their expert's average gradient — and adds a conflict elimination loss to reroute them away. The method is a drop-in regularization plug-in compatible with any MoE-based LVLM.

Problem¶

MoE routers in LVLMs (e.g., MoE-LLaVA) dispatch tokens based on learned routing scores, but offer no mechanism to prevent tokens with opposing parameter-update directions from being co-assigned to the same expert. Prior approaches address data interference at the sample level (via instruction embeddings or task labels), which cannot resolve intra-sample token-level conflicts (e.g., visual vs. text tokens within the same example pulling an expert's weights in opposing directions).

Method¶

Conflicting Token Identification: After a forward pass, STGC performs a backward pass through only the expert layers (no parameter update) to compute per-token gradients $g_n = \nabla_{\theta_{e_i}} L_n$. It computes the mean expert gradient $g_{\text{mean}}$ and flags token $t_n$ as conflicting when $\cos(g_n, g_{\text{mean}}) < \tau$ (default $\tau = 0$). Partial-parameter gradients are used as a proxy to reduce GPU memory overhead.

Conflict Elimination Loss (LCEL): For each identified conflicting token $t_n$ currently in expert $e_{id}$, the router logits are inverted ($z' = -z$), softmax is recomputed, and a cross-entropy loss penalizes high routing probability to the current expert, pushing the router toward other experts: $$\mathcal{L}_{\text{CEL}} = \frac{1}{N_{\text{all}} \cdot E} \sum_n \sum_i \log(p'_{\text{moe}}(t_n)_i) \cdot q_{\text{moe}}(t_n)_i$$

Total loss: $\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{main}} + \alpha \mathcal{L}_{\text{aux}} + \beta \mathcal{L}_{\text{CEL}}$, where $\mathcal{L}_{\text{aux}}$ is the standard load-balancing loss. Applied only during the instruction-tuning stage.

Key Contributions¶

Token-level gradient-based definition of conflicting tokens, replacing embedding- or task-label-based sample-level interference proxies.
Conflict Elimination Loss (LCEL) that optimizes the router to reroute conflicting tokens without updating expert weights during identification.
Plug-in design: STGC integrates into any MoE-LVLM without architectural changes.
Statistical verification via "gradient consistency" metric (mean pairwise cosine similarity of token gradients within an expert), showing STGC raises this metric across all layers during training.

Results¶

All comparisons use MoE-LLaVA as the primary baseline (4Top2 configuration, activating 3.6B parameters with Phi-2 backbone + SigLIP encoder).

vs. MoE-LLaVA† (3.6B act., 665K data): STGC improves average benchmark score from 62.9 → 63.5; VisWiz +1.4%, POPE +0.4%, MME +16.3, MMB +1.7%.
vs. MoE-LLaVA† (3.6B act., 1021K data): Average 63.8 → 64.9 (+1.1%); VisWiz +5.9%, POPE +2.2%, MMB +0.6%; VQAv2 80.0%, GQA 63.0%.
vs. LLaVA-1.5 7B (dense, 6.7B act.): With only 3.6B activated parameters, STGC surpasses on VQAv2 (+1.5%), GQA (+1.0%), MMB (+7.6%), MM-Vet (+10.2%).
Plug-in ablation (Table 2): Consistent gains across 4Top1 (S-1.6B), 4Top2 (S-1.6B), 4Top2 (P-2.7B), and 4Top2† (P-2.7B) configurations.
vs. alternative routing strategies (Table 4): Embedding-based (+57.3 avg), feature-based (+57.2 avg), task-based (+48.9 avg) all underperform STGC (+58.0 avg) over the MoE-LLaVA baseline (57.3 avg).
Scalability: Performance gain grows from +0.6% avg at 665K data to +1.1% avg at 1021K data, consistent with the hypothesis that larger, more diverse data amplifies gradient conflict.

Limitations¶

During training, each token only passes one expert per MoE layer per forward pass; its gradient therefore only reflects conflict relative to that expert's average gradient, not relative to all other experts. Identifying the globally optimal expert for a conflicting token is not addressed.
The conflict elimination loss can only push tokens away from their current expert, not deterministically toward the best alternative expert.
Gradient storage overhead per expert layer necessitates using partial parameter gradients as a proxy rather than full gradients.
Experiments are confined to MoE-LLaVA variants; generalization to other MoE-LVLM architectures (e.g., DeepSeekMoE-based VLMs) is not empirically validated.

Relevance to Vision-Language Models¶

STGC directly addresses a training inefficiency specific to MoE-based LVLMs that arises from the heterogeneity of visual and textual tokens — a problem unique to the multimodal setting where sample-level routing heuristics are insufficient. For researchers tracking VLMs, this work establishes that gradient-level analysis of token routing is a more faithful proxy for data interference than embedding similarity or task identity, with direct implications for how MoE capacity should be managed in models that process mixed modalities. The plug-in design and consistent gains across scales make STGC a practically relevant contribution to efficient LVLM training, connecting to the broader effort to match dense-model quality at sparse-model inference cost.