[논문 리뷰] Chasing Sparsity in Vision Transformers: An End-to-End Exploration

Tianlong Chen, Yu Cheng외 4명|arXiv (Cornell University)|2021. 06. 08.

Advanced Neural Network Applications참고 문헌 95인용 수 85

한 줄 요약

이 논문은 동적 희소 하위망, 구조적 희소성, 데이터-아키텍처 공동 희소성을 통해 정확도를 유지하면서 학습 메모리와 추론 비용을 줄이기 위해 Vision Transformers (ViTs)에 대한 엔드-투-엔드 희소 학습을 도입한다.

ABSTRACT

Vision transformers (ViTs) have recently received explosive popularity, but their enormous model sizes and training costs remain daunting. Conventional post-training pruning often incurs higher training budgets. In contrast, this paper aims to trim down both the training memory overhead and the inference complexity, without sacrificing the achievable accuracy. We carry out the first-of-its-kind comprehensive exploration, on taking a unified approach of integrating sparsity in ViTs "from end to end". Specifically, instead of training full ViTs, we dynamically extract and train sparse subnetworks, while sticking to a fixed small parameter budget. Our approach jointly optimizes model parameters and explores connectivity throughout training, ending up with one sparse network as the final output. The approach is seamlessly extended from unstructured to structured sparsity, the latter by considering to guide the prune-and-grow of self-attention heads inside ViTs. We further co-explore data and architecture sparsity for additional efficiency gains by plugging in a novel learnable token selector to adaptively determine the currently most vital patches. Extensive results on ImageNet with diverse ViT backbones validate the effectiveness of our proposals which obtain significantly reduced computational cost and almost unimpaired generalization. Perhaps most surprisingly, we find that the proposed sparse (co-)training can sometimes improve the ViT accuracy rather than compromising it, making sparsity a tantalizing "free lunch". For example, our sparsified DeiT-Small at (5%, 50%) sparsity for (data, architecture), improves 0.28% top-1 accuracy, and meanwhile enjoys 49.32% FLOPs and 4.40% running time savings. Our codes are available at https://github.com/VITA-Group/SViTE.

연구 동기 및 목표

• ViTs에서 엔드-투-엔드 희소성을 motivate하고 가능하게 하여 학습 메모리와 추론 비용을 줄인다.
• 고정된 파라미터 예산에서 정확도를 유지하거나 개선하는 희소 ViT 학습 방법을 개발한다.
• 하드웨어 효율성을 위해 비정형에서 구조적 형태로 희소성을 확장한다.
• 데이터 희소성과 아키텍처 희소성을 공동 탐색하여 추가적인 효율성을 확보한다.
• DeiT 백본 전반에 걸쳐 ImageNet에서 FLOPs와 대기 시간 절감이 큰 효과를 나타냄을 입증한다.

제안 방법

• 고정된 파라미터 예산으로 희소 ViT를 학습하기 위한 Sparse Vision Transformer Exploration (SViTE)을 도입한다.
• 하드웨어 친화적 희소성을 가이드하는 prune-and-grow 방식으로 self-attention 헤드를 이끄는 Structured Sparse ViT Exploration (S2ViTE)로 확장한다.
• 정보 토큰을 선택하고 모델을 희소화하는 것을 함께 수행하는 Sparse Vision Transformer Co-Exploration (SViTE+)를 추가한다.
• 헤드 중요도를 점수화하기 위한 Taylor-expansion 기반 프록시와 pruning을 위한 MLP 뉴런 점수화를 위한 L1 노름을 사용한다.
• Gumbel-Softmax 및 straight-through 트릭을 갖춘 학습 가능 토큰 선택기를 도입하여 상위-k 정보를 포함하는 patches를 선택하고 데이터 희소성을 가능하게 한다.

실험 결과

연구 질문

• RQ1 엔드-투-엔드 희소 학습이 정확도 손실 없이 상당한 FLOPs 및 파라미터 감소를 이룰 수 있는가?
• RQ2 구조적 희소성(예: 주의 헤드 가지치기)이 비정형 희소성에 비해 하드웨어 친화적 이점을 제공하는가?
• RQ3 데이터 희소성(토큰 선택)과 아키텍처 희소성의 공동 탐색이 성능을 해치지 않으면서 추가적인 효율성 향상을 가져올 수 있는가?
• RQ4 희소한 ViT 모델이 암묵적 정규화로 인해 일반화가 더 잘되거나 특정 희소성 regimes에서 정확도가 향상될 수 있는가?
• RQ5 ImageNet에서 DeiT-Tiny/Small/Base 백본에 걸친 희소성 전략의 성능은 어떠하며, 작은 희소 네트워크가 더 큰 밀집 네트워크를 능가하는가?

주요 결과

• SViTE는 Backbone과 희소성에 따라 25.56%–57.50% 범위의 FLOPs 감소를 달성하며, 정확도 손실은 거의 없고(일반적으로 0.5% 이내) 나타난다.
• SViTE+는 토큰 선택과 함께 DeiT-Small에서 데이터 희소성 50%, 모델 희소성 5%에서 최대 4.40%의 정확도 개선을 달성하고 FLOPs를 49.32% 감소시키며 실행 시간도 4.40% 감소시킨다.
• S2ViTE는 구조적 희소성이 비구조적 변종과 일치하거나 이를 능가하며, 밀집 베이스라인을 능가하면서도 실행 시간 감소를 크게 달성한다(예: 최대 24.70% 감소 보고).
• 데이터 희소성은 규제자 역할을 하며; SViTE+-Small은 토큰을 최대 10%까지 줄이고 이에 상응하는 실행 시간 및 FLOPs 절감과 때로는 정확도 향상을 얻는다.
• 구조적 희소성(S2ViTE)은 여러 설정에서 SSP를 능가할 수 있으며, 40% 구조적 희소성의 S2ViTE-Base는 밀집 DeiT-Base보다 정확도에서 최대 1.24% 높고 FLOPs를 약 34% 감소시킨다.
• ImageNet-1K의 DeiT-Tiny/Small/Base 전반에 걸쳐 제안된 방법들이 일관된 효율성 향상을 제공하고 정확도도 경쟁력이 있거나 개선되며, 더 작은 희소 네트워크가 더 큰 밀집 상대보다 더 나은 성능을 보일 수 있다.