ReSTR: Convolution-free Referring Image Segmentation Using Transformers

In plain words

Referring image segmentation means this: you give a model a photo and a sentence that describes one object in it, and it draws the outline of that object. ReSTR does this with transformers alone, with no convolutional layers.

The photo is cut into a grid of small patches, and the sentence is split into words. Each side is first read by its own transformer, so that parts of the same input can relate to each other over long distances. The two sides are then mixed by attention, so each word can look at every patch and each patch can look at every word. A special learned token ends up acting as a classifier tailored to the sentence: it scores each patch for whether it belongs to the described thing. A small decoder turns that coarse patch map into a pixel-level mask.

The authors report better scores than earlier methods on four public benchmarks, though not on every split.

TL;DR

ReSTR is a transformer-only model for referring image segmentation. It encodes the image (ViT-B/16 at 480 px) and the sentence (GloVe-initialised, 6 layers) separately, fuses them in two transformer stages that end in a learned classifier token, and decodes a coarse patch map to pixels §3, §4.1. It has the best reported IoU on most splits of four benchmarks, with no CRF post-processing tab. 1. Several implementation details differ from the equations, so read Paper vs code before reimplementing.

Problem

Earlier referring-segmentation models encode the image and the sentence with CNNs and RNNs, then fuse them by concatenation followed by convolution §1. The authors argue that convolution and recurrence have trouble with long-range relations between entities in the sentence, and that concatenation-convolution is a fixed, hand-designed way to fuse the two modalities §1. ReSTR replaces both parts with transformers §1.

Method

The model has three parts §3, fig. 2:

1. Feature extraction §3.1. The image is split into non-overlapping 16×16 patches, each linearly projected, with a learned 1-D position embedding added §3.1. A ViT encoder produces patch features z_v. Each transformer block is pre-norm: layer norm, multi-head self-attention, a residual; then layer norm, an MLP, a residual eq. 1 eq. 2. Attention is softmax(q kᵀ / √D_h) with D_h = D/k eq. 5, heads are concatenated and projected eq. 3, and each head outputs A v eq. 4. The sentence is embedded with GloVe, given sinusoidal position encodings, and passed through a transformer encoder to give z_l §3.1.
2. Multimodal fusion §3.2. Both feature sets are layer-normalised and linearly projected to a common width D. A visual-linguistic encoder runs over the concatenated tokens and returns updated visual and language features eq. 7. A linguistic-seed encoder then runs over the updated language tokens plus a learned seed token e_s, initialised randomly §3.2. The output seed token, e_s', is the adaptive classifier eq. 8. The authors argue that the seed should read only language and not the image directly, because background patches would otherwise corrupt the classifier §3.2. Their ablation supports this choice §4.3 tab. 3.
3. Coarse-to-fine decoding §3.3. Each visual token is scored by a dot product with e_s', scaled by 1/√D and passed through a sigmoid eq. 9. The visual tokens are then weighted by these scores eq. 10. The decoder concatenates the visual tokens with the weighted tokens and runs K = log₂P upsampling blocks, which gives K = 4 for P = 16 §3.3 §4.1. Each block doubles the resolution and halves the channels, and a final linear layer outputs one channel. At inference only the pixel map is used §3.3.

Training uses two binary cross-entropy terms §3.3. Patch targets come from average-pooling the ground-truth mask over each patch and thresholding at τ eq. 11. The total loss is λ times the patch loss plus the pixel loss eq. 12. The paper has no algorithm box. The pseudocode below is our restatement of §3 and §4.1, not a published algorithm.

# Training step for one batch  [§3, §4.1]  (restated, not an algorithm box)
# Input:  image (B, 3, 480, 480); words (B, 20); mask (B, 480, 480)
# Output: scalar loss  [eq. 12]

