PaperlayerView exact revision on arXiv

Source-linked AI summary

O1 Embedder: Let Retrievers Think Before Action

Ruiran Yan, Zheng Liu, Defu Lian

arXiv:2502.07555v2cs.CL

TL;DR

Dense retrievers remain challenged by unseen-domain and reasoning-intensive tasks. O1 Embedder generates retrieval-oriented thoughts and jointly learns thought generation and dense retrieval; the authors report improved performance across in-domain and out-of-domain evaluations. The approach also depends on thoughts that do not simply reproduce the target document.

  • Problem

    Dense retrieval models struggle with zero-shot tasks in unseen scenarios and with complex relationships that require intensive reasoning.

    Source: §1 · p. 1

  • Method

    O1 Embedder uses LLM-explored, retrieval-committee-refined thoughts and jointly fine-tunes a pretrained LLM for thought generation and contrastive dense retrieval.

    Source: §1 · p. 2

  • Results

    O1 Embedder significantly outperforms existing methods, with substantial retrieval improvements across in-domain and out-of-domain evaluations.

    Source: §5 · p. 8

  • Takeaways & Limitations

    The authors conclude that O1 Embedder improves retrieval by combining slow-thinking with embedding capabilities, including on complex retrieval tasks.

    Source: §5 · p. 8

    A target document or its rephrasing cannot serve as a thought, because that would circumvent training and ultimately collapse the embedding task.

    Source: §3.2 · p. 3

Abstract

from arXiv · show

The growing power of large language models (LLMs) has revolutionized how people access and utilize information. Notably, the LLMs excel at performing fine-grained data representation, which facilitates precise retrieval of information. They also generate high-quality answers based on external references, enabling the production of useful knowledge. The recent introduction of reasoning models, like OpenAI O1 and DeepSeek R1, marks another leap forward, highlighting LLMs' ability to think progressively before delivering final answers. This breakthrough significantly improves the ability to address complex tasks, e.g., coding and math proofs. Inspired by this progress, we aim to develop similar capabilities for retrieval models, which hold great promise for tackling critical challenges in the field, including multi-task retrieval, zero-shot retrieval, and tasks requiring intensive reasoning of complex relationships. With this motivation, we propose a novel approach called O1 Embedder, which generates useful thoughts for the input query before making retrieval for the target documents. To realize this objective, we conquer two technical difficulties. First, we design a data synthesis workflow, creating training signals for O1 Embedder by generating initial thoughts from an LLM-expert and subsequently refining them using a retrieval committee. Second, we optimize the training process, enabling a pre-trained model to be jointly fine-tuned to generate retrieval thoughts via behavior cloning and perform dense retrieval through contrastive learning. Our approach is evaluated by comprehensive experiments, where substantial improvements are achieved across 12 popular datasets, spanning both in-domain and out-of-domain scenarios. These results highlight O1 Embedder's remarkable accuracy and generalizability, paving the way for the development of next-generation IR foundation models.

1 Introduction

Dense retrieval still struggles with unseen domains and complex relationships that require reasoning, motivating O1 Embedder to think through a query before producing its retrieval embedding. The paper introduces retrieval-oriented thought synthesis and joint training, and reports improvements across retrieval tasks and out-of-domain settings.

1 Introduction

  • Existing dense retrievers struggle with zero-shot retrieval in unseen domains and with complex relationships that cannot be identified directly from semantic meaning.The paper points to specialized medical or legal retrieval, code retrieval, and multi-hop evidence retrieval as examples.Source: §1 · p. 1
  • O1 Embedder generates query thoughts that uncover hidden information needs, then uses the thoughts with the query to produce a discriminative embedding.The resulting embedding targets relevant documents that are challenging to identify using the query alone.Source: §1 · pp. 1–2
  • An LLM explores initial query thoughts, and a retrieval committee scores their usefulness for retrieving the target document before majority voting selects a golden thought.This exploration-refinement workflow creates long-form thought training data.Source: §1 · p. 2
  • The multi-task method fine-tunes a pre-trained LLM with supervised fine-tuning for thought generation and contrastive learning for retrieval embeddings.Loss and workflow optimizations unify the tasks cost-efficiently.Source: §1 · p. 2
  • O1 Embedder improves over existing methods across retrieval tasks, especially reasoning-intensive tasks, generalizes to out-of-domain scenarios, and maintains performance across LLM backbones.The experiments span Llama, Mistral, and Qwen backbones.Source: §1 · p. 2

2 Related Work

Dense retrieval has advanced through pretrained models and training strategies, while general-purpose retrievers increasingly use LLM backbones. Existing methods still face zero-shot and reasoning-intensive retrieval challenges; this section also situates the paper’s thinking-before-retrieval workflow alongside LLM reasoning approaches.

