论文

ReToken: 一个 Token 提升视觉语言模型的视觉检索能力

ReToken: 一个 Token 提升视觉语言模型的视觉检索能力

长视觉上下文 对视觉语言模型构成挑战:随着干扰物数量增加,性能下降,且一次性处理所有 Token 在 GPU 内存限制下计算不可行。 我们提出 ReToken,一个可学习的嵌入作为显式检索目标,从预填充的视觉 KV 缓存中选择一组稀疏的查询相关视觉 Token。仅在小规模图像问答数据集上训练,ReToken 在图像和视频基准上均取得持续改进: - 在 Visual Haystacks 上,Qwen3VL-8B 提升 13.4 点,InternVL3.5 提升 12.4 点(相对 20%); - 在 LVBench 上,零样本迁移到长视频,Qwen3VL-8B 提升 8.0 点。 得益于轻量设计,训练和长视频推理均可在单个 H100 上完成。

论文精读

TL;DR ReToken 用一个可学习的检索 token 从视觉 KV 缓存中稀疏选取相关 token,在图像和视频长上下文任务上带来 10+ 点提升,训练与推理仅需单张 H100。

问题

问题背景 VLM 处理长视觉上下文(如图像集、长视频)时,性能随干扰 token 数量增加而下降,且一次性处理所有 token 面临 GPU 内存瓶颈,因此检索出少量查询相关视觉 token 成为关键。

现有方法局限 当前 VLM 的注意力机制在长视觉上下文中不可靠:文本-视觉注意力与视觉 token 相关性的关联很弱(见原文 Fig. 1),无法直接用于选择 token。现有方案要么暴力处理全部 token,导致计算量过大;要么依赖独立检索模型(如基于 CLIP)预先筛选,增加系统复杂度和延迟,且无法复用预填充的 KV 缓存,每次查询需重新编码整个上下文,推理效率低。

技术挑战与重要性 视觉 token 缺乏显式的语义键,使基于注意力的软选择容易被大量无关 token 淹没,这是视觉检索的内在难点。同时 VLM 输入规模快速增长(长视频、多图像文档),而 GPU 显存增长有限,高效选出稀疏相关 token 成为刚需。该问题的解决可直接推动多模态 RAG、长视频问答、视觉导航等应用,业界对低成本、端到端的视觉检索方案需求迫切。

行业类比 类似大规模监控视频分析中快速定位关键帧,要求在不完全解码视频的情况下精准筛选事件,对内存和延迟要求严苛。

核心洞察

  • ReToken 将视觉检索转化为学习单个可训练的嵌入,显式地选择相关视觉 token,而非依赖 VLM 内部天然不可靠的注意力信号,这从根本上改变了对齐查询和视觉内容的方式。现有方法常试图利用 LLM 的注意力进行 token 筛选或压缩,但实验发现 VLM 中文本-视觉注意力与相关性弱相关,ReToken 通过端到端学习一个明确的检索目标,直接优化对相关 token 的选择,绕开了这一瓶颈,为长上下文理解提供了更可靠的检索机制。
  • ReToken 通过在小型图像 QA 数据集上训练单个 token,即可实现零样本迁移至长视频任务,并以极低的计算成本(单 H100)在长上下文推理中取得显著提升,这挑战了长视频理解必须依赖大规模视频预训练或高计算资源的主流做法。以往的长视频模型往往需要昂贵的视频训练或复杂的 token 压缩机制,而 ReToken 证明了简单的检索 token 在 KV 缓存上的操作即可赋予现成 VLM 强大的长上下文能力,为实际部署提供了轻量且高效的方案。

方法

方法:ReToken(单一检索令牌)

输入:一个视觉-语言模型(VLM)已编码的长视觉序列(图像集或长视频帧),其视觉编码器的键值缓存(KV Cache)被预先计算并缓存;以及一个文本查询(问题)。

