[Paper Review] ContrastMask: Contrastive Learning to Segment Every Thing

Xuehui Wang, Kai Zhao+4 more|arXiv (Cornell University)|Mar 18, 2022

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

ContrastMask proposes a unified pixel-level contrastive learning framework that leverages both annotated masks from base categories and pseudo-masks from novel categories to train a class-agnostic mask segmentation model. By sharing foreground and background queries across base and novel categories and using a query-sharing contrastive loss, it improves foreground-background feature discrimination, achieving state-of-the-art performance with 39.8 mAP on novel categories under partial supervision on COCO.

ABSTRACT

Partially-supervised instance segmentation is a task which requests segmenting objects from novel unseen categories via learning on limited seen categories with annotated masks thus eliminating demands of heavy annotation burden. The key to addressing this task is to build an effective class-agnostic mask segmentation model. Unlike previous methods that learn such models only on seen categories, in this paper, we propose a new method, named ContrastMask, which learns a mask segmentation model on both seen and unseen categories under a unified pixel-level contrastive learning framework. In this framework, annotated masks of seen categories and pseudo masks of unseen categories serve as a prior for contrastive learning, where features from the mask regions (foreground) are pulled together, and are contrasted against those from the background, and vice versa. Through this framework, feature discrimination between foreground and background is largely improved, facilitating learning of the class-agnostic mask segmentation model. Exhaustive experiments on the COCO dataset demonstrate the superiority of our method, which outperforms previous state-of-the-arts.

Motivation & Objective

• To address the challenge of partially-supervised instance segmentation where only base categories have mask annotations, while novel categories have only box-level annotations.
• To improve feature discrimination between foreground and background by leveraging both base and novel category data during training.
• To bridge the generalization gap between base and novel categories by unifying contrastive learning across both categories.
• To reduce annotation burden by enabling effective model learning from limited mask annotations and pseudo-masks from novel categories.
• To develop a robust, class-agnostic mask segmentation model that generalizes well to unseen object categories.

Proposed method

• A query-sharing pixel-level contrastive loss is introduced, where shared foreground and background queries are computed by averaging features within and outside mask regions across both base and novel categories.
• Pseudo-masks for novel categories are generated using Class Activation Maps (CAM), which serve as a region prior for defining positive and negative keys.
• A special sampling strategy selects keys from within and outside mask regions, ensuring that features from the same region are pulled together and contrasted against features from the opposite region.
• The contrastive loss is applied per proposal, with shared queries enabling consistent feature learning across base and novel categories.
• Features from the contrastive learning head are fused into a class-agnostic mask head for final instance segmentation.
• High-confidence CAM regions are used as priors to improve robustness against noisy pseudo-masks.

Experimental results

Research questions

• RQ1 Can a unified contrastive learning framework improve mask segmentation performance by jointly training on base and novel categories?
• RQ2 Does sharing foreground and background queries across base and novel categories enhance feature discrimination between foreground and background?
• RQ3 Can pseudo-masks from CAMs effectively serve as a prior for contrastive learning in the absence of ground-truth masks for novel categories?
• RQ4 How does the query-sharing strategy compare to per-proposal query learning in terms of segmentation performance?
• RQ5 To what extent does using both base and novel category data improve generalization to novel categories?

Key findings

• ContrastMask achieves 39.8 mAP on novel categories using ResNeXt-101-FPN as the backbone, significantly outperforming prior state-of-the-art methods.
• The ablation study confirms that query-sharing is essential, as removing it reduces mAP from 35.1 to 32.7, demonstrating its critical role in feature alignment.
• Using only base categories for contrastive learning yields 33.5 mAP on base categories, while using all categories (base + novel) improves to 35.1 mAP, showing the benefit of including novel category data.
• The method improves mAP by 1.4 points when ground-truth masks are used instead of CAM-based pseudo-masks, indicating that pseudo-mask quality limits performance.
• The query-sharing strategy reduces the negative impact of noisy CAMs, maintaining strong performance even with imperfect pseudo-masks.
• ContrastMask generalizes effectively to novel categories without requiring additional mask annotations, making it suitable for real-world low-resource scenarios.