README.md

---
license: mit
datasets:
- TIGER-Lab/MMEB-train
language:
- en
base_model:
- Qwen/Qwen3-VL-2B-Instruct
- Qwen/Qwen3-VL-4B-Instruct
library_name: transformers
tags:
- Retrieval
- Multimodal
- Embedding
- Chain-of-Thought
- Reinforcement-Learning
pipeline_tag: image-text-to-text
---

<div align="center">

<h1>UniME-R1-2B: Learning from Failures for Unified Multimodal Retrieval</h1>

<a href="#">Zelong Sun*</a>,</span>
<a href="#">Jun Wang*</a>,</span>
<a href="#">Kaicheng Yang</a>,</span>
<a href="#">Tiancheng Gu</a>,</span>
Ziyong Feng,</span>
<a href="#">Zhiwu Lu</a></span>

Glint Lab

[![GitHub](https://img.shields.io/badge/⭐-GitHub-black?logo=github)](https://github.com/deepglint/UniME-R1)
[![Paper](https://img.shields.io/badge/📄-Paper-b31b1b.svg)](https://arxiv.org/abs/2608.06060)
[![Model](https://img.shields.io/badge/🤗-UniME--R1_Models-yellow)](https://huggingface.co/DeepGlint-AI)
</div>

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.

<div align="center">
  <img src="figures/intro.png" width="95%">
</div>

- **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.

<div align="center">
  <img src="figures/method.png" width="100%">
</div>

## 🧱 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 |
| **Adviser** (`adviser/`) | Qwen3-VL-4B-Instruct | Full merged weights (bf16) | Outputs 5 structured XML fields; served with vLLM |

```
├── 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:

| Field | Role |
|-------|------|
| `<rerank_think>` | Candidate-by-candidate relevance analysis |
| `<rerank_list>` | Predicted candidate ordering (JSON array, 0-based) |
| `<rerank_judge>` | Path decision: best-match ID, or `-1` if none matches |
| `<cot_focus>` | Retrieval-failure diagnosis — which discriminative cues are confused |
| `<cot_answer>` | Concise refined query text for re-retrieval |

## 🚀 Quick Start

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

```bash
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)

```bash
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:

```bash
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.

<div align="center">
  <img src="figures/mmeb-v2.png" width="95%">
</div>

### 🌈 Zero-shot General Retrieval

<div align="center">
  <img src="figures/zero-shot.png" width="90%">
</div>

## 🖊️ Citation

If you find this repository useful, please use the following BibTeX entry for citation.

```bibtex
@misc{unimer1,
      title={Learning from Failures: Retrieval-Centric CoT via Hard Negatives for Unified Multimodal Retrieval}, 
      author={Zelong Sun and Jun Wang and Kaicheng Yang and Tiancheng Gu and Ziyong Feng and Zhiwu Lu},
      year={2026},
      eprint={2608.06060},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2608.06060}, 
}
```

<div align="center">
⭐ Don't forget to star this repository if you find it helpful!
</div>