[논문 리뷰] Scalable Vision Transformers with Hierarchical Pooling

Zizheng Pan, Bohan Zhuang외 3명|arXiv (Cornell University)|2021. 03. 19.

Image Processing Techniques and Applications인용 수 12

한 줄 요약

이 논문은 계층적 풀링을 통해 시퀀스 길이와 계산 비용을 점진적으로 감소시켜 확장 가능한 비전 트랜스포머 아키텍처인 계층적 시각 트랜스포머(HVT)를 제안한다. 클래스 토큰을 패치 토큰에 대한 평균 풀링으로 대체함으로써 HVT는 계산 비용이 유사한 조건에서 ImageNet에서 DeiT보다 3.03% 높은 정확도를 달성하며, 계산 비용 증가 없이 깊이, 너비, 해상도, 패치 크기의 효율적 확장을 가능하게 한다.

ABSTRACT

The recently proposed Visual image Transformers (ViT) with pure attention have achieved promising performance on image recognition tasks, such as image classification. However, the routine of the current ViT model is to maintain a full-length patch sequence during inference, which is redundant and lacks hierarchical representation. To this end, we propose a Hierarchical Visual Transformer (HVT) which progressively pools visual tokens to shrink the sequence length and hence reduces the computational cost, analogous to the feature maps downsampling in Convolutional Neural Networks (CNNs). It brings a great benefit that we can increase the model capacity by scaling dimensions of depth/width/resolution/patch size without introducing extra computational complexity due to the reduced sequence length. Moreover, we empirically find that the average pooled visual tokens contain more discriminative information than the single class token. To demonstrate the improved scalability of our HVT, we conduct extensive experiments on the image classification task. With comparable FLOPs, our HVT outperforms the competitive baselines on ImageNet and CIFAR-100 datasets. Code is available at https://github.com/MonashAI/HVT

연구 동기 및 목표

• 표준 비전 트랜스포머(ViT)가 추론 전반에 걸쳐 전체 길이의 패치 시퀀스를 유지함으로써 높은 계산 비용과 계층적 표현의 부재 문제를 해결한다.
• 진행적 풀링을 통해 시퀀스 길이를 줄임으로써 계산 비용을 감소시키고, 모델의 깊이, 너비, 해상도, 패치 크기의 효율적 확장을 가능하게 한다.
• 평균 풀링을 통해 생성된 패치 토큰의 평균값이 단일 학습 가능한 클래스 토큰보다 더 구분력 있는 정보를 담고 있음을 입증함으로써, 분류에 있어 후자의 의존성을 제거할 수 있다.
• 계산 비용 증가 없이 절약된 FLOPs를 활용해 모델 용량을 증가시킴으로써 이미지 분류 벤치마크에서 최신 성능을 달성한다.

제안 방법

• 다중 스테이지에서 시퀀스 길이를 감소시키는 계층적 풀링 메커니즘을 도입하여, CNN의 공간적 다운샘플링과 유사하게 작동한다.
• ViT 블록을 스테이지로 분할하고 각 스테이지 후에 평균 풀링 또는 최대 풀링을 적용하여 시퀀스 길이와 계산 부담을 점진적으로 감소시킨다.
• 학습 가능한 클래스 토큰을 모든 패치 토큰에 대한 전역 평균 풀링으로 대체하여 최종 예측 벡터를 생성한다.
• 깊이, 너비, 입력 해상도, 패치 크기를 별도로 확장할 수 있도록 아키텍처를 설계하여 절약된 FLOPs를 활용해 용량을 향상시킨다.
• 일부 실험에서는 패치 크기를 줄인 공간적 다운샘플링을 위해 2D 풀링 변형을 사용하며, 이에 따라 헤드 및 블록 스케줄링을 신중히 조정한다.
• 다양한 수의 어텐션 헤드, 입력 해상도, 패치 크기를 가진 모델을 훈련시켜 확장성과 성능 간의 상충 관계를 평가한다.

Figure 1 : Performance comparisons on ImageNet. With comparable GFLOPs (1.25 vs. 1.39), our proposed Scale HVT-Ti-4 surpasses DeiT-Ti by 3.03% in Top-1 accuracy.

실험 결과

연구 질문

• RQ1 비전 트랜스포머에서 계층적 풀링을 적용하면 시퀀스 길이와 계산 비용을 줄일 수 있을까, 동시에 정확도는 유지되거나 향상될 수 있을까?
• RQ2 클래스 토큰을 패치 토큰에 대한 평균 풀링으로 대체하면 전용 클래스 토큰을 사용할 때보다 더 나은 분류 성능을 달성할 수 있을까?
• RQ3 계층적 풀링을 통해 FLOPs가 유사한 조건에서 모델 용량을 깊이, 너비, 해상도, 패치 크기 측면에서 얼마나 효과적으로 확장할 수 있을까?
• RQ4 ImageNet과 CIFAR-100에서 DeiT와 같은 강력한 베이스라인과 유사한 FLOP 제약 조건 하에서 HVT의 성능은 어떻게 비교될 수 있을까?
• RQ5 2D 풀링은 모델 성능에 어떤 영향을 미치며, 적절한 아키텍처 조정을 통해 추가적인 향상을 이끌 수 있을까?

주요 결과

• FLOPs가 유사한 조건(1.25 vs. 1.39 GFLOPs)에서 HVT-Ti-4는 ImageNet에서 DeiT-Ti보다 3.03% 높은 정확도를 기록한다.
• CIFAR-100에서 패치 크기가 8인 HVT-S-4는 77.29%의 정확도를 달성하여 32패치 버전 대비 9.14% 향상된 성능을 보였다.
• 입력 해상도를 160에서 384로 늘리면 정확도가 73.84%에서 76.31%로 상승하여 고해상도 표현의 이점이 입증되었다.
• 16개의 어텐션 헤드를 가진 HVT-S-4는 3헤드 버전보다 6.79%p 높은 75.43%의 정확도를 기록했다.
• 2D 풀링을 적용한 HVT-S-2는 CIFAR-100에서 77.58%의 정확도를 기록하여 FLOPs가 거의 동일한 DeiT-S(71.99%)를 능가했다.
• 평균 풀링을 통해 생성된 시각적 토큰은 항상 클래스 토큰보다 더 높은 성능을 보였으며, 이는 그들이 더 풍부한 구분 능력을 지닌다는 것을 시사한다.

Figure 2 : Overview of the proposed Hierarchical Visual Transformer. To reduce the redundancy in the full-length patch sequence and construct a hierarchical representation, we propose to progressively pool visual tokens to shrink the sequence length. To this end, we partition the ViT [ 11 ] blocks i