Skip-Attention: Improving Vision Transformers by Paying Less Attention

ICLR 2024

 

Introduction

 

기존 ViT에서 인접한 layer 사이의 attention correlation을 비교 해봄

생각보다 비슷한 correlation을 가지고 있는 layer들이 많이 있음

그럼 굳이 모든 layer에서 self attention을 해야될 필요가 있을까?

 

본 논문은 skip attention 컨셉을 사용하여, 기존의 불필요한 computation cost를 줄이는데 집중

또한 여러 transformer 구조에 쉽게 적용 가능한 장점

 

실제로 같은 computation cost를 사용했을 때 여러 task에서 좋은 성능을 보여줌

 

Proposed method

 

Centered Kernel Alignment (CKA) 사용하여 vector 사이의 유사도 측정

class token의 attention map같은 경우 layer 3~10에서는 비슷한 값을 보여줌

심지어 self-attention 결과도 비슷한 양상을 보임

 

본 논문에서는 과감하게 MSA 부분을 제거하고 이를 대체하는 함수로 교체

해당 함수는 convolution 컨셉을 활용, global 보다는 local 영역을 봄

 

다만 ECA를 통해 channel 사이의 attention 고려

특이한 점은 SkipAT function을 통과 하기 전에 class token에 해당하는 정보는 빼고 진행 한 후 최종 결과에 다시 concat 하여 사용

해당 모듈은 correlation이 높은 layer 3~8에서만 적용함

 

실험 결과

 

여러 transformer 계열에 적용 시 성능 향상을 보임

 

(a) 마지막 layer 4개에 대해서 cls token attention map을 보면 기존 ViT 보다 object에 대해 attention이 더 잘됨

(b) 실제 CKA를 보면 기존보다 layer 사이의 correlation이 많이 사라짐 → representation 성능이 더 좋아짐