Source-linked AI summary
OPDSearch+: On-Policy Distillation with RL Refinement for Search-Augmented Reasoning
Qinglin Ye, Zhiyuan Gu, Jingjie Xia, Yiheng Zhang, Kaiyan Zhao, Shunchao Zheng, Yuhang Mu, Wenchao Du, Yiming Wang
TL;DR
Small language models struggle with search-augmented reasoning, while existing OPD approaches require costly trajectory construction or task-specific teacher training and can be unstable. OPDSearch+ distills a frozen instruct teacher on live student search trajectories before RL refinement, achieving 0.4402 mean EM across seven QA benchmarks and strong multi-hop gains.
Problem
Existing OPD methods face expensive dynamic multi-turn trajectory collection, costly or unstable task-specific teacher training, and distributional misalignment with off-the-shelf teachers.
Method
OPDSearch+ applies per-position forward-KL distillation from a frozen off-the-shelf teacher to student-generated live-search trajectories, followed by RL refinement.
Results
0.4402 mean EM across seven QA benchmarks outperforms all 3B RL baselines, with gains of 13.1% on HotpotQA and 8.5% on 2WikiMultihopQA.
Takeaways & Limitations
The teacher reshapes the student policy into a stronger initialization, enabling subsequent RL to reach solutions unavailable to RL alone from scratch.
Takeaways & Limitations
The reverse-KL surrogate can produce gradients prone to entropy collapse because its high-coefficient tokens are frequently sampled.
Abstract
from arXiv · showhide
Search-augmented reasoning remains difficult for small language models. On-policy distillation (OPD) from trained teachers offers a promising direction, but suffers from two issues: (1) high-quality multi-turn search trajectories depend on dynamic retriever responses, making SFT data prohibitively expensive to collect at scale; (2) task-specifically trained teachers incur substantial training cost, while directly applying OPD with an off-the-shelf teacher without task-specific fine-tuning constrains the student to the teacher's performance ceiling and suffers from severe training instability. We propose OPDSearch+, the first distillation paradigm that requires no teacher fine-tuning for search-augmented reasoning. We investigate the role of a frozen off-the-shelf instruct model as the teacher in on-policy distillation, and reveal a key insight: the teacher reshapes the student's policy distribution so that subsequent RL converges to a superior solution that RL alone cannot reach. In stage one, the student interacts with a live search engine and is distilled via a per-position forward KL objective, transferring reasoning decomposition and evidence integration skills without any task-specific teacher training. In stage two, RL refines the distilled student from a richer behavioral foundation, achieving performance that RL alone cannot reach from scratch. Across seven QA benchmarks, OPDSearch+ with a 3B model consistently outperforms all prior 3B RL baselines, achieving gains of 13.1% on HotpotQA and 8.5% on 2WikiMultihopQA.
1 Introduction
OPDSearch+ avoids task-specific teacher training by distilling a frozen teacher on live search trajectories, then uses RL to refine the resulting student policy. Across seven QA benchmarks, the 3B system outperforms prior 3B RL baselines, especially on multi-hop tasks.
- Motivation: Existing approaches face costly dynamic search-trajectory construction and expensive, unstable task-specific teacher training.Direct OPD with an off-the-shelf teacher also suffers from teacher–student distributional misalignment.
- Method: OPDSearch+ uses a frozen off-the-shelf teacher to distill student-generated live-search trajectories without task-specific teacher training.The teacher supplies token-level supervision for reasoning decomposition and evidence integration.
- Method: The two-stage pipeline applies RL after distillation, refining the student from a richer behavioral foundation than RL from scratch.The paper frames distillation as reshaping the student policy for subsequent RL.
- Results: 13.1% and 8.5% gains are reported on HotpotQA and 2WikiMultihopQA, respectively, over the best 3B baseline AutoRefine-Base.These are the paper’s highlighted multi-hop improvements.
- Results: 0.4402 mean EM across seven QA benchmarks exceeds all 3B RL baselines, including GiGPO-Instruct at 0.421.The result is reported for the 3B OPDSearch+ model.
2 Related Work
Related work spans RL for search-augmented reasoning, on-policy and offline distillation, and hybrid KD–RL methods. OPDSearch+ is situated among approaches addressing retrieval, reasoning supervision, and training-distribution mismatch.
- RL for Search-Augmented Reasoning: Search-R1 and subsequent methods train search agents with outcome-only rewards, while process-reward methods add step-level supervision with additional engineering.Examples include GPT-generated sub-questions and contrastive group construction.
- Knowledge Distillation for Reasoning LLMs: Offline distillation suffers from train–test distribution mismatch because students train on teacher trajectories but infer from their own distributions.On-policy distillation addresses this by scoring student-generated samples with the teacher.
- Knowledge Distillation for Reasoning LLMs: Recent OPD work studies how KL direction and entropy adaptation affect training stability under on-policy distribution mismatch.The cited methods include reverse-KL, entropy-adaptive mixing, and analyses identifying forward KL as important for stability.
- KD-RL Hybrid Methods: KD–RL hybrids combine distillation with reinforcement learning, but the cited methods operate in static text-generation settings rather than interactive retrieval.Examples include advantage-weighted distillation, proximal OPD, teacher-context generation, and self-conditioned KL.
3 Method
OPDSearch+ uses forward-KL on-policy distillation over student-generated, live-search trajectories, then combines distillation with RL-oriented refinement and regularization. Its design provides token-level supervision, handles retrieval-induced distribution shift, and favors forward-KL stability over reverse-KL alternatives.
- Implicit Supervision: Teacher-favored tokens receive amplified gradients up to Rmax, providing implicit query-quality supervision without a dedicated retrieval reward.This per-token weighting can distinguish teacher-preferred entity-specific queries from low-value question copying.
- Distribution Shift: On-policy distillation reduces retrieval-context mismatch because prefixes are induced by the student interacting with the retriever rather than pre-collected teacher trajectories.Search results depend on prior queries, so policy changes can cascade into different retrieved passages and later reasoning contexts.
- On-Policy Search Distillation: OPDSearch+ distills a frozen teacher on student-generated trajectories that interleave reasoning, queries, retrieved passages, and answers.The student interacts with a live search engine, while retrieved passages are environment-provided rather than model-generated.
- RL Refinement: The objective combines forward-KL distillation with KL regularization against the frozen initial student checkpoint, while teacher-derived token rewards motivate subsequent RL refinement.The teacher signal is dense and stable but reflects teacher preferences rather than task-specific outcomes.
- Forward-KL Objective: Forward KL is estimated with teacher-to-student importance ratios on model-generated token positions, excluding retrieved-passage tokens through state masking.The ratio is clipped between ϵ and Rmax, limiting importance-weight magnitude while introducing controlled bias relative to the exact gradient.
- Objective Choice: Reverse-KL surrogates can collapse because their unbounded high-coefficient terms are frequently sampled when the student over-assigns probability relative to the teacher.Forward-KL instead caps high ratios on tokens that are rarely sampled under the student policy, reducing the frequency of high-magnitude contributions.
4 Experiments
Experiments evaluate OPDSearch+ across seven QA benchmarks, compare forward-KL with alternative objectives, and examine training dynamics and policy changes. The results show strong performance, stable distillation, and behavioral changes that support subsequent RL.
- Main results: OPDSearch+ with a 3B model achieves mean EM 0.4402, surpassing GiGPO-Instruct at 0.421 and all other 3B methods.
- Single-hop QA: OPDSearch+ exceeds AutoRefine-Base on single-hop QA with SH-Avg 0.520 versus 0.512.Its individual scores are NQ 0.4852, TriviaQA 0.6226, and PopQA 0.4513.
- Multi-hop QA: +13.1% on HotpotQA and +8.5% on 2WikiMultihopQA are the largest gains over AutoRefine-Base, while MuSiQue is best among 3B methods.Bamboogle reaches 0.4560, second only to GiGPO’s 0.641 on its 125-sample set.
- Training dynamics: Distillation loss decreases from ∼0.66 to ∼0.20, while query formulation becomes progressively more entity-specific and decomposed.Training remains stable through 501 steps without oscillations or divergence.
- Objective comparison: Forward-KL achieves the highest validation EM and continues improving through 400 steps, whereas reverse-KL alternatives either collapse or plateau below it.Rev-KL Adaptive avoids collapse but reaches lower EM, 0.470 versus 0.493 for forward-KL.
- Policy reshaping: OPD expands behavioral diversity and controls policy drift, with entropy 1.35 versus 0.82 for the base model and KL ∼0.2 versus pure RL’s KL >0.8 before collapse.Average search turns decrease from 3.5 to 2.7, while forward-KL gradient norms decrease from 1.1 to 0.7 over 418 steps.
5 Discussion
Discussion ablations examine teacher scale, training order, and comparisons with offline SFT and pure RL. They support sequential OPD→RL training as a complementary pipeline in which OPD establishes search behaviors and RL refines answer accuracy.
- Complementarity with RL: OPD and RL contribute complementary capabilities: OPD establishes search behaviors, while RL refines answer accuracy.The sequential pipeline reaches 0.4402, whereas joint optimization reaches 0.3660 because of objective interference.
- Teacher scale and two-stage training: OPD→RL achieves HotpotQA 0.458 and 2Wiki 0.426, exceeding pure RL by +32.1% and +35.7%, respectively.Pure RL remains competitive on single-hop QA with SH-Avg 0.514 but underperforms on multi-hop tasks.
- Offline SFT comparison: OPD→RL reaches 0.4402, exceeding Offline SFT→RL at 0.4185 by +5.2%.Pure OPD at 0.3655 also exceeds Offline SFT at 0.3264 by +12.0% without trajectory filtering or data construction.
- Teacher scale and two-stage training: OPD→RL reaches 0.4402, outperforming pure OPD at 0.3655 by +20.4%.The result supports using OPD as a behavioral foundation and RL as a later refinement stage.
- Teacher scale and two-stage training: The two-stage pipeline with a 7B teacher reaches 0.4205, exceeding pure OPD with the larger 14B teacher.
6 Conclusion
OPDSearch+ is an on-policy distillation framework for search-augmented reasoning that uses a frozen off-the-shelf teacher without task-specific training. Its 3B two-stage OPD→RL pipeline achieves mean EM 0.4402 and particularly strong multi-hop results.
- OPDSearch+ uses a frozen off-the-shelf teacher and requires no task-specific training.
- The 3B OPD→RL pipeline achieves mean EM 0.4402 and outperforms all 3B baselines, with particularly strong multi-hop gains.