[논문 리뷰] AdaViT: Adaptive Tokens for Efficient Vision Transformer

Hongxu Yin, Arash Vahdat외 4명|arXiv (Cornell University)|2021. 12. 14.

CCD and CMOS Imaging Sensors인용 수 16

한 줄 요약

A-ViT는 입력 이미지의 복잡도에 따라 추론 중 처리되는 토큰 수를 동적으로 줄이는 적응형 토큰 계산 메커니즘을 제안한다. 기존 네트워크 파라미터를 재사용하여 각 토큰의 정지 확률을 학습하고, 분포 기반 사전 정규화를 적용함으로써, 아키텍처 변경 없이도 DeiT-Tiny의 경우 62⁠% 높은 처리량과 DeiT-Small의 경우 38⁠% 높은 처리량을 달성하면서 정확도 손실은 0.3⁠%에 불과하여 하드웨어에 최적화된 효율적인 추론을 가능하게 한다.

ABSTRACT

We introduce A-ViT, a method that adaptively adjusts the inference cost of vision transformer (ViT) for images of different complexity. A-ViT achieves this by automatically reducing the number of tokens in vision transformers that are processed in the network as inference proceeds. We reformulate Adaptive Computation Time (ACT) for this task, extending halting to discard redundant spatial tokens. The appealing architectural properties of vision transformers enables our adaptive token reduction mechanism to speed up inference without modifying the network architecture or inference hardware. We demonstrate that A-ViT requires no extra parameters or sub-network for halting, as we base the learning of adaptive halting on the original network parameters. We further introduce distributional prior regularization that stabilizes training compared to prior ACT approaches. On the image classification task (ImageNet1K), we show that our proposed A-ViT yields high efficacy in filtering informative spatial features and cutting down on the overall compute. The proposed method improves the throughput of DeiT-Tiny by 62% and DeiT-Small by 38% with only 0.3% accuracy drop, outperforming prior art by a large margin. Project page at https://a-vit.github.io/

연구 동기 및 목표

• 입력에 따라 변하는 추론을 가능하게 하여 이미지의 복잡도에 따라 계산 비용을 감소시키는 것.
• 추가 파라미터나 네트워크 아키텍처 변경 없이도, 정보가 적은 공간적 토큰의 처리를 적절히 정지시키는 메커니즘을 설계하는 것.
• 새로운 분포 기반 사전을 도입하여 학습을 안정화하고, 정지 깊이를 목표 계산 예산에 맞추는 것.
• 표준 하드웨어에서 모델 정확도를 유지하면서도 상당한 처리량 향상을 달성하는 것.

제안 방법

• A-ViT는 기존의 적응형 계산 시간(Adaptive Computation Time, ACT)을 토큰 수준로 확장하여, 각 토큰이 정지 조건을 충족하면 더 이상 처리하지 않도록 허용한다.
• 각 토큰의 정지 확률은 각 트랜스포머 블록의 마지막 밀집층에서 유도된 단일 뉴런을 통해 계산되며, 추가 파라미터나 계산 없이도 기존 파라미터를 재사용한다.
• 기대 정지 깊이를 목표로 향하도록 이끄는 분포 기반 사전 정규화를 도입하여 학습의 안정성과 수렴성을 향상시킨다.
• 정보가 많은 토큰들만 깊이 처리되는 동적이고 공간적으로 적응형 계산을 가능하게 하여 FLOPs와 추론 시간을 감소시킨다.
• 보조 네트워크나 최적화가 필요 없이, 휴면 손실과 분포 기반 사전을 함께 엔드 투 엔드로 학습한다.

실험 결과

연구 질문

• RQ1 아키텍처 변경 없이도 토큰 수준에서의 적응형 계산이 Vision Transformers의 추론 효율성을 향상시킬 수 있는가?
• RQ2 기존 네트워크 파라미터만을 사용하여 정지 결정을 효율적으로 학습할 수 있는가?
• RQ3 분포 기반 사전이 학습 안정성과 목표 계산 예산 향한 수렴을 향상시킬 수 있는가?
• RQ4 적응형 토큰 정지가 의미론적 이미지 콘텐츠와 인간의 인지와 어느 정도 관련이 있는가?

주요 결과

• A-ViT는 ImageNet1K에서 DeiT-Tiny의 처리량을 62⁠% 향상시키고 DeiT-Small의 경우 38⁠% 향상시키며, 정확도 손실은 단지 0.3⁠%에 그친다.
• 세밀한 공간적 적응형 토큰 제거 덕분에 기존 ACT보다 FLOPs를 25⁠% 감소시킨다.
• 분포 기반 사전 정규화를 통해 목표 정지 깊이에 더 빠르게 수렴하고, 휴면 손실만 사용하는 경우 대비 6.4⁠%의 정확도 향상을 달성한다.
• 정지 계산에 단일 벡터 요소만을 사용할 경우 DeiT-Tiny의 정확도 손실은 0.08⁠%, DeiT-Small의 경우 0.04⁠%에 불과하여 최소한의 파라미터 오버헤드를 확인한다.
• 정지 학습을 위한 보조 네트워크를 추가하면 정확도 향상은 0.06⁠%에 그치지만 처리량은 12.6⁠% 감소시켜 실용성이 떨어진다.