PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yao Xiao, Reuben Tan, Zhen Zhu, Yuqun Wu, Jianfeng Gao, Derek Hoiem

arXiv:2607.28627v1cs.CVcs.AIcs.LG

TL;DR

Long visual contexts make it difficult and often infeasible for vision-language models to identify the few relevant visual tokens among many distractors. ReToken uses one learnable retrieval token to select query-relevant tokens, improving image and long-video benchmark performance, including an 8.0-point gain on LVBench with Qwen3VL-8B.

  • Problem

    Vision-language models need to retrieve a small relevant subset from long visual contexts, but conventional text–visual attention provides an unreliable relevance signal and full-context processing can exceed GPU memory.

  • Method

    ReToken adds one learnable retrieval token that scores visual frames using projected mean value vectors, enabling sparse selection from a pre-filled visual KV cache.

  • Results

    ReToken produces consistent gains across image and video benchmarks, including an 8.0-point improvement on LVBench with Qwen3VL-8B after zero-shot transfer from multi-image QA training.

  • Takeaways & Limitations

    ReToken offers a lightweight practical approach to scalable long-context multimodal reasoning, with training and long-video inference fitting on a single H100.

  • Takeaways & Limitations

    ReToken requires two forward passes and slightly more memory and response time, while its multi-image QA training may limit modeling of temporal structure in videos.

Abstract

from arXiv · show

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

1 Introduction

Long visual contexts reduce VLM reliability and can exceed GPU memory, making selective retrieval of relevant visual tokens necessary. ReToken addresses this with a single learnable retrieval embedding and improves image and long-video performance while remaining lightweight.

  • Value-space retrieval: Recall@1 rises from 65.7 to 78.0 on Qwen3VL and from 78.8 to 83.8 on InternVL3.5 when matching target phrases against average visual value projections rather than keys.The controlled setting uses two images; the finding motivates retrieval in value space.
  • Method: RETOKEN is a single learnable embedding appended to the question that scores frames by cosine similarity between projected embeddings and final-layer mean value vectors.Training uses class-balanced binary cross-entropy against ground-truth relevance labels, with the VLM frozen by default.
  • Results: 13.4 points improve Qwen3VL-8B and 12.4 points improve InternVL3.5 on Visual Haystacks, corresponding to over 20% relative gain.These gains are reported across the two listed VLM baselines.
  • Results: 8.0 points improve Qwen3VL-8B on LVBench after RETOKEN transfers zero-shot from multi-image QA to long-video understanding.The result demonstrates transfer across visual modalities without long-video training.

2 Related Work

Prior visual-retrieval systems use external embedding retrievers, while long-context methods compress or retrieve visual content using query-independent or query-dependent selection. ReToken is most closely related to learnable-token modules that aggregate visual features for vision-language models.

  • Visual Retrieval: Visual retrieval commonly aligns visual and language spaces with dual encoders or adapts a VLM into a universal embedder.Examples include CLIP, Perception Encoder, SigLIP2, E5-V, and LamRA.
  • Visual Retrieval: External retrievers select relevant images before the VLM re-encodes them for answer generation.
  • Long Context Understanding: Memory and token-compression methods bound long visual contexts independently of the query by maintaining fixed-size state or shortening the visual prefix.They merge, evict, prune, merge, or summarize visual tokens to reduce attention and KV-cache cost.
  • Long Context Understanding: Retrieval-based methods select a query-relevant subset at inference time to preserve fine-grained evidence.Goldfish retrieves top-K video clips by caption-query similarity, while Video-RAG adds visually aligned auxiliary text from ASR.
  • Learnable Tokens for Visual Aggregation: Learnable-token methods use a small set of queries to aggregate or retrieve visual content into fixed-size representations for a frozen LLM.Q-Former, Perceiver Resampler, and text-conditioned InstructBLIP exemplify this design.

3 Approach

