[Paper Review] CATrans: Context and Affinity Transformer for Few-Shot Segmentation

Shan Zhang, Tianyi Wu+2 more|arXiv (Cornell University)|Apr 27, 2022

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

CATrans proposes a novel Context and Affinity Transformer for few-shot image segmentation, integrating Relation-guided Context (RCT) and Relation-guided Affinity (RAT) modules to enhance cross-image feature matching. By leveraging both global context from support features and attention-aware affinity with self-affinity regularization, CATrans achieves state-of-the-art performance, outperforming prior methods by up to +2.8% mIoU on PASCAL-5i and +2.1% on COCO-20i with ResNet-101.

ABSTRACT

Few-shot segmentation (FSS) aims to segment novel categories given scarce annotated support images. The crux of FSS is how to aggregate dense correlations between support and query images for query segmentation while being robust to the large variations in appearance and context. To this end, previous Transformer-based methods explore global consensus either on context similarity or affinity map between support-query pairs. In this work, we effectively integrate the context and affinity information via the proposed novel Context and Affinity Transformer (CATrans) in a hierarchical architecture. Specifically, the Relation-guided Context Transformer (RCT) propagates context information from support to query images conditioned on more informative support features. Based on the observation that a huge feature distinction between support and query pairs brings barriers for context knowledge transfer, the Relation-guided Affinity Transformer (RAT) measures attention-aware affinity as auxiliary information for FSS, in which the self-affinity is responsible for more reliable cross-affinity. We conduct experiments to demonstrate the effectiveness of the proposed model, outperforming the state-of-the-art methods.

Motivation & Objective

• To address the challenge of large intra-class appearance and geometric variations in few-shot segmentation (FSS) by improving feature matching between support and query images.
• To overcome limitations in prior Transformer-based FSS methods that rely solely on context or affinity, by jointly modeling both for more robust representation learning.
• To enhance context transfer from support to query images using a Relation-guided Context Transformer (RCT) that leverages discriminative support features.
• To improve cross-affinity estimation by incorporating self-affinity as an auxiliary signal, reducing noise in pixel-wise correspondences.
• To design a hierarchical architecture that fuses multi-scale context and affinity maps for improved segmentation accuracy.

Proposed method

• The Relation-guided Context Transformer (RCT) propagates semantic context from support images to query images using a transformer block conditioned on high-level support features, improving feature alignment under appearance variance.
• The Relation-guided Affinity Transformer (RAT) computes attention-aware cross-affinity between support and query features, using self-affinity maps as a regularizer to suppress noisy correlations.
• CATrans integrates RCT and RAT in a hierarchical architecture, stacking multi-level correlation maps to capture both local and global relationships across feature hierarchies.
• High-resolution query features are concatenated with low-resolution correlation maps to guide the decoder, enhancing spatial precision in segmentation.
• The model uses a dual-branch feature encoder (e.g., ResNet or Swin-T) followed by shared attention mechanisms in RCT and RAT to maintain parameter efficiency.
• The framework is trained end-to-end with cross-entropy and Dice loss, optimizing for pixel-wise segmentation on query images given few-shot support masks.

Experimental results

Research questions

• RQ1 How can context and affinity information be effectively combined in a Transformer-based architecture to improve few-shot segmentation under large appearance variations?
• RQ2 Can self-affinity maps improve the reliability of cross-affinity estimation in support-query matching for FSS?
• RQ3 Does hierarchical fusion of multi-scale context and affinity representations lead to better segmentation performance than single-scale or non-hierarchical designs?
• RQ4 To what extent does enhancing support feature representation through relation-guided context propagation improve generalization in few-shot settings?
• RQ5 How do the proposed RCT and RAT modules compare to existing attention-based matching mechanisms in terms of robustness and accuracy?

Key findings

• CATrans achieves 67.6% mIoU on the 1-shot PASCAL-5i benchmark using Swin-T as the backbone, surpassing the previous SOTA DGPNet by +2.8%.
• On COCO-20i, CATrans reaches 60.1% mIoU in the 5-shot setting with Swin-T, outperforming DGPNet by +2.0%.
• Ablation studies show that combining both RCT and RAT modules yields a +5% mIoU gain over the baseline, with RCT alone improving performance by up to +7.7% on 1-shot PASCAL-5i.
• Multi-scale feature fusion improves mIoU by 1.4% on PASCAL-5i and 1.8% on COCO-20i, demonstrating the benefit of high-resolution guidance in the decoder.
• The model maintains strong performance across different backbone architectures, with consistent gains using ResNet-50 and ResNet-101, indicating robustness to model capacity.
• Qualitative results show that RAT reduces false positives caused by shape and color similarity, while RCT improves context-aware segmentation in challenging cases with large appearance variations.