Token Merging: Your ViT But Faster

추가 훈련 없이 토큰 병합으로 속도 높이기

 

Github

arXiv

 

 

Abstract

추가 훈련 없이 ViT 모델의 처리량을 늘릴 수 있는 토큰 병합(ToMe) 소개.

ToMe는 경량 매칭 알고리즘을 사용하여 유사한 토큰을 점진적으로 결합한다.

 

 

 

Introduction

컴퓨팅 비용을 줄이기 위해 Swin Transformer와 같은 방법들이 나왔지만 vanilla ViT는 여전히 바람직한 특성 몇 가지를 갖고 있다.

• 간단한 행렬 곱셈으로 구성
• MAE와 같은 자체 감독 훈련 가능
• 데이터에 대한 가정이 부족하여 다양한 양식에 적용 가능하고 확장성이 있음

 

더 빠른 모델을 위해 런타임 시 토큰을 정리하는 token pruning이 제안되었지만 정보 손실의 우려가 있고 모델의 재교육해야 함.

 

본 논문에서는 토큰을 제거하는 대신 결합하는 Token Merging(ToMe)을 제안한다.

매칭 알고리즘으로 인해 pruning 보다 더 정확하고 빠르며, 모델의 훈련 여부와 관계없이 작동한다.

 

 

 

Token Merging

본 연구의 목표는 토큰 병합 모듈을 ViT에 삽입하여 추가 훈련 없이 처리량을 늘리는 것이다.

---

Strategy

Transformer의 L개의 각 블록에서 r개의 토큰을 줄여 총 rL개의 토큰을 줄인다.

r을 통해 정확도-속도의 trade off가 가능.

 

다음 그림과 같이 attention과 MLP 사이에 토큰 병합을 적용하여 attention에서 병합 항목을 결정할 수 있도록 한다.

---

Token Similarity

토큰 병합 전 '유사'의 의미를 정의해야 하는데, transformer의 중간 feature는 과도하게 매개변수화되기 때문에 단순히 feature 사이의 거리를 비교하는 것은 최적이 아니다.

 

여기서는 QKV self attention의 Key를 사용하는데, 이는 key가 내적 유사성에 사용되기 위해 각 토큰의 정보를 이미 요약하고 있기 때문이다.

Attention

 

최종적으로 key의 cosine distance를 사용.

---

Bipartite Soft Matching

각 레이어에서 수천 개의 토큰에 대해  수행 가능한 효율적인 솔루션이 필요하다. 

솔루션은 병렬화 할 수 있고 점진적이어야 한다.

(클러스터링은 너무 잘? 묶여서 네트워크에 부정적 영향을 끼칠 수 있어 매칭을 사용한다고 함.)

 

1. 토큰을 두 세트로 나눔
2. A의 각 토큰에서 B의 토큰 중 가장 유사한 토큰과의 edge를 그림
3. 가장 유사성이 높은 r개의 연결만 선택
4. 병합

실제로 이 'bipartite soft matching'은 무작위 토큰을 드롭하는 것만큼이나 빠르고 코드 몇 줄만 추가하면 가능하다.

 

병합 방법과 토큰 세트 선택 방법 비교

---

Tracking Token Size

토큰의 병합은 attention의 softmax 과정에 영향을 줄 수 있는데, 하나의 key를 가지고 있지만 실제로는 여러 개의 패치를 포함하고 있기 때문이다.

 

이는 proportional attention 이라고 명명한 간단한 변경으로 해결한다.

Proportional attention은 각 토큰의 크기(잠재적 패치 수)를 포함하는 벡터 s를 attention 과정에서 가중치로 반영하는 것이다.

---

Training with Merging

훈련에서는 복잡한 gradient trick을 사용하지 않고 average pooling을 사용하는 것처럼 병합된 토큰을 통해 역전파한다.

 

 

 

Experiments