[논문 리뷰] Skip-Attention: Improving Vision Transformers by Paying Less Attention

Shashanka Venkataramanan, Amir Hossein Ghodrati외 3명|arXiv (Cornell University)|2023. 01. 05.

CCD and CMOS Imaging Sensors인용 수 6

한 줄 요약

Skip-Attention (SkipAt)은 더 깊은 레이어에서 비용이 많이 드는 어텐션 계산을 줄이기 위해 이전 레이어의 자기어텐션 계산을 재사용함으로써 비전 트랜스포머(Vision Transformers, ViTs)의 성능을 향상시킵니다. 깊이-wise 컨벌루션에 영감을 받은 경량 파rametric 함수를 도입함으로써, 이미지 분류, 세그멘테이션, 노이즈 제거, 비디오 작업 등 다양한 작업에서 기준 ViT 대비 높은 처리량과 향상된 정확도를 달성하며, ADE20K에서 최대 40%의 속도 향상과 26% 빠른 자기지도 학습 전훈을 기록합니다.

ABSTRACT

This work aims to improve the efficiency of vision transformers (ViT). While ViTs use computationally expensive self-attention operations in every layer, we identify that these operations are highly correlated across layers -- a key redundancy that causes unnecessary computations. Based on this observation, we propose SkipAt, a method to reuse self-attention computation from preceding layers to approximate attention at one or more subsequent layers. To ensure that reusing self-attention blocks across layers does not degrade the performance, we introduce a simple parametric function, which outperforms the baseline transformer's performance while running computationally faster. We show the effectiveness of our method in image classification and self-supervised learning on ImageNet-1K, semantic segmentation on ADE20K, image denoising on SIDD, and video denoising on DAVIS. We achieve improved throughput at the same-or-higher accuracy levels in all these tasks.

연구 동기 및 목표

• 비전 트랜스포머(ViTs)에서 시퀀스 길이에 따라 제곱적으로 증가하는 자기어텐션 계산의 높은 계산 비용을 해결합니다.
• 클래스 토큰을 중심으로 다중헤드 자기어텐션(MSA) 출력이 레이어 간에 높은 상관관계를 보이며, 이는 중복성을 시사합니다.
• 더 깊은 레이어에서 비용이 많이 드는 MSA 계산을 이전 레이어의 표현을 재사용함으로써 대체할 수 있는 플러그인 모듈을 개발합니다.
• 성능을 유지하면서 효율성을 향상시키기 위해 성능을 보존하는 경량 파arametric 함수를 설계하여 더 빠른 추론과 학습을 가능하게 합니다.
• 이식성 있는 성능을 입증하기 위해 이미지 분류, 세그멘테이션, 이미지 및 비디오 노이즈 제거, 자기지도 학습 등 다양한 비전 작업에 대해 일반화 능력을 입증합니다.

제안 방법

• 더 깊은 레이어에서 MSA 블록을 이전 레이어(예: 레이어 3–8)의 어텐션 특징을 재사용하는 방식으로 대체하는 플러그인 모듈인 SkipAt을 제안합니다.
• 재사용된 특징을 대상 레이어의 표현 공간에 적합시키기 위해 학습 가능한 파arametric 함수 Φ를 도입합니다.
• 표현력과 효율성의 균형을 고려해, 기본적으로 커널 크기 5×5와 채널 확장 비율 2를 가진 깊이-wise 컨벌루션(DwC)을 Φ 함수로 사용합니다.
• 분류 및 조밀한 예측 작업 모두에 스킵 메커니즘을 적용하여 이미지 수준 및 픽셀 수준 출력을 지원합니다.
• 기본 ViT와 동일한 데이터와 학습 프로토콜을 사용하여 표준 목표 함수로 엔드 투 엔드로 모델을 훈련합니다.
• 모든 ViT 아키텍처(소형, 베이스, 라지 버전 포함)에 적용 가능하도록 설계하여 다양한 데이터셋과 작업에 적용 가능합니다.

Figure 2 : Attention correlation . Mean of the attention heads from the CLS token of a pretrained ViT-T/16 at different layers from the validation set of ImageNet-1K. Numbers below each attention map indicates the cosine similarity of $A_{l}^{\texttt{[CLS]}}$ with $A_{l-1}^{\texttt{[CLS]}}$ .

실험 결과

연구 질문

• RQ1 성능 저하 없이 더 깊은 ViT 레이어의 자기어텐션 계산을 이전 레이어의 특징 재사용으로 얼마나 잘 근사할 수 있는가?
• RQ2 간단한 파arametric 함수가 직접 재사용 또는 기본 ViT 대비 성능 향상에 효과적으로 기여할 수 있는가?
• RQ3 SkipAt은 조밀한 예측 및 비디오 이해를 포함한 다양한 비전 작업에서 추론 속도와 FLOPs에 어떤 영향을 미치는가?
• RQ4 SkipAt은 자기지도 학습 전훈에서 성능을 유지하거나 향상시키며, 정확도 손실 없이 전훈 시간을 단축시킬 수 있는가?
• RQ5 정확도와 처리량의 균형을 고려할 때 파arametric 함수의 최적 설정(커널 크기, 채널 확장 비율)은 무엇인가?

주요 결과

• ADE20K 세그멘테이션 벤치마크에서 SkipAt은 기준 ViT-T/16 대비 0.6% 높은 정확도를 기록하며 추론 속도 40% 향상을 달성합니다.
• ImageNet-1K에서 SkipAt은 67.7%의 top-1 정확도를 기록하며 초당 6.9K 장(6.9×10³ img/sec)의 처리량을 확보하여, 기준 ViT 대비 정확도 1.9% 향상과 처리량 19% 향상을 기록합니다.
• DAVIS에서 비디오 노이즈 제거 작업에서 SkipAt은 최고 성능(SOTA) 방법과 동일한 PSNR(35.16)를 기록했으며, FLOPs는 17% 줄여(77.1 GFLOPs 대비 93.2 GFLOPs) 효율성을 확보합니다.
• ImageNet-1K에서 자기지도 학습 전훈 시간을 26% 단축시키며, 다운스트림 정확도에 손실 없이 표현 학습의 효율성을 입증합니다.
• 최적의 SkipAt 설정은 커널 크기 5×5의 깊이-wise 컨벌루션과 채널 확장 비율 2를 사용하며, ImageNet-1K에서 67.7%의 top-1 정확도와 6.9K img/sec의 처리량을 달성합니다.
• 제거 분석 결과, 특징 재사용에 항등 함수를 사용할 경우 정확도가 4.7% 감소하지만, SkipAt의 파arametric 함수는 정확도와 효율성 모두에서 모든 기준 대비 일관되게 뛰어난 성능을 보입니다.

Figure 3 : CKA analysis of $A^{[\texttt{CLS}]}$ and $Z^{\text{MSA}}$ across different layers of pretrained ViT-T/16 on the validation set of Imagenet-1K. Vanilla ViT-T/16 has high correlation across both attention maps (layer 3 to 10) and $Z^{\text{MSA}}$ (layer 2 to 8)