본문 바로가기

Test-time Adaptation

Beyond Text-Visual Attention: Exploiting Visual Cues forEffective Token Pruning in VLMs

논문 주소: https://arxiv.org/pdf/2412.01818
깃허브: https://github.com/Theia-4869/VisPruner

 

Abstract

 Large Vision-Language Models (LVLMs)는 일반적으로 텍스트 토큰보다 훨씬 더 많은 visual tokens를 포함하고 있으며, 이는 상당한 계산 부담을 초래합니다. 최근 이러한 문제를 해결하기 위해 language model 초기 단계에서 visual tokenspruning하려는 노력이 이루어져 왔습니다. 대부분의 기존 연구들은 textvisual tokens 사이의 attention scores를 사용하여 visual tokens의 중요도를 평가합니다. 하지만 본 연구에서 language model 내의 text-visual attention을 먼저 분석하였으며, 이 점수가 토큰 pruning을 위한 이상적인 지표가 아닐 수 있음을 발견했습니다.

 

 이러한 분석을 바탕으로, 본 연구에서는 LVLMs에서 더욱 효과적인 토큰 pruning을 위해 visual cues를 활용하는 plug-and-play 방식인 VisPruner를 제안합니다. 구체적으로, 먼저 visual attention을 사용하여 제한된 수의 중요한 토큰들을 선택합니다. 그런 다음, 남은 토큰들 중에서 similarity를 기반으로 중복된 토큰들을 제거합니다. 초기에 선택된 중요한 토큰들과 함께 다양한(diverse) 토큰들을 유지함으로써, 입력 이미지의 시각적 정보를 최대한 보존합니다.

 

 실험 결과는 VisPruner가 다양한 VLM architecturesreduction ratios 전반에서 강력한 성능을 유지하며, text-visual attention에 기반한 기존 방법들을 크게 능가함을 입증합니다. 특히, 별도의 훈련 없이도 VisPruner는 유사한 성능을 유지하면서 LLaVA-1.5-7BFLOPs91%, inference latency을 75%까지 줄일 수 있습니다. 

 

1. Introduction

 

최근 LLMs의 발전과 함께, 그 강력한 추론 능력을 multi-modal tasks로 확장하려는 수많은 노력이 이어졌고, 이는 VLMs의 등장으로 이어졌습니다. 기존의 VLMs는 일반적으로 시각적 입력을 시퀀스 표현으로 변환하는 visual encoder (예: CLIP), 질문을 처리하고 답변을 생성하는 language model (예: Llama 2), 그리고 language model이 시각적 토큰을 입력으로 받아들일 수 있게 하는 modal alignment module (예: MLP, Q-Former)로 구성됩니다.

 

 이러한 접근 방식은 visual question answering과 같은 멀티모달 태스크에서 놀라운 성능을 달성했지만, 입력 시퀀스의 대다수를 차지하는 visual tokensVLM 추론의 계산 비용을 크게 증가시킵니다. 예를 들어, LLaVA-1.5는 336 × 336 이미지를 576개의 visual tokens로 변환하는 반면, 텍스트 토큰은 대개 100개 미만입니다. 고해상도 입력이나 비디오 스트림은 더 긴 시각적 토큰 시퀀스를 생성하여 계산 부하를 더욱 가중시킵니다.

 

 VLM의 추론 비용을 줄이기 위해, FastVlanguage model 내에서 visual tokens를 조기에 pruning하는 시도를 처음으로 제안했습니다. 이 연구는 language model의 깊은 레이어에서 상당수의 visual tokens가 텍스트 토큰보다 현저히 적은 attention을 받는다는 점을 확인하고, 레이어 2 이후에 이러한 low-attention visual tokens를 대거 제거합니다. 이후 많은 연구가 이 패러다임을 따라 text-visual attention에 기반하여 visual tokenspruning해 왔습니다. 비록 VLM 내의 visual tokens가 중복적이라 하더라도, language model 내부의 text-visual attention이 토큰 pruning을 위한 이상적인 지표가 아닐 수 있음을 발견했습니다. 그림 1에서 보듯, text-visual attention 기반의 pruning은 위치 편향(positional bias) 문제로 인해 덜 유용한 정보가 담긴 이미지 하단부(또는 무의미한 paddings)를 주로 보존합니다. 반면, 중요한 시각 정보가 있는 중앙 영역의 토큰들이 제거되어 모델이 정답을 맞히기 어렵게 만듭니다.

 

