UniME-R1: Learning from Failures for Unified Multimodal Embedding

Learning from Failures: Retrieval-Centric CoT via Hard Negatives for Unified Multimodal Retrieval

Zelong Sun*, Jun Wang*, Kaicheng Yang^†, Tiancheng Gu, Ziyong Feng, Zhiwu Lu^‡

📢 Latest News

• 2025/08/06: ✨ We release the model weights in 🤗 HuggingFace.
• 2025/08/06: ✨ We release the evaluation, training and serving code of UniME-R1.

✅ Release Status

• Release the paper of UniME-R1
• Release model weights (UniME-R1-2B / UniME-R1-4B)
• Release the evaluation scripts
• Release the Embedder training code & Adviser GRPO training pipeline

💡 Highlights

Existing Reasoner–Embedder methods generate Chain-of-Thought (CoT) from the query alone — they explain what the query describes, but not what the retriever misunderstands. UniME-R1 argues that effective retrieval reasoning should be grounded in retrieval feedback.

• Retrieval-Centric CoT (RC-CoT) — The Adviser analyzes the actual 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.

Generic CoT broadly enriches query semantics, whereas RC-CoT diagnoses retrieval failures and corrects the retrieval direction with targeted evidence.

• 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. This avoids unnecessary re-retrieval for already-solvable queries.
• 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.

Overview of UniME-R1. A dual-mode embedder performs direct and RC-CoT-enhanced retrieval; a retrieval-aware adviser examines the initial top-k candidates and adaptively routes each query to reranking or RC-CoT-enhanced full-corpus re-retrieval.

🛠️ Implementation

📦 Environment

UniME-R1 uses two separate conda environments — keep them isolated to avoid dependency conflicts (the RL stack pins different versions of torch / vLLM).

Prerequisites: Python >= 3.10, NVIDIA GPUs with CUDA (8 × 80GB for multi-node training; single GPU fine for evaluation). GPU wheels for torch / flash-attn must match your CUDA version — install those first.

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

1. Embedder environment — vlm2vec (Embedder training & evaluation)

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

2. RL environment — unime_rl (Adviser GRPO training & Adviser evaluation)

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

# Install torch matching your CUDA first (must also be compatible with vLLM), 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 RL/requirements.txt
cd RL && pip install -e . && cd ..     # install the vendored verl package

📊 Model Download

Pre-trained checkpoints are hosted on HuggingFace. Download them under models/:

cd models
# UniME-R1-2B: Embedder (Qwen3-VL-2B) + Adviser (Qwen3-VL-4B)
huggingface-cli download DeepGlint-AI/UniME-R1-2B --local-dir UniME-R1-2B
# UniME-R1-4B: Embedder (Qwen3-VL-4B) + Adviser (Qwen3-VL-4B)
huggingface-cli download DeepGlint-AI/UniME-R1-4B --local-dir UniME-R1-4B

Model	Embedder Backbone	Adviser Backbone	MMEB-V2 Overall	HuggingFace
UniME-R1-2B	Qwen3-VL-2B	Qwen3-VL-4B	69.9
UniME-R1-4B	Qwen3-VL-4B	Qwen3-VL-4B	70.3

Each release contains two sub-directories: adviser/ (full merged Qwen3-VL-4B weights, vLLM-ready) and embedder/ (LoRA DoRA adapter + new_token_embeddings.pt for the <dis_emb> / <gen_emb> tokens). The Adviser is shared across the two releases.

🗂️ Project Structure

UniME-R1/
├── Embedder/               # Embedder training & evaluation framework
│   ├── train.py            # Training entry point
│   ├── eval.py             # Distributed evaluation (+ per-modality eval_*.py)
│   ├── convert_model.py    # Weight format conversion utility
│   ├── src/                # Source code (model, data, loss, trainer, GradCache)
│   └── shell/              # Unified eval.sh + train scripts + YAML configs
├── RL/                     # GRPO reinforcement learning (vendored verl)
│   ├── eval/               # Unified Adviser evaluation
│   │   ├── adviser_eval.py # Rerank + CoT iterative-retrieval evaluation
│   │   ├── eval.sh         # eval entrypoint + per-modality YAML configs
│   │   └── report_score.py
│   ├── train/              # Adviser GRPO config, launchers & Embedder serving
│   │   ├── adviser_grpo.yaml
│   │   ├── run_adviser_grpo.sh (+ start_ray_{head,worker}.sh)
│   │   └── start_embedder.sh + serve_embedder.py  # frozen-Embedder API
│   └── verl/               # Vendored verl framework (Apache 2.0, Bytedance)
├── models/                 # Checkpoints (download from HuggingFace)
├── dataset/                # Training/eval data (see Data Preparation)
└── requirements.txt

🚀 Quick Start

This repository supports three workflows: Embedder-only retrieval, full Embedder–Adviser inference (rerank + RC-CoT), and training (Embedder SFT + Adviser GRPO).

🔍 1. Embedder-only Evaluation (direct <dis_emb> retrieval)

conda activate vlm2vec
cd Embedder

# Pick a modality: image | video | visdoc | uvrb | image_caption | all
bash shell/eval/eval.sh image "../models/UniME-R1-2B/embedder"

Supported benchmarks:

Config	Benchmark	Tasks
image.yaml	MMEB-V2	Classification, VQA, Image retrieval, Visual grounding
video.yaml	MMEB-V2	Video classification, Retrieval, Moment retrieval, Video QA
visdoc.yaml	MMEB-V2	ViDoRe / VisRAG document retrieval
uvrb.yaml	Zero-shot	UVRB (Universal Video Retrieval Benchmark)
image_caption_retrieval.yaml	Zero-shot	Flickr30K, COCO2014, ShareGPT4V, Urban1k

