[Paper Review] Learning Equivariant Segmentation with Instance-Unique Querying

Wenguan Wang, James Liang+1 more|arXiv (Cornell University)|Oct 3, 2022

Colorectal Cancer Screening and Detection22 citations

TL;DR

The paper introduces a training framework that enhances query-based instance segmentation by enforcing dataset-level instance uniqueness and transformation-equivariant learning of query embeddings and features, leading to significant AP gains without changing inference.

ABSTRACT

Prevalent state-of-the-art instance segmentation methods fall into a query-based scheme, in which instance masks are derived by querying the image feature using a set of instance-aware embeddings. In this work, we devise a new training framework that boosts query-based models through discriminative query embedding learning. It explores two essential properties, namely dataset-level uniqueness and transformation equivariance, of the relation between queries and instances. First, our algorithm uses the queries to retrieve the corresponding instances from the whole training dataset, instead of only searching within individual scenes. As querying instances across scenes is more challenging, the segmenters are forced to learn more discriminative queries for effective instance separation. Second, our algorithm encourages both image (instance) representations and queries to be equivariant against geometric transformations, leading to more robust, instance-query matching. On top of four famous, query-based models ($i.e.,$ CondInst, SOLOv2, SOTR, and Mask2Former), our training algorithm provides significant performance gains ($e.g.,$ +1.6 - 3.2 AP) on COCO dataset. In addition, our algorithm promotes the performance of SOLOv2 by 2.7 AP, on LVISv1 dataset.

Motivation & Objective

• Motivate improving discriminative power of instance queries beyond intra-scene training.
• Promote cross-scene (dataset-wide) instance discrimination to distinguish all dataset instances.
• Enforce transformation equivariance to make queries and features robust to geometric changes.
• Demonstrate that equivariance regularization yields gains without architectural changes or slower inference.

Proposed method

• Define a dense feature extractor f to produce image embeddings I and a query creator h to generate N instance-aware queries {q_n}.
• Train with intra-scene mask loss L_intra_mask and introduce inter-scene mask loss L_inter_mask that forces queries to mismatch in other images, using an external memory and sparse, instance-balanced sampling.
• Add an equivariance loss L_equi that enforces f(g(I)) ≈ g(f(I)) and {q_n^g, I^g} to align with transformed ground-truth masks g(M_sigma(n)).
• Combine L_intra_mask, L_inter_mask, and L_equi to form a cross-scene, transformation-equivariant training objective that can be plugged into existing query-based methods.
• Use focal loss for L_inter_mask to handle many easy negatives, and a combination of dice/focal losses for L_equi depending on base method.

Experimental results

Research questions

• RQ1 Can cross-scene (dataset-level) querying improve discriminability of instance queries beyond intra-scene training?
• RQ2 Does enforcing transformation equivariance on features and queries lead to more robust instance-query matching than standard augmentation?
• RQ3 What gains in AP can be achieved on COCO and LVIS when applying the proposed framework to existing query-based models?
• RQ4 Is the proposed training framework architecture- and inference-speed-agnostic for mainstream query-based segmenters?

Key findings

• Applied to CondInst, SOLOv2, SOTR, and Mask2Former across backbones (ResNet/Swin), the method yields AP gains of +1.6 to +3.2 on COCO and +2.7 AP on LVISv1 for SOLOv2.
• On COCO test-dev, reported gains include up to +3.2 AP for certain setups and notable improvements in AP_S, AP_M, and AP_L across methods (e.g., CondInst and SOTR variants).
• Ours with SOTR-Res50 achieves AP 42.2, AP50 61.9, AP75 43.9, AP_S 11.0, AP_M 60.5, AP_L 73.5; with Res101 yields AP 42.6, AP50 64.1, AP75 45.8, AP_S 11.2, AP_M 61.2, AP_L 75.3.