Source-linked AI summary
On Token's Dilemma: Dynamic MoE with Drift-Aware Token Assignment for Continual Learning of Large Vision Language Models
Chongyang Zhao, Mingsong Li, Haodong Lu, Dong Gong
TL;DR
MoE-based continual instruction tuning can forget prior knowledge when newly added routers redirect old-task tokens to new experts. LLaVA-DyMoE analyzes this token-level routing-drift and applies drift-aware assignment with complementary routing regularization, achieving over a 7% MFN gain and a 12% reduction in forgetting versus baselines.
Problem
Dynamic MoE expansion still suffers forgetting because newly trained routers can attract old-task tokens to new experts, corrupting established routing for prior tasks.
Method
LLaVA-DyMoE combines Token Assignment Guidance, which redirects ambiguous and old tokens away from new experts, with routing regularization that separates expert groups and promotes new-expert specialization.
Results
Over a 7% gain in mean final accuracy and a 12% reduction in forgetting were achieved compared with baseline methods on the CoIN benchmark.
Takeaways & Limitations
LLaVA-DyMoE mitigates routing-drift-induced forgetting and remains complementary to existing MCIT paradigms for further gains.
Abstract
from arXiv · showhide
Multimodal Continual Instruction Tuning aims to continually enhance Large Vision Language Models (LVLMs) by learning from new data without forgetting previously acquired knowledge. Mixture of Experts (MoE) architectures naturally facilitate this by incrementally adding new experts and expanding routers while keeping the existing ones frozen. However, despite expert isolation, MoE-based continual learners still suffer from forgetting due to routing-drift: old-task tokens become mistakenly attracted to newly added experts, degrading performance on prior tasks. We analyze the failure mode at the token level and reveal the token's dilemma: ambiguous and old tokens in new-task data offer minimal learning benefit yet induce forgetting when routed to new experts, due to their ambiguous routing assignment during training. Motivated by this, we propose LLaVA-DyMoE, a dynamic MoE framework that incrementally expands the MoE with drift-aware token assignment. We characterize token types via their routing score distributions and apply targeted regularization. Specifically, a token-level assignment guidance steers ambiguous and old tokens away from new experts to preserve established routing patterns and alleviate routing-drift, while complementary routing score regularizations enforce expert-group separation and promote new-expert specialization. Extensive experiments demonstrate that our LLaVA-DyMoE effectively mitigates routing-drift-induced forgetting, achieving over a 7% gain in mean final accuracy and a 12% reduction in forgetting compared to baselines. The project page is https://zhaoc5.github.io/DyMoE.
1. Introduction
LVLMs need continual instruction tuning to absorb new task knowledge while retaining prior capabilities. LLaVA-DyMoE addresses routing-drift in expanding MoE models through token-aware assignment and complementary routing regularization.
- Continual instruction tuning adapts LVLMs to emerging tasks while preserving performance on previously learned tasks.
- MoE architectures support continual learning through modular expert allocation and parameter isolation, but existing approaches can still suffer forgetting from routing instability.Task-specific routing may require heavy computation and become unreliable for diverse, complex tasks.
- Routing-drift occurs when newly trained router parameters attract old-task tokens to new experts, corrupting established routing policies despite frozen old experts.
- LLaVA-DyMoE combines Token Assignment Guidance with Routing Score Regularization to redirect ambiguous tokens away from new experts and promote expert-group separation.The framework targets token-level assignment while encouraging new-expert specialization.
- Over 7% MFN gain and 12% forgetting reduction were achieved compared with baseline methods on the CoIN benchmark.
- LLaVA-DyMoE is compatible with data-based and task-specific-routing MCIT paradigms for further performance gains.
2. Related Work
Related work in continual vision-language learning addresses catastrophic forgetting through continual instruction tuning, modular components, and parameter or data-based protection strategies. MoELoRA introduced the CoIN benchmark, while later methods expanded components or regularized updates to retain prior knowledge.
- Continual learning methods train models on non-stationary task sequences while targeting catastrophic-forgetting mitigation.
- Continual instruction tuning for LVLMs avoids repeatedly retraining models from scratch as new vision-language tasks arrive.
- MoELoRA introduced the CoIN benchmark and used a Mixture-of-Experts framework with LoRA experts for continual vision-language learning.
- SEFE and ProgLoRA incrementally add components while using regularization or progressive strategies to retain prior knowledge during continual updates.
3. Methodology: LLaVA-DyMoE with Drift-Aware Token Assignment
LLaVA-DyMoE incrementally expands a LoRA-based MoE while using token-level routing analysis and two-fold regularization to reduce routing drift during continual LVLM instruction tuning.
- Dynamic MoE Layers: LLaVA-DyMoE uses LoRA modules as MoE experts within LLaVA’s language backbone, with a router assigning each token to its top-K experts.Experts are incrementally added for new tasks while previously existing parameters remain frozen.
- Dynamic MoE Layers: Dynamic expansion freezes old experts and trains newly added experts and router parameters, but old-task tokens can be mis-routed to new experts and cause forgetting.The expanded router allows new-task tokens to reach both old frozen and new trainable expert groups.
- Token’s Dilemma: Controlled experiments categorize new-task tokens as new, old, or ambiguous according to the relative dominance of new- versus old-group routing scores.The analysis evaluates how each token type affects new-task learning and old-task forgetting.
- Token’s Dilemma: Ambiguous and old tokens provide little new-task benefit yet can induce forgetting when routed to new experts, exposing a token-level assignment dilemma.Ambiguous tokens are difficult to identify because their affinities to the two expert groups are similar.
- LLaVA-DyMoE Regularization: The framework combines token-routing guidance with complementary regularization that promotes exclusive group routing and new-expert specialization during training.The approach targets routing drift using token routing scores and is designed to work alongside other continual-learning methods.
- Token Assignment Guidance: Token Assignment Guidance identifies token types from routing-score ambiguity and sends tokens to the new group only when they are non-ambiguous and new-dominant.The relative score difference Drel and threshold τ determine ambiguity, while group-wise confidence determines dominance.
4. Experiments
Experiments evaluate LLaVA-DyMoE on CoIN’s eight VQA tasks and show that token-level routing regularization improves accuracy while reducing forgetting. Ablations and complementary-data studies support the method’s component contributions and compatibility.
- Experimental Setup: The evaluation uses CoIN’s eight VQA tasks, with 569k training samples and 261k test samples.
- Main Results: Static continual-training methods show significant forgetting and inferior accuracy because prior knowledge is repeatedly overwritten.
- Main Results: IncMoELoRA reduces overwriting through per-task expert expansion but suffers ambiguous token-expert assignments that activate newly added experts.
- Main Results: 7.35%, 8.20%, and 12.00% improvements are reported on MFN, MAA, and BWT, respectively, over the compared methods.
- Ablation Study: TAG prevents ambiguous tokens from contributing to new-task learning, while Lexc and Lspe further improve separation and new-task specialization.
- Complementary Techniques: LLaVA-DyMoE with a small replay buffer achieves a BWT of -1.55%, compared with ProgLoRA’s -6.59%, while attaining higher mean final accuracy.
5. Conclusion
The conclusion attributes forgetting in dynamic MoE expansion to token-level routing-drift and presents LLaVA-DyMoE as a drift-aware remedy. It also identifies scalability to larger models and more realistic scenarios as future work.
- Ambiguous and old tokens in new-task data offer minimal learning value yet corrupt old-task routing when left unguided.
- LLaVA-DyMoE combines Token Assignment Guidance with routing losses that enforce expert-group separation and mitigate routing-drift.
- The method is complementary to existing MCIT paradigms because it targets the routing mechanism.
- Future work should investigate scalability on larger-scale models and in more realistic scenarios.
A.1. More Implementation Details
The implementation uses LLaVA-v1.5 with frozen existing components and adds trainable LoRA experts as new tasks arrive. The default configuration specifies the expert count, rank, and routing width.
- The backbone is pre-trained, instruction-untuned LLaVA-v1.5, using Vicuna and CLIP ViT-L/14 for language and visual features.
- Only newly added modules are trainable during continual learning; all other components remain frozen.
- The default setting adds 16 rank-4 LoRA experts for each new task and uses K = 16 for top-K routing.
A.2. Details of Datasets
The CoIN benchmark contains eight VQA tasks spanning varied sample sizes, styles, and domains. The supplied dataset details include task-specific modality and evaluation information, while Table 6 addresses varying task orders.
- Benchmark Composition: CoIN includes ScienceQA, TextVQA, ImageNet, GQA, VizWiz, RefCOCO, VQAv2, and OCR-VQA.
- Task Details: ScienceQA evaluates reasoning over integrated visual and textual information using image-text and text-only samples.
- Task Details: TextVQA focuses on text recognition in visual question answering, while VizWiz targets assistive visual question answering for visually impaired users.
- Evaluation Variations: Table 6 reports performance under varying task orders.
B.1. Experiments on Different Task Orders
LLaVA-DyMoE remains stable across task orderings and instruction-template variants, while absorbing new knowledge without assuming a fixed incoming order.
- Task-order robustness: LLaVA-DyMoE exhibits excellent stability across Origin, Reverse, and Alphabet task orderings, with forgetting remaining consistently low.
- Task-order robustness: The method absorbs effective new knowledge into new routers and experts without assumptions about incoming data order.
- Instruction-template robustness: Experiments evaluate performance under Origin, Diverse, and 10Type instruction templates from the CoIN benchmark.
B.3. Ablation Studies on MoE Configurations
Ablations show that LLaVA-DyMoE is robust across MoE configurations, consistently outperforming the baseline across top-K, expert number, and expert capacity choices.
- Top-K and expert number: LLaVA-DyMoE consistently outperforms IncMoELoRA across different top-K values and numbers of experts.
- Expert capacity: Across expert capacities with LoRA ranks 1, 2, and 4, LLaVA-DyMoE consistently outperforms the baseline.
- Overall robustness: Overall, the method remains robust to variations in top-K, expert number, and expert capacity.
B.4. Experiments with Different Backbone Sizes
Additional experiments examine scalability, routing compatibility, replay compatibility, pruning, and expert activation patterns, with results supporting broad robustness of LLaVA-DyMoE.
- Backbone scaling: With a larger 13B LLaVA backbone, LLaVA-DyMoE reaches 60.39% MFN while maintaining a forgetting rate of -4.64%.
- Replay compatibility: The method achieves competitive or better performance than ProgLoRA across replay buffer sizes of 200, 500, and 1000.
- Task-level routing compatibility: LLaVA-DyMoE provides additive gains when integrated with task-specific routing methods.
- Expert pruning: LLaVA-DyMoE remains robust after pruning one-eighth or one-quarter of the lowest-frequency experts, though naive pruning slightly reduces performance.
- Expert activation: Across layers, experts show varying activation strengths and diverse utilization patterns on the eight CoIN tasks.
B.9. Qualitative Result Examples
Qualitative comparisons show that LLaVA-DyMoE retains finer-grained prior-task knowledge than IncMoELoRA after final-task training, while both models remain challenged by complex scenes.
- Baseline errors: IncMoELoRA may regress to coarse or incorrect labels, such as predicting “Dog” for “Bernese mountain dog” or “Otter” for “Sloth bear.”
- Qualitative comparisons: LLaVA-DyMoE retains fine-grained knowledge that IncMoELoRA often forgets on ScienceQA and ImageNet examples.
- Remaining challenges: Both models remain challenged by complex scenes containing small objects, including an ambiguous Mouse case.
- Efficiency: LLaVA-DyMoE incurs only a 4.4% training-time overhead, increasing training time from 74.18 to 77.43 minutes, while inference efficiency remains unaffected.
- Efficiency and ethics: The paper emphasizes parameter and inference efficiency through a sparse MoE architecture, while noting that safety alignment and fairness remain future considerations.