[TopV] Compatible Token Pruning with Inference Time Optimization for Fast and Low-Memory VLM

CVPR 2025

2025   ·   llm-side   ·   efficient-vlm

LLM CVPR 2025

Cheng Yang, Yang Sui, … Xia Hu, Bo Yuan · Rutgers / Rice / Utah / CSU Fullerton

arXiv

한 줄 요약. FastV·LLaVA-PruMerge 같은 LLM-side 가지치기는 attention 점수로 토큰 중요도를 매기는데, 이는 ① 탐욕적 휴리스틱이고 ② FlashAttention과 호환 안 됨(attention 행렬을 명시적으로 계산해야 해 FlashAttention 이득이 사라짐) ③ KV cache와 호환 안 됨(디코딩 스텝마다 자르면 KV cache 크기가 가변→전체 캐시를 저장해야 해 메모리 절감 무효). TopV는 가지치기를 최적화 문제로 정의 — vision-aware cost(Feature Similarity + Relative Spatial Distance + Absolute Central Distance)로 각 토큰의 기여도를 재고 Sinkhorn 알고리즘으로 푼 뒤 복구(recovery). prefilling 단계에서 단 1회 잘라 그대로 디코딩 → FlashAttention·KV cache 모두 호환하며 메모리까지 줄인다. training-free. InternVL2에서 정확도 +1.2%·visual FLOPs 47%↓·동적 메모리 61%↓·추론 2.1×.

배경

LVLM은 시각 토큰이 입력의 대부분을 차지한다(OCR에서 LLaVA-1.5는 576토큰=입력의 87%, InternVL2는 해상도에 따라 256~1792토큰=최대 95%). “시각 토큰은 attention을 덜 받는다”는 관찰로 LLM-side 가지치기가 나왔지만, attention 점수에 의존하는 방식엔 세 약점이 있다.

  • 탐욕적 휴리스틱 — FastV·PruMerge는 “attention 높음 = 중요”라고 가정하지만, 이는 각 토큰의 이후 층 기여도를 정확히 못 잡는다. (VTW는 아예 특정 층 뒤 시각 토큰을 전부 버려 OCR·캡션에 취약.)
  • FlashAttention 비호환 — FlashAttention은 softmax와 V 곱을 융합해 attention 행렬을 내놓지 않는다. attention 점수를 쓰려면 이를 명시적으로 계산해야 해 가속 이득이 사라진다.
  • KV cache 비호환 — FastV처럼 디코딩 스텝마다 자르면 KV cache 크기가 가변이라, prefill의 전체 캐시를 저장해야 해 메모리 절감이 무효가 된다.

attention 점수에 기대지 말고, 토큰의 이후 층 기여도를 직접 최적화하고, prefill에서 한 번만 잘라 FlashAttention·KV cache를 그대로 살릴 수 없을까?

Figure 1. TopV 파이프라인. prefilling 단계의 한 층(L_i)에서 시각 토큰을 source·target으로 두고, vision-aware cost(특징 유사도·상대 공간 거리·절대 중심 거리)로 기여 행렬을 만들어 Sinkhorn으로 최적화 → 중요 토큰을 남기고(복구 포함) 한 번만 가지치기.

핵심 아이디어

① 최적화 기반 중요도 (attention 아님)

토큰 가지치기를 "이후 층에 가장 큰 영향을 주는 토큰을 남기는" 최적화 문제로 정식화. vision-aware cost function = Feature Similarity + Relative Spatial Distance + Absolute Central Distance로 각 source 토큰의 중요도를 측정하고 Sinkhorn 알고리즘으로 기여(contribution) 행렬을 푼다. target 토큰은 Post-LN 출력(Position 3)으로 잡는다.

② prefill 1회 + 복구 → FlashAttn·KV 호환

가지치기를 prefilling에서 단 1회 수행하고 줄인 집합을 디코딩 내내 유지 → KV cache가 고정돼 메모리 절감. attention 행렬을 명시적으로 계산하지 않아 FlashAttention과 호환. 가지친 뒤 token recovery로 정보 손실을 보완한다.

  • 학습 — 추가 학습·fine-tuning 없는 training-free. LLaVA·InternVL2에 바로 적용.

적용·평가

항목 내용
적용 모델 LLaVA-1.5 · InternVL2(2B 등)
데이터셋 AI2D · ScienceQA · MMMU · MMBench · POPE · Nocaps · OCRBench · OK-VQA 등
Task Image Understanding (VQA · OCR · Captioning 포함)
대표 결과 InternVL2: 정확도 +1.2% · visual FLOPs 47%↓ · 동적 메모리 61%↓ · 추론 2.1× / LLaVA: +0.39% · FLOPs 50%↓ · 메모리 49%↓ · 1.68×

결과

정량

Table 3. VTW와 비교(LLaVA-1.5-7B, 비슷한 FLOPs). 디테일이 중요한 Nocaps(103.52 vs 58.37)·OCRBench(31 vs 5.1)·OK-VQA(52.64 vs 39.88)에서 TopV가 크게 앞선다.
  • 디테일 task에 강함 — “특정 층 뒤 전부 버리는” VTW는 Nocaps·OCRBench·OK-VQA에서 붕괴하지만, TopV는 비슷한 FLOPs에서 이를 수 배 끌어올린다(중요 토큰을 정확히 남기기 때문).
  • 효율 — InternVL2에서 visual FLOPs 47%↓·메모리 61%↓·2.1×, LLaVA에서 50%↓·49%↓·1.68×.

정성

Figure 5. 중요 토큰 시각화(빨강=모델이 주목하는 영역). 왼쪽 FastV 대비 오른쪽 TopV가 답에 필요한 영역(글자 등)을 더 정확히 남긴다.
  • 무엇을 남기나(Fig 5). FastV가 고른 토큰과 비교해 TopV는 답에 필요한 영역(예: 이미지 속 글자)에 토큰을 더 집중시킨다 — 최적화 기반 중요도의 효과.
  • 메모리 거동(Fig 3). FastV는 디코딩마다 KV cache가 들쭉날쭉(톱니) 늘지만, TopV는 prefill 1회 가지치기라 메모리가 낮고 평탄하게 유지된다.

한 줄 정리 & 의의

  • “attention 점수” 대신 “최적화”로 자르고, prefill에서 한 번만. vision-aware cost + Sinkhorn으로 중요도를 풀어, FlashAttention·KV cache를 깨지 않고 메모리까지 줄이는 LLM-side 가지치기.
  • 차별점. FastV(attention·스텝마다·FlashAttn/KV 깨짐)·VTW(전부 버림·OCR 취약)·PruMerge와 달리, 최적화 기반 + prefill 1회로 호환성과 메모리를 동시에 잡는다.
  • 위치. LLM — LLM prefilling 단계에서 줄인다. → Efficient VLM 개요