PolyFormer: Referring Image Segmentation as Sequential Polygon Generation

In plain words

Give a model a photo and a sentence such as "the dog on the left bed", and it should outline exactly that object. Many systems colour in the object one pixel at a time. PolyFormer instead draws the outline as a short list of corner points, the way a person would trace a shape with a few clicks, and fills the shape in afterwards.

The model reads the image and the sentence together and writes out the points one at a time, each point informed by the ones before it. The points are real-valued positions, not rounded to a coarse grid, so small placement errors are avoided. The same model also outputs a box around the object, so it can answer "where is it" as well as "what is its outline".

TL;DR

Use PolyFormer if you need object masks from text queries and want a sparse, editable outline rather than a dense mask. The authors report better mask and box accuracy than the mask-based and sequence-based methods they compare with on RefCOCO, RefCOCO+ and RefCOCOg. The numbers are their own, from one training recipe, and the repository is missing parts of the evaluation.

Problem

Referring image segmentation needs a mask for the object a sentence names. Most instance-segmentation models predict each pixel independently with a dense convolutional output, which ignores the structure among the outputs §1. A sparse polygon is cheaper and matches common annotation formats, but earlier polygon models did not work well §1§2. Earlier sequence-to-sequence models for vision quantise coordinates into discrete bins and classify them, which adds quantisation error and treats localisation as classification §1§3.4.

Method

PolyFormer is a sequence-to-sequence model §3.1fig. 2. The output sequence holds the box corners first, then the vertices of one or more polygons, with separator and end tokens in between §3.2fig. 3. Polygons are ordered clockwise and start at the vertex closest to the image origin §3.2.

Inputs. An image is encoded by a Swin transformer, using its fourth stage §3.3. The expression is encoded by BERT §3.3. Both feature sets are projected to the same width with linear layers eq. 1 and concatenated into one sequence. A multi-modal transformer encoder processes that sequence §3.3.

Decoder. A transformer decoder with self-attention and cross-attention produces one state per output token §3.4. Each coordinate token is embedded by bilinear interpolation between the four nearest entries of a learned 2D codebook of 64 x 64 bins §3.4eq. 2. The codebook is what makes the input embedding continuous.

Heads. A linear class head predicts the token type: coordinate, separator or end §3.4eq. 3. A three-layer MLP with ReLU, followed by a sigmoid, predicts the continuous (x, y) of the next coordinate §3.4eq. 4. There is no quantisation on the output side §3.4.

Training. Each target token gets a loss §3.5eq. 5. Coordinate tokens get an L1 regression loss, weighted by lambda_box for the two box tokens and by lambda_poly for polygon vertices. Every token gets a label-smoothed cross-entropy term on its type, weighted by lambda_cls. Training uses polygon augmentation: with probability 0.5 the dense contour is resampled at a coarser step, which gives several sparse polygons for one object §3.5fig. 5.

Inference. Decoding starts from the start token. At each step the class head picks the token type. Coordinates are read from the regression head, and the result is fed back as the next input. Decoding stops at the end token §3.5. The mask is built from the predicted polygons §3.5. The code takes their union eval_utils.py.

Multi-task. The same model is trained on referring expression comprehension (boxes) and referring image segmentation (polygons), with box and polygon tokens in one sequence §3.2§4.2.

# Target sequence for one (image, expression) pair  [§3.2] [fig. 3]
# Input:  box as two corners, polygons as lists of (x, y) in [0, 1]
# Output: tokens with a type each: 0 coordinate, 1 separator, 2 end

def build_target(box, polygons):
    # Order every polygon clockwise and start it at the vertex nearest
    # the image origin  [§3.2]
    polygons = [order_polygon(poly) for poly in polygons]
    # Sort polygons by the distance of their start vertex to the origin
    # [§3.2]
    polygons = sort_by_start_distance(polygons)
    # Box corners come first, two coordinate tokens  [§3.2]
    tokens = [box[0], box[1]]
    types = [COORD, COORD]
    for index, poly in enumerate(polygons):
        # Separator before every polygon except the first  [§3.2]
        if index > 0:
            tokens.append(SEP)
            types.append(SEP)
        # Vertices of this polygon, one coordinate token each  [§3.2]
        tokens += poly
        types += [COORD] * len(poly)
    # End token closes the sequence  [§3.2]
    tokens.append(EOS)
    types.append(EOS)
    return tokens, types

# Forward pass for teacher-forced training  [§3.3] [§3.4] [eq. 1-4]
# Input:  image (512, 512, 3); expression token ids; target tokens
# Output: class logits (T, 3); regressed coordinates (T, 2)

