Source-linked AI summary

Boosting Multimodal Large Language Models with Visual Tokens Withdrawal for Rapid Inference

Zhihang Lin, Mingbao Lin, Luxi Lin, Rongrong Ji

arXiv:2405.05803v3cs.CVcs.AI

TL;DR

MLLM inference is expensive because large models process additional visual tokens, and existing reduction methods face flexibility, importance-selection, caching, or attention-implementation constraints. VTW withdraws vision tokens after a selected layer, using KL divergence to choose that layer based on early information migration. Across diverse multimodal tasks, it reduces FLOPs by over 40% without compromising performance, although complex tasks may require more migration layers and yield smaller reductions.

  • Problem

    MLLMs incur high inference costs from billions of parameters and visual tokens, while existing token-reduction methods have practical limitations involving flexibility, importance metrics, KV Cache, and Flash-attention.

  • Method

    VTW preserves vision tokens through early layers, withdraws them after a selected layer, and selects the first layer meeting a KL-divergence criterion.

  • Results

    VTW reduces over 40% FLOPs across diverse multimodal benchmarks without compromising performance and produces imperceptible answer differences in multimodal chatbots.

  • Takeaways & Limitations

    Vision tokens can be removed in deep MLLM layers after visual information migrates to subsequent text tokens, enabling rapid inference across the evaluated multimodal tasks.

  • Takeaways & Limitations

    More complex tasks may need additional layers for information migration, making FLOP reductions relatively marginal compared with simpler tasks.

Abstract

from arXiv · show

Multimodal large language models (MLLMs) demand considerable computations for inference due to the extensive parameters and the additional input tokens needed for visual information representation. Herein, we introduce Visual Tokens Withdrawal (VTW), a plug-and-play module to boost MLLMs for rapid inference. Our approach is inspired by two intriguing phenomena we have observed: (1) the attention sink phenomenon that is prevalent in LLMs also persists in MLLMs, suggesting that initial tokens and nearest tokens receive the majority of attention, while middle vision tokens garner minimal attention in deep layers; (2) the presence of information migration, which implies that visual information is transferred to subsequent text tokens within the first few layers of MLLMs. As per our findings, we conclude that vision tokens are unnecessary in the deep layers of MLLMs. Thus, we strategically withdraw them at a certain layer, enabling only text tokens to engage in subsequent layers. To pinpoint the ideal layer for VTW, we initially analyze a limited set of tiny datasets and choose the first layer that meets the Kullback-Leibler divergence criterion. Our VTW approach can cut computational overhead by over 40\% across diverse multimodal tasks while maintaining performance.

Introduction

MLLMs face high inference costs because large models process long sequences of visual tokens, while existing token-reduction methods can lose task-relevant information or limit practical acceleration. VTW withdraws all vision tokens after a selected layer, using observed information migration to reduce computation while preserving performance.

  • Motivation: MLLM inference is costly because model computation grows quadratically with input length, and high-resolution images add many vision tokens.This cost hinders real-time use.
  • Limitations of Existing Methods: Existing token-reduction methods may use fixed vision tokens across questions, causing flexibility problems and a reported 160.4-point accuracy drop on MME.They also rely on importance metrics whose adequacy has limited theoretical support.
  • Limitations of Existing Methods: FastV can limit KV-cache memory benefits because retained vision tokens vary during decoding and can conflict with Flash-attention’s inability to retain the required attention scores.These issues reduce its practical usefulness for broad MLLM acceleration.
  • Visual Tokens Withdrawal: VTW preserves all vision tokens through early layers, then withdraws them at a specific layer so only text tokens participate in deeper computation.The approach avoids selecting individual important vision tokens and is designed to remain compatible with KV Cache.
  • Layer Selection: VTW selects the first layer satisfying a KL-divergence criterion by comparing standard decoding with visual-token-withdrawal decoding across layers on a small dataset subset.This provides an operational rule for choosing the withdrawal layer.
  • Results: Over 40% FLOPs are reduced across eight multimodal benchmarks without compromising performance, and chatbot answers show imperceptible differences.The evaluated tasks include visual question answering, hallucination evaluation, visual reasoning, and video understanding.

Related Work

MLLMs extend language models with visual and other modalities but incur substantial token and computation costs. Related token-reduction work prunes or merges tokens to reduce inference time and memory usage.

  • Multimodal Large Language Models: MLLMs combine language decoders with visual or other modality inputs, but their token counts and quadratic computation create substantial inference costs.LLaVA uses 576 vision tokens for a 336 × 336 image, with higher-resolution inputs increasing the burden.
  • Token Reduction: Token-reduction methods decrease inference time and memory usage by pruning tokens or merging non-critical tokens.Examples include EViT and Evo-ViT for fusion, ToMe for soft matching, and Token Pooling for clustering.

Methodology

