SPViT: 소프트 토큰 프루닝을 통한 더 빠른 비전 트랜스포머 실현

최근, 비전 트랜스포머(Vision Transformer, ViT)는 컴퓨터 비전 분야에서 지속적으로 새로운 기준을 세우고 있지만, 높은 계산 및 메모리 비용으로 인해 산업 생산에서의 확산이 어려웠습니다. 하드웨어 효율성을 위한 전통적인 모델 압축 패러다임인 프루닝(pruning)은 다양한 DNN 구조에 널리 적용되어 왔습니다. 그럼에도 불구하고, ViT 구조에 대한 독립적인 프루닝 방법론은 여전히 불분명합니다. 이에 우리는 구조적 특성, ViT 내부 데이터 패턴, 그리고 관련 엣지 디바이스 배포를 고려하여 입력 토큰의 희소성을 활용하고 계산을 고려한 소프트 프루닝(soft pruning) 프레임워크를 제안합니다. 이 프레임워크는 플래튼(flatten) 및 CNN 유형 구조 모두에 적용할 수 있는 기본 트랜스포머(vanilla Transformers) 위에 설정될 수 있으며, 예를 들어 풀링 기반 ViT(Pooling-based ViT, PiT)와 같은 구조에서도 사용 가능합니다.구체적으로, 우리는 적응형 인스턴스별 토큰 선택을 위한 경량 모듈인 동적 주의 기반 다중 헤드 토큰 선택기(dynamic attention-based multi-head token selector)를 설계하였습니다. 또한, 선택기 모듈이 생성한 정보량이 적은 토큰들을 완전히 버리지 않고 후속 계산에 참여하는 패키지 토큰(package token)으로 통합하는 소프트 프루닝 기술을 소개하였습니다. 우리의 프레임워크는 제안된 계산을 고려한 학습 전략(computation-aware training strategy)을 통해 특정 엣지 디바이스의 정확도와 계산 제약 조건 사이의 균형을 맞춥니다.실험 결과는 우리의 프레임워크가 이미지 분류에서 유사한 성능을 유지하면서 ViT의 계산 비용을 크게 줄이는 것을 보여줍니다. 더욱이, 우리의 프레임워크는 모바일 디바이스와 FPGA의 자원 사양을 충족시키면서 실시간 실행까지 보장할 수 있습니다. 예를 들어, ImageNet 데이터셋에서 DeiT-T 모델의 지연 시간(latency)을 26ms로 줄여 기존 연구보다 26%~41% 우수한 성능을 달성하였으며, 최상위 1등급(top-1 accuracy) 정확도는 0.25%~4% 더 높아졌습니다.