Images Speak in Images: A Generalist Painter for In-Context Visual Learning

In plain words

Most computer-vision models are built for one job: one network that outputs depth, another that outputs masks, another that cleans up noise. Language models can do many jobs from a text prompt, but vision outputs have no such common format. Painter's answer is to make every output a picture. A depth map becomes a coloured image, a segmentation becomes coloured regions, and a set of keypoints becomes coloured dots. Then the model only has to paint in missing parts of pictures, which is one kind of task for every job.

Training shows the model two input-and-answer pairs from the same job side by side, with part of the answer hidden. Afterwards, a new picture plus one solved example from the same job tells the model which job to do, and it paints the answer. Nothing task-specific is added to the network. The authors report that this works on several standard tasks, sometimes beating specialist models and sometimes trailing them.

TL;DR

Painter writes the outputs of seven dense tasks as RGB images and trains one masked-image model on stitched pairs from the same task. A solved pair placed in front of the query acts as a visual prompt, so new tasks need no new head. Author-reported, it beats the best specialist on NYUv2 depth (RMSE 0.288 vs 0.330) and trails them on most others tab. 1.

Problem

Vision tasks differ in their output types, so there is no shared format that can be used as a task prompt the way language instructions are in NLP §1. Earlier generalists discretise continuous outputs into tokens, which adds quantisation error §1, §2. The concurrent visual-prompting work MAE-VQGAN showed in-context vision but reported no standard benchmark results §2.

Method

Output spaces as images §3.1

Every task output is redefined as a three-channel image with the same spatial size as the input, so pixel positions still line up with the input.

• Depth. Depth in metres, defined on [0, 10], maps to an 8-bit integer by a linear scale, and the same value is written to all three channels. At inference the three channels are averaged and the linear map is inverted §3.1.
• Semantic segmentation. With L classes, each class index is written as a three-digit number in base b = ceil(L^(1/3)), with margin m = floor(256 / b). Its digits are the three channel values. For ADE-20K (150 classes plus background), b = 6 and m = 42 §3.1. The exact colour formula differs between the main text and the appendix, see Pitfalls and Paper vs code.
• Keypoints. The 17 keypoint classes are encoded in two channels as small squares coloured the same way as segmentation classes. Localisation uses one channel with a Gaussian heatmap per keypoint, peak value 255 at the keypoint §3.1.
• Panoptic segmentation. It is split into semantic segmentation and class-agnostic instance segmentation. Instances get a colour from the location of their centre, using 16 × 20 × 20 blocks as in SOLO §3.1. At inference each colour is a kernel that is thresholded by distance, and each instance takes the majority semantic class inside its mask §3.1. Panoptic merging follows Panoptic FPN app. A.
• Restoration (denoising, deraining, low-light enhancement). Input and output are already RGB, so no conversion is needed §3.1.

Adding a task needs only new pairs in the training set, with no change to the architecture or loss §1, §3.1.

Masked image modelling §3.2

• Training sample. Two pairs of the same task, each already augmented separately, are concatenated into one sample §3.2. The code stacks them vertically, which gives 896 × 448 for 448 × 448 pairs §4.1 pairdataset.py. The second pair is drawn from the same task type pairdataset.py.
• Masking. Output patches are hidden with a learned mask token, using block-wise masking at a 75% ratio that the authors found to work well §3.2. The code's masking is described in Paper vs code.
• Encoder. The encoder is a ViT-L with 24 blocks, as the paper describes it §3.2. It processes the input and the masked output in parallel. After three blocks the two streams are merged, which saves nearly half the compute §3.2, app. B. The outputs of four evenly spaced blocks are concatenated §3.2.
• Decoder. A three-layer head maps each patch's features back to its 16 × 16 × 3 pixels: a linear layer, a 3 × 3 convolution and another linear layer §3.2, app. B.
• Loss. Smooth-L1 on the masked pixels only §3.2. L1 and L2 were worse in the ablation app. B tab. S2.

# Painter training step  [§3.2]
# Input:  two (input, output) pairs of one task, each (H, W, 3)
# Output: scalar loss on the masked output patches

