PaperlayerView exact revision on arXiv

Source-linked AI summary

Not All Attention Heads Contribute to Critical Visual Token Selection: Head-Aware Pruning Matters More

Chaofang Ma, Lin Jiang, Carol Jingyi Li, Xingyu Liu, Zeyu Li, Jiang Xu, Wei Zhang

arXiv:2608.25332v1cs.CV

TL;DR

VLM inference becomes costly as visual-token counts grow, and conventional all-head pruning may include noise because critical-token localization is concentrated in a few heads. ProViP progressively removes redundant tokens before and during LLM reasoning using head-aware pruning, achieving competitive task performance and improved inference efficiency, particularly under aggressive pruning.

  • Problem

    Visual-token proliferation increases VLM memory, computation, and latency, while conventional pruning assumes all attention heads contribute equally to token-importance estimation.

  • Method

    ProViP is a training-free progressive framework that removes redundant visual tokens before LLM reasoning and further prunes them during reasoning with head-aware selection.

  • Results

    Experiments show improved inference efficiency while maintaining task performance, particularly under aggressive pruning ratios.

  • Takeaways & Limitations

    Head-aware token-importance estimation supports efficient VLM inference by selectively using heads rather than aggregating all attention heads.

  • Takeaways & Limitations

    ProViP requires intermediate results and explicit attention computation, restricting direct use with closed-source black-box VLMs and partially offsetting acceleration benefits.

Abstract

from arXiv · show

Vision-Language Models (VLMs) have exhibited impressive performance across diverse visual scenarios. However, this success comes at the cost of explosive growth in visual tokens, which imposes substantial memory and computational overhead during inference, ultimately increasing latency. To improve VLM inference efficiency, a typical class of visual token pruning methods estimates token importance by aggregating attention scores across all heads in the pruning layer of the Large Language Model (LLM) backbone and prunes tokens based on aggregated scores. However, in this paper, we reveal a compelling phenomenon: the capability to pinpoint critical visual tokens is concentrated within a small fraction of heads. Aggregation exclusively on these heads can improve task performance. Inspired by this observation, we propose ProViP, a training-free progressive visual token pruning framework. ProViP first removes redundant visual tokens based on the embedding similarity of input tokens before reasoning of the LLM backbone, and then further prunes tokens during reasoning via head-aware pruning. Experiments demonstrate that ProViP delivers outstanding task performance and inference efficiency. For instance, when applied to LLaVA-1.5-7B, ProViP retains 95.9% of the original performance and achieves 1.62x inference speedup under an 88.9% pruning ratio.

1 Introduction

VLMs face rising inference costs from proliferating visual tokens, while conventional pruning aggregates attention across all heads. ProViP instead combines pre-reasoning redundancy removal with head-aware pruning during reasoning.

  • Increasing visual-token counts impose substantial memory and computational overhead, increasing VLM inference latency.
  • Visual token pruning retains task-relevant tokens to reduce memory and computational burdens while preserving performance.
  • All-head attention aggregation implicitly assumes equal head contributions, although critical-token localization is concentrated in a small fraction of heads.
  • Restricting aggregation to visual heads improves shallow-layer pruning performance, while gains diminish in deeper layers.
  • ProViP is a training-free progressive framework that removes redundant tokens before LLM reasoning and applies head-aware pruning during reasoning.
  • Experiments report strong task performance and inference efficiency across multiple benchmarks.

2 Related Work

VLM pipelines convert visual inputs into visual tokens, align them with textual embeddings, and process the combined sequence through an LLM backbone. Related pruning methods exploit redundancy because not all visual tokens are task-relevant.

  • A typical VLM pipeline transforms raw visual inputs into visual tokens that are concatenated with textual tokens for LLM processing.
  • Visual encoders extract semantic features and map them into the LLM input-embedding space for dimensional and semantic alignment.
  • Visual token pruning improves efficiency because not all visual tokens carry task-relevant information.
  • Prior methods remove redundant visual tokens according to differing pruning principles.

3 Observation: Not All Heads Matter

The paper argues that visual-token localization is concentrated in a subset of heads rather than distributed equally across the LLM. Experiments compare head-selection strategies across pruning layers and show the clearest advantage in shallow layers.

  • All-head aggregation assumes equal head contributions, whereas visual-token localization may be concentrated in a small subset of visual heads.
  • Visual heads focus attention on instruction-relevant visual tokens, while non-visual heads disperse attention or select irrelevant regions.
  • The heuristic extracts semantic anchor tokens, computes each head’s cross-modal attention variance, and ranks heads by confidence.
  • Layer-wise experiments on LLaVA-1.5-7B compare All Attention Heads with Visual Heads Only using MME and POPE.
  • A 13.7 F1 score gap favors Visual Heads Only over All Attention Heads when pruning at layer 7 on POPE.
  • The advantage is evident in shallow layers but diminishes in middle and deep layers as performance converges.

