基于 Transformers 的推测式 RAG 实现

大型语言模型（LLM）在需要最新且精确信息的知识密集型任务中往往难以胜任。这正是检索增强生成（RAG）发挥作用的地方：它将 LLM 的生成能力与外部知识库相结合，以提升准确性与相关性。

尽管有效，传统 RAG 系统在处理冗长复杂文档时面临挑战，导致延迟增加，且有时准确性下降。为应对这些问题，推测式 RAG的概念应运而生，成为一种颇具前景的解决方案。下面让我们进一步了解。

什么是推测式 RAG？

推测式检索增强生成（RAG）是一种创新框架，通过将生成过程划分为两个明确步骤，从而提升 RAG 系统的效率与准确性：

1. 起草：第一步使用一个更小、更专业的语言模型，即 RAG 起草器（RAG Drafter）。该模型从不同的检索文档子集生成多个答案草稿。每个子集都会精心挑选，以体现多元视角，尽量减少冗余并提升草稿质量。
2. 验证：由起草器生成的草稿会交由更大的通用型语言模型（RAG 验证器，RAG Verifier）进行评估。该模型依据草稿所提供的推理过程，专注于筛选出最准确、最可靠的答案。验证过程借助通用型模型的语言建模能力，确保最终输出既准确又具备上下文相关性。

让我们通过下图更好地理解推测式 RAG：

来源：Wang 等，2024

我们按步骤解读该图。图中展示了几种利用外部信息帮助 LLM 回答问题的方法，重点介绍一种新方法——推测式 RAG。先看第一行的方法：

• (a) 标准 RAG：最简单的方式。所有相关文档（图中的 1、2…）与问题（Q）一起直接提供给 LLM（“通用型 LM”），以生成答案（A）。
• (b) 自反式 RAG：LLM 被训练去分析文档及其自身的答案，生成诸如“Need”“Relevant”“Supported”等标签，以改进其推理。
• (c) 纠错式 RAG：在将文档交给 LLM 之前，使用独立模型（“NLI 模型”）检查检索到的文档是否确与问题相关。

现在，重点看看(d) 推测式 RAG：

1. 专业起草器：为避免给大模型造成负担，一个更小的“起草器”先行处理。它结合问题（Q）与文档（1、2…）生成多个答案草稿（α1、α2、α3…）。每个草稿基于不同的文档子集。
2. 理由说明：起草器还会给出为何为各草稿选择特定文档的解释（β1、β2、β3…）。可将其视为“展示解题过程”。
3. 评估：此时由主力 LLM（“通用型 LM”）接手。它同时查看每个草稿及其解释，并据此挑选最佳答案（A）。

本质上，推测式 RAG 就像一个专家团队：起草器提出多种方案，通用型 LLM 作最终决策。这样既更高效，也有望更准确。

推测式 RAG 的优势

相比传统 RAG，推测式 RAG 具有以下关键优势：

1. 更高的准确性

通过使用专门的 RAG 起草器从多样的文档子集生成多个草稿，推测式 RAG 提高了最终输出的准确性。该方法允许系统纳入多元视角，降低产生错误或偏见答案的可能性。实验中，推测式 RAG 在 PubHealth 等基准上较传统 RAG 的准确率最高提升 12.97%。

2. 更低的延迟

传统 RAG 的一大挑战是处理冗长复杂文档带来的延迟上升。推测式 RAG 将起草工作卸载给更小、更高效的模型，并可并行生成草稿，从而缩短整体处理时间；在某些情况下，延迟可降低至多 51%。

不同数据集上推测式 RAG 与标准 RAG 的延迟对比。来源：Wang 等，2024

3. 资源效率更高

推测式 RAG 通过将起草任务分配给更小的模型来优化资源使用，相比传统 RAG 中的大型通用模型，计算开销更低。此番分工不仅加速生成流程，也降低总体算力负担，使之更为高效。

4. 可扩展且灵活

