[논문 리뷰] VSA: Learning Varied-Size Window Attention in Vision Transformers

Qiming Zhang, Yufei Xu외 2명|arXiv (Cornell University)|2022. 04. 18.

Advanced Neural Network Applications인용 수 4

한 줄 요약

이 논문은 시각 Transformer에서 고정 크기 창을 대체하기 위해 데이터 기반으로 적응 가능한 창 구성으로 학습 가능한 메커니즘인 Varied-Size Window Attention (VSA)를 제안한다. 각 어텐션 헤드별로 최적의 창 크기와 위치를 예측하는 창 회귀 모듈을 사용함으로써 VSA는 장거리 모델링을 향상시키고 다양한 물체 크기에서의 특징 표현을 개선하며, 거의 무시할 만한 계산 부담으로 성능을 향상시킨다. Swin-T에서 ImageNet에서 1.1%의 정확도 향상을 기록했으며, 더 큰 입력 이미지(예: 480×480)를 사용할 경우 최대 1.9%까지 성능 향상을 기록한다.

ABSTRACT

Attention within windows has been widely explored in vision transformers to balance the performance, computation complexity, and memory footprint. However, current models adopt a hand-crafted fixed-size window design, which restricts their capacity of modeling long-term dependencies and adapting to objects of different sizes. To address this drawback, we propose extbf{V}aried- extbf{S}ize Window extbf{A}ttention (VSA) to learn adaptive window configurations from data. Specifically, based on the tokens within each default window, VSA employs a window regression module to predict the size and location of the target window, i.e., the attention area where the key and value tokens are sampled. By adopting VSA independently for each attention head, it can model long-term dependencies, capture rich context from diverse windows, and promote information exchange among overlapped windows. VSA is an easy-to-implement module that can replace the window attention in state-of-the-art representative models with minor modifications and negligible extra computational cost while improving their performance by a large margin, e.g., 1.1\% for Swin-T on ImageNet classification. In addition, the performance gain increases when using larger images for training and test. Experimental results on more downstream tasks, including object detection, instance segmentation, and semantic segmentation, further demonstrate the superiority of VSA over the vanilla window attention in dealing with objects of different sizes. The code will be released https://github.com/ViTAE-Transformer/ViTAE-VSA.

연구 동기 및 목표

• 시각 Transformer에서 고정 크기 창 어텐션의 한계를 해결하기 위해, 이는 장거리 상관관계 모델링을 제한하고 다양한 크기의 물체에 대한 적응을 제한한다.
• 입력 데이터에 기반해 각 어텐션 헤드별로 최적의 창 크기와 위치를 동적으로 학습할 수 있도록 하기 위해.
• 계산 비용이나 메모리 사용량을 크게 증가시키지 않으면서도 비전 작업 성능을 향상시키기 위해.
• 이격된 창 메커니즘에 대한 의존도를 줄이기 위해, 겹치는 적응형 창을 통해 정보 교환을 가능하게 하기 위해.
• 분류, 검출, 세그멘테이션을 포함한 다양한 작업에 대해 일반화 성능을 입증하기 위해, 특히 고해상도 입력에서의 성능을 강조하기 위해.

제안 방법

• VSA는 기본 고정 크기 창 내의 토큰을 기반으로 목표 어텐션 창의 크기와 공간적 위치를 예측하는 창 회귀 모듈을 도입한다.
• 그런 다음 예측된 목표 창에서 키와 밸류 토큰을 샘플링하고, 쿼리는 기본 창에서 유지함으로써 가변적인 수신 필드를 가능하게 한다.
• 각 어텐션 헤드는 독립적으로 VSA를 적용하여 헤드 간 다양한 창 구성이 다중 척도의 맥락을 포착할 수 있도록 한다.
• 이 모듈은 플러그 앤 플레이 방식으로, Swin-T와 같은 모델의 표준 창 어텐션을 거의 변화 없이 대체할 수 있으며, 거의 무시할 만한 FLOPs 부담을 유발한다.
• 겹치는 창을 통해 창 간 특징 교환을 촉진함으로써, 이격된 창 메커니즘의 필요성을 줄인다.
• 창 회귀는 모델의 나머지 부분과 함께 엔드 투 엔드로 학습되며, 물체 크기와 공간적 레이아웃에 대한 데이터 기반 적응을 가능하게 한다.

실험 결과

연구 질문

• RQ1 학습 가능한 창 메커니즘이 고정 크기 창 제약을 초월해 시각 Transformer의 장거리 모델링을 향상시킬 수 있는가?
• RQ2 적응형 창 크기 및 위치 예측이 다양한 크기의 물체에 대해 더 나은 특징 표현을 이끌어낼 수 있는가?
• RQ3 VSA는 계산 비용을 크게 증가시키지 않고도 이미지 분류 및 조밀 예측 작업 성능을 향상시킬 수 있는가?
• RQ4 고정 크기 창이 큰 물체를 포착하기 어려운 더 큰 입력 해상도에서 VSA는 어떻게 성능을 발휘하는가?
• RQ5 겹치는 창을 통한 효과적인 창 간 정보 교환을 통해 VSA는 이격된 창 메커니즘의 필요성을 줄이거나 제거할 수 있는가?

주요 결과

• Swin-T에서 ImageNet에서 1.1%의 정확도 향상을 기록했으며, 더 큰 이미지(예: 480×480)로 훈련 및 테스트할 경우 성능 향상은 1.9%까지 증가한다.
• 객체 검출, 인스턴스 세그멘테이션, 세분화 세그멘테이션과 같은 후행 작업에서 VSA는 표준 창 어텐션을 항상 능가하며, 특히 큰 물체에서 두드러진 성능 향상을 보였다.
• VSA를 적용한 모델는 t-SNE 시각화에서 더 나은 특징 군집화를 보이며, 더 구분력 있고 의미 있는 의미론적 표현을 나타낸다.
• VSA는 겹치는 창 간 효과적인 정보 교환을 가능하게 하여, Swin의 VSA 변형에서는 이격된 창 메커니즘이 불필요해진다.
• 최적화되지 않은 PyTorch 샘플링 연산으로 인해 12–17%의 속도 저하가 발생하지만, GPU 메모리 사용량은 2%만 증가하여 매우 낮은 오버헤드를 보였다.
• 시각화 결과 VSA는 고정 크기 창보다 더 나은 물체 커버리지와 더 높은 변동성을 보이는, 맥락 인식 창을 생성함을 확인했으며, 주목할 만한 영역에서 창 크기와 위치의 다양성이 높았다.