CATP: Cross-Attention Token Pruning for Accuracy Preserved Multimodal Model Inference

1. Problem Statement

The paper addresses the computational inefficiency and accuracy degradation associated with Large Multimodal Models (LMMs), specifically focusing on the BLIP-2 architecture.

• Computational Bottleneck: In BLIP-2, the Large Language Model (LLM) Decoder accounts for over 87% of the total parameters (2.7B out of 3.1B), dominating inference time and cost.
• Quadratic Complexity: The attention mechanism in transformers has quadratic time and space complexity relative to input sequence length. Reducing the number of query tokens (input to the LLM) is a viable path to efficiency.
• Failure of Existing Methods: State-of-the-art (SOTA) pruning strategies designed for single-modality models (e.g., magnitude-based pruning or self-attention probability pruning like SpAtten) fail when applied to multimodal models. When tested on BLIP-2, these methods cause a catastrophic drop in accuracy (e.g., dropping from ~52% to <10% on VQA tasks), failing to preserve the semantic relationship between visual and textual inputs.

2. Methodology: Cross-Attention Token Pruning (CATP)

The authors propose CATP, a post-training token pruning method specifically designed for multimodal models. The core insight is that the importance of a query token (text input to the LLM) should be determined by its cross-attention relationship with image tokens (visual input), rather than just self-attention or magnitude.

Core Mechanism

1. Cross-Attention Utilization: CATP leverages the cross-attention layers within the Q-Former (the bridge between the image encoder and LLM in BLIP-2). These layers compute the relevance between image patches and query tokens.
2. The Voting Strategy:
   • Simply summing cross-attention probabilities is ineffective because they are normalized via Softmax (summing to 1).
   • Instead, CATP employs a voting mechanism. For each query token, every image token "votes" based on the strength of their cross-attention probability.
   • Ranking: For a specific image token and attention head, the image token assigns points to the query token based on its rank in cross-attention probability. If a query token has the $n$-th highest probability, it receives $L_0 - n$ points (where $L_0$ is the total number of query tokens).
   • Aggregation: Points are accumulated across all attention heads and selected layers. The query tokens with the lowest total importance scores are pruned.
3. Layer Selection: The method can utilize cross-attention probabilities from all layers or specific layers (e.g., the first or last) to compute importance scores.

Advanced Variants

• Weighted Voting: To further refine pruning, the method incorporates image-token importance. Image tokens with higher self-attention scores (indicating they are more significant visually) are assigned higher weights in the voting process.
• Layer Importance Analysis: Experiments show that not all layers contribute equally. The first and last cross-attention layers provide the most valuable information for pruning decisions, outperforming middle layers.

3. Key Contributions

• Novel Pruning Paradigm: Introduces the first token pruning method for multimodal models that relies on cross-attention rather than self-attention, acknowledging the unique dependency between modalities.
• Voting Algorithm: Proposes a rank-based voting scheme to convert normalized cross-attention probabilities into actionable importance scores, solving the issue of probability summation constraints.
• Accuracy Preservation: Demonstrates that leveraging cross-modal relationships allows for aggressive token reduction without the severe accuracy collapse seen in single-modality pruning methods.
• Layer and Token Weighting: Identifies that specific layers (start/end) and image token weights significantly impact pruning efficacy, offering a roadmap for further optimization.

4. Experimental Results

The method was evaluated on the BLIP-2 (blip2-opt-2.7b) model using a Visual Question Answering (VQA) dataset (10% sampling).

• Baseline Comparison:
  • L2-norm Baseline: Retaining 1/2 of tokens resulted in 9.50% accuracy.
  • Self-Attention Baseline (SpAtten): Retaining 1/2 of tokens resulted in 44.22% accuracy.
  • CATP (All Layers): Retaining 1/2 of tokens achieved 41.14% accuracy.
  • CATP (First Layer Only): Retaining 1/2 of tokens achieved 44.17% accuracy.
• Performance Gains:
  • CATP showed up to 6.6x higher accuracy compared to the L2-norm baseline.
  • CATP showed up to 12.1x higher accuracy compared to the Self-attention baseline in aggressive pruning scenarios (e.g., keeping only 1/8 of tokens).
• Weighted Voting Impact: Incorporating image-token importance weighting further improved accuracy, particularly at aggressive pruning ratios (e.g., keeping 1/8 tokens improved from 23.27% to 29.30%).
• Layer Sensitivity: Using the first or last cross-attention layer alone yielded up to 3.2x higher accuracy compared to using middle layers.

5. Significance

• Bridging the Efficiency-Accuracy Gap: CATP successfully resolves the trade-off between computational efficiency and model precision in multimodal inference. It enables significant reduction in inference costs (by reducing query sequence length) while maintaining high task performance.
• Modality-Aware Pruning: The work highlights that pruning strategies for multimodal models cannot simply be adapted from NLP models; they must explicitly account for the interaction between different modalities (vision and text).
• Scalability: By reducing the computational load of the LLM decoder (the most expensive component), CATP makes large-scale multimodal models more feasible for deployment in resource-constrained environments.

In conclusion, CATP represents a significant step forward in efficient multimodal inference, proving that cross-attention mechanisms are the key to identifying and preserving the most critical information in multimodal inputs.