Token Pruning in Multimodal Large Language Models: Are We Solving the Right Problem?

https://arxiv.org/abs/2502.11501

Token Pruning in Multimodal Large Language Models: Are We Solving the Right Problem?

이번엔 멀티모달이라 그렇게 땡기는건 아니지만 그래도 이쪽 분야의 토큰 프루닝도 알면 좋을테니...

visual token pruning 연구가 정말 중요한 토큰을 잘 고르고 있는지에 대해 놓치고 있고, 기존 attention 기반 pruning이 random, pooling보다 못할 수 있다는 점을 실험적으로 보여줬습니다. 

 

기존 pruning 방법들은 bias를 가져 random이나 pooling보다 못 함을 보여준다. 

최신 token에 attention이 쏠리는 편향을 보여주고 있음 = position bias

지속적으로 기존에 있던 pruning 방법들의 문제를 보여줍니다. 

그리고 토큰 수를 더 줄여도 결국 random의 성능이 가장 높은 것을 보인다. 

vanilla FastV보면 살아 남은 것이 엄청 뭉쳐있는 것을 볼 수 있네요... 

RdfCOCO는 위치 정보가 매우 중요한데 여기서도 기존 prunning 방법의 성능 저하가 매우 크다. 

SparseVLM이나 FastV 계열 방법은 text와 visual token 사이의 attention을 활용함 -> 이 것은 task가 언어 정보에 강하게 의존할 때만 유리함 

Visual Haystack 벤치마크(text와 visual alignment를 확인)를 봤을 때 SparseVLM은 성능을 유지하는 것을 볼 수 있음 

Task 유형	적합한 pruning 방향
일반 visual perception / spatial task	language guidance보다 spatial coverage, redundancy 보존이 중요
text-heavy retrieval / object matching task	text-guided pruning이 효과적
OCR / document / phrase-grounded task	language-conditioned visual selection이 중요할 가능성 큼

visual token에는 두가지 기준이 있음

Importance : 출력 Y를 예측하는데 중요한 token인가. - SQA, TextVQA 같은 knowledge/reasoning task

Redundancy / Uniqueness : 다른 토큰과 중복되지 않고 구조 정보를 보존하는가. - MME, POPE 같은 perception task 

MME, POPE 같은 perception-dominant task에서는 중요도가 낮을 때 redundancy first pruning이 좋음. 반대 케이스에서는 또 다른 프루닝 방식이 좋았음 

FLOPs가 실제 속도 향상을 잘 설명하지 못한다고도 저자는 주장합니다. 

정비례하지 않고, 오히려 역전되는 경우도 있음 

 

문제의식	MLLM은 이미지/비디오 입력에서 수백~수천 개의 visual token을 사용하기 때문에 inference latency와 KV cache 비용이 큼. 기존 연구들은 token pruning으로 이를 줄이려 했지만, 정말 중요한 token을 잘 고르고 있는지, 평가가 올바른지는 충분히 검증하지 않음.
핵심 주장	기존 attention-based visual token pruning은 근본적으로 불완전할 수 있음. 특히 FastV, SparseVLM 같은 방법이 Random selection이나 단순 Pooling보다 낮은 성능을 보이는 경우가 많음.
비교한 방법	FastV, SparseVLM, MustDrop 같은 기존 pruning 방법과 Random token selection, Pooling 같은 단순 baseline을 비교함.
사용 모델	LLaVA-1.5-7B/13B, LLaVA-Next-7B, Qwen2-VL-7B/72B 등 대표 MLLM을 사용함.
사용 데이터셋	GQA, MMBench, MME, POPE, ScienceQA, TextVQA, VizWiz, RefCOCO, Visual Haystack 등 다양한 visual understanding, grounding, retrieval task를 평가함.
주요 결과 1: 단순 baseline의 강함	LLaVA-1.5-7B에서 576개 visual token 중 144개만 남기는 경우, Pooling은 평균 96.4%, Random은 95.0% 성능을 유지했지만, Vanilla FastV는 89.8%, SparseVLM은 93.5%에 그침. 즉, 잘 설계된 pruning보다 단순한 spatially uniform 방식이 더 안정적임.
주요 결과 2: Position Bias	FastV처럼 attention score를 기준으로 token을 고르면 visual sequence 뒤쪽 token이 과도하게 선택됨. 이는 실제 중요도라기보다 위치 편향(position bias) 때문일 수 있음.
해결 실험: Window FastV	저자들은 local window 안에서 token을 고르는 Window FastV를 제안함. 이 방식은 token이 이미지 전체에 더 균일하게 남도록 하며, Vanilla FastV보다 성능 저하가 작음.
주요 결과 3: Spatial task에서 취약	RefCOCO grounding task에서는 대부분의 pruning 방법이 크게 성능 저하를 보임. 특히 spatial localization이 필요한 task에서는 token pruning이 위험하며, spatial uniformity가 중요함.
주요 결과 4: Language guidance의 조건부 효과	text-guided pruning은 항상 좋은 것이 아니라, Visual Haystack처럼 text와 visual object matching이 중요한 task에서 특히 효과적임. 일반 visual perception task에서는 language guidance보다 spatial coverage가 더 중요할 수 있음.
주요 결과 5: Importance vs. Redundancy	pruning 기준은 단순히 “중요한 token”만 보면 안 됨. Perception task에서는 token 간 중복을 줄이면서 구조를 보존하는 redundancy-first가 유리하고, SQA/TextVQA 같은 reasoning task에서는 예측에 중요한 token을 남기는 importance-first가 유리함.
주요 결과 6: FLOPs 평가의 한계	token 수나 FLOPs 감소율은 실제 latency를 잘 설명하지 못함. SparseVLM은 FastV와 유사한 FLOPs를 보이지만 실제 latency는 더 느림. FlashAttention 호환성, pruning layer 위치, selection overhead가 실제 속도에 큰 영향을 줌.
주요 결과 7: Training-aware compression	Qwen2-VL처럼 학습 단계에서 이미 patch merging/token compression을 수행한 모델은 inference-time pruning을 다르게 평가해야 함. training-aware compression을 고려하면 Qwen2-VL은 큰 token reduction에서도 vanilla와 거의 유사한 성능을 유지함.
논문의 결론	MLLM token pruning은 단순히 attention score가 낮은 token을 제거하는 문제가 아님. spatial uniformity, task-specific language dependency, importance-redundancy trade-off, hardware-aware latency, training-aware compression을 함께 고려해야 함.
한계	실험이 LLaVA, LLaVA-Next, Qwen2-VL 중심으로 수행되어 더 다양한 MLLM architecture와 scale에서 검증이 필요함.
연구적 시사점	앞으로 token pruning 논문은 Random/Pooling 같은 단순 baseline을 반드시 포함해야 하며, FLOPs가 아니라 실제 latency와 hardware compatibility를 보고해야 함. 또한 attention score 기반 pruning보다 spatially uniform하고 task-adaptive한 token reduction이 더 유망함.