PaperlayerView exact revision on arXiv

Source-linked AI summary

A Unified Generative Retriever for Knowledge-Intensive Language Tasks via Prompt Learning

Jiangui Chen, Ruqing Zhang, Jiafeng Guo, Maarten de Rijke, Yiqun Liu, Yixing Fan, Xueqi Cheng

arXiv:2304.14856v1cs.IR

TL;DR

KILTs need relevant contexts at different granularities, but universal retrievers can be too coarse and task-specific ones may generalize poorly. UGR combines generative n-gram identifiers with task prompts in one retriever, and reports better retrieval than prevailing baselines across benchmark settings.

  • Problem

    Universal document retrievers can return contexts that are too coarse for KILTs, while task-specific retrievers may generalize poorly to out-of-domain data and unseen tasks.

    Source: §1 · p. 2

  • Method

    UGR uses one generative retriever with n-gram identifiers for contexts at different granularities and task-specific prompts to learn multiple retrieval tasks.

    Source: §1 · p. 2

  • Results

    UGR outperforms prevailing baselines on in-domain, out-of-domain, and unseen-task retrieval, and supports new state-of-the-art downstream results on multiple datasets with existing readers.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Integrating knowledge from different retrieval tasks in UGR has the potential to save time and computational resources in academic and industrial environments.

    Source: §6 · p. 9

    UGR needs a complex scoring function to address repeated identifiers; broader retrieval applications across multiple corpora and modalities remain a future direction.

    Source: §6 · p. 9

Abstract

from arXiv · show

Knowledge-intensive language tasks (KILTs) benefit from retrieving high-quality relevant contexts from large external knowledge corpora. Learning task-specific retrievers that return relevant contexts at an appropriate level of semantic granularity, such as a document retriever, passage retriever, sentence retriever, and entity retriever, may help to achieve better performance on the end-to-end task. But a task-specific retriever usually has poor generalization ability to new domains and tasks, and it may be costly to deploy a variety of specialised retrievers in practice. We propose a unified generative retriever (UGR) that combines task-specific effectiveness with robust performance over different retrieval tasks in KILTs. To achieve this goal, we make two major contributions: (i) To unify different retrieval tasks into a single generative form, we introduce an n-gram-based identifier for relevant contexts at different levels of granularity in KILTs. And (ii) to address different retrieval tasks with a single model, we employ a prompt learning strategy and investigate three methods to design prompt tokens for each task. In this way, the proposed UGR model can not only share common knowledge across tasks for better generalization, but also perform different retrieval tasks effectively by distinguishing task-specific characteristics. We train UGR on a heterogeneous set of retrieval corpora with well-designed prompts in a supervised and multi-task fashion. Experimental results on the KILT benchmark demonstrate the effectiveness of UGR on in-domain datasets, out-of-domain datasets, and unseen tasks.

1 INTRODUCTION

KILTs need relevant contexts at different levels of granularity, while existing universal and task-specific retrievers each have drawbacks. UGR uses a single generative model with n-gram identifiers and task prompts; experiments report improved retrieval across benchmark settings.

1 INTRODUCTION

  • KILTs need retrieved contexts at task-appropriate granularities, but universal document retrievers can be too coarse, while task-specific retrievers may generalize poorly and be costly to deploy.Retrieval supports a search-and-reader pipeline in which the reader uses retrieved contexts to perform the end task.Source: §1 · pp. 1–2
  • UGR is a single generative retriever intended to share knowledge across tasks while returning contexts at the granularity each retrieval task requires.Given a query, the model generates identifiers for relevant contexts.Source: §1 · p. 2
  • N-gram-based identifiers use important n-grams from contexts as possible keys for retrieving documents, passages, sentences, and entities.The method samples n-grams using re-normalized BERT [CLS]-token attention weights; constrained beam search with an FM-Index [12] ensures generated n-grams occur in the corpus, and interactive scoring breaks ties.Source: §1 · p. 2
  • Supervised multi-task training uses task-specific prompts and a Seq2Seq objective to retain task specifications while learning across the four retrieval tasks.The training mixture contains the four tasks, and multi-task supervision is intended to encourage shared knowledge rather than overemphasizing task-specific knowledge.Source: §1 · p. 2
  • UGR outperforms prevailing baselines on in-domain, out-of-domain, and unseen-task retrieval, and its contexts support new state-of-the-art downstream results on multiple datasets with existing readers.The evaluation uses eleven KILT datasets spanning four retrieval granularities.Source: §1 · p. 3

