[논문 리뷰] Vision Transformer Architecture Search

Xiu Su, Shan You외 7명|arXiv (Cornell University)|2021. 06. 25.

Advanced Neural Network Applications인용 수 10

한 줄 요약

이 논문은 공유 가중치와 개인화된 클래스 토큰을 사용하는 슈퍼트랜스포머를 활용해 다양한 아키텍처를 효율적으로 탐색하는 비전 트랜스포머를 위한 신경망 아키텍처 탐색 방법 ViTAS를 제안한다. 하드웨어 제약 조건 하에 깊이, 너비, 시퀀스 길이, 어텐션 헤드를 종합적으로 탐색함으로써 ViTAS는 1.3G FLOPs에서 ImageNet에서 74.7%의 top-1 정확도를 달성하여 기준 ViT보다 2.5% 높은 성능을 보였다.

ABSTRACT

Recently, transformers have shown great superiority in solving computer vision tasks by modeling images as a sequence of manually-split patches with self-attention mechanism. However, current architectures of vision transformers (ViTs) are simply inherited from natural language processing (NLP) tasks and have not been sufficiently investigated and optimized. In this paper, we make a further step by examining the intrinsic structure of transformers for vision tasks and propose an architecture search method, dubbed ViTAS, to search for the optimal architecture with similar hardware budgets. Concretely, we design a new effective yet efficient weight sharing paradigm for ViTs, such that architectures with different token embedding, sequence size, number of heads, width, and depth can be derived from a single super-transformer. Moreover, to cater for the variance of distinct architectures, we introduce extit{private} class token and self-attention maps in the super-transformer. In addition, to adapt the searching for different budgets, we propose to search the sampling probability of identity operation. Experimental results show that our ViTAS attains excellent results compared to existing pure transformer architectures. For example, with $1.3$G FLOPs budget, our searched architecture achieves $74.7\%$ top-$1$ accuracy on ImageNet and is $2.5\%$ superior than the current baseline ViT architecture. Code is available at \url{this https URL}.

연구 동기 및 목표

• 비전 특화 최적화 없이 NLP 모델에서 유래된 비전 트랜스포머(ViTs)의 비최적 설계 문제를 해결하기 위해.
• 고정된 하드웨어 예산 하에서 깊이, 너비, 시퀀스 길이, 어テン션 헤드 수 등 여러 아키텍처 차원을 효율적으로 탐색하기 위해.
• 다양한 비전 트랜스포머 아키텍처를 하나의 슈퍼트랜스포머로 표현할 수 있도록 가중치 공유 메커니즘을 설계하기 위해.
• 슈퍼네트워크 내에서 아키텍처 변동성을 처리하기 위해 개인화된 클래스 토큰과 어텐션 맵을 도입하기 위해.
• 신뢰도 있는 정확도를 확보하면서도 다양한 FLOP 예산에 적응하기 위해 정체성 연산에 대한 샘플링 확률을 학습하는 방식으로 탐색을 조정하기 위해.

제안 방법

• 모든 아키텍처 변형 간에 공유된 어텐션 및 피드포워드 가중치를 갖는 슈퍼트랜스포머를 설계하기 위해.
• 각 아키텍처 변형에 대해 아키텍처 다양성을 처리하기 위해 개인화된 클래스 토큰과 자기 어텐션 맵을 도입하기 위해.
• 정체성 연산에 대한 학습 가능한 샘플링 확률을 사용하는 미분 가능 아키텍처 탐색 전략을 구현하기 위해.
• 토큰 임bedding 크기, 시퀀스 길이, 어텐션 헤드 수, 너비, 깊이의 변형을 포함하는 탐색 공간을 사용하기 위해.
• 정확도와 FLOP 예산을 균형 잡는 데 초점을 맞춘 미분 가능 탐색 목적 함수를 사용해 슈퍼네트워크를 최적화하기 위해.
• 성능을 유지하면서도 탐색 비용을 줄이기 위해 조기 정지와 프루닝을 적용하기 위해.

실험 결과

연구 질문

• RQ1 깊이, 너비, 시퀀스 길이가 다른 다양한 비전 트랜스포머 아키텍처를 효과적으로 표현할 수 있는 통합 슈퍼네트워크가 가능한가?
• RQ2 개인화된 클래스 토큰과 어텐션 맵의 포함이 공유 아키텍처 탐색 공간 내에서 성능 향상에 어떻게 기여하는가?
• RQ3 학습 가능한 샘플링 확률을 활용한 미분 가능 아키텍처 탐색이 고정된 FLOP 예산 하에서 ViT 정확도 향상에 얼마나 기여하는가?
• RQ4 ViTAS는 ImageNet에서 기존의 비전 트랜스포머 아키텍처와 비교해 정확도와 효율성 측면에서 어떻게 다른가?
• RQ5 제안된 가중치 공유 기법이 성능을 저하시키지 않으면서도 효율적이고 효과적인 아키텍처 탐색을 가능하게 하는가?

주요 결과

• ViTAS는 1.3G FLOPs 예산 하에서 ImageNet에서 74.7%의 top-1 정확도를 달성하여 기준 ViT보다 2.5% 높은 성능을 보였다.
• 같은 하드웨어 제약 조건 하에서 기존의 순수 트랜스포머 모델보다 검색된 아키텍처가 뛰어난 정확도를 보였다.
• 개인화된 클래스 토큰과 어텐션 맵의 사용은 슈퍼네트워크 내에서 다양한 아키텍처 간 간섭을 줄여 성능 향상에 크게 기여했다.
• 적응형 샘플링 확률을 활용한 미분 가능 탐색은 최소한의 계산 오버헤드로 아키텍처 공간을 효율적으로 탐색할 수 있게 했다.
• 가중치 공유 메커니즘은 깊이, 너비, 시퀀스 길이, 어텐션 헤드 수를 포함한 다양한 아키텍처 차원에 걸쳐 하나의 슈퍼네트워크 내에서 효과적인 아키텍처 탐색을 가능하게 했다.
• 제시된 예산 하에서 순수 트랜스포머 모델 중에서 ViTAS는 이미지넷에서 최고 성능을 기록했다.