[논문 리뷰] ATS: Adaptive Token Sampling For Efficient Vision Transformers

Mohsen Fayyaz, Soroush Abbasi Kouhpayegani외 5명|arXiv (Cornell University)|2021. 11. 30.

Advanced Neural Network Applications참고 문헌 32인용 수 8

한 줄 요약

이 논문은 입력 이미지의 복잡도에 따라 동적으로 토큰 수를 줄이는 파라미터가 없는, 미분 가능한 모듈인 Adaptive Token Sampling(ATS)을 제안한다. 각 이미지의 중요도를 평가하고 샘플링함으로써, ATS는 ImageNet에서 정확도를 유지하면서도 GFLOPs를 37% 감소시킨다. 이는 기존 모델에 즉각적으로 통합할 수 있는 효율성 향상 기능을 제공한다.

ABSTRACT

While state-of-the-art vision transformer models achieve promising results for image classification, they are computationally very expensive and require many GFLOPs. Although the GFLOPs of a vision transformer can be decreased by reducing the number of tokens in the network, there is no setting that is optimal for all input images. In this work, we, therefore, introduce a differentiable parameter-free Adaptive Token Sampling (ATS) module, which can be plugged into any existing vision transformer architecture. ATS empowers vision transformers by scoring and adaptively sampling significant tokens. As a result, the number of tokens is not anymore static but it varies for each input image. By integrating ATS as an additional layer within current transformer blocks, we can convert them into much more efficient vision transformers with an adaptive number of tokens. Since ATS is a parameter-free module, it can be added to off-the-shelf pretrained vision transformers as a plug-and-play module, thus reducing their GFLOPs without any additional training. However, due to its differentiable design, one can also train a vision transformer equipped with ATS. We evaluate our module on the ImageNet dataset by adding it to multiple state-of-the-art vision transformers. Our evaluations show that the proposed module improves the state-of-the-art by reducing the computational cost (GFLOPs) by 37% while preserving the accuracy.

연구 동기 및 목표

• 비전 트랜스포머의 높은 계산 비용(추론에 많은 GFLOPs가 필요)을 해결하기 위해.
• 다양한 입력 이미지에 대해 최적화되지 않은 고정된 토큰 수를 가진 비전 트랜스포머의 한계를 극복하기 위해.
• 입력 이미지마다 동적으로 토큰 수를 줄일 수 있는 파라미터가 없고, 미분 가능한 모듈을 개발하기 위해.
• 재학습 없이도 기존 비전 트랜스포머에 통합 가능하면서도, 동시에 엔드 투 엔드 미세조정을 지원하기 위해.
• 분류 정확도를 훼손하지 않으면서도 상당한 계산 효율성 향상을 달성하기 위해.

제안 방법

• ATS는 클래스 토큰의 어텐션 맵에서 토큰의 중요도를 순위 매기는 미분 가능한 스코링 메커니즘을 도입한다.
• 이 모듈은 스코어 기반으로 가장 중요한 토큰들만 적응적으로 샘플링하여, 각 이미지의 총 토큰 수를 줄인다.
• ATS는 표준 트랜스포머 블록 내에 추가 레이어로 삽입되며, 고정된 토큰 처리 파이프라인을 대체한다.
• 이 모듈은 파라미터가 없어, 미세조정 없이 사전학습된 비전 트랜스포머에 직접 배포할 수 있다.
• 미분 가능한 성질 덕분에 원하는 경우 엔드 투 엔드 학습을 지원하여 토큰 샘플링 과정을 최적화할 수 있다.
• 이 방법은 어떤 비전 트랜스포머 아키텍처와도 호환되며, 플러그 앤 플레이 모듈로 적용할 수 있다.

실험 결과

연구 질문

• RQ1 파라미터가 없고, 미분 가능한 모듈이 입력 이미지의 복잡도에 따라 비전 트랜스포머의 토큰 수를 동적으로 줄일 수 있는가?
• RQ2 적응형 토큰 샘플링이 정확도 저하 없이 상당한 계산 절감을 이끌 수 있는가?
• RQ3 제안된 모듈이 재학습 없이도 기존 사전학습된 비전 트랜스포머에 배포 가능한가?
• RQ4 ATS를 적용한 비전 트랜스포머의 성능은 원본 대비 GFLOPs와 정확도 측면에서 어떻게 비교되는가?
• RQ5 ATS를 포함한 전체 모델의 엔드 투 엔드 학습이 효율성과 정확도를 추가로 향상시킬 수 있는가?

주요 결과

• ATS는 ImageNet 데이터셋에서 비전 트랜스포머의 계산 비용을 GFLOPs 기준으로 37% 감소시킨다.
• ATS를 장착한 비전 트랜스포머의 정확도는 원본 대비 변화가 없다.
• 이 모듈은 추가 학습 없이도 기존 사전학습된 비전 트랜스포머에 즉각적으로 플러그 앤 플레이 구성 요소로 적용할 수 있다.
• ATS의 미분 가능한 설계 덕분에 엔드 투 엔드 미세조정이 가능하여 샘플링 과정의 추가 최적화가 가능하다.
• 적응형 토큰 샘플링 메커니즘은 각 이미지에서 가장 정보가 풍부한 토큰들만 성공적으로 식별하고 유지하여 효율성을 향상시킨다.
• 이 방법은 상당히 감소된 추론 복잡도를 유지하면서도 최신 기술 수준의 성능을 유지도한다.