On Token's Dilemma: Dynamic MoE with Drift-Aware Token Assignment for Continual Learning of Large Vision Language Models¶
🕒 Published (v1): 2026-03-29 02:30 UTC · Source: Arxiv · Venue: CVPR 2026 · link
Why this paper was selected
CVPR 2026; drift-aware dynamic MoE token assignment for continual multimodal learning
Ask a follow-up
Open an assistant pre-loaded with this paper's context.
💬 Ask ChatGPT✦ Ask Claude
TL;DR¶
Continual instruction tuning of Large Vision-Language Models with Mixture-of-Experts (MoE) still suffers catastrophic forgetting even when old experts are frozen, because newly trained router parameters attract old-task tokens to new experts—a phenomenon the authors term routing-drift. This paper introduces LLaVA-DyMoE, which characterizes each token's routing ambiguity and applies Token Assignment Guidance (TAG) plus Routing Score Regularization (RSR) to steer problematic tokens away from new experts during training. On the CoIN benchmark the method achieves >7% gain in mean final accuracy and a 12% reduction in forgetting over baselines.
Problem¶
Dynamic MoE expansion for Multimodal Continual Instruction Tuning (MCIT) freezes old expert weights but still trains new experts and router parameters. The updated router distorts the assignment of old-task tokens, routing them to new experts that were never trained on those patterns—causing inference-time mis-routing and catastrophic forgetting. Prior remedies rely on task-level routing predictors that sacrifice token-level MoE flexibility or require heavy auxiliary computation for task identification.
Method¶
Token taxonomy via routing scores. For each token's representation \(h\) at layer \(l\), the router produces logits \(s_t = [s_{t-1}; s_{t,\text{new}}]\). Group-level confidence is extracted as \(c_{\text{old}} = \max(s_{t-1})\), \(c_{\text{new}} = \max(s_{t,\text{new}})\). Relative ambiguity is: $\(D_{\text{rel}} = \frac{|c_{\text{new}} - c_{\text{old}}|}{\max(|c_{\text{new}}|, |c_{\text{old}}|) + \epsilon}\)$ Tokens with \(D_{\text{rel}} \le \tau\) are ambiguous; those with \(c_{\text{new}} > c_{\text{old}}\) and \(D_{\text{rel}} > \tau\) are new; the remainder are old.
Token Assignment Guidance (TAG). A binary mask \(M_{\text{new}} = \mathbf{1}[(c_{\text{new}} > c_{\text{old}}) \wedge (D_{\text{rel}} > \tau)]\) gates each token to its expert group. Tokens failing either condition receive logit \(-\infty\) for new experts, forcing them to old frozen experts. This eliminates gradient pollution of the new router from ambiguous and old tokens without discarding those tokens entirely.
Routing Score Regularization (RSR). Two complementary soft losses act on the collective routing weights \(g_{\text{old}} = \sum_{i \in S_{t-1}} w_i\) and \(g_{\text{new}} = \sum_{i \in S_{t,\text{new}}} w_i\): - Exclusivity loss \(\mathcal{L}_{\text{exc}} = g_{\text{old}} \cdot g_{\text{new}}\): minimizing the product penalizes split routing across old and new groups. - Specialization loss \(\mathcal{L}_{\text{spe}} = -y \log g_{\text{new}} - (1-y)\log(1-g_{\text{new}})\) (BCE) with soft target \(y = 1 - \tilde{g}_{\text{old}}\): encourages new experts to specialize when old experts are inactive.
Total objective: \(\mathcal{L} = \mathcal{L}_{\text{NTP}} + \lambda \mathcal{L}_{\text{aux}} + \alpha(\mathcal{L}_{\text{exc}} + \mathcal{L}_{\text{spe}})\), where \(\mathcal{L}_{\text{NTP}}\) is the autoregressive instruction-tuning loss and \(\mathcal{L}_{\text{aux}}\) is standard MoE load-balancing applied only to new experts. TAG operates only during training; inference is unmodified.
Key Contributions¶
- Identification and analysis of the token's dilemma: ambiguous and old tokens in new-task data yield minimal new-task learning benefit yet induce routing-drift and forgetting when routed to new experts.
- Controlled two-task experiments isolating each token type's causal effect on forgetting and plasticity.
- LLaVA-DyMoE: a two-fold regularization (TAG + RSR) that addresses routing-drift at the token level, imposes no inference-time overhead, and is orthogonal to task-level routing methods and data-based rehearsal approaches.
Results¶
- On the CoIN benchmark (8 sequential VQA tasks: SQA, TextVQA, VQAv2, VizWiz, GQA, OCR, ImageNet, REF):
- >7% improvement in Mean Final Accuracy (MFN) over baseline methods.
- ~12% reduction in forgetting compared to baselines.
- Baselines compared: LoRA, MoELoRA, EWC, LwF, IncLoRA, O-LoRA, IncMoELoRA (all shown in Figure 1 radar chart; per-task and aggregate numbers available).
Limitations¶
- Requires a hyperparameter \(\tau\) to threshold ambiguity; sensitivity analysis is not detailed in the provided text.
- Evaluation is restricted to the CoIN benchmark with VQA-style tasks; generalization to more diverse MCIT settings (e.g., generation, grounding, embodied) is not demonstrated.
- The method assumes tasks arrive sequentially with clear boundaries; task-boundary-free or online streaming settings are not addressed.
- Combining with task-level routing methods for further gains is claimed to be compatible but not quantitatively validated in the excerpted text.
Relevance to Vision-Language Models¶
LLaVA-DyMoE directly addresses a critical bottleneck for deploying LVLMs in real-world settings where new instruction-following capabilities must be added post-deployment without costly full retraining. By operating at the token level within the MoE routing mechanism of LLaVA, it provides a principled and lightweight solution compatible with existing PEFT paradigms such as LoRA. The routing-drift analysis bridges the plasticity–stability dilemma from classical continual learning to the specific architecture of transformer-based VLMs, offering insights transferable to any MoE-based LVLM (e.g., Mixtral-style vision models). This work is particularly relevant for researchers exploring scalable, modular VLM adaptation in multi-task and lifelong-learning regimes.