[논문 리뷰] WeakTr: Exploring Plain Vision Transformer for Weakly-supervised Semantic Segmentation

Lianghui Zhu, Yingyue Li외 5명|arXiv (Cornell University)|2023. 04. 03.

Advanced Neural Network Applications인용 수 16

한 줄 요약

WeakTr은 적응형 주의 융합으로 엔드-투-엔드에서 고품질 CAM을 생성하는 평범한 Vision Transformer를 활용하며 온라인 재훈련을 위한 그래디언트 클리핑 디코더를 사용해 VOC 2012와 COCO 2014에서 최첨단 WSSS 결과를 달성합니다.

ABSTRACT

Transformer has been very successful in various computer vision tasks and understanding the working mechanism of transformer is important. As touchstones, weakly-supervised semantic segmentation (WSSS) and class activation map (CAM) are useful tasks for analyzing vision transformers (ViT). Based on the plain ViT pre-trained with ImageNet classification, we find that multi-layer, multi-head self-attention maps can provide rich and diverse information for weakly-supervised semantic segmentation and CAM generation, e.g., different attention heads of ViT focus on different image areas and object categories. Thus we propose a novel method to end-to-end estimate the importance of attention heads, where the self-attention maps are adaptively fused for high-quality CAM results that tend to have more complete objects. Besides, we propose a ViT-based gradient clipping decoder for online retraining with the CAM results efficiently and effectively. Furthermore, the gradient clipping decoder can make good use of the knowledge in large-scale pre-trained ViT and has a scalable ability. The proposed plain Transformer-based Weakly-supervised learning method (WeakTr) obtains the superior WSSS performance on standard benchmarks, i.e., 78.5% mIoU on the val set of PASCAL VOC 2012 and 51.1% mIoU on the val set of COCO 2014. Source code and checkpoints are available at https://github.com/hustvl/WeakTr.

연구 동기 및 목표

• 합성 인덕션 바이어스 없이 평범한 ViT를 사용해 WSSS를 위한 CAM 품질 향상을 고무한다.
• ViT 헤드를 가중하기 위한 적응형 주의 융합 모듈을 제안하여 CAM 생성을 개선한다.
• 분류 신호를 통해 CAM 품질을 최적화하는 엔드-투-엔드 CAM 학습 전략을 도입한다.
• 그래디언트 클리핑 디코더를 사용하는 온라인 재훈련 접근법을 개발하여 분할 모델을 효율적으로 업데이트한다.
• VOC 2012 및 COCO 2014 벤치마크에서 최첨단 WSSS 성능을 입증한다.

제안 방법

• C 클래스 토큰과 N^2 패치 토큰을 입력으로 하는 일반적인 ViT 백본을 트랜스포머 인코더에 사용한다.
• 패치 토큰에 대한 컨볼루션으로 거친 CAM을 생성한 뒤, 자기 주의 맵의 적응형 주의 융합으로 이를 정제한다.
• 어텐션 맵을 풀링하고 FFN을 통과시켜 W'를 얻어 동적인 헤드 가중치 W를 계산한 다음, 교차 어텐션 및 패치 어텐션 맵에 가중치를 적용하여 CAM_fine을 형성한다.
• 헤드 가중치를 감독하기 위해 L = L_Fine-CAM + L_CLS-token + L_Coarse-CAM인 결합 손실로 엔드-투-엔드 학습한다.
• 전역/로컬 그래디언트 통계에 기반하여 그래디언트 흐름을 제약하여 세분화 네트를 업데이트하는 온라인 재훈련용 그래디언트 클리핑 디코더를 도입한다.
• 추론 중에는 CRF를 적용해 세분화 맵을 다듬는다.

실험 결과

연구 질문

• RQ1 평범한 Vision Transformer의 자체 어텐션 맵을 적응적으로 융합하여 WSSS용 고품질 CAM을 어떻게 생성할 수 있는가?
• RQ2 CAM 정제 단계를 거치지 않고도 적응형 헤드 가중치를 통한 엔드-투-엔드 CAM 학습이 의사 라벨 품질을 향상시킬 수 있는가?
• RQ3 그래디언트 클리핑 디코더를 이용한 온라인 재훈련이 전통적인 CAM 정제 파이프라인에 비해 WSSS의 효율성과 정확도를 향상시키는가?
• RQ4 ViT 백본을 사용한 표준 WSSS 벤치마크(VOC 2012 및 COCO 2014)에서 WeakTr의 성능 영향은 어떠한가?
• RQ5 CAM 품질과 최종 분할의 mIoU 측면에서 WeakTr은 최신 WSSS 방법과 어떻게 비교되는가?

주요 결과

• WeakTr는 VOC 2012 val 및 COCO 2014 val 벤치마크에서 최첨단 WSSS 결과를 달성한다.
• VOC 2012 val에서 ViT-S를 사용한 WeakTr은 78.4% mIoU 및 테스트에서 79.0%를 달성하여 이전 방법들보다 우수하다.
• VOC 2012 train에서 CAM 개선(Fine-CAM)은 이전 CAM 방법들(e.g., MCTformer, ViT-PCM)보다 우수하다.
• 그래디언트 클리핑 디코더를 이용한 온라인 재훈련은 전체 학습 시간을 크게 절약(약 2.6x 빠름)하고 높은 mIoU를 유지한다.
• 적응형 주의 융합(AAF)은 CRF 후처리 사용 시 특히 평균 합계 집계보다 CAM 정밀도/재현율 및 mIoU가 더 높다.
• 아블레이션 연구는 그래디언트 패치 크기와 클리핑 시작 임계값이 최종 성능에 영향을 주며 제안된 디코더에서 의미 있는 이득을 보인다.