2.1 Dense Retrieval

  • Dense retrieval expanded from fine-tuned pretrained models to multilingual and multimodal models, with adaptation, hard-negative mining, larger batches, and cross-encoder distillation improving performance.The related-work discussion also describes LLM-powered retrievers and general-purpose systems.Source: §2.1 · p. 2
  • Despite these advances, existing retrievers remain challenged by zero-shot retrieval in new domains and tasks requiring intensive reasoning to identify semantic relationships.Figure 1 depicts O1 Embedder generating thoughts about a query before producing its dense-retrieval embedding.Source: §2.1 · p. 2, Figure 1 · p. 2

2.2 LLMs’ Reasoning Ability

  • LLM reasoning methods include stepwise Chain-of-Thought, majority-vote Self-Consistency, tree-structured ToT, and DAG-structured GoT, alongside test-time scaling.The paper describes reasoning-capable LLMs as generating structured reasoning before a final solution.Source: §2.2 · p. 3
  • Although LLM reasoning can help address traditional retrieval challenges, embedding models have primarily focused on discriminative representations, leaving their reasoning capabilities largely unexplored.Source: §2.2 · p. 3

3 Method

The method formulates retrieval around thought-augmented embeddings, synthesizes query-thought-document training data, and jointly trains thought generation and dense retrieval.

3 Method

  • The method first formulates O1 Embedder, then develops thought-data production and multi-task training.Source: §3 · p. 3

3.1 Problem Formulation

  • O1 Embedder addresses zero-shot and complex retrieval limits by generating query thoughts and combining them with the query in a unified embedding optimized for retrieval.Source: §3.1 · p. 3

3.2 Data Production

  • An exploration-refinement pipeline generates candidate thoughts, filters them through retriever comparisons and majority voting, and creates query-thought-document triplets.Relevant documents or their rephrasings are unsuitable as thoughts because they would circumvent and collapse embedding training.Source: §3.2 · pp. 3–4, Figure 2 · p. 4

3.3 Multi-Task Training

  • A pre-trained LLM is fine-tuned for thought generation through behavior cloning and for document discrimination through contrastive learning.Source: §3.3 · p. 4

3.3.1 Behavior cloning.

  • Behavior cloning fine-tunes the model on query-thought-document triplets with next-token prediction loss over thought tokens.Source: §3.3.1 · p. 4

3.3.2 Contrastive learning.

  • A dedicated <emb> token separates embedding from generation, while contrastive losses train original and thought-augmented query embeddings together.The negative samples include in-batch and hard negatives introduced by a pre-trained embedder.Source: §3.3.2 · p. 4

3.3.3 Joint training.

  • Because parallel tasks constrain batch size through GPU memory, shared sample encodings let generation add almost no memory and permit larger batches.Source: §3.3.3 · pp. 4–5, Figure 3 · p. 5

3.4 Retrieval

  • At retrieval, multiple query thoughts are independently encoded, mean-pooled into a thought-augmented embedding, and used to retrieve top-N documents by similarity.Source: §3.4 · p. 5

4 Experiment

The experiments assess O1 Embedder against established retrievers on in-domain and out-of-domain tasks, including tests of its thinking mechanism. The reported results show gains across these settings, with the strongest overall result summarized as improved retrieval performance from thinking.

4.1.1 Datasets.

  • Training uses thought-augmented MS MARCO queries; evaluation covers three in-domain and nine out-of-domain datasets, with no task-specific fine-tuning for out-of-domain tests.The out-of-domain suite uses nDCG@10; MS MARCO-related tasks use MRR@10 and Recall@1k.Source: §4.1.1 · p. 5

4.1.2 Baselines.

  • Baselines include BM25, BERT-based retrievers, and LLM-based RepLLaMA and Promptriever; several models with extra training data or evaluation-task overlap are excluded.The included LLM baselines share the Llama-2-7B backbone used by the default O1 Embedder setting.Source: §4.1.2 · pp. 5–6

4.1.3 Implementation Details.

  • The primary evaluation uses a Llama-2-7B backbone, LoRA fine-tuning, 15 hard negatives per query, and up to 256 tokens for thought generation.Candidate thoughts are generated with Llama-3.1-70B-Instruct and refined using a committee of four retrievers.Source: §4.1.3 · p. 6

4.2.1 In-domain Performance.

  • On MS MARCO, O1 Embedder reaches MRR@10 43.1 (+1.9% improvement) and Recall@1k 99.5; on DL’19 and DL’20, it scores nDCG@10 75.3 (+1.0% improvement) and 74.4 (2.1% improvement).The authors report that it outperforms all other models across the in-domain metrics and datasets.Source: §4.2.1 · p. 6, Table 1 · p. 6

4.2.2 O.O.D. Performance.

  • Across nine out-of-distribution datasets, O1 Embedder averages a 2.3% improvement over the baseline and leads except on FiQA, where it trails CPT-XL 175B.The authors highlight complex-reasoning tasks such as HotPotQA and CosQA.Source: §4.2.2 · p. 6, Table 2 · p. 7
  • The thinking mechanism improves NQ by 3.9% and HotPotQA by 3.0%, but gains are smaller on TREC-Covid, FiQA, and SciFact.The authors attribute the smaller specialized-domain improvements to potentially noisy or hallucinated generated content and identify thought refinement as future work.Source: §4.2.2 · pp. 6–7

