Source-linked AI summary
A Neural Corpus Indexer for Document Retrieval
Yujing Wang, Yingyan Hou, Haonan Wang, Ziming Miao, Shibin Wu, Hao Sun, Qi Chen, Yuqing Xia, Chengmin Chi, Guoshuai Zhao, Zheng Liu, Xing Xie, Hao Allen Sun, Weiwei Deng, Qi Zhang, Mao Yang
TL;DR
Existing retrieval systems make their indexes difficult to optimize directly for the final retrieval target, motivating an end-to-end alternative. NCI is a sequence-to-sequence model that generates relevant document identifiers using tailored decoding and indexing techniques. It outperforms reported baselines on NQ320k and TriviaQA, while large-scale deployment remains constrained by model capacity, inference speed, and index updates.
Problem
Existing index-retrieve pipelines are difficult to optimize directly for the final retrieval target, while retrieval recall is crucial because online ranking can process only limited candidate sets.
Method
NCI unifies training and indexing in a sequence-to-sequence network that encodes queries and generates relevant document identifiers, using query generation, semantic identifiers, a prefix-aware weight-adaptive decoder, and consistency regularization.
Results
+21.4% relative enhancement for Recall@1 on NQ320k and +16.8% relative enhancement for R-Precision on TriviaQA, compared to the best baseline method.
Takeaways & Limitations
NCI demonstrates that an end-to-end differentiable document-retrieval model can outperform inverted-index and dense-retrieval solutions.
Takeaways & Limitations
Large-scale deployment is constrained by the need for greater model capacity, faster inference, and practical updates when new documents are added.
Abstract
from arXiv · showhide
Current state-of-the-art document retrieval solutions mainly follow an index-retrieve paradigm, where the index is hard to be directly optimized for the final retrieval target. In this paper, we aim to show that an end-to-end deep neural network unifying training and indexing stages can significantly improve the recall performance of traditional methods. To this end, we propose Neural Corpus Indexer (NCI), a sequence-to-sequence network that generates relevant document identifiers directly for a designated query. To optimize the recall performance of NCI, we invent a prefix-aware weight-adaptive decoder architecture, and leverage tailored techniques including query generation, semantic document identifiers, and consistency-based regularization. Empirical studies demonstrated the superiority of NCI on two commonly used academic benchmarks, achieving +21.4% and +16.8% relative enhancement for Recall@1 on NQ320k dataset and R-Precision on TriviaQA dataset, respectively, compared to the best baseline method.
1 Introduction
Because ranking is costly and only a limited candidate set can be reranked online, document-retrieval recall is crucial. NCI replaces conventional indexing with an end-to-end sequence-to-sequence model that directly generates relevant document identifiers, achieving strong benchmark improvements while aiming toward unified differentiable search systems.
- Motivation: Document retrieval must achieve high recall because costly ranking models can process only hundreds or thousands of candidates per query online.The retrieval stage supplies candidates for subsequent relevance scoring, making its recall important to overall search effectiveness.
- Existing approaches: Traditional term-based and semantic-based retrieval use inverted indexes or dense representations with ANN search, but each has documented limitations.Term matching misses differently worded semantic equivalents, while ANN methods rely on similarity functions and assumptions about the representation space.
- NCI: NCI is an end-to-end sequence-to-sequence model that encodes a query and generates identifiers for relevant documents, retrieving top candidates through beam search.Training can use ground-truth and augmented query-document pairs, while query generation supplies additional training pairs.
- NCI design: NCI combines query generation, semantic document identifiers, a prefix-aware weight-adaptive decoder, and consistency-based regularization to improve retrieval effectiveness.The semantic identifiers are generated hierarchically, and the regularization is intended to mitigate over-fitting.
- Findings: NCI significantly outperforms inverted-index and dense-retrieval solutions and can achieve better ranking results than ANN-based approaches without an explicit ranking model.The authors also describe reduced dependence among retrieval-system modules and easier deployment and maintenance as potential system-design benefits.
2 Related work
Related work includes sparse retrieval with inverted indexes, dense retrieval with vector similarity and ANN search, autoregressive retrieval that generates identifiers or words, and pretrained-language-model approaches. These lines differ in representation, indexing, and inference mechanisms.
- Sparse retrieval: Sparse retrieval builds inverted indexes using term-matching metrics such as TF-IDF, query likelihood, or BM25.Neural extensions augment documents, learn sparse representations, or estimate context-aware term importance.
- Dense retrieval: Dense retrieval represents queries and documents as dense vectors and uses inner-product or cosine similarity, often with ANN indexes at inference time.Pretrained models such as BERT and RoBERTa are used to obtain representations.
- Autoregressive retrieval: Autoregressive retrieval uses end-to-end generation for tasks including entity linking, direct document-identifier generation, and relevant-word generation.DSI generates document identifiers corresponding to queries, while SEAL generates relevant words using passage n-grams as identifiers.
- Pre-trained language models: Pretrained-language-model retrieval computes document representations offline, represents the query online, and scores query-document pairs with efficient crossing operations.The described serving pattern separates offline document processing from online relevance scoring.
3 Neural corpus indexer
NCI is an end-to-end sequence-to-sequence retriever that generates semantic document identifiers directly from queries. Its design combines document-aware identifiers, generated query-document training pairs, and a prefix-aware decoder with consistency regularization.
- 3 Neural corpus indexer: NCI takes a query as input and generates the most relevant document identifier with a sequence-to-sequence neural network.The encoder produces a query representation, while the decoder outputs document identifiers; inference retrieves top candidates through beam search.
- 3.1 Representing document with semantic identifiers: Hierarchical k-means organizes documents into a tree so semantically similar documents receive nearby identifiers.BERT representations are recursively clustered, and each document receives a deterministic routing path from the root to a leaf.
- 3.2 Query generation: Query generation binds document content to identifiers by augmenting training with generated queries and document-query examples.The query-generation transformer uses document terms as input; random sampling is used to produce diverse queries, while document terms also serve as additional queries.
- 3.3 Prefix-aware weight-adaptive decoder: The PAWA decoder predicts identifier tokens using weights adapted to the query and the preceding identifier prefix.Different identifier positions and prefixes can have different semantic meanings, so the decoder conditions token classification on prefix information.
- 3.4 Training and inference: NCI trains with sequence-to-sequence cross-entropy plus consistency-based regularization over training and augmented query-document pairs.The regularization compares decoder representations from independent dropout passes, and the combined objective is optimized end-to-end by gradient descent.
4 Experiments
Experiments evaluate NCI on NQ320k and TriviaQA using standard retrieval metrics, component ablations, model-capacity analyses, semantic-identifier analyses, and efficiency measurements. NCI improves retrieval across benchmarks, while each major design component contributes to performance and practical deployment remains feasible in some near-real-time settings.
- Datasets and metrics: Experiments use NQ320k and TriviaQA, evaluating Recall@N, MRR, and R-Precision for document retrieval.Recall@N measures whether the desired document appears among the top-N candidates; MRR measures the rank of the first relevant document; R-Precision is precision after retrieving R documents.
- Benchmark results: 72.78% Recall@1 on NQ320k lets NCI outperform SEAL by 21.4% after fine-tuning DocT5Query for the dataset.Across the reported comparisons, NCI also improves over SEAL by 17.6% for Recall@1, 10.0% for Recall@10, 3.2% for Recall@100, and 14.9% for MRR@100 on NQ320k.
- Ablation study: Query generation substantially boosts performance, while using document contents as queries helps NCI learn document semantics.The ablation discussion attributes the query-generation benefit to training with augmented queries, which helps the model understand document meanings.
- Ablation study: Removing the PAWA decoder, semantic identifiers, consistency regularization, or constrained beam search reduces performance across the evaluated metrics.The reported analyses associate these drops with the need for a task-specific decoder, semantic priors, resistance to decoder over-fitting, and explicit valid-identifier constraints, respectively.
- Model and decoder analyses: Increasing model capacity accelerates convergence; the small model has lower recall, while base and large models become similar after sufficient training.The large model is reported to be slightly higher at convergence.
- Model and decoder analyses: A four-layer PAWA adapter performs best among the tested layer counts, whereas larger adapters degrade performance, consistent with over-fitting.Performance improves from zero to four layers, decreases at six layers, and drops significantly at eight layers.
- Semantic-identifier analysis: Similar queries produce close retrieved-document distributions, while distinct query groups produce non-overlapping collections and separate embedding clusters.The analysis uses semantic-identifier prefixes, retrieval probabilities, and t-SNE visualizations of BERT-based document embeddings.
- Efficiency analysis: NCI has latency and throughput reported as affordable for some near-real-time applications without re-ranking, with latency comparable to DSI and SEAL at matched settings.BM25 is faster on CPU but has much lower recall metrics.
5 Limitation & Future Works
NCI still faces deployment limitations at large scale, especially in model capacity, inference speed, and updating the model-based index when documents are added.
- NCI requires substantially greater model capacity to extend from the current implementation to web-scale retrieval.
- NCI’s inference speed must improve before it can serve online queries in real time.
- Updating the model-based index when new documents are added remains difficult.
- The proposed remedies include sparsely-gated Mixture of Expert architectures, semantic clustering, model compression, and a hybrid traditional index for new documents.
6 Conclusion
The paper introduces NCI as an end-to-end neural retrieval paradigm that unifies training and indexing, directly generating relevant document identifiers and improving retrieval performance on two benchmarks.
- NCI unifies document-retrieval training and indexing within an end-to-end deep neural network.
- NCI directly generates identifiers of relevant documents for an input query and can be optimized with augmented query-document pairs.
- NCI uses a prefix-aware weight-adaptive decoder to optimize recall and ranking performance.
- NCI demonstrates outstanding performance over state-of-the-art solutions on the NQ320k and TriviaQA datasets.
A Related work
The paper situates its approach within traditional document retrieval and ranking, then reviews related research on end-to-end retrieval models that avoid explicit indexes.
- Traditional web search uses document retrieval to select relevant candidates before applying document ranking for more precise scoring.
- The retrieval stage requires balancing efficiency against recall, while ranking uses deeper models and richer features for query-document scoring.
- The related-work discussion motivates examining alternatives to the conventional retrieval-and-ranking framework.
A.1 Document retrieval
Prior document-retrieval research includes sparse term matching, dense vector search, neural ranking, and pretrained-language-model approaches, each emphasizing different efficiency and semantic-representation trade-offs.
- Sparse retrieval builds inverted indexes using term-matching metrics such as TF-IDF, query likelihood, and BM25.
- Dense retrieval represents queries and documents as vectors and compares them using inner product or cosine similarity.
- Neural ranking models include representation-based approaches that compare query and document embeddings in a shared semantic space.
- Single-tower BERT ranking can achieve strong performance, but its computational cost limits industrial-scale web-search deployment.
- Interaction-based ranking uses a crossing layer with simple similarity functions such as dot product or a single feed-forward layer for efficiency.
- Pretraining research has proposed paragraph-level tasks including ICT, BFS, and WLP because standard masked-language-model loss is not naturally fitted to embedding-based retrieval.
A.4 End-to-end retrieval
End-to-end retrieval methods avoid separate retrieval and reranking stages, but existing approaches still face effectiveness and efficiency limitations. The paper motivates a new indexing paradigm that directly predicts document identifiers.
- The index-retrieve paradigm separately optimizes document retrieval and reranking, making retrieval potentially sub-optimal for the overall system.
- ColBERT and COIL avoid reranking but are not clearly better than BM25 indexer plus BERT reranker and remain insufficiently efficient for industrial use.
- NCI-related model-based indexes generate document identifiers directly, replacing explicit retrieval indexes with learned prediction.
- DSI uses an encoder-decoder architecture for relevant docid generation, whereas SEAL uses passage n-grams with an FM-Index.
- SEAL’s approach is difficult to use because enumerating all passage n-grams for retrieval is challenging.
B.1 Dataset processing
The experiments use NQ320k and TriviaQA, with Wikipedia-based query-document data processed into cleaned article text and semantic identifiers. Hierarchical k-means supplies the identifiers, while constrained beam search supports retrieval.
- NQ320k uses natural-language questions and Wikipedia articles in HTML format.
- NQ320k processing removes useless HTML tag tokens and extracts each article’s title, abstract, and content with regular expressions.
- TriviaQA contains 78k query-document pairs from the Wikipedia domain.
- Hierarchical k-means is applied to article embeddings to produce semantic identifiers for each article.
- Constrained beam search uses a prefix tree containing valid identifiers to return documents with the highest probabilities.
B.4 Baselines
The baselines include sparse lexical retrieval, dense representation search, and query-expansion methods. They use BM25, ANN or brute-force similarity, and DocT5Query-generated questions to retrieve documents.
- The baseline descriptions rely mainly on official open-source implementations.
- BM25 calculates query-document similarity from contributions of individual query terms, weighted by term importance and term-document correlation.
- BM25 plus DocT5Query appends document-generated questions to documents before indexing them for bag-of-words BM25 retrieval.
- BERT plus ANN uses CoCondenser query and document representations with Faiss’s Flat Index method.
- BERT plus BruteForce separately encodes queries and documents with CoCondenser and computes cosine similarity for each pair.
C More Experimental Results
Additional experiments examine regularization strength and document the work’s social-impact and privacy considerations. The selected regularization value is α = 0.15, and the stated datasets contain no privacy information.
- α = 0.15 generally achieves better performance at convergence and is selected as NCI’s default regularization value.
- The regularization study compares α values of 0, 0.1, 0.15, 0.2, and 0.3 on NQ320k and TriviaQA.
- The authors state that the work aims to unify learning and indexing with an end-to-end deep neural network.
- The authors foresee potential positive social impacts from encouraging unified retrieval and reranking frameworks.
- The NQ and TriviaQA datasets are described as containing no privacy information.