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

AI summary

Overview

Research area: Computer vision and multimodal machine learning — specifically long-context visual retrieval in vision-language models (VLMs).

Technical level: Intermediate. The paper is readable with a general understanding of transformers and attention, but the core argument depends on distinguishing query-key attention scores from value projections.

Scope: The paper proposes ReToken, a single learnable embedding that is trained as an explicit retrieval target so a frozen VLM can pick a small set of query-relevant images or video frames out of a much larger visual context.

What This Paper Is About

Vision-language models struggle when the visual input is long: only a handful of images or frames matter for any given question, and feeding the whole context into the model can exceed GPU memory. The standard fix is to retrieve the relevant frames first, but the paper shows that the attention scores VLMs naturally produce are a poor retrieval signal — for Qwen3VL-8B they reach only 5.1% average recall@1 across layers on QAEgo4D Test-MC, and only 63.3% recall@1 in the simplest two-image case on Visual Haystacks. The goal is to build a lightweight, learnable retrieval mechanism that works inside the VLM itself, rather than relying on a separate external retriever or on unreliable attention weights.

Key Contributions

1. A diagnostic finding about retrieval space. The authors show that retrieval scores computed in the value space, rather than the conventional query-key space, provide a substantially stronger signal for retrieving visual information. In a controlled two-image setting with a target phrase, value-space pooling raises recall@1 from 65.7 to 78.0 on Qwen3VL and from 78.8 to 83.8 on InternVL3.5.

2. ReToken, a single learnable retrieval token. A single learnable embedding is appended to the question and trained explicitly as a retrieval target, scoring each frame by the cosine similarity between its projected embedding and the frame's mean value vector at the final layer. The token and one projection matrix are the only added parameters, and the VLM is frozen by default.

3. Consistent gains across image and video benchmarks. ReToken improves Qwen3VL-8B by 13.4 points and InternVL3.5 by 12.4 points on Visual Haystacks (over 20% relative gain), and 8.0 points on LVBench with Qwen3VL-8B.

4. Zero-shot image-to-video transfer. Although trained only on multi-image QA data, ReToken transfers to long-video benchmarks without video training, and both training and long-video inference fit on a single H100.

Main Findings

• Attention-based retrieval is weak in VLMs. On QAEgo4D Test-MC, the attention-based retriever (following ReKV) averages 5.1% recall@1 across layers for Qwen3VL-8B, and on Visual Haystacks lands at 63.3% recall@1 even in the two-image case.

• Value-space pooling beats query-key scoring, but only with precise retrieval text. With the full question sentence, value pooling (62.6 on Qwen3VL, 75.6 on InternVL3.5) does not beat query-key scoring (63.3, 78.5). With a target phrase naming the entity, value pooling reaches 78.0 and 83.8 versus 65.7 and 78.8.

• ReToken dominates retrieval baselines at large context. On Visual Haystacks with K=1 and C=50, ReToken reaches 64.7 recall and 72.0 accuracy, versus 1.8 recall / 51.2 accuracy for ReKV, 3.1 / 51.5 for a chain-of-thought phrase retriever, 20.8 / 60.7 for SigLIP2, 58.6 accuracy for the standard full-context pipeline, and 80.7 for the ground-truth-cache oracle.

• Gains grow with context size. Freezing Qwen3VL-8B, ReToken adds +3.9 points at C=2 and +13.4 points at C=50 (72.0 vs 58.6); with InternVL3.5-8B the largest gain is +12.4 points at C=50 (69.7 vs 57.3). At C=1 the two are identical, since retrieval is skipped when C ≤ K.

• Beats external baselines at every context size. ReToken surpasses Gemini-1.5 Pro (62.8 at C=50), SigLIP2 (58.7), MIRAGE (63.6), and REN with DINO·SigLIP2 (68.3) at C=50; InternVL2 reports overflow or out-of-memory errors from C=20 onward.

• Strong zero-shot video transfer at tight budgets. On QAEgo4D Test-MC with K=1, ReToken scores 49.6 versus 42.8 for uniform sampling and 42.6 for ReKV; at K=16 it reaches 60.0 and at K=32, 60.6.

• Benefit scales with video length. ReToken gives no gain on the Short split of Video-MME, because those videos are only about 2 minutes (60 frames at 0.5 FPS), which falls below the retrieval input budget of 100 frames. On LVBench it moves Qwen3VL-8B from 40.6 to 48.6 (+8.0) and Video-MME from 65.1 to 67.1 (+2.0).

• Value scoring is the right choice. At C=50, K=1, training ReToken against average image keys yields 59.4 recall / 70.6 accuracy, while average image values yield 64.7 / 72.0.

• Two-pass inference is essential for ReToken. At C=50, K=1, a single-pass variant scores 50.4 while the two-pass version reaches 72.0. ReKV behaves in the opposite way (51.2 single-pass versus 50.4 two-pass).

