💻 computer science

ReLoop-UME: Recurrent Depth with Learnable Retrieval Registers for Universal Multimodal Embedding

ReLoop-UME introduces a novel universal multimodal embedding architecture that enhances retrieval performance and efficiency by executing early layers once, recurrently reusing a parameter-shared block with learnable retrieval registers to accumulate evidence across depth, and applying final mapping layers only after the last loop, thereby achieving superior speed and accuracy compared to existing methods.

Original authors: Shijie Wang, Xiangzhao Hao, Yueti Li, Guangyu Cao, Xinyu Tang, Haiyun Guo

Published 2026-08-03

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Shijie Wang, Xiangzhao Hao, Yueti Li, Guangyu Cao, Xinyu Tang, Haiyun Guo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find a specific needle in a massive, chaotic haystack. But this isn't just any haystack; it's a "multimodal" one, meaning the hay is made of text, the needles are pictures, and sometimes the whole thing is a video. In the world of artificial intelligence, this is called Universal Multimodal Embedding (UME). Think of it as a super-smart librarian who can take a question written in words, a photo of a cat, or a clip of a movie, and instantly translate them all into a single, secret "language of similarity." Once everything is translated into this secret code, the librarian can instantly tell you which items are related, even if they look completely different on the surface.

For a long time, these AI librarians worked like sprinters: they would look at the input once, run a single lap through their brain, and spit out the answer. This was fast, but sometimes the answer wasn't very precise, especially for tricky questions. Recently, some researchers tried to make the librarian "think harder" by forcing them to write out a long, step-by-step reasoning note before giving the answer. While this helped, it was like asking the librarian to write a whole novel just to find a book; it took way too long and slowed everything down. The big question became: Is there a way to make the librarian think deeper without making them write a novel?

Enter ReLoop-UME, a new approach that suggests the answer lies not in writing more, but in re-reading the right parts. The researchers discovered that an AI's brain doesn't work in a straight line where every layer does the same thing. Instead, they found that the early layers are like a "context setter" (getting the vibe), the middle layers are the "detective work" (finding the clues), and the final layers are just "packaging" (writing the final label).

The paper proposes a clever trick: instead of making the AI write a long reasoning note (which is slow and messy), they let the AI loop through the "detective work" section of its brain multiple times. To make sure the AI doesn't forget what it found in the first loop, they added a special tool called Learnable Retrieval Registers. Imagine these as sticky notes that the AI can stick to its own forehead. As the AI loops through the detective section, it updates these sticky notes with new clues, accumulating evidence without ever needing to generate new words or slow down the process.

The results are impressive. On a massive test called MMEB-V2, which includes thousands of tasks involving images, videos, and documents, this new method consistently found better matches than previous models. It was 44.9 times faster than the previous "reasoning" models that tried to write out their thoughts, and 1.5 times faster than another fast model, while still being more accurate. The researchers suggest that by focusing the extra computing power exactly where the "detective work" happens and using those sticky notes to hold the clues, they created a system that is both smarter and significantly quicker. It's like giving a detective a magnifying glass and a notepad, rather than forcing them to write a diary entry before they can point out the suspect.

Technical Summary: ReLoop-UME

Problem Statement

Universal Multimodal Embedding (UME) aims to map heterogeneous inputs (text, images, videos, visual documents) into a shared embedding space for retrieval. Existing approaches generally fall into two categories:

Single Forward Encoding: Models like VLM2Vec and GME produce embeddings in one pass. While efficient, their fixed computation may be insufficient for complex query intents and fine-grained cross-modal matching.
Token-Expanded Reasoning: Recent methods (e.g., UME-R1, TRACE, PLUME) introduce additional computation via explicit reasoning tokens (Chain-of-Thought) or latent autoregressive states. While these improve retrieval accuracy, they suffer from significant drawbacks:
- Latency: Serial generation increases inference time.
- Dependency: The final embedding relies on generated intermediate states, which can introduce errors or hallucinations.
- Resource Growth: These methods expand the token sequence length, increasing memory and computational overhead.

The central question posed by the authors is: Can useful computation be expanded along model depth while keeping the token workspace fixed?

Methodology: ReLoop-UME

The authors propose ReLoop-UME, a framework that allocates additional computation to specific layers of the model depth rather than extending the token sequence. The method is built on three core components:

1. Layer-wise Analysis of Retrieval Formation

Through an analysis of independently trained single-forward UME models, the authors observe a non-uniform progression of retrieval discriminability across layers:

Prefix Understanding Stage (Early Layers): Contextualizes multimodal inputs; positive and negative similarities remain entangled.
Retrieval Formation Stage (Middle-to-Late Layers): A contiguous interval where the separation between positive and negative query–candidate similarities increases most significantly. This is where retrieval-discriminative features are formed.
Embedding Mapping Stage (Final Layers): Maps the formed features into the final output embedding space.

This analysis identifies the Retrieval Formation Stage as the optimal target for additional depth.

2. Localized Recurrent Refinement

Instead of generating new tokens, ReLoop-UME executes the model in three phases:

Prefix: The input passes through the early "Understanding" layers once.
Recurrent Block: The input is then processed by a parameter-shared block corresponding to the "Retrieval Formation" layers. This block is applied $T$ times (recurrent depth).
Suffix: After the final loop, the refined hidden states pass through the "Mapping" layers to produce the final embedding.

This approach concentrates extra computation specifically on the layers responsible for establishing query–candidate discrimination without duplicating backbone parameters or extending the token sequence.

3. Learnable Retrieval Registers

To manage state across recurrent loops without expanding the token sequence, the authors introduce Learnable Retrieval Registers.

A small set of trainable states ( $M$ registers) is appended to the multimodal token sequence.
These registers persist through every loop, attending to the complete input and exchanging evidence among themselves.
They serve as a fixed-size workspace for accumulating and consolidating retrieval evidence across loops.
The final register acts as the stable readout for the embedding.
Unlike reasoning tokens, these registers are not generated step-by-step; they are fixed in size and do not grow with recurrent depth.

Key Contributions

Stage-wise Account of Retrieval Formation: The paper provides a layer-wise analysis revealing that retrieval geometry forms non-uniformly, identifying a specific "Retrieval Formation" interval as the target for depth expansion.
Localized Recurrent Refinement: ReLoop-UME couples localized recurrence with persistent retrieval states. It repeatedly applies a parameter-shared retrieval-forming block while using Learnable Retrieval Registers to maintain a fixed-size workspace for cross-loop evidence accumulation.
Efficiency-Accuracy Tradeoff: The method achieves stronger retrieval performance with significantly lower latency compared to token-expanded reasoning methods, avoiding the serial generation bottleneck.

Experimental Results

The authors evaluated ReLoop-UME on the MMEB-V2 benchmark (78 tasks across image, video, and visual documents) and performed zero-shot transfer to MRMR.

Performance:
- On MMEB-V2, the 2B model achieved an All score of 63.2, outperforming strong baselines like PLUME (61.6) and UME-R1 (60.1).
- The 7B model reached 65.9, surpassing UME-R1 (64.5).
- Gains were most pronounced in Image and Visual Document tasks; Video performance showed mixed results, likely due to the fixed frame sampling strategy.
- Zero-shot transfer to MRMR showed competitive results, particularly in Knowledge retrieval (71.6 vs. 67.4 for Ops-MM-Embedding).
Efficiency:
- ReLoop-UME is 44.9× faster than UME-R1 and 1.5× faster than PLUME.
- Compared to the single-forward baseline VLM2Vec-V2, ReLoop-UME introduces only a 1.3× latency overhead while improving the All score by 5.2 points.
Ablation Studies:
- Stage Localization: Recurrence applied only to the identified formation interval (layers 17–26 for Qwen2-VL-2B) yielded the best results. Full-stack recurrence or recurrence only in the final layer performed worse.
- Registers: Removing the registers dropped performance, confirming their role in preserving evidence. Increasing the register count beyond 5 yielded diminishing returns.
- Depth: Performance saturated at a moderate recurrent depth ( $T=4$ ); increasing to $T=8$ slightly degraded performance.

Significance and Claims

The paper claims that ReLoop-UME offers a more efficient route to universal multimodal embeddings by shifting the paradigm from token expansion to depth expansion.

Mechanism: By localizing recurrence to the specific layers where retrieval discrimination is formed, the model strengthens query–candidate separation without the latency and memory costs of autoregressive token generation.
State Management: The Learnable Retrieval Registers provide a mechanism to accumulate evidence across loops without increasing the sequence length or requiring explicit reasoning trajectories.
Generalization: The approach is shown to be effective across different model backbones (Qwen2-VL, Qwen3-VL, Qwen3.5) and scales, suggesting the "localization and looping" principle is robust to architectural variations.

The authors conclude that while the method significantly improves the accuracy-efficiency frontier, limitations remain in scenarios requiring dense temporal evidence (e.g., specific video actions not captured in sampled frames) or fine-grained local details in visual documents, where the global recurrent refinement may dilute specific modifiers.