Scaling Open-Vocabulary Image Segmentation with Image-Level Labels

In plain words

Most segmentation models only know the labels they were trained with, such as the 20 or 150 classes of a benchmark. Models that understand arbitrary words, such as CLIP and ALIGN, can say what is in an image but not where each thing is at the pixel level. This paper asks how to get both: segment an image into regions for any text you type, using cheap image-caption data rather than expensive pixel labels.

The idea is to first split the image into a small set of candidate regions, each with its own feature vector. The regions are learned from class-agnostic masks, but they carry no class name. Then the model matches the words of a caption to these regions, so "dog" and "grass" in a caption each latch onto a few regions. At test time, a class name is embedded with the same text encoder and each pixel takes the name whose best-matching region covers it.

Because the region step does not depend on a fixed label set, the same model can be asked about new names and tested on other datasets without retraining. The headline claim is that this beats earlier open-vocabulary segmentation models on PASCAL, and generalises better than a model trained on pixel labels from COCO.

TL;DR

OpenSeg represents an image as N mask proposals with features, instead of one vector or a dense per-pixel map. It is trained with a Dice loss on class-agnostic COCO masks plus a caption grounding loss, and can add caption-only data through teacher pseudo-masks §3.3. The paper reports 72.2 mIoU on PASCAL-20 with EfficientNet-B7 and prompt engineering tab. 2, but these are author-reported numbers and the baselines are partly copied from other papers. There is no official code.

Problem

Per-pixel open-vocabulary models such as LSeg need pixel-level class labels, which are costly to collect at scale and limited to the labelled vocabulary §1 §2. Image-level models such as CLIP and ALIGN are trained on image-caption pairs, which are cheap, but they represent an image with a single vector, so location is lost §1 fig. 3. The paper argues the missing piece is a visual grouping step that turns pixels into regions before text alignment §1.

Method

Overview fig. 3: an image becomes N = 128 mask proposals, each with a feature vector, so Z lies in R^(N x D) §3 app. E. The model has three parts:

1. Mask proposals learned from class-agnostic COCO masks §3.1.
2. Region-word alignment learned from image-caption pairs with a grounding loss §3.2.
3. Optional caption-only data, labelled with pseudo masks from a teacher model §3.3.

At inference, the region features are compared with class-name embeddings and combined with the masks to give per-pixel labels §3.4.

3.1 Learning segmentation masks

• Features: an FPN fused to P2 resolution, following MuST §3.1. A convolution and fully connected layers produce F_s, to which a learnable position embedding PE is added §3.1.
• Mask decoder: N randomly initialised queries q^0 attend to F_s + PE with cross-attention, in the style of DETR, Max-DeepLab and MaskFormer §3.1. The module is stacked T = 3 times app. C. Adding self-attention among queries did not help app. C.
• Masks: each query gives one soft mask by a sigmoid of the dot product between the query and F_s + PE §3.1. No code is available, so the placement of PE is taken from the text alone.
• Labelled masks: class-agnostic COCO panoptic masks, M per image §4.1.
• Loss: Dice coefficient between predicted and labelled masks eq. 1. Since N > M, each labelled mask is matched to its best predicted mask, and the remaining proposals are not penalised §3.1. There is no "no object" class, so predicted masks may also cover unlabelled regions §1.

\mathcal{L_S} = \frac{1}{M} \sum_{j=1}^{M} \left( 1 - \max_i \mathrm{Dice}(s_i, s_j^{l}) \right) eq. 1

3.2 Visual-semantic alignment with masks

• Caption filtering: keep only nouns and adjectives. Each kept word is dropped with probability 1 - kp, with kp = 0.5 §3.2 §4.1.
• Region features: F_z uses the same architecture as F_s §3.2. Each region is pooled from the image features by its mask, without normalisation in the stated formula §3.2.
• Word features: a frozen pre-trained text encoder, the BERT-Large of ALIGN §4.1.
• The grounding score uses a softmax over regions with a learnable temperature tau eq. 2, and a cosine similarity between regions and words. The image-caption score G averages these over the kept words eq. 3.
• Grounding loss: a symmetric softmax over the mini-batch, in both the image-to-caption and caption-to-image directions eq. 4.
• Total loss: L = L_G + alpha * L_S eq. 5. The value of alpha is not stated.

