Abstract
Lifelong Concept Learning (LCL) in text-to-image models faces the dual challenges of catastrophic forgetting and unsustainable parameter growth. We present $R^2MoE$ (Redundancy-Removal Mixture of Experts), a framework that utilizes a mixture-of-experts architecture coupled with routing distillation and layer-wise expert elimination. Our approach achieves an 87.8% reduction in forgetting and uses 63.3% fewer parameters than existing state-of-the-art methods, enabling efficient and scalable personalization for generative models.
Introduction
Large-scale generative models like Stable Diffusion have revolutionized image synthesis, yet personalizing these models for specific user-defined concepts remains a significant challenge. This is particularly true in Lifelong Concept Learning (LCL) scenarios, where a model must sequentially learn an endless stream of new visual concepts. Existing methods typically suffer from two primary issues: Catastrophic Forgetting, where new learning overwrites previous knowledge, and Parameter Expansion, where the model size grows linearly with the number of concepts, eventually exceeding the capacity of consumer-grade hardware.
Methodology
The $R^2MoE$ framework addresses these challenges by treating the model as a collection of specialized "experts" implemented as Low-Rank Adaptation (LoRA) modules. These experts are managed by a gating network that dynamically routes prompts to the appropriate module.
The framework introduces three core innovations:
- Routing Distillation: To combat forgetting, the authors identified that the gating network often loses the "map" of which expert corresponds to which concept. By employing knowledge distillation, the new gating network is trained to retain the routing logic of its predecessor, ensuring old concepts remain accessible.
- Redundant Layer-wise Expert Elimination (LRER): Recognizing that not every layer requires a unique expert for every concept (e.g., early layers handling shared textures), $R^2MoE$ calculates an importance score for each new expert. Experts falling below a specific threshold are eliminated, and the model reuses existing parameters, significantly curbing growth.
- Hierarchical Local Attention-Guided Inference (HLAG): To prevent feature entanglement in multi-concept scenes, the model uses GPT-4 to generate spatial bounding boxes and the Segment Anything Model (SAM) to create precise masks. This ensures that experts only influence their designated regions.
Key Results
The researchers evaluated $R^2MoE$ on the DreamBooth and CustomConcept101 benchmarks, simulating a sequential lifelong learning environment.
Compared to the previous state-of-the-art method, CIDM, $R^2MoE$ demonstrated substantial improvements:
- Parameter Efficiency: On the 101-concept benchmark, $R^2MoE$ required only 15.2M parameters, a 63.3% reduction compared to CIDM (41.4M) and an even larger gap compared to C-LoRA (80.8M).
- Forgetting Rate: The method achieved a near-zero forgetting rate of 0.19%, significantly lower than CIDM's 8.56%.
- Image Fidelity: $R^2MoE$ outperformed existing methods in Image Alignment by 4.3% to 10.4%, maintaining high quality even as the number of learned concepts increased.
Conclusion
$R^2MoE$ provides a scalable solution for lifelong personalization in generative models. By intelligently managing expert redundancy and preserving routing logic, it allows for the continuous addition of new concepts without the typical penalties of forgetting or excessive model size. Future work may focus on integrating the layout-guidance capabilities of GPT-4 and SAM directly into the diffusion model to reduce inference latency.