4.2.3 Impact of Thought.

  • With thinking enabled, O1 Embedder gains MRR@10 +1.4% on MS MARCO, nDCG@10 +1.6% on DL’19 and +2.1% on DL’20, and 2.3 points over O1 Embedder w/o T in zero-shot average nDCG@10.The no-thinking variant achieves results comparable to RepLLaMA.Source: Table 1 · p. 6, §4.2.3 · p. 7

4.2.4 Summary for main result.

  • O1 Embedder significantly outperforms BERT- and LLM-based models in-domain, generalizes across out-of-domain tasks, and improves retrieval over the vanilla embedding model through thinking.The authors report the generalization advantage as especially pronounced on challenging problems.Source: §4.2.4 · p. 7

4.3 Extended Analysis

The extended analysis tests whether joint training helps a retriever use generated thoughts, then examines robustness across architectures and model sizes and illustrates a multi-hop case.

4.3.1 Joint training.

  • Jointly trained O1 Embedder improves retrieval across all datasets, while adding GPT-4o-mini thoughts to RepLLaMA yields a slight 0.1 decrease in overall performance.The authors attribute RepLLaMA's weaker results to potential noise in generated thoughts, particularly in specialized domains.Source: Table 3 · p. 7, §4.3.1 · p. 8

4.3.2 Robustness.

  • Across Llama, Mistral, and Qwen backbones, O1 Embedder maintains strong in-domain and out-of-domain retrieval performance.The analysis evaluates different pretrained architectures and reports effectiveness regardless of differences in pretrained capabilities.Source: §4.3.2 · p. 8, Table 4 · p. 8
  • As Qwen2.5 models shrink, effectiveness drops across evaluated datasets; 3B performs similarly to 7B, while 1.5B with thought is comparable to 7B RepLLaMA.The authors partly attribute larger models' advantage to better generalization from training on more extensive pretraining corpora.Source: §4.3.2 · p. 8

4.3.3 Case Study.

  • For a multi-hop query about Martin Luther King’s age upon college admission, the generated thought surfaces that he entered Morehouse at 15 in 1944.The authors say the thought also adds college and attendance-period context relevant to retrieval.Source: Table 5 · p. 8, §4.3.3 · p. 8

5 Conclusion

O1 Embedder combines slow-thinking before retrieval with joint training of thinking and embedding, and the authors report improved performance across in-domain and out-of-domain benchmarks.

5 Conclusion

  • O1 Embedder performs slow-thinking before retrieval to uncover query information needs and identify semantic relevance for complex retrieval tasks.The approach is intended to help with relevant documents that are difficult to identify from the query alone.Source: §5 · p. 8
  • A tailored data-production method generates long-form thoughts optimized for retrieval utility, while multi-task training establishes thinking and embedding capabilities.Source: §5 · p. 8
  • Across popular evaluation benchmarks, O1 Embedder significantly outperforms existing methods in both in-domain and out-of-domain retrieval.The authors identify multi-round reasoning, lightweight distillation, and other retrieval tasks as future directions.Source: §5 · p. 8

A Prompts

The appendix shows teacher-model prompts and responses used for training-data generation, including an instruction added to reduce query duplication.

A Prompts

  • Table 6 presents an example of the teacher model’s prompt and response for data generation.The table omits long documentation for presentation.Source: Table 6 · p. 12
  • For in-context examples, the authors randomly select three query–positive-document pairs from the training set.Source: A Prompts · p. 12
  • Because many generated responses duplicated the query, the authors added an instruction not to repeat it.Source: A Prompts · p. 12

B Detailed Scores of Different Base Models

The appendix compares zero-shot scores across base models and reports that model scale matters more for overall performance than differences among same-scale models.

B Detailed Scores of Different Base Models

C Algorithm

Algorithm 1 builds query–thought–document training examples by generating candidate thoughts and selecting them with a retrieval committee. The resulting enhanced dataset retains the original queries and their positive documents alongside the selected thoughts.

C Algorithm

  • The data-production inputs include a query–document dataset, a teacher model, retrieval models, an example count, and a candidate-thought count.Source: C Algorithm · p. 12
  • For each query–document pair, the algorithm samples examples, generates candidate thoughts, scores them against the target document with retrieval models, and selects a final thought by majority voting.The selected thought is added to the enhanced dataset with the original query and its positive document.Source: C Algorithm · pp. 12–13
  • The algorithm’s prompt is constructed using sampled examples and specific instructions before candidate thoughts are generated.Source: C Algorithm · p. 12

D Additional Case analyses & Attention

The attention analysis examines how the <emb> token attends to generated-thought tokens. Figure 4 highlights patterns shared with the positive document, including “Japan,” which is absent from the original query.

D Additional Case analyses & Attention

E Complex Retrieval Example

A CosQA example illustrates O1 Embedder’s response to a code-retrieval query about checking whether a directory is writable. Its generated thought combines a direct solution with an additional directory-validation approach.

E Complex Retrieval Example

Loading 2502.07555v2…