DeepGlint-AI / UniME-R1-2B

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:

├── adviser/                       # Qwen3-VL-4B Adviser (full weights, vLLM-ready)
│   ├── model-0000{1,2}-of-00002.safetensors
│   ├── config.json / generation_config.json
│   ├── preprocessor_config.json / video_preprocessor_config.json
│   └── tokenizer.json / chat_template.jinja
└── embedder/                      # Qwen3-VL-2B Embedder (LoRA adapter)
    ├── adapter_config.json
    ├── adapter_model.safetensors
    ├── new_token_embeddings.pt    # trained <dis_emb>/<gen_emb> embeddings
    ├── config.json
    └── preprocessor_config.json / video_preprocessor_config.json

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 .

The Adviser always emits five structured fields:

🚀 Quick Start

git clone https://github.com/deepglint/UniME-R1.git
cd UniME-R1

conda create -n vlm2vec python=3.10 -y
conda activate vlm2vec

# Install torch matching your CUDA first, e.g.:
# pip install torch==2.5.1 torchvision --index-url https://download.pytorch.org/whl/cu121
# pip install flash-attn==2.7.3 --no-build-isolation

pip install -r Embedder/requirements.txt

🔍 Embedder-only evaluation (direct <dis_emb> retrieval)

cd Embedder
bash shell/eval/eval.sh image "../models/UniME-R1-2B/embedder"

🎯 Full Adviser evaluation (rerank + RC-CoT iterative retrieval)

Serve the Adviser via vLLM, then run the unified evaluation:

vllm serve models/UniME-R1-2B/adviser --tensor-parallel-size 8 --port 9000

cd RL/eval
export EMBEDDER_CHECKPOINT="../../models/UniME-R1-2B/embedder"
export ADVISER_MODEL="Qwen3-VL-4B-Instruct"
export ADVISER_URL="http://127.0.0.1:9000/v1"
bash eval.sh image      # image | visdoc | video | uvrb | image_caption

📊 Results

🏆 MMEB-V2

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.

@misc{unime-r1,
      title  = {Learning from Failures: Retrieval-Centric CoT via Hard Negatives for Unified Multimodal Retrieval},
      author = {Sun, Zelong and Wang, Jun and Yang, Kaicheng and Gu, Tiancheng and Feng, Ziyong and Lu, Zhiwu},
      year   = {2025},
      url    = {https://github.com/deepglint/UniME-R1}
}