GroupViT: Semantic Segmentation Emerges from Text Supervision

In plain words

Teaching a computer to label every pixel of a picture usually needs people to draw outlines around objects. GroupViT asks whether the model can learn where the objects are from only the captions that come with web images, such as "a dog runs across the grass". Nobody draws any outlines.

The model works in stages. It starts with small image patches and, at each stage, learns to merge similar pieces into fewer, larger regions. The merging is driven only by the caption-matching training signal: the picture as a whole must match its caption. Regions that tend to be named the same way end up together without anyone telling the model what they are.

Once trained, each final region gets a label by comparing its representation with the text embedding of a label name, such as "cat" or "road". Because the label names are only text, the same model can be pointed at a new list of classes without retraining. The results trail models trained with full pixel labels on some datasets, and the paper describes where the approach is weak.

TL;DR

• A hierarchical ViT groups image patches into arbitrary-shaped segments through learned group tokens and a differentiable hard assignment, trained with image-text contrastive losses only §3.1 §3.2.
• A multi-label variant adds caption nouns in short templates as extra positives, which the authors report as a key driver of grouping quality §3.2 tab. 1.
• Author-reported: zero-shot mIoU of 52.3 on PASCAL VOC 2012 and 22.4 on PASCAL Context, on a model trained on about 26M public image-text pairs §4.4 tab. 5.
• Worth reading if you need label-free or open-vocabulary segmentation. Check the Pitfalls first: the paper and the official code differ in the loss, the assignment, the data and the inference thresholds.

Problem

Semantic segmentation with deep networks is usually trained with per-pixel labels, which are expensive. The model also only recognises the classes it was labelled with §1. Earlier text-supervised work either used a pretrained detector or still needed segmentation labels for some categories, or relied on large private datasets §2. The authors ask whether a segmentation model can be learned from image-text pairs alone and then transferred to any label vocabulary without fine-tuning §1.

Method

GroupViT has two parts: an image encoder that groups patches, and a text encoder. Both are trained jointly with contrastive losses §3 fig. 2.

Image encoder §3.1. The image is split into non-overlapping patches that are linearly projected to tokens §3.1 Architecture. The Transformer layers are split into grouping stages. In each stage, learnable group tokens are concatenated with the segment tokens and passed through Transformer layers §3.1 Multi-stage Grouping. At the end of the stage, a Grouping Block assigns each segment token to one group token and merges the segments of each group into a new, larger segment §3.1 Grouping Block eq. 3-5 fig. 2b. After the last stage, a final Transformer runs over the segment tokens, and the segment tokens are averaged and projected to give the image embedding §3.1 eq. 2.

Grouping Block §3.1. Similarity between each group token and each segment token is computed as a scaled dot product, and a softmax is taken over the groups for every segment eq. 3. Gumbel(0, 1) noise, one sample per group, is added before the softmax eq. 3 §3.1. The assignment is then made hard (one-hot over groups) while keeping the gradient of the soft version, the straight-through estimator eq. 4 §3.1. Each new segment is the group token plus a projected, hard-weighted average of the segments assigned to it eq. 5. The paper reports that this hard assignment works better than the soft one §3.1 tab. 1.

Image-text losses §3.2. Each image has one caption. Matched pairs in a batch are positives and all other pairs are negatives. The text is embedded with a Transformer and the representation of its end-of-sentence token is used §3.2. The image-text loss is a symmetric InfoNCE over the batch with a learnable temperature §3.2 eq. 6.

Multi-label loss §3.2 fig. 3. For each caption, K nouns are drawn at random and each noun is put into several handcrafted sentence templates, giving K extra texts per image. Each image then has K positives in the batch. The loss adds two terms, one from the image side and one from the prompt side eq. 7. The total objective is the image-text loss plus the multi-label loss eq. 8. The paper sets K = 3 §4.1.

Zero-shot segmentation §3.3 fig. 4. The image is passed through the image encoder without the final average pooling, which yields one embedding per output segment. Each label name is turned into prompt sentences and embedded. Each segment is then assigned to the label with the highest normalised similarity §3.3. To map segments back to pixels, the paper multiplies the stage assignment matrices §3.3. Background is handled by thresholding the softmax-normalised similarity, with the thresholds in §4.1.

# GroupViT image encoder, two grouping stages  [§3.1] [§4.1]
# Input:  image (3, 224, 224)
# Output: image embedding (d,)