The methodology analyzes how attention is distributed across multimodal input tokens, uses information migration to motivate withdrawing vision tokens, and removes them after a selected layer while retaining text-token computation.

  • Architecture: MLLMs combine a vision encoder, cross-modal projector, and pretrained LLM to process visual and textual inputs.The vision encoder produces vision tokens, which the projector maps into the LLM’s text embedding space.
  • Input representation: The input sequence contains system, vision, instruction, and output tokens, with vision tokens comprising most input tokens and dominating computation.Prior token-reduction methods reduce vision tokens but can have limited flexibility, incomplete importance metrics, and KV-cache incompatibility.
  • Attention analysis: Attention analysis shows that middle vision tokens receive minimal deep-layer attention: at least 80% goes to 35 system tokens, while 576 vision tokens receive only 5%.Attention to system tokens rises in initial layers, while attention to vision tokens decreases; increasing output length further shifts attention toward output tokens.
  • Information migration: Information migration occurs because instruction and output tokens attend to preceding vision tokens and absorb visual and textual information through causal self-attention.This explains why later output tokens increasingly attend to text tokens rather than vision tokens.
  • Visual Tokens Withdrawal: The method withdraws all vision tokens at layer K, preserving normal computation before K and leaving only text tokens in deeper layers.The approach avoids selecting individual important tokens and is designed to remain compatible with KV caching.

Experimentation

Across multimodal tasks, VTW generally preserves or improves performance while reducing computation, memory, and latency. Ablations show that withdrawing vision tokens only after sufficient processing and selecting the layer by KL divergence are important for maintaining accuracy.

  • Multimodal Chatbot: VTW preserves chatbot response quality while accelerating responses despite removing vision tokens in deep layers.The VTW chatbot produces correct and similar outcomes to the standard chatbot.
  • Visual Question Answering: VTW outperforms FastV on AI2D and SQA Image while achieving lossless acceleration with nearly half the baseline FLOPs.The method uses visual data in shallow layers and avoids excessive attention to irrelevant information.
  • Visual Reasoning: VTW achieves comparable or superior performance to the baseline and FastV on MMMU Val and MMB EN with fewer FLOPs.
  • Hallucination Evaluation: VTW matches baseline performance on POPE and avoids the hallucination increase observed when FastV removes vision tokens after the second layer.The comparison suggests that premature visual-token removal is harmful in hallucination evaluation.
  • Video Understanding: VTW remains comparable to the baseline and FastV on video question answering tasks while using fewer FLOPs.
  • Efficiency: VTW reduces GPU memory overhead by 35%, cuts FLOPs to nearly half, and lowers per-sample latency to approximately 0.63× the baseline.FastV retains all image-token KV cache during the first forward, producing higher peak GPU memory.
  • Withdrawal-Layer Ablation: Early withdrawal degrades performance, whereas withdrawing vision tokens beyond a specific layer yields performance similar to the baseline.This supports the claim that vision tokens become unnecessary in deep layers.
  • Withdrawal-Layer Selection: When KL divergence converges beyond a threshold, VTW accuracy also converges to the baseline, supporting KL divergence as the withdrawal-layer criterion.

Conclusion

The paper introduces VTW, a plug-and-play strategy that withdraws vision tokens at selected layers to accelerate MLLM inference. Experiments across multimodal tasks and chatbots validate faster inference without compromising performance.

  • Conclusion: VTW withdraws vision tokens at specific MLLM layers based on attention sinks and early information migration.The method is presented as a plug-and-play module for faster inference.
  • Conclusion: Experiments across multimodal tasks and chatbots validate VTW’s ability to accelerate inference without compromising performance.

More Multimodal Chatbot Results

VTW reduces chatbot inference cost while preserving response quality across varied multimodal tasks. The reported examples span question answering, captioning, OCR, and visual reasoning.

  • Over 40% FLOPs are reduced for multimodal chatbots while response quality remains comparable to standard decoding.The reported withdrawal layers are 16 for LLaVA-1.5-7B and LLaVA-NeXT-7B, and 20 for LLaVA-1.5-13B.

More Results on Downstream Task

VTW generalizes to fine-grained downstream tasks by preserving LISA’s segmentation ability on Reasonseg.

  • VTW does not degrade LISA’s segmentation ability, demonstrating generalization to the fine-grained Reasonseg task.The evaluation uses LISA, which produces segmentation masks while retaining multimodal language-model capabilities.

Pseudo Code of Visual Tokens Withdrawal

The VTW procedure searches for a withdrawal layer using a subset of the target dataset, compares inference outcomes with and without VTW, and evaluates the model using the selected layer.

  • VTW samples a subset of the target dataset as input for selecting the withdrawal layer.
  • The selected layer k is used to withdraw vision tokens during final evaluation on the target dataset.
  • The procedure runs inference without VTW and with VTW before calculating their divergence.These inference results are obtained using the main paper’s equations, followed by a divergence calculation.

Results on More Kinds of MLLMs

VTW also applies to MLLMs with architectures unlike the LLaVA series. Across Qwen-VL2 and Intern-VL2, it reduces vision-token inference cost with minimal performance degradation.

  • Nearly 40% lower vision-token inference cost is achieved with minimal performance degradation across Qwen-VL2 and Intern-VL2.These models are open-source MLLMs whose architectures differ substantially from the LLaVA series.

Results on the MME’s subtasks

On MME’s 14 subtasks, VTW generally matches baseline performance while showing marginal degradation on position, poster, and OCR. This supports using text tokens alone in deep layers despite removing vision tokens there.

  • MME subtask results: VTW demonstrates comparable performance to the baseline across most MME subtasks, including existence, count, color, commonsense reasoning, and artwork.MME evaluates perception and cognition across 14 subtasks.
  • MME subtask results: VTW shows marginal performance degradation on the position, poster, and OCR subtasks.
  • MME subtask results: Despite using only text tokens in deep layers, VTW remains on par with normal inference using vision and text tokens throughout the model.
Loading 2405.05803v3…