[논문 리뷰] Unified Visual Transformer Compression

Shixing Yu, Tianlong Chen외 6명|arXiv (Cornell University)|2022. 03. 15.

Advanced Neural Network Applications인용 수 28

한 줄 요약

UVC는 예산 하에서 ViT의 가지치기, 블록 건너뛰기 및 지식 증류를 공동으로 최적화하여 ImageNet에서 DeiT 및 T2T-ViT 백본에 대해 FLOPs를 크게 감소시키고 정확도 손실은 최소로 유지합니다.

ABSTRACT

Vision transformers (ViTs) have gained popularity recently. Even without customized image operators such as convolutions, ViTs can yield competitive performance when properly trained on massive data. However, the computational overhead of ViTs remains prohibitive, due to stacking multi-head self-attention modules and else. Compared to the vast literature and prevailing success in compressing convolutional neural networks, the study of Vision Transformer compression has also just emerged, and existing works focused on one or two aspects of compression. This paper proposes a unified ViT compression framework that seamlessly assembles three effective techniques: pruning, layer skipping, and knowledge distillation. We formulate a budget-constrained, end-to-end optimization framework, targeting jointly learning model weights, layer-wise pruning ratios/masks, and skip configurations, under a distillation loss. The optimization problem is then solved using the primal-dual algorithm. Experiments are conducted with several ViT variants, e.g. DeiT and T2T-ViT backbones on the ImageNet dataset, and our approach consistently outperforms recent competitors. For example, DeiT-Tiny can be trimmed down to 50\% of the original FLOPs almost without losing accuracy. Codes are available online:~\url{https://github.com/VITA-Group/UVC}.

연구 동기 및 목표

• 정확도를 희생하지 않으면서 계산 비용을 줄여 비전 트랜스포머의 효율적 배치를 촉진한다.
• 단일 엔드-투-엔드 최적화에서 여러 압축 기법을 결합한 통합 프레임워크를 제안한다.
• 전역 자원 예산 하에서 가지치기, 건너뛰기, 증류의 자동 균형 조정을 가능하게 한다.
• ImageNet에서 DeiT 및 T2T-ViT와 같은 인기 있는 ViT 변형에 대해 UVC를 평가하고 최첨단 가지치기 방법과 비교한다.

제안 방법

• ViT 압축을 제약된 엔드-투-엔드 최적화 문제로 공식화한다.
• 증류 손실을 포함해 네트워크 가중치, 층별 가지치기 마스크/비율 및 건너뛰기 구성을 공동으로 학습한다.
• 각 트랜스포머 블록 내 주의 헤드와 MLP 뉴런을 가지치기하기 위해 이중 수준 그룹 희소성을 사용한다.
• 트랜스포머 블록 간의 건너뛰기 구성을 위한 이진 게이팅 변수를 도입하여 블록 드롭을 가능하게 한다.
• 프라이멀-듀얼 알고리즘으로 미니맥스 최적화를 해결한다.
• DeiT 및 T2T-ViT 백본에 대한 참조 구현 및 실험 결과를 제공한다.

실험 결과

연구 질문

• RQ1 ViT 블록 내 및 블록 간 건너뛰기 구성을 통해 자원 예산 하에서 가지치기를 공동으로 최적화할 수 있는가?
• RQ2 가지치기, 건너뛰기 및 지식 증류를 통합하는 것이 단일 방법의 ViT 압축보다 성능이 우수한가?
• RQ3 다양한 FLOPs 예산에서 DeiT 및 T2T-ViT에 대해 정확도/지연 상충이 어떻게 나타나는가?
• RQ4 ImageNet에서 UVC와 기존 패치 감소 및 가지치기 방법은 어떻게 비교되는가?

주요 결과

• UVC는 DeiT-Tiny에서 약 50%의 FLOPs 감소를 달성하면서 최소한의 정확도 손실(증류/비증류 설정에서 각각 0.3%/0.9%)을 보인다.
• UVC는 동일하거나 더 큰 FLOPs 감소에서 여러 현대 ViT 압축 방법(SViTE, PatchSlimming, HVT, IA-RED 2 등)보다 정확도 면에서 우수하다.
• DeiT-Base의 경우 약 45%의 FLOPs 감소를 달성하며 허용 가능한 정확도 감소를 보이고, 일부 기준선이 안정적으로 학습하기 어려운 경우에서도 경쟁력 있는 결과를 제공한다.
• T2T-ViT-14에 적용했을 때 상당한 FLOPs 절감(예: 44–60% 범위)과 소폭의 정확도 저하를 달성하여 일부 패치 기반 감소 방법보다 우수하다.
• 절단/가지치기/증류를 공동으로 최적화하면 단일 구성 요소나 순차 적용보다 더 나은 트레이드오프를 얻을 수 있다.
• UVC의 코드가 공개되어 있다.
• (본문 주석) UVC는 ImageNet에서 여러 ViT 변형에서 강력한 성능을 보인다.