Stable Diffusion Models are Secretly Good at Visual In-Context Learning

2025-08-13

AI Summary

This work investigates the visual in-context learning (V-ICL) capability of pre-trained Stable Diffusion models without fine-tuning or additional training data. To this end, we propose a plug-and-play attention recalibration mechanism that dynamically reconstructs contextual relationships between example prompts and the query within self-attention layers, enabling rapid adaptation across diverse vision tasks. Our approach is the first to uncover the latent, general-purpose V-ICL capacity inherently embedded in diffusion-based generative models and supports multi-prompt ensembling to enhance generalization. We validate our method on six representative vision benchmarks; notably, on foreground segmentation using Pascal-5i, it achieves an mIoU of 62.4%, outperforming Visual Prompting and IMProv by +8.9% and +3.2%, respectively, and substantially surpassing existing zero-shot transfer methods.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Generation

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Large language models (LLM) in natural language processing (NLP) have demonstrated great potential for in-context learning (ICL) -- the ability to leverage a few sets of example prompts to adapt to various tasks without having to explicitly update the model weights. ICL has recently been explored for computer vision tasks with promising early outcomes. These approaches involve specialized training and/or additional data that complicate the process and limit its generalizability. In this work, we show that off-the-shelf Stable Diffusion models can be repurposed for visual in-context learning (V-ICL). Specifically, we formulate an in-place attention re-computation within the self-attention layers of the Stable Diffusion architecture that explicitly incorporates context between the query and example prompts. Without any additional fine-tuning, we show that this repurposed Stable Diffusion model is able to adapt to six different tasks: foreground segmentation, single object detection, semantic segmentation, keypoint detection, edge detection, and colorization. For example, the proposed approach improves the mean intersection over union (mIoU) for the foreground segmentation task on Pascal-5i dataset by 8.9% and 3.2% over recent methods such as Visual Prompting and IMProv, respectively. Additionally, we show that the proposed method is able to effectively leverage multiple prompts through ensembling to infer the task better and further improve the performance.

Problem

• Repurpose Stable Diffusion for visual in-context learning
• Enable multi-task adaptation without fine-tuning
• Improve performance on six vision tasks

Innovation

• Repurposes Stable Diffusion for visual in-context learning
• Uses in-place attention recomputation in self-attention layers
• Achieves multi-task adaptation without fine-tuning