TRACE: Task-Adaptive Reasoning and Representation Learning for Universal Multimodal Retrieval

Xiangzhao Hao; Shijie Wang; Tianyu Yang; Tianyue Wang; Haiyun Guo; Jinqiao Wang

cs.CVMar 3, 2026

TRACE: Task-Adaptive Reasoning and Representation Learning for Universal Multimodal Retrieval

Authors: Xiangzhao Hao, Shijie Wang, Tianyu Yang, Tianyue Wang, Haiyun Guo, Jinqiao Wang

Organizations: Institute of Automation, Chinese Academy of Sciences · University of Chinese Academy of Sciences

Abstract

Universal Multimodal Retrieval requires unified embedding models capable of interpreting diverse user intents, ranging from simple keywords to complex compositional instructions. While Multimodal Large Language Models (MLLMs) possess strong reasoning capabilities, prevailing adaptations confine them to static encoders, underutilizing their generative potential. This encoder-only paradigm struggles with complex intents that demand logical deduction rather than superficial pattern matching. To address this, we introduce TRACE (Task-adaptive Reasoning And Compressing Embeddings). TRACE unifies generative reasoning with discriminative representation learning. It first generates a structured Chain-of-Thought (CoT) to explicitly reason about the query, and subsequently compresses this reasoning trace into a compact embedding via a dedicated token. To train this framework, we construct M-BEIR-CoT, a large-scale dataset featuring a difficulty-aware routing strategy. Experiments on the M-BEIR benchmark establish TRACE as the new state-of-the-art. Crucially, TRACE demonstrates a learned implicit routing behavior. It autonomously activates reasoning for complex queries while bypassing it for simpler ones, achieving an optimal balance between retrieval accuracy and inference throughput. Furthermore, by internalizing the deductive process, TRACE exhibits remarkable zero-shot transferability to unseen domains and novel constraints.

Figures & tables

Figure 1: The TRACE Framework. TRACE learns a query-dependent inference strategy. (a) For simple queries, it implicitly bypasses the reasoning stage and directly extracts features to maintain high efficiency. (b) For complex queries, it automatically activates the task-adaptive reasoning process. The model generates an explicit reasoning trace [ 44 ] to resolve semantic ambiguities before compressing this context into the final representation. (c) Performance comparison on the M-BEIR benchmark [ 43 ] demonstrates the effectiveness of TRACE, particularly on reasoning-intensive tasks.

Figure 2: The construction pipeline of the M-BEIR-CoT dataset. The process operates in three phases: (1) Query Complexity Assessment: An advanced MLLM assesses query difficulty, routing simple queries to a direct path (generating only <|emb|> ) and complex queries to a reasoning path (generating CoT + <|emb|> ). (2) Task-Specific CoT Generation: We design specialized prompts for diverse tasks ( e.g . , captioning, text edit, VQA) to generate structured reasoning traces enclosed in <reasoning> tags. (3) Dual Filtering & Curation: To ensure data quality, we apply a coarse-to-fine strategy. We first use rule-based filtering to verify formats and lengths, followed by model-based filtering to ensure semantic consistency between the generated text and ground-truth targets.

Figure 3: Illustration of the TRACE architecture. The model processes a multimodal query through a frozen vision encoder and a trainable projector. The LLM acts as a unified reasoner and encoder. It first generates a Chain-of-Thought (CoT) [ 44 ] to interpret the intent and then compresses the semantics into a learnable <|emb|> token. The final query feature is extracted from the hidden state immediately preceding <|emb|> . During training, the model is optimized jointly using Cross-Entropy (CE) loss for reasoning generation and InfoNCE loss [ 32 ] for embedding alignment.

	$q^{t}\to c^{i}$			$q^{t}\to c^{t}$	$q^{t}\to(c^{i},c^{t})$		$q^{i}\to c^{t}$			$q^{i}\to c^{i}$	$(q^{i},q^{t})\to c^{t}$		$(q^{i},q^{t})\to c^{i}$		$(q^{i},q^{t})\to(c^{i},c^{t})$