\sigma(\mathbf{x})_i = \frac{e^{x_i / \tau}}{\sum_{j} e^{x_j / \tau}} eq. 2

G(I_b, C_b) = \frac{1}{K} \sum_{j=1}^{K} \sum_{i=1}^{N} \sigma\big(g(\mathbf{z}, w_j)\big)_i \cdot \langle z_i, w_j \rangle eq. 3

\mathcal{L_G} = -\frac{1}{|B|} \sum_{b=1}^{|B|} \Big( \log \sigma\big(G(\mathbf{I}, C_b)\big)_b + \log \sigma\big(G(I_b, \mathbf{C})\big)_b \Big) eq. 4

\mathcal{L} = \mathcal{L_G} + \alpha \mathcal{L_S} eq. 5

The paper states that alpha = 0 (no mask loss) gives poor performance, so the class-agnostic mask annotations matter §3.2.

3.3 Learning from caption-only data

• Teacher: trained on a segmentation dataset with the mask loss only §3.3.
• Pseudo labels: the teacher annotates a large image-text dataset (Localized Narratives) with pseudo masks §3.3 §4.1.
• Student: trained on a mix of human and pseudo labels, with the segmentation loss applied to all examples §3.3 app. G.

3.4 Inference

• Each class name is embedded with the text encoder. Multi-word names keep all their word embeddings, giving K word embeddings §3.4.
• Region logits are cosine similarities between words and regions, in R^(K x N) §3.4.
• Pixel logits are the region logits multiplied by the masks §3.4.
• Each pixel takes the argmax over classes eq. 6.
• Ensembling: several text queries per class (synonyms, subcategories, plurals), taking the maximum score §4.3 app. I.

pred[i, j] = \argmax_k \, \mathbf{y}[k, i, j] eq. 6

The following pseudocode follows the text. The learning-rate schedule, optimiser and loss weights are in the Hyperparameters table.

# OpenSeg training step  [§3.1] [§3.2] [eq. 5]
# Input:  images (B, 3, H, W); captions: B lists of words
#         labelled_masks: per image, M class-agnostic masks (M, H, W)
# Output: scalar loss

def training_step(images, captions, labelled_masks):
    # Features for mask prediction (F_s) and for region pooling (F_z)  [§3.1]
    feats_s = mask_features(images)           # (B, H, W, D), D = 640
    feats_z = align_features(images)          # (B, H, W, D), same design
    # N learned queries cross-attend to F_s + PE, stacked T = 3 times  [app. C]
    queries = mask_decoder(feats_s + pos_embed)       # (B, N, D)
    # Sigmoid of the dot product gives one soft mask per query  [§3.1]
    pred_masks = sigmoid(queries @ (feats_s + pos_embed))  # (B, N, H, W)

    # Mean over labelled masks of (1 - best Dice)  [eq. 1]
    # Averaging over images is not stated; this averages over B and M
    dice = dice_matrix(pred_masks, labelled_masks)   # (B, N, M)
    seg_loss = mean(1 - max(dice, axis=N))

    # Mask pooling: region = sum over pixels of mask * feature  [§3.2]
    regions = einsum(pred_masks, feats_z)     # (B, N, D), not normalised
    # Word features from a frozen BERT-Large text encoder  [§4.1]
    # Nouns and adjectives only; each word dropped with prob. 1 - kp  [§3.2]
    words = text_encoder(kept_words(captions))   # list of (K_b, D)

    # Grounding loss in both directions over the batch  [eq. 4]
    grounding_loss = 0
    for index b in batch:
        # Caption b against all images, and image b against all captions
        to_captions = [image_caption_score(regions[i], words[b])
                       for i in batch]
        to_images = [image_caption_score(regions[b], words[j])
                     for j in batch]
        grounding_loss -= log_softmax(to_captions)[b]
        grounding_loss -= log_softmax(to_images)[b]
    grounding_loss = grounding_loss / batch_size

    # Total loss; alpha is not stated in the paper  [eq. 5]
    return grounding_loss + alpha * seg_loss


