Visual Prompting via Image Inpainting

In plain words

Large language models can often do a new task when you show them a couple of worked examples in the prompt, with no retraining. This paper asks whether the same trick works for images. The answer it proposes is to paste the example input and output images, plus the new input image, into one larger picture laid out as a grid. Then it hides the empty cell in the corner and asks a picture model to fill it in. The model has only been trained to fill in holes in pictures, yet the filled-in cell often shows the answer to the new task, such as the cut-out object or the colored version of a photo.

The picture model is a standard masked autoencoder that was pretrained on figures taken from computer-vision papers. The authors argue that these figures already look like grids of related images, so the training data matches the prompt format. The method needs no labels for the new task and no change to the model's weights. The results are reasonable on simple image-to-image jobs and well below specialised, supervised systems.

TL;DR

• Visual prompting: a grid image holds one or two input-output examples and a query; a masked image model inpaints the empty bottom-right cell §3.2 §4.2.
• The model, MAE-VQGAN, is a ViT-L masked autoencoder whose decoder predicts VQGAN codebook tokens (1024 entries) instead of pixels §3.1.
• Pretraining uses 88,645 arXiv computer-vision figures; the paper reports that ImageNet-only pretraining performs much worse on these tasks §3.3 §4.2 tab. 1.
• No weight updates at test time. Author-reported results are competitive with other pretrained inpainting models but not with supervised task-specific models §5 §4.4.

Problem

Adapting a pretrained vision model to a new task usually means fine-tuning it on labelled data or changing its head. The paper asks whether a single pretrained model can perform new image-to-image tasks from a few input-output examples given at test time, with no fine-tuning and no change to the model §1. Two gaps motivate the work. Standard inpainting models are trained on natural images and may not handle images that stitch several unrelated pictures together §3.3. Earlier few-shot methods assume a labelled base-class set and task-specific architectures, which this approach does not §2.

Method

Overview

The approach has two parts. First, an inpainting model f predicts masked regions of an image. Second, a hand-written function g builds a visual prompt from the examples and the query, so that the masked region of the prompt is the answer §3.2.

The inpainting model is a masked autoencoder (MAE) on a ViT backbone. Its output is a softmax over a fixed VQGAN codebook of 1024 visual tokens for each patch, trained with cross-entropy on the masked patches only. The VQGAN codebook is an ImageNet-pretrained, frozen encoder and decoder §3.1 models_mae.py models_mae.py.

Formally, the paper writes:

y = f(x, m) §3.1

where x is an image and m a binary mask of the region to fill. For a prompt, the paper writes:

[x_{vp}, m] = g(S, x_q), \qquad y_{vp} = f(x_{vp}, m) §3.2

where S = {(x_i, y_i)} are the example pairs, x_q is the query, x_vp is the visual prompt, and the answer y_q is the part of y_vp under m §3.2. The paper states these equations without numbers.

Pretraining (MAE-VQGAN)

# MAE-VQGAN pretraining step  [§3.1] [code: models_mae.py#L198-L208]
# Input:  images (B, 3, 224, 224)
# Output: scalar loss

def pretrain_step(images, mask_ratio=0.75):
    # Ground-truth tokens: frozen VQGAN encoder, one code per 16x16 patch
    # [§3.1] [code: models_mae.py#L204-L205]
    targets = vqgan_codebook_indices(images)   # (B, 196)
    # Hide a random subset of patches; the encoder sees only the rest
    # [§3.1] [code: models_mae.py#L123-L148]
    visible, hidden_mask, restore = random_mask(images, mask_ratio)
    # Decoder predicts a distribution over the 1024 codes for every patch
    # [§3.1] [code: models_mae.py#L172-L196]
    logits = decoder(encoder(visible), restore)   # (B, 196, 1024)
    # Cross-entropy averaged over the hidden patches only
    # [§3.1] [code: models_mae.py#L206-L207]
    loss = cross_entropy(logits, targets, hidden_mask)
    return loss