def training_step(pair_a, pair_b, mask):
    # Stack the two pairs vertically, input over output  [§3.2]
    # Stacking direction is from the code, not the paper
    # [code: Painter/data/pairdataset.py#L100-L104]
    inputs = stack_rows(pair_a.input, pair_b.input)     # (2H, W, 3)
    targets = stack_rows(pair_a.output, pair_b.output)  # (2H, W, 3)

    # Embed both stacks with one patch embedding  [§3.2]
    input_tokens = patch_embed(inputs)                  # (B, h, w, C)
    target_tokens = patch_embed(targets)                # (B, h, w, C)

    # Replace masked output tokens with one learned mask token  [§3.2]
    target_tokens = where(mask, mask_token, target_tokens)

    # Run the blocks on both streams; merge them after block 3  [§3.2]
    # The paper says the features are added; the code averages them
    # [code: Painter/models_painter.py#L408-L415]
    tokens = concat_batch(input_tokens, target_tokens)  # (2B, h, w, C)
    kept = []
    for block_index, block in enumerate(blocks):
        tokens = block(tokens)
        if block_index == 2:
            tokens = (tokens[:B] + tokens[B:]) * 0.5
        # Keep the normalised outputs of blocks 6, 12, 18 and 24  [§3.2]
        # [code: Painter/models_painter.py#L416-L417]
        if block_index in (5, 11, 17, 23):
            kept.append(norm(tokens))
    features = concat_channels(kept)                    # (B, h, w, 4C)

    # Decode each patch back to its pixels  [§3.2]
    prediction = decoder(features)                      # (B, 2H, W, 3)

    # Smooth-L1 on masked output pixels, averaged over the masked count  [§3.2]
    # [code: Painter/models_painter.py#L459-L461]
    loss = smooth_l1(prediction, targets)[mask].mean()
    return loss

In-context inference §3.3

The prompt pair (an input and its solved output) is stacked above the query image. The query's output half is masked, and the model paints it. The authors' demo fills the prompt's output half with a dummy copy, which is never read painter_inference_demo.py. The prompt is chosen by one of three rules: random, searched, or learned §3.3, §4.3.

# In-context inference  [§3.3]
# Input:  a solved prompt pair and a query image
# Output: predicted output for the query, at its original size

def predict(model, prompt_in, prompt_out, query):
    # Stack the prompt pair above the query  [§3.3]
    # [code: Painter/util/painter_inference_demo.py#L111-L114]
    inputs = stack_rows(prompt_in, query)               # (2H, W, 3)
    targets = stack_rows(prompt_out, prompt_out)        # bottom half unused

    # Mask the query's output half, which the model paints  [§3.3]
    # [code: Painter/util/painter_inference_demo.py#L56-L58]
    mask = zeros_like(inputs_patches)
    mask[bottom_half] = 1

    # Return the painted bottom half, resized to the input's size
    # [code: Painter/util/painter_inference_demo.py#L65-L70]
    painted = model(inputs, targets, mask)[bottom_half]
    return resize(painted, original_size)

Learned prompts are the one alternative to picking an example from the training set. The prompt is treated as learnable tensors, the whole model is frozen, and the training loss is minimised over the prompt §3.3, §4.3.

Training setup §4.1

Training uses AdamW, cosine learning rate with warm-up, drop path 0.1 and a light augmentation: random resized crop and horizontal flip. The eight task types are sampled with weights 0.1, 0.2, 0.15, 0.25, 0.2, 0.15, 0.05 and 0.05 for NYUv2 depth, ADE-20K segmentation, COCO class-agnostic instances, COCO semantic segmentation, COCO keypoints, SIDD denoising, five deraining sets and LoL enhancement §4.1. The values are listed in the Hyperparameters table.

Hyperparameters

Paper values come from §4.1 and §3.2. Code values come from the official training script and the defaults in main_train.py. Where they differ, the Paper vs code section explains why.

Results

Author-reported, using ViT-L and the joint model tab. 1. Specialists are the best listed for each task. Generalists are the other rows of that table.

Bold marks where the paper says Painter is best or beats the listed generalist. The paper's caption gives the generalist backbones: Unified-IO XL with 2925M parameters, UViM with ViT-L, Pix2Seq v2 with ViT-B plus a transformer decoder tab. 1.

Joint vs separate training tab. 2. A model trained on all tasks does better than one trained per task on most metrics. The exceptions are keypoint AP (72.5 separate vs 72.1 joint) and SIDD SSIM (0.956 vs 0.954). Separate models get the same number of iterations as the joint model for that task.

Prompt choice tab. 3. Random, searched and learned prompts are close.

Few-shot segmentation tab. 4. On open-vocabulary FSS-1000, mIoU is 62.3 for Painter and 58.3 for MAE-VQGAN, the concurrent visual-prompting work §4.4.

Ablations tab. S2, on ADE-20K with a 3K-iteration schedule app. B.

Paper vs code

Code: Painter/ in the official repository at commit cfec845, which the paper links from its author block. The same repository also contains SegGPT code; that is a separate paper, not on this card, and none of it is used here.

