PaperlayerView exact revision on arXiv

Source-linked AI summary

Gumbel Reranking: Differentiable End-to-End Reranker Optimization

Siyuan Huang, Zhiyuan Ma, Jintao Du, Changhua Meng, Weiqiang Wang, Jingwen Leng, Minyi Guo, Zhouhan Lin

arXiv:2502.11116v2cs.CLcs.IR

TL;DR

Scarce labeled query-document pairs and limitations of distillation make end-to-end reranker training challenging. Gumbel Reranking learns a stochastic document-wise Top-k attention mask and optimizes the RAG language loss; experiments report gains across settings, including a 10.4% improvement in Recall@5 on HotpotQA for indirectly relevant documents.

  • Problem

    Reranker fine-tuning is challenged by scarce labeled query-document pairs, while distillation objectives may not directly minimize generation loss and can overlook document interdependencies.

    Source: §1 · p. 1

  • Method

    Gumbel Reranking treats reranking as learning a stochastic document-wise Top-k attention mask, using differentiable sampling to optimize the reranker on the RAG system's overall language-modeling loss.

    Source: §1 · pp. 1–2

  • Results

    G-Rerank improves reranker performance across settings, especially in multi-hop QA, and reports a 10.4% improvement in Recall@5 on HotpotQA for indirectly relevant documents.

    Source: §1 · p. 2, §5 · p. 9

  • Takeaways & Limitations

    The framework is suited to multi-hop QA because subset sampling can capture inter-document relationships and complete evidence chains.

    Source: §3.4 · p. 5, §4.3 · p. 8

    Training relies on parallel encoding/decoding during pre-filling, limiting direct adoption in conventional autoregressive LLMs, although the trained reranker can be adapted to other decoder-only models.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

RAG reranker fine-tuning is hindered by scarce labeled query-document pairs and by distillation objectives that can misalign training with inference or miss document interdependencies. Gumbel Reranking learns a stochastic document-wise Top-k attention mask and optimizes the RAG language loss, with reported gains across settings and a 10.4% Recall@5 improvement on HotpotQA.

  • Scarce labeled query-document pairs complicate downstream reranker fine-tuning, while distillation can misalign objectives and overlook candidate-document interdependencies.Perplexity-based methods evaluate documents in isolation, a limitation for multi-hop reasoning.Source: §1 · p. 1
  • G-Rerank reframes reranking as learning a stochastic, document-wise Top-k attention mask over retrieved candidates.The selected subset is represented through attention computation rather than learned only as a direct ranking output.Source: §1 · p. 1, Figure 1 · p. 2
  • The Gumbel Trick and Relaxed Top-k techniques make the mask differentiable, enabling optimization of reranker sampling weights against overall RAG language-modeling loss.G-Rerank also accounts for interdependencies among candidate documents.Source: §1 · p. 2
  • 10.4% improvement in Recall@5 on HotpotQA is reported for distinguishing indirectly relevant documents.The approach also achieves consistent improvements across mining, reranking, and generation settings.Source: §1 · p. 2

2 Related Works

Prior retrieval and reranking methods adapt models for downstream tasks but leave gaps in generation-objective alignment and modeling relationships among retrieved documents. Related work also explores differentiable subset sampling and soft attention masks as ways to model document selection.

3 Methodology

The method treats reranking as document selection through masked attention, then relaxes discrete top-k selection so the reranker can be trained end to end on the RAG language loss. Its pipeline preserves candidate-document independence during pre-filling and models document interdependencies through subset sampling.

3.1 Problem Setting

  • The setting provides queries and retrieved candidates but lacks relevance annotations, and the non-differentiable top-k operation blocks language-loss gradients to the reranker.Source: §3.1 · p. 3

3.2 Viewing Reranker as Attention Mask

  • Reranking is simulated by masking attention to retain the top-k document set; independent pre-filling and shared positional encodings prevent cross-document leakage and position bias.The top-k documents are treated as a set, so their order is not important.Source: §3.2 · pp. 3–4

3.3 Differentiable Masked Attention

  • Gumbel-Softmax and repeated relaxed top-k sampling turn the hard mask into a soft sampled-subset mask, allowing final language loss to optimize the reranker.Source: §3.3 · p. 4

3.4 Gumbel Reranking Pipeline

  • G-Rerank scores candidates, samples a top-k mask, computes generation loss through masked attention, and minimizes that loss; experiments fix LLM parameters while training the reranker.The framework models complete evidence subsets rather than evaluating each candidate independently.Source: §3.4 · pp. 4–5, Figure 2 · p. 5

