UME-R1: Exploring Reasoning-Driven Generative Multimodal Embeddings

AI summary

Overview

Research area: Multimodal embeddings and multimodal large language models (MLLMs), specifically combining chain-of-thought reasoning and reinforcement learning with verifiable reward (RLVR) for retrieval-style embedding tasks.

Technical level: Advanced. The paper assumes familiarity with contrastive learning (InfoNCE), GRPO-style reinforcement learning, and MLLM training pipelines.

Scope: The paper introduces UME-R1, a two-stage training framework that lets a multimodal embedding model produce either conventional discriminative embeddings or newly proposed reasoning-driven generative embeddings, and evaluates it on the MMEB-V2 benchmark (78 tasks across video, image, and visual documents).

What This Paper Is About

Existing MLLM-based multimodal embedding models are discriminative: they encode an input and read out the final hidden state as the embedding without generating any new tokens, so they cannot exploit the reasoning-driven generation paradigm that has boosted other MLLM tasks. The authors ask how to make an embedding model behave generatively, and whether reinforcement learning (which normally needs verifiable answers, absent in retrieval) can improve such embeddings. Their goal is a single model that can flexibly emit either a discriminative embedding or a reasoning-driven generative embedding produced after a chain-of-thought and a summary.

Key Contributions

1. A cold-start SFT dataset with CoT annotations. The authors sample 50,000 instances from each of the 20 in-distribution datasets in MMEB, add training instances from LLaVA-Hound, ViDoRe, and VisRAG to cover video and visual-document modalities, yielding 1.76 million pairs. CoT rationales for both query and target are generated by the pure-thinking model GLM-4.1V-Thinking, then filtered for repetitive text, reasoning longer than 8,192 tokens, and malformed <think>...</think><answer> output, leaving 1.46 million cold-start SFT pairs. A separate 11,136-pair set is balanced-sampled for RL, prioritizing instances not in the SFT data.

2. UME-R1, a framework for switching between embedding types. A two-stage training strategy: a cold-start SFT stage that combines discriminative contrastive loss, generative contrastive loss over reasoning-augmented representations, and a next-token prediction loss; then an RL stage. The model uses a template with a <disc_emb> token (whose hidden state gives the discriminative embedding) and a generated <gen_emb> token (whose hidden state gives the reasoning-driven generative embedding).

3. First rule-based RL applied to multimodal embeddings. Because embeddings lack standard best answers like math problems, the authors design a reward combining a format reward with an embedding reward that multiplies a ranking term (how many positive-target samples fall in the top-G of positives plus negatives) by a similarity-gap term (average positive similarity minus average negative similarity), and train with GRPO.

4. State-of-the-art results and analysis on MMEB-V2. UME-R1 outperforms conventional discriminative embedding models across 78 tasks, and analyses of an oracle upper bound and pass@k indicate substantial remaining headroom. Datasets, models, and code are released at https://github.com/XMUDeepLIT/UME-R1.

Main Findings

• Generative embeddings beat discriminative ones. Compared with VLM2Vec-V2, UME-R1 achieves an overall improvement of 2.1 while using only two-thirds of VLM2Vec-V2's training data. Compared with DUME (a discriminative-only model trained on the same data and settings), UME-R1 raises total scores by 4.1 for images, 9.0 for videos, and 11.1 for visual documents.

• The two embedding types are complementary. Under an oracle that picks the better embedding mode per test instance, UME-R1-2B and UME-R1-7B gain 4.3 and 3.6 overall points respectively over the generative-only setting — with UME-R1-7B reaching an overall 68.1 and UME-R1-2B reaching 64.4.

• RL improves generative embeddings. Removing the RL stage costs 1.4 points on images, 1.0 on videos, 0.4 on visual documents, and 1.0 overall (2B model). Removing the similarity-gap reward costs 1.4 / 1.0 / 0.3 / 0.9, and removing the ranking reward costs 0.6 / 0.4 / 0.6 / 0.5 on the same four columns. A fixed-threshold reward of 0.5 mainly helps video tasks and yields limited gains elsewhere.

• Generative training also helps discriminative embeddings. Under identical training settings, adding reasoning-driven generative embedding training and the next-token prediction objective improves the discriminative embedding score by 3 points overall, with a 7.5-point gain on visual document tasks; the RL stage adds a further 0.4 points overall without degrading discriminative performance.

• Repeated sampling improves coverage. Using unbiased pass@k estimation on four randomly selected test sets (128 sampled examples each) from image and video modalities, both UME-R1-2B and UME-R1-7B improve coverage with repeated sampling, indicating inference-time scaling potential.

• Self-generated reasoning beats an external reasoner. With the 9B GLM-4.1V-Thinking as an external reasoning model (DUME+Gen), improvements of 19.7 and 3.9 appear on K700 and MSVD, but a 12.3-point drop occurs on CIRR; UME-R1 consistently outperforms DUME+Gen.