推测式 RAG 框架适应性强，可广泛应用于多种知识密集型任务，而无需大量调优。这种可扩展性使其非常适合从问答到复杂文档分析等多种应用场景。

如何实现推测式 RAG

为在基础问答任务中实现一个简化版推测式 RAG，我们将使用一个小型数据集，以及由两个模型构成的简易方案——小模型用于起草，大模型用于验证。

为便于演示，我们使用Hugging Face Transformers 库。我们将展示如何用小型语言模型生成多个草稿，并用更大型号验证并选出最佳草稿。

步骤 1：安装所需库

首先安装必要库。您需要在从 HuggingFace 安装 transformers 库后重启会话。请确保已在 Colab 或其他编辑器中以密钥方式添加了 HuggingFace 凭证。

!pip install transformers torch datasets技术支持

步骤 2：加载数据集

本示例将使用 Hugging Face datasets 库中的SQuAD（斯坦福问答数据集）。

from datasets import load_dataset
# Load the SQuAD dataset
dataset = load_dataset("squad", split="train[:100]")  # Using a small subset技术支持

步骤 3：初始化模型

我们将使用 Hugging Face 库中的两个预训练模型。小模型（distilbert-base-uncased-distilled-squad）充当 RAG 起草器；大模型（bert-large-uncased-whole-word-masking-finetuned-squad）充当 RAG 验证器。

随后，我们将基于所选模型构建起草器与验证器管线。

from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline

# Initialize the smaller model (RAG Drafter)
drafter_model_name = "distilbert-base-uncased-distilled-squad"
drafter_model = AutoModelForQuestionAnswering.from_pretrained(drafter_model_name)
drafter_tokenizer = AutoTokenizer.from_pretrained(drafter_model_name)

# Initialize the larger model (RAG Verifier)
verifier_model_name = "bert-large-uncased-whole-word-masking-finetuned-squad"
verifier_model = AutoModelForQuestionAnswering.from_pretrained(verifier_model_name)
verifier_tokenizer = AutoTokenizer.from_pretrained(verifier_model_name)

# Set up pipelines
drafter_pipeline = pipeline("question-answering", model=drafter_model, tokenizer=drafter_tokenizer)
verifier_pipeline = pipeline("question-answering", model=verifier_model, tokenizer=verifier_tokenizer)技术支持

步骤 4：定义起草函数

接下来，使用预定义的起草器管线，通过小模型生成多个草稿。

def generate_drafts(question, context, num_drafts=3):
   drafts = []
   for _ in range(num_drafts):
       draft = Drafter_pipeline(question=question, context=context)
       drafts.append(draft)
   return drafts技术支持

步骤 5：定义验证函数

现在我们来定义一个 verify_drafts() 函数，利用更大模型基于置信度分数对草稿进行验证并选择最佳草稿。

其工作方式是：使用偏移映射将每个草稿的起止字符位置映射到分词后的上下文中的对应 token 位置。验证器模型随后基于这些位置为每个草稿打分，分数反映模型对该答案位置正确性的置信度。选择置信度最高的草稿作为最佳答案，以确保最终输出既准确又可靠。

该函数对 token 对齐与打分的精细处理，是推测式 RAG 性能提升的关键。

def verify_drafts(question, context, drafts):
    best_draft = None
    highest_score = 0

    # Tokenize the context using the verifier's tokenizer, keeping track of offsets
    inputs = verifier_tokenizer(question, context, return_tensors="pt", return_offsets_mapping=True)
    offset_mapping = inputs['offset_mapping'][0]  # This will give us the character-to-token mapping
    input_ids = inputs['input_ids'][0]

    for draft in drafts:
        start_char = draft['start']
        end_char = draft['end']

        # Find the corresponding token positions using offset mapping
        start_index = None
        end_index = None

        for idx, (start, end) in enumerate(offset_mapping):
            if start_index is None and start_char >= start and start_char < end:
                start_index = idx
            if end_index is None and end_char > start and end_char <= end:
                end_index = idx
            if start_index is not None and end_index is not None:
                break

        # Ensure indices were found and are within bounds
        if (start_index is None or end_index is None or
                start_index >= len(input_ids) or end_index >= len(input_ids)):
            print(f"Draft skipped: Out of bounds or no matching tokens. "
                  f"Start Index: {start_index}, End Index: {end_index}")
            continue

        # Get the confidence score using the larger model
        outputs = verifier_model(input_ids=input_ids.unsqueeze(0))
        score = (outputs.start_logits[0, start_index].item() +
                 outputs.end_logits[0, end_index].item())

        if score > highest_score:
            highest_score = score
            best_draft = draft

    if best_draft is None:
        print("No valid draft found after verification.")

    return best_draft技术支持

