[논문 리뷰] Few-Shot Segmentation via Cycle-Consistent Transformer

Gengwei Zhang, Guoliang Kang외 2명|arXiv (Cornell University)|2021. 06. 04.

Domain Adaptation and Few-Shot Learning참고 문헌 44인용 수 31

한 줄 요약

CyCTR는 사이클-일치(transformer)를 도입하여 픽셀 단위의 지원 정보를 쿼리 피처로 융합하고 소샷 세그멘테이션을 수행하며, Pascal-5i 및 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.

연구 동기 및 목표

• 소샷 세그멘테이션에서 각 쿼리 픽셀에 대해 픽셀 수준의 지원 정보를 활용하는 것을 동기로 삼는다.
• 교차 이미지 어텐션 중 해로운 지원 특징을 걸러내기 위한 사이클-일치 어텐션 메커니즘을 개발한다.
• 자체 정렬 및 교차 정렬 Transformer 블록을 통해 픽셀 단위의 지원 특징을 쿼리 특징으로 집계하는 CyCTR를 제안한다.
• 표준 소샷 세그멘테이션 벤치마크(Pascal-5i, COCO-20i)에서 최첨단 성능을 입증한다.

제안 방법

• 인코더당 두 개의 트랜스포머 블록으로 CyCTR를 도입한다: 자체 정렬 블록(쿼리 컨텍스트)과 교차 정렬 블록(쿼리–지원 어텐션).
• 교차 어텐션 중 사이클-일치 어텐션을 구현하여 사이클-일치하지 않는 지원 픽셀을 억제한다(식 5).
• 친화도 A = QK^T / sqrt(d)을 계산하고 합산 중 비-사이클-일치 지원 토큰을 0으로 만드는 사이클-일치 바이어스 B를 적용한다(식 3–5).
• 확장 가능한 교차 어텐션을 위해 K-샷 설정을 다루는 샘플링 전략을 사용하여 전경/배경 토큰(N_fg 및 N_s)을 선택한다.
• 공유 백본(ImageNet-사전학습 ResNet)을 위에 CyCTR를 얹어 사용하고, 사전 지도 맵과 전역 지원 특징을 포함한 뒤 분류기 헤드를 적용한다.
• Dice 손실 및 기반 지원 기반 분할 맵에서 보조 손실로 학습하고, AdamW 최적화를 사용한다.

실험 결과

연구 질문

• RQ1 지원 이미지와 쿼리 이미지 간의 픽셀-단위 교차 어텐션이 프로토타입 기반 방법을 넘어선 소샷 세그멘테이션을 개선할 수 있는가?
• RQ2 사이클-일치 어텐션을 도입하면 해로운 지원 픽셀을 효과적으로 필터링하고 유익한 배경 픽셀을 보존하는가?
• RQ3 CyCTR은 표준 벤치마크(Pascal-5i, COCO-20i)에서 1샷과 5샷 설정에서 어떤 성능을 보이는가?
• RQ4 인코더 깊이, 숨겨진 차원, 샘플링 전략이 성능과 효율성에 어떤 영향을 미치는가?

주요 결과

• CyCTR는 1샷 및 5샷 설정에서 Pascal-5i와 COCO-20i에서 최첨단 성능을 달성한다.
• Pascal-5i에서 ResNet-50로 1-shot mIoU = 64.0, 5-shot mIoU = 69.3(평균).
• Pascal-5i에서 ResNet-101로 1-shot mIoU = 63.7, 5-shot mIoU = 67.4(평균).
• COCO-20i에서 ResNet-50로 1-shot mIoU = 40.3, 5-shot mIoU = 41.1(평균).
• 사이클-일치 어텐션은 기본 교차 어텐션 및 베이스라인 대비 현저한 향상을 보이며(주요 변형에 대해 mIoU에서 최대 약 0.6–0.9%의 이점이 Ablation 결과에서 나타남).
• 시행착오 대비 CyCTR는 평가 분할에서 mv(mIoU)와 FB-IoU의 차이가 큰 편인 이전 방법들보다 더 큰 이득을 제공한다(예: 1-shot에서 ResNet-101의 Pascal-5i FB-IoU 73.0%, 5-shot에서 75.4%).