ReToken retrieves question-relevant visual tokens using a single learnable retrieval embedding and value-space frame representations. It enables sparse, retrieve-then-answer inference over cached visual features while keeping the VLM frozen by default.

  • Problem Setup: Long videos can make full-context inference infeasible, so ReToken retrieves K relevant frames from F frames before answer generation.The selected frames provide visual tokens sufficient to answer the question.
  • Limitations of Attention-Based Retrieval: Attention-based retrieval can misidentify relevant tokens because attention is optimized for next-token prediction rather than retrieval.Typical VLM training data rarely requires selecting among partially relevant visual inputs, worsening this mismatch.
  • Why Values Carry Retrieval Signal: Value projections provide a stronger retrieval signal because they encode content propagated to attending tokens, whereas query-key scores determine only aggregation weights.Pooling values within each frame makes the representation more sensitive to retrieval text than key features.
  • ReToken Retrieval Score: ReToken appends one learnable embedding to the question and scores frames by cosine similarity between its projected output and pooled frame-value features.The only added parameters are the retrieval embedding Xr and a final-layer projection Wr.
  • Inference: At inference, ReToken encodes the video once, retrieves visual KV-cache subsets for each question, and generates answers using about K ¨ M{F visual tokens instead of all M.This retrieve-then-answer design is especially efficient when multiple questions share the same video.

4 Experiments

Experiments show that ReToken improves visual and long-video retrieval, especially under tight retrieval budgets and growing context lengths, while adding modest per-question overhead. Its gains are strongest when relevant evidence is localized, whereas performance is weaker for temporally structured reasoning.

  • Retrieval performance: RETOKEN substantially outperforms all baselines and retains most of the GT Cache upper-bound accuracy at C = 50.This advantage holds across the reported retrieval comparisons, while other methods fail at large context size.
  • KV-cache behavior: 86.5 at C = 2 versus 80.7 at C = 50 shows that GT Cache accuracy degrades with more context despite using the same ground-truth relevant image.Previous irrelevant images contaminate image tokens through attention, while partial tuning encourages cleaner stored KV-cache representations.
  • Video retrieval: 6.8 points is RETOKEN’s improvement over uniform sampling on QAEgo4DTest-MC with K = 1.Its advantage is most pronounced under tight retrieval budgets; the lead narrows as K grows because neighboring video frames provide complementary context.
  • Long-video understanding: 8.0 points is RETOKEN’s improvement on LVBench, with benefits increasing as video length grows.It provides no gain on Video-MME’s Short split because those videos fall below the retrieval input budget and do not trigger retrieval.
  • Runtime and memory: 14.7 seconds per video is the approximate encoding cost, while RETOKEN adds roughly 0.4 seconds to per-question retrieval and answering.Encoding creates a persistent KV cache shared across questions and can be performed as preprocessing.
  • Error analysis: 15.4 for key information retrieval and 10.5 for entity recognition are RETOKEN’s strongest LVBench task-type gains.Gains shrink when relevance depends on cross-frame or temporal structure rather than localized per-frame content.

5 Conclusion

ReToken improves visual retrieval in pretrained vision-language models by exploiting stronger text-aligned signals in the value space, with gains transferring zero-shot from multi-image training to images and long videos. Its limitations include two-pass computation and limited training data, while future work targets temporally compositional and displacement-aware retrieval.

  • Conclusion: ReToken improves visual retrieval by leveraging stronger text-aligned signals in pretrained VLM value spaces.The method is a learnable token trained as an explicit retrieval target.
  • Conclusion: ReToken delivers significant gains on image and long-video benchmarks despite training only on multi-image data.The approach transfers zero-shot from images to videos.
  • Limitations: ReToken requires a two-pass forward and attends to more visual context in early layers, slightly increasing memory requirements and response time.Both added costs arise from the retrieval design.
  • Limitations: Training data is limited to multi-image QA, and video training may improve temporal structure capture and video understanding.The limitation concerns the training distribution rather than the inference benchmarks.
  • Future Work: Future retrieval could score consecutive frame sets or model temporally offset queries instead of matching each query independently to per-frame value means.These directions address information emerging from temporal combinations and temporal displacement.

A Prompt Templates

The prompt templates distinguish all visual, retrieved visual, question, retrieval, and answer tokens, then use separate formats for retrieval-trigger training, QA preservation, and two-stage inference. Retrieval first processes all visual tokens with a <Retrieval> token, after which answering uses only retrieved visual tokens.

  • A Prompt Templates: RETOKEN prompt templates color-code all vision, retrieved vision, question, <Retrieval>, and answer tokens.These labels identify the content-bearing tokens used during training and inference.
  • A.1 Original Prompt Template: The original prompt feeds All Vision Tokens alongside Question Tokens to the vision-language model.Its format is <system><|vision_start|>All Vision Tokens<|vision_end|>Question Tokens<|im_end|><|im_start|>assistant:.
  • A.2 Training Prompt Template: Training Prompt 1 appends <Retrieval> after Question Tokens to teach when retrieval should be invoked.This format is used in both frozen and partially fine-tuned settings.
  • A.2 Training Prompt Template: Training Prompt 2 preserves original QA abilities and is added only when the LLM is partially fine-tuned.It retains the original prompt structure with All Vision Tokens and Question Tokens.
  • A.3 Inference Pipeline: Inference uses a two-stage retrieve-then-answer pipeline.Stage 1 retrieves query-relevant information; Stage 2 generates the final answer from Retrieved Vision Tokens.
  • A.3 Inference Pipeline: At most K1 frames are attended during Stage 1, where <Retrieval> aggregates query-relevant information at the final layer.Stage 1 uses All Vision Tokens and Question Tokens before triggering retrieval.
  • A.3 Inference Pipeline: Stage 2 supplies only Retrieved Vision Tokens to generate the final answer.Its prompt replaces All Vision Tokens with Retrieved Vision Tokens before answer generation.

B Additional Details · B.1 Training Data Filtering

ReToken training data is filtered because pretrained Qwen3VL can exploit memorized QA pairs instead of retrieving target images. Three filters produce a 70,686-example dataset while leaving evaluation data untouched and using only frozen-Qwen3VL-8B training signals.

  • B.1 Training Data Filtering: Qwen3VL often has lower generation loss with distractor images than with the target image alone, indicating memorization can bypass retrieval.The original MIRAGE set includes VQA datasets many of which Qwen3VL saw during pretraining.
  • B.1 Training Data Filtering: The MIRAGE-augmented dataset combines images from different QA pairs into multi-image inputs and is filtered to restore the intended retrieval task.The filtering addresses examples where memorized QA pairs undermine retrieval.
  • B.1 Training Data Filtering: The memorization filter discards examples whose target-only generation loss exceeds the loss for the full image set.Such examples can be answered without retrieving the target image.
  • B.1 Training Data Filtering: The difficulty filter ranks remaining examples by average attention between question query tokens and target-image key tokens, retaining those with higher scores.A controlled experiment finds that easier retrieval examples have higher attention scores on the target image.
  • B.1 Training Data Filtering: 70,686 examples remain after filtering in the final training dataset.Filtering is applied only to the MIRAGE training split.
  • B.1 Training Data Filtering: Visual Haystacks evaluation and video datasets remain untouched, with filtering signals computed only on training examples using frozen Qwen3VL-8B.The procedure has no access to evaluation labels or examples.

B.2 Chain-of-Thought (CoT) Retrieval Baseline

The CoT retrieval baseline uses query rewriting to convert a question into a short search phrase for locating visual evidence. Its prompt requires returning only a search phrase of no more than five words.

  • The baseline is a CoT prompt template for query rewriting.
  • The prompt asks what phrase to search for to find visual evidence needed to answer the question.It requires responding with the search phrase only, limited to no more than 5 words.

B.3 Key or Value? Training Comparison · C Design Choice

With the VLM frozen and only ReToken trained, the Value variant shows a clear training advantage: faster retrieval-loss convergence and a larger gap between relevant and irrelevant images.

  • B.3 Key or Value? Training Comparison: Value converges much faster in retrieval loss than the alternative variant.Figure 6 compares training dynamics for the retrieval objective.
  • B.3 Key or Value? Training Comparison: Value learns a substantially larger retrieval score gap between relevant and irrelevant images.The score gap is measured across training steps.
  • C Design Choice: The training comparison evaluates retrieval loss as one criterion for choosing between Key and Value.The Value variant is favored by its faster retrieval-loss convergence.
  • C Design Choice: The training comparison also evaluates the retrieval score gap between relevant and irrelevant images.The Value variant is favored by learning a larger separation on this measure.
  • C Design Choice: Together, the reported training curves identify Value as the preferred design choice.Both retrieval-loss convergence and relevant-versus-irrelevant score separation favor Value.
  • B.3 Key or Value? Training Comparison: Value yields a clear advantage when the VLM is frozen and only ReToken is trained.The comparison tracks retrieval loss and retrieval score gap across training steps.

C.1 Multiple Tokens

Training three ReTokens sequentially performs on par overall with a single token, likely because all three embeddings receive supervision through one averaged retrieval score and do not specialize. Specializing tokens by domain is left for future work.

  • Training and scoring: Three ReTokens are appended sequentially, and their retrieval logits are averaged after cosine similarity with the averaged image value feature.The retrieval score uses each projected ReToken against the averaged image value feature.
  • Results: Performance is on par overall with a single token, with small fluctuations in both directions.The comparison is reported for Qwen3VL-8B in Table 13.
  • Analysis: Because all three embeddings are supervised through one averaged score, the multi-token design does not encourage specialization.The tokens therefore likely converge on similar solutions.
  • Future direction: Empowering different tokens to specialize for distinct domains is suggested as a promising direction beyond this work’s scope.

C.2 Query Composition or visual Summarization

ReToken can summarize both visual inputs and queries because it attends to images and query tokens. Ablations show that allowing image attention improves retrieval over skipping visual tokens, although the latter retains decent performance.

  • Query Composition or visual Summarization: ReToken attends to both images and queries, enabling visual and query summarization.It is appended to the input and can attend to both modalities.
  • Query Composition or visual Summarization: The ablation restricts ReToken to query-token attention while caching vision-token keys and values in a separate stream.A second stream processes only the query and ReToken, preventing direct visual-token attention.
  • Query Composition or visual Summarization: Allowing ReToken to attend to images yields better retrieval than skipping visual tokens entirely.Skipping visual tokens still produces decent performance, but underperforms image-attending ReToken; the VLM is frozen.
Loading 2607.28627v1…