Gumbel Reranking: Differentiable End-to-End Reranker Optimization

2025/02/16 by Huang, Siyuan, Ma, Zhiyuan, Du, Jintao +5 · 2 citations
#Computation and Language (cs.CL) #FOS: Computer and information sciences #Information Retrieval (cs.IR)

paper · doi:10.48550/arxiv.2502.11116

Abstract

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.

Cited by

• CLaRa: Bridging Retrieval and Generation with Continuous Latent Reasoning
• MiRAGE: Misconception Detection with Retrieval-Guided Multi-Stage Reasoning and Ensemble Fusion

Related

• Beyond Sequential Reranking: Reranker-Guided Search Improves Reasoning Intensive Retrieval
• ProRank: Prompt Warmup via Reinforcement Learning for Small Language Models Reranking
• GroupRank: A Groupwise Reranking Paradigm Driven by Reinforcement Learning
• ERank: Fusing Supervised Fine-Tuning and Reinforcement Learning for Effective and Efficient Text Reranking
• Drowning in Documents: Consequences of Scaling Reranker Inference