核心模块

  1. 可学习检索令牌
    引入一个额外的可训练嵌入向量 [RET],作为专用的检索探针。该令牌不携带任何语义内容,仅被训练用于衡量视觉令牌与查询的相关性。

  2. 稀疏令牌选择
    在每一层或某个特定层,将 [RET] 作为查询(Query),对缓存的视觉键(Key)进行注意力计算,得到每个视觉令牌的得分。然后使用 top-k 筛选出得分最高的 k 个视觉令牌,其余令牌被丢弃(或掩码)。k 为固定超参,通常远小于总视觉令牌数,保证计算效率。

  3. 下游生成
    仅将选中的稀疏视觉令牌(其值向量)连同文本查询一起送入 VLM 的语言解码器,生成最终答案。训练时,模型端到端优化,梯度通过 top-k 操作的直通估计器(straight-through estimator)回传至 [RET] 嵌入。

训练策略

  • 仅使用小型图像问答数据集(如 TextVQA、GQA),冻结 VLM 所有参数,只训练 [RET] 嵌入和可能的少量适配参数。
  • 损失函数为标准语言建模损失(下一个令牌预测),但选择操作迫使 [RET] 学会检索真正支撑答案的视觉证据。
  • 训练后,[RET] 可直接用于长视频等场景,无需重新训练,展现出强大的零样本迁移能力。

输出:针对查询生成的自然语言答案,同时计算开销仅限于选中的稀疏令牌。

差异点:与依赖模型原生注意力(不可靠)或使用额外检索器(需额外计算和标注)的方法不同,ReToken 将检索功能内化为一个可学习的令牌,复用 VLM 已有的键值缓存进行高效稀疏选择,训练成本和推理延迟均极低。

实验

实验设计

ReToken 的训练仅使用一个小规模图像问答数据集(未公开具体名称),训练成本极低,在单张 H100 GPU 上即可完成。评估在 Visual Haystacks 和 LVBench 两个基准上进行,前者包含图像和视频的多跳检索任务,后者为长视频理解。基线模型包括原始的 Qwen3VL-8B 和 InternVL3.5,以及基于注意力信号的检索基线(如思维链式检索)。ReToken 在 Visual Haystacks 上直接测试,并零样本迁移至 LVBench,验证其泛化能力。

关键发现

在 Visual Haystacks 上,ReToken 为 Qwen3VL-8B 带来 +13.4 点(>20% 相对提升),InternVL3.5 获得 +12.4 点 的显著增益。即使从未见过长视频数据,在 LVBench 上也实现了 +8.0 点 的提升。这表明一个可学习的检索 token 能有效从预填充的视觉 KV 缓存中筛选出查询相关的稀疏 token 集合,解决了原生 VLM 注意力信号不可靠的问题。轻量设计使长视频推理不会超出单卡显存限制,实用性强。

基线对比解读

与依赖文本-视觉注意力权重的检索方法相比,ReToken 的优势在于它不是简单利用模型的内部注意力分布(实验显示该分布与相关性弱相关),而是显式训练一个专门的检索目标。这种有监督的稀疏选择策略比无监督的注意力修剪更鲁棒。在同等计算预算下,ReToken 避免了处理全部视觉 token 带来的性能衰退和显存爆炸,同时保持了比直接截断或均匀采样更高的查全率。该方案可作为一种通用插件,与现有 VLM 框架兼容。

行业影响

落地场景

ReToken 的核心能力是在长视觉上下文中高效选取与查询相关的 token,适用于任何需要处理大量图像/视频帧且仅小部分内容与问题相关的场景。典型产品包括:

  • 长视频问答与检索系统(如视频会议摘要、监控视频事件回溯、在线教育视频内容定位)
  • 多图商品搜索与内容审核(如电商平台“以图搜图”、社交媒体虚假视觉信息筛查)
  • 具身智能与自动驾驶数据处理(从长时间录制中快速检索关键传感器帧)

商业价值

