Nüwa: Mending the Spatial Integrity Torn by VLM Token Pruning

TL;DR

A two-stage token pruning framework that enables efficient feature aggregation while maintaining spatial integrity is proposed that achieves SOTA performance on multiple VQA benchmarks and yields substantial improvements on visual grounding tasks.

Abstract

Vision token pruning has proven to be an effective acceleration technique for the efficient Vision Language Model (VLM). However, existing pruning methods demonstrate excellent performance preservation in visual question answering (VQA) and suffer substantial degradation on visual grounding (VG) tasks. Our analysis of the VLM's processing pipeline reveals that strategies utilizing global semantic similarity and attention scores lose the global spatial reference frame, which is derived from the interactions of tokens' positional information. Motivated by these findings, we propose $\text{Nüwa}$, a two-stage token pruning framework that enables efficient feature aggregation while maintaining spatial integrity. In the first stage, after the vision encoder, we apply three operations, namely separation, alignment, and aggregation, which are inspired by swarm intelligence algorithms to retain information-rich global spatial anchors. In the second stage, within the LLM, we perform text-guided pruning to retain task-relevant visual tokens. Extensive experiments demonstrate that $\text{Nüwa}$ achieves SOTA performance on multiple VQA benchmarks (from 94% to 95%) and yields substantial improvements on visual grounding tasks (from 7% to 47%).

Figures (17)

• Figure 1: Nüwa Performance On VQA and VG tasks, preserving 95% and 47% under 88.9% reduction of vision tokens. (a) Our Nüwa outperforms current efficient VLMs on 10 VQA benchmarks; (b) On 3 visual grounding benchmarks, Nüwa also achieves SOTA results.
• Figure 2: The left panel contrasts our Nüwa framework with prior token pruning methods. (a) Pruning at the vision encoder stage; (b) Text-guided pruning within the LLM; (c) Our two-stage approach: initial spatial-aware pruning via local aggregation that preserves global anchors in the vision encoder, followed by text-guided refinement in the LLM.
• Figure 3: (a) to (d) show different types of attention flows (First row) and gradient-weighted attention flows (Second row), where A-to-B means the degree of attention A pays to B. (e) shows the differences in Last-to-Vision attention maps across different tasks. VLMs exhibit a two-stage visual processing pipeline, with task-independent multimodal interactions in early layers and task-specific processing in middle layers.
• Figure 4: Visualization of VLM's Two-Stage Vision Tokens Processing: (a) Layer-wise Analysis of VAE and OCC Metrics; (b) Layer-wise Instance Heatmap Visualization. Both demonstrate fine-grained feature extraction at the mid-stage.
• Figure 5: Sketch of different Position Embedding Strategies. RPME retains the spatial integrity.