Technical Summary: Reason What Matters (ReWAM)

Problem Statement

Universal Multimodal Embeddings (UME) aim to learn unified representations across modalities to support diverse retrieval tasks. While recent methods integrate Chain-of-Thought (CoT) reasoning into UME to handle complex, compositional, and fine-grained retrieval tasks, two critical limitations hinder their deployment at corpus scale:

1. Coarse Credit Assignment: Existing methods often employ Group Relative Policy Optimization (GRPO) with retrieval-based rewards. This approach assigns a single scalar advantage to all CoT tokens within a trajectory. Consequently, it fails to distinguish between reasoning steps that are genuinely supported by input evidence and those that are hallucinated or irrelevant. It also does not explicitly identify the specific distinctions that separate positive targets from hard negatives. This lack of granularity limits performance gains, particularly on Out-of-Distribution (OOD) tasks.
2. Inference Latency: Current approaches typically generate a complete CoT trace before producing an embedding, even when a partial trace already contains sufficient evidence for retrieval. This autoregressive generation of full traces introduces substantial latency that scales with both trace length and corpus size, making explicit reasoning impractical for large-scale systems.

Methodology: Reason What Matters (ReWAM)

ReWAM is a retrieval-grounded reasoning framework designed to jointly optimize reasoning quality and inference efficiency. It builds upon a decoupled reasoner–embedder architecture, where a fixed CoT-conditioned embedder provides a stable retrieval signal. The framework introduces two core components:

1. Retrieval-aware Self-Distillation (RASD)

RASD addresses the coarse credit assignment problem by constructing "privileged guidance" to refine token-level supervision.

• Privileged Information Construction: A multimodal analyzer contrasts the positive target against retrieved hard negatives to extract two types of privileged information: (1) Grounded Evidence (facts supported by the multimodal input) and (2) Decision Boundaries (distinctions separating the target from confusable candidates).
• Claim Verification: The analyzer verifies sampled CoT trajectories against this privileged information to identify supported claims and conflicting statements, forming "Contrastive Evidence Privileged Information" (CEPI).
• Token-Level Reweighting: An on-policy self-teacher uses the CEPI to rescore the sampled trajectories. It computes the log-likelihood difference for each token with and without the privileged guidance. This difference serves as a stop-gradient signal to generate token-specific weights.
• Objective: These weights modulate the trajectory-level advantage derived from retrieval rewards. This ensures that policy updates focus on reasoning steps that are both input-supported and discriminative, rather than uniformly reinforcing the entire trajectory.

2. Retrieval-adaptive Inference (RAI)

RAI addresses inference latency by adaptively controlling the length of the CoT and accelerating token generation.

• Retrieval Confidence Head: A lightweight MLP head predicts the "remaining retrieval utility" of a partial CoT. It estimates how much additional retrieval gain can be achieved by continuing the reasoning trace versus stopping.
• Early Truncation: During inference, if the predicted remaining utility falls below a threshold, the CoT generation is truncated immediately, avoiding unproductive reasoning steps.
• Speculative Decoding: To accelerate the generation of useful continuations, ReWAM employs speculative decoding. A semi-autoregressive draft model proposes multiple tokens in parallel, which are verified by the frozen target reasoner. This reduces the per-token decoding cost while maintaining the quality of the explicit CoT.

Key Contributions

1. ReWAM Framework: A novel retrieval-grounded framework that bridges the gap between retrieval quality and inference efficiency for universal multimodal embeddings, establishing explicit reasoning as a scalable paradigm for massive corpora.
2. Retrieval-aware Self-Distillation (RASD): A method for token-level credit assignment that leverages retrieval evidence verification to distinguish between input-supported reasoning and hallucinations, refining the training signal beyond scalar trajectory rewards.
3. Retrieval-adaptive Inference (RAI): An adaptive inference strategy that combines trajectory-level truncation (guided by predicted remaining utility) with token-level acceleration (via speculative decoding) to significantly reduce latency without sacrificing retrieval performance.

Experimental Results

Extensive experiments were conducted on the MMEB-V2 (general-purpose multimodal retrieval) and MRMR (reasoning-intensive retrieval) benchmarks.

• Retrieval Performance: ReWAM achieves state-of-the-art performance across both benchmarks.
  • On MMEB-V2, ReWAM (Qwen2-VL-2B) achieves an overall score of 66.5, outperforming the previous best explicit-CoT method (Embed-RL) and surpassing the latent-reasoning method PLUME by 4.9 points while retaining explicit CoT interpretability.
  • On MRMR, ReWAM achieves top scores in reasoning-intensive tasks, leading by significant margins (e.g., +3.1 points over RIME, +6.5 points over Embed-RL) and demonstrating strong performance in domain-specific tasks like Mathematics, Physics, and Engineering.
• Inference Efficiency: ReWAM delivers substantial throughput improvements.
  • It achieves up to 5× the throughput of competitive explicit-CoT UME methods (e.g., Embed-RL).
  • It surpasses the efficiency-oriented latent-reasoning method PLUME in both retrieval performance and throughput.
  • Ablation studies confirm that the retrieval confidence head reduces CoT length by 26–36% while increasing throughput by an additional 26–64% over speculative decoding alone.

Significance

The paper argues that explicit reasoning has often been dismissed as too costly for large-scale deployment due to latency and coarse optimization signals. ReWAM demonstrates that by grounding reasoning in retrieval feedback (via RASD) and adaptively controlling inference (via RAI), explicit CoT can be made both highly effective and efficient. The results suggest that reasoning-enhanced UME is practical for real-world systems operating over massive corpora, offering a path to high-quality retrieval without sacrificing the interpretability and compositional reasoning capabilities of explicit text traces.