TOKEN MERGING: YOUR VIT BUT FASTER

[논문리뷰]

ICLR 2023 notable top 5%

---

ABSTRACT

• 훈련할 필요 없이 기존 ViT 모델의 처리량을 늘릴수 있는 간단한 방법
• 일반적이고 가벼운 알고리즘을 사용하여 유사한 토큰들을 점진적으로 합친다.
• ToMe(Token Merging)은 training동안 쉽게 적용이 가능하다.

---

Introduction

• Transformer의 토큰을 Run-time에 Pruning하여 더 빠른 모델을 가능하게 하는 분야가 등장
• Token Pruning의 대부분은 Training 속도를 높이기 위해 적용할 수 없다. 
• 따라서 Token을 Pruning하는 것 보다 더 나은 방식인 combine을 적용하여 Token Merging 방법을 제안한다.
• 본 논문에서의 custome matching algorithm은 pruning만큼 빠르면서 더 정확하다. 게다가 본 논문의 방법은 training 중이거나 training 중이 아니더라도 가능하다. (훈련 여부와 상관없음)

---

RELATED WORK

Efficient Transformers

• 본 논문에서는 기존의 ViT model의 token들을 merging함으로써 속도를 높이는 것에 초점을 두었다.

Token Reduction

• Transformer는 token의 수에 상관없이 작동할 수 있기 때문에, 최근 많은 연구에서 이러한 token들을 pruning하는 시도가 많이 있었다.
• 하지만, 이러한 방법들은 training을 필요로 하지만, 우리의 방법은 training이 없이 사용이 가능하다.
• 또한, 대부분의 pruning work들은 동적(dynamic)하다. 즉, 토큰의 수는 image 또는 문장마다 다르다. 
  • 이러한 방법은 accuracy에는 도움이 되지만, 토큰 수가 다른 샘플은 더이상 배치할 수 없기 때문에 실용성이 없다.
• 이러한 문제들을 해결하기 위해 대부분의 pruning 논문은 토큰을 제거하는 대신 훈련 중에 mask를 적용하여 pruning속도를 저하시킨다. 반면에, 본 논문의 방법은 inference와 training 모두에 적용되어 두 경우 모두 실제 속도 향상에 달성이 가능하다.

Combining Tokens

• Token을 제거하는 작업은 많지만, combine하는 작업은 거의 없다.
• Token Polling(2021)이 가장 우리의 방법과 유사하다. 그러나, 느린 k-means 접근법을 사용하였다.
• 지금까지 Training없이 토큰을 결합할때 speed-accuracy trade-off를 제공하는데 성공한 접근 방식이 없었다.

---

TOKEN MERGING

• 논문의 목표는 기존의 ViT에 token merging module을 삽입하는 것이다.
• 중복되는 tokens들을 합침으로써 반드시 train할 필요는 없지만 연산량을 증가시킬 수 있기를 희망한다.

Strategy

• $ r $ 은 token들 양이다.
• Transformer의 각 블록에서 token을 병합하여 layer당 $ r $개를 줄인다.
• 네트워크의 $ L $ 블록을 통해, 점진적으로 $ rL $ token을 합친다. 토큰 수가 적을수록 정확도는 낮지만 처리량은 높아지기 때문에 $ r $을 변경하면 speed - accuracy의 밸런스를 맞출 수 있다.
• 중요한 것은 이미지의 내용에 상관업이 $ rL $ 토큰을 줄인다는 것이다. 

• 본 논문은 token merging step을 attention과 MLP 사이에 배치함으로써 token에서 정보를 흘려줄 수 있으며, attention의 결과를 통해 어떤 대상을 병할할지를 결정 할 수 있다. -> 두 가지 모두 정확도를 높일 수 있다.

Token Similarity

• Similar token을 병합하기 전에, 우선 'similar'하다는 기준을 내려야한다.
• 두 Token의 특징 사이의 distance가 작은 경우에 유사하다고 할 수 있지만, 이것은 반드시 최적은 아니다.
• Transformer는 overparameterized 되어있다. 이는 features에 중요하지 않은 노이즈를 포함할 가능성이 있음을 의미한다.
• 본 논문에서는 어떤 것이 유사한지를 결정하기 위해 각 Token의 Key 사이에 dot product similarity metric을 사용한다.

Biparite Soft Matching

• Token similarity가 결정되고 전체 수를 $ r $만큼 줄이기 위해 일치시킬 Token을 결정하는 빠른 방법이 필요하다.
• 논문에서는 다음과 같은 목표로 알고리즘을 구현하였다.
  • 1) 병렬화할 수 없는 반복적인 것들을 피한다
  • 2) 병합이 만드는 변화가 점진적이기를
• 논문에서는 다음과 같은 순서로 알고리즘을 구현하였다.
  • 1) Token들을 같은 사이즈의 A와 B 그룹으로 나눈다.
  • 2) A의 각 토큰에서 B의 가장 유사한 토큰으로 선을 그린다.
  • 3) $ r $개의 가장 유사하게 유지한다.
  • 4) 연결된 Token들을 병합한다.
  • 5) 두 그룹을 다시 합친다.
• 위 방식처럼 수행하게 되면 A의 각 토큰에는 하나의 엣지가 있기 때문에 4단계에서 연결된 구성 요소를 찾는 것은 간단하다.
• 위와 같은 방식을 "bipartite soft matching"이라 하고, 이는 Token을 무작위로 떨어뜨리는 것과 거의 같은 속도이다.

Tracking Token Size

• Token이 병합되면 더 이상 하나의 입력 Patch를 나타내지 않게 되는 것이다. 
  • 이렇게 되면 softmax attention의 결과를 바꿀수 있다. 예를들어 같은 key를 가진 두 개의 토큰을 병합한다면, 그 key는 softmax term에 덜 영향을 받을 것이다.

Proprotinal attention

• $ s $는 각 토큰의 크기(토큰이 나타내는 patch 수)를 포함하는 행 벡터이다.

---

IMAGE EXPERIMENTS

• ImageNet-1k을 4가지 다른 방법으로 학습하였다.
  • AugReg, MAE, SWAG, DeiT
• 논문의 방법으로 off-the-shelf(미리 학습된 모델)을 돌리거나, MAE와 DeiT의 경우, 논문의 방법을 적용하여 Training을 시킴.

Token Similarity

• attention 후 attention key를 사용하여 병합을 수행하는것이 훨씬 정확하다.
• Table 1b와 같이 Token 거리를 측정하는 것이 cosine similarity가 가장 좋다. 이후 효율성을 위해 Attention Head를 연결하는 대신 K 평균을 수행한다.

Algorithmic Choices

• 병합할 Token을 결정한 후에는 Token 크기 $ s $로 weighted averaging을 수행하여 결합한다.
• 두 개의 세트 A, B를 분할할 때 교대로 Token을 할당하는 것이 가장 잘 작동한다.

Proportional Attention.

• off-the-shelf MAE 모델을 제외하고 모든 모델에 대해 Proportional attention을 사용한다.

Comparing Matching Algorithms

• Table2에 본 논문의 방법인 bipartite matching과 다른 token reduction algorithm을 비교하였다.
• 논문의 bipartite matching은 greedy한 접근법의 정확성과 pruning 속도를 가지고 있다.

Selecting a Merging Schedule.

• 기본적으로 layer당 Token을 병합한다.
• 이 방법의 최적함을 테스트하기 위해 실험을 한 결과 일정하게 병합된 토큰이 증가함에 따라 최적에 가깝다는 것을 발견하였다.

 

이외 다른 실험적인 세팅은 스킵

---

COMPARISON TO OTHER WORKS