Source-linked AI summary
On-Policy Distillation Meets Off-Policy GRPO: Training Compact Instruction-Following Rerankers
Vignesh Prabhakar, Jialing Pan, Anil Babu Ankisettipalli
TL;DR
Offline distillation limits compact rerankers to fixed teacher rankings, motivating a method that combines off-policy teacher strengthening with on-policy student exploration. The resulting reward-based distillation performs best across MAIR tasks and transfers across student architectures, within the study’s evaluated scope.
Problem
Offline imitation on fixed examples restricts compact reranker supervision to the teacher’s observed ranking space, motivating student-driven exploration under distribution shift.
Method
A 4B teacher is strengthened with off-policy GRPO, then a 1B student samples its own rankings and learns from soft teacher-derived rewards using on-policy GRPO.
Results
Across all 126 MAIR tasks, the proposed method obtains the highest task-macro point estimates among evaluated distillation variants and also improves three architecturally distinct student backbones.
Takeaways & Limitations
The results support combining student-driven sampling with permutation-level teacher rewards rather than offline pairwise transfer or on-policy teacher-distribution matching.
Takeaways & Limitations
The study is limited to English-language, text-only reranking, uses one strengthened teacher family, and has mostly single-seed ablations.
Abstract
from arXiv · showhide
Compact instruction-following rerankers are attractive for deployment, but conventional distillation pipelines typically train students by offline imitation of teacher outputs on a fixed set of examples, constraining supervision to the teacher's observed ranking space. We revisit reranker distillation through the lens of reinforcement learning. We propose a two-stage framework combining off-policy teacher optimization with on-policy student distillation. In Stage 1, a 4B teacher reranker is strengthened with off-policy GRPO using LLM-judge feedback on 88K instruction-following examples. In Stage 2, a compact 1B student samples rankings from its own policy and receives soft teacher-derived rewards on those rankings, coupling student exploration with knowledge transfer. Our strongest gains appear under distribution shift. On MAIR-11, the original 11-subset, 869-query evaluation, the proposed student reaches 0.7670 nDCG@6, outperforming offline listwise KD by +4.6 points. Controlled comparisons against offline pairwise RankNet KD and on-policy GKD show that neither changing the offline distillation objective nor moving teacher-distribution matching on-policy reproduces the performance of reward-based on-policy distillation over student-sampled rankings. The advantage persists on MAIR-Full: across all 126 tasks and 9,356 queries, the proposed method obtains the highest task-macro point estimates among the evaluated distillation variants, reaching 0.6808 nDCG@6 and 0.7865 MRR@6. It also exceeds two released 7B RL-trained rerankers on the comparable MAIR-11 evaluation, while the same Stage 2 training procedure consistently improves three architecturally distinct alternative student backbones. On the 9,861-query validation benchmark, the resulting 1B reranker achieves 0.7624 nDCG@6 while providing a favorable quality-efficiency tradeoff relative to larger alternatives.
1 Introduction
The paper addresses compact instruction-following reranking by replacing fixed-distribution imitation with student-driven, reward-based on-policy distillation. Its two-stage framework strengthens a teacher and transfers capability across distribution shifts and architectures.
- Existing compact-reranker distillation typically imitates teacher labels or scores offline on fixed examples, limiting supervision to the teacher’s observed ranking space.
- The proposed pipeline strengthens a 4B teacher with off-policy GRPO and distills it into a 1B student using on-policy GRPO with soft teacher rewards.
- The student samples rankings from its own policy, receiving teacher supervision as soft rewards over ranking permutations rather than token-level distributions.
- +4.6 nDCG@6 points separate the proposed method from offline listwise KD on MAIR-11, with controls ruling out offline pairwise and on-policy distribution matching as equivalent explanations.
- Across MAIR tasks, alternative student architectures, and larger released rerankers, the method shows broad OOD and architectural generalization.
- 0.7624 nDCG@6 is achieved by the resulting 1B reranker on the 9,861-query validation benchmark while reducing latency relative to 4B alternatives.
2 Related Work
Prior work spans instruction-following reranking, prompted large-model ranking, reinforcement-learning alignment, and offline retrieval distillation. The paper positions its contribution at their intersection by applying reward-based on-policy learning to compact ranking models.
- Instruction-following rerankers condition ranking decisions on natural-language criteria, extending neural reranking beyond implicit query intent.
- Prompted and open-source LLM rerankers address ranking at inference time, whereas this work trains a compact 1B reranker for deployment.
- RL methods such as GRPO optimize ranking behavior, while prior offline-RL perspectives emphasize the distributional coverage limits of fixed data.
- Retrieval distillation commonly transfers fixed teacher scores offline, although soft supervision can sometimes let students outperform teachers through regularization.
- The paper instead treats distillation as policy learning: students sample rankings from their own policies and receive teacher-derived rewards.
- The method combines on-policy teacher-distribution distillation with reward-signal optimization, two directions previously treated separately in reranking.
3 Method
The method uses two GRPO stages over Plackett–Luce rankings: teacher optimization with judge rewards, followed by student-policy exploration evaluated with soft teacher rewards. Stabilization terms are secondary to this policy-gradient signal.
- Model naming conventions: ZeRank-2 is the 4B teacher backbone, Base-1B is the Llama-Nemotron-Rerank-1B-v2 backbone, and Distilled-1B is the Stage 2 student.
- Stage 1 — Teacher Training (off-policy GRPO): Stage 1 trains the teacher with off-policy GRPO using sampled rankings and LLM-judge rewards normalized into group-relative advantages.
- Stage 2 — Student Distillation (on-policy GRPO): Stage 2 samples rankings from the student’s own Plackett–Luce policy parameterized by candidate scores.
- Stage 2 — Student Distillation (on-policy GRPO): The fixed teacher converts scalar utilities into rank-normalized soft relevances and scores each student-sampled permutation with nDCG@6.
- Stage 2 — Student Distillation (on-policy GRPO): Teacher rewards are transformed into group-normalized advantages, which weight the student’s policy-gradient update.
- Stage 2 — Student Distillation (on-policy GRPO): The teacher evaluates student-generated rankings rather than generating the rankings used for student training.
- Stage 2 — Student Distillation (on-policy GRPO): KL and entropy regularizers use first-selection marginals, avoiding direct regularization over the full combinatorial permutation distribution.
- 3.4 Hypothesis: Why On-Policy Distillation Generalizes: The central hypothesis is that student-policy supervision expands effective training coverage, producing a larger advantage under distribution shift.
4 Experimental Setup
The experiments use broad instruction-following retrieval data, fixed candidate pools, multiple reranking baselines, and complementary MAIR evaluations. Training combines off-policy teacher optimization with on-policy student distillation, using judge-validated rewards and standard ranking metrics.
- Datasets: The training split contains 88,755 examples, while validation contains 9,861 queries across web, code, mathematics, news, and multi-hop retrieval.
- Datasets: MAIR-11 covers 11 heterogeneous subsets and 869 queries for controlled ablations and external-model comparisons.
- Datasets: MAIR-Full evaluates all 126 MAIR tasks and 9,356 queries using task-macro averages because task sizes vary substantially.
- Evaluation protocol: Candidate pools are frozen and reused across systems, while model inputs are truncated to 512 tokens to prevent retrieval variation from confounding comparisons.
- Models and baselines: The study compares a 4B ZeRank-2 teacher and 1B Llama-Nemotron student across supervised, offline, off-policy, on-policy, pairwise, and external reranking baselines.
- Training: Stage 1 uses off-policy GRPO, while Stage 2 uses on-policy GRPO with teacher-derived soft rewards, KL regularization, entropy regularization, and three training epochs.
- Reward validation: The Stage 1 judge agrees with human relevance annotations on 92.6% of examples, with Cohen’s κ = 0.84, while five judges show high inter-judge agreement.
- Metrics: Primary metrics are nDCG@6 and MRR@6, with query-micro averages for validation and MAIR-11 and task-macro averages for MAIR-Full.
5 Main Results
The proposed reward-based on-policy student performs especially well under distribution shift, outperforming offline and distribution-matching distillation controls on MAIR-11 and MAIR-Full. It also exceeds larger released RL-trained rerankers, while repairing concentrated teacher failures rather than uniformly surpassing the teacher.
- Validation benchmark: 0.7624 nDCG@6 is achieved by Distilled-1B on the 9,861-query validation benchmark, the best overall nDCG@6 among evaluated models.
- MAIR-11: 0.7670 nDCG@6 is achieved by A3 on MAIR-11, outperforming offline listwise KD by 4.6 points.
- Teacher behavior: Teacher GRPO reaches 0.7422 nDCG@6 on validation but 0.6880 on MAIR-11, with concentrated failures on LitSearch and Touche.
- Teacher behavior: A3 does not uniformly outperform its teacher; its aggregate improvement partly comes from repairing the teacher’s concentrated out-of-distribution failures.
- MAIR-11: A3 outperforms off-policy sampling, offline listwise KD, on-policy GKD, RankNet pairwise KD, and distillation from the unaligned teacher on MAIR-11.
- External comparisons: 0.7670 nDCG@6 and 0.8289 MRR@6 allow A3 to exceed Rank-R1-7B and REARANK-7B on the comparable MAIR-11 evaluation.
- MAIR-Full: 0.6808 nDCG@6 and 0.7865 MRR@6 are A3’s highest task-macro point estimates among evaluated distillation variants on MAIR-Full.
6 Ablation Study
The ablations identify student-driven on-policy distillation with soft teacher rewards as the primary source of robust OOD gains, while teacher strengthening is necessary and auxiliary regularizers are secondary. Across validation and MAIR evaluations, reward-based on-policy distillation outperforms alternative sampling and supervision choices, with robustness–fit tradeoffs under distribution shift.
- Sampling policy and supervision form: +4.6 nDCG@6 points over offline listwise KD on MAIR-11, reward-based on-policy distillation outperforming both offline pairwise KD and on-policy teacher-distribution matching.The four-way comparison reports validation nDCG@6 values of 0.7608, 0.7514, 0.7476, and 0.7624 for A1, A10, A9, and A3; MAIR-11 values are 0.7212, 0.7416, 0.7386, and 0.7670.
- Where supervision is evaluated: Teacher-driven off-policy sampling consistently degrades performance versus student sampling, isolating student exploration as the key mechanism when teacher-derived supervision is held fixed.The comparison between A2 and A3 supports evaluating supervision on the student’s own sampled ranking space.
- Teacher strengthening: Stage 1 teacher strengthening is necessary: on-policy learning from the unaligned teacher improves over alternatives but remains well below A3 on MAIR-11.The result attributes the full two-stage gain to combining student-side on-policy learning with prior instruction-aware teacher optimization.
- Soft rewards: Hard-label supervision achieves the best validation performance among core variants but substantially reduces MAIR-11 performance, exposing an in-distribution-fit versus OOD-robustness tradeoff.Soft rewards preserve relative preference information, whereas hard labels fit the in-distribution benchmark more aggressively.
- Auxiliary terms and robustness: Removing KL or entropy leaves performance close to the full objective, while hyperparameter sweeps remain nearly flat on validation and KL has modest OOD sensitivity.These findings position policy-gradient training with soft teacher rewards as dominant and KL and entropy as secondary regularizers.
7 Efficiency Analysis
The efficiency analysis compares single-query latency under fixed hardware and finds that the distilled 1B reranker offers a favorable quality–efficiency tradeoff relative to larger rerankers.
- Table 4 compares single-query inference latency under a fixed hardware setup.
- 27.0ms to 9.2ms: the distilled 1B student reduces mean latency relative to the teacher while improving validation performance.
- The distilled 1B reranker remains competitive with or stronger than evaluated baselines on ranking quality while being substantially faster than 4B rerankers.
8 Discussion and Conclusion
The discussion attributes the strongest distillation behavior to student-driven sampling with permutation-level teacher rewards, and reports broad performance and architectural generalization. It also distinguishes the contributions of teacher strengthening and soft reward supervision.
- Student-driven sampling with permutation-level teacher rewards produces stronger OOD distillation behavior than on-policy teacher-distribution matching or offline pairwise distillation.The A2–A3 comparison isolates the role of evaluating supervision on the student’s own ranking distribution.
- Across all 126 MAIR tasks, A3 obtains the highest task-macro point estimates among the evaluated distillation variants.
- A3 exceeds the evaluated 7B Rank-R1 and REARANK models on comparable MAIR-11 evaluations despite its substantially smaller size.
- Teacher strengthening and soft reward supervision are complementary: unaligned-teacher distillation remains below the full method under distribution shift, while hard-label training generalizes less robustly.
- The same Stage 2 procedure improves BGE-Reranker-v2-Gemma, MXBAI-Rerank-Large-v1, and RankZephyr-7B, supporting transfer across three distinct student architectures.
9 Limitations
The study is bounded to English-language, text-only reranking and relies on LLM-judge rewards with limited human validation. Teacher-side generality and some ablation comparisons remain constrained.
- The study focuses on English-language, text-only reranking.
- LLM-judge rewards may inherit judge bias and calibration error, despite validation against human relevance annotations.That validation contains only 188 examples.
- Most controlled ablations use a single random seed, so small differences among closely clustered objective variants require cautious interpretation.Additional multi-seed evidence is provided for A3, A4, and A5.
- Teacher-side architectural and capability generality remains open because all experiments use the same strengthened ZeRank-2 teacher.The study establishes student-side architectural generality across three student families.
- The Stage 1 teacher is not uniformly robust: aggregate MAIR-11 degradation is concentrated in LitSearch and Touche.
10 Ethical Considerations
The paper identifies bias, misranking, and judge-calibration risks in instruction-following reranking, while reporting broad empirical robustness and cross-architecture gains. These findings are bounded by the English text-only evaluation setting and a single strengthened teacher family.
- Ethical considerations: Judge-derived rewards may inherit calibration errors or systematic preferences, potentially amplifying bias or under-ranking relevant material for some domains or groups.The paper partially mitigates this risk through comparisons with other LLM evaluators and human annotations, but does not eliminate all bias.
- MAIR-Full results: A3 achieves the highest task-macro point estimates among evaluated distillation variants across all 126 MAIR tasks.MAIR-Full contains 9,356 queries, extending the broader distribution-shift evaluation beyond MAIR-11.
- Distribution-shift analysis: The Stage 1 teacher is stronger than Base-1B on 7 of 11 MAIR-11 subsets but exhibits concentrated weaknesses on farther-shifted domains.The largest reported weakness is on LitSearch, rather than a uniform degradation across subsets.
- Cross-architecture generalization: Stage 2 improves all three alternative student backbones on validation and preserves these gains under MAIR-11 distribution shift.The experiments use the same strengthened ZeRank-2 teacher across architectures.
- Robustness: Validation performance remains nearly flat across hyperparameter sweeps, while MAIR-11 shows modest sensitivity, especially to the KL coefficient.The sweep results suggest the reported gains are not dependent on brittle tuning around one configuration.
F.2 MAIR-11 Results
On MAIR-11, the method remains broadly stable across hyperparameter settings, with somewhat greater sensitivity to the KL coefficient than on validation. The benchmark contains 869 queries across 11 subsets.
- KL sensitivity: Larger KL values improve MAIR-11 performance, although the KL term is characterized as an auxiliary stabilizer rather than the main learning source.MAIR-11 is more sensitive to KL than the validation set.
- Evaluation setting: MAIR-11 contains 869 queries across 11 subsets.The evaluation is reported as a distribution-shift benchmark in the surrounding results analysis.
- Hyperparameter robustness: Validation performance is nearly flat across the tested hyperparameter sweeps, while MAIR-11 shows only modest variation.Red stars in Figure 5 mark the default configuration used for the main results.
G Per-Dataset Validation Breakdown Across All Baselines
Across validation datasets and comparisons with larger rerankers, A3 offers a quality-efficiency tradeoff rather than uniform dominance. Its main advantage over offline KD emerges under distribution shift, while objective variants remain closely clustered.
- Per-dataset results: A3 remains competitive across all eight validation datasets and obtains the highest or tied-highest point estimate on WebQA.Jina is strongest on MS MARCO and Robust04, while Cohere Rerank v4.0-fast is strongest on MetaMath and LeetCode.
- Cross-model validation: 0.762 nDCG@6 is achieved by A3 across the full 9,861-query validation benchmark, versus 0.750 for Rank-R1-7B and 0.748 for REARANK-7B.A3 uses a substantially smaller model, while the larger baselines remain stronger on some individual subsets.
- Objective components: A3, A4, and A5 remain tightly clustered, with no consistent winner across MAIR-11 subsets.This supports treating policy-gradient learning with soft teacher rewards as primary, with KL and entropy as auxiliary stabilizers.
- Paired effects: A3’s paired advantage over A1 is small on validation but substantially larger on MAIR-11.The validation difference is directionally reliable but practically negligible according to the reported paired analysis.
- Qualitative analysis: A3 often improves ranking by distinguishing among highly similar candidates when the relevant passage is already present.Many observed failures involve plausible candidate sets whose best instruction- or claim-aligned passage is ranked below distractors.
- Distribution shift: 4.6 points is A3’s nDCG@6 gain over A1 on MAIR-11, increasing from 0.7212 to 0.7670.MRR@6 also increases from 0.7860 to 0.8289, a 4.3-point gain.
- Judge validation: Judge outputs show strong inter-judge consistency and 92.6% agreement with human labels on 174 of 188 manually labeled examples.The reported Cohen’s κ is 0.84, while the paper still treats judge outputs as approximations rather than substitutes for human annotation.