Methods	VN	COCO	F200K	WebQA	EDIS	WebQA	VN	COCO	F200K	NIGHTS	OVEN	InfoS	FIQ	CIRR	OVEN	InfoS	Avg.
	R@5	R@5	R@10	R@5	R@5	R@5	R@5	R@5	R@10	R@5	R@5	R@5	R@10	R@5	R@5	R@5
CLIP-L [ 35 ]	43.3	61.1	6.6	36.2	43.3	45.1	41.3	79.0	7.7	26.1	24.2	20.5	7.0	13.2	38.8	26.4	32.5
SigLIP [ 50 ]	30.1	75.7	36.5	39.8	27.0	43.5	30.8	88.2	34.2	28.9	29.7	25.1	14.4	22.7	41.7	27.4	37.2
BLIP [ 21 ]	16.4	74.4	15.9	44.9	26.8	20.3	17.2	83.2	19.9	27.4	16.1	10.2	2.3	10.6	27.4	16.6	26.8
BLIP2 [ 20 ]	16.7	63.8	14.0	38.6	26.9	24.5	15.0	80.0	14.2	25.4	12.2	5.5	4.4	11.8	27.3	15.8	24.8

Table 1: Comparison on the M-BEIR test set. We report Recall@5 for the majority of datasets, with the exception of FashionIQ (FIQ) [ 45 ] and Fashion200K (F200K) [ 13 ] where Recall@10 is utilized. The first row indicates the specific retrieval task type (e.g.

q^{t}\to c^{i}

for text-to-image retrieval). Abbreviations: VN (VisualNews) [ 25 ] , InfoS (InfoSeek) [ 6 ] . The best results are highlighted in bold , and the second-best are underlined . The Improvement row specifically reports the performance gains of TRACE compared to LamRA-Ret.

Method	MSCOCO (Simple)		CIRR (Complex)
Method	QPS $\uparrow$	R@5 $\uparrow$	QPS $\uparrow$	R@5 $\uparrow$
Direct Embedding	15.68	87.40	12.15	53.06
Always CoT (Forced)	4.45	63.90	3.42	54.73
TRACE (Ours)	8.25	89.10	6.48	57.03

Table 2: Efficiency vs. Accuracy Trade-off. We compare standard direct embedding, forced reasoning (Always CoT), and our adaptive approach (TRACE). QPS (Queries Per Second) is measured on a single NVIDIA H20 GPU during the online encoding phase.

	$q^{t}\to c^{i}$			$q^{i}\to c^{t}$			$(q^{i},q^{t})\to c^{i}$		$q^{\text{dialog}}\to c^{i}$	$(q^{i}\oplus q^{t})\to c^{i}$		ITM
Methods	Share4V	Urban ∗	Flickr	Share4V	Urban ∗	Flickr	CIRCO ∗	GeneCIS ∗	VisD ∗	VIST	MT-FIQ ∗	CC-Neg	Sugar-Crepe ∗
	R@1	R@1	R@1	R@1	R@1	R@1	MAP@5	R@1	R@1	R@1	R@5	Acc.	Acc.
CLIP-L [ 35 ]	84.0	52.8	67.3	81.8	68.7	87.2	4.0	13.3	23.7	0.6	17.7	66.7	73.0
Long-CLIP-L [ 51 ]	95.6	86.1	76.1	95.8	82.7	89.3	5.7	16.3	37.9	1.1	18.5	76.3	80.9
UniIR-CLIP [ 43 ]	85.8	75.0	78.7	84.1	78.4	94.2	12.5	16.8	26.8	0.6	39.4	79.9	80.3
E5-V [ 17 ]	86.7	84.0	79.5	84.0	82.4	88.2	24.8	18.5	54.6	10.0	19.2	83.2	84.7

Table 3: Comparison with unseen datasets (Zero-shot). The columns indicate task types, covering text query (

q^{t}

), image query (

q^{i}

), dialog query (

q^{\text{dialog}}

), and multi-interleaved queries. Abbreviations: Share4V (ShareGPT4V) [ 5 ] , Urban (Urban-1k) [ 51 ] , VIST (Visual Storytelling) [ 10 ] , VisD (Visual Dialog) [ 8 ] , and MT-FIQ (Multi-round FashionIQ) [ 45 ] .

*

