[24’] LLaVA-PruMerge: Adaptive Token Reduction for Efficient Large Multimodal Models

🔍 Abstract

• Problem
  • 최근의 LMM은 High-resolution Image Token을 Transformer 구조에서 처리하기 위해 굉장히 많은 계산량을 필요로 한다.
• Insight
  • CLIP의 마지막 직전 Layer에서 [CLS] Token과의 Attention Score를 보면, Image Token 중 중요한 Token이 어떤 것인지 알 수 있다.
  • 또한 Attention Score가 높은 Token은 굉장히 sparse해서, dynamic selection을 통해 중요한 Token만 남길 수 있다.
• Method
  • 따라서, 이를 활용해 중요한 Image Token만을 남기고, 나머지 Token을 제거하는 방법을 제안한다.
  • 구체적으로는, PruMerge라는 방법을 제안하며, Clustering을 통해 주변 정보를 통합하는 과정을 거친다.
  • 추가로 PruMerge+을 제안하며, 여기서는 미처 고려하지 못한 Image Token을 Grid 형식으로 더 Sampling하여 계산량은 약간 늘리면서 성능을 향상시키도록 했다.
• Result
  • LLaVA-1.5에 적용하여 Visual Token의 수를 1/14로 줄였으며, 성능은 약간 저하되었다.
  • 이 방식은 LLM에 들어가는 Input부터 Visual Token의 수가 줄어드는 것이지만 Training-free로도 작동할 수 있으며, 더 높은 성능을 위해서는 LoRA Fine-tuning이 필요하다.

1. PruMerge

PruMerge는 중요한 Image Token만을 남기고, 나머지 Token을 제거하는 방법을 제안한다. 이를 위해 CLIP의 마지막 직전 Layer에서 [CLS] Token과의 Attention Score를 보면, Image Token 중 중요한 Token이 어떤 것인지 알 수 있다. 또한 Attention Score가 높은 Token은 굉장히 sparse해서, dynamic selection을 통해 중요한 Token만 남길 수 있다. 이러한 아이디어로 제안된 PruMerge는 다음과 같은 과정을 거친다.

1. Adaptive Important Token Selection (AITS)
   • CLIP의 마지막 직전 Layer에서 Image Token과 [CLS] Token의 Attention Score를 계산한다.
   • 이때 높은 Attention Score를 가진 Token 순서대로 정렬하여 중요한 Image Token을 선택한다.
2. Token Supplement (TS)
   • 중요한 Image Token을 선택한 후, 각 Token의 Representational Ability을 높이기 위해 Clustering을 통해 주변 정보를 통합한다.
3. Weighted Token Merging
   • Clustering을 통해 주변 정보를 통합할 때, 각 정보의 중요도를 의미하는 Weight은 Attention Score를 이용한다.
   • 이를 통해 $n$개의 Token이 $m$개로 Reduction되며, 이를 LLaVA와 같은 LMM의 Visual Input으로 사용한다.

각 과정을 자세히 살펴보자.

1.1. Adaptive Important Token Selection (AITS)

AITS는 CLIP의 마지막 직전 Layer에서 Image Token과 [CLS] Token의 Attention Score를 계산한다. Attention Value를 보면 알 수 있듯, Image Token 중 중요한 Token은 sparse하게 나타난다. 따라서 이를 dynamic selection을 통해 중요한 Token만 남길 수 있다. 이때는 IQR(Interquartile Range)를 이용하여 Outlier Detection으로 중요한 Token을 선택한다. 즉, IQR은 Q3과 Q1의 차이로, Attention 값이 Q3 + 1.5 $\times$ IQR 이상인 Token을 중요한 Token으로 선택한다.

Visual Token Sampling에서는 두 가지를 관찰할 수 있었다. (1) Image Complexity가 높을수록 중요한 Token이 많다. (2) 대개 Image Token 중 중요한 Token은 실제로 Spatial하게 중요한 Content와 Align하는 경향이 있다.

1.2. Token Supplement (TS)

만약 Large Object가 있는 이미지가 있다고 가정하자. 이때, 굉장히 공격적인 Token Pruning을 통해 매우 중요한 Token만 남겼다면, 이러한 중요한 Token만으로는 Object의 전체적인 정보를 담기 어렵다. 따라서 Clustering을 통해 주변 정보를 통합하는 과정이 필요하다. 이를 위해 Token Supplement (TS)를 제안한다. 유사도를 계산할 때에는 key vector를 사용한다.

Integration 시에는 [CLS] Attention 값 $\textbf{a}[i]$를 이용하여 Weighted Average를 계산한다. 이때, Attention Score가 높은 Token일수록 Weight가 높아지게 된다.

1.3. Full Algorithm

전체 알고리즘을 살펴보면 위와 같다. PruMerge+*의 경우 **spatially uniform sampling을 수행하여 계산량은 약간 늘리면서 성능을 향상시키는 추가 방법론이다.

2. Experiments

2.1. Main Results

PruMerge의 경우 Performance Gap이 꽤나 있는데, PruMerge+의 경우 이를 줄일 수 있었다. 어떻게 보면, 이는 CLIP Model의 Attention만으로는 충분한 정보를 얻기 어렵다는 것을 보여준다. 추가로 저자들은 Video Benchmark에 대해서도 실험을 진행했는데, 이는 생략한다.

2.2. Ablation Study

추가 Ablation Study를 통해 확인한 결과는 다음과 같다.

• Token Sampling Strategy (AITS): Random Sampling의 경우 성능 하락이 크다. Spatially Uniform Sampling의 경우 성능 하락이 Random Sampling보다는 적지만, LLaVA-PruMerge에 비해서는 성능이 낮다.
• Token Supplement (TS): AITS에서 TS를 사용하여 성능이 향상되었다. 특히, SQA의 경우 오히려 기존 모델보다 성능이 높아졌다.
• LoRA Fine-tuning: Fine-tuning을 통해 성능을 높일 수 있다.

💡 Summary

해당 논문의 내용을 간단히 요약하면 다음과 같다.

• CLIP의 [CLS] Token과의 Attention Score를 통해 LMM에 중요한 Image Token을 선택할 수 있다고 주장하나, 이 방법으로 수행한 PruMerge에는 한계가 있으며 Spatial Uniform Sampling을 추가한 PruMerge+에서만 성능이 유지됨

📃 Reference

• [24’] LLaVA-PruMerge: Adaptive Token Reduction for Efficient Large Multimodal Models

Multimodal 카테고리 내 다른 글 보러가기