步骤 6：运行并评估推测式 RAG 流程

现在，将我们的推测式 RAG 流程应用于数据集中的 10 个示例问题。

correct = 0
total = 10  # Evaluate on 10 samples for simplicity

for i in range(total):
    sample = dataset[i]
    question = sample['question']
    context = sample['context']
    drafts = generate_drafts(question, context)
    best_answer = verify_drafts(question, context, drafts)

    print(f"Q: {question}")

    if best_answer is not None:
        print(f"A: {best_answer['answer']}\n")
        # For simplicity, compare with the first answer (gold) provided in the dataset
        if best_answer['answer'].lower() in sample['answers']['text'][0].lower():
            correct += 1
    else:
        print("No valid draft found.\n")

accuracy = correct / total * 100
print(f"Accuracy: {accuracy}%")技术支持

该代码片段在数据集的 10 个样本上评估简化版推测式 RAG 的准确率。对于每个样本，它基于给定问题与上下文生成多个答案草稿，并通过验证函数选出最佳草稿。若找到有效答案，则将其与数据集提供的标准答案进行对比；如匹配，则correct 计数加一。

最后，代码计算并打印正确答案占总样本的百分比。本代码在该简化数据集上得到 90% 的准确率。

推测式 RAG 的挑战

尽管推测式 RAG 表现良好，但仍面临一些挑战：

• 实现复杂度高：
• 需要协调多个模型的配合。
• 集成到现有系统的复杂度增加。
• 训练开销：
• 训练 RAG 起草器需要额外步骤。
• 需要额外的资源与时间。
• 使部署更复杂，尤其在需要快速上线的环境中。
• 对高质量检索的依赖：
• 效果高度依赖检索文档的质量。
• 糟糕的检索会导致次优草稿。

推测式 RAG 的应用

推测式 RAG 适用于对准确、高效信息检索要求较高的众多领域。以下是一些潜在应用：

• 知识密集型问答：
• 在准确性与相关性至关重要的任务中，推测式 RAG 表现出色。
• 该框架可借助多样文档子集，为复杂查询生成准确且证据充分的答案。
• 实时信息检索：
• 在需要最新信息的场景中，推测式 RAG 可提升响应速度与可靠性。
• 该框架在保持准确性的同时降低延迟，非常适合实时应用。
• 医学与法律文本分析：
• 可用于对复杂文档的准确解读至关重要的领域。
• 通过生成多个草稿并基于理由选择最准确者，确保关键决策建立在可靠信息之上。
• 教育工具：
• 可用于为在线学习平台创建准确且具备上下文的内容。
• 确保学习者获得可靠、全面的信息，提升整体学习体验。

结论

推测式 RAG 通过提升准确性、降低延迟、提高资源效率与可扩展性，强化了检索增强生成。我们展示了其基于 Hugging Face Transformers 的实现。

尽管仍有挑战，推测式 RAG 在多个领域的潜力使其值得进一步研究。

若您希望进一步学习 RAG 实现，推荐以下教程：

• 基于 LangGraph 的纠错式 RAG（CRAG）实现
• RankGPT 作为 RAG 重排序智能体（教程）
• 使用 Llama 3.1 8B、Ollama 与 Langchain 的 RAG