4 Experiments

Experiments assess G-Rerank across mining, reranking, generation, indirect-evidence identification, a Gumbel ablation, and a learnable-weight variant. Reported results include gains in indirect-document recall and generation.

4.1 Experimental Setup

  • Experiments use FiD and CEPE-Llama2-7B readers, RankT5-Base and BGE-Base-Reranker, five QA datasets, and four LLM-supervised reranker-training baselines.The datasets cover single-hop and multi-hop QA; FiD is fine-tuned for each task, while CEPE-Llama2-7B is used zero-shot.Source: §4.1 · pp. 5–6

4.2 Main Experiments

  • G-Rerank achieves the best or second-best performance on most datasets; with RankT5, HotpotQA Recall improves by 5.7% in the Reranker Setting.In the Mining Setting, RankT5 improves by 5.3% on HotpotQA.Source: Table 1 · p. 6, §4.2 · p. 7
  • With CEPE-Llama2-7B, G-Rerank gives the best generation performance on NQ and TQA while fine-tuning only the retrieval module and keeping the language model fixed.The reader uses zero-shot capabilities rather than being fine-tuned on downstream datasets.Source: §4.2 · p. 7, Table 3 · p. 7

4.3 Identifying Indirectly Relevant Documents

  • With FiD-Large, G-Rerank improves recall by 10.4% for identifying indirectly relevant HotpotQA documents that are part of the evidence chain but do not contain the answer.The study identifies indirect documents as relevant documents that do not directly contain the final answer.Source: §4.3 · pp. 7–8, Table 2 · p. 7

4.4 Necessity of Gumbel Trick

  • On NQ, EM falls from 46.2 with Gumbel noise to 12.7 without it; without noise, maximum document weight declines and the reranker learns a uniform soft mask.The reported weight trend is shown for the NQ dataset.Source: §4.4 · p. 8, Figure 4 · p. 8

4.5 Learnable Sampling Weights

  • Without a reranker, learnable document sampling weights still yield reliable relevance scores and outperform perplexity- or attention-based scalar metrics in the Mining Setting.The learned weights are initialized to zero.Source: §4.5 · pp. 8–9, Figure 5 · p. 8

5 Conclusion

G-Rerank treats reranking as masked attention and uses Gumbel Trick and Relaxed Top-k to optimize document-wise attention masks. Experiments show improved reranker performance, especially on multi-hop QA tasks.

5 Conclusion

  • G-Rerank reframes reranking as masked attention to enable direct optimization of document-wise attention masks.The framework uses the Gumbel Trick and Relaxed Top-k.Source: §5 · p. 9
  • Experiments across different settings show improved reranker performance, especially on multi-hop QA tasks.Source: §5 · p. 9
  • The method captures document interdependencies and aligns retrieval and generation objectives.Source: §5 · p. 9

6 Limitations

G-Rerank's training applicability is constrained by a need for parallel encoding/decoding during pre-filling, but the trained reranker can be adapted to other decoder-only models. The method also requires tuning Gumbel-Softmax hyperparameters.

6 Limitations

  • Parallel encoding/decoding during pre-filling limits direct adoption in conventional autoregressive LLMs.After training, the reranker is independent of a specific LLM and can be adapted to other decoder-only models.Source: §6 · p. 9
  • Gumbel-Softmax temperature τ and scaling factor κ require tuning, with τ = 0.5 and κ = 1.0 providing robust, stable performance across architectures and datasets.Source: §6 · p. 9
  • The framework primarily constrains training models rather than the final inference model.The reranker requires the specified model capabilities only during training.Source: §6 · p. 9

7 Ethical Considerations

RAG performance depends on training-data and underlying-model quality, and bias in training data can amplify biases in retrieved results and affect downstream applications.

7 Ethical Considerations

  • Improving RAG accuracy does not eliminate inherent risks of biased data or model outputs.Source: §7 · p. 9
  • RAG performance still heavily depends on training-data quality and the underlying models.Source: §7 · p. 9
  • Bias in training data, particularly for domain-specific queries, can amplify biases in retrieved results and impact downstream applications.Source: §7 · p. 9

A Effect of hyper-parameters on the Training Process

The appendix examines how temperature τ and scaling factor κ shape document sampling during training. Temperature controls softmax sharpness and exploration, while κ shifts the influence of reranker scores relative to Gumbel noise.