The image is split into 16x16 patches. Patches are masked at random, not in blocks; the paper does not state a masking ratio, but the code uses 0.75 main_pretrain.py. Only visible patches go through the encoder, as in MAE §4.1. The decoder predicts a 1024-way distribution per patch, and the loss uses ground-truth codes from the frozen VQGAN encoder §3.1.

Building a visual prompt

# Visual prompting with an inpainting model  [§3.2] [§4.2]
# Input:  pairs: list of (example_in, example_out), each a 111x111 tile
#         query: 111x111 image tile
# Output: predicted output tile for the query

def visual_prompt(pairs, query, model):
    # Grid of 2x2 tiles with 1 px white padding: example pair in the top
    # row, query in the bottom-left cell  [§4.2] [code: pascal_dataloader.py#L47]
    canvas = white_canvas(224, 224)
    canvas = place(canvas, pairs[0], row="top")
    canvas = place(canvas, query, cell="bottom_left")
    # Hide the bottom-right cell: 7x7 of the 14x14 patches
    # [code: evaluate/mae_utils.py#L51-L57]
    visible = top_row_and_left_column_patches()
    # Argmax over codes, then VQGAN decode; keep visible pixels as they are
    # [§3.1] [code: evaluate/mae_utils.py#L127-L138]
    codes = argmax(model.predict(canvas, visible))
    decoded = vqgan_decode(codes)            # bilinear resize to 224
    result = paste(canvas, decoded, where=not visible)
    # Answer is the bottom-right cell  [code: evaluate/segmentation_utils.py#L8-L9]
    return crop_bottom_right(result, start=113)

The paper describes the prediction as argmax over the per-token softmax, decoded by the VQGAN decoder §3.1. The evaluation code also pastes the visible patches back into the decoded image, so only the hidden patches come from the decoder mae_utils.py.

Prompt layouts for the main tasks

For the image-to-image tasks in §4.2, the 2x2 grid has the example input and example output side by side in the top row and the query image in the bottom-left cell. The bottom-right cell is the target fig. 1 §4.2. Other tasks use different grids. The synthetic study uses a 3x2 grid with two example pairs and a query §4.3. Ablations on layout and mask colour are in §4.4.

Ensembling

The paper builds several prompts for the same examples and query and averages the predictions, for example with a horizontal and a vertical layout §3.2 §4.4. The text says "simple average" §3.2. The code's segmentation path averages the softmax outputs of two canvases before taking the argmax evaluate_segmentation.py. Other combinations (averaged logits, product, max) are coded but not described in the paper evaluate_segmentation.py.

Dataset: Computer Vision Figures

• Source: arXiv LaTeX sources, 2010 to 2022, the cs.CV partition §3.3.
• Filter: a binary classifier trained on 2,000 manually tagged images, to keep figure-like images with at least one natural image. The classifier's architecture is not stated §3.3.
• Result: 88,645 images from 23,302 papers, split 90% train / 10% validation §3.3. The datasheet is in the supplement §7.
• Training crops: random crops of figures with no further parsing fig. 3 §3.3.

Hyperparameters

Results

Author-reported. Setup: MAE-VQGAN, ViT-L/16 backbone, 224 px input, prompt built from one example pair and a query, evaluated on Pascal-5i (segmentation), Pascal VOC 2012 single-object subset (detection) and 1,000 ImageNet validation images (colorization) §4.2. Baselines: a copy of the example output, and MAE, BEiT and VQGAN pretrained on ImageNet or on the figures dataset §4.1.

Main table tab. 1

Source: tab. 1. Column labels are copied from the paper: detection columns are labelled "Split 1" to "Split 4", while segmentation uses "Split 0" to "Split 3" (see Pitfalls). Metrics: segmentation and detection report mIoU (higher is better); colorization reports MSE and LPIPS (lower is better). The table caption gives the metrics tab. 1.

Supplement mean and standard deviation §6, supp. table: MAE-VQGAN (figures) foreground segmentation 27.17 ± 2.27, single-object detection 25.00 ± 0.47; MAE (figures) 19.57 ± 3.62 and 5.39 ± 0.32; Copy 14.91 ± 1.93 and 12.76 ± 0.54.