2 RELATED WORK

Prior KILT retrieval work includes document-level approaches that can miss finer-grained needs and task-specific models whose generalization may suffer. Generative retrieval and prompt learning provide relevant precedents, but prior generative systems addressed individual retrieval tasks.

2 RELATED WORK

  • KILT retrieval research includes document retrievers that ignore finer-grained context needs and task-specific models that may generalize poorly across tasks.KILT provides a common format and Wikipedia snapshot for five tasks: fact checking, open-domain question answering, slot filling, entity linking, and dialogue.Source: §2 · p. 3
  • Generative retrieval has been proposed to optimize a corpus-level objective and reduce storage and computation costs relative to modularized pipelines.The approach generates short context identifiers rather than full contexts.Source: §2 · p. 3
  • Generative IR maps queries to identifiers for relevant contexts; prior work proposed this formulation but designed systems for individual retrieval tasks.Examples include document, passage, and entity retrieval systems.Source: §2 · p. 3
  • Prompt learning reformulates downstream data to resemble pretraining, and prior instruction-based fine-tuning improved generalization in language models.UGR applies prompt learning to its four retrieval tasks using task-specific prompts.Source: §2 · p. 3

3 OUR APPROACH

UGR formulates document, passage, sentence, and entity retrieval as one generative task that produces context identifiers from queries. It combines n-gram identifiers with task-specific prompts, multi-task training, and constrained inference.

3.1 Retrieval task description

  • UGR directly generates identifiers for relevant contexts, unifying document, passage, sentence, and entity retrieval as a Seq2Seq problem.The formulation uses a query and relevant context identifiers across retrieval tasks.Source: §3.1 · p. 3

3.2 Overview of the approach

  • The approach addresses both unifying retrieval tasks in one generative form and specializing a single model for task differences.The overview also includes the training and inference process.Source: §3.2 · p. 4

3.3 N-gram-based identifiers

  • UGR uses important context n-grams as identifiers, estimating their importance from BERT [CLS]-token attention and sampling them without metadata or human annotation.The design seeks semantic structure, low-cost creation, and identifiers that distinguish contexts; 0.04% of documents collide over the same n-grams in training with ten 10-grams on Wikipedia English, and inference addresses repetition.Source: §3.3 · p. 4

3.4 Prompt engineering

  • Task-specific prompts address interference between retrieval tasks and include manually written discrete prompts, learned continuous prompts, and hybrids combining both.Hybrid prompts add anchor text such as “document” rather than relying only on a learned continuous vector.Source: §3.4 · pp. 4–5
  • UGR uses task-specific prompts in a mixture of retrieval tasks so the model can learn shared knowledge while distinguishing task specifications.The authors describe the prompts as supporting generalization across tasks.Source: §3.4 · p. 5

3.5 Training and inference

  • The training and inference process combines multi-task learning across retrieval corpora with efficient matching of contexts that contain generated n-grams..Source: §3.5 · p. 5

3.5.1 Multi-task training.

  • UGR is built on BART and trained end-to-end with a Seq2Seq likelihood objective that maximizes the probability of output identifier strings using teacher forcing.The parameters are optimized with cross-entropy loss.Source: §3.5.1 · p. 5

3.5.2 Inference process.

  • At inference, an FM-index supports constrained beam search for valid n-grams; candidate contexts are ranked by combining evidence from generated n-grams.The method selects the context with the highest score, addressing collisions among document, passage, and sentence identifiers.Source: §3.5.2 · p. 5

4 EXPERIMENTAL SETUP

The experimental setup evaluates retrieval across KILT datasets and four granularities, separating datasets used in training from held-out datasets. It specifies comparison systems, evaluation metrics, and UGR implementation choices.

4.1 Datasets

4.2 Baselines

  • Comparisons use BM25, GENRE, SEAL, MT-DPR, and BLINK as task-specific retrieval baselines.The reported baseline results are the best corresponding-task results from the original papers.Source: §4.2 · p. 6

4.3 Evaluation metrics

  • R-precision evaluates all four retrieval tasks, while downstream tasks use task-specific metrics and dev-set results.Downstream metrics include Accuracy, Exact Match, ROUGE-L, and F1.Source: §4.3 · p. 6

