[논문 리뷰] Wave-ViT: Unifying Wavelet and Transformers for Visual Representation Learning

Ting Yao, Yingwei Pan외 3명|arXiv (Cornell University)|2022. 07. 11.

Advanced Neural Network Applications인용 수 8

한 줄 요약

Wave-ViT는 자기주의 메커니즘 내에서 파형 변환과 비전 트랜스포머를 통합하는 새로운 웨이블릿 블록을 제안하여, 가역적이고 정보를 유지하는 다운샘플링을 가능하게 한다. 풀링과 같은 비가역적 연산 대신 이산 웨이블릿 변환(DWT)을 사용해 키/값을 다운샘플링하고, 역이산 웨이블릿 변환(IDWT)을 통해 특징 맵을 재구성하고 향상시킴으로써, 최신 기준 모델과 유사한 FLOPs로 이미지 인식, 객체 검출, 세분화 작업에서 뛰어난 성능을 달성한다.

ABSTRACT

Multi-scale Vision Transformer (ViT) has emerged as a powerful backbone for computer vision tasks, while the self-attention computation in Transformer scales quadratically w.r.t. the input patch number. Thus, existing solutions commonly employ down-sampling operations (e.g., average pooling) over keys/values to dramatically reduce the computational cost. In this work, we argue that such over-aggressive down-sampling design is not invertible and inevitably causes information dropping especially for high-frequency components in objects (e.g., texture details). Motivated by the wavelet theory, we construct a new Wavelet Vision Transformer ( extbf{Wave-ViT}) that formulates the invertible down-sampling with wavelet transforms and self-attention learning in a unified way. This proposal enables self-attention learning with lossless down-sampling over keys/values, facilitating the pursuing of a better efficiency-vs-accuracy trade-off. Furthermore, inverse wavelet transforms are leveraged to strengthen self-attention outputs by aggregating local contexts with enlarged receptive field. We validate the superiority of Wave-ViT through extensive experiments over multiple vision tasks (e.g., image recognition, object detection and instance segmentation). Its performances surpass state-of-the-art ViT backbones with comparable FLOPs. Source code is available at \url{https://github.com/YehLi/ImageNetModel}.

연구 동기 및 목표

• 다양한 척도의 비전 트랜스포머에서 비가역적 다운샘플링 연산(예: 평균 풀링)으로 인한 정보 손실 문제를 해결하기 위해.
• 풀링 대체로 가역적인 웨이블릿 기반 다운샘플링을 사용하여 자기주의 계산 중 고주파 수치(예: 질감)를 유지하기 위해.
• 역웨이블릿 변환을 통해 다운샘플된 특징을 재구성함으로써 자기주의 출력에서 국소적 맥락 모델링을 향상시키기 위해.
• 최소한의 계산 오버헤드로 높은 정확도를 유지하면서도 확장성 있고 효율적인 트랜스포머 블록을 설계하기 위해.
• 제안된 웨이블릿 블록의 효과성을 여러 비전 작업, 특히 이미지 인식과 조밀 예측 작업에서 검증하기 위해.

제안 방법

• 이산 웨이블릿 변환(DWT)을 사용해 키와 값을 네 개의 하위대역(LL, LH, HL, HH)으로 분해하여 가역적 다운샘플링을 수행하는 웨이블릿 블록을 도입한다.
• 네 개의 DWT 하위대역을 하나의 다운샘플된 특징 맵으로 통합하고, 주파수 하위대역 간 국소 불변성을 확보하기 위해 3×3 컨볼루션을 적용한다.
• 다운샘플된 키/값과 원본 쿼리에 대해 다중 헤드 자기주의를 수행하여 계산 효율성을 유지한다.
• 다운샘플된 키/값에 대해 역이산 웨이블릿 변환(IDWT)을 적용하여 해상도가 높은 특징 맵을 재구성하고 효과적 수용각을 확대한다.
• 주의 기반 특징 맵과 IDWT를 통해 재구성된 맵을 융합하여 최종 출력을 생성함으로써 국소적 맥락 집약을 향상시킨다.
• 다단계, 다중 척도 트랜스포머 아키텍처에 웨이블릿 블록을 통합하여, 시각적 표현 학습을 위한 새로운 백본인 Wave-ViT를 구축한다.

실험 결과

연구 질문

• RQ1 가역적 웨이블릿 변환이 비전 트랜스포머에서 비가역적 풀링 연산을 효과적으로 대체하여 고주파 이미지 세부 정보를 유지할 수 있는가?
• RQ2 역웨이블릿 변환의 통합이 계산 비용 증가 없이 자기주의 출력의 국소적 맥락 모델링을 향상시키는가?
• RQ3 표준 풀링 기반 자기주의 블록과 비교할 때 제안된 웨이블릿 블록은 다양한 비전 작업에서 정확도와 효율성 측면에서 어떻게 성능을 내는가?
• RQ4 Wave-ViT는 기존 모델과 유사한 FLOPs로 이미지 인식, 객체 검출, 인스턴스 세분화 작업에서 최신 기준 성능을 달성할 수 있는가?
• RQ5 웨이블릿 기반 다운샘플링을 사용할 경우 특징 표현의 강건성이 얼마나 향상되는가? 이는 시각화된 사전 지도를 통해 어떻게 확인되는가?

주요 결과

• Wave-ViT-S는 ADE20K 검증 세트에서 51.5%의 mIoU를 기록하여 최고 경쟁 모델(Swin-S)을 2.0%포인트 이상 앞서며 우수한 성능을 보였다.
• Wave-ViT-B는 ImageNet-1K에서 82.7%의 Top-1 정확도를 달성하여 비가역적 풀링 기반 기준 모델(82.0%)보다 0.7%포인트 높은 성능을 보였다.
• 제거 분석 결과, 평균 풀링을 DWT로 단독으로 대체할 경우 Top-1 정확도가 82.0%에서 82.5%로 향상되어 가역적 다운샘플링의 유용성을 입증하였다.
• DWT 기반 블록에 IDWT를 추가함으로써 성능은 82.7%로 더욱 향상되었고, 계산 비용 증가는 극히 미미하여 확대된 수용각의 효과를 확인하였다.
• Score-CAM 시각화 결과, Wave-ViT-S는 PVTv2-B2보다 더 강건하고 의미적으로 집중된 표현을 학습하고 있음을 확인하였으며, 관련 객체에 대한 주의 집중도가 높았다.
• Wave-ViT는 다양한 작업에 대해 잘 일반화되며, 비교 가능한 FLOPs 조건에서 이미지 인식, 객체 검출, 인스턴스 세분화 작업에서 최신 기준 성능을 달성하였다.