[논문 리뷰] PPT: Token Pruning and Pooling for Efficient Vision Transformers

WU Xin-jian, Fanhu Zeng외 2명|arXiv (Cornell University)|2023. 10. 03.

Advanced Neural Network Applications인용 수 6

한 줄 요약

이 논문은 가중치 학습이 없는 조건에서 계산의 부과적 중복을 줄이기 위해 토큰 프루닝과 풀링을 적응적으로 조합하는 새로운 프레임워크 PPT를 제안한다. 레이어 깊이와 입력에 따라 달라지는 토큰 중요도 분포에 따라 프루닝 또는 풀링을 동적으로 선택함으로써, DeiT-S에서 이미지 분류 성능에 영향을 주지 않고 FLOPs를 37% 이상 감소시키고 처리량을 45% 향상시킨다.

ABSTRACT

Vision Transformers (ViTs) have emerged as powerful models in the field of computer vision, delivering superior performance across various vision tasks. However, the high computational complexity poses a significant barrier to their practical applications in real-world scenarios. Motivated by the fact that not all tokens contribute equally to the final predictions and fewer tokens bring less computational cost, reducing redundant tokens has become a prevailing paradigm for accelerating vision transformers. However, we argue that it is not optimal to either only reduce inattentive redundancy by token pruning, or only reduce duplicative redundancy by token merging. To this end, in this paper we propose a novel acceleration framework, namely token Pruning & Pooling Transformers (PPT), to adaptively tackle these two types of redundancy in different layers. By heuristically integrating both token pruning and token pooling techniques in ViTs without additional trainable parameters, PPT effectively reduces the model complexity while maintaining its predictive accuracy. For example, PPT reduces over 37% FLOPs and improves the throughput by over 45% for DeiT-S without any accuracy drop on the ImageNet dataset. The code is available at https://github.com/xjwu1024/PPT and https://github.com/mindspore-lab/models/

연구 동기 및 목표

• 모든 토큰 간의 주의 메커니즘에 기인한 계산 비효율성 문제를 해결한다.
• 기존 방법이 주로 비주의적 또는 중복적인 토큰 중 하나의 유형에만 초점을 맞추는 데서 비롯되는 한계를 극복하기 위해, 둘 다 동시에 다루는 것을 목표로 한다.
• 레이어와 입력 인스턴스 간에 토큰 프루닝과 풀링을 적응적으로 통합할 수 있는 경량이며 가중치 없는 메커니즘을 개발한다.
• 세부 조정이나 아키텍처의 대대적 개선 없이도 비전 트랜스포머의 효율적인 추론과 더 빠른 학습을 가능하게 한다.
• 최신 토큰 압축 기법들과 비교해 더 나은 정확도-FLOPs 트레이드오프를 달성한다.

제안 방법

• 각 레이어의 토큰 중요도 점수 분포에 기반해 프루닝 또는 풀링을 선택하는 적응형 토큰 압축 모듈을 제안한다.
• 토큰 점수의 분산을 기반으로 한 히우리스틱 기반의 결정 기준을 사용하여, 높은 분산(부분적으로만 중요한 토큰)인 경우는 프루닝, 낮은 분산(유사한 토큰)인 경우는 풀링을 적용한다.
• 학습 가능한 가중치를 추가하지 않아도 되도록 표준 트랜스포머 블록에 모듈을 통합함으로써 사전 학습된 모델과의 호환성을 확보한다.
• 레이어 간에 다른 전략을 적용: 얕은 레이어에서는 유사성이 높기 때문에 풀링을 우선시하고, 깊은 레이어에서는 중요도가 뚜렷하기 때문에 프루닝을 우선시한다.
• 입력 인스턴스마다 국소적인 토큰 점수 분포에 따라 정책을 동적으로 조정함으로써 인스턴스 인식 기반의 압축을 가능하게 한다.
• 프루닝과 풀링 간의 선호도를 제어하기 위해 학습 가능한 임계값 τ를 설정하며, 이는 세부 조정을 통해 최적화하거나 사전 설정된 설정에서 경험적으로 결정한다.

실험 결과

연구 질문

• RQ1 비전 트랜스포머에서 토큰 프루닝과 풀링을 동시에 적용할 경우, 각각을 별도로 사용할 때보다 더 나은 효율성-정확도 트레이드오프를 달성할 수 있는가?
• RQ2 레이어 깊이에 따라 토큰 중복성 패턴이 다르게 나타나는가? 즉, 얕은 레이어에서는 풀링, 깊은 레이어에서는 프루닝이 더 적합한가?
• RQ3 학습이 불필요한 인스턴스 인식 기반 정책 선택 메커니즘이 다양한 입력에서 비주의적 및 중복적인 중복성을 효과적으로 균형 잡을 수 있는가?
• RQ4 결정 임계값 τ의 선택이 적응형 압축 프레임워크의 성능과 내구성에 어떤 영향을 미치는가?
• RQ5 제안된 방법은 세부 조정 없이도 표준 벤치마크인 ImageNet에서 정확도 저하 없이 구현 가능한가?

주요 결과

• PPT는 DeiT-S에서 이미지 분류 성능에 영향을 주지 않고 FLOPs를 37% 이상 감소시키고 처리량을 45% 이상 향상시킨다.
• 고정된 규칙 기반 정책(예: 얕은 레이어에서 풀링, 깊은 레이어에서 프루닝)이나 무작위 선택보다 적응형 선택 전략이 더 뛰어난 성능을 보이며, 특히 낮은 FLOP 수준에서 두드러진다.
• 제안된 프레임워크에서 토큰 풀링만으로도 토큰 프루닝만 사용하는 것보다 더 우수한 성능을 보였지만, 적응형 선택을 통한 두 기법의 조합이 전체적으로 가장 뛰어난 성능을 달성한다.
• 세부 조정을 통해 최적의 결정 임계값 τ는 6×10⁻⁵로 도출되었고, 사전 설정 설정에서는 7×10⁻⁵로 확인되었다. 이는 프루닝과 풀링 간의 선호도를 균형 있게 조절한다.
• 이 방법은 다양한 비전 트랜스포머 아키텍처에서 강력한 일반화 성능을 유지하며, ImageNet에서 최신 기술 대비 뛰어난 효율성-정확도 트레이드오프를 달성한다.
• 이 프레임워크는 추론 가속화에 효과적이며, 추가적인 가중치 없이도 정확도 저하 최소화로 더 빠른 학습을 가능하게 한다.