[논문 리뷰] DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification

Yongming Rao, Wenliang Zhao외 4명|arXiv (Cornell University)|2021. 06. 03.

Advanced Neural Network Applications참고 문헌 31인용 수 310

한 줄 요약

DynamicViT은 경량 예측 모듈과 주의 마스킹을 사용하여 비전 트랜스포머에서 입력 의존적인 계층적 토큰 가지치기를 도입하고, 상당한 FLOPs 감소를 달성하되 정확도 손실은 최소화합니다.

ABSTRACT

Attention is sparse in vision transformers. We observe the final prediction in vision transformers is only based on a subset of most informative tokens, which is sufficient for accurate image recognition. Based on this observation, we propose a dynamic token sparsification framework to prune redundant tokens progressively and dynamically based on the input. Specifically, we devise a lightweight prediction module to estimate the importance score of each token given the current features. The module is added to different layers to prune redundant tokens hierarchically. To optimize the prediction module in an end-to-end manner, we propose an attention masking strategy to differentiably prune a token by blocking its interactions with other tokens. Benefiting from the nature of self-attention, the unstructured sparse tokens are still hardware friendly, which makes our framework easy to achieve actual speed-up. By hierarchically pruning 66% of the input tokens, our method greatly reduces 31%~37% FLOPs and improves the throughput by over 40% while the drop of accuracy is within 0.5% for various vision transformers. Equipped with the dynamic token sparsification framework, DynamicViT models can achieve very competitive complexity/accuracy trade-offs compared to state-of-the-art CNNs and vision transformers on ImageNet. Code is available at https://github.com/raoyongming/DynamicViT

연구 동기 및 목표

• 정보 이미지 패치의 희소성 활용으로 비전 트랜스포머 가속화 동기 부여.
• 토큰을 점진적이고 입력 의존적으로 가지치기하는 동적 토큰 희소화 프레임워크 제안.
• Gumbel-Softmax 및 주의 마스킹을 통한 미분 가능한 가지치기로 끝-to-end 학습 가능한 예측 모듈 개발.
• ImageNet에서 여러 백본 트랜스포머에 걸친 FLOPs 크게 감소 및 처리량 향상 시연

제안 방법

• 여러 트랜스포머 블록에 경량 예측 모듈을 삽입하여 토큰당 중요도 추정.
• 토큰 특징으로부터 로컬-글로벌 임베딩을 계산하여 각 토큰의 Drop/Keep 확률 예측.
• 학습을 위한 미분성 보장을 위해 Gumbel-Softmax로 이진 Keep/Drop 마스크 샘플링.
• 훈련 중 가지치된 토큰과의 상호 작용을 제거하기 위해 Self-attention에서 주의 마스킹 적용, 학습 중 계산은 균일하게 유지.
• 교차 엔트로피 손실, 교사 백본에 대한 증류 손실, KL 발산, 비율 제약 가지치기 손실의 결합으로 학습.
• 추론 시 학습된 점수에 따라 각 스테이지에서 일정 수의 토큰을 가지치고 목표 유지 비율을 달성

실험 결과

연구 질문

• RQ1 비전 트랜스포머를 불필요한 토큰 가지치기로 정확도 손실 없이 가속화할 수 있는가?
• RQ2 트랜스포머 프레임워크 내에서 동적 토큰 가지치기 메커니즘을 끝-to-end로 어떻게 학습시킬 수 있는가?
• RQ3 계층적이고 입력 의존적인 토큰 가지치기가 다른 백본에서 모델 효율성과 정확도에 어떤 영향을 미치는가?

주요 결과

Base Model	Keeping Ratio ρ at each stage (1.0, 0.9, 0.8, 0.7)	ImageNet Acc. (%)	GFLOPs	Throughput (im/s)
DeiT-S [25]	1.0, 0.9, 0.8, 0.7	79.8	4.6	1337.7
DeiT-S [25]	ρ=1.0
DeiT-S [25]		79.8 (-0.0)	4.0 (-14%)	1524.8 (+14%)
DeiT-S [25]		79.6 (-0.2)	3.4 (-27%)	1774.6 (+33%)
DeiT-S [25]		79.3 (-0.5)	2.9 (-37%)	2062.1 (+54%)
LV-ViT-S [16]	1.0, 0.9, 0.8, 0.7	83.3	6.6	993.3
LV-ViT-S [16]		83.3 (-0.0)	5.8 (-12%)	1108.3 (+12%)
LV-ViT-S [16]		83.2 (-0.1)	5.1 (-22%)	1255.6 (+26%)
LV-ViT-S [16]		83.0 (-0.3)	4.6 (-31%)	1417.6 (+43%)
LV-ViT-M [16]	1.0, 0.9, 0.8, 0.7	84.0	12.7	589.5
LV-ViT-M [16]		83.9 (-0.1)	11.1 (-13%)	688.5 (+17%)
LV-ViT-M [16]		83.9 (-0.1)	9.6 (-24%)	791.2 (+34%)
LV-ViT-M [16]		83.8 (-0.2)	8.5 (-33%)	888.2 (+50%)

• 계층적 토큰 희소화는 입력 토큰의 최대 66%를 가지치고, 31%–37%의 FLOPs를 감소시키는 동시에 처리량을 40% 이상 증가시키며 백본 전반에서 약 0.5%의 정확도 하락만을 보인다.
• DynamicViT는 최첨단 CNN 및 비전 트랜스포머 대비 ImageNet에서 컴플렉시티/정확도 트레이드오프에서 경쟁력을 달성한다.
• 동적 가지치기는 합리적으로 동작하여 이미지 중앙과 객체 주변의 토큰을 보존하고 주변 영역은 가지치며, 가지치기가 진행될수록 해석 가능성을 보인다.
• 동적 토큰 희소화는 모델 확장의 폭 확장에 대한 실행 가능한 대안을 제공하며, 같은 수준 또는 더 나은 효율을 달성한다.
• 더 큰 모델(DeiT-B 및 384x384 입력)은 DynamicViT의 이점을 얻으며 FLOPs 감소와 modest한 정확도 하락을 보인다.
• 부분적 또는 무작위 토큰 제거 전략보다 동적 학습 가지치기의 효과를 제거한다는 점에서 타당성을 보인다.