def forward(image, expression, target_tokens):
    # Swin stage-4 features, flattened to visual tokens  [§3.3]
    visual = swin_stage4(image)              # (256, C_v) for 512 input
    # BERT word features for the expression  [§3.3]
    words = bert(expression)                 # (L, 768)
    # Project both to one width and concatenate  [§3.3] [eq. 1]
    memory = concat(proj_v(visual), proj_l(words))
    # Multi-modal encoder: 6 layers (base) or 12 layers (large)
    # [§4.2] [app. B]
    memory = encoder(memory)
    # Embed each shifted coordinate as a bilinear mix of four bins  [§3.4]
    # [eq. 2]; scale x in [0, 1] by 63 so bins span 64 grid positions
    # [code: data/refcoco_dataset.py#L123]
    inputs = embed_bilinear(shift_right(target_tokens), scale=63)
    # Causal self-attention plus cross-attention to the memory  [§3.4]
    states = decoder(inputs, memory)
    # Three-way class head: coordinate, separator, end  [§3.4] [eq. 3]
    class_logits = linear_cls(states)
    # Three-layer MLP with a sigmoid, so outputs lie in (0, 1)  [§3.4]
    # [eq. 4]
    coords = sigmoid(mlp_reg(states))
    return class_logits, coords

# Training loss for one sequence  [§3.5] [eq. 5]
# Input:  class_logits (T, 3), coords (T, 2), target coordinates and types
# Output: scalar loss

def training_loss(class_logits, coords, target, types):
    # L1 regression on coordinate tokens only  [§3.5] [eq. 5]
    total = 0.0
    for t in range(T):
        if types[t] != COORD:
            continue
        # Box corners use lambda_box; polygon vertices use lambda_poly
        # [§3.5] [§4.2]
        weight = lambda_box if t < 2 else lambda_poly
        total += weight * abs(coords[t] - target[t]).sum()
    # Label-smoothed cross-entropy over the three token types, weighted
    # by lambda_cls  [§3.5] [eq. 5] [app. B]
    total += lambda_cls * label_smoothed_ce(class_logits, types, eps=0.1)
    return total

# Polygon augmentation for one training sample  [§3.5] [fig. 5]
# Input:  dense_polygon (contour densified between vertices), polygon
# Output: polygon used as the training target

def augment(dense_polygon, polygon):
    # With probability 0.5 use a sparser polygon  [§3.5] [§4.2]
    if random() < 0.5:
        # Keep every k-th dense point; the paper gives no range for k
        # [§3.5]; the code draws k from 25 to 40  [code:
        # data/refcoco_dataset.py#L99-L102]
        k = random_int(25, 40)
        return dense_polygon[::k]
    return polygon

# Greedy decoding at inference  [§3.5] [code: tasks/refcoco.py#L165-L290]
# Input:  image, expression
# Output: box and polygons in [0, 1], or an empty list

def predict(image, expression):
    memory = encode(image, expression)
    tokens = [BOS]
    outputs = []
    # Stop after 210 steps; the code's cap, not stated in the paper
    for step in range(210):
        class_logits, coords = decode(tokens, memory)
        # Pick the most likely token type, no beam search  [code:
        # tasks/refcoco.py#L230]
        token_type = argmax(class_logits[-1])
        if token_type == COORD or (token_type == EOS and step < 6):
            # Read the coordinate from the regression head, clip at 1
            # [code: tasks/refcoco.py#L235-L237]
            x, y = min(coords[-1], 1.0)
            outputs.append((x, y))
            tokens.append(bin_id(x, y))
        elif token_type == SEP:
            outputs.append(SEP)
            tokens.append(SEP_ID)
        else:
            break
    return outputs

Hyperparameters

Values in the code column come from the training and evaluation scripts. The paper leaves several of them out; those rows say so.

Results

Author-reported. Input 512 x 512 §4.2. Base uses Swin-B and large uses Swin-L, both with a BERT-base text encoder §4.2. Pretraining on referring expression comprehension uses Visual Genome, RefCOCO, RefCOCO+, RefCOCOg and Flickr30k §4.2. Fine-tuning combines the RefCOCO, RefCOCO+ and RefCOCOg training splits, with validation and test images removed; ReferIt uses only its training set §4.2. mIoU is the mean of per-image IoU; oIoU is total intersection over total union §4.1. Prec@0.5 counts a box as correct if its IoU with the ground truth is above 0.5 §4.1.

Referring image segmentation, oIoU tab. 1

Referring image segmentation, mIoU tab. 1

