GEM: A Generative Embedding Model Bridging Reasoning and Retrieval

AI summary

Overview

• Research area: Natural Language Processing / information retrieval, specifically dense retrieval with large language models, reasoning-intensive retrieval, and instruction-following retrieval.
• Technical level: Advanced. The paper assumes familiarity with dense bi-encoders, contrastive (InfoNCE) training, causal language modelling, KV caching, and standard IR metrics such as nDCG@10, MAP@1000 and p-MRR.
• Scope: The paper proposes GEM, a single 4B-parameter model trained from Qwen3-4B-Instruct-2507 that reasons over a query and then encodes that reasoning into a retrieval embedding, and evaluates it on BRIGHT, FollowIR and InstructIR.

What This Paper Is About

Conventional retrievers rank documents by lexical and/or semantic similarity between a query and documents, so they struggle when the relevant document does not share the wording or surface meaning of the query. Users of modern LLMs increasingly express complex information needs as natural-language instructions, and this creates a growing gap between how needs are expressed and how retrievers interpret them. GEM addresses this by letting one model first reason about user intent and relevance criteria, then append a dedicated embedding token to encode the enriched context for retrieval.

Key Contributions

1. A unified generative embedding model. GEM integrates reasoning and embedding within a single model using a generate-then-encode paradigm, jointly trained with contrastive and causal language modelling objectives rather than as a separate reasoner-plus-retriever pipeline.
2. A tailored data generation strategy. The authors synthesise document pairs conditioned on validated reasoning: candidate responses are filtered by an LLM-based relevance classifier, positives are generated to align with the reasoning, and hard negatives share similar topics but contain subtle contradictions, discouraging surface-level matching.
3. Empirical evidence that generation augments retrieval. Experiments on reasoning-intensive (BRIGHT) and instruction-following (FollowIR, InstructIR) retrieval show GEM outperforming its non-reasoning variant and matching baselines built on substantially larger models, with further gains from test-time compute scaling.
4. A claim of novelty. The authors state that, to their knowledge, GEM is the first embedding model to leverage its own generative capabilities to produce reasoning-aligned embeddings.

Main Findings