본 연구에서는 language model 내의 text-visual attention을 심층 분석하여 두 가지를 발견했습니다:

  1. Text-visual attention shift: language modelrotary position embedding이 가진 장기 쇠퇴(long-term decay) 특성으로 인해, [이미지, 텍스트] 구조의 입력에서 뒤에 위치한 텍스트 토큰들은 물리적으로 자신과 더 가까운 visual tokens(즉, 이미지 하단 패치들)에 더 높은 attention을 주는 경향이 있습니다.
  2. Text-visual attention dispersion: Positional Embedding 을 조정하더라도, attention density curve가 높은 엔트로피와 낮은 피크를 보이는 문제(균일 분포에 가까움)가 지속되어 실제로 중요한 visual tokens를 선택하기 어렵게 만듭니다.

이러한 관찰을 바탕으로, VLMs에서 더 효과적인 토큰 pruning을 위해 visual cues를 활용하는 plug-and-play 방식인 VisPruner를 제안합니다. VisPruner는 먼저 visual encoderattention을 기반으로 풍부한 시각 정보를 가진 소수의 중요한 토큰을 선택합니다. 그다음, 보존된 토큰들이 전경(foreground)에만 치우쳐 배경 정보를 놓치지 않도록, similarity를 기반으로 나머지에서 중복 토큰을 제거합니다. 이렇게 남겨진 다양한(diverse) 토큰들은 초기 선택된 중요 토큰들과 함께 입력 이미지의 시각 정보 보존을 극대화합니다. 특히, language model 내부에서 pruning하는 기존 방식과 달리, 본 방식은 language model 입력 전에 visual cues를 활용해 pruning하므로 FlashAttention과 같은 다양한 최적화 기술과 호환되며 더 높은 추론 효율성을 달성합니다.

단순한 설계 덕분에 VisPrunerLLaVA 시리즈 및 Qwen-VL 등 다양한 VLMs에 쉽게 적용할 수 있습니다. 실험 결과, VisPruner는 극한의 reduction ratios 하에서도 높은 성능을 유지함을 입증했습니다. 예를 들어, LLaVA-1.5-7B에서 94.4%의 visual tokens를 제거하여 FLOPs를 95% 이상 줄이면서도 10개 태스크에서 원본 성능의 91.5%를 유지했습니다.

본 연구의 기여도는 다음과 같이 세 가지로 요약됩니다:

  1. language modeltext-visual attention이 토큰 pruning을 위한 이상적인 지표가 아닐 수 있음을 철저히 조사하고 발견했습니다.
  2. 관찰 결과를 바탕으로 visual cues를 활용하여 효과적으로 토큰을 pruning하는 plug-and-play 방식인 VisPruner를 도입했습니다.
  3. 다양한 VLMsVisPruner를 적용하고 광범위한 실험을 통해 극한의 reduction ratios에서도 높은 성능을 유지함을 보였습니다.

2. RELATED WORK

Vision-language models (VLMs)

최근 LLMs의 눈부신 성공은 그 강력한 추론 능력을 멀티모달 이해 태스크로 확장하려는 추세로 이어졌고, 이는 VLMs의 탄생으로 이어졌습니다. 이러한 VLMs는 일반적으로 입력 이미지를 serializing representations으로 변환하는 visual encoder와 텍스트 생성을 위한 language model을 포함합니다. language model이 시각적 표현을 입력으로 받아들일 수 있도록 하기 위해, 이들은 대개 단순한 선형 레이어, MLP projector, 또는 깊은 쿼리 기반 네트워크(query-based network)와 같은 alignment module을 사용하여 시각 및 언어 모달리티를 정렬합니다.

비록 이것이 language model에 시각적 인지 능력을 부여하지만, 긴 visual token sequences의 도입은 계산 부담을 증가시킵니다. 더욱이, 기존 VLMs가 여전히 시각적 결함이나 hallucinations(환각) 문제로 고통받고 있다는 연구 결과들이 있습니다. 이를 완화하기 위해 입력 이미지의 해상도를 높여 성능을 개선하려는 노력이 이루어져 왔으나, 이는 계산 부담을 더욱 악화시킵니다. 자원이 제한된 실제 시나리오에서 VLMs를 응용하기 위해서는 추론 프로세스를 최적화하는 것이 시급한 과제입니다.

Token reduction for VLMs