A Effect of hyper-parameters on the Training Process

  • A small τ sharpens document selection toward the highest-scoring candidate, while a larger τ smooths sampling and promotes exploration and diversity.Larger τ may help avoid local optima and improve generalization during training.Source: A Effect of hyper-parameters on the… · p. 13
  • As τ approaches zero, selection becomes deterministic and focuses on the highest-scoring document, encouraging exploitation and rapid convergence.Source: A Effect of hyper-parameters on the… · p. 13
  • A small κ lets Gumbel noise dominate, increasing exploration, while a large κ gives reranker scores more influence and makes selection more deterministic.An overly large κ may limit exploration and potentially reduce generalization through overfitting.Source: A Effect of hyper-parameters on the… · p. 13
  • When κ is small, the Gumbel noise term dominates the selection process and encourages exploration of varied documents.Source: A Effect of hyper-parameters on the… · p. 13

B LLM-Supervised Baselines

The section describes LLM-supervised reranker baselines and situates the paper’s end-to-end Gumbel subset-sampling strategy among FiD and multi-hop QA approaches.

B LLM-Supervised Baselines

  • Attention Distillation (ADist) uses language-model cross-attention, weighted by value-vector norms, to form document relevance targets that the reranker matches with KL divergence.Source: B LLM-Supervised Baselines · p. 13
  • EMDR2 treats retrieved documents as latent variables and maximizes answer log-likelihood using the language model and reranker distributions.Source: B LLM-Supervised Baselines · p. 13
  • Perplexity Distillation (PDist) trains the reranker toward a language-model-derived document distribution, while LOOP estimates relevance by excluding each document from the top-k set.Source: B LLM-Supervised Baselines · pp. 13–16

C.1 Variants of FiD

  • KG-FiD adds graph-based reranking, FiDO reallocates computation for speed, and FiD-Light targets efficiency; RFiD and MG-FiD combine passage ranking with finer evidence signals.Source: §C.1 · p. 16
  • The experiments use classical FiD to focus on reranker training strategies and their impact on open-domain QA performance.Source: §C.1 · p. 16

C.2 Multi-Hop Question Answering

  • For multi-hop QA, the paper targets reranker training in widely used RAG systems and proposes an end-to-end strategy based on Gumbel subset sampling.Source: §C.2 · p. 16

C.3 Dataset Description

The dataset set spans multi-hop reasoning, long-form search queries, and domain-specific comprehension, with a separate table caption identifying HotpotQA FiD-Base experiments.

C.3 Dataset Description

  • HotpotQA tests reasoning across Wikipedia articles, 2WikiHop adds knowledge-graph navigation, and Musique evaluates compositional reasoning through simpler sub-questions.Source: §C.3 · pp. 16–17
  • Table 4 covers HotpotQA with FiD-Base as reader and marks the best and second-best performances.Source: Table 4 · p. 17
  • Natural Questions contains real-world queries requiring information from long documents, while TextbookQA tests textbook-style understanding of textual and diagrammatic content.Source: §C.3 · p. 17

C.4 Additional Results with FiD-Base

The FiD-Base HotpotQA results favor G-Rerank across most metrics; the section also explains challenges in scoring indirect evidence and defines the scalar relevance baselines.

C.4 Additional Results with FiD-Base

  • G-Rerank outperforms other methods across most metrics on HotpotQA when using FiD-Base as reader.Source: §C.4 · p. 17

D Challenges in Handling Indirectly Relevant Documents with EMDR/PDist

E Irrelevant Document Setting

  • The HotpotQA rule labels relevant documents without the final answer as indirect, but can classify partial-answer or fragmentary evidence as indirect too.Source: E Irrelevant Document Setting · p. 18
  • Training data contain 90,447 queries to 664,247 entries with indirect documents, and development data contain 7,405 to 5,966, supporting evaluation of evidence use.Source: E Irrelevant Document Setting · p. 18

F.1 Scalar Relevance Baselines

  • Lowest-perplexity scoring treats lower language-model perplexity as greater relevance, while attention scoring weights each document’s cross-attention by its value-vector norm.Source: §F.1 · p. 18
  • LOOP scores a document by the language-model performance loss after its removal; a larger leave-one-out score marks greater relevance.Source: §F.1 · pp. 18–19

F.2 Our method

The method replaces a dedicated reranker with learnable document weights that are trained through differentiable masked attention.

F.2 Our method

  • The approach assigns each document a zero-initialized learnable scalar weight and updates the weights using language-model loss through stochastic relaxed top-k selection.Source: §F.2 · p. 19

G URLs and Licenses

Loading 2502.11116v2…