[Paper Review] Few-Shot Segmentation via Cycle-Consistent Transformer

Gengwei Zhang, Guoliang Kang+2 more|arXiv (Cornell University)|Jun 4, 2021

Domain Adaptation and Few-Shot Learning44 references31 citations

TL;DR

CyCTR introduces a cycle-consistent transformer to fuse pixel-wise support features into query features for few-shot segmentation, achieving state-of-the-art results on Pascal-5i and COCO-20i.

ABSTRACT

Few-shot segmentation aims to train a segmentation model that can fast adapt to novel classes with few exemplars. The conventional training paradigm is to learn to make predictions on query images conditioned on the features from support images. Previous methods only utilized the semantic-level prototypes of support images as conditional information. These methods cannot utilize all pixel-wise support information for the query predictions, which is however critical for the segmentation task. In this paper, we focus on utilizing pixel-wise relationships between support and query images to facilitate the few-shot segmentation task. We design a novel Cycle-Consistent TRansformer (CyCTR) module to aggregate pixel-wise support features into query ones. CyCTR performs cross-attention between features from different images, i.e. support and query images. We observe that there may exist unexpected irrelevant pixel-level support features. Directly performing cross-attention may aggregate these features from support to query and bias the query features. Thus, we propose using a novel cycle-consistent attention mechanism to filter out possible harmful support features and encourage query features to attend to the most informative pixels from support images. Experiments on all few-shot segmentation benchmarks demonstrate that our proposed CyCTR leads to remarkable improvement compared to previous state-of-the-art methods. Specifically, on Pascal-$5^i$ and COCO-$20^i$ datasets, we achieve 67.5% and 45.6% mIoU for 5-shot segmentation, outperforming previous state-of-the-art methods by 5.6% and 7.1% respectively.

Motivation & Objective

• Motivate leveraging pixel-level support information for each query pixel in few-shot segmentation.
• Develop a cycle-consistent attention mechanism to filter out harmful support features during cross-image attention.
• Propose CyCTR to aggregate pixel-wise support features into query features via self- and cross-alignment transformer blocks.
• Demonstrate state-of-the-art performance on standard few-shot segmentation benchmarks (Pascal-5i, COCO-20i).

Proposed method

• Introduce CyCTR with two transformer blocks per encoder: a self-alignment block (query context) and a cross-alignment block (query–support attention).
• Implement cycle-consistent attention to suppress cycle-inconsistent support pixels during cross-attention (Equation 5).
• Compute an affinity A = QK^T / sqrt(d) and apply a cycle-consistency bias B to zero out non-cycle-consistent support tokens during aggregation (Equations 3–5).
• Use a sampling strategy to handle K-shot settings by selecting foreground/background tokens (N_fg and N_s) for scalable cross-attention.
• Adopt a shared backbone (ImageNet-pretrained ResNet) with CyCTR on top, including prior maps and global support features, followed by a classifier head.
• Train with Dice loss and auxiliary loss from a grounded support-based segmentation map; employ AdamW optimization.

Experimental results

Research questions

• RQ1 Can pixel-wise cross-attention between support and query images improve few-shot segmentation beyond prototype-based methods?
• RQ2 Does introducing cycle-consistent attention effectively filter harmful support pixels and preserve informative background pixels?
• RQ3 How does CyCTR perform on standard benchmarks (Pascal-5i, COCO-20i) in 1-shot and 5-shot settings?
• RQ4 What is the impact of encoder depth, hidden dimension, and sampling strategy on performance and efficiency?

Key findings

• CyCTR achieves state-of-the-art results on Pascal-5i and COCO-20i across 1-shot and 5-shot settings.
• On Pascal-5i with ResNet-50, 1-shot mIoU = 64.0 and 5-shot mIoU = 69.3 (mean).
• On Pascal-5i with ResNet-101, 1-shot mIoU = 63.7 and 5-shot mIoU = 67.4 (mean).
• On COCO-20i with ResNet-50, 1-shot mIoU = 40.3 and 5-shot mIoU = 41.1 (mean).
• Cycle-consistent attention improves over vanilla cross-attention and baseline by notable margins (ablation results show gains of up to ~0.6–0.9% in mIoU for key variants).
• Compared to prior methods, CyCTR yields larger gains in mIoU and FB-IoU across evaluated splits (e.g., Pascal-5i FB-IoU 73.0% for 1-shot and 75.4% for 5-shot with ResNet-101).