4 Method

ProViP progressively prunes visual tokens before and during LLM reasoning, combining text-guided similarity pruning with layer-aware head selection. It retains selected tokens while aggregating residual information into centroids to reduce redundancy and preserve salient content.

  • Overview: ProViP uses sequential critical textual token selection, pre-LLM pruning, and shallow- and middle-layer pruning during LLM reasoning.The framework is training-free and applies pruning progressively across the VLM pipeline.
  • Pruning during the Base LLM: ProViP uses Visual Heads Only for shallow-layer pruning and All Attention Heads for middle-layer pruning because head-selection gains are stronger in shallow layers.Middle-layer pruning uses a broader attention pool after visual information is sufficiently integrated across heads.
  • Critical Textual Token Selection: Critical textual tokens are selected using textual self-similarity to retain semantically important instruction tokens without relying on visual information.This selection is intended to avoid noise from non-content words and preserve anchors even when they lack direct visual correspondence.
  • Pruning before the Base LLM: Before the LLM, ProViP ranks visual tokens by peak recalibrated similarity to critical textual tokens and retains the top N1 tokens.The cross-modal similarity matrix is mean-adjusted to suppress uniformly high, non-discriminative affinities before max-pooling across textual tokens.
  • Pruning before the Base LLM: Pruned visual-token embeddings are assigned to their nearest retained centroids, reducing information loss from potentially misranked task-relevant tokens.The retained tokens serve as cluster centroids that absorb residual embeddings from pruned tokens.
  • Pruning during the Base LLM: Visual heads are identified using a head-level confidence metric based on attention concentration, and the top Nvis heads guide attention-based token pruning.The pruning layer aggregates cross-modal attention from selected visual heads and preserves visual tokens with the top N2 importance scores.

5 Evaluation

Across LLaVA and Qwen2.5-VL, ProViP preserves task performance under visual-token pruning while improving inference efficiency and generalizing across architectures.

  • Main Results: ProViP retains competitive task performance across pruning ratios of 66.7%, 77.8%, and 88.9% on LLaVA-1.5-7B.At 66.7% token removal, performance drops only 0.7%; ProViP also outperforms all baselines across the three ratios.
  • High-Resolution Results: 95.9% average task performance is achieved on LLaVA-NEXT-7B with 320 retained tokens under high-resolution inputs.Performance degradation on POPE is nearly negligible, and ProViP outperforms ApET and DART by 1.5% and 0.7%, respectively.
  • Efficiency Analysis: 126 ms inference latency represents a 38.2% reduction versus vanilla LLaVA-1.5-7B when ProViP retains 64 tokens.Prefilling time is reduced by 62.2% on a single NVIDIA A40 GPU.
  • Generalization: At a 10% pruning ratio, ProViP improves average task performance by 3.2% over PDrop on Qwen2.5-VL-7B.At a 20% pruning ratio, ProViP reaches 87.0% on POPE and surpasses vanilla Qwen2.5-VL-7B.
  • Ablation Study: Visual Heads Only achieves the best average performance across all tested pruning ratios when pre-LLM pruning is disabled.Its advantage over All Attention Heads grows from 0% under mild pruning to 5.3% under more aggressive settings.

6 Conclusion

ProViP addresses noise from all-head attention aggregation with training-free, head-aware progressive visual-token pruning, improving efficiency while maintaining task performance, especially under aggressive pruning.

  • Conclusion: ProViP selectively aggregates a subset of attention heads to estimate visual-token importance more precisely than conventional all-head aggregation.The framework combines this head-aware pruning with progressive token reduction and is training-free.
  • Conclusion: ProViP improves inference efficiency while maintaining task performance, particularly under aggressive pruning ratios.The findings emphasize head-aware token-importance estimation as an important design consideration for efficient VLMs.

Limitations

ProViP requires access to intermediate LLM states and explicit attention computation, constraining deployability and partially offsetting acceleration benefits.

  • Limitations: Access to intermediate inference results restricts ProViP to open-source architectures and prevents direct application to closed-source black-box VLMs.The limitation applies because pruning occurs before and within the base LLM.
  • Limitations: Explicit attention-score computation at pruning layers prevents ProViP from fully bypassing attention computation.This overhead partially offsets acceleration benefits from high-performance implementations such as FlashAttention.

A Benchmarks