Synthetic data study tab. 2

Colour-aware mIoU on 100 examples per task, from the paper's table:

Few-shot comparison tab. 3

Foreground segmentation mIoU on Pascal-5i. Fine-tuning rows use K labelled examples per target class. MAE-VQGAN uses one prompt example. FWB and CyCTR use the full labelled base-class training set and segmentation-specific architectures and are reported as upper bounds §4.4.

Source: tab. 3. Input resolution of MAE-VQGAN is 111 x 111 for this table tab. 3.

Prompt design §4.4

Foreground segmentation mIoU on Pascal-5i, from the "Prompt Engineering" table in §4.4:

Prompt ensembling §4.4, colour-aware mIoU, cumulative ensembles:

Data effect §4.4, dataset-effect figure: MAE-VQGAN trained on ImageNet alone stays near 5 mIoU for segmentation; training on figures, and on figures plus ImageNet, performs better, with the combination best. Exact values are only in the figure, not in text.

Paper vs code

The code is the authors' repository at commit 6351e0d786c2d25bfdd5fe504153b1dfb345fa81. It was found through the paper's project page (yossigandelsman.github.io/visual_prompt), which links github.com/amirbar/visual_prompting. The repository owner is the first author.

• Batch size and hardware. Paper: batch 48 on 8 Quadro RTX 6000 GPUs §4.1. Code: --batch_size default 64 per process main_pretrain.py; the README's pretraining command uses 64 on 8 V100 GPUs README.md.
• Training schedule. Paper gives no optimizer, learning rate, warmup, mask ratio or epoch count for MAE-VQGAN §4.1. Code and README: AdamW (0.9, 0.95), base LR 1e-4 scaled by batch size, warmup 15 epochs, 1000 epochs, mask ratio 0.75 main_pretrain.py main_pretrain.py README.md. The argparse defaults differ (LR 1e-3, warmup 40, 400 epochs) and the README overrides them.
• Augmentation. Paper: "randomly sample crops" from figures fig. 3. Code: random resized crop with scale 0.2 to 1.0, horizontal flip main_pretrain.py.
• Loss. Paper: cross-entropy over the 1024 codes §3.1. Code: target codes come from the frozen VQGAN encoder applied to the unmasked image, and the loss averages only over hidden patches models_mae.py.
• Decoder size. Not stated in paper. Code: 8 decoder blocks, width 512, 16 heads models_mae.py.
• Mixing ImageNet. Paper reports results for figures plus ImageNet, and the README lists a 3400-epoch "CVF + IN" checkpoint [README.md#L103-L106]. The training script defines --imagenet_percent (default 0.5) main_pretrain.py, but no code in the repository reads it, and the data loader is a single ImageFolder main_pretrain.py. The mixing procedure is therefore not visible in the code.
• Figure-resolution tile. Paper: input 224 px §4.1; few-shot table: MAE-VQGAN at 111 x 111 tab. 3. Code: evaluation places two 111 px tiles with 1 px padding in a 224 px canvas pascal_dataloader.py and resizes masks and images to 111 px evaluate_segmentation.py.
• Mask layout at test time. Paper: bottom-right cell is inpainted §3.2 fig. 1. Code: exactly the bottom-right 7x7 patch block is hidden in the 14x14 patch grid mae_utils.py.
• Metric crop. Paper: mIoU after mapping each pixel to the nearest foreground or background colour §4.2. Code: additionally crops the canvas to start at pixel 113 before computing IoU segmentation_utils.py.
• Ensembling. Paper: simple average of predictions §3.2. Code: average of softmax probabilities of two canvases, then argmax evaluate_segmentation.py. Logit averaging and other combinations are coded but not described in the paper evaluate_segmentation.py.
• Detection post-processing. Paper's supplement: "morphological operations" to get a box §6. Code: pixels are rounded to the nearest of two colours, then the largest connected component whose first pixel is white is boxed, with no morphological step evaluate_segmentation.py box_ops.py.
• Code licence. The datasheet says the CSV and code will be distributed under MIT §7 datasheet. The repository has no LICENSE file, so the code licence is not stated. We therefore quote no code verbatim.
• Dataset. The repository includes figures_dataset/download_links.py and CSV splits, but not the images. The README notes that downloading the arXiv sources takes days README.md. This card publishes no dataset content.

No code-level discrepancy was found for the core inpainting model itself (patch embedding, random masking, decoder, codebook loss) beyond those listed.

Pitfalls

• Code licence is missing. The repository has no LICENSE file, so you cannot copy code verbatim or ship it. Re-implement from the description above, or ask the authors.
• Training recipe is only partly in the paper. Learning rate, schedule, mask ratio, epochs and augmentation are only in the code or README. Use those values, but note they come from code, not the paper.
• Random masks, tile masks at test time. Pretraining hides random patches; testing hides a fixed bottom-right block models_mae.py mae_utils.py. The paper does not discuss this gap.
• Metric crop and padding. Results depend on the 1 px padding, the 111 px tile size and the crop at pixel 113 pascal_dataloader.py segmentation_utils.py. Re-implement these, or your numbers will not match.
• Layout names disagree. The prompt-engineering table says "Horizontal / Vertical" §4.4, while the text talks about top-to-bottom and left-to-right. A commented-out table in the LaTeX source (not in the rendered paper) lists the same runs as Top/Bottom and Left/Right. Its Left/Right black-and-white row averages 31.57, which matches the paper's "Vertical" black-and-white entry. Confirm the convention before reproducing the best layout.
• Split labels disagree. Table 1 numbers detection columns "Split 1" to "Split 4" while segmentation uses "Split 0" to "Split 3" tab. 1. Check which split indices the code uses before comparing.
• Ensembling is underspecified. "Simple average" can mean averaging logits or probabilities, which changes the argmax. The code averages probabilities for the default path evaluate_segmentation.py.
• Detection post-processing differs. The paper mentions morphological operations; the code does not use them box_ops.py. Expect different boxes if you follow the text.
• Thresholds for detection are not stated. Colour rounding for the mask uses a threshold t whose default is (0, 0, 0) evaluate_segmentation.py reasoning_dataloader.py. The paper does not give a threshold.
• Visible patches are copied. The reported output pastes the original visible pixels back in place of the decoder output mae_utils.py. The paper's wording suggests the decoder output is used as-is.
• VQGAN decoding limits fidelity. The output is a VQGAN decode, so out-of-distribution pixels (blur, unusual colours) cannot be produced §4, Limitations. Bilinear resizing to 224 px is part of the decode mae_utils.py.
• Single example is ambiguous. Single-example prompts can be ambiguous; misaligned example pairs degrade results §4, Limitations. More examples help, from 1 to 8 in the authors' experiment §4.4.
• ImageNet mixing is not reproducible from the code alone. See Paper vs code.
• Compute is not stated. The paper gives no GPU-hours. Pretraining runs 1000 epochs in the README, and inference with ensembling runs several forward passes per prompt README.md evaluate_segmentation.py.
• Dataset licensing and redistribution. The figures are excerpts of other papers and are not redistributed by this card. The paper's claim of MIT for the CSV is not verifiable from the repository.
• Results are author-reported and single-task. The paper reports one prompt per task in the main table, with no confidence intervals in the main text tab. 1. Supplement standard deviations exist for two tasks only §6, supp. table.

Lineage

• Builds on: masked autoencoders (MAE, [mae] in the paper) for the encoder and masking §2 §3.1; VQGAN for the discrete codebook §3.1; prompting in NLP (GPT-3, Brown et al.) as the motivation §1; Image Analogies (Hertzmann et al., 2001) as conceptual precursor §2.
• Compared against: Copy baseline; BEiT and VQGAN pretrained on ImageNet or figures; MAE with ImageNet or figures; fine-tuned MAE with 1, 4 or 16 labelled examples; FWB and CyCTR (labelled 1-shot segmentation) tab. 1 tab. 3.
• Superseded by: not known from content already in this wiki.