One Token Can Help! Learning Scalable and Pluggable Virtual Tokens for Retrieval-Augmented Large Language Models

📄 arXiv: 2405.19670v4 📥 PDF

作者: Yutao Zhu, Zhaoheng Huang, Zhicheng Dou, Ji-Rong Wen

分类: cs.CL

发布日期: 2024-05-30 (更新: 2024-12-11)

备注: Accepted by AAAI 2025, repo: https://github.com/DaoD/SPRING/

💡 一句话要点

提出可扩展和可插拔的虚拟Token，用于增强检索式大语言模型，提升效果并保持通用性。

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

关键词: 检索增强生成 大型语言模型 虚拟Token 可插拔 微调 问答系统 知识库

📋 核心要点

  1. 现有RAG方法微调LLM参数以提升性能，但会牺牲LLM的通用生成能力，限制了已部署LLM的适用性。
  2. 提出学习可扩展和可插拔的虚拟Token，仅微调Token嵌入，在不修改LLM参数的前提下增强RAG性能。
  3. 设计训练策略提升方法的可扩展性、灵活性和泛化性，并在12个问答任务上验证了方法的优越性。

📝 摘要（中文）

检索增强生成（RAG）是一种有前景的方法，可以提高大型语言模型（LLM）生成更具事实性、准确性和最新内容的能力。现有方法要么优化提示，以指导LLM利用检索到的信息，要么直接微调LLM以适应RAG场景。虽然微调可以产生更好的性能，但它通常会通过修改LLM的参数来损害其通用生成能力。这种限制在实际应用中提出了挑战，尤其是在LLM已经部署时，因为参数调整可能会影响其原始功能。为了解决这个问题，我们提出了一种新颖的方法，该方法涉及学习可扩展和可插拔的RAG虚拟Token。通过保持LLM的原始参数，仅微调这些可插拔Token的嵌入，我们的方法不仅增强了LLM的性能，还保留了其通用生成能力。此外，我们设计了几种训练策略，以提高我们方法的可扩展性、灵活性和泛化性。在12个问答任务中进行的综合实验证明了我们方法的优越性。

🔬 方法详解

问题定义：现有检索增强生成（RAG）方法，特别是通过微调大型语言模型（LLM）来提升性能的方式，存在破坏LLM原有通用能力的风险。这对于已经部署的LLM来说是一个严重的问题，因为任何参数调整都可能影响其预期的功能。因此，如何在不修改LLM参数的情况下，有效提升RAG的性能，是本文要解决的核心问题。

核心思路：本文的核心思路是引入可扩展和可插拔的虚拟Token。这些Token的嵌入可以被微调，以适应RAG任务，而无需修改LLM本身的参数。这种设计允许在保持LLM原始功能的同时，增强其在RAG场景下的表现。通过仅调整少量Token的嵌入，可以避免对整个LLM进行微调带来的风险。

技术框架：该方法的核心在于学习一组虚拟Token的嵌入表示。这些Token被插入到LLM的输入序列中，与检索到的文档一起作为LLM的上下文。LLM利用这些Token来更好地理解和利用检索到的信息。训练过程中，LLM的参数保持固定，只有虚拟Token的嵌入被更新。整体流程包括：1) 检索相关文档；2) 将文档和虚拟Token插入到LLM的输入序列中；3) LLM生成答案；4) 基于生成答案和真实答案之间的差异，更新虚拟Token的嵌入。

关键创新：最重要的技术创新点在于引入了可插拔的虚拟Token，并设计了相应的训练策略。与直接微调LLM相比，这种方法能够在不影响LLM通用能力的前提下，提升RAG的性能。此外，论文还提出了多种训练策略，以提高虚拟Token的可扩展性、灵活性和泛化性。

关键设计：关键设计包括：1) 虚拟Token的数量：需要根据任务的复杂度和LLM的容量进行调整；2) 损失函数：用于衡量生成答案和真实答案之间的差异，例如交叉熵损失；3) 训练策略：包括如何初始化虚拟Token的嵌入、如何更新嵌入、以及如何防止过拟合等。论文中具体使用的参数设置和网络结构细节未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，该方法在12个问答任务上均取得了优异的性能。与现有RAG方法相比，该方法能够在保持LLM通用能力的同时，显著提升RAG的准确性和效率。具体的性能提升幅度未知，但摘要中明确指出该方法优于现有方法。

🎯 应用场景

该研究成果可广泛应用于各种需要检索增强生成的大语言模型应用场景，例如问答系统、知识库构建、智能客服等。通过即插即用的方式，可以快速提升现有LLM在特定领域的表现，而无需重新训练整个模型，具有很高的实际应用价值和商业潜力。未来，该方法可以进一步扩展到其他类型的LLM和RAG任务中。

📄 摘要（原文）

Retrieval-augmented generation (RAG) is a promising way to improve large language models (LLMs) for generating more factual, accurate, and up-to-date content. Existing methods either optimize prompts to guide LLMs in leveraging retrieved information or directly fine-tune LLMs to adapt to RAG scenarios. Although fine-tuning can yield better performance, it often compromises the LLMs' general generation capabilities by modifying their parameters. This limitation poses challenges in practical applications, especially when LLMs are already deployed, as parameter adjustments may affect their original functionality. To address this, we propose a novel method that involves learning scalable and pluggable virtual tokens for RAG. By maintaining the LLMs' original parameters and fine-tuning only the embeddings of these pluggable tokens, our approach not only enhances LLMs' performance but also preserves their general generation capabilities. Furthermore, we design several training strategies to improve the scalability, flexibility, and generalizability of our method. Comprehensive experiments across 12 question-answering tasks demonstrate the superiority of our approach.