The evaluation spans diverse VLM benchmarks, model architectures, and visual-token pruning baselines. These benchmarks assess reasoning, perception, hallucination, and text-reading capabilities, while baselines use redundancy, attention, similarity, or feature-density criteria.

  • Benchmarks: GQA evaluates real-world visual reasoning and compositional question answering using scene-graph-based questions and metrics for consistency and grounding.Its dataset contains 22 million reasoning questions with functional programs representing question semantics.
  • Benchmarks: MMBench provides a bilingual, objective evaluation across diverse VLM abilities through a larger, quality-controlled question collection.It targets scalability, bias, and coarse-grained evaluation limitations.
  • Benchmarks: MME evaluates VLM perception and cognition across 14 subtasks using manually annotated instruction-answer pairs.Its concise instructions support comparisons without interference from extraneous factors.
  • Benchmarks: POPE measures object hallucination by converting open-ended image descriptions into polling-based closed-ended questions.The formulation aims to provide more stable and robust assessment of object identification.
  • VLM Models: The evaluated model families include LLaVA-1.5-7B, LLaVA-NeXT-7B, and Qwen2.5-VL-7B, with differing visual encoders, resolution strategies, and LLM backbones.LLaVA-NeXT uses dynamic high-resolution AnyRes processing, while Qwen2.5-VL supports arbitrary resolutions and aspect ratios.
  • Baselines: The compared pruning baselines reduce visual redundancy using attention, similarity, clustering, token merging, adaptive allocation, or text-guided relevance.Examples include FastV, DART, LLaVA-PruMerge, MustDrop, PDrop, HiRED, VisionZip, and SparseVLM.

D Implementation Details

The experiments use a fixed software and hardware environment and evaluate layer-wise pruning configurations on MME and POPE with LLaVA-1.5-7B.

  • Environment: All experiments run on an NVIDIA A40 GPU with 40GB memory using Python 3.10, PyTorch 2.5.1, and CUDA 12.1.Transformers versions differ across evaluated model families, and LMMs-Eval is used for Qwen2.5-VL-7B.
  • Evaluation Procedure: The layer-wise evaluation uses LLaVA-1.5-7B on MME and POPE under Visual Heads Only and All Attention Heads configurations.The procedure sets each configuration, evaluates each benchmark, and records task performance.
  • Evaluation Procedure: Algorithm 1 evaluates every configuration on each benchmark and records the resulting task performance.These steps operationalize the layer-wise experiments described in Section 3.

E Observation with Various Visual Heads

Experiments show that visual-head-only aggregation is most beneficial in shallow pruning layers, while adding more heads introduces noise and narrows the performance gap. The section also formalizes ProViP’s computational costs across progressive pruning stages.

  • Observation with Various Visual Heads: Visual-head-only aggregation improves performance over all-head aggregation in shallow layers, but its advantage becomes less pronounced in deeper layers.This layer-dependent pattern motivates restricting attention aggregation to visually selective heads where gains are strongest.
  • Observation with Various Visual Heads: As the number of visual heads increases, the performance gap between Visual Heads Only and All Attention Heads narrows in the same shallow pruning layer.The reported interpretation is that additional heads introduce noise into visual-token importance estimation.
  • Computational Cost: Without pruning, per-layer visual-token computation includes MHA and FFN terms parameterized by Nv, d, and an FFN intermediate size of 4d.The total unpruned cost is accumulated across L layers of the base LLM.
  • Computational Cost: ProViP’s cost analysis models visual-token counts decreasing from Nv to N1 before reasoning, then to N2 and N3 at progressive pruning layers Pi and Pj.The analysis separates MHA and FFN computation from the overhead of retaining task-relevant visual tokens.
  • Computational Cost: ProViP’s retention overhead is analyzed separately for pre-LLM pruning and shallow-layer pruning, while critical textual-token selection contributes negligible overhead because |T| is small.The framework’s total cost combines these retention terms with staged MHA and FFN costs.

H More Ablation Results

Ablation results show that combining pre-LLM redundancy removal with progressive in-LLM pruning consistently outperforms using only pre-LLM pruning, especially at higher pruning ratios.

  • Ablation Results: ProViP consistently outperforms pruning only before the LLM backbone across pruning ratios from 66.7% to 97.2%.The comparison is conducted on LLaVA-1.5-7B using six pruning-ratio settings.
  • Ablation Results: The performance gap increases from 2.3% to 8.0% as the pruning ratio rises.This widening gap highlights the effectiveness of combining the framework’s components.
  • Ablation Results: Visualization examples illustrate how ProViP retains task-relevant visual tokens while maintaining task performance and enhancing efficiency.The examples are presented on POPE with LLaVA-1.5-7B.
Loading 2608.25332v1…