HATS : head-aware token selection for efficient vision transformers

다국어 초록 (Multilingual Abstract)

Given the high computational cost of ViTs, token compression has emerged as an effective technique to eliminate low-importance tokens. A common strategy for estimating token importance is to use the [CLS] attention scores averaged across all heads in each block. Notably, prior studies have shown that the [CLS] token in early blocks often attends to arbitrary regions and local features, whereas in the final block, it focuses on task-relevant tokens. This discrepancy lead to the unintended removal of task-relevant tokens during the compression stages in early blocks. Motivated by this issue, we hypothesized that the head average token scoring approach obscures task-relevant signals. We empirically discover that certain heads in early blocks show high similarity to the final block's attention, capturing key task-relevant representations. Based on this insight, we propose a simple and effective method for token compression called Head-Aware Token Selection (HATS). HATS selects heads within each block whose [CLS] attention distributions are similar to that of the final block and then groups these heads to form a combined [CLS] attention map. Consequently, token importance estimates become more stable, and tokens required for the final prediction are better preserved during compression, without requiring any additional training or architectural modifications. Various experiments demonstrate that HATS achieves more reliable estimation of token importance, mitigates accuracy degradation effectively, and can be applied in an off-the-shelf manner to existing token compression frameworks.

국문 초록 (Abstract)

Vision Transformer(ViT)는 뛰어난 성능에도 불구하고 토큰 수에 따라 연산량이 제곱으로 증가하는 높은 비용 문제가 있어 리소스가 제한된 장치에 배포하기 어렵다. 이러한 한계를 완화하기 위해 토큰 압축 기술이 부상했지만, 토큰의 중요도를 평가하는 방식에 한계가 존재한다. 기존의 일반적인 방법은 각 블록의 모든 어텐션 헤드에서 [CLS] 어텐션 맵을 단순 평균(Head-Wise Average)하여 중요도를 추정한다. 하지만 ViT는 최종 예측을 마지막 블록의 [CLS] 토큰에 의존하며 초기 블록의 [CLS] 어텐션은 마지막 블록과 달리 태스크와 무관한 임의의 영역에도 집중하는 경향이 있다. 이러한 불일치는 초기 압축 단계에서 태스크 수행에 필수적인 토큰이 조기에 제거되는(premature removal) 문제를 야기하여 전반적인 성능 저하를 초래한다. 또한 ViT의 초반 블록에서 각 헤드는 각각 과업에 관한 중요성이 다를 수 있다는 사전 연구가 존재한다. 따라서 단순히 모든 헤드를 평균 내어 중요도 추정에 사용하는 기존 방법론은 각각의 헤드가 가진 과업 중요도를 왜곡시킬 수 있다.
본 연구는 이러한 문제의식에서 출발하여 초기 블록의 헤드 중에도 마지막 블록의 [CLS] 어텐션 패턴과 유사하게 태스크 관련 정보를 포착하는 헤드들이 존재한다는 새로운 관찰에 주목한다. 이 통찰을 바탕으로 본 연구는 Head-Aware Token Selection (HATS)라는 간단하고 효과적인 토큰 압축 전략을 제안한다. HATS는 먼저 코사인 유사도를 사용해 마지막 블록의 [CLS] 어텐션 맵과 가장 유사한 패턴을 보이는 초기 헤드들을 식별(Head Selection)한다. 그 다음 이 선택된 헤드들의 조합을 구성하고 그 평균 맵을 마지막 블록과 다시 비교하여 가장 높은 유사도를 보이는 최적의 헤드 조합을 선택(Head Grouping)한다. 이 2단계 전략을 통해 HATS는 추가적인 학습이나 구조 변경 없이도 초기 블록에서 더 정확하게 토큰 중요도를 추정하고 태스크 관련 토큰이 압축 과정에서 효과적으로 보존되도록 한다. 실험을 통해 HATS는 기존의 다양한 토큰 압축 프레임워크에 즉시 적용(off-the-shelf) 가능하며 정확도 저하를 효과적으로 완화함을 입증한다.

목차 (Table of Contents)

• 1. Introduction 1
• 2. Related work 8
• 3. Head-wise [CLS] Attention Analysis 10
• 3.1. Early block [CLS] Attention Analysis 10
• 3.2. Deeper block [CLS] Attention Analysis 11
• 3.3. Head Combination [CLS] Attention Map 12
• 4. Methods 14
• 4.1. Head-Aware Token Selection 14
• 4.2. Review of [CLS] Attention-based Token Scoring 15
• 4.3. Head Selection 15
• 4.4. Head Grouping 16
• 5. Experiments 19
• 5.1. Implementation Details 19
• 5.2. Off-the-Shelf Integration and Comparison with Existing Token Compression Methods 20
• 5.3. Ablation Studies on HATS 22
• 5.3.1 Random Head Selection 22
• 5.3.2 Impact of Block-Wise Application Scope of HATS 23
• 5.3.3 Choice of Similarity Function 24
• 6. Conclusion and limitations 25
• References 26
• Appendix 33
• Korean Abstract 42