SD-Search: On-Policy Hindsight Self-Distillation for Search-Augmented Reasoning

2026-05-18

AI Summary

This work addresses the challenge in search-augmented reasoning that existing methods, relying solely on trajectory-level rewards and lacking fine-grained step-level supervision, struggle to optimize query quality. The authors propose SD-Search, which introduces, for the first time, an internal self-distillation mechanism from a hindsight perspective within the policy itself. Specifically, it employs an on-policy hindsight self-distillation framework where compact hindsight blocks serve as teacher conditions to guide online token-level Jensen–Shannon divergence-based distillation onto the student policy, complemented by GRPO trajectory rewards. Notably, this approach generates endogenous dense supervision signals without requiring external large models or human annotations, significantly enhancing both query quality and overall reasoning performance—all without introducing additional training stages—thereby demonstrating the efficacy and advantages of endogenous step-level supervision.

Abstract

Search-augmented reasoning agents interleave internal reasoning with calls to an external retriever, and their performance relies on the quality of each issued query. However, under outcome-reward reinforcement learning, every search decision in a rollout shares the same trajectory-level reward, leaving individual queries without step-specific credit. Recent process-supervision approaches address this gap by drawing step-level signals from outside the policy, relying either on a much larger teacher model, or on sub-question annotations produced by a stronger external system. In contrast, we propose SD-Search, which derives step-level supervision from the policy itself through on-policy hindsight self-distillation, requiring neither an external teacher nor additional annotations. In SD-Search, a single model plays two roles that differ only in conditioning: a student that sees only the context available at inference time, and a teacher that additionally conditions on a compact hindsight block summarizing the search queries and final outcomes of a group of rollouts sampled from the same question. Since the teacher knows how each rollout unfolded and which ones succeeded, its query distribution implicitly marks which decisions were worth making, and the student is trained to recover this behavior by minimizing the token-level Jensen--Shannon divergence to the teacher at search-query positions. This layers a dense, step-level signal on top of GRPO's coarse trajectory reward. Crucially, this signal is produced by the policy itself within the standard RL training loop, without external model inference, auxiliary annotation pipeline, or additional training stage.

Problem

• search-augmented reasoning
• step-level credit assignment
• reinforcement learning
• query quality
• outcome-reward

Innovation

• on-policy self-distillation
• search-augmented reasoning
• step-level supervision
• hindsight distillation
• reinforcement learning