One Model Is Enough: Native Retrieval Embeddings from LLM Agent Hidden States

Bo Jiang

cs.CLMar 9, 2026

One Model Is Enough: Native Retrieval Embeddings from LLM Agent Hidden States

Authors: Bo Jiang

Organizations: Temple University

Abstract

LLM agents that retrieve external knowledge typically generate a search query as text, then run a separate embedding model to encode it into a vector. This two-model pipeline adds infrastructure complexity and latency, yet is redundant: the LLM already encodes the full conversational context in its hidden states. We propose equipping LLM agents with native retrieval capability by adding a lightweight projection head that maps hidden states directly into the embedding space, eliminating the need for a separate embedding model. Trained with a combination of alignment, contrastive, and rank distillation losses, our method retains 97% of baseline retrieval quality while enabling the LLM agent to search with its own representations. Experiments on the QReCC conversational search benchmark show competitive Recall@10 and MRR@10 compared to the standard generate-then-encode pipeline, with systematic ablations confirming the contribution of each loss component.

Figures & tables

	Recall	MRR	nDCG	Latency
	@10	@10	@10	(p50)
Baseline	0.637	0.329	0.402	43.5ms
Ours	0.607	0.293	0.367	2.0ms
$\Delta$	–3.0%	–3.6%	–3.5%	21.8 $\times$
95% CI	[–4.7, –1.4]	[–4.7, –2.4]	[–4.5, –2.4]	—

Table 1: Main results on QReCC test set. Our projection head achieves near-parity retrieval quality while eliminating the embedding model. Latency is measured per query on an NVIDIA RTX PRO 6000. McNemar’s test:

\chi^{2}=12.21

p=0.0005

Configuration	$\mathcal{L}_{a}$ / $\mathcal{L}_{c}$ / $\mathcal{L}_{r}$	R@10	MRR
Align only	1.0 / 0.0 / 0.0	0.567	0.273
Contrastive only	0.0 / 1.0 / 0.0	0.498	0.231
Rank distill only	0.0 / 0.0 / 1.0	0.001	0.000
Align + Contra	0.5 / 0.5 / 0.0	0.582	0.281
Align + Rank	0.5 / 0.0 / 0.5	0.579	0.272
All three	0.5 / 0.5 / 0.5	0.595	0.288

Table 2: Loss ablation results.

\mathcal{L}_{a}

: alignment,

\mathcal{L}_{c}

: contrastive,

\mathcal{L}_{r}

: rank distillation. The last row uses extended training (80 epochs, lr

2\times 10^{-4}

Epochs / LR	R@10	MRR	Status
30 / $3\times 10^{-4}$	0.595	0.288	converged
50 / $3\times 10^{-4}$	0.552	0.261	converged
80 / $2\times 10^{-4}$	0.607	0.293	converged
30 / $5\times 10^{-4}$	0.001	0.000	collapsed

Table 3: Effect of training epochs and learning rate. Too-high learning rate causes training collapse; longer training with lower learning rate yields the best results.

Config	$\lambda_{a}$ / $\lambda_{r}$	$K$	R@10	MRR
Align-heavy	1.0 / 0.5	128	0.591	0.280
RD-heavy	0.5 / 1.0	128	0.516	0.236
Top-256	0.5 / 0.5	256	0.584	0.284

Table 4: Rank distillation weight and top-

K

sensitivity. All use

\lambda_{c}=0.0

, 30 epochs, lr

3\times 10^{-4}

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Group	Experiment	R@10	$\Delta$ R@10	95% CI	MRR	$\Delta$ MRR	95% CI
C	recipe_e80_lr2e4	0.607	–0.030	[–0.047, –0.014]	0.293	–0.036	[–0.047, –0.024]
A	All three losses	0.595	–0.042	[–0.058, –0.025]	0.288	–0.040	[–0.052, –0.029]
B	Align-heavy RD	0.591	–0.046	[–0.063, –0.029]	0.280	–0.049	[–0.060, –0.037]
B	RD top-256	0.584	–0.053	[–0.070, –0.037]	0.284	–0.045	[–0.057, –0.033]
A	Align + Contra	0.582	–0.055	[–0.072, –0.038]	0.281	–0.047	[–0.059, –0.036]
A	Align + Rank	0.579	–0.058	[–0.075, –0.041]	0.272	–0.057	[–0.069, –0.045]

Appendix

Table 5: Full results for all 12 configurations across three experiment groups. Group A: loss ablation (30 epochs, lr

3\times 10^{-4}

). Group B: loss weight tuning (30 epochs, lr

3\times 10^{-4}

\lambda_{c}=0

). Group C: training recipe variations. Baseline: Recall@10 = 0.637, MRR@10 = 0.329.

Agentic Retrieval

Agentic retrieval lets AI agents iteratively search, inspect, and verify evidence to answer complex information needs.

See all 93 papers

GEM: A Generative Embedding Model Bridging Reasoning and Retrieval
Aug 13, 2026Zhili Shen, Craig MacdonaldText Embedding Models
University of Glasgow
Latent Abstraction for Retrieval-Augmented Generation
Apr 20, 2026Ha Lan N. T, Minh-Anh Nguyen, Dung D. LeRetrieval-Augmented Generation Adaptive RAG
Center for AI Research, VinUniversity, Vietnam
Retrieval as Reasoning: Self-Evolving Agent-Native Retrieval via LLM-Wiki
May 25, 2026Haoliang Ming, Feifei Li, Xiaoqing Wu +1Multi-Hop QA Retrieval-Augmented Generation
WeChat, Tencent Inc., Beijing, China

Aug 13, 2026J/K move · Enter open · S save

GEM: A Generative Embedding Model Bridging Reasoning and Retrieval

Zhili Shen, Craig Macdonald

University of Glasgow

Text Embedding Models

Open paper AskPDF arXiv

Modern LLMs excel at reasoning and instruction following, enabling users to express complex and diverse information needs. However, conventional retrievers largely rely on surface-level matching between queries and documents, resulting in a growing gap between how users express their needs and how retrievers interpret them. In this paper, we present GEM, a generative embedding model that augments retrieval through its own knowledge by explicitly reasoning about user intent and relevance criteria. GEM unifies generation and embedding within a single model: it first reasons over the query, then appends an embedding token to encode the enriched context for retrieval. Evaluated on reasoning-intensive and instruction-following retrieval tasks, GEM demonstrates the effectiveness of its reasoning-augmented retrieval, outperforming its non-reasoning variant and matching baselines using substantially larger models. Furthermore, GEM's generative nature allows test-time compute scaling via prompting to further enhance retrieval performance. Our code is available at: https://anonymous.4open.science/r/GEM.

Abstract

Agentic Retrieval

Explore similar work

GEM: A Generative Embedding Model Bridging Reasoning and Retrieval

Latent Abstraction for Retrieval-Augmented Generation

Retrieval as Reasoning: Self-Evolving Agent-Native Retrieval via LLM-Wiki