indicates domain shifts from training data. We report R@1 for most datasets, MAP@5 for CIRCO [ 3 ] , R@5 for MT-FIQ, and Accuracy for ITM tasks. The best results are highlighted in bold , and the second-best are underlined . The Improvement row reports the performance gain of TRACE specifically compared to LamRA-Ret.

Pooling Position	R@1	R@5	R@10
Last Token ( <\|emb\|> )	24.46	52.25	63.81
Structure </think>	24.36	53.06	63.93
Pre-Token (Ours)	28.37	57.03	68.30

Table 4: Ablation on feature extraction position. The pre-token strategy provides the optimal semantic bottleneck.

Data Components	R@1	R@5	R@10
Direct Training (None)	24.52	53.15	63.10
<answer> Only	25.99	53.67	64.48
<reasoning> Only	26.38	53.41	64.60
Full CoT	28.37	57.03	68.30

Table 5: Ablation on reasoning data components. Full reasoning traces yield the highest retrieval accuracy.

Query Side	Candidate Side	R@1	R@5	R@10
Full CoT	Full CoT	7.75	18.90	26.50
Full CoT	<answer>	8.61	19.78	26.88
Full CoT	None (Raw)	28.37	57.03	68.30

Table 6: Ablation on asymmetric reasoning. Applying reasoning to the candidate side causes severe performance degradation.

Figure 4: Visualization of Adaptive Activation. (Left) In-Domain Retrieval: TRACE dynamically toggles between a direct path and a reasoning path based on query complexity. (Right) Zero-Shot Generalization: The adaptive behavior effectively transfers to unseen domains and novel constraints.

Panel of Figure 5 — Figure 5: Examples of task-specific prompts used in M-BEIR-CoT.

Rank (r)	Alpha ( $\alpha$ )	R@1	R@5	R@10
16	32	26.67	56.31	66.45
32	64	26.9	56.19	67.29
128	256	28.37	57.03	68.30
128	512	27.48	56.74	67.84
256	512	28.19	56.81	68.02

Table 7: Ablation on LoRA hyperparameters (r and alpha).

$\lambda_{\text{gen}}$	R@1	R@5	R@10
0.0	24.52	53.15	63.10
0.5	26.89	55.92	66.80
1.0	28.37	57.03	68.30
2.0	27.01	55.74	66.91

Table 8: Ablation on loss weights.

Task	Dataset	Domain	# Train	# Dev	# Test	# Pool
$q^{t}\rightarrow c^{i}$	VisualNews	News	99K	20K	20K	542K
	MSCOCO	Misc.	100K	24.8K	24.8K	5K
	Fashion200K	Fashion	15K	1.7K	1.7K	201K
$q^{t}\rightarrow c^{t}$	WebQA	Wiki	16K	1.7K	2.4K	544K
$q^{t}\rightarrow(c^{i},c^{t})$	EDIS	News	26K	3.2K	3.2K	1M
$q^{t}\rightarrow(c^{i},c^{t})$	WebQA	Wiki	17K	1.7K	2.5K	403K

Table 9: Summary of the Simple Subset in M-BEIR-CoT.

Task	Dataset	Domain	# Train	# Dev	# Test	# Pool
$(q^{i},q^{t})\rightarrow c^{i}$	CIRR	Daily	24.9K	-	-	21K
$(q^{i},q^{t})\rightarrow c^{i}$	FashionIQ	Fashion	15.6K	-	-	18K
$(q^{i},q^{t})\rightarrow c^{t}$	InfoSeek	Wiki	241K	3.4K	6.5K	1M
$(q^{i},q^{t})\rightarrow c^{t}$	OVEN	Wiki	294K	11.7K	23.3K	1M

Table 10: Summary of the Reasoning Subset in M-BEIR-CoT.

Dataset	Image Source	Task	Query Format	Candidate Format
ShareGPT4V	SA-1B	$q^{t}\to c^{i}$	<long text>	<image>
ShareGPT4V	SA-1B	$q^{i}\to c^{t}$	<image>	<long text>
Urban-1K	MSCOCO	$q^{t}\to c^{i}$	<long text>	<image>
Urban-1K	MSCOCO	$q^{i}\to c^{t}$	<image>	<long text>
Flickr30K	Flickr	$q^{t}\to c^{i}$	<short text>	<image>
Flickr30K	Flickr	$q^{i}\to c^{t}$	<image>	<short text>

