VLTP: Vision-Language Guided Token Pruning for Task-Oriented Segmentation

📄 arXiv: 2409.08464v2 📥 PDF

作者: Hanning Chen, Yang Ni, Wenjun Huang, Yezi Liu, SungHeon Jeong, Fei Wen, Nathaniel Bastian, Hugo Latapie, Mohsen Imani

分类: cs.CV

发布日期: 2024-09-13 (更新: 2024-11-26)

备注: Accepted at WACV 2025

💡 一句话要点

提出VLTP，利用视觉-语言引导的token剪枝加速面向任务的ViT分割模型。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Vision Transformer Token剪枝 面向任务分割 视觉-语言引导 多模态学习

📋 核心要点

  1. ViT在分割任务中表现出色，但计算成本高昂，图像token剪枝是有效的加速策略。
  2. VLTP利用视觉-语言引导，预测token与任务的相关性，仅保留相关token进行后续处理。
  3. 实验表明，VLTP能显著降低ViT的计算成本，且性能下降可忽略不计。

📝 摘要（中文）

本文提出了一种视觉-语言引导的token剪枝（VLTP）机制，旨在加速基于Vision Transformer（ViT）的分割模型，尤其是在多模态大型语言模型（MLLM）引导的面向任务的分割（TOS）中。作者认为，ViT无需处理所有图像token，只需处理与推理任务相关的token。VLTP设计了一个新的剪枝解码器，将图像token和视觉-语言引导作为输入，预测每个图像token与任务的相关性。只有高相关性的token才会被传递到ViT的更深层。实验表明，VLTP框架可以在不降低性能的情况下将ViT的计算成本降低约25%，或者以仅1%的性能下降为代价，降低约40%的计算成本。

🔬 方法详解

问题定义：论文旨在解决面向任务的分割（TOS）中，基于ViT的分割模型计算成本过高的问题。现有token剪枝方法在TOS任务中表现不佳，因为TOS中每个图像patch的类别并非预定义，而是依赖于特定的输入任务。

核心思路：论文的核心思路是并非所有图像token都对TOS任务有同等的重要性，只有与任务相关的token才需要被ViT的深层处理。通过视觉-语言引导，可以预测每个token与任务的相关性，从而实现有选择性的token剪枝。

技术框架：VLTP框架包含一个ViT编码器、一个视觉-语言引导模块和一个剪枝解码器。ViT编码器提取图像token特征。视觉-语言引导模块利用多模态大型语言模型（MLLM）提供任务相关的上下文信息。剪枝解码器以图像token特征和视觉-语言引导作为输入，预测每个图像token与任务的相关性得分。最后，根据相关性得分对token进行剪枝，只保留高相关性的token传递到ViT的更深层。

关键创新：VLTP的关键创新在于引入了视觉-语言引导的token剪枝机制。与以往的token剪枝方法不同，VLTP能够根据具体的任务动态地调整剪枝策略，从而更有效地降低计算成本，同时保持甚至提升性能。

关键设计：剪枝解码器是一个轻量级的神经网络，例如MLP或Transformer。损失函数的设计目标是使相关token的得分更高，不相关token的得分更低。具体实现中，可以使用交叉熵损失或Dice损失等。视觉-语言引导模块可以使用预训练的MLLM，例如CLIP或BLIP，提取文本描述的特征向量。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，VLTP框架可以在不降低性能的情况下将ViT的计算成本降低约25%，或者以仅1%的性能下降为代价，降低约40%的计算成本。这些结果表明VLTP在加速ViT模型方面具有显著优势，尤其是在面向任务的分割任务中。

🎯 应用场景

VLTP可应用于各种需要高效率和实时性的视觉任务，例如自动驾驶、机器人导航、医学图像分析等。通过降低计算成本，VLTP使得在资源受限的设备上部署复杂的ViT模型成为可能，从而推动相关技术在更广泛的领域得到应用。

📄 摘要（原文）

Vision Transformers (ViTs) have emerged as the backbone of many segmentation models, consistently achieving state-of-the-art (SOTA) performance. However, their success comes at a significant computational cost. Image token pruning is one of the most effective strategies to address this complexity. However, previous approaches fall short when applied to more complex task-oriented segmentation (TOS), where the class of each image patch is not predefined but dependent on the specific input task. This work introduces the Vision Language Guided Token Pruning (VLTP), a novel token pruning mechanism that can accelerate ViT-based segmentation models, particularly for TOS guided by multi-modal large language model (MLLM). We argue that ViT does not need to process every image token through all of its layers -- only the tokens related to reasoning tasks are necessary. We design a new pruning decoder to take both image tokens and vision-language guidance as input to predict the relevance of each image token to the task. Only image tokens with high relevance are passed to deeper layers of the ViT. Experiments show that the VLTP framework reduces the computational costs of ViT by approximately 25% without performance degradation and by around 40% with only a 1% performance drop. The code associated with this study can be found at this URL.