🤖 AI

Learning to Predict Middle-Layer Attention in MLLMs for Visual Token Prunin

This paper proposes Middle-layer Attention Prediction (MAP), a method that dynamically identifies sample-specific optimal attention layers and distills them into a lightweight predictor to prune visual tokens before language model processing, achieving a 3.09x end-to-end speedup while retaining 97.5% of performance on LLaVA-NeXT-7B.

Original authors: Yuyao Sun, Tao Deng, Shuang Li, Deqing Wang, Hao Geng, Minjun Yu

Published 2026-08-10

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuyao Sun, Tao Deng, Shuang Li, Deqing Wang, Hao Geng, Minjun Yu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a super-smart robot how to see and talk about the world at the same time. This is the world of Multimodal Large Language Models (MLLMs). Think of these models as brilliant detectives who can look at a picture and answer questions about it, like "What color is the dog's collar?" or "Why is the person running?" To do this, the robot doesn't just "look" at the image; it breaks the picture down into thousands of tiny digital puzzle pieces called visual tokens. Each token is a small chunk of information about a part of the image.

The problem is, these robots are incredibly hungry. To get a really good look at a picture, they might need to process thousands of these tokens. It's like trying to read a whole encyclopedia to find the answer to a simple question about a cat's hat. It takes a lot of time and energy (computing power), making the robot slow and expensive to run. Scientists have been trying to figure out how to make these robots faster by throwing away the useless puzzle pieces before they even start reading them. This is called visual token pruning. The trick is knowing which pieces to keep and which to toss. If you throw away the wrong piece, the robot gets confused. If you keep too many, it's still slow. The big question has been: How do we know exactly which pieces of the picture matter most for the specific question being asked?

The Problem: One Size Does Not Fit All

The researchers behind this paper, led by Yuyao Sun and Tao Deng, discovered a tricky flaw in how we currently try to help these robots decide what to keep.

Previously, scientists thought that if they looked at the robot's "attention" (a fancy way of saying "what the robot is focusing on") from a specific, middle layer of its brain, they could find the most important picture pieces. It was like saying, "Hey robot, just look at what you were thinking about halfway through your thought process, and keep those picture pieces."

But the authors found that this "one-size-fits-all" approach is broken. Imagine you are looking at a photo of a beach.

If someone asks, "What color is the sand?", the robot's brain lights up on the sand in the middle of its thinking process.
If someone asks, "What color is the sky?", the robot's brain lights up on the sky, but maybe at a different stage of its thinking.

The researchers showed that the "best" layer of the brain to check depends entirely on the specific question. Sometimes the answer is in layer 10, sometimes in layer 20. Using a fixed layer is like trying to use a single key for every door in a building; it works for some, but fails for others.

Furthermore, there was a second, bigger problem. To find out which layer was the "best" one for a specific question, the robot had to actually process the entire image through all those layers first. It's like trying to find the best route to school by driving the whole way there first, just to see which turn was best. By the time you figure out which pieces to keep, you've already wasted all the time and energy you were trying to save!

The Solution: The "Question Contrast" Detective

To solve this, the team proposed a new method called MAP (Middle-layer Attention Prediction). Instead of making the robot do the heavy lifting during the actual conversation, they taught a tiny, super-fast assistant to guess what the robot would have focused on.

Here is how they trained this assistant:

The "Contrast" Trick (QCTS): They came up with a clever way to pick the right "teacher" for each question. They took an image and asked the robot two things:
- Question A: The real question (e.g., "What color is the boat?").
- Question B: A boring, generic question (e.g., "What is in this image?").
They then compared the robot's attention for both questions. The layers where the robot's focus changed the most between the two questions were the ones that really cared about the specific details of the boat. This method, called Question Contrastive Teacher Selection (QCTS), acts like a spotlight, instantly identifying which part of the robot's brain is most excited about the specific question.
The Light-Weight Predictor: Once they knew which layer was the "star" for a specific question, they didn't keep that heavy layer. Instead, they taught a tiny, lightweight predictor to mimic that star layer's behavior. This predictor is so small and fast that it can look at the image and the question before the big robot even starts thinking, and immediately say, "Keep these 5% of the picture pieces, throw away the rest."

The Results: Fast, Light, and Smart

The team tested this new system on several different robot brains (MLLMs) and a wide variety of tricky questions. The results were impressive.

Speed: On one of their test models (LLaVA-NeXT-7B), MAP made the robot 3.09 times faster from start to finish. That's like turning a 10-minute walk into a 3-minute jog.
Efficiency: They managed to throw away 94.4% of the visual tokens (keeping only 5.56% of them) and the robot still got 97.5% of the correct answers it would have gotten with the full picture.
Smart Selection: Unlike other methods that just guess or use a fixed rule, MAP's "diversity" rule made sure that even when keeping very few tokens, it didn't just pick similar-looking pieces. It picked pieces that were different from each other but still relevant to the question, ensuring the robot didn't miss any crucial details.

Why It Matters

This paper doesn't just suggest a cool idea; it provides a working tool that solves a real bottleneck. By proving that the "best" layer changes for every question and that we can predict it without doing the heavy work first, MAP allows these powerful AI models to run much faster on standard computers. It's a bit like giving a librarian a magical index card that tells them exactly which book to pull off the shelf before they even walk into the library, saving them from having to scan every single book on the wall.

