[논문 리뷰] A Unified Pruning Framework for Vision Transformers

Hao Yu, Jianxin Wu|arXiv (Cornell University)|2021. 11. 30.

Advanced Neural Network Applications참고 문헌 43인용 수 5

한 줄 요약

이 논문은 모델 아키텍처를 변경하지 않고 자기주의, 피드포워드 네트워크, 정규화 레이어, 컨볼루션 레이어 등 모든 구성 요소에서 채널을 동시에 정규화하는 통합된 구조적 프루닝 프레임워크인 UP-ViTs를 제안한다. 효율적인 중요도 평가 모듈과 점진적 프루닝 전략을 사용함으로써 UP-ViTs는 압축된 모델에서 최신 기술 수준의 정확도를 달성한다. 예를 들어, UP-DeiT-T는 원본 DeiT-T와 동일한 FLOPs를 유지하면서 ImageNet에서 75.79%의 top-1 정확도를 달성한다.

ABSTRACT

Recently, vision transformer (ViT) and its variants have achieved promising performances in various computer vision tasks. Yet the high computational costs and training data requirements of ViTs limit their application in resource-constrained settings. Model compression is an effective method to speed up deep learning models, but the research of compressing ViTs has been less explored. Many previous works concentrate on reducing the number of tokens. However, this line of attack breaks down the spatial structure of ViTs and is hard to be generalized into downstream tasks. In this paper, we design a unified framework for structural pruning of both ViTs and its variants, namely UP-ViTs. Our method focuses on pruning all ViTs components while maintaining the consistency of the model structure. Abundant experimental results show that our method can achieve high accuracy on compressed ViTs and variants, e.g., UP-DeiT-T achieves 75.79% accuracy on ImageNet, which outperforms the vanilla DeiT-T by 3.59% with the same computational cost. UP-PVTv2-B0 improves the accuracy of PVTv2-B0 by 4.83% for ImageNet classification. Meanwhile, UP-ViTs maintains the consistency of the token representation and gains consistent improvements on object detection tasks.

연구 동기 및 목표

• 모델의 구조적 일관성을 유지하는 효과적이고 일반화 가능한 비전 트랜스포머 프루닝 방법의 부족을 해결한다.
• 토큰 수준의 프루닝이 공간적 구조를 손상시키고 후행 작업으로의 이식성에 악영향을 미치는 한계를 극복한다.
• 기본 ViT와 변형 모델(예: 컨볼루션 또는 계층적 설계를 가진 모델 포함) 모두에 적용 가능한 통합 프레임워크를 개발한다.
• 높은 정확도와 낮은 계산량을 동시에 확보하면서도 객체 검출 및 기타 조밀한 예측 작업으로의 일반화 능력이 뛰어난 압축된 ViT를 가능하게 한다.
• 모델의 통합성을 유지하면서 모든 구성 요소를 동시에 프루닝함으로써 기존의 프루닝 방법보다 뛰어난 성능을 달성한다.

제안 방법

• 기울기 기반 크기 추정을 사용하여 사전 훈련된 ViT 구성 요소의 각 필터의 중요도 점수를 계산하는 효율적인 평가 모듈을 설계한다.
• 모든 구성 요소(다중 헤드 어텐션, 피드포워드 네트워크, 정규화 레이어, 컨볼루션 포함)에서 동시에 균일하게 불필요한 채널을 제거하는 구조적 프루닝을 적용한다.
• 어텐션 헤드의 구조를 유지하고 표현 일관성을 확보하는 데 중점을 둔 다중 헤드 자기주의 레이어의 새로운 채널 제거 전략을 제안한다.
• 성능 유지 목적으로 반복적으로 가장 중요도가 낮은 블록을 제거하고 하이브리드 블록으로 대체하는 점진적 블록 프루닝 방법을 도입한다.
• 프루닝 후 정확도를 유지하기 위해 피넷 훈련 중 지식 정복을 사용하며, 교사 모델은 원본 아키텍처와 일치시킨다.
• DeiT, PVTv2, Swin Transformer를 포함한 여러 ViT 변형 모델에 걸쳐 통합 프루닝 파이프라인을 적용하여, 프루닝 후에도 아키텍처 일관성을 확보한다.

실험 결과

연구 질문

• RQ1 비전 트랜스포머의 모든 구성 요소를 아키텍처 설계를 변경하지 않고 동시에 압축할 수 있는 통합 프루닝 프레임워크를 설계할 수 있는가?
• RQ2 FFN, MSA, 정규화, 컨볼루션을 포함한 모든 레이어에서의 구조적 프루닝이 토큰 수준의 프루닝과 비교해 정확도와 추론 효율성에 어떤 영향을 미치는가?
• RQ3 프루닝 중 원본 토큰 표현을 유지하면 객체 검출과 같은 후행 작업으로의 이식성에 유리한가?
• RQ4 ViT의 너비와 깊이를 압축할 때, 한 번에 프루닝하는 것과 점진적 프루닝 전략 중 어느 것이 더 효과적인가?
• RQ5 기존 최신 기술 수준의 압축된 ViT보다 더 넓은 아키텍처 및 작업 범위에서 일반화 능력이 뛰어난가?

주요 결과

• UP-DeiT-T는 동일한 FLOPs를 유지하면서 ImageNet에서 75.79%의 top-1 정확도를 달성하여 원본 DeiT-T보다 3.59% 높은 성능을 보였다.
• UP-PVTv2-B0는 동일한 계산 비용에서 원본 PVTv2-B0보다 ImageNet 분류에서 4.83% 높은 정확도를 기록했다.
• UP-DeiT-T를 10개 블록으로 압축할 때, 점진적 블록 프루닝이 한 번에 블록 제거하는 것보다 높은 정확도(73.78%)를 달성했다.
• 컴ponent 1(너비)을 한 번에 프루닝할 경우 66.95%의 정확도를 기록했으며, 이는 동일한 작업에서 점진적 프루닝(65.05%)보다 높은 성능을 보였다.
• 프루닝된 UP-ViTs 모델들은 토큰 표현을 유지함으로써 객체 검출과 같은 후행 작업 전반에서 일관된 성능 향상을 보였다.
• 프레임워크는 DeiT, PVTv2, Swin Transformer를 포함한 다양한 ViT 변형 모델에 효과적으로 일반화되어 있으며, 광범위한 적용 가능성과 강건성을 입증했다.