ReToken 从两点直接贡献商业回报:

  1. 显著降低计算成本:仅需单张 H100 即可完成训练与长视频推理,避免了全量视觉 token 送入 VLM 造成的 GPU 内存爆炸,让长上下文 VLM 应用不再依赖高成本多卡集群,使中小型业务也有能力部署复杂视觉理解服务。
  2. 提升用户体验与效率:在 Visual Haystacks 上对 Qwen3VL-8B 提升 13.4 点(>20% 相对提升),在长视频基准 LVBench 上零样本迁移提升 8.0 点。更准确的检索意味着更快的回答响应、更少的人工审核介入,直接改善用户留存与任务完成率。

与现有产品/工作流的接口

ReToken 作为一个轻量可学习 token,可以无侵入地插入现有 VLM 推理管线:

  • 前置步骤:对输入图像/视频帧提取视觉特征并填入 KV cache(已有成熟工具如 vLLM、SGLang 支持)。
  • 核心模块:将 ReToken 作为额外的 learnable query 与 KV cache 交互,产生稀疏选择 mask,仅保留相关 token 送入后续 LLM 生成阶段。模块仅增加极少参数,可单独加载或与基础 VLM 联合微调。
  • 集成路径:可通过自定义推理引擎算子或模型加载 adapter 的方式嵌入。例如,在 FastAPI 视觉搜索服务中,对每个请求先并行执行 ReToken 选择,再将浓缩后的 token 序列传给主 VLM,无需改动模型主体架构与训练流程。

具体落地用例

  • 电商平台相似商品检索:用户上传一张服饰图片,平台需从千万商品库中找出同款或近似款。现有管道往往采用两阶段(粗筛 Embedding 检索 + 精排 VLM 对比),但粗筛召回数仍可能数千,全量送入 VLM 耗时且昂贵。可将粗筛结果的图像批量转为 KV cache,用 ReToken 快速筛选出与查询最相关的视觉区域,仅对这部分 token 做精排,单次检索延迟和成本大幅下降。
  • 长视频内容审核:UGC 平台需检测数小时直播回放中是否出现违规片段。传统方案逐帧分析对 GPU 压力极大。引入 ReToken 后,可先将整段视频关键帧编码并预缓存,审核请求(如“是否有暴力画面?”)到来时,用 ReToken 定位高相关帧的 token 子集,仅对这部分进行违规分类,使得实时级响应成为可能,同时保持审核准确度。

局限

  • **训练数据与泛化性受限**:ReToken 仅在一个小型图文问答数据集上进行训练,训练数据的多样性和规模有限,可能限制了检索 token 在更复杂或分布外视觉上下文中的泛化能力。论文未探索在不同领域(如医学影像、遥感)上的迁移效果,也未讨论当视觉内容高度相似或干扰物占比极大时,注意力信号是否会失效。对于需要精确像素级理解或细粒度区域定位的任务,基于 token 的稀疏检索可能因丢失空间信息而表现下降。
  • **对底层 VLM 架构的依赖**:该方法通过提取文本到视觉 token 的注意力作为相关性信号,其有效性高度依赖于底层 VLM 的注意力质量。论文仅在 Qwen3VL-8B 和 InternVL3.5 上验证,未测试其他架构(如 LLaVA 系列、BLIP-2 等),这些模型的跨模态注意力分布可能不同,导致 ReToken 的普适性存疑。此外,如果基础 VLM 本身长上下文能力较弱,检索信号可能不可靠,从而限制该方法的下限。
  • **与重训/压缩方法的对比不足**:ReToken 作为轻量级即插即用模块,避免了对整个 VLM 进行高成本长上下文微调,但实验未与 token 合并(如 ToMe)、关键帧选择或针对检索任务微调 VLM 等强基线进行充分比较。这些方法可能通过端到端训练获得更高的检索精度,而 ReToken 的固定单一 token 设计可能难以捕捉查询与视觉内容间的复杂交互,在需要多步骤推理或比较多个视觉源的场景中可能不如可学习的压缩器或交叉注意力机制。
论文Yao Xiao2026-07-30原文

相关内容