UniME-R1 is an Embedder–Adviser framework that learns to reason over retrieved candidates (not the query alone) and generate Retrieval-Centric Chain-of-Thought (RC-CoT) to correct retrieval failures. This repository ships the 2B-scale pair: a Qwen3-VL-2B dual-mode Embedder and the Qwen3-VL-4B retrieval-aware Adviser.
💡 Highlights
Retrieval-Centric CoT (RC-CoT) — The Adviser analyzes the actual top-k retrieved candidates to diagnose model-specific confusion, then emits <cot_focus> (which discriminative cues are missing) and <cot_answer> (a concise corrected query) to redirect retrieval.
Dual-Mode Embedder — A single VLM backbone produces two embeddings via special tokens: <dis_emb> for fast discriminative retrieval and <gen_emb> for RC-CoT-enhanced re-retrieval. Candidates are encoded once with <dis_emb> and reused across both paths — no candidate-side CoT, no index rebuilding.
Adaptive Rerank-or-Retrieve — The Adviser predicts whether a match exists in the top-k set. If yes, it reranks; if not, it appends RC-CoT to the query and re-retrieves over the full corpus.
Retrieval-Oriented GRPO — The Adviser is optimized with a 4-component reward (format / NDCG rerank / CoT-embedding quality / judge decision) that calls a frozen Embedder API to score the Adviser's CoT against mined hard negatives, so the RL signal reflects real end-to-end retrieval quality.
🧱 Model Components
This release contains two sub-directories:
Component
Backbone
Format
Notes
Embedder (embedder/)
Qwen3-VL-2B-Instruct
LoRA (DoRA, r=16, α=64) + new_token_embeddings.pt
Adds <dis_emb> / <gen_emb> tokens; pooling at the special-token position
The Embedder is released as a PEFT/LoRA adapter — load it on top of Qwen/Qwen3-VL-2B-Instruct. The two special tokens (<dis_emb>=151670, <gen_emb>=151669) and their embeddings are provided via added_tokens.json + new_token_embeddings.pt.
UniME-R1 achieves the best overall performance at both model scales. Notably, the 2B model already outperforms all medium-size (4B–7B) baselines, indicating the gains stem from the framework rather than model scale alone.
🌈 Zero-shot General Retrieval
🖊️ Citation
If you find this repository useful, please use the following BibTeX entry for citation.
bibtex
1@misc{unimer1,
2 title={Learning from Failures: Retrieval-Centric CoT via Hard Negatives for Unified Multimodal Retrieval},
3 author={Zelong Sun and Jun Wang and Kaicheng Yang and Tiancheng Gu and Ziyong Feng and Zhiwu Lu},
4 year={2026},
5 eprint={2608.06060},
6 archivePrefix={arXiv},
7 primaryClass={cs.CV},
8 url={https://arxiv.org/abs/2608.06060},
9}
⭐ Don't forget to star this repository if you find it helpful!
