Visual Prompting via Image Inpainting

Given an input image $x\in R^{H\times W\times 3}$ and a binary mask $m\in \{0,1{\}}^{H\times W}$, the goal of an inpainting function $f$ is to synthesize a new image $y\in R^{H\times W\times 3}$, with the masked regions filled:

To implement $f$ with a neural network, it is necessary to consider design choices like the network architecture, how to train it, and whether it outputs a distribution over possible completions or pixels. We propose the MAE-VQGAN model, which combines ideas from MAE mae and VQGAN esser2021taming. Following the design of MAE, the model is based on ViT vaswani2017attention; dosovitskiy2020image and it is trained via masked auto-encoding by randomly masking image patches and then applying $f$ to reconstruct the image from the non-masked parts. MAE-VQGAN models the distribution $p_{\theta }(z_i|x,m)$, where $z_i\in V$ is a visual token from a VQGAN vocabulary $V$ that corresponds to the $i^{th}$ ViT patch. For simplicity, we use a fixed ImageNet pretrained VQGAN codebook.²²2Using a publicly available checkpoint from https://github.com/CompVis/taming-transformers Unlike MAE which directly predicts pixels, MAE-VQGAN assigns probabilities to visual tokens via a softmax layer, which is better suited for capturing ambiguities. During training, we obtain ground truth visual tokens by mapping the image to visual tokens indices using the VQGAN encoder. The model is trained using cross entropy loss.

Let $\hat{z}=({\hat{z}}_1,...,{\hat{z}}_k)$ be the ordered set of predicted visual tokens. To obtain ${\hat{z}}_i$, we use argmax:

Then, to decode the visual tokens to pixels, we apply VQGAN decoder to $\hat{z}$ to obtain $y$.

To prompt an inpainting model, we construct a visual prompt, a grid-like image composed of task input-output example(s), and a new query image. The model then has to inpaint the rest of the image such that it is consistent with the task defined in the examples (see Figure 1).

Let $S=\left\{\right(x_i,y_i){\}}_{i=1}^n$ be the set of input-output examples where $x_i$ is an image and $y_i$ is a function of $x_i$ (e.g $y_i$ is a segmentation mask). We assume $n$ is small (one or few examples). Then, given $S$ and a new input query $x_q$, the goal is to predict the corresponding label $y_q$. To prompt the inpainting model discussed in Section 3.1, we need to define a function $g$ that maps the examples set $S$ and query image $x_q$ to a new image and a mask:

The image $x_{vp}$ is the visual prompt and the mask $m$ defines the masked region $f$ has to predict. For a given task, there might exist multiple implementations of $g$ that can be considered. The goal of the inpainting model is to reason about the visual prompt $x_{vp}$, and output a plausible completion without performing any additional training:

To obtain $y_q$, we just take the part of $y_{vp}$ corresponding to the mask $m$.

Visual Prompt Engineering. For the visual prompting to work, $g$ should output a good visual prompt, composed of the examples $S$ and query image $x_q$. Therefore, $g$ has to determine where and how to embed the inputs in the visual prompt image, considering the nature of the completion task. All the functions $g$ used in this work were hard-coded and manually engineered. In most cases, $g$ stacks the examples and image query horizontally by creating an image grid of $(n+1)\times 2$ cells, where the $i^{th}$ example is placed in the $i^{th}$ row, and the image query is in the last row. The grid has a fixed size, and therefore before populating it the input-output example pair(s) and query are first resized. Another consideration is how to draw every $(x_i,y_i)$ pair. For example, if $y_i$ is a segmentation mask, we can choose to use different colors to draw it. In Section 4.4, we describe different prompt design choices and their effect on the results.

Visual Prompt Ensembling. There could be multiple options to define $g$. The idea in prompt ensembling, inspired by NLP jiang2020can; prompt_ensemble, is to construct multiple different prompts, apply the inpainting model $f$ on each prompt individually to obtain a set of predictions. The final prediction can be determined, for example, via majority voting, or weighted average. For simplicity, here we use a simple average.

