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

Core Idea§

ReToken introduces a single learnable embedding $r \in \mathbb{R}^d$ that acts as an explicit retrieval target for vision-language models (VLMs). During inference, the visual encoder pre-fills all image/video tokens and caches their keys and values. Given a text query, ReToken scores each cached visual key $k_i$ via an attention-style relevance function:

$$s_i = \frac{r^\top k_i}{\sqrt{d}}$$

The top-$k$ most relevant visual tokens are then selected, and only their KV pairs are fed into the language model alongside the text query. This sparse selection drastically reduces the context length, making long-video inference feasible on a single H100 while also improving accuracy by removing distractors.

Training§

Despite its effectiveness, ReToken is trained on only a small image-QA dataset. The model is trained end-to-end by backpropagating through the sparse selection using a Gumbel-Softmax relaxation of the top-$k$ operation:

$$\mathcal{L} = -\log p_\theta(y \mid x_q, \{v_i\}_{i \in \mathcal{I}})$$

where $\mathcal{I}$ is the sampled index set. The retrieval token $r$ is the only new parameter; the underlying VLM remains frozen or only lightly fine-tuned.

Implementation Sketch§

import torch
import torch.nn.functional as F

class ReTokenSelector(nn.Module):
    def __init__(self, d_model, k):
        super().__init__()
        self.retoken = nn.Parameter(torch.randn(1, 1, d_model))
        self.k = k

    def forward(self, visual_keys, visual_vals):
        # visual_keys: (B, N, d)
        scores = torch.matmul(self.retoken, visual_keys.transpose(-2, -1)) / math.sqrt(d)
        scores = scores.squeeze(1)  # (B, N)
        # Hard top-k for inference; use Gumbel-Softmax during training
        topk_idx = scores.topk(self.k, dim=-1).indices
        selected_keys = torch.gather(visual_keys, 1, topk_idx.unsqueeze(-1).expand(-1, -1, d))
        selected_vals = torch.gather(visual_vals, 1, topk_idx.unsqueeze(-1).expand(-1, -1, d))
        return selected_keys, selected_vals

Empirical Gains§

On the Visual Haystacks benchmark, ReToken improves Qwen3VL-8B by 13.4 points and InternVL3.5 by 12.4 points (over 20% relative improvement). It also transfers zero-shot to long-video retrieval in LVBench, gaining 8.0 points with Qwen3VL-8B. The method is lightweight—both training and long-video inference fit on a single H100—making it a practical drop-in enhancement for existing VLMs.

Why It Works§

The key insight is that retrieval and generation can be decoupled: instead of forcing the language model to attend over a massive KV cache, a dedicated token learns to identify where the answer is. This reduces the problem to sparse attention over a handful of relevant features, mitigating the distractors that confuse standard VLMs in long-context settings. The approach is agnostic to the specific VLM backbone and requires no architectural changes beyond the added token.