PuMer: Pruning and Merging Tokens for Efficient Vision Language Models

TL;DR

PuMer tackles the computational bottleneck of cross-modal attention in vision-language transformers by introducing text-informed image token pruning and modality-aware token merging. It deploys non-parametric token reducers across cross-modal layers to cascade token reductions, achieving up to 2x throughput and up to 51% memory reduction with less than 1% accuracy loss on ViLT and METER across multiple VL tasks. The method requires no additional learnable parameters and integrates with standard fine-tuning workflows, optionally benefiting from a distillation loss. By leveraging text guidance for pruning and modality-aware merging, PuMer complements other efficiency techniques and demonstrates robust, practical gains in VL model deployment.

Abstract

Large-scale vision language (VL) models use Transformers to perform cross-modal interactions between the input text and image. These cross-modal interactions are computationally expensive and memory-intensive due to the quadratic complexity of processing the input image and text. We present PuMer: a token reduction framework that uses text-informed Pruning and modality-aware Merging strategies to progressively reduce the tokens of input image and text, improving model inference speed and reducing memory footprint. PuMer learns to keep salient image tokens related to the input text and merges similar textual and visual tokens by adding lightweight token reducer modules at several cross-modal layers in the VL model. Training PuMer is mostly the same as finetuning the original VL model but faster. Our evaluation for two vision language models on four downstream VL tasks shows PuMer increases inference throughput by up to 2x and reduces memory footprint by over 50% while incurring less than a 1% accuracy drop.

Figures (5)

• Figure 1: PuMer applies token reduction to VL models via pruning and merging. PuMer makes VL models run faster by text-informed image pruning to remove text-irrelevant image tokens and modality-aware merging to compress similar input tokens.
• Figure 2: General architecture of vision language models. The input image is projected into many small image patches ("tokens") that are processed by the image encoder. The cross-modal attention between text and image tokens has quadratic time complexity, which is computationally expensive. Both ViLT and METER models follow this pattern.
• Figure 3: PuMer applies token reducers in the cross-modal layers of a VL model. Each token reducer is non-parametric and uses text-informed pruning and modality-aware merging to reduce image and text tokens.
• Figure 4: Illustration of merging by bipartite matching. In this example, there are 8 tokens, $\mathbf{E}$ consists of token $t_1$, $t_3$, $t_5$ and $t_7$, $\mathbf{O}$ has $t_2, t_4, t_6, t_8$. Assume for $t_1, t_3, t_5, t_7$ in $\mathbf{E}$, the most similar tokens in $\mathbf{O}$ are $t_4, t_4, t_6, t_6$ respectively, and $t_3-t_4, t_7-t_6, t_5-t_6$ are the edges (darker and thicker lines mean larger similarity values) with top-$r^{\prime}$ ($r^{\prime}=3$) most similarity, then we merge $(t_3, t_4)$ into one token $t^m_{34}$, $(t_5, t_7, t_6)$ into one token $t^m_{567}$, and keep $t_1, t_2, t_8$, in this case, we reduce three (3/8=37.5%) tokens.
• Figure 5: Comparing PuMer with DynamicViT and ToMe for the ViLT model on the VQAv2 dataset. Setting different pruning and merging ratios for DynamicViT and ToMe gives different inference throughput and accuracy numbers. Right and top lines are better trade-offs.