• Costs are higher for generative embeddings. DUME requires 1487 H20 GPU-hours for fine-tuning, while UME-R1 uses 2336 H20 GPU-hours for SFT plus 1344 H20 GPU-hours for RL. On a single L40s GPU with vLLM, generative embeddings run at 1.48 (CIRR), 1.14 (FashIQ), 0.50 (K700), and 1.10 (MSVD) samples/s versus 20.0, 19.1, 1.59, and 28.0 samples/s for discriminative embeddings; the gap narrows as input length grows.

Methodology in Plain English

The authors take a standard multimodal embedding setup — queries, positive targets, and in-batch negatives, optimized with an InfoNCE contrastive loss — and augment it with reasoning. First, they use a thinking-capable model (GLM-4.1V-Thinking) to write step-by-step reasoning and a short summary for every query and target in a large training corpus, then keep only well-formed, non-repetitive, reasonably short examples.

Training happens in two stages. In stage one, the model is shown a template asking it to reason inside <think> tags, summarize inside <answer> tags, and then emit a <gen_emb> token. The loss has three parts: the usual discriminative contrastive loss on the <disc_emb> token, a contrastive loss on the generative embedding computed from reasoning-augmented representations, and a next-token prediction loss so the model actually learns to produce the reasoning and summary text at inference time.

In stage two, the model is tuned with GRPO, a reinforcement learning method that compares a group of sampled outputs against their group mean rather than using a learned value function. The reward has two parts: a format reward of 1 or 0 depending on whether the output follows the template, and an embedding reward that multiplies how often positive-target samples rank in the top-G of all positives and negatives by the average similarity gap between positives and negatives. This design avoids the zero-gradient problem of a fixed similarity threshold, which the authors show is less effective.

Evaluation uses MMEB-V2 (9 meta-tasks, 78 tasks) with Hit@1 for image and video tasks and NDCG@5 for visual document tasks; inference uses greedy search with up to 8,192 newly generated tokens.

Why This Matters

Impact on research. This is the first work to treat multimodal embedding as a generation problem, and the first to apply rule-based RL with verifiable reward to embedding tasks, which normally lack ground-truth answers. It also reframes embedding improvement as an inference-time compute scaling problem rather than purely a model-size problem, and shows discriminative and generative embeddings are complementary rather than competing.

Real-world applications:

• Multimodal retrieval systems that search across images, video, and visual documents (slides, PDFs, scanned pages) with a single model.
• Retrieval-augmented generation pipelines that need retrievers which can explain and reason about why a result matches.
• Video search and moment retrieval, where the model must ground a query to a specific segment.
• Content moderation or classification pipelines where a written reasoning trace makes retrieval decisions auditable and interpretable.

Industry relevance. The results use open 2B and 7B Qwen2-VL backbones and an open benchmark, so the approach is directly reproducible at modest scale. The measured GPU-hour and throughput costs also give practitioners concrete numbers for deciding when reasoning-driven retrieval is worth the latency and compute.

Future Directions

1. Adaptive routing. The oracle analysis shows gains from choosing between discriminative and generative embeddings per instance, but designing a practical router to make that choice remains an open problem stated in the limitations.

2. Better RL data and strategies. The authors note room for improvement through harder negative examples for RL or scaling up training instances, and suggest adaptive thresholds for the reward may help.

3. Inference-time scaling. The pass@k curves suggest repeated sampling improves embedding quality, but the paper calls for explicitly exploring inference-time scaling techniques to further enhance generative embedding quality.

4. Cost reduction. Generative embeddings require long CoT generation, raising training and inference costs; reducing this overhead while retaining the quality gains is an open engineering question.

Target Audience

Researchers and engineers working on multimodal retrieval, multimodal embedding models, and MLLM post-training, particularly those interested in applying reinforcement learning to tasks without verifiable ground-truth answers. It is also relevant to practitioners building multimodal RAG or cross-modal search systems who need to weigh accuracy against inference latency, and to readers tracking the intersection of reasoning models and representation learning.

Authors’ abstract

The remarkable success of multimodal large language models (MLLMs) has driven advances in multimodal embeddings, yet existing models remain inherently discriminative, limiting their ability to benefit from reasoning-driven generation paradigm. In this work, we pioneer the exploration of generative embeddings, unifying embedding tasks within a generative paradigm. We propose UME-R1, a universal multimodal embedding framework consisting of a two-stage training strategy: a cold-start supervised fine-tuning equips the model with reasoning capabilities and enables it to generate both discriminative and generative embeddings; a subsequent reinforcement learning enhances reasoning and further optimizes generative embedding quality. This pioneering work reveals four key insights: 1) generative embeddings unlock substantial performance gains over conventional discriminative embeddings by leveraging the powerful generative reasoning capabilities of MLLMs; 2) discriminative and generative embeddings are complementary, whose combined oracle performance far exceeding that of either alone; 3) RL can effectively enhance generative embeddings, establishing a scalable optimization paradigm.; 4) repeated sampling at inference boosts downstream task coverage (pass@k), highlighting the inference-time scalability potential of generative embeddings. Evaluated on the MMEB-V2 benchmark across 78 tasks spanning video, image, and visual documents, UME-R1 significantly outperforms conventional discriminative embedding models and offers a foundation for more interpretable, reasoning-driven generative multimodal embeddings. Our code, models, and datasets will be publicly available at https://github.com/XMUDeepLIT/UME-R1.