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

Authors: Yao Xiao¹, Reuben Tan², Zhen Zhu^1,3, Yuqun Wu¹, Jianfeng Gao², Derek Hoiem^1
1University of Illinois at Urbana-Champaign, ²Microsoft Research, ³Google DeepMind

Attention tells you where a VLM looks, so shouldn't it tell you which frames matter? It doesn't: the retrieval signal lives in the value space, and a single learned token is enough to read it out.

• Cheap to run: the video is encoded once into a persistent KV cache, so each new question pays only for retrieval and decoding.
• Cheap to train: one token, one projection, 1 H100 is enough.
• Results promising: over 20% relative gain for both Qwen3VL-8B and InternVL3.5-8B.

git clone https://github.com/avaxiao/ReToken.git
cd ReToken
conda create -n retoken python=3.11 -y
conda activate retoken

pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124

pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl

pip install -r requirements.txt
conda install -c conda-forge ffmpeg -y
pip install -e .

Versions are pinned deliberately: torch==2.6.0/cu124 is what the flash-attn wheel above is built against, and transformers==4.57.3 satisfies both backbones at once (Qwen3-VL needs >=4.57, InternVL3.5 remote code needs <5). Installing unpinned pulls a much newer torch and the flash-attn build fails.

Follow the Visual Haystacks instructions to obtain the two pieces:

• the VHs QA files — visual_haystack_<N>.json for each haystack size N
• the COCO images the haystacks are drawn from (train2017 / val2017)

Then set the paths in vhs_benchmark/conf/config.yaml:

basic:
  image_root: /path/to/datasets/coco
  test_file_base: /path/to/datasets/VHs_qa/single_needle/VHs_large
  output_root: /path/to/ReToken/vhs_benchmark/output_single_needle

Also set basic.image_counts for the haystack sizes to sweep.

This is the experiment that motivates the paper, and it is cheap to reproduce (~5 min per cell on one 80 GB GPU). To reproduce it, run the 2×2 grid for either backbone:

# Qwen3-VL-8B — swap solver=internvl_3_5_rekv for InternVL3.5-8B.
# The two overrides have the same name and nesting on both backbones.
for VV in False True; do
  for NQ in False True; do
    python vhs_benchmark/main_multiple_topK.py solver=qwenvl_3_rekv \
        solver.retrieve_with_V_V=$VV \
        solver.use_needle_query=$NQ \
        solver.flag_save_name=rekv_vv${VV}_needle${NQ} \
        'solver.topk=[1]' 'basic.image_counts=[2]'
  done
done

Note:

• solver.flag_save_name must differ per cell — it is part of the output path, and runs resume by skipping samples that already exist, so a shared name would make later cells silently reuse the first one's results.
• retrieve_with_V_V: score in key space (False, original ReKV: query·key) or in value space (True, Value·Value).
• use_needle_query: retrieve using the full question (False) or using only the target phrase (True).

Expected retrieval recall at image_count=2, topk=1:

python vhs_benchmark/main_multiple_topK.py solver=qwenvl_3_retoken

Each solver config in conf/solver carries only the knobs that change between experiments:

A solver implements generate(...), and optionally generate_multi_topk(...) to encode the images once and sweep several topk values in one pass. Everything that must match across backbones — seeding, shuffling, resume, result format — lives in the base class.

DATA_ROOT=/path/to/datasets   # wherever you keep benchmark data
huggingface-cli download Becomebright/QAEgo4D-MC-test --repo-type dataset --local-dir $DATA_ROOT/qaego4d
huggingface-cli download lmms-lab/LVBench             --repo-type dataset --local-dir $DATA_ROOT/lvbench
huggingface-cli download lmms-lab/Video-MME           --repo-type dataset --local-dir $DATA_ROOT/videomme

Unpack any video archives so the frames are readable by decord, then place the annotations and videos where the loaders expect them:

QAEgo4D is the exception to --data_root: its annotation stores an absolute video_path for every entry instead of deriving one from the root. Point those at your copy before the first run, e.g.

import json
d = json.load(open('test_mc.json'))
for e in d:
    e['video_path'] = f"/path/to/datasets/qaego4d/videos/{e['video_id']}.mp4"
json.dump(d, open('test_mc.json', 'w'))

Then pass the root to each run with --data_root /path/to/datasets. It defaults to the DEFAULT_DATA_ROOT literal near the top of video_qa_ReToken/datasets.py, which the table above resolves against:

DEFAULT_DATA_ROOT = '/path/to/datasets'   # default for --data_root

The relative layout above lives beside it in the same file as DATASET_ANNO_RELPATH / DATASET_VIDEO_RELPATH; edit those if your directories differ. The annotation path can also be set per run with --anno_path.