def group_vit_encode(image):
    # Split the image into non-overlapping patches and embed each  [§3.1]
    segments = patch_embed(image)                  # (N, 384), N = 196
    # Stage 1: 64 learnable group tokens, then 6 Transformer layers  [§4.1]
    group_tokens = learned_groups[0]               # (64, 384)
    segments, group_tokens = transformer(segments, group_tokens, depth=6)
    # Grouping Block: merge segments into 64 new segments  [§3.1] [eq. 3-5]
    segments = grouping_block(group_tokens, segments)  # (64, 384)
    # Stage 2: 8 group tokens, then 3 Transformer layers  [§4.1]
    group_tokens = learned_groups[1]               # (8, 384)
    segments, group_tokens = transformer(segments, group_tokens, depth=3)
    # Grouping Block: merge into 8 final segments  [§3.1] [§4.1]
    segments = grouping_block(group_tokens, segments)  # (8, 384)
    # Final Transformer over the segment tokens  [§3.1]
    segments = transformer(segments, depth=3)      # (8, 384)
    # Average the segments, then project to the shared space  [§3.1] [eq. 2]
    return mlp(mean(segments, axis=0))

# Grouping Block  [§3.1] [eq. 3-5]
# Input:  group_tokens (M_l, d), segments (M_prev, d)
# Output: new_segments (M_l, d)

def grouping_block(group_tokens, segments):
    # Scaled similarity between each group and each segment  [eq. 3]
    scores = W_q(group_tokens) @ W_k(segments).T   # (M_l, M_prev)
    # One Gumbel(0, 1) sample per group, added to its row  [eq. 3] [§3.1]
    noise = sample_gumbel(M_l)                     # (M_l,)
    # Softmax over groups, computed for each segment  [eq. 3]
    soft = softmax(scores + noise[:, None], axis=0)  # (M_l, M_prev)
    # Hard one-hot assignment with the soft gradient  [eq. 4] [§3.1]
    hard = one_hot(argmax(soft, axis=0)).T + soft - stop_gradient(soft)
    # Each group: its token plus the hard-weighted mean of its segments  [eq. 5]
    weights = hard / sum(hard, axis=1, keepdims=True)
    merged = W_o(weights @ W_v(segments))          # (M_l, d)
    return group_tokens + merged

# Image-text contrastive loss, symmetric  [§3.2] [eq. 6]
# Input:  image_emb (B, d), text_emb (B, d), tau (learnable scalar)
# Output: scalar loss

def contrastive_loss(image_emb, text_emb, tau):
    # L2-normalise both sides  [§3.2]
    image_emb = normalise(image_emb)
    text_emb = normalise(text_emb)
    # Similarity of every image with every text in the batch  [§3.2]
    logits = image_emb @ text_emb.T / tau          # (B, B)
    # Matched pairs sit on the diagonal; both directions are summed  [eq. 6]
    targets = arange(B)
    loss_i2t = cross_entropy(logits, targets)
    loss_t2i = cross_entropy(logits.T, targets)
    return loss_i2t + loss_t2i

# Multi-label contrastive loss with K prompts per image  [§3.2] [eq. 7]
# Input:  image_emb (B, d), prompt_emb (B, K, d), tau (scalar)
# Output: scalar loss

def multi_label_loss(image_emb, prompt_emb, tau):
    # Flatten the K prompts of every caption into one list  [§3.2]
    flat_prompts = reshape(prompt_emb, (B * K, d))
    # Image-to-prompts: the K positives of each image  [eq. 7]
    logits = image_emb @ flat_prompts.T / tau      # (B, B * K)
    loss_i2p = -mean(logsumexp(logits masked to own K) - logsumexp(logits))
    # Prompts-to-image: each prompt against all images  [eq. 7]
    logits_p = flat_prompts @ image_emb.T / tau    # (B * K, B)
    loss_p2i = cross_entropy(logits_p, own_image_index)
    return loss_i2p + loss_p2i

# Zero-shot semantic segmentation  [§3.3] [§4.1]
# Input:  image (3, H, W), label_names (C,) as text
# Output: label per segment, then per pixel

def zero_shot_segment(image, label_names):
    # Forward without final pooling: one embedding per segment  [§3.3]
    segment_emb = group_vit_segments(image)        # (8, d)
    # Turn each label name into prompt sentences and average them  [§3.3]
    text_emb = embed_prompted_labels(label_names)  # (C, d)
    # Softmax over labels for each segment  [§3.3]
    probs = softmax(segment_emb @ text_emb.T, axis=-1)
    # Segments with a low top probability are background  [§4.1]
    # Thresholds: 0.9 on VOC 2012, 0.5 on PASCAL Context  [§4.1]
    # Map segments to pixels with the product of stage assignments  [§3.3]
    return pixel_labels