• Mask ratio. The paper gives 75% block-wise masking §3.2. The training script sets 784 masked patches out of 1568 (56 × 28) per stitched sample train_painter_vit_large.sh, which is 50%. The argument --mask_ratio (default 0.5) is parsed but never read main_train.py. The paper's 75% cannot be matched to the released default.
• Mask location. The paper masks the task output image §3.2. The code generates the mask over the whole stitched grid and does not restrict it to the query's half, so masked patches can also fall in the prompt's output pairdataset.py masking_generator.py.
• Merge operation. The paper says features are added §3.2. The code averages them models_painter.py.
• Architecture. The paper calls the encoder a vanilla ViT §3.2. The code uses window attention with a 14 × 14 window in 16 of 24 blocks, global attention in the other 8, and relative position embeddings models_painter.py. It also adds learned stream tokens to the input and the masked target models_painter.py models_painter.py. None of this is in the paper.
• Decoder head. The paper describes a linear layer, a 3 × 3 convolution and a linear layer §3.2. The code puts a LayerNorm and a GELU between the convolution and the final 1 × 1 convolution models_painter.py.
• Loss. The paper gives smooth-L1 with no β §3.2. The code uses β = 0.01, divides by the masked count plus 0.01, and zeroes the loss weight of a sample when its unmasked output is nearly black models_painter.py.
• Weight decay and warm-up. The paper gives 0.05 and 5.4K warm-up iterations §4.1. The script uses the default 0.1 and one warm-up epoch of 15 main_train.py train_painter_vit_large.sh.
• Batch and length. The paper gives batch 2048 and 54K iterations §4.1. The script gives 256 per node and 15 epochs train_painter_vit_large.sh. Node count and dataset size are not stated, so the two cannot be reconciled.
• Augmentation. Colour jitter and the second crop are not in the paper main_train.py main_train.py.
• Semantic colour scheme. The main text gives floor(l / b²) × m §3.1. The appendix pseudocode and the released palette use 255 minus that value app. A, fig. S1 gen_color_ade20k_sem.py. The released colour decoder skips label 0, which is painted black gen_color_ade20k_sem.py.
• Keypoint colour scheme. The paper encodes classes as 9 × 9 squares with the segmentation scheme §3.1. The pose pipeline uses a two-channel palette with 5 levels per channel and a step of 51 custom_transform.py.
• Panoptic thresholds. The paper follows Panoptic FPN without numbers app. A. The code uses overlap 0.5, stuff area 8192 px and instance score 0.55 COCOPanoEvaluatorCustom.py.
• Depth units. The paper inverts a linear map to metres §3.1. The inference script writes values scaled by 10000 before averaging painter_inference_depth.py. From the training scale, the saved integer is therefore in millimetres (inferred from code, not stated).
• Matches. Instance block counts (16 × 20 × 20) §3.1 and the COCO colour base of 7 app. A match the code.

Pitfalls

• Two colour formulas. Implement the appendix and code version (255 minus the digits times the margin, black for ignore), not the main text's floor(l / b²) × m. Only the code version matches what the released model expects.
• Mask ratio and location. The 75% figure cannot be reproduced from the release. Use 784 of 1568 patches as the working default, and expect masked patches in the prompt half too.
• Pair order is not stated. The code stacks inputs above outputs and may draw the same pair twice as prompt and sample pairdataset.py.
• Sampling weights do not sum to one. The eight weights add up to 1.15. The sampler normalises them, so each type's share is its weight divided by 1.15 (NYUv2 depth 0.1 gives about 8.7%) pairdataset.py.
• Resolution and decoding. Inference resizes to 448 × 448 without keeping the aspect ratio, then resizes back: nearest neighbour for keypoints and instances, bilinear for depth and segmentation painter_inference_demo.py painter_inference_depth.py. The paper does not say how inputs are resized.
• Segmentation decoding. Each predicted pixel takes the class whose palette colour is nearest in L1 distance app. A, fig. S2. Blurred predictions therefore give mixed boundaries.
• Depth resolution. Each channel is an 8-bit value over [0, 10] m, so one level is about 3.9 cm before the three channels are averaged §3.1.
• Compute. A stitched sample has 1568 patches, twice a single 448 × 448 image's 784 §4.1. Merging after three blocks removes half the work for the rest of the encoder §3.2, app. B.
• Prompt tuning is under-specified. The number of steps, learning rate and choice of starting pair for learned and searched prompts are not stated §3.3, §4.3. The searched prompt needs a full pass over the training set §3.3.
• Panoptic output needs extra thresholds. Instance score, overlap and stuff-area thresholds come from the code, not the paper COCOPanoEvaluatorCustom.py.
• No variance reported. The paper gives no variance, error bars or repeated runs. Prompt choice gaps are small (ADE-20K mIoU 49.3 random vs 49.9 searched tab. 3), so they may be within run-to-run noise.
• Equations are inline. The paper has no numbered equations, so this card cites sections instead.

Lineage

• Builds on: masked image modelling as in BEiT, SimMIM and MAE §3.2; the vision transformer §3.2; SOLO for instance colouring §3.1; Panoptic FPN for merging app. A; the stitched two-pair prompt idea is the paper's own §3.2.
• Generalist comparators: Pix2Seq, Pix2Seq v2, Unified-IO, OFA and UViM §2 tab. 1.
• Visual-prompting comparator: MAE-VQGAN, cited as a concurrent work that does in-context inpainting §2 tab. 4.