[24’] HiRED: Attention-Guided Token Dropping for Efficient Inference of High-Resolution Vision-Language Models in Resource-Constrained Environments

🔍 Abstract

• Problem
  • 최근의 많은 LMM(Large Multimodal Models)은 High-Resolution Image를 처리하기 위하여 위와 같이 이미지를 조각내어 여러 번 CLIP과 같은 Vision Encoder를 통과시키는 Dynamic Partitioning을 사용한다. 그러나 이러한 방법은 Visual Token의 수를 늘려 계산량을 굉장히 증가시키는 문제가 있다.
• Insight
  • 저자들은 LLM이 일부 Visual Token에만 LLM이 High Attention Score를 가지는 것을 관찰하였다. 따라서, 이러한 상위 Token들만 남겨도 성능 저하가 적을 것으로 예상하였다.
  • 그러나, 이러한 Attention을 계산하기 위해서는 반드시 LLM을 통과해야만 한다. 저자들은 Inference Efficiency를 향상시키기 위한 목적을 가지고 있으므로, LLM 없이도 이러한 상위 Token을 예측할 수 있는 방법을 찾아야 했다.
  • 이 과정에서 저자들은 CLIP의 Layer-wise Attention을 분석하여 CLIP의 Early Layer는 공간적으로 Subject가 있는 Patch, Late Layer는 Semantic하게 중요한 Patch에 Attention을 준다고 해석하였다. 따라서 이들을 결합하여 상위 토큰을 예측하여 사용하였다.
• Method
  • 저자들은 resource constraint 환경에서 Inference Efficiency를 향상시키기 위해 HiRED(High-Resolution Early Dropping)이라는 방법을 제안한다.
  • HiRED는 training-free method로, Vision Encoder 이후 LLM Input으로 들어가는 Visual Tokens 중 Attention을 기반으로 중요하지 않은 Token을 Dynamic Dropping하는 방법이다. 저자들은 다른 방법들과 달리 LLM Input 자체에 대한 Dropping을 수행하여 연산량을 획기적으로 줄였다는 점을 강조한다.
• Result
  • 저자들은 다양한 LMM 대상 Benchmark에서 약간의 성능 저하와 함께 기존의 20~40% 정도의 Visual Token만으로 Inference를 수행할 수 있음을 보여준다.

1. Key Insight

1.1. LMM Attention

Visual Token Sparsity. 저자들은 LMM의 Attention Matrix를 관찰하였다. 그 결과, 일부 Visual Token만이 LLM의 Attention을 많이 받는 것을 확인할 수 있었다. 따라서, 이러한 Token들만 남겨도 성능 저하가 적을 것으로 예상하였다. 그러나 이는 명확하지 않은 분석인데, Attention Sink 개념에 따르면 상위 Token이라고 해서 반드시 중요한 정보를 담고 있는 것은 아니기 때문이다. 따라서, 너무 적은 상위 Token만을 사용하면 성능이 급격히 저하될 수 있다.

Sub-images with Different Context Amounts. 저자들은 Dynamic Partitioning을 통해 High-Resolution Image를 처리하는 LMM을 관찰하였다. 각 Sub-image마다 다른 Context가 포함되어 있고, 따라서 질문에 따라 Sub-image의 중요도가 다를 것이다. 위 Table에서는 Abstract 부분에서 설명한 Figure 1에 대해 상위 10% Attention이 속하는 비율을 나타내는데, Figure 1에서는 Road Sign에 대해 물었기 때문에 Road Sign이 포함된 Sub-4의 비율이 높게 나타난다.

1.2. CLIP [CLS] Attention

그렇다면 LMM의 Attention Matrix를 분석하여, 중요한 Token만 남기면 성능 저하 없이 적은 토큰으로 Inference를 할 수 있음을 추측할 수 있다. 그러나, 이러한 Attention을 계산하기 위해서는 반드시 LMM을 통과해야만 한다. 따라서, Inference Efficiency를 향상시키기 위해서는 LLM 없이도 중요한 Token을 예측할 수 있는 방법이 필요하다. 저자들은 이러한 방법으로 CLIP의 Layer-wise Attention을 분석하였다. 일반적으로 [CLS]는 Semantic Information을 담기 위해 노력하며, 따라서 Query가 [CLS]일 때의 Attention Matrix는 Semantic Information을 담고 있을 것이다. 이를 Layer-wise로 해석해보면 다음과 같다.