def image_caption_score(regions, words):
    # Cosine similarity between N regions and K words  [§3.2]
    similarity = cosine(regions, words)       # (N, K)
    # Softmax over regions for each word, temperature tau  [eq. 2]
    weights = softmax(similarity / tau, axis=N)   # (N, K)
    # Weighted similarity, averaged over the K words  [eq. 3]
    return mean(sum(weights * similarity, axis=N))

# OpenSeg inference  [§3.4] [eq. 6]
# Input:  image (3, H, W); class_names: list of K names
# Output: label map (H, W)

def predict(image, class_names):
    # Mask proposals and pooled region features for the image  [§3.4]
    masks = mask_proposals(image)             # (N, H, W), sigmoid
    regions = pool_regions(masks, image)      # (N, D)
    # One embedding per class name; multi-word names keep all words  [§3.4]
    words = text_embeddings(class_names)      # (K, D)
    # Cosine similarity between each word and each region  [§3.4]
    region_logits = cosine(words, regions)    # (K, N)
    # Per-pixel logits: region logits times masks, summed over regions  [§3.4]
    pixel_logits = region_logits @ masks.flatten(1)   # (K, H*W)
    # Argmax over the class dimension  [eq. 6]
    return argmax(pixel_logits, axis=K).reshape(H, W)

Hyperparameters

Results

Author-reported. Setups differ between rows: backbone, training data and baselines are as stated per row.

Zero-shot evaluation on the PASCAL-Context and PASCAL-20 benchmarks tab. 2 (mIoU, PC = PASCAL-Context, A = ADE20k, number = class count):

Mask proposal recall at IoU 0.5 / 0.7 / 0.9, 128 proposals per image tab. 1:

