Dynamic Retrieval-Augmented Generation

📄 arXiv: 2312.08976v2 📥 PDF

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

分类: cs.SE, cs.LG

发布日期: 2023-12-14 (更新: 2024-02-20)

备注: 10 pages

💡 一句话要点

提出动态检索增强生成（DRAG）方法，提升代码生成任务中大语言模型的准确性与效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 检索增强生成 代码生成 实体嵌入 知识注入 大语言模型

📋 核心要点

  1. 现有检索增强方法直接将检索信息附加到请求，限制了上下文窗口大小，影响生成质量。
  2. DRAG方法通过实体增强生成，将检索实体的压缩嵌入注入模型，扩展了可利用的上下文信息。
  3. 在代码生成任务上，DRAG模型在多个数据集上超越了现有基线，验证了其有效性。

📝 摘要（中文）

当前最先进的大语言模型在生成高质量文本和封装广泛的世界知识方面表现出色。然而，这些模型经常出现幻觉，并且缺乏局部相关的实际数据。检索增强方法被引入以克服这些问题并提供更准确的响应。通常，检索到的信息只是简单地附加到主请求中，从而限制了模型的上下文窗口大小。我们提出了一种基于实体增强生成的新型动态检索增强生成（DRAG）方法，该方法将检索到的实体的压缩嵌入注入到生成模型中。所提出的流程是为代码生成任务开发的，但可以转移到某些自然语言处理领域。为了训练模型，我们收集并发布了一个新的项目级代码生成数据集。我们将其与公开可用的数据集一起用于评估。我们的方法实现了多个目标：（1）消除了上下文窗口的长度限制，节省了提示大小；（2）允许大量扩展可用于上下文的检索实体数量；（3）缓解了拼写错误或无法找到相关实体名称的问题。这使得该模型能够以很大的优势击败所有基线（GPT-3.5除外）。

🔬 方法详解

问题定义：现有检索增强的代码生成方法受限于上下文窗口大小，无法有效利用大量检索到的信息。直接拼接检索结果会增加prompt长度，降低生成效率，并且容易受到检索错误实体的影响。

核心思路：DRAG的核心在于将检索到的实体信息进行压缩表示，并通过嵌入的方式注入到生成模型中，从而避免了直接拼接文本带来的上下文长度限制。这种方法允许模型利用更多的检索信息，同时减轻了拼写错误和实体识别失败的影响。

技术框架：DRAG包含以下主要步骤：1) 实体检索：从知识库或文档中检索与输入相关的实体。2) 实体嵌入：将检索到的实体转换为低维嵌入向量。3) 动态注入：将实体嵌入动态地注入到生成模型的中间层，影响模型的生成过程。4) 代码生成：利用注入了实体信息的模型生成代码。

关键创新：DRAG的关键创新在于动态地将检索到的实体信息以嵌入的形式注入到生成模型中，而不是简单地将检索结果拼接在输入端。这种方法突破了上下文窗口的限制，允许模型利用更多的外部知识，并且对检索错误具有一定的鲁棒性。

关键设计：DRAG的具体实现细节包括：1) 实体嵌入模型的选择（例如，预训练的知识图谱嵌入模型）。2) 嵌入向量的维度。3) 嵌入向量注入到生成模型的具体位置（例如，Transformer的哪一层）。4) 训练过程中使用的损失函数，可能包括生成损失和实体对齐损失等。

📊 实验亮点

DRAG模型在代码生成任务上取得了显著的性能提升，在多个数据集上超越了现有基线模型（除GPT-3.5外）。实验结果表明，DRAG能够有效利用检索到的信息，提高代码生成的准确性和效率，尤其是在需要大量外部知识的任务中表现出色。

🎯 应用场景

DRAG方法不仅适用于代码生成，还可以扩展到其他需要利用外部知识的自然语言处理任务，例如问答系统、文本摘要和机器翻译。该方法能够提升模型在知识密集型任务中的准确性和可靠性，具有广泛的应用前景。

📄 摘要（原文）

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.