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

2026-04-26

AI Summary

This work addresses the substantial computational overhead of vision-language models caused by processing long visual sequences. To mitigate this, the authors propose a two-stage adaptive visual token pruning framework. In the first stage, a learnable module removes redundant tokens immediately after the visual encoder to alleviate attention sink issues. In the second stage, task-relevant tokens are preserved within intermediate layers of the large language model based on text-to-vision attention scores. The method retains only 5.5% of the original visual tokens while maintaining approximately 95% of the baseline performance, achieving a 3.2× speedup in inference. This approach significantly reduces computational burden without compromising the balance between accuracy and efficiency.

Technology Category

• Computer Vision: Large Vision Models
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Vision-Language Models
• Token Pruning
• Attention Mechanism
• Computational Efficiency
• Visual Tokens

Innovation

• token pruning
• vision-language models
• attention mechanism
• learnable pruning
• inference acceleration