HiRED: Attention-Guided Token Dropping for Efficient Inference of High-Resolution Vision-Language Models

📄 arXiv: 2408.10945v3 📥 PDF

作者: Kazi Hasan Ibn Arif, JinYi Yoon, Dimitrios S. Nikolopoulos, Hans Vandierendonck, Deepu John, Bo Ji

分类: cs.CV, cs.AI

发布日期: 2024-08-20 (更新: 2024-12-25)

备注: Accepted in AAAI 2025

🔗 代码/项目: GITHUB


💡 一句话要点

HiRED:一种用于高效推理高分辨率视觉-语言模型的注意力引导Token丢弃方法

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

关键词: 视觉-语言模型 高分辨率图像 Token丢弃 注意力机制 高效推理 模型优化 资源受限设备

📋 核心要点

  1. 高分辨率VLM生成大量视觉tokens,导致计算成本高昂,尤其是在资源受限的GPU上。
  2. HiRED利用ViT中CLS token的注意力,为图像分区智能分配token预算,实现高效推理。
  3. 实验表明,HiRED在保持或提高准确性的同时,显著提升了吞吐量、降低了延迟并节省了内存。

📝 摘要(中文)

高分辨率视觉-语言模型(VLMs)通过保留详细的图像信息来提高多模态任务的准确性,被广泛应用。然而,由于需要编码高分辨率图像输入的多个分区,这些模型通常会生成过多的视觉tokens。通过多个Transformer网络处理如此大量的视觉tokens带来了巨大的计算挑战,特别是对于资源受限的商用GPU。为了解决这个问题,我们提出了高分辨率早期丢弃(HiRED),这是一种即插即用的token丢弃方法,旨在在固定的token预算内运行。HiRED利用视觉Transformer(ViT)中CLS token的注意力来评估图像分区的视觉内容,并相应地为每个分区分配最佳token预算。然后,选择每个分区内预算中最具信息的视觉tokens,并将其传递给后续的大型语言模型(LLM)。实验表明,与现有的token丢弃方法相比,HiRED实现了卓越的准确性和性能。在LLaVA-Next-7B上,HiRED-20%(即20%的token预算)在NVIDIA TESLA P40(24 GB)上进行单次推理时,token生成吞吐量提高了4.7倍,响应延迟降低了78%,并节省了14%的GPU内存。对于更大的批处理大小(例如4),HiRED-20%通过减少30%的内存使用来防止内存不足错误,同时保持吞吐量和延迟优势。

🔬 方法详解

问题定义:高分辨率视觉-语言模型为了捕捉图像的细节信息,需要处理大量的视觉tokens。然而,这导致了巨大的计算开销,尤其是在资源有限的设备上,成为了一个瓶颈。现有的token丢弃方法可能无法有效地平衡准确性和效率,或者无法充分利用图像的结构信息。

核心思路:HiRED的核心思路是根据图像不同区域的重要性动态地分配token预算。它利用视觉Transformer (ViT) 中CLS token的注意力机制来评估每个图像分区的视觉内容,并据此确定每个分区应该保留多少个token。这种方法能够优先保留信息量大的区域的tokens,从而在有限的token预算下最大化模型的性能。

技术框架:HiRED是一个即插即用的模块,可以集成到现有的视觉-语言模型中。其主要流程如下:1) 输入高分辨率图像,并将其划分为多个分区。2) 使用视觉Transformer (ViT) 提取每个分区的视觉tokens。3) 利用ViT中CLS token的注意力权重来评估每个分区的视觉重要性。4) 根据每个分区的重要性,为其分配一个token预算。5) 在每个分区内,选择最重要的tokens(例如,基于注意力权重)。6) 将选择的tokens传递给后续的大型语言模型 (LLM) 进行处理。

关键创新:HiRED的关键创新在于使用CLS token的注意力来指导token丢弃过程。与传统的token丢弃方法(例如随机丢弃或基于固定阈值的丢弃)相比,HiRED能够更智能地选择要保留的tokens,从而更好地保留图像的关键信息。此外,HiRED的设计使其易于集成到现有的VLM架构中,无需进行大量的模型修改。

关键设计:HiRED的关键设计包括:1) 使用ViT的CLS token注意力作为分区重要性的度量。CLS token通常被认为包含了图像的全局信息,因此其注意力权重可以反映每个分区对整体图像理解的贡献。2) 设计了一个token预算分配策略,根据分区的重要性动态地调整每个分区允许保留的token数量。3) 使用简单的token选择方法(例如,基于注意力权重的排序)来选择每个分区内最重要的tokens。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

HiRED-20%(即20%的token预算)在LLaVA-Next-7B模型上,使用NVIDIA TESLA P40(24 GB)进行单次推理时,token生成吞吐量提高了4.7倍,响应延迟降低了78%,并节省了14%的GPU内存。对于更大的批处理大小(例如4),HiRED-20%通过减少30%的内存使用来防止内存不足错误,同时保持吞吐量和延迟优势。这些结果表明HiRED在提高VLM的效率方面具有显著的优势。

🎯 应用场景

HiRED可应用于各种需要处理高分辨率图像的视觉-语言任务,例如图像描述、视觉问答、图像编辑和机器人导航。通过降低计算成本和内存需求,HiRED使得这些任务能够在资源受限的设备上运行,从而扩展了VLM的应用范围。此外,HiRED还可以用于提高VLM的推理速度,从而改善用户体验。

📄 摘要(原文)

High-resolution Vision-Language Models (VLMs) are widely used in multimodal tasks to enhance accuracy by preserving detailed image information. However, these models often generate an excessive number of visual tokens due to the need to encode multiple partitions of a high-resolution image input. Processing such a large number of visual tokens through multiple transformer networks poses significant computational challenges, particularly for resource-constrained commodity GPUs. To address this challenge, we propose High-Resolution Early Dropping (HiRED), a plug-and-play token-dropping method designed to operate within a fixed token budget. HiRED leverages the attention of CLS token in the vision transformer (ViT) to assess the visual content of the image partitions and allocate an optimal token budget for each partition accordingly. The most informative visual tokens from each partition within the allocated budget are then selected and passed to the subsequent Large Language Model (LLM). We showed that HiRED achieves superior accuracy and performance, compared to existing token-dropping methods. Empirically, HiRED-20% (i.e., a 20% token budget) on LLaVA-Next-7B achieves a 4.7x increase in token generation throughput, reduces response latency by 78%, and saves 14% of GPU memory for single inference on an NVIDIA TESLA P40 (24 GB). For larger batch sizes (e.g., 4), HiRED-20% prevents out-of-memory errors by cutting memory usage by 30%, while preserving throughput and latency benefits. Code - https://github.com/hasanar1f/HiRED