Hyperparameters

Not stated in the paper: the output size of the projection head, the initial temperature, the Gumbel temperature, the optimiser betas and epsilon, the pretraining augmentations, and the number of GPUs. The code fills some of these in (see Paper vs code).

Results

Author-reported. Zero-shot mIoU on the validation splits, with a ViT-S GroupViT trained on CC12M and filtered YFCC, 2-stage §4.4 tab. 5.

The fully supervised rows fine-tune a segmentation head with pixel labels §4.4. The paper calls GroupViT's VOC result "approaching" the fully supervised DeiT row without any pixel labels §4.4.

Ablations (CC12M, 1-stage, VOC 2012 mask mIoU) §4.2.

Additional results (appendix). ImageNet, text-only training app. C.1 tab. C.1:

Mask probing, Jaccard similarity between the best mask (attention head or group) and the ground truth app. C.2 tab. C.2: Random 23.6; DeiT 24.6; MoCo (ImageNet) 28.2; DINO (ImageNet) 45.9; DINO (CC12M+YFCC) 41.8; CLIP (CC12M+YFCC) 28.6; GroupViT 51.8.

Oracle and COCO app. C.3 app. C.4 tab. C.3: with ground-truth matching of groups to classes, mIoU would be 73.7 on VOC 2012, 54.6 on PASCAL Context and 54.0 on COCO. The model's own mIoU is 52.3, 22.4 and 24.3.

RedCaps replacing YFCC app. C.5 tab. C.4: VOC 2012 50.8, PASCAL Context 23.7, COCO 27.5, against 52.3, 22.4 and 24.3 for CC+YFCC.

Paper vs code

Official code: NVlabs/GroupViT at commit 13b786155a1dfffe4703f40d028c92be58e1178d. The repository is linked from the paper's abstract [abstract]. It sits under the NVlabs GitHub organisation, and the author list gives NVIDIA affiliations [title block]. The code is under the NVIDIA Source Code License, which allows non-commercial research or evaluation use only LICENSE. Nothing from the code is quoted here; the citations describe behaviour only.

