TREK: Distill to Explore, Reinforce to Refine

2026-07-06 · Source: Takara TLDR - Daily AI Papers · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Robotics & Autonomous Systems · Depth: Expert, medium

Summary

TREK (Teacher-Routed Exploration via Forward KL) is a new staged procedure designed to enhance Group Relative Policy Optimization (GRPO) for large language models, particularly on challenging reasoning tasks where GRPO typically stalls. Unlike traditional distillation for imitation, TREK utilizes distillation to expand the student model's exploration support by incorporating verified output trajectories. This method is highly versatile, compatible with external black-box, white-box, or self-context teachers. The TREK process identifies prompts with low student pass rates, generates verified candidate solutions from a proposal source, selects the top-r proposals based on student likelihood, and then uses a brief forward-KL phase to integrate these verified solution modes into the student's policy before resuming standard GRPO refinement. On mathematical reasoning, TREK with DeepSeek-V4 proposals improved Qwen3-8B's AIME 2025 score from 36.9 to 40.3 and AIME 2024 from 47.9 to 51.1 (avg@16). For agentic tasks, it raised ALFWorld success from 75.8 to 82.8 and ScienceWorld from 12.5 to 26.7.

Key takeaway

For machine learning engineers optimizing LLMs on complex reasoning or agentic tasks, TREK offers a robust approach to overcome policy stagnation. You should consider implementing its staged distillation and forward-KL phase to expand your model's exploration capabilities, especially on prompts where your current policy struggles. This method significantly improves performance on benchmarks like AIME and ALFWorld, accelerating convergence on difficult problems.

Key insights

TREK uses distillation for exploration, not imitation, to expand LLM policy support on hard reasoning tasks, outperforming GRPO.

Principles

• Distillation can expand policy exploration.
• Verified trajectories improve student support.
• Adaptive sampling targets hard prompts.

Method

TREK identifies low-pass-rate prompts, queries a proposal source for verified solutions, ranks them by student likelihood, applies a forward-KL phase to integrate top-r modes, then resumes GRPO.

In practice

• Use DeepSeek-V4 as a proposal source.
• Apply TREK to Qwen3 models.
• Improve agentic task success rates.

Topics

• Large Language Models
• Reinforcement Learning
• Knowledge Distillation
• Policy Optimization
• Mathematical Reasoning
• Agentic AI

Code references

• RLHFlow/Reinforce-Ada
• MiuLab/PairDistill

Best for: Research Scientist, AI Engineer, AI Scientist, Machine Learning Engineer, NLP Engineer

Related on AIssential

• DeferMem: Query-Time Evidence Distillation via Reinforcement Learning for Long-Term Memory QA — Query-time evidence distillation via reinforcement learning significantly improves LLM agent long-term memory QA efficiency and accuracy.
• Knowledge Distillation from Large Reasoning Models to Compact Student Models: A Case Study on the John O Bryan Mathematics Competition — Chain-of-Thought distillation significantly improves compact models' mathematical reasoning, with response length being a critical factor.
• MENTOR: Reinforcement Learning via Flexible Teacher-Optimized Rewards for Tool-Use Distillation — Flexible, process-aware reward structures in RL effectively distill LLM tool-use into SLMs, improving out-of-domain generalization over rigid methods.
• Thinking with Reasoning Skills: Fewer Tokens, More Accuracy — Distilling and retrieving reusable reasoning skills significantly improves LLM efficiency and performance.
• Extreme Region Policy Distillation — Decoupling RL optimization into aggressive signal extraction and constrained distillation improves both sample and KL efficiency.
• DeferMem: Query-Time Evidence Distillation via Reinforcement Learning for Long-Term Memory QA — DeferMem improves LLM long-term memory QA by decoupling retrieval from query-conditioned evidence distillation using RL.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by Takara TLDR - Daily AI Papers.