4.4 Implementation details

  • UGR uses a transformer encoder-decoder with 406M parameters, 1024-dimensional hidden states, 12 layers, and 16 attention heads.The feed-forward layer size is 4096.Source: §4.4 · p. 6
  • N-gram identifiers use n=10, with 10 n-grams for document and passage retrieval and five for sentence retrieval.Entity retrieval instead uses the entity name as its sole identifier; ρ is set to 0.01.Source: §4.4 · p. 6
  • UGR compares discrete, continuous, and hybrid prompts, using six prompt tokens for continuous and hybrid variants.Training initializes from BART_large, uses Adam at 3e−5, and first trains the prompt encoder before fixing it to train the generative model.Source: §4.4 · p. 6
  • Inference uses constrained beam search over a Wikipedia FM-index, decoding for 10 timesteps with 15 beams.The setup specifies α=2.0, β=0.8, and g=5.Source: §4.4 · p. 6

5 EXPERIMENTAL RESULTS

The experiments assess UGR against retrieval baselines, on unseen tasks, across identifier designs, with downstream readers, and for memory and inference cost.

5.1.1 In-domain performance.

  • BART_sp beats BART_mt, while BART_hp exceeds BART_sp, supporting task-specific prompts over simple mixed-data fine-tuning.The authors attribute hybrid prompts’ advantage over discrete and continuous prompts to using natural language to control continuous prompt tokens.Source: §5.1.1 · p. 7

5.1 Evaluation on in-domain and out-of-domain datasets

  • UGR with hybrid prompts significantly outperforms all baselines in-domain and beats them on every out-of-domain dataset.Tables 4 and 5 report R-precision for the four retrieval tasks in-domain and out-of-domain, respectively.Source: §5.1.1 · p. 7, Table 4 · p. 7, §5.1.2 · p. 7

5.2 Adaptability to unseen tasks

  • Zero-shot UGR is competitive with BM25; few-shot UGR_hp performs strongly and beats prior SOTA on some datasets, while NQ document retrieval is comparable to full-supervision SOTA.The unseen-task evaluation trains on three retrieval tasks and tests on the held-out fourth; few-shot training uses 1,000 instances per task.Source: §5.2 · pp. 7–8, Table 6 · p. 8

5.3 Analysis of n-gram-based identifiers

  • On selected DR, PR, and SR datasets, important n-gram sampling beats random sampling by a large margin; ER requires no n-gram sampling.The comparison covers selected datasets, rather than every dataset for each retrieval task.Source: §5.3 · p. 8, Table 7 · p. 8
  • On FEV, with qualitatively similar PR and SR observations, 5- or 15-grams and 15 identifiers can hurt performance, so length and count require a trade-off.Short n-grams can repeat across identifiers; longer n-grams risk generation errors, and using 15 identifiers introduces noise relative to 10.Source: §5.3 · pp. 8–9, Figure 3 · p. 8

5.4 Downstream performance

  • UGR_hp+FiD significantly beats prior SOTA across datasets, while documents reduce slot-filling performance about 4.51% versus passages.FiD receives the query and the first five contexts retrieved by each model.Source: §5.4 · p. 9, Table 8 · p. 9

5.5 Memory and inference efficiency

  • Generative retrievers use more parameters but less memory and run faster than traditional models; UGR costs more memory and time than GENRE yet matches SEAL’s inference time.The authors attribute the traditional models’ larger footprint to storing dense corpus representations and UGR’s extra cost over GENRE to its FM-index.Source: §5.5 · p. 9, Table 9 · p. 9

6 CONCLUSION

UGR formulates multi-granularity retrieval for KILTs as conditional generation, using n-gram identifiers and task-description prompt tokens in one model. KILT results demonstrated the method’s superiority.

6 CONCLUSION

  • Empirical results on KILT demonstrated the superiority of UGR.Source: §6 · p. 9
  • UGR formulates retrieval as conditional generation and uses n-gram identifiers for relevant contexts at different granularities.Source: §6 · p. 9
  • UGR maps task descriptions to a few prompt tokens to preserve task specifications across retrieval tasks.Source: §6 · p. 9
  • Integrating knowledge from different retrieval tasks in UGR has the potential to save time and computational resources, but identifier repetition requires a complex scoring function.Source: §6 · p. 9
  • Future work could explore efficient semantic identifiers and unified retrieval across multiple corpora and modalities.Source: §6 · p. 9
Loading 2304.14856v1…