The images produced by $g$ are by construction not natural. Specifically, these images have a grid-like figure structure that stitches together images coming from different distributions, like natural images and segmentation masks. Therefore, a model trained on a standard dataset (e.g., ImageNet ILSVRC15) might struggle to process these grid-like images. To mitigate the domain gap, we collected a new dataset.

The Computer Vision Figures (Figures) dataset consists of $88,645$ images that more closely resemble the structure of our visual prompts. The dataset was collected from Arxiv, the open-access web archive for scholarly articles from a variety of academic fields. Arxiv sources are publicly available to download starting from $2010$. We downloaded all paper sources from $2010$ to $2022$ and selected the Computer-Vision partition “cs.CV” sources, as they contain images that more closely resemble a grid structure, as shown in Figure 3. To remove unrelated source images like graphs or charts, we manually tagged $2000$ images and trained a binary image classifier to assign a high score to source images in a figure-like structure with at least one natural image. We then used the classifier over the entire data to keep only the most informative source images, coming from $23,302$ different papers. We randomly partitioned $90\%$ of the data to train and left the rest for validation. We include a datasheet with more information in the Supplementary Material.

To study the effect of model choice on prompting results, we experiment using different models, including MAE-VQGAN (see Section 3.1) and several other inpainting models briefly described below.

VQGAN esser2021taming is an autoregressive transformer model used for inpainting and image generation. Visual tokens are predicted sequentially, line-by-line, and the model is trained using cross-entropy loss. The VQGAN model codebook is used to encode visual tokens, and it is trained beforehand using perceptual loss johnson2016perceptual and GAN loss goodfellow2014generative. We train it on ImageNet and our Figures dataset, following hyperparams in esser2021taming, and use a pretrained codebook with a vocabulary of size $|V|=1024$.

BEiT bao2021beit is a masked auto-encoder. The model maps each input $16\times 16$ patch to a visual token from a d-VAE ramesh2021zero vocabulary of size $8192$. To encode each visual token, the image is first resized to $112\times 112$ and then mapped to $196$ tokens. We use the publicly available BEiT large model, pretrained on ImageNet-21k. We also pretrain a large BEiT model on Figures for $1000$ epochs.

MAE mae. Similar to BEiT, MAEs attempt to reconstruct a masked input image. Unlike in BEiT, the model directly regresses pixels and it is trained with $l2$ loss. During pretraining, only non-masked tokens are fed into the encoder, which results in a faster training time. We use a publicly released checkpoint pretrained on ImageNet, and pretrain another model for $1000$ epochs on our dataset.

Copy Example. This simple baseline simply replicates the first example label as the output.

Implementation Details. All the models we describe are large transformer-based models vaswani2017attention; dosovitskiy2020image, with patch size 16$\times$16, embedding dim $1024$, $24$ layers, and $16$ heads. For training, we used a machine with 8 Quadro RTX 6000 GPUs, with a batch size of $48$. The input image size is $224\times 224$.

We quantitatively evaluate the inpainting models described above on computer vision tasks.

Visual Prompt. Given one example pair and a query image, we structure the prompt in the same fashion for all tasks. We construct a grid of $2$$\times$$2$ sub-images, where the example pair is embedded in the first row, and the query image appears in the bottom left cell. See the example in Figure 1.

Computer vision tasks. We evaluate the inpainting models on standard image to image tasks like Foreground Segmentation, Single Object Detection and Colorization.

• [leftmargin=*]

• Foreground Segmentation. The goal is to binary-segment the query image to Foreground and Background. The example is an image and corresponding binary segmentation mask. The query is a new image, and the goal is to complete a corresponding segmentation mask. We use the Pascal-5i pascal_5i dataset, which is comprised of 4 different image splits where every split contains between 346 and 725 images and associated segmentation masks. For each class, the data contains a few image-mask pairs, together with held-out image queries. For every image query, we choose one random example pair. To evaluate, every pixel in the completed image is first mapped to the nearest Foreground or Background color. Finally, we report the mean IOU (mIOU) metric.

