[논문 리뷰] Shunted Self-Attention via Multi-Scale Token Aggregation

Sucheng Ren, Daquan Zhou외 3명|arXiv (Cornell University)|2021. 11. 30.

Advanced Neural Network Applications참고 문헌 34인용 수 10

한 줄 요약

이 논문은 시각 Transformer가 한 개의 어텐션 레이어에서 미세한 해상도와 굵은 해상도의 특징을 동시에 모델링할 수 있도록 하는 새로운 다중 척도 토큰 집약 메커니즘인 셰운티드 싱글 어텐션(Shunted Self-Attention, SSA)을 제안한다. SSA는 토큰을 선택적으로 병합함으로써 기능을 집약한다. 이는 계산 비용을 줄이고 이미지 분류, 객체 검출, 영역 세분화 작업에서 성능을 향상시키며, SOTA인 포칼 트랜스포머(Focal Transformer)의 절반 크기의 모델 크기와 FLOPs로 ImageNet에서 84.0%의 Top-1 정확도를 달성한다. 또한 유사한 효율성 제약 조건 하에서 COCO에서 1.3 mAP 높은 성능을 기록하고, ADE20K에서 2.9 mIOU 향상을 이룬다.

ABSTRACT

Recent Vision Transformer~(ViT) models have demonstrated encouraging results across various computer vision tasks, thanks to their competence in modeling long-range dependencies of image patches or tokens via self-attention. These models, however, usually designate the similar receptive fields of each token feature within each layer. Such a constraint inevitably limits the ability of each self-attention layer in capturing multi-scale features, thereby leading to performance degradation in handling images with multiple objects of different scales. To address this issue, we propose a novel and generic strategy, termed shunted self-attention~(SSA), that allows ViTs to model the attentions at hybrid scales per attention layer. The key idea of SSA is to inject heterogeneous receptive field sizes into tokens: before computing the self-attention matrix, it selectively merges tokens to represent larger object features while keeping certain tokens to preserve fine-grained features. This novel merging scheme enables the self-attention to learn relationships between objects with different sizes and simultaneously reduces the token numbers and the computational cost. Extensive experiments across various tasks demonstrate the superiority of SSA. Specifically, the SSA-based transformer achieves 84.0\% Top-1 accuracy and outperforms the state-of-the-art Focal Transformer on ImageNet with only half of the model size and computation cost, and surpasses Focal Transformer by 1.3 mAP on COCO and 2.9 mIOU on ADE20K under similar parameter and computation cost. Code has been released at https://github.com/OliverRensu/Shunted-Transformer.

연구 동기 및 목표

• 각 어텐션 레이어의 토큰들 간에 균일하고 정적인 수신 영역을 가진 시각 Transformer가 다중 척도 특징을 포착하는 데에 한계가 있음.
• 다양한 척도의 객체를 포함한 이미지에서 미세한 세부 정보를 유지하면서 계산 비용과 메모리 소비를 줄임.
• 단일 자기어텐션 레이어가 동시에 큰 객체와 작은 객체를 주시할 수 있도록 하여, 후속 비전 작업의 성능 향상.
• 기존 ViT 아키텍처에 대한 대규모 수정 없이도 적용 가능한 일반적이고 즉시 사용 가능한 메커니즘 설계.

제안 방법

• 어텐션 계산을 두 개의 브랜치로 분할하는 셰운티드 싱글 어텐션 메커니즘을 도입: 하나는 병합된 굵은 해상도 토큰을 위한 것이고, 다른 하나는 병합되지 않은 미세한 해상도 토큰을 위한 것이다.
• 특징 유사성과 공간적 맥락을 기반으로 공간적으로 인접한 토큰을 선택적으로 병합하는 학습 가능한 토큰 집약 함수를 적용하여 큰 객체를 위한 더 큰 수신 영역을 형성한다.
• 粗모양 및 미세한 특징에 대해 별도의 어텐션 맵을 유지함으로써, 동일한 레이어 내에서 다양한 척도 간의 관계를 학습할 수 있도록 한다.
• 이웃 토큰의 맥락을 통합함으로써 국소적 특징 표현을 향상시키는 디테일 전용 피드포워드 네트워크를 통합하여 특징 정제를 향상시킨다.
• 공간 해상도를 유지하고 초기 특징 추출 과정에서의 정보 손실을 줄이기 위해 다중 척도 패치 임베딩 헤드를 사용한다.
• 병합된 토큰은 시퀀스 길이와 FLOPs를 줄이고, 병합되지 않은 토큰은 고해상도 세부 정보를 유지하는 혼합 계산 전략을 적용한다.

실험 결과

연구 질문

• RQ1 히어archical 특징 맵에 의존하지 않고도 단일 자기어텐션 레이어가 다중 척도 특징을 효과적으로 모델링할 수 있는가?
• RQ2 선택적 토큰 병합이 객체 검출 및 영역 세분화와 같은 조밀한 예측 작업에서 성능 향상과 함께 계산 비용 절감에 기여하는가?
• RQ3 다른 척도에 대해 별도의 브랜치나 다운샘플링 전략을 사용하는 모델보다 통합된 어텐션 메커니즘이 더 우수한 성능을 낼 수 있는가?
• RQ4 표준 컨볼루션 또는 선형 감소 방식과 비교해 볼 때 제안된 토큰 집약 함수는 정확도와 효율성 측면에서 어떻게 성능을 내는가?
• RQ5 피드포워드 레이어에 국소 세부 정보를 통합할 경우, ViT 기반 모델의 특징 표현 능력은 어느 정도 향상되는가?

주요 결과

• 셰운티드 트랜스포머는 모델 파라미터와 FLOPs를 SOTA 포칼 트랜스포머의 50%로 줄였음에도 불구하고 ImageNet에서 84.0%의 Top-1 정확도를 달성한다.
• COCO 객체 검출에서, 유사한 파라미터 및 계산 예산 조건 하에서 포칼 트랜스포머를 1.3 mAP 높이기 위해 초월한다.
• ADE20K 영역 세분화에서, 유사한 효율성 제약 조건 하에서 포칼 트랜스포머보다 mIOU를 2.9 포인트 향상시켰다.
• SegFormer 프레임워크와 함께 사용했을 때, 제안된 백본은 ADE20K에서 48.3%의 mIoU를 기록했으며, MiT-B2보다 1.8 mIoU 높고, 파라미터 수는 25.1M(비교 대상 27.5M)으로 더 적다.
• 디테일 전용 피드포워드 레이어는 ViT에서 표준 피드포워드 레이어 대비 0.2% 향상된 Top-1 정확도를 기록했고, 셰운티드 트랜스포머에서는 0.3% 향상되어 국소 세부 정보 유지에 효과적임을 입증한다.
• 제안된 패치 임베딩 헤드는 비중첩 패치 임베딩 대비 1.4% 향상된 ImageNet 정확도를 기록했고, 겹침 패치 임베딩 대비 0.3% 향상되었으며, FLOPs 증가 폭은 최소한이었다.