Image Inpainting Models are Few-Shot Learners (Given the Right Data)
Anonymous Author(s)
AffiliationAddress email
Abstract
How does one adapt a pre-trained visual model to novel downstream tasks withouttask-specific finetuning or any model modification ?
Taking inspiration from prompting in NLP systems, this paper investigates visual prompting: given input-outputimage example(s) of a new task at test time and a new input image, the goal is toautomatically produce the correct output image, consistent with the proposed task.
We show that posing this problem as a simple image inpainting task – literally justfilling in a hole in a concatenated visual prompt image – turns out to be surprisinglyeffective, given that the inpainting algorithm has been trained on the right data.
We train masked auto-encoding models on a new dataset that we curated - 88kunlabeled figures from academic papers sources on Arxiv.
We apply visual prompting to these pretrained models and demonstrate results on various downstreamtasks, including foreground segmentation, single object detection, colorization,edge detection, etc.
All our code, models, and dataset will be made available.
Introduction
In the past few years, self-supervised learning has gained popularity in computer vision and naturallanguage processing (NLP).
The growing capacity of modern deep learning models made themprone to overfitting when trained on relatively small labeled datasets.
Self-supervised learningprovided a solution to this problem by generating “free labels” for any dataset, without the need formanual annotation, addressing the data hunger in these high-capacity deep learning models.
However,features learned via self-supervision are not “ready for use” – they typically need to be adapted for agiven downstream task by fine-tuning on some labeled dataset.
Could this fine-tuning be avoided?
In NLP, prompting [5] has recently emerged as a way to employ a model for a new task without anyadditional training.
A common way of task-prompting for a specific language understanding task attest time is to provide the trained model with input corresponding to example(s) of the target tasktogether with the query.
E.g., typing the following input prompt:
Je suis désolé
I’m sorry
J’adore la glacewill prompt the model [5] to perform the task of French-to-English translation, returning:
I love ice cream
Can this idea of test-time task prompting be generalized to the visual domain?
That is, instead of the each trained model serves its .g
[Figure]
Output
Inpainting
Model f
SupportImage
SupportResult
QueryImage xyx q vp C oncatenate into single image x
[Figure]
segmentation, detection, classification), can we have a single general model that is able to perform awide range of user-specified tasks without any fine-tuning (i.e., weight modification) ?
In this paper we take a step toward this goal by demonstrating that large-capacity inpainting models,when trained on the right data, can be surprisingly effective tools for visual prompting .
As shown in
Figure 1, we define each task by constructing a new grid-like image that contains an input-outputexample(s) of the task and a novel query (green border).
The input-output example describes the task,and the image query defines a new input.
The model then produces the result by simply inpaintingthe rest of the image (red border).
This setting is most similar to the classic Image Analogies [22]formulation, but is less constrained: instead of explicitly defining the A, A’, and B images separately,we simply concatenate them into a single image with a hole (hence, visual prompting is not exactlyan analogy since there is no implied left-to-right ordering).
Our goal is also not dissimilar from theaims of meta-learning and few-shot learning methods, except that we make no distinction betweentasks and example pairs.
The only requirement of our formulation is that the tasks must be defined interms of image-to-image translations, which is a very large subset of vision problems.
To obtain data that fits our framework, we utilize a domain that spans across a variety of computervision tasks - figures and infographics from computer vision articles available on Arxiv.
We builda large dataset of 88k figures, many of which contain grids of images and their corresponding taskresults (e.g. images and their segmentation masks/stylized versions/edges, etc.).
We then train largecapacity inpainting models to predict randomly masked patches from figures given other patchesfrom the same figure.
Our main contributions are as follows.
First, we present a simple yet surprisingly powerful generalapproach for visual prompting.
We show that various computer vision tasks can be treated as gridinpainting problems, given a few examples of task inputs and outputs and a query image.
Second,we provide a new dataset that allows a model to learn such grid structures without any labeling, taskdescriptions, or any additional information about the grid structure.
Finally, we show that, whileusing our new dataset for training is essential, adding more generic images from other sources (e.g.
ImageNet) further improves the results.
Related Work
Natural Image Inpainting .
Filling empty regions in an image has been widely explored for naturalimages.
Earlier methods used data from the input image itself for inpainting [14, 4, 11, 3, 47], whereaslater works utilized datasets of images as source data [19, 37, 55, 30, 31].
Recent methods haveattempted to apply transformers to visual synthesis tasks [8, 57, 15, 56, 7].
Due to the exponentiallylarge number of completion options for a single output patch, these approaches rely on a discrete latentcodebook [50, 39, 15] which serves as a smaller yet expressive vocabulary.
To tackle the multimodalnature of synthesis, different approaches have been proposed to model the distribution over possiblecompletions [15, 56, 7].
For example, [15, 56] proposed to synthesize images line-by-line using anautoregressive model and [7] have proposed iterative parallel decoding.
While the standard inpaintingtask typically aims to complete blank parts in natural images, our focus is on completing grid-likevisual prompts, which require reasoning across multiple images to infer the completion.
Hole-filling as a Pretext task.
Recent work has shown that self-supervised pretraining can generatepowerful representations for transfer learning, even outperforming its supervised counterparts onchallenging vision benchmarks [49, 18, 9, 21, 6, 10, 35, 16].
Pathak et al.
proposed contextencoding [37], where the goal is to predict an image region given its context.
Based on the recentsuccess of Vision Transformers (ViTs) [13], multiple works have proposed the Masked image AutoEncoding (MAE) [1, 20, 53] as a self-supervised pretext task for Vision Transformers (ViTs).
The goalof MAEs is to reconstruct the image given a small subset of input patches.
After pretraining MAEs onunlabeled data, they produce representations that transfer well when fine-tuned on downstream tasks.
Here, we use visual prompting to adapt these models to downstream tasks without any finetuning.
Few-Shot Learning.
In this setting, the algorithm is trained on a labeled dataset of base classes,from which it should transfer to a set of novel classes given only a few training examples (like 10 or30) [36, 26, 32, 52, 54, 46, 58, 2].
Unlike Few-Shot approaches, here we do not assume access to alarge training set of base-classes, and our architecture is not task-specific.
Our approach is Few-Shotonly in the sense that we construct a visual prompt that contains one or two task examples.
Image Analogies.
Hertzmann et.
al.
[22] proposed the framework of Image Analogies for texturesynthesis, where the algorithm is given a pair of training images and a query image.
The goalof the learner is to synthesize a new image conditioned on the query, following the relationshipinferred from the training pair.
Other works have used analogies in style transfer [48], and as asupervised image synthesis task [40].
Predicting the correct completion was previously modeledas a classification problem [42, 23], and other works have explored analogies in the context oflearning different transformations between pairs of images [34, 45].
Unlike these approaches, we useinpainting MAE models that learn from data, without assuming any predefined analogies structure.
Prompting in NLP.
With the recent success of large unsupervised language models [43, 12], Brownet al.
[5] presented how a variety of NLP problems can be reformulated to a text completion problemgiven a predefined prompt, which can be used to solve different tasks without any finetuning.
Prompting was shown to be a useful tool for solving various NLP tasks and benchmarks [38, 5].
More recently different approaches to prompting have emerged including Prompt Engineering [5, 33],
Prompt Ensembling [24], and Prompt Prefix Tuning [29, 27].
Inspired by the success of prompting in
NLP, we aim to study prompting in computer vision where prompting hasn’t been widely explored.
Visual Prompting using Image Inpainting
We turn to describe how to perform visual prompting using Image Inpainting models.
In Section 3.1,we describe Inpainting Models.
We then proceed to discuss visual prompting and propose differentways to create visual prompts in Section 3.2.
An illustration of the approach is included in Figure 1.
3.1 Inpainting
Given an input image x ∈ R H × W × 3 and a binary mask m ∈ { 0 , 1 } H × W , the goal of an inpainting function f is to synthesize a new image y ∈ R H
× W × 3 , with the masked regions filled:
[Equation]
To implement f with a neural network, it is necessary to consider design choices like the networkarchitecture, 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 [20] and VQGAN [15].
Following the design of MAE, the model is based on ViT [51, 13] and it is trained via masked autoencoding, by randomly masking image patches then applying f to reconstruct the image from thenon-masked parts.
MAE-VQGAN models the distribution p θ ( z i | x, m ) , where z i ∈ V is a visual tokenfrom a VQGAN vocabulary V that corresponds to the i th ViT patch.
Unlike MAE which directlypredicts pixels, this allows the model to assign probabilities to visual tokens, which is better suited forcapturing ambiguities.
During training, we obtain ground truth visual tokens by mapping the imageto visual tokens indices using the VQGAN encoder.
The model is trained using cross entropy loss.
Let ˆ z
=
(ˆ z 1 , ..., ˆ z k ) be the ordered set of predicted visual tokens.
To obtain ˆ z i , we use argmax:
[Equation]
Then, to decode the visual tokens to pixels, we apply VQGAN decoder to ˆ z to obtain y .
3.2 Prompting Inpainting Models
To prompt an inpainting model, we construct a visual prompt , a grid-like image composed of supportexamples, and a new query image.
The model then has to inpaint the rest of the image such thatit is consistent with the task defined in the support, applied over the image query (see Figure 1).
Let S =
{ ( x i , y i ) } ni =1 be the set of support 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 support set S and query image x q to a new image and a mask:
[Equation]
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:
[Equation]
To obtain y q , we just take the part of y vp corresponding to the mask m .
For the visual prompting to work, g should output a good visual prompt, composed of the supports Sand query image x q .
Therefore, g has to determine where and how to embed the inputs in the visualprompt image, considering the nature of the completion task.
All the functions g used in this workwere hard-coded and manually engineered.
In most cases, g stacks the support pairs and image queryhorizontally by creating an image grid of ( n + 1) × 2 cells, where the i th support pair is placed in thei th row, and the image query is in the last row.
Another consideration is how to draw every supportpair ( x i , y i ) .
For example, if y i is a segmentation mask, we can choose to use different colors to drawit.
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 [24, 28], is to construct multiple different prompts, apply the inpainting modelf on each prompt individually to obtain a set of predictions
3.3 The Computer Vision Figures Dataset
[Figure]
The images produced by g are by construction not natural.
Specifically, these images have a grid-likefigure structure that stitches together images coming from different distributions, like natural imagesand segmentation masks.
Therefore, a model trained on a standard dataset (e.g., ImageNet [41]) mightstruggle to process these grid-like images.
To mitigate the domain gap, we collect a new dataset.
The Computer Vision Figures (Figures) dataset consists of 88 , 645 images that more closely resemblethe structure of our visual prompts.
The dataset was collected from Arxiv, the open-access webarchive for scholarly articles from a variety of academic fields.
Arxiv sources are publicly availableto download starting from 2010 .
We downloaded all paper sources from 2010 to 2022 and selectedthe Computer-Vision partition “cs.CV” sources, as they contain images that more closely resemblea grid structure, as shown in Figure 2.
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 tosource images in a figure-like structure with at least one natural image.
We then used the classifierover the entire data to keep only the most informative source images, coming from 23 , 302 differentpapers.
We randomly partitioned 90% of the data to train and left the rest for validation.
We includeadditional examples and provide more information about the dataset in the Supplementary Material.
Experiments and Results
To study visual prompting, we pretrain different models (see Section 4.1) on ImageNet and onthe Figures dataset, then evaluate the models using different prompts on simple downstream computervision tasks (see Section 4.2).
Using a synthetic dataset, we assess how the choice of model and dataaffect the success of visual prompting in Section 4.3, and explore different prompting design choicesin Section 4.4.
We include additional experiments and visualizations in the Supplementary Material.
4.1 Models and Baselines
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 [15] 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-entropyloss.
The VQGAN model codebook is used to encode visual tokens, and it is trained beforehandusing perceptual loss [25] and GAN loss [17].
We train it on ImageNet and our Figures dataset,following hyperparams in [15], and use a pretrained codebook with a vocabulary of size | V | = 1024 .
BEiT [1] is a masked auto-encoding model.
The model maps each input 16 × 16 patch to a visualtoken from a d-VAE [39] vocabulary of size 8192 .
To encode each visual token, the image is firstresized to 112 × 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.
[Table]
MAE [20].
Similar to BEiT, MAEs attempt to reconstruct a masked input image.
Unlike in BEiT, themodel directly regresses pixels and it is trained with l 2 loss.
During pretraining, only non-maskedtokens are fed into the encoder, which results in a faster training time.
We use a publicly releasedcheckpoint pretrained on ImageNet, and pretrain another model for 1000 epochs on our dataset.
Copy Support.
This simple baseline simply replicates the first support example label as the output.
Implementation Details.
All the models we describe are large transformer-based models [51, 13],with patch size 16 × 16, embedding dim 1024 , 24 layers, and 16 heads.
For training, we used amachine with 8 Quadro RTX 6000 GPUs, with a batch size of 48 .
The input image size is 224 x 224 .
4.2 Downstream Computer Vision Tasks
We evaluate the inpainting models described above on computer vision tasks, as described below.
Visual Prompt.
Given one support pair and a query image, we structure the prompt in the samefashion for all tasks.
We construct a grid of 2 × 2 sub-images, where the support pair is embedded inthe first row, and the query image appears in the bottom left cell.
See the example in Figure 1.
Foreground Segmentation.
The goal is to binary-segment the query image to Foreground and
Background.
The support example is an image and corresponding binary segmentation mask.
Thequery is a new image, and the goal is to complete a corresponding segmentation mask.
We use the
Pascal-5i [44] dataset, which is comprised of 4 different image splits where every split containsbetween 346 and 725 images and associated segmentation masks.
For each class, the data containsa few support image-mask pairs, together with held-out image queries.
For every image query, wechoose one random support pair.
To evaluate, every pixel in the completed image is first mapped tothe 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-segmentthe object that appears in the query image.
However, this task is more challenging than Foreground
Segmentation because the support mask is obtained from a bounding box which is more coarse than asegmentation mask.
We use the Pascal VOC 2012 dataset using images and their associated detectionboxes.
For simplicity, we use Pascal annotations to include only images with a single object andfilter out trivial images that have an object covering more than 50% of the image.
We randomlyselect a support and image query of the same object class and repeat the process with 4 differentrandom seeds.
For evaluation, we follow a similar process as in Foreground Segmentation to obtaina binary segmentation mask.
Then we keep the connected component with the largest area usingmorphological operations and draw a bounding box around it.
Finally, we report the mIOU results.
Colorization.
The goal is to map a gray-scale image to a color image.
The support part contains a grayscaled image and the corresponding color image, as shown in Figures 1 and 3.
We randomly sampled1000 image pairs for support and image query from ImageNet [41] validation set and converted themto gray-scale to obtain gray-scale and color version for each image.
We report the MSE loss.
[Figure] [Figure]
Results.
We include quantitative results in Table
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 andgenerates much sharper images than the MAE.
We find that VQGAN struggles to output accurateresults, likely due to sequential decoding.
The BEiT model is outperformed by MAE, most likelybecause its training process is less sample efficient.
For more results, see the Supplementary Material.
4.3 Synthetic Data Study
[Figure]
To assess the compositional prediction capabilities of inpainting models, we created a set of 3 simplesynthetic tasks and 3 of their combinations, and evaluated each model on 100 examples per task.
Visual Prompt.
Given two support pairs and a query image, we structure the prompt in the samefashion for all tasks.
We construct a grid of 3 × 2 sub-images, where the support pairs are embeddedin the first two rows, and the query image in the bottom left cell.
We include examples in Figure 4.
Change prediction.
For each support pair ( x i , y i ) , x i is an image of a colored shape, and y i isthe same image with an introduced change.
The change can be either in color, shape, size or acombination of two changes.
We describe the tasks in more detail in the Supplementary Material.
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 pixelpredictions 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 4.
Without training on the Figures dataset, inpainting models fail to generalize to these previously unseentasks.
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 VQGANmodel utilizes sequential decoding and therefore lacks context, which leads to poor performance.
The
[Table]
(cid:20)(cid:21)(cid:19)(cid:18)(cid:28)(cid:27)(cid:26)(cid:25)(cid:25)(cid:29)(cid:24)(cid:23)(cid:28)(cid:25)(cid:22)(cid:21)(cid:24)(cid:27)
[Table] [Figure] [Figure]
4.4 Analysis
[Figure]
More supports, better results.
We study how increasing the number of supports in the visualprompt affects the results.
Intuitively, we expect that including more support examples should reduceambiguities 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 supports and animage query.
We randomly choose different numbers of supports and randomize the placements.
Theresults in Figure 5 confirm that using more support 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 6.a-c).
The model generatesplausible completions when changing the prompt layout (e.g. horizontal order vs. vertical order) andwhen changing the mask colors (black and white vs. purple yellow).
The mIOU results in Table 3indicate that the model performs better with a vertical layout and when the segmentation mask colorsare 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 6.d-e).
Segmentation mIOU results on Pascal- using different prompt layouts and color.
[Table] [Table]
[-#$%&’()&*+",$-".&/+$0$12(+3$4$5," ! 8#$93’+(6",$-".&/+0$12(+3$4$5," ! 6#$%&’()&*+",$-".&/+$0$:3,,&;$4$</’=, !>#$?@3’"A3$"++3*+(&*$B"=$&C$B"’73>$ ="+62$C&’$!" # !3#$?@3’"A3$"++3*+(&*$B"=$&C$B"’73>$="+
-]In every line, the result is based on an ensemble of all previous prompts
[Table]
[-C&’$!
-][Table]
Segmentation
Prompt Ensembling.
Inspired by Prompt Ensembling in NLP [28], given the same support pair andimage query, we construct multiple different visual prompts (e.g, horizontal and vertical layouts, see
Figure 6.a-b).
We then average the completion results.
The results in Table 3 on the Synthetic Studytasks demonstrate that utilizing multiple prompts can lead to improved, and more stable performance.
Dataset effect.
We evaluate the effect of pretraining on a larger and more diverse dataset.
We comparetraining on ImageNet only, Figures only, and a combination of the two.
We report the mIOU resultson Pascal 5i for Foreground Segmentation in Figure 7.
The MAE-VQGAN trained on ImageNetachieves 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.
Limitations.
For visual prompting to work, the inpainting models still require training on the
Computer Vision Figures.
However, our initial experiments suggest that combining it with morenatural images leads to improved performance.
Another limitation with the approach we propose forvisual prompting is ambiguities in the task definition (See Figure 8).
Ambiguities can potentially beavoided by providing multiple support examples.
Lastly, to synthesize pixels we rely on the VQGANdecoder, which was pretrained on ImageNet.
This leads to naturally looking images, and thus themodel fails to synthesize out-of-distribution completions, as in the "image blurring task" in Figure 8.
Conclusion
This paper proposes an approach to visually prompt inpainting models.
We show that inpaintingmodels trained using a Masked Auto-Encoding objective on an unlabeled Computer Vision Figuresdataset can be prompted to solve computer vision tasks without fine-tuning.
Therefore, in Computer
Vision, just like in NLP, Masked Auto-Encoders can behave as Few-Shot learners.