Metrics: Hit@K, NDCG@10, Precision@10, Recall@10, F1, MAP, MRR.

🎯 2. Full Adviser Evaluation (rerank + RC-CoT iterative retrieval)

First, serve the Adviser model via vLLM:

# Use 8 GPUs (tensor parallel); adjust --tensor-parallel-size to your setup
vllm serve models/UniME-R1-2B/adviser \
    --tensor-parallel-size 8 \
    --port 9000

Then run the unified Adviser evaluation:

conda activate unime_rl
cd RL/eval

# Point to the matching Embedder checkpoint and the Adviser URL
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"

# Pick a modality: image | visdoc | video | uvrb | image_caption
bash eval.sh image

The Adviser pipeline runs:

1. The Embedder retrieves top-k candidates using <dis_emb> embeddings
2. The Adviser analyzes each candidate, reranks them, and judges correctness
3. If no correct match is found (<rerank_judge> == -1), the Adviser generates RC-CoT (<cot_focus>, <cot_answer>)
4. The Embedder re-encodes the query with RC-CoT via <gen_emb> and re-retrieves
5. Steps 2–4 repeat for max_rounds iterations (default 1; UniME-R1++ uses 2)

🏋️ 3. Training

Embedder (GradCache contrastive learning + hard negatives)

conda activate vlm2vec
export MODEL_NAME="Qwen/Qwen3-VL-2B-Instruct"   # or a local path
export EXP_NAME="unime_embedder_qwen3vl_2b"

cd Embedder
bash shell/train/train_v2-qwen3vl-2B_fullv2_multinode.sh

For multi-node training, set NNODES, NPROC_PER_NODE, MASTER_ADDR, MASTER_PORT, NODE_RANK, and WANDB_API_KEY (or WANDB_DISABLED=true).

Key Embedder arguments: --model_type qwen3_vl --pooling special_token --use_special_emb_tokens True --emb_mode both --num_hardneg K --grad_cache True --lora True. See Embedder/README.md for the full list.

Adviser GRPO (frozen Embedder + verl)

conda activate unime_rl
cd RL

# 1. Serve the frozen Embedder for dynamic CoT-guided query re-encoding
bash train/start_embedder.sh

# 2. Multi-node: start the Ray cluster
bash train/start_ray_head.sh        # on head node
bash train/start_ray_worker.sh      # on each worker node

# 3. Launch GRPO training on the head node
bash train/run_adviser_grpo.sh

Note: The frozen-Embedder serving scripts (train/start_embedder.sh + train/serve_embedder.py) are bundled. The offline RL data-prep scripts (sample construction, target-embedding precompute) are not bundled; the GRPO config RL/train/adviser_grpo.yaml expects pre-computed artifacts (*.jsonl + sample_embeddings.pt) already present under dataset/Adviser/RL/. See RL/README.md for the expected data layout.

📊 Results

🏆 MMEB-V2

UniME-R1 achieves the best overall performance at both model scales, reaching 69.9 (2B) and 70.3 (4B). Notably, the 2B model already outperforms all medium-size (4B–7B) baselines, indicating that the gains stem from the framework rather than model scale alone.

🌈 Zero-shot General Retrieval

UniME-R1 also generalizes well to diverse cross-modal retrieval tasks (Flickr30K, COCO2014, ShareGPT4V, Urban1K, UVRB).

📈 Inference Efficiency

On MMEB-V1 (3,600 queries, 111,384 candidates), UniME-R1 encodes each candidate once with <dis_emb> at 0.01 s/candidate — 27× faster than candidate-side CoT methods (Embed-RL: 0.27 s). Query-side overhead increases modestly from 0.28 s to 0.34 s, since the rerank-or-retrieve mechanism avoids unnecessary full-corpus retrieval when the target is already in the initial top-k.

📦 Data Preparation

Training data is hosted on HuggingFace: . It is not bundled in this repo due to size — download it under dataset/:

huggingface-cli download DeepGlint-AI/UniME-R1 --repo-type dataset --local-dir dataset

This provides both dataset/Adviser/ (Adviser SFT + GRPO artifacts) and dataset/Embedder/ (Embedder training parquets). All UniME-R1 training data is derived from the MMEB-V2 training set, plus visual-document data (ViDoRe, VisRAG) and video data (LLaVA-Hound, MSRVTT, MSVD).

Stage	Data	Volume
Embedder <dis_emb>	MMEB-V2 train (all retained examples)	full corpus
Embedder <gen_emb>	RC-CoT-augmented subset of MMEB-V2	1.73M samples
Adviser SFT	structured 5-field annotations	643K samples
Adviser GRPO	embedder-aligned failure cases	13K samples

Please refer to the dataset/ layout in the repository and the paper's Appendix A for data sources, modality-balanced sampling, hard-negative mining, and the expected format of the GRPO artifacts.

👏 Acknowledgements

This project builds upon and adapts code from:

• verl (Bytedance) — RL training framework
• VLM2Vec — Multimodal embedding framework and MMEB benchmark
• Transformers (HuggingFace) — Model implementations
• Qwen-VL (Alibaba) — Vision-language model backbones

📄 License

This project is licensed under the MIT License. See LICENSE for details. The RL/verl/ directory contains a vendored fork of verl (Apache 2.0, © Bytedance Ltd.).

🖊️ 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}
}

⭐ Don't forget to star this repository if you find it helpful!