VLM 추론을 최적화하는 한 가지 방법은 입력 시퀀스의 대다수를 차지하는 visual tokens를 줄이는 것입니다. 몇몇 연구들이 language model에서의 token reduction을 탐구해 왔습니다. 텍스트와 비교했을 때, 이미지 정보는 더 높은 중복성(redundancy)을 갖는 경향이 있어 VLMs를 위한 visual token reduction은 더욱 합리적이고 효과적입니다. visual tokens를 더 콤팩트한 표현으로 압축하려는 다양한 노력이 있었으나, 이러한 방법들은 추가적인 훈련을 요구하여 계산 부담과 훈련 비용을 증가시킵니다.

본 연구에서 주로 training-free VLM inference optimization에 집중하며, 이는 일반적으로 중요도가 낮은 토큰을 제거함으로써 달성됩니다. FastVVLMs에서 visual tokens의 중복성과 비효율성을 처음으로 식별하고, language model의 두 번째 레이어 이후에 attention scores가 낮은 visual tokens를 제거하는 간단한 방법을 제안했습니다. SparseVLM은 텍스트 프롬프트에서 방해 요소(distractions)를 제거하고 더 정확한 text attention을 사용하여 visual tokens를 점진적으로 희소화(sparsify)합니다. 이후 점점 더 많은 연구가 visual tokens의 중요도를 평가하기 위해 language model 내부의 text-visual attention에 의존하는 방식을 선택하고 있습니다. 그러나 심층적인 조사 끝에, 이러한 text-visual attention이 토큰 pruning을 위한 이상적인 지표가 아닐 수 있음을 발견했습니다. 이전 연구들과 달리,  더 효과적인 토큰 pruning을 위해 visual cues를 활용할 것을 제안합니다. 특히, 일부 동시기 연구들도 

유사한 결론을 공유하고 있습니다.

 

3. TEXT-VISUAL ATTENTION INVESTIGATION

대부분의 이전 연구들은 language model 내의 text-visual attention을 활용하여 pruning을 위한 visual tokens의 중요성을 평가합니다. 이 섹션에서 심층 조사를 수행하고, 이 attention이 토큰 pruning을 위한 이상적인 지표가 아닐 수 있음을 제안합니다. 섹션 3.1에서는 VLMsattention mechanismspreliminaries으로 소개합니다. 이후 섹션 3.2와 3.3에서는 text-visual attention에서 관찰되는 두 가지 현상인 attention shiftattention dispersion을 분석하며, 이들은 각각 positional biasdensity distribution의 영향을 받습니다.

3.1. Preliminaries

기존 VLMs는 일반적으로 visual encoderlanguage model이라는 두 가지 핵심 구성 요소로 이루어지며, 둘 다Transformer 기반 아키텍처입니다. 둘 다 self-attention mechanism에 의존하지만, 구체적인 구현에는 미묘한 차이가 있습니다.

 

visual encoder (예: CLIP)는 global attention mechanism을 사용합니다. 여기서 패치 단위의 이미지 $X$는 세 개의 가중치 행렬 $W_Q, W_K, W_V$를 통해 각각 query $Q$, key $K$, value $V$로 변환됩니다:

 

$$Q = XW_Q, K = XW_K, V = XW_V \quad                                  (1)$$
 

그 후, scaled dot-product attention이 다음과 같이 계산됩니다:

$$A = \text{softmax} \left( \frac{QK^T}{\sqrt{d}} \right), O = AV \quad (2)$$

 

여기서 $A$의 첫 번째 행을 [CLS] attention이라고 부릅니다.

 

반면, language model (예: Llama 2)은 각 토큰이 이전 토큰(과거)에만 주목할 수 있고 이후 토큰(미래)에는 주목할 수 없는 causal self-attention을 활용합니다. relative rotary position embedding이 포함된 causal attention은 다음과 같이 계산됩니다:

 

$$A = \text{softmax} \left( \frac{(R_\theta Q)(R_\theta K)^T + M}{\sqrt{d}} \right), O = AV \quad (3)$$
 

여기서 $R_\theta$block-diagonal rotation matrix이고, $M$은 각 토큰이 자기 자신과 이전 토큰에만 주목하도록 보장하는 lower triangular causal mask입니다.마지막 텍스트 토큰이 visual tokens로부터 받는 attentionlast attention이라고 부르며, 분석을 위해 여러 헤드의 multi-head attention을 평균화합니다.

3.2. Text-visual attention shift

 

VLMstext-visual attention을 분석하기 위해,  먼저 LLaVA-mix665k 데이터에서 $N$개의 이미지-텍스트 쌍을 무작위로 샘플링합니다($N=1,000$, 모델은 LLaVA-1.5-7B 사용).

 