• Single Object Detection. Similarly to Foreground Segmentation, the goal here is to binary-segment the object that appears in the query image. However, this task is more challenging than Foreground Segmentation because the example mask is obtained from a bounding box which is more coarse than a segmentation mask. We use the Pascal VOC 2012 dataset using images and their associated detection boxes. For simplicity, we use Pascal annotations to include only images with a single object and filter out trivial images that have an object covering more than $50\%$ of the image. We randomly select an example pair and image query of the same object class and repeat the process with $4$ different random seeds. For evaluation, we follow a similar process as in Foreground Segmentation to obtain a binary segmentation mask. Then we keep the connected component with the largest area using morphological operations and draw a bounding box around it. We report the mIOU results.

• Colorization. The goal is to map a gray-scale image to a color image. The example pair is a gray-scaled image and the corresponding color image, as shown in Figures 1 and 4. We randomly sampled $1000$ example pairs and image query from ImageNet ILSVRC15 validation set and converted them to gray-scale to obtain gray-scale and color version for each image. We report the MSE loss and LPIPS zhang2018unreasonable.

Results. We include quantitative results in Table 1, and qualitative completion results in Figure 4. Training on the Figures dataset improves the results for most models in all the downstream tasks. MAE-VQGAN outperforms the other models by a large margin for detection and segmentation and generates much sharper images than the MAE. We find that VQGAN struggles to output accurate results, likely due to sequential decoding. The BEiT model is outperformed by MAE, most likely because its training process is less sample efficient. For more results, see the Supplementary Material.

To assess the compositional prediction capabilities of inpainting models, we created a set of 3 simple synthetic tasks and 3 of their combinations, and evaluated each model on $100$ examples per task.

Visual Prompt. Given two example pairs and a query image, we structure the prompt in the same fashion for all tasks. We construct a grid of $3\times 2$ sub-images, where the example pairs are embedded in the first two rows, and the query image in the bottom left cell. We include examples in Figure 5.

Change prediction tasks. Each example pair is an image of a colored shape, and a corresponding image with an introduced change. The change can be either in color, shape, size or a combination of two changes. Next, we describe each individual task in more detail.

• [leftmargin=*]

• Resize. Each example pair contains an image of a circle, and a corresponding image with the circle smaller in size. The goal is to predict the image with the resized version given image query.

• Shape. Here every example pair is an image with circle, and a corresponding image with a rectangle. Both are similar in size and appear in the same location. The goal is to predict the image with rectangle, given a new image query.

• Color. Each example pair contains an image of a circle appearing in the same location, with the color changed from green to blue. Given a new image query, the goal is to predict the corresponding image with the circle colored in blue.

Evaluation. We map each predicted pixel to its nearest neighbor color from a predefined set of options: black, white, blue, or green. We then measure and report the color-aware mIOU, by considering pixel predictions that appear in the ground-truth shape color as foreground and treat the rest as background.

Results. The results are presented in Table 2, for MAE-VQGAN prediction examples see Figure 5. Without training on the Figures dataset, inpainting models fail to generalize to these previously unseen tasks. The performance of all models increases when they are trained on the Figures dataset. Yet, the same models struggle with combinations of tasks due to the increasing complexity. The VQGAN model utilizes sequential decoding and therefore lacks context, which leads to poor performance. The MAE model outperforms MAE-VQGAN on color, and BEiT performs poorly in size. These models rely on pretrained codebooks (VQGAN and dVAE) that are likely not geared towards these tasks.

Comparison to finetuning and Few-Shot baselines.We include a comparison to baselines that utilize $K=\{1,4,16\}$ training examples for each target class. For completeness, we also include the results of FWB nguyen2019feature and CyCTR zhang2021few, classic 1-shot baselines, which we view as an upper-bound of our approach. FWB and CyCTR utilize a fully labeled base classes train set (2086 to 5883 on different Pascal 5i splits). Additionally, their architecture was designed for the foreground segmentation task (e.g, they operate in higher resolution). The results in Table 3 indicate that the Visual Prompting results of MAE-VQGAN trained on Figures are significantly superior to standard finetuning baselines of MAEs pretrained on unlabeled ImageNet. FWB nguyen2019feature and CyCTR zhang2021few outperform Visual Prompting, mainly because they pretrain on a large tagged base classes dataset and utilize architectures that are specific to image segmentation.

