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

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

Cell Image Analysis Techniques인용 수 4

한 줄 요약

이 논문은 입력 이미지 콘텐츠에 따라 동적으로 토큰 수를 줄이는 미분 가능하고 파라미터가 없는 모듈인 Adaptive Token Sampling (ATS)을 제안한다. 분류 토큰의 어텐션 웨이트를 사용해 역변환 샘플링을 통해 중요도를 평가하고 유의미한 패치를 샘플링함으로써, ImageNet, Kinetics-400, Kinetics-600에서 정확도 손실 없이 GFLOPs를 최대 2배까지 감소시킨다. 이는 사전 훈련된 모델에 즉시 통합할 수 있는 효율성 향상 솔루션을 제공한다.

ABSTRACT

While state-of-the-art vision transformer models achieve promising results in image classification, they are computationally 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 Sampler (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 constant anymore and varies for each input image. By integrating ATS as an additional layer within the 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 the off-the-shelf pre-trained vision transformers as a plug and play module, thus reducing their GFLOPs without any additional training. Moreover, due to its differentiable design, one can also train a vision transformer equipped with ATS. We evaluate the efficiency of our module in both image and video classification tasks by adding it to multiple SOTA vision transformers. Our proposed module improves the SOTA by reducing their computational costs (GFLOPs) by 2X, while preserving their accuracy on the ImageNet, Kinetics-400, and Kinetics-600 datasets.

연구 동기 및 목표

• 시퀀스 길이에 따라 제곱적으로 증가하는 비전 트랜스포머의 높은 계산 비용을 해결하기 위해.
• 고정 비율의 프루닝을 피하고 입력 이미지에 따라 다를 수 있는 동적이고 콘텐츠 인식형 토큰 감소를 가능하게 하기 위해.
• 재훈련 없이도 사전 훈련된 모델에 통합할 수 있는 파라미터가 없고, 미분 가능한 모듈을 설계하기 위해.
• 이미지 및 비디오 분류 벤치마크에서 정확도를 유지하거나 향상시키면서 GFLOPs를 줄이기 위해.
• 아키텍처 재훈련 없이도 엣지 디바이스에서 효율적인 추론을 위한 플러그 앤 플레이 솔루션을 제공하기 위해.

제안 방법

• ATS는 분류 토큰의 어텐션 웨이트를 사용해 모든 입력 토큰의 중요도 점수를 계산한다.
• 이 점수들의 누적 분포를 기반으로 역변환 샘플링을 적용하여 유의미한 토큰의 부분집합을 선택한다.
• 선택된 토큰들은 어텐션 웨이트와 값의 가중 조합을 사용해 부드럽게 다운샘플링되어 특징 표현을 유지한다.
• 이 모듈은 미분 가능하고 파라미터가 없어, 추가 훈련 없이도 어떤 비전 트랜스포머 블록에나 통합할 수 있다.
• 기존의 사전 훈련된 모델에 플러그 앤 플레이 레이어로 삽입하거나, 브레인보드와 함께 공동 훈련할 수 있다.
• 여러 스테이지에 걸쳐 점진적인 토큰 프루닝이 가능하며, 복잡도에 따라 이미지당 토큰 수가 변동된다.

실험 결과

연구 질문

• RQ1 입력 콘텐츠에 따라 동적으로 토큰 수를 줄일 수 있는 파라미터가 없고, 미분 가능한 모듈이 비전 트랜스포머의 토큰 수를 줄일 수 있는가?
• RQ2 고정 비율 프루닝과 비교해 복잡성 비용을 줄이면서도 정확도를 유지할 수 있는가?
• RQ3 ATS는 사전 훈련된 비전 트랜스포머에 피니튜닝 없이 플러그 앤 플레이 모듈로 적용할 수 있는가?
• RQ4 여러 스테이지에 걸쳐 점진적인 토큰 샘플링이 효율성과 정확도에 어떤 영향을 미치는가?
• RQ5 학습 가능한 점수 네트워크보다 어텐션 기반 중요도 점수 평가가 효율성과 일반화 성능 측면에서 뛰어나다고 할 수 있는가?

주요 결과

• ATS는 DeiT-S 및 기타 최신 비전 트랜스포머에서 GFLOPs를 최대 2배까지 감소시키면서 ImageNet에서 상위 1위 정확도를 유지하거나 향상시킨다.
• DeiT-S+ATS 모델은 단지 2.9 GFLOPs로 79.7%의 상위 1위 정확도를 달성하여 동일한 FLOP 수준에서 기준 모델인 DeiT-S(78.9%)를 초월한다.
• 후기 스테이지(예: 스테이지 3)에 ATS를 통합할 경우, 초기 스테이지(73.1%)보다 더 높은 정확도(78.5%)를 기록한다. 이는 더 신뢰할 수 있는 어텐션 점수 때문이며.
• 스테이지 3~11에 걸쳐 ATS를 다중 스테이지로 통합할 경우, 2.6 GFLOPs에서 79.2%의 상위 1위 정확도를 달성하여 점진적인 효율성 향상을 입증한다.
• EViT-DeiT-S에 ATS를 추가하면 GFLOPs를 3.0에서 2.5로 감소시키며 정확도 손실 없이도 성능 향상을 보이며, 정적 프루닝보다 호환성과 우수성을 입증한다.
• 시각화 결과는 모델이 복잡하고 혼잡한 이미지에서는 더 많은 토큰을 샘플링하고, 단순한 이미지에서는 더 적은 토큰을 선택하며, 주목할 만한 영역에 더 집중된 어텐션을 보임을 확인한다.