ArXiv

ReToken: One Token to Improve Vision-Language Models for Visual Retrieval

Authors
Yao Xiao, Reuben Tan, Zhen Zhu...
Categories
cs.CV, cs.AI, cs.LG
arXiv
https://arxiv.org/abs/2607.28627v1
PDF
https://arxiv.org/pdf/2607.28627v1

Brief

ReToken addresses degraded performance of vision–language models under long visual context by learning a single embedding that retrieves a sparse, query-relevant subset of visual tokens from a precomputed KV cache. Trained on a small image–QA dataset, it yields large retrieval improvements (e.g., +13.4 pts for Qwen3VL-8B on Visual Haystacks, +12.4 pts for InternVL3.5, and +8.0 pts zero-shot on LVBench) while remaining compact enough for single-H100 training and inference.

Why it matters

ReToken is a single learnable embedding trained as an explicit retrieval target that selects a sparse set of query-relevant visual tokens from a pre-filled visual KV cache to handle long visual context.

Key details

  • Empirical gains: on Visual Haystacks ReToken improves Qwen3VL-8B by 13.4 points and InternVL3.5 by 12.4 points (>20% relative); on LVBench it transfers zero-shot to long video for an 8.0-point gain with Qwen3VL-8B.
  • ReToken is lightweight—trained on only a small image–QA dataset and designed so both training and long-video inference fit on a single H100; code is released at https://github.com/avaxiao/ReToken.
Source evidence

Abstract

Long visual context poses a challenge for vision-language models: performance degrades as the number of distractors grows, and processing all tokens at once is computationally infeasible under GPU memory constraints. We present ReToken, a single learnable embedding trained as an explicit retrieval target that selects a sparse set of query-relevant visual tokens from a pre-filled visual KV cache. Trained on only a small image-QA dataset, ReToken yields consistent gains across image and video benchmarks: on Visual Haystacks it improves Qwen3VL-8B by 13.4 points and InternVL3.5 by 12.4 points (>20% relative), and on LVBench it transfers zero-shot to long video for an 8.0-point gain with Qwen3VL-8B. Thanks to its lightweight design, both training and long-video inference fit on a single H100. Code is available at: https://github.com/avaxiao/ReToken

Comment: Code: https://github.com/avaxiao/ReToken