[논문 리뷰] Discrete Representations Strengthen Vision Transformer Robustness

Chengzhi Mao, Lu Jiang외 4명|arXiv (Cornell University)|2021. 11. 20.

Advanced Neural Network Applications참고 문헌 41인용 수 10

한 줄 요약

이 논문은 표준 연속 픽셀 토큰을 벡터 양자화된 VQ-GAN 인코더에서 유래한 이산 토큰으로 대체함으로써 강건성을 향상시키는 Dr. ViT라는 비전 트랜스포머 변종을 제안한다. 이산 토큰과 잔여 픽셀 임베딩을 조합함으로써 모델은 더 더이상의 불변성과 형태 인식 능력을 갖춘 표현을 학습하게 되어, ImageNet 정확도를 훼손하지 않은 채 일곱 개인 ImageNet 강건성 벤치마크에서 최대 12% 향상된 성능을 달성한다.

ABSTRACT

Vision Transformer (ViT) is emerging as the state-of-the-art architecture for image recognition. While recent studies suggest that ViTs are more robust than their convolutional counterparts, our experiments find that ViTs trained on ImageNet are overly reliant on local textures and fail to make adequate use of shape information. ViTs thus have difficulties generalizing to out-of-distribution, real-world data. To address this deficiency, we present a simple and effective architecture modification to ViT's input layer by adding discrete tokens produced by a vector-quantized encoder. Different from the standard continuous pixel tokens, discrete tokens are invariant under small perturbations and contain less information individually, which promote ViTs to learn global information that is invariant. Experimental results demonstrate that adding discrete representation on four architecture variants strengthens ViT robustness by up to 12% across seven ImageNet robustness benchmarks while maintaining the performance on ImageNet.

연구 동기 및 목표

• 비전 트랜스포머의 국소 텍스처에 대한 과도한 의존성과 분포 이탈에 대한 낮은 일반화 능력을 해결하기 위해.
• ImageNet에서의 도메인 내 성능을 훼손하지 않으면서 강건성을 향상시키기 위해.
• 이산 토큰화가 비전 트랜스포머가 전역적이고 형태 불변 특징을 학습하는 데 기여할 수 있는지 탐색하기 위해.
• 모든 비전 트랜스포머 변종에서 강건성을 향상시키는 플러그 앤 플레이 입력층 수정 방법을 개발하기 위해.
• 이미 생성에 사용된 바이어스된 이산 표현이 분류에서 일반화 성능을 향상시킬 수 있는지 증명하기 위해.

제안 방법

• 표준 연속 픽셀 임베딩을 VQ-GAN 인코더에서 유도된 이산 토큰으로 대체하여, 이미지 패치를 코드북 인덱스로 양자화한다.
• 학습된 코드북을 사용해 이미지 패치를 이산 시각 단어로 매핑함으로써 구조적 및 형태 정보를 유지한다.
• 이산 토큰 임베딩과 원본 픽셀 임베딩을 연결하여 국소 세부 정보를 유지하며, 특히 작은 객체에 대해 유리하다.
• 표준 최적화(Adam, 코즈인 디케이)와 데이터 증강(RandAug, Mixup)을 사용해 이산 및 연속 표현을 모두 활용해 ViT를 훈련한다.
• 작은 모델에는 경량 VQ-인코더를, ViT-B에는 더 큰 인코더를 사용하며, 코드북 크기는 각각 1,024(ImageNet) 및 8,192(ImageNet-21K)이다.
• ViT-B의 최종 특징에 대해 평균 풀링을 적용하고, 더 높은 해상도 입력(384×384, 512×512)에서 미세조정하여 일반화 성능을 향상시킨다.

실험 결과

연구 질문

• RQ1 이산 토큰화가 비전 트랜스포머의 분포 이탈에 대한 강건성 향상에 기여할 수 있는가?
• RQ2 연속 픽셀 토큰을 이산 표현으로 대체함으로써 비전 트랜스포머가 더 형태 불변 특징을 학습하도록 유도할 수 있는가?
• RQ3 이산 및 연속 토큰의 조합이 도메인 내 및 도메인 외 벤치마크 성능에 어떤 영향을 미치는가?
• RQ4 이 방법은 아키텍처 변경 없이 다양한 비전 트랜스포머 아키텍처에 적용 가능한가?
• RQ5 이산 토큰의 사용이 국소 텍스처가 아닌 전역적 맥락을 반영하는 더 나은 어텐션 패턴을 유도하는가?

주요 결과

• Dr. ViT는 Stylized-ImageNet에서 표준 ViT 대비 최대 12% 향상된 강건성을 보이며, ImageNet-Sketch 및 ImageNet-C에서는 최대 10% 향상된다.
• 이 방법은 데이터세트 특화 데이터 증강을 사용하지 않아도 일곱 개인 ImageNet 강건성 벤치마크 중 네 개에서 최신 기술(SOTA) 성능을 달성한다.
• 이산 토큰을 추가함으로써 어텐션 패턴이 향상되어, 모델이 국소 텍스처보다 전역적 구조에 더 민감해졌음을 시각화 비교를 통해 입증했다.
• 이산 토큰 전용 버전은 작은 모델(ViT-Ti)에서는 성능이 열악하지만, 픽셀 임베딩과 조합하면 성능 향상이 뚜렷하다.
• ViT-B를 300 에포크 대신 800 에포크 동안 훈련하면 ImageNet에서 추가로 0.1~1%의 정확도 향상이 나타나, 이산 토큰과 함께 장기간 훈련이 유리함을 시사한다.
• 이 접근법은 기존의 강건성 기법들과 수직적이고, 최소한의 수정으로 어떤 ViT 기반 아키텍처에나 통합할 수 있다.