[Paper Review] Few-shot Segmentation with Optimal Transport Matching and Message Flow

Weide Liu, Chi Zhang+3 more|arXiv (Cornell University)|Aug 19, 2021

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper proposes CMNet, a few-shot segmentation model that uses partial optimal transport matching to establish dense, constrained correspondences between query and support features, while employing a message flow module to enhance local and global feature representations. The method achieves state-of-the-art performance on PASCAL VOC 2012, MS COCO, and FSS-1000 under both standard and weakly supervised settings.

ABSTRACT

We tackle the challenging task of few-shot segmentation in this work. It is essential for few-shot semantic segmentation to fully utilize the support information. Previous methods typically adopt masked average pooling over the support feature to extract the support clues as a global vector, usually dominated by the salient part and lost certain essential clues. In this work, we argue that every support pixel's information is desired to be transferred to all query pixels and propose a Correspondence Matching Network (CMNet) with an Optimal Transport Matching module to mine out the correspondence between the query and support images. Besides, it is critical to fully utilize both local and global information from the annotated support images. To this end, we propose a Message Flow module to propagate the message along the inner-flow inside the same image and cross-flow between support and query images, which greatly helps enhance the local feature representations. Experiments on PASCAL VOC 2012, MS COCO, and FSS-1000 datasets show that our network achieves new state-of-the-art few-shot segmentation performance.

Motivation & Objective

• Address the limitation of previous few-shot segmentation methods that rely on global average pooling of support features, which discards spatial structure and underutilizes less discriminative parts.
• Overcome the many-to-one matching problem in prior methods by enabling a constrained many-to-many correspondence between query and support features.
• Enhance local feature representations by introducing a message flow module that propagates information both within images (inner-flow) and across images (cross-flow).
• Improve few-shot segmentation performance under weak supervision, such as bounding box annotations, by maintaining high accuracy with limited annotation types.
• Ensure robustness to ambiguous or similar-looking objects by validating failure cases and analyzing model limitations.

Proposed method

• Propose a Partial Optimal Transport (POT) matching module that enforces a maximum matching flow based on the object size from ground truth, enabling more balanced and comprehensive correspondence between query and support features.
• Introduce a Message Flow module that models query and support features as graphs, propagating messages along edges to refine local representations through inner-flow (within-image) and cross-flow (between-image) mechanisms.
• Integrate positional encoding into the feature representation to improve spatial alignment and correspondence accuracy.
• Use a two-branch encoder-decoder architecture with cross-attention and feature refinement to guide segmentation predictions using the matched and propagated features.
• Apply the POT module after a prior mask to refine predictions, avoiding overfitting to noisy or incorrect correspondences from full optimal transport.
• Validate the method under weakly supervised settings by replacing pixel-level support annotations with bounding boxes, demonstrating robustness to coarser supervision.

Experimental results

Research questions

• RQ1 Can a constrained many-to-many matching mechanism, such as partial optimal transport, improve correspondence learning in few-shot segmentation compared to global pooling or unconstrained matching?
• RQ2 How does message propagation across local and global feature spaces enhance segmentation accuracy, especially for less discriminative object parts?
• RQ3 Does incorporating positional encoding improve spatial alignment and segmentation performance in few-shot settings?
• RQ4 Can the proposed method maintain high performance when support annotations are reduced to bounding boxes instead of pixel-level masks?
• RQ5 What are the failure modes of the model, and how do they relate to visual similarity or lack of semantic context in query images?

Key findings

• CMNet achieves new state-of-the-art performance on PASCAL VOC 2012, with 62.5% mIoU under the 1-shot setting, outperforming prior methods including PGNet (61.7% mIoU).
• On the MS COCO dataset, the method achieves the highest performance among all previous state-of-the-art methods in both 1-shot and 5-shot settings.
• On the FSS-1000 benchmark, CMNet sets a new SOTA result in both 1-shot and 5-shot evaluation protocols.
• The ablation study confirms that positional encoding improves performance by 0.4 mIoU, increasing from 62.1% to 62.5% on PASCAL VOC 2012 under 1-shot evaluation.
• Using full optimal transport matching degrades performance compared to partial optimal transport, indicating that constrained matching avoids incorrect correspondences, such as matching background regions.
• In weakly supervised settings with bounding box support, CMNet outperforms existing methods on both PASCAL VOC 2012 and FSS-1000, demonstrating strong generalization under reduced supervision.