Mean IoU and grounding mIoU on the ADE20k and PASCAL-Context sets with ResNet-101 fig. 5 (the figure's bottom table; "grounding" uses only the ground-truth classes of each image as queries):

Ablations (COCO training unless noted, mIoU on A-847 / PC-459 / A-150 / PC-59):

Further results in the appendix: cross-dataset recall on MSeg datasets app. D; batch size app. F; mask loss on caption data app. G; keep probability app. H; ensembling and prompt engineering, where COCO + Loc. Narr. goes from 6.8 to 7.6 mIoU on A-847 app. I.

Headline comparison: the 72.2 PAS-20 result is 19.9 above the 52.3 LSeg ViT-L figure §4.3 tab. 2. The paper also reports OpenSeg with ResNet-101 above LSeg+ by 2.7 mIoU on PC-459 and 1.9 mIoU on A-847 §1.

Paper vs code

No official code.

Searched for an official repository in this order:

• The arXiv abstract page: comment "Accepted at ECCV 2022", no code or project link.
• The paper's LaTeX source: no code URL. The appendix says the authors will open-source their modified class-name list app. I; that is a promise, not a release.
• GitHub repository search for "OpenSeg" and for the title. Nothing is owned by Google Research or by an author. The "openseg-group" repositories belong to a different project (OCNet/OCRNet), so they were not used.

Pitfalls

• Loss weight alpha is not stated. The total loss is L_G + alpha * L_S eq. 5, but no value is given, and the text only says alpha = 0 is much worse §3.2. Expect to tune it.
• Dice loss wording. The text says to "maximize" the Dice of the best match §3.1, but eq. 1 is a minimisation of 1 - max Dice. Implement the minimised form. The Dice formula itself is cited to V-Net and not written out; whether it is computed on soft sigmoid outputs or on thresholded masks is not stated.
• Mask pooling is not normalised. The pooling formula is a plain sum over pixels weighted by the mask §3.2. Whether the sum is divided by the mask area is not stated, and the choice changes the scale of region features for large and small masks.
• Averaging over images in eq. 1. The equation averages over labelled masks per image. How the result is averaged over the batch is not stated.
• Keep probability conflict. The main text sets kp = 0.5 §4.1. The appendix says nouns are kept with probability 0.75 app. H. The ablation row kp = 0.75 (6.8, 11.2, 24.8, 45.9) matches the headline COCO + Loc. Narr. result, while kp = 0.5 gives 7.0, 10.7, 25.0, 46.0. Which value produced the headline numbers cannot be settled from the text.
• Backbone-init claim does not match its table. The text says the NoisyStudent initialisation is worse "on all benchmarks" by less than 0.5 mIoU §4.4. In tab. 3 its PC-59 score (46.9) is 1.0 above the ALIGN score (45.9). The table labels the scratch run "Random init.", while the text says "from scratch".
• Loss over local or global batch. By default the contrastive loss is computed per core over 16 examples §4.1. The "sync" variant over the global batch is better on A-847, PC-459 and A-150 but worse on PC-59 by 1.1 app. F. Pick one deliberately; the strongest model uses sync §4.3.
• Baselines are partly copied. The LSeg numbers in tab. 2 come from the LSeg paper, with a different backbone and setup, not from a re-run here. The 19.9 gain is the prompt-engineered EfficientNet-B7 row against that copied ViT-L number §4.3.
• Grounding mIoU is an oracle setting. It uses only the ground-truth classes of each image as queries §4.3. It is not comparable with mIoU, which uses a fixed label set §4.3.
• Class-name curation is part of the gains. The ensembling step uses hand-made synonym lists and context additions, such as "drinking glass" for "glass" and "ceiling fan, floor fan" for "fan" §4.3 app. I. The reported 2.6 mIoU average gain §4.3 comes from that curation, not from the model alone.
• Argmax over possibly negative scores. Pixel logits are cosine values times sigmoid masks §3.4, so they can be negative. The paper does not say how pixels with low or negative scores are handled, such as a background class (for PASCAL VOC the background is assigned by a rule, see below).
• PASCAL VOC background rule. Pixels predicted as PASCAL-Context classes that are not in PASCAL VOC are set to background §4.1.
• Unlabelled proposals are not penalised. Only the best match per labelled mask is trained, so proposals can overlap and cover unlabelled regions §3.1 fig. 4. The paper calls this a feature, but downstream code must decide how to resolve overlaps.
• Unlisted pseudo-label details. The teacher architecture, pseudo-label thresholds and any filtering are not stated §3.3. Only the mixing of human and pseudo labels is described.
• Cost. Training uses 64 TPU v3 cores with a 1024 global batch for 30k to 60k steps, and the strongest model uses an ALIGN-initialised EfficientNet-B7 §4.1. Initialising from scratch needs 180k steps and a retuned learning rate of 0.32 §4.4.
• Appendix letters. The LaTeX source uses \begin{appendix} but does not define it, so the section letters A to J are assumed from source order. Check them against the compiled PDF.
• Unexplained count. The text says LSeg+ uses "134 per-pixel class name supervision" §4.3; what the 134 counts is not stated.
• Limitations. The model cannot segment concepts without a segmentation proposal, and the authors note that generalisation of mask prediction to unseen concepts is an open question app. B. They also warn that biases from the training distribution call for study before deployment app. A.

Lineage

• Builds on: CLIP and ALIGN for image-text alignment, and the ALIGN image encoder used to initialise the backbone §4.1. Mask queries and cross-attention follow DETR, Max-DeepLab and MaskFormer §3.1; DETR has a card at DETR. The teacher pseudo-label recipe follows MuST §3.3.
• Compared against: LSeg, SPNet, ZS3Net, GroupViT, ALIGN with proposals and LSeg+ (this paper's re-implementation) §4.3 tab. 2.
• Concurrent work noted by the authors: Zabari and Hoshen, DenseCLIP and Xu et al. §2.
• Superseded by: not known from content already in this wiki.