[论文解读] Fine-tune the Entire RAG Architecture (including DPR retriever) for Question-Answering

Shamane Siriwardhana, Rivindu Weerasekera+2 more|arXiv (Cornell University)|Jun 22, 2021

Topic Modeling参考文献 7被引用 4

一句话总结

本文提出对整个检索增强生成（RAG）架构（包括密集段落检索器DPR）进行端到端微调，通过联合训练问题编码器和段落编码器实现。该方法通过并行处理解决了训练过程中重新编码和重新索引知识库的工程挑战，与原始RAG相比，在SQuAD上的精确匹配指标提升了12%（40.02% vs. 28.12%）。

ABSTRACT

In this paper, we illustrate how to fine-tune the entire Retrieval Augment Generation (RAG) architecture in an end-to-end manner. We highlighted the main engineering challenges that needed to be addressed to achieve this objective. We also compare how end-to-end RAG architecture outperforms the original RAG architecture for the task of question answering. We have open-sourced our implementation in the HuggingFace Transformers library.

研究动机与目标

• 实现对完整RAG架构（包括DPR检索器）的端到端训练，该功能在原始Hugging Face实现中此前未被支持。
• 解决在训练过程中更新知识库嵌入和重新索引的工程挑战，这些操作计算成本高且耗时。
• 评估联合微调问题编码器和段落编码器是否相比冻结段落编码器能提升问答性能。
• 开源实现以促进更广泛的研究使用和可复现性。

提出的方法

• 扩展了Hugging Face Transformers库，以支持RAG模型的端到端训练，包括DPR段落编码器。
• 采用双BERT架构：一个用于编码问题，一个用于编码段落，两者均进行联合微调。
• 使用Python多进程和HuggingFace Datasets实现知识库的并行重新编码与重新索引，避免阻塞主训练循环。
• 利用FAISS实现高效的近似最近邻搜索，并在每次重新编码步骤后更新向量索引。
• 通过文档检索概率p(z|x)的梯度反向传播，同时更新问题编码器和段落编码器。
• 采用周期性更新策略——每N个训练步骤执行一次重新编码和重新索引——以平衡性能与计算成本。

实验结果

研究问题

• RQ1 与冻结段落编码器相比，RAG中DPR检索器的端到端微调是否能提升问答性能？
• RQ2 在RAG训练过程中维护动态更新的知识库面临哪些关键工程挑战？
• RQ3 知识库重新编码与重新索引的频率如何影响模型性能与训练效率？
• RQ4 联合训练问题编码器和段落编码器是否能提升领域适应性和泛化能力？

主要发现

• 对包含DPR检索器的RAG架构进行端到端微调，使SQuAD数据集上的精确匹配指标从原始RAG的28.12%提升至40.02%。
• 性能提升表明，对问题编码器和段落编码器进行联合优化可显著提升检索与生成质量。
• 在训练过程中重新编码和重新索引知识库计算成本较高，但通过并行处理使其在合理训练周期内可行。
• 研究证实，在训练过程中更新段落编码器能增强查询与相关段落之间的对齐，尤其在领域特定或非维基百科类设置中表现更优。
• 开源的实现使研究人员能够基于Hugging Face Transformers和PyTorch Lightning复现并扩展该方法。