[논문 리뷰] Compressing Vision Transformers for Low-Resource Visual Learning

Eric Youn, Sai Mitheran J.외 2명|arXiv (Cornell University)|2023. 09. 05.

Advanced Neural Network ApplicationsComputer Science인용 수 3

한 줄 요약

이 논문은 NVIDIA Jetson Nano(4GB RAM)와 같은 저자원 엣지 디바이스에 효율적인 배포를 가능하게 하기 위해 구조적 프루닝, 지식 증류, 그리고 트레이닝 후 양자화를 조합한 하이브리드 압축 프레임워크를 제안한다. 이 방법은 최소한의 지연과 메모리 사용량으로 스위니 트랜스포머 수준의 정확도를 달성하여 UAV 기반 재난 모니터링과 같은 자원 제약 상황에서 실시간 세그멘테이션을 가능하게 한다.

ABSTRACT

Vision transformer (ViT) and its variants have swept through visual learning leaderboards and offer state-of-the-art accuracy in tasks such as image classification, object detection, and semantic segmentation by attending to different parts of the visual input and capturing long-range spatial dependencies. However, these models are large and computation-heavy. For instance, the recently proposed ViT-B model has 86M parameters making it impractical for deployment on resource-constrained devices. As a result, their deployment on mobile and edge scenarios is limited. In our work, we aim to take a step toward bringing vision transformers to the edge by utilizing popular model compression techniques such as distillation, pruning, and quantization. Our chosen application environment is an unmanned aerial vehicle (UAV) that is battery-powered and memory-constrained, carrying a single-board computer on the scale of an NVIDIA Jetson Nano with 4GB of RAM. On the other hand, the UAV requires high accuracy close to that of state-of-the-art ViTs to ensure safe object avoidance in autonomous navigation, or correct localization of humans in search-and-rescue. Inference latency should also be minimized given the application requirements. Hence, our target is to enable rapid inference of a vision transformer on an NVIDIA Jetson Nano (4GB) with minimal accuracy loss. This allows us to deploy ViTs on resource-constrained devices, opening up new possibilities in surveillance, environmental monitoring, etc. Our implementation is made available at https://github.com/chensy7/efficient-vit.

연구 동기 및 목표

• NVIDIA Jetson Nano(4GB RAM)를 탑재한 UAV와 같은 메모리 제약이 있고 배터리로 구동되는 엣지 디바이스에 비전 트랜스포머를 배포할 수 있도록 하는 것.
• 모델 크기와 추론 지연을 최소화하면서도 스위니 트랜스포머 수준의 높은 세그멘테이션 정확도를 유지하는 것.
• 엣지 배포를 위한 구조적 프루닝, 증류, 그리고 트레이닝 후 양자화를 조합한 실용적인 종단 간 압축 파이프라인을 개발하는 것.
• LPCV 데이터셋을 사용하여 엄격한 메모리 및 지연 제약 조건 하에서 실제 재난 현장 해석을 평가하는 것.

제안 방법

• 계산 그래프의 구조를 변경하지 않으면서도 여분의 어텐션 헤드와 선형 레이어의 뉴런을 제거하여 파라미터 수를 줄이는 구조적 프루닝을 적용한다.
• 다단계 증류 전략을 사용: 예측 헤드에서 출력 로짓 지식 전달, 초기 레이어에서 특징 모방, 더 깊은 레이어에서 생성 기반 증류.
• 모델 크기와 메모리 프로필을 줄이기 위해 트레이닝 후 양자화(PTQ)를 적용하지만, 파이토치에서 양자화된 CUDA 연산에 대한 지원이 제한되어 있어 추가 성능 향상에 제약이 있다.
• 프루닝 이후 반복적 피니팅을 통해 증류를 조합하여 성능 안정성과 일반화 능력을 향상시킨다.
• 더 강력한 스위니-v2-T 모델을 교사로 사용하여 지식 전달을 위해 더 작은 모바일비트 학생 네트워크를 지도한다.
• 주요 평가 지표로 평균 교차율(mIoU)을 사용하여 LPCV 데이터셋에서 압축된 모델을 평가한다.

Figure 1: Our framework for edge-oriented, rapid on-device disaster scene parsing

실험 결과

연구 질문

• RQ1 어떤 어텐션 헤드와 선형 레이어 가중치를 구조적으로 프루닝하는 것이 엣지 디바이스에서 정확도 저하 없이 모델 크기를 효과적으로 줄일 수 있는가?
• RQ2 출력 로짓, 특징, 생성 기반의 다단계 감독을 통한 지식 증류는 대규모 교사 모델에서 작고 컴act한 학생 ViT로 성능을 효과적으로 전달할 수 있는가?
• RQ3 Jetson Nano에서 메모리 사용량과 지연을 줄이기 위해 트레이닝 후 양자화를 얼마나 효과적으로 적용할 수 있으며, 세그멘테이션 정확도를 유지할 수 있는가?
• RQ4 클래스 불균형과 세밀한 시각적 차이(예: 빙하 유출과 흙길)는 실제 재난 현장에서 압축된 모델의 성능에 어떤 영향을 미치는가?
• RQ5 일괄 프루닝 대비 반복적 프루닝과 피니팅을 통해 압축 효율성이 향상되는가?

주요 결과

• 최종적으로 압축된 모델은 시스템 RAM이 4GB 뿐인 LPCV 데이터셋에서 스위니 트랜스포머 수준의 정확도를 달성하여 엣지 배포 조건을 충족시켰다.
• 어텐션 헤드를 4개에서 2개로 줄이는 프루닝은 성능에 거의 영향을 주지 않으면서도 파라미터 수를 크게 줄였다.
• 트레이닝 후 양자화는 메모리 사용량을 크게 줄였지만, 파이토치의 양자화된 CUDA 연산에 대한 제한된 지원으로 인해 추가 성능 향상에 제약이 있었다.
• 주요 클래스에서는 뛰어난 성능를 보였지만, 빈도가 낮고 세밀한 클래스인 빙하 유출과 흙길과 같은 클래스에서는 어려움을 겪어 클래스 불균형 문제를 드러냈다.
• 반복적 프루닝은 일괄 프루닝 대비 유의미한 성능 향상은 보였지만, 이 특정 아키텍처와 압축 목표에 대해서는 유의미한 이점이 없었다.
• 이 프레임워크는 Jetson Nano에서 실시간 세그멘테이션을 성공적으로 구현하여 UAV 기반 재난 모니터링의 실현 가능성을 입증했다.

Figure 2: Samples from the LPCV dataset (CVPR 2023 challenge)