본 연구에서는 FastV를 따라 language model의 layer 2 이후에 text-visual attention에 기반하여 75%의 visual tokenspruning하고, 각 위치의 토큰이 보존되도록 선택되는 빈도와 평균 attention을 시각화했습니다. 그림 2a에 나타난 바와 같이, 선택 빈도와 수신된 attention 사이에는 명확한 양의 상관관계가 있으며, text-visual attention의 뚜렷한 shift 현상이 관찰됩니다. 텍스트 토큰은 rotary position embedding의 장기 쇠퇴(long-term decay) 특성으로 인해 인덱스가 높은 visual tokens에 집중하는 경향이 있습니다. language model의 입력 시퀀스에서 텍스트 토큰은 visual tokens 뒤에 위치하므로, 자신과 물리적으로 더 가까운(즉, 나중에 위치한) visual tokens에 대해 더 높은 attention을 보입니다. 그림 2b에서 이 shiftlanguage model 레이어 인덱스 사이의 관계를 보여주는데, 이러한 positional bias는 첫 번째 레이어부터 나타나며 얕은 레이어에서 더 두드러집니다.

여기서 자연스러운 의문이 생깁니다: 내재적인 positional bias가 존재하는 상황에서, 높은 attention scores가 반드시 더 풍부한 시각 정보를 나타내는가? 이를 확인하기 위해 각 사분위수 위치의 visual tokens를 별도로 보존하고 벤치마크 성능을 평가했습니다(그림 2c). 놀랍게도, 모델의 얕은 레이어에서는 텍스트 attention을 가장 많이 받은 하단 토큰들이 아니라, 이미지의 **중앙 영역(central region)**에 위치한 visual tokens가 가장 높은 성능을 유지했습니다. 이러한 격차는 모델의 10번 레이어를 넘어서며 positional bias가 약해짐에 따라 줄어듭니다. 그러나 추론 효율성을 최적화하려면 language model 내에서 가능한 한 일찍 pruning해야 하므로, text-visual attention shift 현상이 pruning 후 성능에 부정적인 영향을 미친다는 것을 알 수 있습니다.

 

3.3. Text-visual attention dispersion

 

Positional Embedding으로 인한 attention shift의 부정적인 영향을 고려할 때, Positional Embedding 쇠퇴가 없는 text-visual attentionpruning의 근거로 사용하는 아이디어가 떠오를 수 있습니다. 불행히도, Positional Embedding 을 제거한 후에도 attention dispersion(주의 분산)이라 부르는 또 다른 현상을 발견했습니다. 모든 visual tokens에 동일한 Positional Embedding 을 적용한 equi last attention의 누적 분포 곡선을 그림 3a에 그렸습니다. [CLS] attention은 상위 20%의 visual tokens가 약 80%의 주의를 흡수하며 더 집중되어 있는 반면, text-visual attention은 Positional Embedding 제거 여부와 상관없이 더 많은 visual tokens에 널리 퍼져 있습니다.

그림 3b의 밀도 분포에서도 [CLS] attention의 극단적인 피크와 달리, last attention의 밀도는 높은 엔트로피와 낮은 피크를 보이며 균일 분포(uniform distribution)에 가깝게 나타납니다. 고도로 집중된 attention은 모델이 중요한 토큰을 높은 확신으로 식별하고 있음을 의미하는 반면, 균일한 분포는 attention scores에 기반해 중요한 토큰을 선택하기가 더 어렵다는 것을 암시합니다.

이를 검증하기 위해 다양한 attention 기반 pruning 성능을 평가했습니다(그림 3c). Positional Embedding 에 의한 장기 쇠퇴를 제거한 후 모든 벤치마크에서 성능이 향상되어 positional bias의 악영향이 재확인되었습니다. 그러나 주목할 점은, last attention 기반 pruning이 여전히 일부 벤치마크(특히 환각 평가인 POPE)에서 random pruning보다 성능이 뒤처진다는 것입니다.  이 격차를 attention dispersion 때문으로 보고 있으며, 이는 attention만으로는 풍부한 시각 정보를 가진 중요한 토큰을 식별하기 어렵게 만듭니다. 반면 [CLS] attention 기반의 pruning은 모든 벤치마크에서 일관되게 최적의 결과를 얻었으며, 이는 VLMs에서 더 효과적인 토큰 pruning을 위해 visual cues를 활용하게 된 동기가 되었습니다.

 

4. EXPLOITING VISUAL CUES FOR TOKEN PRUNING

 

