Source-linked AI summary
CorpusBrain: Pre-train a Generative Retrieval Model for Knowledge-Intensive Language Tasks
Jiangui Chen, Ruqing Zhang, Jiafeng Guo, Yiqun Liu, Yixing Fan, Xueqi Cheng
TL;DR
KILT retrieval methods often require task-specific models and substantial annotated data, limiting their practicality across applications and low-resource settings. CorpusBrain pre-trains one generative retriever to encode corpus information in its parameters, and reports retrieval gains, stronger downstream performance, and effectiveness with little supervision.
Problem
Task-specific generative retrieval models can require separate models and substantial application-specific annotated data, limiting practicality across applications and low-data settings.
Source: §3.1 · p. 3Method
CorpusBrain pre-trains one generative retriever that samples document contexts as pseudo-queries and generates source or destination document identifiers using hyperlinks, without an additional index.
Source: §6 · p. 9Results
Takeaways & Limitations
CorpusBrain works well in zero- and low-resource settings, including when fine-tuned with very little supervision.
Source: §1 · p. 2, §5.4 · p. 8CorpusBrain underperforms other methods especially on QA datasets such as NQ, HoPo, and TQA; the authors suggest declarative pseudo-queries may differ from questions.
Source: §5.1 · p. 7
Abstract
from arXiv · showhide
Knowledge-intensive language tasks (KILT) usually require a large body of information to provide correct answers. A popular paradigm to solve this problem is to combine a search system with a machine reader, where the former retrieves supporting evidences and the latter examines them to produce answers. Recently, the reader component has witnessed significant advances with the help of large-scale pre-trained generative models. Meanwhile most existing solutions in the search component rely on the traditional ``index-retrieve-then-rank'' pipeline, which suffers from large memory footprint and difficulty in end-to-end optimization. Inspired by recent efforts in constructing model-based IR models, we propose to replace the traditional multi-step search pipeline with a novel single-step generative model, which can dramatically simplify the search process and be optimized in an end-to-end manner. We show that a strong generative retrieval model can be learned with a set of adequately designed pre-training tasks, and be adopted to improve a variety of downstream KILT tasks with further fine-tuning. We name the pre-trained generative retrieval model as CorpusBrain as all information about the corpus is encoded in its parameters without the need of constructing additional index. Empirical results show that CorpusBrain can significantly outperform strong baselines for the retrieval task on the KILT benchmark and establish new state-of-the-art downstream performances. We also show that CorpusBrain works well under zero- and low-resource settings.
1 INTRODUCTION
CorpusBrain replaces multi-stage search with a pre-trained generative retriever that encodes corpus information in its parameters. The authors report retrieval improvements on KILT and stronger downstream performance, including in low-resource settings.
1 INTRODUCTION
- KILT search commonly uses index-retrieve-then-rank pipelines, whose components are difficult to optimize end-to-end and whose indexes increase memory and computational costs.Errors can also accumulate across the sequential components.Source: §1 · pp. 1–2
- CorpusBrain pre-trains a single-step generative model to encode corpus information in its parameters, then generates relevant document identifiers from a query.Its pre-training tasks use sampled sentences or paragraphs and hyperlinks to construct pseudo query-document pairs without additional human supervision.Source: §1 · p. 2
- CorpusBrain improves retrieval over strong baselines on KILT and advances downstream performance, including with little fine-tuning supervision.The model is pre-trained on English Wikipedia and fine-tuned on eleven KILT datasets spanning five tasks.Source: §1 · p. 2
2 RELATED WORK
The related work covers traditional pipeline retrieval, emerging model-based IR, and the KILT benchmark's shared setting for knowledge-intensive tasks.
2 RELATED WORK
- The section situates these systems within work on pipeline IR, model-based IR, and knowledge-intensive language tasks.These are the three lines of related work reviewed in the section.Source: §2 · p. 2
2.1 Traditional Pipeline IR Framework
- Traditional IR separates indexing, retrieval, and ranking; sparse methods use term-based indexes, while dense methods retrieve via document representations and approximate nearest-neighbor search.Ranking approaches include vector-space, probabilistic, learning-to-rank, and neural models; the pipeline has end-to-end optimization and memory drawbacks.Source: §2.1 · pp. 2–3
2.2 Model-based IR Approaches
- Model-based IR consolidates search components, but earlier generative approaches target specific tasks and require substantial annotated data.Examples include generating entity names for entity linking, evidence sentence identifiers for fact checking, or alternative document identifiers.Source: §2.2 · p. 3
2.3 Knowledge-Intensive Language Tasks
- KILT provides a common interface for eleven datasets across five tasks, grounded in the same Wikipedia snapshot.The benchmark facilitates task-to-task comparisons despite differing dataset formats, assumptions, and knowledge sources.Source: §2.3 · p. 3
3 OUR APPROACH
This section introduces CorpusBrain and outlines its motivation, architecture, and pre-training tasks.
3 OUR APPROACH
- The section presents CorpusBrain as the paper's proposed generative retrieval model.It first introduces the motivation for the method.Source: §3 · p. 3
- The section next describes the model architecture.Architecture is one of the topics covered after the motivation.Source: §3 · p. 3
- The section also describes the pre-training tasks used by the approach.The tasks are introduced following the architecture overview.Source: §3 · p. 3
3.1 Motivation
The motivation is to replace task-specific generative retrieval and traditional search pipelines with a general-purpose model that can serve diverse KILT tasks. CorpusBrain is pre-trained to generate document identifiers without an additional index.
3.1 Motivation
- Traditional search pipelines separate retrieval stages, whereas model-based IR aims to consolidate them for end-to-end optimization and lower memory and computation costs.The cited vision describes a single consolidated model as an alternative to the index-retrieve-then-rank framework.Source: §3.1 · p. 3
- Earlier generative retrieval systems target specific KILT tasks, requiring separate models and substantial application-specific annotations that are unsuitable for zero- or few-shot settings.Separate specialized models may also impose considerable memory or computation overhead.Source: §3.1 · p. 3
- CorpusBrain is designed as a general-purpose retriever, with pre-training objectives for varied query-document granularity, variable numbers of relevant documents, and inter-document similarity.It uses an encoder-decoder to generate document identifiers through a standard sequence-to-sequence objective.Source: §3.1 · p. 3
- ISS uses a sampled sentence and LPS an early paragraph as input; HIP uses anchor context to predict a destination-page identifier.The figure also shows ISS and LPS predicting the source document identifier and linked destination identifiers.Source: Figure 2 · p. 4, §3.1 · p. 4
3.2 Model Architecture
CorpusBrain uses a Transformer encoder-decoder to represent each query and generate identifiers for relevant documents. Constrained beam search keeps generated identifiers within the corpus’s valid document set.
3.2 Model Architecture
- A Transformer encoder-decoder pairs a bidirectional Query Encoder with an Identifier Decoder that sequentially generates relevant-document identifiers.Document identifiers may be unique titles or URLs.Source: §3.2 · p. 4, §3.2.2 · p. 4
3.2.1 Query Encoder.
- The Query Encoder maps the input query to hidden vectors that form its query representation.The encoder is described as capturing the query’s essential topics.Source: §3.2.1 · p. 4
3.2.2 Identifier Decoder.
- Constrained beam search uses a prefix tree to restrict generated identifiers to valid document identifiers from the corpus.Each tree node’s children specify valid continuations of the prefix leading to that node.Source: §3.2.2 · p. 4
3.3 Pre-training Tasks
The pre-training tasks form query–document-identifier pairs from corpus documents using sentence, paragraph, and hyperlink signals. They share a sequence-to-sequence objective optimized end-to-end.
3.3 Pre-training Tasks
- Each corpus document is represented by a unique identifier, its sentences, and anchors linked to destination-page identifiers.These document elements provide the material for constructing the pre-training examples.Source: §3.3 · p. 4
3.3.1 Inner Sentence Selection (ISS).
- ISS samples an inner sentence as a pseudo query and trains generation of its document identifier alongside identifiers of sampled linked pages.This supports sentence-level queries and dynamic prediction of relevant documents.Source: §3.3.1 · p. 4
3.3.2 Lead Paragraph Selection (LPS).
- LPS uses leading paragraphs as pseudo queries for paragraph-level tasks and shares ISS’s target of the document and sampled linked-page identifiers.The leading paragraphs are treated as potentially higher-quality summaries than the remaining paragraphs.Source: §3.3.2 · p. 4
3.3.3 Hyperlink Identifier Prediction (HIP).
- HIP uses an anchor’s sentence and its adjacent sentences as a pseudo query, targeting the identifier of the page linked by that anchor.The surrounding context compensates for anchor texts that are too short to carry enough semantics.Source: §3.3.3 · p. 5
3.3.4 Learning Objective.
- All three tasks use a standard sequence-to-sequence cross-entropy objective, with all model parameters optimized end-to-end.The objective is defined over the pre-training dataset.Source: §3.3.4 · p. 5
4 EXPERIMENTAL SETTINGS
Experiments evaluate retrieval on eleven KILT datasets spanning five knowledge-intensive task types, using English Wikipedia for pre-training. The evaluation asks models to provide Wikipedia evidence pages for input queries.
4 EXPERIMENTAL SETTINGS
- The experimental-settings section introduces the pre-training corpora and downstream KILT datasets used in the experiments.The settings are organized around the corpus and downstream-task descriptions.Source: §4 · p. 5, §4.1 · p. 5
4.1.1 Pre-training Corpus.
- The pre-training corpus is the 2019/08/01 English Wikipedia dump, selected for public availability and its large collection of documents.The dump was pre-processed by Petroni et al. [38].Source: §4.1.1 · p. 5
4.1.2 Downstream Tasks.
- The KILT evaluation covers eleven datasets spanning five knowledge-intensive task types.The retrieval task requires providing Wikipedia pages as evidence for each input query.Source: §4.1.2 · p. 5
4.2 Baselines
The comparisons include both traditional information-retrieval systems and model-based retrieval systems, including generative retrieval baselines.
4.2 Baselines
- The baseline set is divided into traditional IR models and model-based IR models.The listed baselines include sparse and dense retrieval methods.Source: §4.2 · p. 5, §4.2.1 · p. 5
4.2.1 Traditional IR Models.
4.2.2 Generative Retrieval Models.
- Generative retrieval baselines include fine-tuned BART [28], T5 [39], SEAL [2], and GENRE [9].Their reported settings include zero-shot, task-specific, or multi-task fine-tuning, depending on the model.Source: §4.2.2 · p. 5
4.3 Evaluation Metrics
The section defines R-precision for downstream KILT retrieval and specifies which splits are evaluated in each setting.
4.3 Evaluation Metrics
- R-precision (%) is the retrieval metric used for downstream KILT evaluation.It is calculated as r/R, where R is the number of Wikipedia pages in a provenance set and r is the number of relevant pages among the top-R retrieved pages.Source: §4.3 · pp. 5–6
- The paper also reports average performance across all eleven datasets.Source: §4.3 · p. 6
- Zero-shot and full fine-tuning results use both test and dev sets, while zero- and low-resource results use dev only.The KILT leaderboard limits how often test performance can be submitted.Source: §4.3 · p. 6
4.4 Implementation Details
CorpusBrain uses a BART-large-like encoder-decoder architecture and Wikipedia titles as document identifiers. The implementation builds pseudo-training pairs from Wikipedia and fine-tunes on formatted KILT retrieval data.
4.4.1 Model Architecture.
- The 406M-parameter model uses 12-layer encoder and decoder stacks, 1024 hidden units, 4096-unit feedforward layers, and 16 attention heads.CorpusBrain and baseline BART use the same architecture for fair comparison.Source: §4.4.1 · p. 6
- Wikipedia titles identify documents; each article yields 18 pseudo-pairs across ISS, LPS, and HIP.ISS uses the lead two sentences and up to ten sampled sentences; LPS uses the top three paragraphs; HIP uses three sampled hyperlink anchors as input.Source: §4.4.2 · p. 6
4.4.2 Pre-training Process.
- CorpusBrain is initialized from the official BART checkpoint and trained with Adam, a 3e-5 learning rate, and linear decay after 10% warmup.Training uses batches of 8192 tokens on two NVIDIA Tesla V100 32GB GPUs.Source: §4.4.2 · p. 6
4.4.3 Fine-tuning Process.
- KILT inputs remain unchanged, while relevant page titles are concatenated with [SEP] as targets for fine-tuning.Fine-tuning uses 4096-token batches; inference uses constrained beam search with 10 beams and up to 15 decoding steps. Strategies include task-specific, KILT multitask, and KILT-plus-BLINK multitask training.Source: §4.4.3 · p. 6
5 EXPERIMENTAL RESULTS
The experiments compare CorpusBrain across KILT baselines, tuning strategies, pre-training tasks, resource settings, efficiency, and a retrieval case study. Overall, the reported results favor CorpusBrain, while the authors note weaker performance on some question-answering datasets.
5.1 Baseline Comparison
- Generative retrieval models significantly outperform traditional IR models across datasets, and CorpusBrain zero-shot is comparable to most fully fine-tuned traditional baselines.On zsRE, CorpusBrain zero-shot and DPR achieve 94.84% and 28.9% R-precision, respectively; the authors also report CorpusBrain underperforms particularly on QA datasets such as NQ, HoPo, and TQA.Source: §5.1 · pp. 6–7
- CorpusBrain_mt+BLINK ranks first on all 11 dev sets and on 8 of 11 test sets, placing second on the other 3 test sets.Among baselines applicable to all 11 tasks, it performs best on both dev and test; Tables 2 and 3 cover the test and dev results, respectively.Source: Table 3 · p. 7, Table 2 · p. 7, §5.1 · p. 7
5.2 Impact of Fine-tuning Strategies
- CorpusBrain outperforms BART across KILT tasks under both task-specific and multitask fine-tuning, and multitask fine-tuning performs better than task-specific tuning.The authors report faster convergence than BART under both strategies.Source: §5.2 · pp. 7–8, Figure 3 · p. 8
5.3 Impact of Pre-training Tasks
- ISS generally performs best among the individual pre-training tasks, while HIP performs best on entity-linking tasks; combining all three gives the best performance.The authors attribute HIP's entity-linking advantage to its capture of inter-page relations and note that zero-shot entity linking remains poor without entity-boundary tokens during pre-training.Source: §5.3 · p. 8, Table 4 · p. 8
5.4 Zero- and Low-Resource Settings
- With the same limited supervision, CorpusBrain outperforms BART on all five datasets and achieves competitive results with fully supervised RAG; on zsRE, 100 examples yield 98.28% R-precision versus GENRE’s 94.84%.With zero resources, CorpusBrain outperforms RAG on FEV (71.69% vs. 63.50%) and zsRS (90.95% vs. 65.36%).Source: §5.4 · p. 8
5.5 Memory and Inference Efficiency
- CorpusBrain uses 34 times less memory than DPR and 20 times less than RAG, and the authors report reduced retrieval inference time.CorpusBrain stores corpus information in model parameters, while dense retrieval computes relevance over corpus representations; its inference time is directly proportional to beam size, with limited constrained-decoding overhead.Source: §5.5 · p. 9, Table 6 · p. 9
5.6 Case Study
- In the ELI5 example, CorpusBrain_mt+BLINK generates correct and potentially relevant titles, including “legal liability,” while GENRE's beams miss another relevant document.The example compares the models' top-five beams against “False advertising” and “Media regulation” as the ground-truth relevant titles.Source: §5.6 · p. 9, Table 5 · p. 9
6 CONCLUSION
CorpusBrain encodes corpus information in a pre-trained generative retrieval model, using designed pre-training tasks to support adaptation across KILT tasks. The paper reports retrieval improvements over strong baselines and strong zero- and low-resource performance, while identifying broader retrieval settings and end-to-end KILT as future work.
6 CONCLUSION
- CorpusBrain encodes corpus information in its parameters and uses pre-training tasks that emphasize different aspects of query-document semantics.The model can be adapted to downstream KILT tasks without constructing an additional index.Source: §6 · p. 9
- CorpusBrain improved retrieval over strong baselines on KILT and achieved strong performance in zero- and low-resource settings.Source: §6 · p. 9
- Future work includes testing other document identifiers and downstream IR tasks, and designing an end-to-end generative KILT system.Source: §6 · p. 9