Entry point: video_qa_ReToken/run_multi_eval.py (CLI only — it delegates to the modules below).

--dataset and --strategies derive their choices from DATASET_LOADERS and STRATEGIES, so adding either is one registry entry plus one function.

# QAEgo4D
python -m video_qa_ReToken.run_multi_eval \
    --save_dir results/qaego4d_retoken_K1 \
    --model_path AvaXiao/ReToken-Qwen3VL-8B \
    --data_root /path/to/datasets \
    --dataset qaego4d --sample_fps 0.5 \
    --n_local 30000 --early_layer_budget 256 \
    --encode_chunk_size 128 --retrieve_size 1 \
    --strategies gt_frames evenly_sample retoken

# LVBench
python -m video_qa_ReToken.run_multi_eval \
    --save_dir results/lvbench_retoken \
    --model_path AvaXiao/ReToken-Qwen3VL-8B \
    --data_root /path/to/datasets \
    --dataset lvbench --sample_fps 0.5 \
    --n_local 30000 --early_layer_budget 256 \
    --encode_chunk_size 128 --retrieve_size 100 \
    --strategies gt_frames evenly_sample retoken

# Video-MME -- no gt_frames here: the benchmark ships no temporal annotations,
# so gt_frames strategy is rejected (see DATASETS_WITHOUT_TEMPORAL_GT in datasets.py).
python -m video_qa_ReToken.run_multi_eval \
    --save_dir results/videomme_retoken \
    --model_path AvaXiao/ReToken-Qwen3VL-8B \
    --data_root /path/to/datasets \
    --dataset videomme --sample_fps 0.5 \
    --n_local 30000 --early_layer_budget 256 \
    --encode_chunk_size 128 --retrieve_size 100 \
    --strategies evenly_sample retoken

Key arguments:

Only the retrieval token trains; the backbone stays frozen.

Training needs two pieces: the label file (which conversation, which images, which one is the needle) and the images it refers to.

1. Images. The label file indexes the MIRAGE training set — download and process the images following the instructions there. ReToken uses five of those sources:

<data_root>/
├── coco/train2017        # ~703k refs
├── vg/VG_100K, VG_100K_2 # ~157k
├── ocr_vqa/images        # ~133k
├── gqa/images            # ~124k
└── textvqa/train_images  #  ~39k

MIRAGE lists further sources (sam, share_textvqa, slidevqa, webqa) that ReToken's label file never references — you do not need them.

2. Label file. Download it here and place it in qwen-vl-finetune/qwenvl/data/label_file/, or point annotation_path in qwen-vl-finetune/qwenvl/data/__init__.py at wherever you keep it. The dataset key used by the released script is retoken_simple_single_conv_filtered.

3. Set the paths. In qwen-vl-finetune/qwenvl/data/__init__.py (and, for the InternVL run, retoken/train_internvl_retoken.py):

DEFAULT_DATA_ROOT = "/path/to/datasets"   # default for --data_root

The launcher forwards it: DATA_ROOT=/path/to/datasets bash ./scripts/retoken_sft_qwen3_8b.sh passes --data_root through to training.

Image paths in the label file are relative to this, e.g. coco/train2017/000000319364.jpg, so it must be the directory holding the folders above. The label file itself is found automatically inside the repo.

Checkpoints go to <repo>/checkpoints/ by default; pass OUTPUT_DIR=/somewhere to the launcher to put them elsewhere.

cd qwen-vl-finetune
bash scripts/retoken_sft_qwen3_8b.sh

Trainer: qwen-vl-finetune/qwenvl/train/train_qwen.py

Launcher: qwen-vl-finetune/scripts/retoken_sft_qwen3_8b.sh

cd InternVL-ReToken
bash scripts/train_internvl3_5_8b_retoken.sh

Trainer: retoken/train_internvl_retoken.py

Launcher: InternVL-ReToken/scripts/train_internvl3_5_8b_retoken.sh

If you find ReToken useful, please consider citing:

@article{xiao2026retoken,
  title={ReToken: One Token to Improve Vision-Language Models for Visual Retrieval},
  author={Xiao, Yao and Tan, Reuben and Zhu, Zhen and Wu, Yuqun and Gao, Jianfeng and Hoiem, Derek},
  journal={arXiv preprint arXiv:2607.28627},
  year={2026}
}

This work is built upon ReKV, Visual Haystacks, Qwen-VL and InternVL. Thanks for their excellent works.

License: This project is released under the MIT License.

Note that some of the software to download and install for this project is subject to separate copyright notices and license terms, which use is subject to the terms and conditions under which they are made available.