위에서 분석한 VLMsattention 조사 결과를 바탕으로, 더욱 효과적인 토큰 pruning을 위해 visual cues를 활용하는 VisPruner를 제안합니다. 먼저 섹션 4.1에서 visual encoderattention을 기반으로 풍부한 정보를 가진 중요한 토큰의 일부를 선택합니다. 그다음 섹션 4.2에서 보완책으로서 나머지 토큰들 중 similarity를 기반으로 또 다른 다양한(diverse) 토큰 세트를 유지합니다. 이 두 부분은 visual tokens를 크게 줄인 후에도 모델의 높은 성능을 함께 유지합니다. VisPruner의 전체 설계는 그림 4에 설명되어 있습니다.

 

4.1. Attention-based important tokens

시각적 입력은 본질적으로 중복성(redundancy)을 가지고 있으며, attention score는 토큰의 중요성과 정보 콘텐츠를 평가하는 데 널리 사용되는 지표입니다. 그러나 앞서 분석한 바와 같이, language model 내부의 attentionshiftdispersion 현상으로 인해 실제로 중요한 visual tokens를 효과적으로 식별하지 못합니다.

따라서 VisPruner에서는 visual encoderattention을 활용하여 중요한 토큰을 선택합니다. 구체적으로, 식 (2)에서 계산된 attention matrix $A$에 대해 먼저 attention heads 간의 평균을 구합니다. 그 후, [CLS] token을 포함하는 CLIP과 같은 visual encoders의 경우, $A$의 첫 번째 행을 $a_v \in \mathbb{R}^n$으로 취합니다. 이는 [CLS] token에서 다른 patch tokens로 향하는 attention을 나타내며, 여기서 $n$visual tokens의 길이입니다.

[CLS] token이 없는 모델(예: SigLIP)의 경우, 각 patch token이 다른 모든 토큰으로부터 받는 평균 attention으로 $A$의 행들을 직접 평균냅니다. 시각적 attention $a_v$를 얻은 후, 계산 예산 $R$을 기반으로 동적 임계값 $\tau$를 도출할 수 있습니다:

$$\tau = \min \{ \tau \mid |\{a \in a_v \mid a \ge \tau\}| \le n \times (1 - R \cdot r) \}$$

여기서 $r$은 전체 계산 예산 내에서 중요한 토큰의 비율을 나타내는 hyperparameter입니다. 실제로는 $\tau$$a_v$$(1 - R \cdot r)$ percentile로 설정합니다. 이 임계값을 통해  모든 visual tokens 중에서 풍부한 정보를 가진 중요한 부분 $Z_v^{imp}$를 선택할 수 있습니다:

$$Z_v^{imp} = \{ z_v^i \in Z_v \mid a_v^i \ge \tau \}$$
 

4.2. Similarity-based diverse tokens

visual encodersattention이 종종 전경(foreground) 객체에 더 집중된다는 점을 고려할 때, pruning 후 배경 관련 질문에서 성능이 저하되는 것을 방지하기 위해  보완책으로서 다양한(diverse) 토큰 세트도 유지합니다.

남은 토큰들 중에서, 먼저 각 토큰 사이의 cosine similarity를 다양성(diversity)의 척도로 계산합니다. 그런 다음, 남은 토큰의 수가 $R \cdot (1 - r)$ 요건을 충족할 때까지 다른 토큰과 가장 높은 similarity를 가진 토큰들을 점진적으로 제거하며, 이는 알고리즘 1에 구현되어 있습니다. 이러한 다양한 토큰 $Z_v^{div}$는 입력 이미지의 더 많은 영역을 더 잘 커버하고 중요한 토큰을 보완하여, 입력 이미지의 시각 정보를 가능한 한 많이 보존합니다.

4.3. Inference with visual pruning

attention-based important visual tokens $Z_v^{imp}$와 similarity-based diverse visual tokens $Z_v^{div}$를 결합하고, 이를 이미지 패치 시퀀스의 원래 위치에 따라 정렬합니다. 병합된 visual tokensmultimodal projector $g$를 통과한 후 language instruction $H_q$와 결합되어 language model $f_\phi$에 입력되어 응답 $X_a$를 생성합니다:

$$H_v = g(Z_v^{imp} \oplus Z_v^{div}), \quad X_a = f_\phi(H_v \oplus H_q)$$
 

여기서 $\oplus$concatenation 연산을 나타냅니다. 중요한 토큰과 다양한 토큰 사이의 보완적 설계 덕분에  VisPruner는 높은 reduction ratios에서도 상당한 성능을 유지할 수 있으며 VLMs의 추론 효율성을 크게 향상시킵니다.