[논문 리뷰] Token Merging: Your ViT But Faster

Daniel Bolya, Cheng-Yang Fu외 4명|arXiv (Cornell University)|2022. 10. 17.

Generative Adversarial Networks and Image Synthesis인용 수 62

한 줄 요약

Token Merging (ToMe)은 빠른 이분 매칭 프로세스로 유사한 토큰을 병합하여 Vision Transformers의 처리량을 약 2x(비디오의 경우 2.2x)까지 향상시키고, 정확도 손실은 최소화하며, 이미지, 비디오 및 오디오 작업에서 학습 여부에 관계없이 사용할 수 있습니다.

ABSTRACT

We introduce Token Merging (ToMe), a simple method to increase the throughput of existing ViT models without needing to train. ToMe gradually combines similar tokens in a transformer using a general and light-weight matching algorithm that is as fast as pruning while being more accurate. Off-the-shelf, ToMe can 2x the throughput of state-of-the-art ViT-L @ 512 and ViT-H @ 518 models on images and 2.2x the throughput of ViT-L on video with only a 0.2-0.3% accuracy drop in each case. ToMe can also easily be applied during training, improving in practice training speed up to 2x for MAE fine-tuning on video. Training with ToMe further minimizes accuracy drop, leading to 2x the throughput of ViT-B on audio for only a 0.4% mAP drop. Qualitatively, we find that ToMe merges object parts into one token, even over multiple frames of video. Overall, ToMe's accuracy and speed are competitive with state-of-the-art on images, video, and audio.

연구 동기 및 목표

• 무거운 재훈련이나 아키텍처 변경 없이 더 빠른 ViT 추론과 학습을 촉진합니다.
• 토큰 수를 줄여 처리량을 증가시키는 간단한 토큰 병합 모듈을 도입합니다.
• 이미지, 비디오, 오디오에 걸쳐 경쟁력 있는 정확도를 유지하면서 적용 가능성을 확보합니다.
• 기성형 모델에 적합한 경량의 병렬 가능 병합 알고리즘을 제공합니다.

제안 방법

• 각 트랜스포머 블록의 attention과 MLP 분기 사이에 토큰 병합 단계를 삽입합니다.
• 토큰 후보를 병합하기 위해 attention 키 간의 점곱 유사도를 사용하여 토큰 유사성을 정의합니다.
• 토큰 크기로 가중 평균하여 토큰을 병합하는 빠른 이분 소프트 매칭을 사용합니다.
• 병합된 토큰이 여러 입력 패치를 나타내도록 비례적 주의(attention)를 적용합니다.
• 원한다면 M e ToMe를 학습시켜 병합을 풀링으로 간주하고 병합된 토큰을 통해 역전파합니다.
• 레이어 전체에서 제거되는 토큰의 총량을 제어하기 위해 상수 및 감소형 병합 스케줄을 탐색합니다.

실험 결과

연구 질문

• RQ1 가볍게 설계된 토큰 병합 모듈이 재학습 없이 ViT 처리량을 충분히 크게 증가시킬 수 있는가?
• RQ2 토큰을 어떤 방식으로 매칭하고 병합해야 이미지, 비디오, 오디오 모듈에서 정보 내용을 보존할 수 있는가?
• RQ3 병합 스케줄과 특징 선택이 서로 다른 ViT 모델과 사전 학습 방식에서 속도 향상 및 정확도에 어떤 영향을 주는가?
• RQ4 ToMe가 추론 전용과 학습 중 모두에서 효과적이며 다양한 모달리티에서 적용 가능한가?

주요 결과

• ToMe는 이미지에서 ViT-L@512 및 ViT-H@518의 처리량을 대략 두 배로 증가시키고, 약 0.2–0.3%의 정확도 손실을 보입니다.
• 비디오에서 ToMe는 학습 없이 ViT-L의 처리량을 약 2.2x까지 달성하며 정확도 손실은 0.2–0.3%에 그칩니다.
• ToMe로 학습 시 비디오에서 MAE 파인 튜닝의 학습 속도 증가가 약 2x까지 관찰됩니다.
• 오디오의 경우 MAE 사전 학습을 사용했을 때 ViT-B에서 약 2x의 처리량 증가와 약 0.4%의 mAP 감소를 달성합니다.
• ToMe는 프레임 간 객체 부분을 병합하고 비디오에서 부분 추적을 수행할 수 있으며, 배경과 전경 정보를 병합하되 큰 정확도 저하 없이 작동합니다.
• ToMe는 추가 학습 트릭이나 새로운 매개변수 없이 이미지, 비디오, 오디오에서 최첨단 방법과 경쟁력이 있습니다.