• Partial tuning helps images but hurts video. Tuning layers 1–3 gives the best image result (75.0 accuracy at C=50, K=1, versus 72.0 frozen), but on QAEgo4D Test-MC with K=16 partial tuning drops ReToken from 60.0 to 58.0, which the authors attribute to domain shift between image training data and video benchmarks.

• Partial tuning also produces a cleaner KV cache. The gap between re-encoding the ground-truth image and reusing its fused KV cache shrinks from 7.1 points (frozen, C=50) to 4.9 points when tuning layers 1–3 with both generation and retrieval losses.

• Error analysis shows localized evidence wins, dispersed evidence loses. On LVBench with K=100, ReToken gains +15.4 on key information retrieval, +10.5 on entity recognition, +5.5 on reasoning, +3.6 on temporal grounding, +1.9 on event understanding, and loses −5.2 on summarization.

• Modest runtime overhead. Encoding dominates at roughly 14.7 seconds per video (done once). ReToken adds about 0.4 seconds per question over uniform sampling (0.519 vs 0.081 seconds at K=16), with identical answer-stage cost and similar peak GPU memory (66.6 GB vs 65.1 GB retrieval, 65.4 GB answer).

Methodology in Plain English

The authors start by testing whether a VLM's own attention can find the relevant image among distractors. It cannot, so they look for a better place to compare text against images. Inside each attention layer, the keys determine how attention is distributed, while the values carry the content that actually gets passed along. Averaging the value vectors within each frame therefore produces a representation of what that frame contributes — a better target for a text query. The catch is that averaging all question tokens to form the query introduces noise and can erase the advantage, so the query needs to be precise.

Their fix is to learn that precise query. They add one trainable embedding, called ReToken, to the end of the question, plus a single projection matrix. In the final layer, the ReToken output is projected and compared by cosine similarity to each frame's average value vector, producing one score per frame. Those scores are trained against ground-truth relevant/irrelevant labels using a class-balanced binary cross-entropy loss, which averages the positive and negative terms separately so that the many irrelevant images do not swamp the loss. Only the token and projection are updated; the VLM stays frozen in the default setting. A partial-tuning variant also unfreezes the earliest layers and adds a standard next-token generation loss to keep the model's answering ability intact.

At inference the pipeline is two-pass: the video is encoded once into a persistent KV cache; a first pass with the ReToken appended computes which frames are relevant; a second pass generates the answer using only the cached KV of those frames. For long videos, early layers restrict ReToken's attention to the top K′=256 frames before ranking, and older KV states can be offloaded to CPU.

Why This Matters

The work reframes long-context multimodal reasoning as a retrieval problem that can be solved inside the model rather than with a separate retriever bolted on. Because the retriever and the generator share the same forward pass, there is no re-encoding step, and the authors show the approach works with a minimal parameter footprint on a single H100. It also provides a concrete diagnostic — attention scores are not the retrieval signal people assume they are — that other long-context multimodal work can build on. These claims are all as reported in the paper; the paper does not report deployment-scale or production validation.

Real-world applications:

• Long video archives (security footage, sports, or lecture recordings) where a natural-language question must be answered from hours of footage without loading everything into memory.
• Multi-image document or slide-deck question answering, where only one page in a large set is relevant to a query.
• Egocentric and wearable-camera assistants that reason over recorded daily activity at low frame rates.
• Any setting that asks many questions about the same video, since encoding happens once and each additional question costs only about 0.4 extra seconds of retrieval.

Industry relevance: The efficiency profile matters for anyone serving VLMs under GPU memory limits, and the zero-shot image-to-video transfer means a model can be adapted for video use without collecting video training data. The paper states that code is available at https://github.com/avaxiao/ReToken.

Future Directions

• Doing retrieval in the value space at every layer. ReToken is trained and scored only at the final layer, which is why inference needs a two-pass pipeline; extending supervision to intermediate layers is an obvious follow-up.
• Closing the image-to-video gap for partial tuning. Tuning early layers improved image results but reduced video accuracy, which the authors attribute to domain shift; training on video data or mitigating that shift is left open.
• Handling dispersed evidence. ReToken loses 5.2 points on LVBench summarization questions, where relevance is spread across a whole video rather than localized; the paper did not report a method for that regime.
• Understanding the distorted KV cache. The ground-truth-cache accuracy falls from 86.5 at C=2 to 80.7 at C=50 even though the same image is cached, indicating that images absorb distracting information from earlier context; the paper notes a cleaner cache helps Visual Haystacks but may hurt video understanding, which requires connecting adjacent frames.

Target Audience

Researchers and engineers working on long-context vision-language models, multimodal retrieval, and video question answering will get the most from this paper, particularly those who have tried attention-based KV selection such as ReKV and found it lacking. It is also relevant to practitioners who need query-relevant frame selection on a single GPU with limited memory. Readers without a background in transformer attention internals will need to work through Section 3.1 to follow the key-versus-value argument, but the experimental results and error analysis are accessible without that background.

Authors’ 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