The text states the margins over the previous best: 2.48 (RefCOCO val), 5.40 (RefCOCO+ val) and 4.52 (RefCOCOg val) for PolyFormer-L §1. These equal the differences between the rows above. The PolyFormer-L mIoU values match those stated in the introduction §1.

Referring expression comprehension, Prec@0.5 tab. 2

Zero-shot referring video segmentation, Ref-DAVIS17 tab. 3

PolyFormer runs frame by frame, with no fine-tuning on video §4.3. The dagger rows were trained on image datasets only; ReferFormer was trained on image and video data tab. 3.

Ablations (RefCOCO / RefCOCO+ / RefCOCOg val, mIoU, base model)

Sequence construction tab. 5:

Decoder tab. 4:

Coordinate codebook size, base model tab. 6 app. C:

The authors say the order of polygon vertices is the largest single gain (+12.43 from random order), followed by augmentation (+3.72) and multi-task training (+3.07) §4.4. The regression decoder adds +1.67 to +2.70 over the classification decoder §4.4 tab. 4.

Paper vs code

Official code: amazon-science/polygon-transformer. The paper names its project page in the title block, and that page links the repository. The owner, amazon-science, is the organisation of the authors' employer (AWS AI Labs, per the affiliations). Code commit 69fc728b2ec6a2b3595ec34db64074badcb19151.

Discrepancies, each checked against the code at that commit:

• Batch size. Paper: 160 pretrain, 128 fine-tune §4.2. Code, with 8 GPUs and per-GPU batches: base pretrain 20 x 8 x 8 = 1280 [code: run_scripts/pretrain/pretrain_polyformer_b.sh#L28-L29]; base fine-tune 16 x 8 x 8 = 1024 train_polyformer_b.sh; large 8 x 8 x 8 = 512 in both train_polyformer_l.sh pretrain_polyformer_l.sh. This assumes fairseq's --batch-size is per GPU; I did not verify that by running.
• Heads. Paper: 12 attention heads app. B. Code: 12 for base, 16 for large polyformer.py [code: models/polyformer/polyformer.py#L148].
• Pretraining loss weights. Not in the paper. Code: lambda_box 1 and lambda_cls 0 pretrain_polyformer_b.sh. With lambda_cls 0, the classification term is zero in pretraining, so the class head gets no gradient there. This follows from the code [code: criterions/label_smoothed_cross_entropy.py#L293]; I did not run it.
• Loss normalisation. Paper: the total loss is the sum of per-token terms §3.5. Code: the L1 term is averaged over coordinate values in the batch [code: criterions/label_smoothed_cross_entropy.py#L307-L309]; the class term is summed over tokens and divided by batch size [code: criterions/label_smoothed_cross_entropy.py#L293]; fairseq then divides all gradients by the criterion's sample size, which is the batch size [code: criterions/label_smoothed_cross_entropy.py#L223] [code: fairseq/fairseq/trainer.py#L831-L844]. Read literally, the class term is scaled by the batch size twice. This is inferred from reading the code; I did not run it.
• Label smoothing. Paper: label-smoothed cross-entropy with factor 0.1 app. B, no detail. Code: the class head has three outputs, so smoothing is spread over three types, with eps_i = 0.1 / 2 [code: criterions/label_smoothed_cross_entropy.py#L115]. Smoothing over a vocabulary would give a different loss.
• Input prompt. Paper: the expression alone is the text input §3.1. Code wraps it: which region does the text " ... " describe? [code: data/refcoco_dataset.py#L163].
• Image resize. Paper: images resized to 512 x 512 §4.2. Code: the shorter side is resized to 512 but both sides are capped at 512, so every image becomes 512 x 512 without keeping its aspect ratio [code: utils/transforms.py#L118-L122] refcoco_dataset.py.
• Augmentation step. Paper: "an interval randomly sampled from a fixed range" §3.5; range not given. Code: integer from 25 to 40, applied to the densified contour, with the first vertex always kept refcoco_dataset.py. The pretraining dataset class has no augmentation step [code: data/refcoco_pretrain_dataset.py].
• Decoding. Paper: autoregressive, no decoding details §3.5. Code: greedy, at most 210 steps, and an end token before step 6 is read as a coordinate [code: tasks/refcoco.py#L169-L170] refcoco.py [code: tasks/refcoco.py#L235]. The evaluation scripts pass generator flags such as --no-repeat-ngram-size=3 evaluate_polyformer_b_refcoco.sh, but eval_step calls task.inference_step directly eval_utils.py.
• Missing evaluation code. Ref-DAVIS17 zero-shot and ReferIt results have no scripts in the repository; the README lists only RefCOCO, RefCOCO+ and RefCOCOg evaluation README.md.
• Classification baseline. Table 4 compares against classification over bins tab. 4. I found no bin-classification coordinate head in the repository.
• Data preparation. Pretraining and fine-tuning data come from files that the README says to obtain from other projects, including SeqTR's instances.json and the refer annotations README.md. Those files are not in the repository.
• Licence. The repository's LICENSE is Apache-2.0 LICENSE. The project page says its design is adapted from Nerfies under CC BY-SA 4.0; that covers the web page, not the code. Vendored folders (for example fairseq/, refer/) carry their own licence files and are not covered by this card's code_license.

Pitfalls

• Coordinate grid. Normalised coordinates are multiplied by 63, not 64, before flooring and ceiling refcoco_dataset.py [code: data/refcoco_dataset.py#L129]. 64 bins span 64 grid positions, with spacing 1/63 of the normalised range. Using 64 as the scale shifts every bin.
• Normalisation reference. Polygon and box coordinates are divided by the original image width and height, not by 512 refcoco_dataset.py refcoco_dataset.py. Boxes are scaled back by the original size at evaluation refcoco.py. Ground-truth masks are scored at 512 x 512 refcoco_dataset.py.
• Mask rasterisation. Predicted vertices are cast to integer pixel positions before rasterising eval_utils.py. Rasterising sub-pixel vertices gives different IoU. Ground-truth masks are resized to 512 x 512 with nearest-neighbour interpolation refcoco_dataset.py. A polygon that fails to rasterise is counted as an empty mask eval_utils.py.
• Regression output range. The sigmoid bounds the output to (0, 1), but at inference only the upper bound is clipped refcoco.py. Clip both ends if you reimplement decoding.
• Final regression layer starts at zero. Its weights and bias are initialised to zero, so the initial output is 0.5 for every coordinate [code: models/polyformer/unify_transformer.py#L1052-L1054]. The paper does not say so.
• Half precision on the text path. The embedding LayerNorm casts its input to float16 on the non-code branch unify_transformer.py. Running in float32 without --fp16 may not match the released runs.
• Decoding cost. inference_step passes the whole prefix to the decoder at every step and keeps no incremental state refcoco.py. Generation cost grows roughly quadratically with output length.
• Augmentation. The target polygon for an object changes between epochs, so its vertex count is not fixed. The ablation reports +3.72 mIoU on RefCOCO from augmentation tab. 5.
• Box tokens by position. The loss picks box tokens by position: the first two targets of each sequence label_smoothed_cross_entropy.py. Any change to the sequence layout must keep the box first.
• Separator and end targets. Separator and end positions carry a [0, 0] placeholder in the regression target refcoco_dataset.py. They are excluded from the regression loss by token type, not by value.
• Hyperparameters outside the paper. Warm-up, weight decay, dropout, drop-path, normalisation and the flags --scale-attn, --scale-fc, --scale-heads, --add-type-embedding and --disable-entangle are set only in the scripts train_polyformer_b.sh [code: run_scripts/finetune/train_polyformer_b.sh#L93-L97]. The paper mentions none of them.
• Training data. Fine-tuning combines three splits and removes images that appear in validation or test create_finetuning_data.py create_finetuning_data.py. The fine-tuning script validates on RefCOCO val each epoch and keeps the best checkpoint by score [code: run_scripts/finetune/train_polyformer_b.sh#L18] [code: run_scripts/finetune/train_polyformer_b.sh#L89]. Any validation number from these runs is therefore selected on that split.
• Compute. Eight GPUs in every script, 20 pretraining epochs and 100 fine-tuning epochs at 512 x 512 train_polyformer_b.sh. The paper gives no memory or wall-clock figures.

Lineage

• Builds on: the Swin transformer for the visual encoder [/papers/2103.14030], and OFA, the sequence-to-sequence framework whose code the repository says it is based on [/papers/2202.03052] README.md.
• Compared against: SeqTR, LAVT, RefTr, CRIS, VLT, ReSTR, STEP, BRINet, CMPC, LSCM, CMPC+, MCN, EFN, BUSNet, CGAN, LTS, UniTAB, MDETR, UNTIER-L, VILLA-L, OFA-B and OFA-L in tables 1 and 2 §4.3. Ref-DAVIS17 comparison uses ReferFormer, CMSA+RNN, URVOS and CITD tab. 3. Among these, only OFA has a card in this wiki.
• Differs from quantised-coordinate sequence models such as those the paper cites for classification-based localisation §1§3.4: the decoder regresses coordinates directly.