Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Dynamic Retrieval-Augmented Generation

Anton Shapkin, Denis Litvinov, Yaroslav Zharov, Egor Bogomolov, Timur Galimzyanov, Timofey Bryksin

TL;DR

This work addresses hallucination and context-window limits in large language models by introducing Dynamic Retrieval-Augmented Generation (DRAG), which encodes retrieved entities as dense embeddings and integrates them into a generator's vocabulary to predict entity names as single tokens. DRAG comprises an embedder that produces entity embeddings and a generator that consumes an extended vocabulary, enabling dynamic retrieval without bleeding prompt length; the approach is trained end-to-end or with selective connecters, and uses cross-attention to fuse retrieved context. The authors validate DRAG across Python code generation, NL2Bash, and Spider SQL tasks, showing consistent improvements over context-promoting baselines and, in many cases, outperforming larger models, with a newly collected repository-level code generation dataset. The work contributes a scalable, entity-aware retrieval mechanism and a sizable dataset to support repository-level code generation research, with potential applicability to other domains requiring precise entity usage, such as legal or medical QA.

Abstract

Current state-of-the-art large language models are effective in generating high-quality text and encapsulating a broad spectrum of world knowledge. These models, however, often hallucinate and lack locally relevant factual data. Retrieval-augmented approaches were introduced to overcome these problems and provide more accurate responses. Typically, the retrieved information is simply appended to the main request, restricting the context window size of the model. We propose a novel approach for the Dynamic Retrieval-Augmented Generation (DRAG), based on the entity-augmented generation, which injects compressed embeddings of the retrieved entities into the generative model. The proposed pipeline was developed for code-generation tasks, yet can be transferred to some domains of natural language processing. To train the model, we collect and publish a new project-level code generation dataset. We use it for the evaluation along with publicly available datasets. Our approach achieves several targets: (1) lifting the length limitations of the context window, saving on the prompt size; (2) allowing huge expansion of the number of retrieval entities available for the context; (3) alleviating the problem of misspelling or failing to find relevant entity names. This allows the model to beat all baselines (except GPT-3.5) with a strong margin.

Dynamic Retrieval-Augmented Generation

TL;DR

This work addresses hallucination and context-window limits in large language models by introducing Dynamic Retrieval-Augmented Generation (DRAG), which encodes retrieved entities as dense embeddings and integrates them into a generator's vocabulary to predict entity names as single tokens. DRAG comprises an embedder that produces entity embeddings and a generator that consumes an extended vocabulary, enabling dynamic retrieval without bleeding prompt length; the approach is trained end-to-end or with selective connecters, and uses cross-attention to fuse retrieved context. The authors validate DRAG across Python code generation, NL2Bash, and Spider SQL tasks, showing consistent improvements over context-promoting baselines and, in many cases, outperforming larger models, with a newly collected repository-level code generation dataset. The work contributes a scalable, entity-aware retrieval mechanism and a sizable dataset to support repository-level code generation research, with potential applicability to other domains requiring precise entity usage, such as legal or medical QA.

Abstract

Current state-of-the-art large language models are effective in generating high-quality text and encapsulating a broad spectrum of world knowledge. These models, however, often hallucinate and lack locally relevant factual data. Retrieval-augmented approaches were introduced to overcome these problems and provide more accurate responses. Typically, the retrieved information is simply appended to the main request, restricting the context window size of the model. We propose a novel approach for the Dynamic Retrieval-Augmented Generation (DRAG), based on the entity-augmented generation, which injects compressed embeddings of the retrieved entities into the generative model. The proposed pipeline was developed for code-generation tasks, yet can be transferred to some domains of natural language processing. To train the model, we collect and publish a new project-level code generation dataset. We use it for the evaluation along with publicly available datasets. Our approach achieves several targets: (1) lifting the length limitations of the context window, saving on the prompt size; (2) allowing huge expansion of the number of retrieval entities available for the context; (3) alleviating the problem of misspelling or failing to find relevant entity names. This allows the model to beat all baselines (except GPT-3.5) with a strong margin.
Paper Structure (35 sections, 1 equation, 1 figure, 4 tables)

This paper contains 35 sections, 1 equation, 1 figure, 4 tables.

Table of Contents

  1. Introduction
  2. Method
  3. Embedder
  4. Generator
  5. Training and Decoding
  6. Related Work
  7. Experiments
  8. Baselines
  9. Context filtering methods
  10. No
  11. All
  12. k-NN
  13. Oracle
  14. RepoCoder
  15. Context utilization methods
  16. ...and 20 more sections

Figures (1)

  • Figure 1: The proposed method architecture. The embedder model produces embeddings of all the documents associated with the sample using cross-attention with the input representation from the generator model. These embeddings extend the generator vocabulary, allowing the generator to decide at each step whether to generate a regular token or to insert an entity.