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

TL;DR

This paper demonstrates that off-the-shelf Stable Diffusion models can be repurposed for visual in-context learning across multiple tasks without additional training, by modifying attention mechanisms to incorporate context.

Contribution

The authors introduce a novel in-place attention re-computation in Stable Diffusion that enables effective visual in-context learning without fine-tuning.

Findings

Improves foreground segmentation mIoU by 8.9% on Pascal-5i.

Effectively leverages multiple prompts for better task inference.

Achieves competitive results across six diverse vision tasks.

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.