UtVAA: Ultra-tiny Vision Transformer with Affix Attention for Mobile Image Classification

기술 요약: UtVAA: 모바일 이미지 분류를 위한 Affix Attention 기반의 초소형 비전 트랜스포머

문제 정의
비전 트랜스포머(ViT)는 이미지 분류에서 합성곱 신경망(CNN)보다 우수한 표현 능력을 입증해 왔으나, 이차적(quadratic) 자기 주의 집중(self-attention) 복잡도와 방대한 파라미터 수로 인해 자원이 제한된 모바일 및 엣지 기기에 배포하는 데 어려움이 있다. 기존의 경량 아키텍처(예: MobileNet, ShuffleNet)는 종종 효과적인 다중 스케일 특징 추출과 다양한 시나리오에서의 일관된 높은 정확도를 유지하는 데 어려움을 겪는다. 또한, 가지치기(pruning)나 양자화(quantization)와 같은 비-아키텍처적 압축 방법은 특징 표현력을 저하시키는 경우가 많다. 엄격한 계산 예산 내에서 강력한 판별 성능을 유지하면서도 '초소형(ultra-tiny)' 트랜스포머 모델(100만 개 미만의 파라미터)을 체계적으로 설계하는 데에는 결정적인 격차가 존재한다.

방법론
본 논문은 극심한 자원 제약 조건 하에서 전역적 문맥 모델링(global contextual modeling)과 지역적 특징 추출 사이의 균형을 맞추도록 설계된 계층적 아키텍처인 UtVAA(Affix Attention 기반 초소형 비전 트랜스포머)를 제안한다. 이 프레임워크는 세 가지 핵심 구성 요소를 통합한다:

1. Dilated Bottleneck Blocks (확장된 병목 블록): 파라미터 수를 늘리지 않고 수용 영역(receptive field)을 효율적으로 확장하기 위해, 저자들은 두 가지 변형(잔차 연결이 있는 버전과 없는 버전)의 Dilated Bottleneck 블록을 도입했다. 이들은 확장된 depthwise 및 pointwise 컨볼루션을 활용한다. 특히 잔차 변형은 최적화 안정성과 기울기 소실 문제를 해결한다. 이 설계는 표준 컨볼루션에 비해 계산 복잡도를 약 8~9배 줄이면서도 다중 스케일 문맥을 보존한다.
2. Affix Attention Block (Affix Attention 블록): 이는 선형 자기 주의 집중에서 흔히 발생하는 공간 정보 손실을 완화하기 위해 설계된 본 아키텍처의 새로운 핵심이다. 이는 순차적 구조가 아닌 병렬 구조를 채택하며, 다음으로 구성된다:
   • Local Representation Block (지역 표현 블록): 미세한 특징을 포착하기 위해 depthwise separable 컨볼루션을 사용한다.
   • Global Representation Block (전역 표현 블록): 패치 수에 대해 선형 스케일링을 구현하는 선형 트랜스포머 공식을 사용하여, 표준 self-attention의 이차적 비용을 피하면서 장거리 의존성을 모델링한다.
   • Coordinate Attention (좌표 주의 집중): 높이와 너비 축을 따라 각각 전역 풀링을 수행하여 방향성 정보를 포착함으로써 공간적 의존성을 인코딩한다.
   • Ternary Fusion (삼항 융합): 전역 블록, 좌표 주의 집중, 그리고 잔차 연결의 출력을 결직(concatenate)한 후, 1×1 컨볼루션을 통해 지역 및 전역 표현을 통합하는 경량 전략이다.
3. Scalable Variants (확장 가능한 변형들): 아키텍처는 Tiny(204.67K 파라미터), Medium(609.36K), Large(858.13K)의 세 가지 크기로 구현되어, 특정 하드웨어 제약에 따른 유연한 배포를 가능하게 한다.

주요 기여