def training_step(image, words, mask):
    # 1. Vision tokens: 16 px patches, 30 x 30 = 900 tokens  [§3.1] [§4.3]
    visual = vision_encoder(image)              # (B, 900, 768)
    # 2. Word tokens: GloVe, sinusoidal positions, 6 layers  [§3.1] [§4.1]
    language = language_encoder(words)          # (B, 20, 300)

    # 3. Normalise each stream, then project both to width D  [§3.2]
    visual_in = project_v(normalize(visual))    # (B, 900, D)
    language_in = project_l(normalize(language))  # (B, 20, D)

    # 4. Visual-linguistic encoder on the joint sequence  [§3.2] [eq. 7]
    joint = transformers_vl(concat(visual_in, language_in))
    visual_fused = joint[:, :900]               # z_v'
    language_fused = joint[:, 900:]             # z_l'

    # 5. Linguistic-seed encoder: the learned seed reads z_l' only
    #    [§3.2] [eq. 8]
    seed_out = transformers_ls(concat(language_fused, seed))
    classifier = seed_out[:, -1]                # e_s', (B, D)

    # 6. Patch scores: dot product scaled by 1/sqrt(D), then sigmoid
    #    [§3.3] [eq. 9]
    patch_prob = sigmoid(einsum(visual_fused, classifier) / sqrt(D))

    # 7. Weight each visual token by its patch probability  [§3.3] [eq. 10]
    masked = visual_fused * patch_prob[:, :, None]

    # 8. Decode to full resolution: 4 upsampling blocks, 30 -> 480  [§3.3]
    #    pixel output activation is not stated  [§3.3]
    pixel_logits = decoder(concat(visual, masked))   # (B, 480, 480)

    # 9. Patch targets: average the mask per 16 x 16 patch, threshold at
    #    tau = 0.8  [§3.3] [eq. 11] [§4.1]
    patch_target = avg_pool(mask, 16) > 0.8

    # 10. Loss: patch term weighted by lambda = 0.1  [§4.1] [eq. 12]
    loss_patch = bce(patch_prob, patch_target)
    loss_pixel = bce_with_logits(pixel_logits, mask)
    return 0.1 * loss_patch + loss_pixel

Key equations, all cited above: eq. 9 patch scores, eq. 11 patch labels, eq. 12 total loss.

Hyperparameters

Results

Author-reported. Setup: four benchmarks (ReferIt, UNC, UNC+, Gref); ViT-B/16 pre-trained on ImageNet-21K; comparison against 13 earlier methods; no CRF post-processing for ReSTR §4.1 tab. 1. Scores are cumulative IoU in percent.

Expression-length analysis: on Gref ReSTR is best in three of four length groups. The 1–5 word group is the exception, where ACM is ahead by about 1.2 points app. A tab. 2. Gref scores by length group (1–5 / 6–7 / 8–10 / 11–20): ReSTR 58.72 / 53.47 / 53.96 / 51.91, ACM 59.92 / 52.94 / 49.56 / 46.21 tab. 2.

Gref ablation, IoU on val tab. 4:

Fusion-encoder variants, Gref val tab. 3: vanilla joint encoder 51.27; independent encoder 45.89; the proposed CME 52.81; CME with weight sharing 52.79. Cost tab. 5: ReSTR-CME has 122.87M parameters and 52.29G MACs at 320 × 320, against ACM's 232.78M and 124.68G with 51.93 IoU on Gref val. Note that ACM and the other rows of tab. 5 use CRF post-processing, which ReSTR does not.

Paper vs code

The paper has no code link, but its project page links the repository at github.com/SouthFlame/restr, and the repository README names the authors README.md. The commit used is 371f102d…. The repository has no LICENSE file, so it may be cited but not quoted.

