Gumbel Reranking: Differentiable End-to-End Reranker Optimization
This work addresses a bottleneck in RAG systems for information retrieval by improving reranker training, though it is incremental as it builds on existing distillation and attention methods.
The paper tackles the challenge of fine-tuning rerankers in RAG systems due to scarce annotated data, by proposing Gumbel Reranking, an end-to-end training framework that reframes reranking as an attention-mask problem, resulting in a 10.4% improvement in recall on HotpotQA for identifying indirectly relevant documents.
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.