Improving In-Context Learning in Diffusion Models with Visual Context-Modulated Prompts

Tianqi Chen,Yongfei Liu,Zhendong Wang,Jianbo Yuan,Quanzeng You,Hongxia Yang,Mingyuan Zhou

2023-12-03

Abstract:In light of the remarkable success of in-context learning in large language models, its potential extension to the vision domain, particularly with visual foundation models like Stable Diffusion, has sparked considerable interest. Existing approaches in visual in-context learning frequently face hurdles such as expensive pretraining, limiting frameworks, inadequate visual comprehension, and limited adaptability to new tasks. In response to these challenges, we introduce improved Prompt Diffusion (iPromptDiff) in this study. iPromptDiff integrates an end-to-end trained vision encoder that converts visual context into an embedding vector. This vector is subsequently used to modulate the token embeddings of text prompts. We show that a diffusion-based vision foundation model, when equipped with this visual context-modulated text guidance and a standard ControlNet structure, exhibits versatility and robustness across a variety of training tasks and excels in in-context learning for novel vision tasks, such as normal-to-image or image-to-line transformations. The effectiveness of these capabilities relies heavily on a deep visual understanding, which is achieved through relevant visual demonstrations processed by our proposed in-context learning architecture.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper aims to address the challenges faced in in-context learning within the visual domain, particularly in achieving efficient in-context learning in visual foundation models such as Stable Diffusion. Specifically, existing methods face several major challenges in visual in-context learning:

1. **Expensive Pre-training**: Some methods require training from scratch, which increases the cost of pre-training.
2. **Restrictive Problem Formulation**: For example, some methods require specific input structures (such as pairwise image grids), limiting their flexibility.
3. **Limited Visual Understanding**: Existing methods may lack sufficient understanding when dealing with complex visual tasks.
4. **Insufficient Out-of-Domain Generalization**: Existing models perform poorly when handling unseen tasks.

To address these challenges, the paper proposes an improved version of the Prompt Diffusion method, called iPromptDiff. iPromptDiff enhances the understanding of visual context by introducing an end-to-end trained visual encoder and effectively integrating it with textual prompts, thereby improving the model's performance on new visual tasks. This method not only reduces overfitting to textual prompts but also enhances the model's generalization ability on out-of-domain tasks. Experimental results show that iPromptDiff performs excellently across various visual tasks, particularly demonstrating significant advantages in handling unseen tasks.