用可微方法优化文档重排序,提升RAG系统相关性判断能力
Gumbel Reranking: Differentiable End-to-End Reranker Optimization
- 将重排序转为可微的Top-k注意力掩码学习,避免训练推理不一致
- 在HotpotQA上对间接相关文档的召回率提升10.4%
- 适合需要精准文档筛选的问答与检索系统
RAG系统依赖重排序模型识别相关文档,但因标注的查询-文档对稀缺,微调困难。现有基于知识蒸馏的方法存在训练-推理错位问题,且难以捕捉候选文档间的相互依赖关系。为此,本文将重排序过程重构为注意力掩码问题,提出Gumbel Reranking——一种旨在最小化训练-推理差距的端到端重排序器优化框架。该方法利用Gumbel技巧与松弛Top-k采样,学习随机的、逐文档的Top-k注意力掩码,实现端到端优化,通过最小化整体语言损失完成训练。在多种设置下的实验表明,性能持续提升,尤其在HotpotQA上对间接相关文档的召回率提高10.4%。
原文摘要 · Abstract (English)
RAG systems rely on rerankers to identify relevant documents. However, fine-tuning these models remains challenging due to the scarcity of annotated query-document pairs. Existing distillation-based approaches suffer from training-inference misalignment and fail to capture interdependencies among candidate documents. To overcome these limitations, we reframe the reranking process as an attention-mask problem and propose Gumbel Reranking, an end-to-end training framework for rerankers aimed at minimizing the training-inference gap. In our approach, reranker optimization is reformulated as learning a stochastic, document-wise Top-$k$ attention mask using the Gumbel Trick and Relaxed Top-$k$ Sampling. This formulation enables end-to-end optimization by minimizing the overall language loss. Experiments across various settings consistently demonstrate performance gains, including a 10.4\% improvement in recall on HotpotQA for distinguishing indirectly relevant documents.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。