[논문 리뷰] Less is More: Pay Less Attention in Vision Transformers

Zizheng Pan, Bohan Zhuang외 3명|arXiv (Cornell University)|2021. 05. 29.

Advanced Neural Network Applications인용 수 8

한 줄 요약

이 논문은 계층적 비전 트랜스포머(LIT)를 제안하며, 계산 비용을 줄이면서도 성능을 유지하기 위해 초기 자기주의 어텐션 레이어를 다층 퍼셉트론(MLP)으로 대체한다. 자기주의 어텐션을 더 깊은 레이어에서만 적용하고, 학습 가능한 비틀림형 토큰 병합(DTM) 모듈을 도입함으로써, 기존 비전 트랜스포머에 비해 상당히 감소된 FLOPs와 메모리 사용량으로 이미지 분류, 객체 검출, 인스턴스 세그멘테이션에서 최신 기준 성능을 달성한다.

ABSTRACT

Transformers have become one of the dominant architectures in deep learning, particularly as a powerful alternative to convolutional neural networks (CNNs) in computer vision. However, Transformer training and inference in previous works can be prohibitively expensive due to the quadratic complexity of self-attention over a long sequence of representations, especially for high-resolution dense prediction tasks. To this end, we present a novel Less attention vIsion Transformer (LIT), building upon the fact that the early self-attention layers in Transformers still focus on local patterns and bring minor benefits in recent hierarchical vision Transformers. Specifically, we propose a hierarchical Transformer where we use pure multi-layer perceptrons (MLPs) to encode rich local patterns in the early stages while applying self-attention modules to capture longer dependencies in deeper layers. Moreover, we further propose a learned deformable token merging module to adaptively fuse informative patches in a non-uniform manner. The proposed LIT achieves promising performance on image recognition tasks, including image classification, object detection and instance segmentation, serving as a strong backbone for many vision tasks. Code is available at: https://github.com/zhuang-group/LIT

연구 동기 및 목표

• 계층적 비전 트랜스포머(HVT)의 초기 단계에서 자기주의 어텐션의 높은 계산 비용과 메모리 소비 문제를 해결하기 위해, 특히 고해상도 특징 맵에서 발생하는 문제를 다루기 위해.
• 자기주의 어テン션의 계산 복잡도가 제곱형임을 감안할 때, HVT의 초기 자기주의 어텐션 레이어가 최종 성능에 의미 있는 기여를 하는지 조사하기 위해.
• 표준 균일한 패치 병합 방식보다 더 적응적이고 정보량이 풍부한 토큰 병합 메커니즘을 설계하여 계층적 표현 학습을 향상시키기 위해.
• 밀도 예측 작업에서 성능을 희생시키지 않은 채 모델 정확도와 추론 효율성 사이의 더 나은 트레이드오프를 달성하기 위해.

제안 방법

• 고해상도 특징에서의 비용이 큰 자기주의 어텐션을 피하기 위해, 계층적 비전 트랜스포머의 첫 두 스테이지에 순수한 MLP 블록을 도입하여 국소 패턴을 캡처한다.
• 표준 다중헤드 자기주의 어텐션(MSA)을 시퀀스 길이가 감소한 더 깊은 스테이지에서만 적용하여 장거리 의존성 모델링을 효율적으로 수행한다.
• 학습 가능한 공간 오프셋을 통해 국소 영역에서 정보가 풍부한 패치를 적응적으로 샘플링하고 병합하는 비틀림형 토큰 병합(DTM) 모듈을 제안한다. 이는 기하학적 변환 모델링을 향상시킨다.
• 고정된 균일한 패치 병합 방식보다는 학습된 비균일 병합 전략을 사용하여 주목할 만한 공간 영역에 집중함으로써 표준 컨볼루션 또는 선형 패치 병합보다 성능을 높인다.
• PVT와 Swin 트랜스포머와의 공정한 비교를 위해 일관된 아키텍처 하이퍼파라미터를 설계하며, 패치 크기와 채널 차원을 동일하게 유지한다.
• 성능 및 효율성 향상을 검증하기 위해 ImageNet, COCO, ADE20K 벤치마크에서 LIT를 훈련하고 평가한다.

실험 결과

연구 질문

• RQ1 고해상도 특징에서 자기주의 어텐션의 계산 복잡도가 제곱형임을 감안할 때, 비전 트랜스포머의 초기 자기주의 어텐션 레이어가 성능 향상에 의미 있는 기여를 하는가?
• RQ2 초기 자기주의 어텐션 레이어를 MLP로 대체함으로써 FLOPs와 메모리 사용량을 크게 줄일 수 있으며, 이와 동시에 성능을 유지하거나 향상시킬 수 있는가?
• RQ3 학습된 비틀림형 토큰 병합 메커니즘이 계층적 비전 트랜스포머에서 표준 고정된 균일한 패치 병합 방식보다 특징 표현 품질을 향상시킬 수 있는가?
• RQ4 자기주의 어텐션을 더 깊은 레이어에서만 선택적으로 적용함으로써 비전 트랜스포머 설계에서 성능-효율성 최적점(sweet spot)을 도달할 수 있는가?
• RQ5 다양한 비전 벤치마크에서 제안된 LIT 아키텍처가 정확도, FLOPs, 메모리 소비 측면에서 최신 기준 비전 트랜스포머와 비교해 어떻게 성능을 내는가?

주요 결과

• LIT는 뿌리 모델보다 더 적은 파라미터를 사용하면서도 1.8G FLOPs로만 85.7%의 Top-1 정확도를 기록하여 ImageNet-1K 분류에서 최신 기준 성능을 달성한다.
• COCO 객체 검출 및 인스턴스 세그멘테이션에서 LIT-B는 각각 46.8 AP와 42.3 마스크 AP를 기록하며, Swin-B와 PVT-B와 비슷하거나 이를 초월하는 성능을 보이며 더 낮은 FLOPs를 사용한다.
• 비틀림형 토큰 병합(DTM) 모듈은 객체의 형태와 기하학적 변환에 따라 적응적으로 공간 오프셋을 학습하여, 표준 병합 방식보다 특징 표현 품질을 향상시킨다.
• 시각화 결과에 따르면, 표준 HVT의 초기 MSA 레이어는 오직 작은 국소 영역에만 주목함을 보이며, 초기 단계에서 자기주의 어텐션의 기여가 미미함을 시사한다.
• Swin-Tiny와 PVT-S에 비해 최대 30%까지 FLOPs를 감소시키면서도 정확도를 유지하거나 향상시키며, 강력한 효율성-정확도 트레이드오프를 입증한다.
• 절단 분석 결과, 초기 자기주의 어텐션을 MLP로 대체함으로써 성능 저하 없이 뚜렷한 효율성 향상이 이루어지며, LIT의 핵심 설계 원칙을 검증한다.