1. Affix Attention Block: 선형 self-attention과 coordinate attention을 사용하여 전역 및 지역 위치 정보를 공동으로 인코딩하고, 이를 ternary fusion 전략을 통해 통합하는 새로운 어텐션 메커니즘이다.
2. Dilated Bottleneck Blocks: 낮은 FLOPs를 유지하면서 확장된 컨볼루션을 사용하여 수용 영역을 확장하고, 잔차 연결을 통해 최적화 안정성을 유지하는 두 가지 효율적인 블록 설계이다.
3. Ultra-tiny Architecture: 이러한 구성 요소들을 UtVAA로 체계적으로 통합하여, 100만 개 미만의 파라미터로 경쟁력 있는 정확도를 달est하였다.
4. Empirical Validation (경험적 검증): 표준 벤치마크(CIFAR-10, CIFAR-100) 및 실제 농업 데이터셋(PlantVillage-Tomato, SLIF-Tomato)에 대한 종합적인 평가를 수행하였다.

실험 결과
저자들은 UtVAA를 최첨단 경량 CNN(MobileNet-V2, ShuffleNet-V2, SqueezeNet) 및 ViT(MobileViT, DeiT, PiT)와 비교 평가하였다.

• CIFAR 벤치마크: UtVAA-Large는 CIFAR-10(93.34%)과 CIFAR-100(70.97%) 모두에서 가장 높은 Top-1 정확도를 달성하였으며, 비교된 모든 방법론을 능가하는 동시에 대규모 ViT보다 현저히 적은 파라미터 수를 유지했다. UtVAA-Tiny는 단 204.67K의 파라미터와 53.95M FLOPs로 CIFAR-10에서 90.46%를 기록하여, 우수한 정확도-복잡도 트레이드오프를 입증했다.
• 식물 질병 데이터셋: PlantVillage-Tomato 데이터셋에서 UtVAA-Tiny는 더 적은 파라미터로 MobileNet-V2(99.45%)와 MobileViT(99.28%)를 제치고 99.61%의 Top-1 정확도를 달성했다. 더 까다로운 현장 데이터인 SLIF-Tomato에서는 99.66%의 정확도를 기록하며 모든 베이스라인을 앞질렀다.
• 효율성: UtVAA-Tiny는 테스트된 모든 모델(트랜스포머 기반 베이스라인 포함) 중 가장 낮은 추론 지연 시간(PlantVillage에서 8.41ms, SLIF에서 7.45ms)을 보여 실시간 엣지 배포에 적합함을 확인했다.
• Ablation Studies (절제 연구): Coordinate attention 모듈을 제거했을 때 주목할 만한 정확도 저하(CIFAR-100에서 70.97%에서 67.91%로)가 발생하였으며, 이는 공간적 특징 표현에 대한 해당 모듈의 기여를 입증한다.

의의 및 주장
본 논문은 UtVAA가 판별 성능의 상당한 손실 없이 트랜스포머 기반 비전 모델을 체계적으로 초소형 아키텍처로 재설계할 수 있음을 보여준다고 주장한다. 계산 비용이 많이 드는 self-attention을 제안된 Affix Attention으로 대체하고 dilated bottleneck을 활용함으로써, 모델은 전역 문맥 모델링과 계산 효율성 사이의 균형을 달성한다. 저자들은 이러한 접근 방식이 사후적(post-hoc) 압축에 대한 원칙적인 아키텍처 대안을 제공하며, 메모리와 FLOPs가 엄격히 제한되는 모바일 및 엣지 기기에서 고성능 비전 트랜스포머를 실행 가능하게 만든다고 단언한다. 이 연구는 이러한 초소형 설계가 정밀 농업과 같이 자원이 제한된 하드웨어에서의 실시간 추론이 중요한 응용 분야에 특히 효과적임을 시사한다. 향후 연구 과제로 마이크로컨트롤러를 위한 하드웨어 인식 최적화와 객체 탐지 및 시맨틱 세그멘테이션과 같은 밀집 예측 작업으로의 아키텍처 확장을 식별하였다.