Table 11: Summary of the Unseen Dataset.

Task	Dataset	Instruction
$q^{t}\to c^{i}$	VisualNews	Identify the news-related image in line with the described event.
		Display an image that best captures the following caption from the news.
		Based on the caption, provide the most fitting image for the news story.
		I want you to retrieve an image of this news caption.
	MSCOCO	Find me an everyday image that matches the given caption.
		Identify the image showcasing the described everyday scene.

Table 12: Summary of the M-BEIR instructions.

Figure 6: Qualitative visualization of Universal Multimodal Retrieval tasks performed by TRACE.

Multimodal IR

Multimodal information retrieval finds relevant images, videos, documents, audio, or other media using queries that may combine different modalities.

See all 122 papers

PLUME: Latent Reasoning Based Universal Multimodal Embedding
Apr 2, 2026Chenwei He, Xiangzhao Hao, Tianyu Yang +6Multimodal Embedding Multimodal IR
Southeast University · Institute of Automation, Chinese Academy of Sciences · University of Chinese Academy of Sciences
Reason What Matters: Retrieval-Grounded Reasoning for Universal Multimodal Embeddings
Sep 14, 2026Mingzhou Jiang, Peixi Wu, Hang Cheng +7Multimodal IR Efficient Multimodal Inference
Tsinghua Shenzhen International Graduate School, Tsinghua University · School of Artificial Intelligence and Data Science, USTC · Kuaishou Technology +1
Think When Needed: Adaptive Reasoning-Driven Multimodal Embeddings with a Dual-LoRA Architecture
May 14, 2026Longxiang Zhang, Weilong Dai, Guanghao Zhang +2Efficient Multimodal Inference Multimodal Embedding
Alibaba Group

Apr 2, 2026J/K move · Enter open · S save

PLUME: Latent Reasoning Based Universal Multimodal Embedding

Chenwei He, Xiangzhao Hao, Tianyu Yang, Yuxiang Ma+5

Southeast University · Institute of Automation, Chinese Academy of Sciences · University of Chinese Academy of Sciences

Multimodal Embedding Multimodal IR Efficient Multimodal Inference Continuous Latent Reasoning

Open paper AskPDF arXiv

Universal multimodal embedding (UME) maps heterogeneous inputs into a shared retrieval space with a single model. Recent approaches improve UME by generating explicit chain-of-thought (CoT) rationales before extracting embeddings, enabling multimodal large language models to better infer complex query intent. However, explicit CoT incurs substantial inference overhead and can compress rich multimodal evidence into a narrow textual bottleneck. We propose PLUME, a latent reasoning framework that advances UME by replacing verbalized CoT with a short autoregressive rollout of continuous latent states. To support diverse multimodal queries, PLUME further introduces a semantic-anchor-guided transition adapter that steers latent rollout along different reasoning trajectories under the same fixed computation budget. To stabilize training, PLUME adopts a progressive explicit-to-latent curriculum that uses verbalized reasoning only as a temporary training scaffold and gradually transfers this behavior into hidden-state computation, eliminating explicit CoT at inference. On the 78-task MMEB-v2 benchmark, PLUME outperforms strong explicit-CoT UME baselines while reducing reasoning from hundreds of generated tokens to fewer than 10 latent steps, delivering over 30x faster inference. PLUME is especially well suited to retrieval settings where relevant evidence is dense, structurally complex, and difficult to organize through verbalized intermediate rationales, such as video and visual document retrieval. These results show that structured latent computation can preserve the benefits of intermediate reasoning without the overhead of explicit rationale generation, providing a stronger and more efficient paradigm for practical retrieval systems.

Abstract

Multimodal IR

Explore similar work

PLUME: Latent Reasoning Based Universal Multimodal Embedding

Reason What Matters: Retrieval-Grounded Reasoning for Universal Multimodal Embeddings

Think When Needed: Adaptive Reasoning-Driven Multimodal Embeddings with a Dual-LoRA Architecture