Source-linked AI summary
TREK: Distill to Explore, Reinforce to Refine
Yuanda Xu, Zhengze Zhou, Kayhan Behdin, Jelena Markovic-Voronov, Hejian Sang, Xiaomin Li, Wenhui Zhu, Xinchen Du, Aida Rahmattalabi, Ran He, Sen Na, Zhipeng Wang, Alborz Geramifard
TL;DR
GRPO can stall when hard prompts require solution modes outside the student’s on-policy support. TREK routes verified proposals to those prompts for forward-KL support expansion before returning to GRPO, improving mathematical performance across Qwen3 scales and accelerating hard agentic-task training.
Problem
GRPO is limited on hard prompts when the student rarely samples useful trajectories beyond its current exploration support.
Method
TREK routes low-pass-rate prompts for verified proposal generation, retains student-likely candidates, consolidates them with short forward-KL training, then resumes on-policy GRPO.
Results
TREK improves every tested Qwen3 scale on AIME 2024 and 2025 and accelerates hardest agentic-task training, with ALFWorld exceeding 60% by step 20 while GRPO remains below 50%.
Takeaways & Limitations
Verified output trajectories let TREK expand exploration support using external black-box or white-box teachers, or the same model with additional inference-time context.
Takeaways & Limitations
TREK depends on verifier quality, adds teacher and verifier calls, and ranks proposals with a reachability proxy sensitive to verbosity and surface form.
Abstract
from arXiv · showhide
Group Relative Policy Optimization (GRPO) is effective when the current policy already samples useful reasoning trajectories, but it stalls on hard prompts whose correct solution modes lie outside the student's on-policy support. We propose TREK (Teacher-Routed Exploration via Forward KL), a simple staged procedure that uses distillation not for imitation but for exploration support expansion. A key advantage of TREK is its generality: because it only consumes verified output trajectories, it can use an external black-box teacher, a white-box teacher, or the same model given additional inference-time context, and it can efficiently identify which hard-prompt samples are most worth consolidating even when teacher internals are unavailable. TREK first identifies prompts where the unaided student has very low pass rate, queries a proposal source to produce verified candidate solutions, keeps the top-$r$ proposals ranked by current student likelihood, applies a short forward-KL phase to pull those verified modes into the student's support, and then returns to standard on-policy GRPO refinement. On mathematical reasoning, TREK with DeepSeek-V4 proposals improves Qwen3 models across all tested scales on AIME 2024 and AIME 2025; for Qwen3-8B, it improves AIME 2025 from 36.9 to 40.3 and AIME 2024 from 47.9 to 51.1 (avg@16), while the self-context variant reaches 38.5 and 49.6 without an external teacher. On agentic tasks, TREK raises ALFWorld success rate from 75.8 to 82.8 and ScienceWorld success rate from 12.5 to 26.7; notably, on the hardest task types, TREK achieves high success rates early in training while unaided GRPO requires substantially more optimization steps to reach comparable levels.
1 Introduction
GRPO’s on-policy exploration can stall on hard prompts when useful solution modes lie outside the student’s current support, even when verifiers can score correct solutions. TREK addresses this by treating distillation as targeted exploration support expansion: it routes verified proposals from teachers or context-augmented models into a brief forward-KL phase before standard GRPO refinement.
- Motivation: GRPO is effective when the student already assigns meaningful probability to useful solution modes, but hard prompts can yield repeated, structurally similar wrong trajectories.The limitation arises from inadequate exploration beyond the student’s current support, not necessarily from unavailable verification.
- Motivation: The bottleneck on hard prompts is exploration coverage: larger rollout groups or sharper relative advantages may still search within a narrow region of the student’s support.The verifier is described as able to score a correct solution if one appears.
- TREK: TREK treats a teacher or context-augmented model as a proposal mechanism that discovers verified solutions the unaided student currently misses, rather than as a universal behavior target.Possible context includes execution feedback, reflection, longer reasoning budgets, and failure lessons.
- TREK: TREK identifies hard prompts by unaided student pass rate, invokes proposals selectively, retains the top-r verified trajectories by student likelihood, applies short forward-KL warm-starting, and returns to ordinary GRPO rollouts.The procedure separates prompt-level proposal availability from trajectory-level reachability.
- Contributions: TREK’s output-only proposal interface supports targeted exploration expansion with black-box teachers, white-box teachers, and same-model context augmentation.This framing broadens distillation beyond imitation or token-level credit shaping.
2 Method
TREK treats distillation as exploration support expansion: it routes hard prompts to verified proposal generation, consolidates student-proximal solution modes with a short forward-KL update, then resumes GRPO refinement. The procedure supports external teachers and same-model additional-context proposals while requiring only verified trajectories.
- Motivation: GRPO can reinforce only modes the student already samples, so TREK targets prompts where unaided student pass rates are low and useful trajectories remain undiscovered.A hard prompt becomes eligible for proposal learning only when the proposal source later produces at least one verified trajectory.
- Hard-prompt routing: TREK queries proposal sources only for low-pass-rate prompts, while prompts above the cutoff remain in ordinary GRPO.Routing is refreshed as training proceeds, and the refresh interval trades off routing accuracy against overhead.
- Proposal sources: Proposal sources may be larger models or the same student with extra context, but retained targets must remain expressible as ordinary student outputs.Examples of additional context include verifier-guided retry, self-consistency, search, reflection, environment interaction, failure lessons, or longer reasoning budgets.
- Proposal selection: TREK discards verifier-failing proposals and retains the top-r verified trajectories ranked by increasing student reachability distance.The trimmed length-normalized NLL measures reachability: smaller dS(yT | x) indicates a verified trajectory closer to the student’s current support.
- Staged optimization: The staged schedule is hard-prompt mining → proposal selection → short forward-KL warm-start → ordinary GRPO refinement.Explicit teacher probabilities are unnecessary; in black-box or context-only settings, the forward-KL phase becomes teacher-forced negative log-likelihood on retained verified samples.
3 Experiments
Experiments show that TREK consistently improves hard-prompt reasoning and agentic performance over direct GRPO by expanding the student’s support with verified proposals. Gains are largest where unaided exploration is weakest, while forward-KL consolidation and early training efficiency distinguish TREK from alternatives.
- Mathematical reasoning: TREK improves every tested Qwen3 scale over direct GRPO on AIME 2024 and AIME 2025 using verified DeepSeek-V4 proposals.Gains are +5.4/+2.8 points for Qwen3-1.7B, +3.2/+3.4 for Qwen3-8B, and +6.4/+4.9 for Qwen3-14B on AIME 2024/2025.
- Mathematical reasoning: 49.6/38.5 and 50.2/42.7 are reached by Qwen3-8B and Qwen3-14B self-context TREK on AIME 2024/2025, respectively.The same-model proposal source improves over direct GRPO at every tested scale without requiring an external teacher.
- Ablations: Forward-KL consolidation supports support expansion on hard prompts more effectively than OPD-style supervision.The reported analysis attributes this to forward KL’s explicit missing-support penalty, whereas OPD tends to match surface form rather than expand mode coverage.
- Agentic tasks: 82.8 and 26.7 are TREK’s final success rates on ALFWorld and ScienceWorld with DeepSeek-V4 proposals, versus direct-GRPO baselines of 75.8 and 12.5.Self-context reaches 80.4 on ALFWorld and 23.4 on ScienceWorld.
- Agentic tasks: +19.1 and +12.6 points are TREK’s gains on ALFWorld Heat & Place and Examine in Light, the lowest-baseline task types.Pick Two & Place and Cool & Place gain +7.9 and +7.5, while Pick & Place moves −3.0 and Clean & Place gains +7.7 to reach 100.0%.
- Training efficiency: Approximately 85% can eventually be reached by direct GRPO, but TREK reaches strong performance on hardest task types in a fraction of the required steps.Its primary advantage is training efficiency rather than a higher asymptotic ceiling, because verified proposals are consolidated early.
4 Related Work
TREK differs from prior teacher-informed RL, distillation, and self-improvement methods by using verified proposals to recover exploration on prompts outside the student’s reliable support. Its prompt-level routing, black-box compatibility, and forward-KL consolidation target missing reward-bearing modes rather than improving supervision on already reachable trajectories.
- Teacher-informed RL and distillation: TREK uses teacher or self-context rollouts as verified proposals for missing-support prompts, unlike prior methods that mainly improve supervision or credit assignment on student-sampled behaviors.It filters proposals by prompt difficulty and student-proximal trajectory reachability before brief consolidation.
- Teacher-informed RL and distillation: TREK routes at the prompt level toward very low-pass-rate prompts, requires only verified output trajectories, and directly optimizes mode coverage with forward-KL consolidation.These distinctions separate it from within-group sample routing, teacher-logit-dependent methods, and token-level supervision.
- On-policy distillation: OPD-style and token-teachability methods allocate supervision within student-reachable behavior, whereas TREK addresses cases where the student does not reliably reach reward-bearing modes.The cited OPD design space includes context distillation and black-box OPD, but the paper positions these methods as complementary.
- Classical distillation: Classical distillation establishes output-level transfer without teacher internals, while TREK selectively uses verified trajectories to expand support rather than treating teacher sequences as general-purpose imitation targets.This preserves compatibility with black-box proposal sources while changing the objective and setting.
- Teacher-guided RL and exploration: Teacher-guided RL methods modulate reward-driven updates or token-level credit assignment, whereas TREK prioritizes exploration recovery when no reward-bearing student trajectories are sampled.Inference-time mechanisms such as search, reflection, and execution feedback are used by TREK as training-time proposal generators filtered for reachability.
5 Discussion and Limitations
TREK’s practical effectiveness depends on verifier quality and a reachability proxy that is useful but imperfect. Although verified-only consolidation is conservative, it is used because unverified teacher trajectories can make consolidation unstable when they lie far outside the student’s support.
- 5 Discussion and Limitations: Stronger verifier checks provide more informative proposal selection, making filtering quality a central limitation of verifier-based training.TREK’s proposal filtering inherits the verifier’s strengths and weaknesses.
- 5 Discussion and Limitations: Trimmed NLL ranks proposals effectively but misses fine-grained learnability and remains sensitive to verbosity and surface form.The proxy therefore provides a practical reachability estimate rather than a complete learnability measure.
- 5 Discussion and Limitations: TREK adds teacher queries and verifier calls beyond ordinary GRPO, motivating more adaptive scheduling to reduce overhead.This additional computation is a practical cost of the staged procedure.
- 5 Discussion and Limitations: Verified-only consolidation is conservative, since failed teacher trajectories may still reveal useful reasoning patterns, decompositions, or subgoals.The restriction is primarily practical: at the studied student scales, unverified trajectories are often too far outside current support for stable consolidation.
6 Conclusion
TREK frames distillation as exploration support expansion and applies it through a routed training procedure for hard prompts. Its core design separates deciding when extra proposals are justified from selecting verifier-passing successes the student is most likely to internalize.
- TREK reframes distillation as exploration support expansion for hard prompts.The approach instantiates this view as a routed training procedure.
- The procedure first determines where extra proposal generation is justified at the prompt level.
- It then selects verifier-passing successes based on trajectory-level proximity to the current student.The selection favors proposals the student is most likely to internalize.
A MATH-500 Results
On MATH-500, direct GRPO is near-saturated at roughly 88–90% for Qwen3-8B and 14B, leaving limited headroom for support expansion. TREK nevertheless improves over direct GRPO at every tested scale, with modest gains relative to harder AIME benchmarks.
- MATH-500 results: TREK improves over direct GRPO at every Qwen3 scale on MATH-500, although gains are modest.Direct GRPO is already near-saturated, limiting the available headroom for support-expansion proposals.
- MATH-500 results: Roughly 88–90%: direct GRPO performance on MATH-500 at the 8B and 14B scales is already near saturation.The saturated regime provides little discriminative signal, so the MATH-500 ablation breakdown is omitted.
- MATH-500 results: Verified DeepSeek-V4 rollouts are selected as proposal trajectories, consolidated into Qwen3 through short forward-KL learning, and followed by GRPO refinement.The matched self-context variant instead uses the deployment model with additional inference-time context as the proposal source.
B Experimental Setup Details
The experimental setup fixes routing and selection parameters across training rounds, evaluates math on held-out benchmarks, and uses task-specific configurations for ALFWorld. The pipeline integrates rollouts, proposal generation, verification, and forward-KL updates within a shared distributed training stack.
- Routing and selection: Each round estimates student pass rate from K = 16 unaided rollouts, gates hard prompts at τlow = 1/8, queries M = 4 proposals, and retains r = 2.Retained proposals are verifier-passing candidates ranked by trimmed length-normalized NLL, with quantile trims (α, β) = (0.10, 0.02).
- Math configuration: The math evaluations use the DAPO-Math-17K training pool for verifier-based post-training and evaluate only on held-out benchmarks.The configuration covers MATH-500 and AIME evaluations.
- Math configuration: GRPO hyperparameters remain shared across direct-GRPO, proposal-learning, and refinement comparisons, while TREK adds the specified K, M, and r routing parameters.This keeps the general math post-training protocol consistent across comparisons.
- ALFWorld configuration: ALFWorld uses text-only observations, and final-policy evaluation omits extra planning context and admissible-action filtering.The setup is summarized as an ALFWorld-specific implementation configuration.
- Training infrastructure: The training pipeline runs on HybridFlow/VERL, sharing checkpoints, decoding settings, and verifier outputs across rollouts, proposal generation, verification, and forward-KL updates.These components operate inside a single distributed loop, with selected prompts sent to a proposal worker.
B.1 Optional Pipelined Scheduling for Faster Proposal Learning
TREK can use a delayed two-queue pipeline so ordinary GRPO proceeds while teacher proposals for hard prompts are generated asynchronously. Verified proposals are later rescored, filtered, and distilled only when ready, with several mechanisms reducing overhead.
- Delayed two-queue pipeline: Ordinary GRPO updates proceed immediately for prompts above the hard-prompt cutoff, while low-pass-rate prompts enter a teacher-proposal queue.The pipeline avoids blocking the current student update while proposal rollouts are pending.
- Delayed two-queue pipeline: At round t + k, the current student rescoring verifier-passing trajectories selects the top-r proposals before the forward-KL phase.Current-policy rescoring prevents stale rankings after the prompt was first identified as hard.
- Staged training path: Cached hard prompts receive λLFKL only after verified proposals are ready, while subsequent GRPO uses later on-policy student rollouts rather than teacher trajectories.The staged loss describes a prompt’s training path across delayed updates, not a single update using both trajectory sources.
- Overhead reduction: Overhead is reduced by asynchronous proposal generation, early stopping after r verifier-passing proposals up to budget M, and skipping distillation when cached prompts are no longer hard.Prompts that cease to be hard can instead continue with ordinary GRPO.
C Self-Context Variant: Failure-Lesson Memory
The self-context variant expands proposal support by prompting the same Qwen3 checkpoint with a precomputed memory of lessons distilled from its verifier-rejected math attempts. It retains only verifier-passing proposals, consolidates them through student-likelihood ranking and forward-KL, and excludes the memory from deployment evaluation.
- Memory construction: The method builds a reusable failure-lesson memory by summarizing recurring error patterns from verifier-rejected unaided Qwen3 math rollouts before training.The memory is distilled once from a large rollout pool into compact, transferable rules.
- Proposal generation: At proposal generation, the memory is prepended to each math problem, while only verifier-passing trajectories are retained.The retained proposals undergo the same top-r student-likelihood selection and short forward-KL consolidation as the DeepSeek-V4 variant.
- Failure lessons: The full memory contains roughly forty rules, including checks for requested quantities, hidden constraints, candidate substitution, extraneous roots, and zero denominators.Only seven representative rules are shown to illustrate the memory’s granularity and style.
- Failure lessons: The lessons also cover ordered-versus-unordered counting, complement counting, and verifying gcd(a, m) = 1 before applying Fermat/Euler modular arithmetic.When gcd(a, m) ≠ 1, the memory recommends computing the cycle of a modulo m directly.