1. Early Layer: Spatially Subject-Patch에 Attention을 주는 것을 확인할 수 있다. 즉, Main Content에 집중한다.
2. Late Layer: Semantically Important Patch에 Attention을 주는 것을 확인할 수 있다. 즉, Informative Area에 집중한다. 여기에는 Background처럼 보이지만 Global Context를 담고 있는 Register Token도 포함된다.

따라서 Main Content를 담고 있는 Early Layer와 Informative Area를 담고 있는 Late Layer를 적절히 사용하면, 중요한 상위 Token을 예측할 수 있을 것이다.

2. Method: HiRED

HiRED(High-Resolution Early Dropping)은 크게 (1) Token Budget Distribution과 (2) Token Dropping 과정으로 나뉜다.

Token Budget Distribution. 먼저, 정해진 Token Budget $b_ {total}$이 있다고 가정하자. 이 중에서 Full Image와 Sub Image를 처리하는 데 필요한 Token Budget 비율은 $\alpha \in [0, 1]$로, hyperparameter로 정한다. 즉, Full Image는 $b_ {full} = \alpha \cdot b_ {total}$만큼의 Token Budget을 사용하고, Sub Image는 $b_ {sub} = (1 - \alpha) \cdot b_ {total}$만큼의 Token Budget을 사용한다. 이제 Sub Image들 간의 Budget을 나누어야 한다. 여기서 저자들은 CLIP의 Early Layer를 사용한다. 보통 Main Content를 담고 있는 Early Layer는 Spatially Subject-Patch에 Attention을 주기 때문에, Sub Image들 간의 Main Content를 비교하여 Budget을 분배한다. 자세히는, Full Image에서 [CLS]와의 Attention을 각 Partition마다 Sum으로 계산하고 이 비율에 따라 Budget을 분배한다.

Token Dropping. 이제 Sub Image들 간의 Budget이 정해졌다. 이제 Sub Image들 각각의 Visual Tokens 중 Attention Score가 낮은 Token을 Dynamic Dropping한다. 이때에는 각 Sub Image의 Late Layer의 [CLS]와의 Attention을 사용하여 상위 Token만을 남겨둔다. 지금까지의 과정을 알고리즘으로 정리하면 다음과 같다.

3. Result

3.1. Main Results

Performance와 Efficiency는 각각 위와 같다. 기존의 Efficient LLM Method에 비해 훨씬 적은 Visual Token을 사용하여 더 높은 성능을 냈다. 일부의 경우에는 Baseline LMM보다 높은 성능이 나오기도 했다. 해당 벤치마크는 SQA, POPE였고, 따라서 일부 문제에서는 오히려 핵심적인 Visual Token만을 사용하는 것이 성능에 도움이 되는 것으로 보인다.

3.2. Ablation Study

Budget Distribution. Budget을 효율적으로 분배하는 것이 성능 향상에 영향을 미친다. 한편, Full Image와 Sub Image 간의 Budget 비율 $\alpha$는 상당히 Robust했다.

Token Dropping. Token Dropping 시 하나의 Head만 보는 것보다 모든 Head를 보는 것이 성능에 도움이 된다. 이는 Multi-Head Attention의 특성을 반영한 결과이다. 또한, Vision Transformer도 Spatial Information이 포함되어 있고, LMM과 직접적으로 연결되는 Late Layer를 사용하는 것이 성능에 도움이 된다. Main object를 잡는 Early Layer만 사용하는 것은 성능이 떨어지는 것으로 나타났다.

Application to Low-resolution Image. Low-resolution Image에 대해서도 적용해보았다. 이때에도 상당히 효과적이었으며, 상위 Token만을 사용하여도 성능이 높게 나타났다.

💡 Summary

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

• Visual Token 중 일부만 LMM Attention이 높은 것을 발견하고, 이를 CLIP의 [CLS] Token Attention으로 예측하는 Heuristic한 방법을 제안
• 적은 Visual Token을 사용하여 Inference Efficiency는 향상시키면서 성능은 약간 저하됨; 일부에서는 오히려 적은 Visual Token을 사용하는 것이 성능이 높게 나타남

📃 Reference

• [24’] HiRED: Attention-Guided Token Dropping for Efficient Inference of High-Resolution Vision-Language Models in Resource-Constrained Environments

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