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

学习预测多模态大模型（MLLMs）中的中间层注意力以实现视觉 Token 剪枝

Abstract: Multimodal large language models (MLLMs) achieve strong performance across diverse vision-language tasks, but their efficiency is limited by the cost of processing numerous visual tokens. Visual token pruning can reduce this cost, but requires accurate token importance estimates. 摘要： 多模态大模型（MLLMs）在各种视觉-语言任务中表现出色，但其效率受到处理大量视觉 Token 所需成本的限制。视觉 Token 剪枝可以降低这一成本，但前提是需要准确估计 Token 的重要性。

Recent studies have demonstrated that text-to-vision attention from middle language model layers can effectively guide visual token pruning, typically using attention from a predefined middle layer to select the visual tokens to retain. Two problems therefore remain. 近期研究表明，来自语言模型中间层的“文本到视觉”注意力可以有效指导视觉 Token 剪枝，通常做法是利用预定义的中间层注意力来选择保留的视觉 Token。然而，目前仍存在两个问题。

First, our analysis shows that the layer whose attention is most responsive to the question varies substantially across samples, making a fixed layer suboptimal. Second, obtaining attention from the appropriate middle layer requires processing numerous visual tokens through several language model layers, by which point considerable computation has already been spent. 首先，我们的分析显示，对问题响应最敏感的注意力层在不同样本间差异巨大，这使得固定层策略并非最优。其次，从合适的中间层获取注意力需要先通过多个语言模型层处理大量视觉 Token，而此时已经消耗了相当多的计算资源。

To address both problems, we propose Middle-layer Attention Prediction (MAP), which uses Question Contrastive Teacher Selection to identify a sample-specific teacher layer by contrasting attention under the original and reference questions, and distills attention from the selected layer into a lightweight predictor that estimates visual token importance from multi-modal input features. 为了解决这两个问题，我们提出了中间层注意力预测（MAP）方法。该方法利用“问题对比教师选择”（Question Contrastive Teacher Selection）机制，通过对比原始问题和参考问题下的注意力来识别样本特定的教师层，并将所选层的注意力蒸馏到一个轻量级预测器中，从而根据多模态输入特征估计视觉 Token 的重要性。

During inference, MAP combines the predicted importance scores with a diversity criterion to prune visual tokens before the first language model layer. Thus, MAP requires no attention maps for pruning and remains compatible with existing inference acceleration techniques. 在推理过程中，MAP 将预测的重要性分数与多样性准则相结合，在进入第一个语言模型层之前对视觉 Token 进行剪枝。因此，MAP 在剪枝时无需依赖注意力图，并能与现有的推理加速技术保持兼容。

Across ten benchmarks on LLaVA-NeXT-7B, MAP retains 97.5% of the unpruned model performance with only 5.56% of the visual tokens, yielding a 3.09x end-to-end speedup. 在 LLaVA-NeXT-7B 的十项基准测试中，MAP 仅使用 5.56% 的视觉 Token 就保留了未剪枝模型 97.5% 的性能，实现了 3.09 倍的端到端加速。