P u M er: Pruning and Merging Tokens for Efficient Vision Language Models (2023)

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.

## introduction large-scale vision language models @xcite @xcite @xcite have shown substantial progress on many vision language tasks such as visual question answering, natural language visual reasoning, and visual entailment. however, state-of-the-art language and vision models are memory intensive and computationally expensive because they use multi-layer self-attention between many language and vision input tokens (small image patches) with quadratic complexity. this inefficiency limits highthroughput cloud deployments and makes it infeasible to run on resource-constrained devices. the key source of inefficiency in deep vl models is that these models need to process the entire input image and text tokens over all the layers. our intuition is that the input image contains redundant information and only parts of the image (salient regions, referred by the text) are required and related to the end task. for example, in figure 1 , most of the image content (the four persons, field) is not needed except for the bottom-center soccer region to answer the visual question "what sport are they playing?". this paper advocates using the correlations between image and text modalities to reduce tokens for vl problems. in the vision-only or text-only domains, researchers have shown that reducing image or text tokens can improve the model computational complexity through pruning @xcite @xcite @xcite ) that learns to remove non-salient image or text tokens for a given task; or merging @xcite @xcite that groups semantically similar tokens. using either reduction method in isolation is not sufficient for a vl problem setting since i) salient image tokens are different given different text inputs, ii) pruning alone causes big information loss, hurting the performance, iii) merging tokens irrespective of their modality confuses the vl models since text and image token representations cannot be perfectly aligned to the same semantic space. in this paper, we design a lightweight and effective framework that integrates these token reduction strategies into vl models. we introduce pumer, a token reduction framework that consists of pruning-and-merging operations to gradually reduce image tokens that are not related to text and merge image and text tokens respective to their modality. in particular, we design (i) text-informed image token pruning to remove image tokens that are irrelevant to text and are unimportant to the vl task predictions (removing tokens that describe persons and field for the second question in the figure 1 example); (ii) modality-aware token merging to merge semantically redundant tokens for text and image tokens modality independently (combining the image tokens describing each person for the first question in figure 1 ). we keep the remaining tokens that are neither pruned nor merged. at the core of pumer is a set of lightweight non-parametric token reducers that decide which image tokens are pruned and merged as the vl model forward computation proceeds. to reduce abrupt image information loss and improve computational efficiency, we scatter the token reducers at different cross-modal layers in the vl model and reduce the tokens in a cascaded fashion. fewer tokens are pruned and merged in earlier layers. pumer is easy to train since the token reducers contain no parameters and add little overhead. the training procedure is almost the same as finetuning the original vl models, except that we add a knowledge distillation loss that further reduces the accuracy gap compared to finetuned models. though we focus on inference efficiency, pumer makes vl models run faster for both training and inference because text and image tokens are reduced in the forward computation. we evaluate pumer over two recent vl models vilt @xcite and meter @xcite across five vision language tasks: text-image retrieval tasks (including image-to-text and text-toimage retrieval) @xcite , visual question answering (vqav2; @xcite , natural language visual reasoning (nlvr2; @xcite , and visual entailment (snli-ve; @xcite . compared to baselines, pumer improves the model inference throughput by 1.7x∼2.1x and reduces memory footprint by 38%∼50% with minimal (less than 1%) accuracy loss. our analysis validates that both text-informed image pruning and modality-aware token merging contribute to the token reduction effectiveness of pumer. ## related work token reduction in nlp and vision. prior work in data pruning @xcite @xcite focus on single-modality models by either pruning input text or image alone. dynamicvit @xcite and a-vit (yin et al., 2022) both progressively remove the uninformative content and keep salient regions in the input image. this type of pruning does not apply to language and vision tasks where the salient regions depend on the input text. our work shows different input texts lead to pruning different image regions even for the same input image. power-bert @xcite speeds up the inference of text-based transformers like bert @xcite by removing the input text tokens, which are not the main computation bottlenecks for most vision and language tasks. another line of work seeks to reduce input tokens by merging tokens. spvit @xcite and evit @xcite select uninformative image tokens and combine them into one token. and evit also requires expensive pretraining. groupvit @xcite combines image tokens via cross-attention to find similar objects for semantic segmentation. recently, tome @xcite , tokenlearner @xcite and tokenpooling @xcite combine tokens without pruning and achieves better speedup versus accuracy trade-offs. our method is inspired by token pruning and merging works but integrates them into a token reduction framework suitable for vl models. our key difference is to leverage the relationships between textual and visual tokens to remove and combine tokens. our experiments (section 5) show improvements over these lines of work. efficient vision language models. many techniques have focused on model pruning @xcite @xcite , dynamic computation by early exiting @xcite @xcite @xcite or designing small and efficient vl models @xcite . combining these orthogonal optimizations with our token reduction method could further accelerate the inference in vl models. 12891 3 background and overview vision language models. figure 2 shows the backbone of a vl model consisting of a text encoder, an image encoder, and a cross-modal encoder. the input sentence (e.g. a question or a statement) is first tokenized as text tokens and fed to the text encoder to create contextualized text representations. similarly, the input image is projected into many small image patches, referred to as "image tokens", that are further contextualized by the image encoder. finally, the cross-modal encoder takes the concatenated text and image tokens and fuses information between image and text modalities via transformer-style @xcite cross-attention interactions. text encoder image encoder image text cross-modal encoder text tokens, image tokens image patches text tokens what sport are they playing? for many vl tasks, the number of tokens of the input image is an order of magnitude more than that of the input text -a visual question can have at most a dozen tokens but the associated image consists of a hundred image tokens. for example, for an image with a resolution of 384x384 and a patch size of 16, the number of tokens is (384/16) 2 = 576. token reduction for efficiency. in this paper, we focus on reducing image tokens to improve computational efficiency of the model through pruning and merging. however, naively removing a large percentage of the image tokens inside the crossmodal layers may cause abrupt image information loss, as the vl model is trained to build representations of the full image for the downstream task. for example, if the soccer region in figure 1 gets pruned, the vl model is unlikely to output the answer "soccer" for the question "what sport are they playing?". on the other hand, simply merging image tokens without text guidance can lead to suboptimal performance. for example, merging the image regions of the background field and soccer in figure 1 does not contribute to answering the visual question "how many people are playing?". the next section describes our text-informed token reduction approach. the basic building blocks of pumer are lightweight non-parametric token reducers that reduce image and text tokens in a cascaded manner to mitigate the information loss and improve the computational efficiency of a vl model. ## pumer: text-informed token reduction framework given a vl cross-modal encoder, pumer progressively reduces image tokens going through the cross-modal encoder (depicted in figure 3 ). pumer uses lightweight token reducers with no learnable parameters, adding them in different layers of the cross-modal encoder to predict which image tokens are removed or merged. token reducers. for an n-layer cross-modal encoder, after the first f (@xmath0, a token reducer first removes k% of the image tokens at any layer ℓ between f and n guided by the text information. the tokens removed in layer ℓ are not used in subsequent layers. then the token reducer merges r% and t% of the image and text tokens respectively in layer ℓ. we scatter the token reducers across the cross-modal layers to achieve a better accuracy and efficiency trade-off. intuitively, reducing at early layers in the cross-modal encoder will have higher inference efficiency but may have bigger performance loss and vice versa. we study this trade-off in more detail in section 6.2. the token reduction algorithm is described in algorithm 1. each token reducer consists of two sequential non-parametric modules: first, a text-informed pruner (tip) prunes image tokens that are not related to the accompanying text (section 4.1); second, a modality-aware merger (mam) reduces tokens by merging similar tokens within the image or text modality (section 4.2). these two steps reduce the image and text tokens to benefit the computational efficiency, while not losing the accuracy. note that if we only apply text-informed pruning to the images without merging, to achieve similar efficiency gains, we need to set a larger pruning ratio which will ## conclusion large vision language models have been effective at visual reasoning tasks due to their complex crossmodal interactions between the text and image tokens. these cross-modal interactions are computationally expensive because all image and text tokens are processed in many layers. we introduce a token reduction framework -pumer that uses text-informed image pruning and modality-aware merging techniques to effectively reduce the image and text tokens inside cross-modal layers. pumer progressively removes the redundant image and text information and makes vl models run faster with minimal task performance drop. pumer is