The authors show that with this method, we can have our cake and eat it too: we can keep the robot's intelligence high while making it significantly faster and less expensive to run. It suggests that the future of AI isn't just about making bigger, heavier brains, but about teaching them to be smarter about what they pay attention to.

Technical Summary: Learning to Predict Middle-Layer Attention in MLLMs for Visual Token Pruning

Problem Statement

Multimodal Large Language Models (MLLMs) have demonstrated strong capabilities in vision-language tasks but suffer from high computational costs due to the processing of numerous visual tokens. While visual token pruning offers a solution to improve inference efficiency, existing methods face two critical limitations:

Suboptimal Fixed Layers: Recent studies suggest that text-to-vision attention from middle language model layers effectively guides token pruning. However, analysis reveals that the specific layer providing the most question-responsive attention varies significantly across different samples. A fixed-layer strategy is therefore suboptimal.
Inference Efficiency Dilemma: Extracting attention from the appropriate middle layer requires passing numerous visual tokens through several language model layers before pruning can occur. This incurs substantial computation before any tokens are removed, negating potential efficiency gains. Additionally, explicitly materializing attention maps during inference can limit compatibility with optimized attention kernels (e.g., FlashAttention).

Methodology: Middle-layer Attention Prediction (MAP)

To address these issues, the authors propose MAP, a framework that distills sample-specific middle-layer attention into a lightweight predictor operating before the first language model layer. MAP consists of two phases:

1. Offline Training

Question Contrastive Teacher Selection (QCTS): Instead of using a fixed layer, MAP dynamically selects a "teacher" layer for each training sample. It compares the text-to-vision attention distribution generated by the input question against a generic reference question ("What is shown in this image?") on the same image. The layer exhibiting the largest Jensen–Shannon divergence between these two distributions is selected as the teacher, as it best captures regions relevant to the specific input question.
Cross-Modal Attention Distillation: A lightweight predictor is trained to reproduce the selected teacher's attention distribution. The predictor takes multimodal input representations (visual tokens from the vision encoder and text tokens from the embedding layer) available before LLM processing. It uses modality-specific MLPs and projection layers to estimate attention scores, minimizing the Kullback-Leibler (KL) divergence between the predicted distribution and the teacher's distribution. The original MLLM remains frozen during this process.

2. Inference-Time Token Pruning

Prediction and Selection: During inference, the predictor estimates the text-to-vision attention distribution directly from multimodal features. These scores serve as visual token importance.
Diversity-Aware Pruning: To ensure the retained tokens are not only relevant but also complementary, MAP combines the predicted importance scores with a diversity criterion. It greedily selects tokens, starting with the top $K_0$ by importance, and iteratively adds candidates that maximize the orthogonal residual relative to the current subspace of selected embeddings.
Execution: Pruning occurs before the first language model layer. The selected tokens are restored to their original order and combined with question tokens, allowing subsequent processing to utilize only the compact token set. This approach avoids reliance on LLM attention maps, maintaining compatibility with optimized kernels.

Key Contributions

Analysis of Layer Variability: The authors demonstrate that while middle-layer text-to-vision attention is effective for pruning, the optimal layer is sample-dependent, rendering fixed-layer strategies suboptimal.
MAP Framework: They propose a method that distills sample-specific attention into a predictor, enabling efficient pruning before LLM processing without incurring the cost of traversing multiple layers.
QCTS Mechanism: They introduce Question Contrastive Teacher Selection, which identifies the most question-responsive layer for each sample by contrasting attention distributions against a reference question, providing tailored supervision for the predictor.
Empirical Validation: Extensive experiments across three MLLM backbones (LLaVA-1.5-7B, LLaVA-NeXT-7B, Qwen2.5-VL-7B) and ten benchmarks show that MAP achieves a superior accuracy-efficiency trade-off compared to existing pruning methods.

Experimental Results

The paper reports results on LLaVA-NeXT-7B and LLaVA-1.5-7B across ten benchmarks (including GQA, MME, VQA-v2, etc.):

Performance Retention: On LLaVA-NeXT-7B, MAP retains 97.5% of the unpruned model's performance while using only 5.56% of the visual tokens (160 tokens retained from 2880).
Speedup: At this retention ratio, MAP achieves a 7.44× speedup in the LLM prefill stage and a 3.09× end-to-end speedup.
Comparison: MAP consistently outperforms state-of-the-art methods (e.g., LearnPruner, CDPruner, ZOO-Prune) across various token retention ratios. The performance gap widens as the retention ratio decreases, indicating MAP's robustness under aggressive compression.
Ablation Studies:
- Teacher Selection: QCTS outperforms fixed-layer selection (e.g., Layer 14) and multi-layer fusion, confirming the necessity of sample-specific teacher selection.
- Diversity: Diversity-aware selection provides marginal gains at high token budgets but significantly improves performance under aggressive pruning (e.g., 32 tokens), mitigating redundancy among high-scoring tokens.
- Efficiency: MAP reduces KV cache usage by up to 91.39% on LLaVA-NeXT-7B.

Significance

The paper claims that MAP addresses the fundamental trade-off between pruning effectiveness and inference cost. By predicting the optimal attention signal rather than computing it, MAP enables early token pruning that preserves question-relevant visual evidence without the computational overhead of traversing multiple LLM layers. The method is presented as a generalizable solution that improves the efficiency of MLLMs while maintaining high accuracy, particularly in scenarios requiring aggressive token reduction.