• Strongest single-model result on BRIGHT: GEM reaches an average nDCG@10 of 29.1 on BRIGHT, ahead of ReasonIR-8B (24.4), Qwen3-4B-Instruct (21.4), GritLM-7B (20.9), Promptriever (20.0), BM25 (14.8) and Contriever (11.1) in the single-model setting.
• Large gains on theorem-based tasks: GEM improves the theorem-based subset average nDCG@10 over its embedding-only variant Qwen3-4B-Instruct from 19.8 to 32.0 using the same backbone.
• A documented weak spot: GEM underperforms on the Pony programming language task (2.8), a result the authors attribute to an out-of-domain problem discussed in Appendix B.6; Promptriever (1.7) and Qwen3-4B-Instruct (3.8) show similar weakness.
• Competitive as a pure encoder of others' reasoning: When GEM instead encodes reasoning generated by GPT-4 as a bi-encoder, it achieves an average nDCG@10 of 30.0 on BRIGHT, competitive with ReasonIR-8B (29.9). In that GPT-4-reasoner setting the average response length was approximately 37% longer than GEM's own generations.
• Test-time compute scaling works, then saturates: Prompting GEM to generate roughly n ∈ {64, 256, 512, 1024, 2048} words improves average BRIGHT nDCG@10, peaking at 30.1 with n = 1024, with gains saturating for long generations.
• Encoding latency stays flat with KV cache reuse: Reusing the KV cache from generation keeps per-query encoding time stable across the evaluated response lengths, whereas re-encoding the same sequence without cache reuse shows encoding time increasing sharply for longer sequences.
• Instruction-following gains over the same backbone: GEM achieves p-MRR of +11.7 versus +6.8 for Qwen3-4B-Instruct, and Robustness@10 of 54.8 versus 46.2, with average FollowIR score 25.5 versus 24.0. GEM's p-MRR matches larger models including Promptriever (+11.2) and FollowIR-7B (+12.2).
• Query expansion is not a substitute: HyDE and Query2Doc often improve nDCG and MAP but their effect on p-MRR is inconsistent, and both degrade p-MRR for Promptriever. Directly appending GEM's response to other retrievers can also hurt, as seen for BGE-large (average p-MRR falling from -2.4 to -3.0), RepLLaMA (from -3.1 to -2.3) and ReasonIR-8B (from -0.4 to -2.9).
• Ablation results: Removing hard queries from ReasonIR drops FollowIR p-MRR from +11.7 to +8.5; removing document generation drops BRIGHT nDCG@10 from 29.1 to 25.8 while p-MRR stays at +11.7; removing the generation loss ("GEM w/o generation") drops p-MRR to +9.5 and BRIGHT to 21.0.
• Trade-off from unification: An embedding-only variant trained on identical data (reusing GEM's reasoning) scores +12.5 p-MRR and 30.0 BRIGHT nDCG@10 — slightly higher than unified GEM (+11.7 and 29.1) — but suffers catastrophic forgetting and produces repetitive or mixed-language content.
• Generation ability is reported as retained: On generative tasks reported in Appendix B.1, GEM built on Qwen3-4B-Instruct scores 84.3 on IFEval, 65.1 on MMLU, 48.0 on ARC-Challenge and 82.1 EM on GSM8K; HumanEval Pass@1 is not reported in the provided content.

Methodology in Plain English

Retrieval is reframed as two steps inside one model. Given a query, GEM is given a meta-instruction to reason about what the user wants and what would count as relevant, and it writes out that reasoning. The model then appends a special token, <|embed|>, whose hidden state becomes the query embedding; documents are encoded the same way after being given a short document-side prompt. Unlike earlier work that uses the end-of-sequence token for pooling, GEM reserves EOS for generation and uses its own embedding token, which is never predicted during decoding and so is excluded from the language modelling loss.

Training data is built in two stages. For each query the team samples K = 8 candidate reasonings at temperature 1.0 and keeps only those for which an LLM classifier confirms the original positive document still looks relevant, discarding queries with no surviving response. Documents are then regenerated conditioned on the reasoning, using Llama-3.1-8B-Instruct with greedy decoding to produce one positive and one hard negative per response — the hard negative being topically similar but subtly contradicting the intent or criteria.

The model is trained with a weighted sum of a causal language modelling loss on the responses and an InfoNCE contrastive loss over prompt-plus-response embeddings versus documents, with cosine similarity as the scoring function. The authors set λ_gen = 0.1 and λ_emb = 1.0, τ = 0.02, a training group size of 2 (one positive, one hard negative) with in-batch negatives gathered across GPUs, 500 steps at effective batch size 512, learning rate 1×10⁻⁵ and 50 warmup steps, using FSDP with CPU offloading, gradient checkpointing and bfloat16 mixed precision. Training took approximately 14 hours on 2 NVIDIA H100 GPUs, and data generation ran on a single NVIDIA H100 using vLLM. The final training set has 370,000 samples: 320,000 from Promptriever (260,000 with reasoning plus 60,000 original non-reasoning samples) and 50,000 reasoning samples built from ReasonIR hard queries.

Why This Matters

The paper argues that the gap between how users phrase needs and how retrievers interpret them is a query-understanding problem, and that a retriever can close that gap itself rather than depending on an upstream reasoning model. It also warns that reasoning-augmented retrieval gains may partly reflect extra lexical and semantic overlap rather than genuine understanding, which is precisely what the p-MRR experiments probe.

Real-world applications implied by the work:

• Instruction-following search: Query-specific criteria such as preferring documents that are easier to understand and avoiding technical jargon, the example used in the paper.
• Reasoning-intensive professional search: Domains such as biology, earth science, economics, psychology, robotics, coding and sustainable living, which form the BRIGHT subsets.
• Theorem and math retrieval: Retrieving scientific theorem question answers using either questions or theorems as queries, where GEM's largest improvements appear.
• Constraint-aware retrieval: Queries where the user specifies what not to retrieve or what requirements a document must meet, which the paper highlights as a limitation of query expansion.

Industry relevance: GEM's KV cache reuse during generate-then-encode keeps query-side encoding time stable as reasoning length grows, and its single-model design avoids maintaining separate reasoner and retriever deployments. The authors note that autoregressive decoding remains expensive, so cost is still a practical constraint.

Future Directions

1. Larger backbones and more data. The limitations section states the authors could not replicate experiments with larger backbones (e.g. 7B parameters) or larger-scale training data, and that generalisation experiments are restricted to backbones up to 4B parameters.
2. Measuring and mitigating hallucination impact. Whether hallucinations in generated documents or in GEM's reasoning can bias retrieval, and how to quantify that, is explicitly left uninvestigated, as is noise propagating from imperfect LLM-based filtering.
3. Decoupled training. The paper trains generation and embedding losses on a shared batch and leaves investigation of a decoupled training strategy, potentially with different batch sizes, to future work.
4. Broader retrieval-involving applications. Extending GEM to retrieval-augmented generation and conversational search is named as valuable future research.

Target Audience

Researchers and practitioners working on dense retrieval, LLM-based embedding models, and reasoning-augmented or instruction-following search. It is most useful to readers already comfortable with contrastive retrieval training and IR evaluation metrics, and to engineers considering whether to collapse a separate reasoner-plus-retriever pipeline into one generative embedding model under constrained compute.

Authors’ abstract

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. \zhili{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.