• Assignment logits are scaled. The paper's eq. 3 has no scaling factor. The code scales the dot product by the inverse square root of the head dimension group_vit.py group_vit.py.
• Assignment normalisation has a +1. The paper divides the weighted sum by the sum of weights eq. 5. The code adds assign_eps=1.0 to the denominator, because sum_assign defaults to False and the GroupViT model does not override it group_vit.py group_vit.py group_vit.py.
• Gumbel noise is training only. The code samples Gumbel noise only in training mode, and uses a deterministic hard argmax at inference group_vit.py. The paper describes the Gumbel-softmax without this distinction §3.1.
• Group tokens also depend on the previous stage. The paper says each stage learns its group tokens as parameters §3.1 Architecture. The code adds a projection of the previous stage's group tokens to the learnable tokens group_vit.py.
• The new segment is not eq. 5. The code adds the projected group tokens, not the original ones, and then applies an extra MLP with a residual connection, plus a normalisation and a linear layer when dimensions differ group_vit.py. A cross-attention block runs before the assignment group_vit.py. None of these are in the paper.
• Loss scale. The paper's image-text loss is the sum of the two directions, with no factor eq. 6. The code multiplies the sum by 0.5 multi_label_contrastive.py and does the same for the multi-label term multi_label_contrastive.py.
• Multi-label target. The paper's multi-label term uses a log of the summed positive scores over the total eq. 7. The code uses a soft-target cross-entropy with equal weight 1/K on each of the K positives multi_label_contrastive.py. These are different objectives; the code's is the one that was trained.
• Separate temperatures. The paper describes one learnable temperature §3.2. The code learns a second one for the multi-label term by default multi_label_contrastive.py multi_label_contrastive.py. Both are initialised at 0.07 and clamped at a maximum scale of 100 multi_label_contrastive.py multi_label_contrastive.py.
• Noun selection. The paper says K nouns are chosen §3.2. The code tags nouns and verbs, not only nouns builder.py, draws K at random builder.py, and when fewer than K are found it pads with the original caption builder.py. Templates are drawn at random per noun from the full ImageNet template set builder.py builder.py.
• Label embeddings are averaged over templates. The code averages each label's template embeddings and re-normalises the average multi_label_contrastive.py.
• Background thresholds differ. The paper uses 0.9 for VOC and 0.5 for PASCAL Context §4.1. The config sets 0.95 for VOC and 0.35 for Context [code: segmentation/configs/base/datasets/pascal_voc12.py#L40] [code: segmentation/configs/base/datasets/pascal_context.py#L40]. The inference code also caps the threshold at the largest group score group_vit_seg.py.
• Segmentation inference has more steps than the paper describes. The code assigns each pixel to its group by argmax, keeps only the five classes with the highest average-feature similarity, applies a second softmax, and multiplies by the group-level probabilities group_vit_seg.py. The code also has a comment marking one of these steps as unverified group_vit_seg.py.
• Sliding window. The VOC and Context configs use sliding-window evaluation with 448 x 448 crops and a 224 stride [code: segmentation/configs/base/datasets/pascal_voc12.py#L40]. The paper says each image is resized to a short side of 448 §4.1.
• Training data. The paper says CC12M and filtered YFCC §4.1. The default config trains on GCC3M, GCC12M and YFCC14M default.yml.
• Warm-up. The paper says 5 epochs of warm-up §4.1. The default config uses 2 default.yml.
• Learning-rate scaling. The code scales the learning rate, warm-up and minimum by the total batch size divided by 4096 main_group_vit.py. The paper's 4096 batch with 16 GPUs README.md matches this, but a different GPU count changes the learning rate.
• Augmentation. The config uses DeiT-style augmentation with random resized crops, RandAugment, colour jitter of 0.4 and random erasing with probability 0.25 default.yml. The paper does not state any augmentation §4.1.
• Drop-path. The config sets a stochastic drop-path rate of 0.1 for the image encoder group_vit_gcc_yfcc_30e.yml. The paper does not mention it.
• Text vocabulary. The paper says 49,152 BPE tokens app. A.1. The config sets 49,408 group_vit_gcc_yfcc_30e.yml.
• Compute. The README says the paper's pre-training used 16 V100 GPUs for two days README.md.

Pitfalls

• Text inconsistencies in the paper. §4.2 says the multi-label loss improves results "by 13.1% and 2.6%" for hard and soft assignment, respectively §4.2. Table 1 shows the reverse: 12.0 to 25.1 (+13.1) without hard assignment, and 36.7 to 39.3 (+2.6) with it tab. 1. §4.2 also says the 2-stage model improves CC12M+YFCC by about 7%; Table 3 shows +15.1 mask and +8.0 boundary mIoU tab. 3. Trust the tables.
• Boundary mIoU is the paper's own smoothness measure. The paper introduces it to quantify how smooth the maps are §4.2 tab. 3. It has not been checked against other work here, so compare it only with the paper's own numbers.
• Background handling is the weak point. On PASCAL Context, background classes such as ground, road and wall are grouped correctly but labelled wrongly, because captions rarely name them. Their IoU is about 1.5 app. C.3. An oracle that assigns labels with ground truth reaches much higher numbers tab. C.3. Plan a background rule that does not depend on the text alone.
• Foreground-only inference. The paper predicts foreground classes only and treats low-confidence segments as background §4.1. This is sensitive to the threshold, and the code's threshold differs from the paper's (see Paper vs code).
• Hard assignment matters for the reported results. With soft assignment, the paper reports a large drop in mIoU §4.2 tab. 1. Keep the straight-through estimator in any reimplementation.
• Number of groups matters. Going from 16 to 64 group tokens raises mIoU in the paper's ablation, but 64 is also the value used in the headline model tab. 2 §4.1.
• Costs are hidden. Training uses large batches of image-text pairs, a learning-rate schedule tied to batch size, and mixed precision in the code main_group_vit.py default.yml. The pretraining cost reported by the authors is 16 GPUs for two days README.md.
• Public data only. The paper says its model is trained on public, noisy image-text data, unlike some related work that uses large private datasets §2. The code's data list includes GCC3M, which the paper does not mention default.yml.
• Unstated details to decide yourself. The projection-head output size, the initial temperature, the Gumbel temperature, the optimiser betas and epsilon, and the augmentations are not stated in the paper. The code's choices for them are listed under Paper vs code.

Lineage

• Builds on: Vision Transformer (ViT) for the patch backbone §3.1; CLIP for the contrastive image-text loss, its text encoder and the prompt templates §3.2 §4.1; Slot Attention for the one-iteration-of-grouping analogy §3.1; Gumbel-softmax and the straight-through estimator for the hard assignment §3.1.
• Compared against: CLIP with non-parametric grouping (pixel-wise, K-means, mean-shift, spectral clustering) §4.4 tab. 4; DeiT, DINO and MoCo with fully supervised or self-supervised transfer §4.4 tab. 5; DINO and MoCo for mask probing app. C.2 tab. C.2.