LearnPruner: Rethinking Attention-based Token Pruning in Vision Language Models

TL;DR

LearnPruner introduces a two-stage token pruning method for vision-language models, effectively reducing computational load while maintaining high performance by leveraging insights into attention mechanisms.

Contribution

The paper proposes LearnPruner, a novel token pruning framework that improves efficiency in vision-language models by analyzing and utilizing attention biases in both vision encoders and LLMs.

Findings

Preserves approximately 95% of original performance with only 5.5% of vision tokens.

Achieves 3.2× inference acceleration.

Analyzes attention sink and bias issues in vision encoders and LLMs.

Abstract

Vision-Language Models (VLMs) have recently demonstrated remarkable capabilities in visual understanding and reasoning, but they also impose significant computational burdens due to long visual sequence inputs. Recent works address this issue by pruning unimportant visual tokens, achieving substantial computational reduction while maintaining model performance. The core of token pruning lies in determining token importance, with current approaches primarily relying on attention scores from vision encoders or Large Language Models (LLMs). In this paper, we analyze the effectiveness of attention mechanisms in both vision encoders and LLMs. We find that vision encoders suffer from attention sink, leading to poor focus on informative foreground regions, while in LLMs, although prior studies have identified attention bias toward token positions, text-to-vision attention demonstrates resistance to this bias and enables effective pruning guidance in middle layers. Based on these observations, we propose LearnPruner, a two-stage token pruning framework that first removes redundant vision tokens via a learnable pruning module after the vision encoder, then retains only task-relevant tokens in the LLM's middle layer. Experimental results show that our LearnPruner can preserve approximately 95% of the original performance while using only 5.5% of vision tokens, and achieve 3.2$\times$ inference acceleration, demonstrating a superior accuracy-efficiency trade-off.