MuseumMaker: Continual Style Customization without Catastrophic Forgetting

Chenxi Liu,Gan Sun,Wenqi Liang,Jiahua Dong,Can Qin,Yang Cong
2024-04-29
Abstract:Pre-trained large text-to-image (T2I) models with an appropriate text prompt has attracted growing interests in customized images generation field. However, catastrophic forgetting issue make it hard to continually synthesize new user-provided styles while retaining the satisfying results amongst learned styles. In this paper, we propose MuseumMaker, a method that enables the synthesis of images by following a set of customized styles in a never-end manner, and gradually accumulate these creative artistic works as a Museum. When facing with a new customization style, we develop a style distillation loss module to extract and learn the styles of the training data for new image generation. It can minimize the learning biases caused by content of new training images, and address the catastrophic overfitting issue induced by few-shot images. To deal with catastrophic forgetting amongst past learned styles, we devise a dual regularization for shared-LoRA module to optimize the direction of model update, which could regularize the diffusion model from both weight and feature aspects, respectively. Meanwhile, to further preserve historical knowledge from past styles and address the limited representability of LoRA, we consider a task-wise token learning module where a unique token embedding is learned to denote a new style. As any new user-provided style come, our MuseumMaker can capture the nuances of the new styles while maintaining the details of learned styles. Experimental results on diverse style datasets validate the effectiveness of our proposed MuseumMaker method, showcasing its robustness and versatility across various scenarios.
Computer Vision and Pattern Recognition
What problem does this paper attempt to address?
This paper proposes two main issues in continuous style customization in the Text-to-Image (T2I) model: catastrophic forgetting and catastrophic overfitting. Although existing pre-trained large-scale T2I models can generate customized images based on textual prompts, they tend to forget previously learned styles and result in unsatisfactory generation when synthesizing styles provided by new users. To address this, the paper introduces the MuseumMaker method. MuseumMaker allows the model to continuously follow a series of customized styles to synthesize images and accumulate these creative works in a "museum". For new styles, it uses a style distillation loss module to extract and learn the style of training data, reducing the learning bias caused by new training image content and preventing overfitting. To address the catastrophic forgetting problem of past learned styles, the paper designs a dual-regularization shared LoRA module to optimize the model update direction and regularize the diffusion model from both weights and features. In addition, to preserve historical knowledge and overcome the limited representational capacity of LoRA, the paper proposes a task-level token learning module to learn unique token embeddings for each new style. Experiments with various style datasets demonstrate the effectiveness and adaptability of MuseumMaker, showcasing its robustness and diversity in different scenarios. Through these mechanisms, MuseumMaker is able to capture subtle differences in new style while preserving the details of previously learned styles.