• Patch-label rule. Paper: average-pool the mask per patch, then threshold at τ = 0.8 §3.3 eq. 11. Code: a convolution with all-ones weights sums the pixels, then thresholds at 0.8 torchutils.py. The code therefore labels a patch positive if it contains at least one foreground pixel, while the paper requires that 80% of its pixels are foreground. These give very different positive sets.
• Patch score. Paper: σ(z_v' e_s'ᵀ / √D) eq. 9. Code: a raw dot product of two learned projections, with no sigmoid and no 1/√D fusion_module.py. The projections are not in the paper fusion_module.py. self.scale only sets their initialisation fusion_module.py. The sigmoid enters only through the BCE-with-logits loss train_restr.py.
• Masking. Paper: weight by the sigmoid probability ŷ_p eq. 10. Code: weight by the raw logit fusion_module.py.
• Decoder input. Paper: the decoder input is z_v, the visual features §3.3. Code: it is the projected, normalised visual tokens fusion_module.py fusion_module.py.
• Padding. Code zeroes padded language tokens inside the fusion step fusion_module.py. The paper does not mention this.
• Vision learning rate. Code: 0.1× base LR for the vision encoder restr.py. The paper does not mention this.
• Decoder activation. Code: bilinear 2× upsampling, a 1×1 convolution and LeakyReLU blocks.py. The paper says only "activation function" §3.3.
• Fusion depth. The config sets 2 layers per sub-encoder config.yml. The --tr_n_layers flag is parsed but never read train_restr.py. The default therefore comes from the config, not the flag. Reading 2 per sub-encoder as 4 in total is our inference from the code.
• Input size. Config lists 384 config.yml. Training overrides this with the 480 input size train_restr.py. The paper says 480 §4.1.
• Checkpoint choice. The training loop keeps the checkpoint with the best cumulative IoU on the validation split, and also saves after 90% of steps train_restr.py. The paper does not say which checkpoint it reports.
• Evaluation threshold. Predicted pixels are positive when the logit is at least 1e-9 evaluate.py evaluate.py. That is roughly a sigmoid of 0.5. The paper states no threshold.
• Augmentation. None is applied. A flip is defined but never used referit_dataset_vit.py. The paper states none §4.1.
• Pre-training source. The code loads timm's default vit_base_patch16_384 weights factory.py. We did not verify that these match the ImageNet-21K pre-training named in §4.1.

Matches: learning rate 1e-5, batch 8, 400,000 steps, 40,000 warm-up steps, weight decay 5e-4, λ = 0.1 on the patch term, τ = 0.8, and the loss structure eq. 12 train_restr.py.

Pitfalls

• Patch labels are ambiguous. The paper's averaging rule and the code's summing rule (above) give different positive sets. Decide which one you implement and test it. The paper's text is the stated spec.
• Unstated activations. The paper does not say whether the pixel head is followed by a sigmoid, or which activation the decoder uses. The BCE calls imply logits for the pixel head and a sigmoid on the patch head.
• Attention typo. The multi-head equation repeats SA_1 where it should list SA_1 to SA_k §3.1 eq. 3. Implement k distinct heads. Note that the paper also uses k for the number of heads and K for decoder blocks; they are different.
• Model selection and reporting share a split. The fusion-design choices and the decoder ablations are scored on Gref val §4.3 tab. 3 tab. 4. The headline Gref number is also on val tab. 1. Checkpoint selection also uses val train_restr.py. Expect the reported gains to be optimistic on that split.
• Small margins, no variance. On Gref val, ReSTR beats the next method by 0.08 points (54.48 vs 54.40) tab. 1. The paper reports single runs with no error bars or seed count §4. The code default seed is 1 train_restr.py.
• Claims are broader than the table. The abstract says ReSTR outperforms all existing models on all public benchmarks. Table 1 shows VLT ahead on UNC+ testB (49.36 vs 48.27) tab. 1. The paper's own text makes the same exception §4.2.
• Baselines differ in setup. Several baselines use DenseCRF post-processing and COCO-pre-trained backbones §4.2 tab. 1. ReSTR uses neither. The comparison is therefore not like-for-like, as the authors also note §4.2.
• Compute. Each fusion pass attends over 920 tokens (900 visual plus 20 language) §4.1 §4.3. The paper dropped the 6-layer decoder run for memory §4.3. MACs in tab. 5 are at 320 × 320, while training uses 480 × 480 tab. 5 §4.1.
• Hyperparameter tuning. The authors state that they did not tune on test data, and that the chosen λ and τ are not the best values on val app. B.
• Environment. The repository is pinned to Python 3.7 and PyTorch 1.13.1 restr.yaml README.md. The training script calls wandb.init on every run train_restr.py, so it needs a wandb setup or edits.

Lineage

• Builds on: the ViT and transformer encoder blocks, which the paper cites but this card does not link §3.1.
• Compared against: LSTM-CNN, RMI, DMN, RRN, CMSA, STEP, BRINet, LSCM, CMPC, ACM, BUSNet, LTS and VLT tab. 1. ACM is the closest recent comparison and the strongest prior on the length-group analysis app. A tab. 2.
• Superseded by: not known from content already in this wiki.