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.
@article{arxiv.2502.11116,
title = {Gumbel Reranking: Differentiable End-to-End Reranker Optimization},
author = {Siyuan Huang and Zhiyuan Ma and Jintao Du and Changhua Meng and Weiqiang Wang and Jingwen Leng and Minyi Guo and Zhouhan Lin},
journal= {arXiv preprint arXiv:2502.11116},
year = {2025}
}