Dataset effect. We evaluate the effect of pretraining on a larger and more diverse dataset. We compare training on ImageNet only, Figures only, and a combination of the two. We report the mIOU results on Pascal 5i for Foreground Segmentation in Figure 7. The MAE-VQGAN trained on ImageNet achieves a consistently low $5$ points mIOU. The model trained on the combined dataset performs best, which demonstrates that MAE-VQGAN can benefit from additional amounts of unlabeled images.

More examples, better results. We study how increasing the number of input-output pair examples in the visual prompt affects the results. Intuitively, we expect that including more examples should reduce ambiguities and lead to better results. We use an MAE-VQGAN pretrained on the Figures dataset, and use data from PASCAL-5i. We construct a large grid that can populate up to $8$ examples and an image query. We randomly choose different numbers of examples and randomize the placements. The results in Figure 6 confirm that using more examples leads to better segmentation results.

Prompt Engineering. We explore the effect of constructing different visual prompts for Foreground Segmentation and their corresponding MAE-VQGAN results (see Figure 8.a-b). The model generates plausible completions when changing the prompt layout (e.g. horizontal order vs. vertical order) and when changing the mask colors, texture or using only edges (see Figure 9). The mIOU results in Table 5 indicate that the model performs better with a vertical layout and when the segmentation mask colors are black and white. Interestingly, by analyzing the average attention heads of a masked patch token, we observe that the attention changes following the change in the prompts layout (see Figures 8.d-e).

Prompt Ensembling. Inspired by Prompt Ensembling in NLP prompt_ensemble, given the same example pair and image query, we construct multiple different visual prompts (e.g, horizontal and vertical layouts, see Figure 8.a-b). We then average the completion results. The results in Table 5 on the Synthetic Study tasks demonstrate that utilizing multiple prompts can lead to improved, and more stable performance.

Style/content extrapolation. Inspired by the classic example from Tenenbaum and Freeman tenenbaum2000separating (on the task originally suggested by Hofstadter hofstadter1995fluid), we use MAE-VQGAN to extrapolate letter sequences printed in different fonts (see Figure 10). We find that the model can extrapolate given style and new content (Figure 10a) but that it struggles to extrapolate new content (Figure 10b). The model also struggles to extrapolate more complex letter sequences; the performance deteriorates even if both style and content are given (Figure 10 c-d).

Limitations. The focus of this work is to present a proof of concept that shows it is possible to visually prompt simple image inpainting models trained on noisy, unlabeled data. Specifically, we demonstrate how to pre-train a network once, then prompt it to perform reasonably well on many tasks. The fact that this is possible is surprising and scientifically interesting, although this approach is not competitive with supervised task-specific models. For visual prompting to work, the inpainting models require training on the Computer Vision Figures dataset. However, our initial experiments suggest that it can benefit from training on additional natural image data (see Figure 7). Other limitations include ambiguities in the task definition, reliance on a pretrained VQGAN decoder, and worse performance when the input-output example(s) are not aligned (see examples in Figure 11).

Qualitative model comparison. We include a qualitative comparison of visual prompting results when applied to Foreground Segmentation using different inpainting models (see Figure 13). Compared to other models, MAE-VQGAN outputs more smooth and accurate segmentation results.

Single Object Detection. We include qualitative results of MAE-VQGAN when applied to Single Object Detection (see Figure 14). As mentioned in Section 4.2, the raw output is rounded and postprocessed using morphological operations to convert the foreground segmentation to a box.

The effect of the input-output example. We fix the input query and change the input-output examples and visual prompt MAE-VQGAN (see Figure 15). Different examples could lead to slightly different synthesis results. However, as long as the example is meaningful (e.g, not fully background or foreground) the model results are plausible regardless of the choice of example.

Visual Prompting results on computer vision tasks. We include the mean and standard deviation for each Visual Prompting model on Pascal 5i (see Table 6).