[논문 리뷰] UFO-ViT: High Performance Linear Vision Transformer without Softmax

Jeong-geun Song|arXiv (Cornell University)|2021. 09. 29.

CCD and CMOS Imaging Sensors참고 문헌 46인용 수 10

한 줄 요약

UFO-ViT는 이미지 분류 및 조밀 예측 작업에서 최고 성능을 달성하면서 계산 비용과 메모리 사용을 줄이며 선형 복잡도를 가지는 비전 트랜스포머를 제안한다. 이는 자기주의 어텐션에서 소프트맥스 비선형성을 L2 정규화 기반의 단위 힘 연산으로 대체함으로써 달성된다. 이 방법은 O(N) 복잡도를 유지하면서 다양한 모델 크기에서 높은 정확도를 유지하며, ImageNet1k 및 COCO 벤치마크에서 기존 트랜스포머와 CNN보다 뛰어난 성능을 보인다.

ABSTRACT

Vision transformers have become one of the most important models for computer vision tasks. Although they outperform prior works, they require heavy computational resources on a scale that is quadratic to $N$. This is a major drawback of the traditional self-attention (SA) algorithm. Here, we propose the Unit Force Operated Vision Transformer (UFO-ViT), a novel SA mechanism that has linear complexity. The main approach of this work is to eliminate nonlinearity from the original SA. We factorize the matrix multiplication of the SA mechanism without complicated linear approximation. By modifying only a few lines of code from the original SA, the proposed models outperform most transformer-based models on image classification and dense prediction tasks on most capacity regimes.

연구 동기 및 목표

• 고해상도 입력에서 표준 자기주의 어텐션의 O(N²) 계산 복잡도 문제를 해결하기 위해.
• 성능 저하 없이 자기주의 어텐션의 소프트맥스 비선형성을 제거하여 선형 시간 어텐션 메커니즘을 가능하게 하기 위해.
• 이미지 분류 및 조밀 예측과 같은 다양한 비전 작업에서 높은 정확도를 유지하면서 FLOPs와 메모리 사용을 크게 줄이는 모델을 개발하기 위해.
• 특히 객체 검출 및 세그멘테이션과 같은 조밀 예측 작업에서 해상도가 다양할 경우에도 효율적인 배포를 가능하게 하기 위해.

제안 방법

• 자기주의 어텐션의 소프트맥스 연산을 L2 정규화로 대체하여 행렬 지수화의 제곱 복잡도를 제거하기 위해.
• 행렬 곱셈의 결합 법칙을 적용하여 QK^T를 단위 벡터로 재구성함으로써 어텐션 계산을 O(N) 연산으로 분해하기 위해.
• 쿼리와 키를 단위 초구면에 투영하는 XNorm이라는 정규화 체계를 도입하여 학습을 안정화시키고 초기화에 대한 의존성을 제거하기 위해.
• ResMLP에서 영감을 얻은 애프린 스케일링 모듈을 사용하여 잔차 연결을 조정함으로써 다양한 모델 깊이에서 안정성을 유지하기 위해.
• 해상도에 영향을 받지 않는 아키텍처를 설계하여 모델 가중치가 입력 해상도에 의존하지 않도록 하여 다양한 입력 크기에서 일관된 성능을 확보하기 위해.
• 표준 ViT에서 몇 줄의 코드만 수정하여 구현함으로써 후행 호환성과 통합 용이성을 확보하기 위해.

실험 결과

연구 질문

• RQ1 자기주의 어텐션에서 성능 저하 없이 선형 복잡도를 달성할 수 있는 비전 트랜스포머가 이미지 분류 작업에서 성능을 유지를 할 수 있는가?
• RQ2 자기주의 어텐션에서 소프트맥스를 L2 정규화로 대체할 경우, 기존 트랜스포머 및 이전의 선형 근사 방법 대비 정확도가 유지되거나 향상되는가?
• RQ3 제안된 방법이 고해상도 입력, 특히 객체 검출 및 인스턴스 세그멘테이션과 같은 조밀 예측 작업에서 효율적으로 확장 가능한가?
• RQ4 다양한 입력 해상도와 모델 용량에서 기존 트랜스포머 및 CNN과 비교해 모델의 메모리 소비와 추론 속도는 어떻게 되는가?

주요 결과

• UFO-ViT-S는 단지 39.7M 파라미터로 COCO 인스턴스 세그멘테이션에서 44.6 mAP를 기록하며, ResNet50 및 PVT-Small를 뛰어넘는 성능을 보였다.
• UFO-ViT-M는 56.4M 파라미터로 COCO에서 46.0 mAP를 기록하며, Swin-T와 XCiT-S12/16를 능가했고, Swin-S의 FLOPs의 절반 미만으로 작동했다.
• 모델은 입력 해상도에 따라 선형적으로 메모리 증가를 유지한다: GPU 메모리 사용량은 토큰 수에 비례하여 선형적으로 증가하며, 표준 트랜스포머의 제곱 증가와는 다름.
• UFO-ViT는 DeiT 및 Swin 트랜스포머보다 단일 V100 GPU에서 최대 4배 더 큰 배치 크기를 지원하여 뛰어난 메모리 효율성을 보였다.
• 추론 스루풋은 DeiT보다 뚜렷하게 높았고, 특히 고해상도에서 Swin 트랜스포머와 유사하거나 더 우수한 성능을 보였다.
• UFO-ViT-B는 경계 상자 검출에서는 UFO-ViT-M보다 略로 낮은 성능을 보였지만, 소형 객체 검출 및 인스턴스 세그멘테이션에서는 더 뛰어난 성능을 기록하여 세밀한 작업에 더 나은 특징 표현을 제공함을 시사했다.