TREK: 탐색을 위한 증류, 정제를 위한 강화

TREK: Distill to Explore, Reinforce to Refine

그룹 상대 정책 최적화(GRPO)는 현재 정책이 유용한 추론 경로를 이미 샘플링하는 경우 효과적이지만, 학생 모델의 on-policy 지원 범위를 벗어난 올바른 해결 방식을 가진 어려운 프롬프트에서는 성능이 정체될 수 있습니다. 우리는 TREK(Teacher-Routed Exploration via Forward KL)을 제안합니다. 이는 모방이 아닌 탐색 지원 확장을 위해 증류를 사용하는 간단한 단계별 절차입니다. TREK의 주요 장점은 일반성입니다. 검증된 출력 경로만 사용하므로 외부 블랙박스 교사, 화이트박스 교사 또는 추가적인 추론 시간 컨텍스트가 제공된 동일 모델을 사용할 수 있으며, 교사의 내부 정보가 없더라도 어떤 어려운 프롬프트 샘플이 가장 통합할 가치가 있는지 효율적으로 식별할 수 있습니다. TREK은 먼저 학생 모델이 외부 도움 없이 매우 낮은 합격률을 보이는 프롬프트를 식별하고, 제안 소스를 통해 검증된 후보 솔루션을 생성합니다. 현재 학생 모델의 가능성에 따라 상위 $r$개의 제안을 선택하고, 짧은 forward-KL 단계를 적용하여 이러한 검증된 모드를 학생 모델의 지원 범위로 끌어들입니다. 그런 다음 표준 on-policy GRPO 정제 단계로 돌아갑니다. 수학적 추론에서 DeepSeek-V4 제안을 사용한 TREK은 AIME 2024 및 AIME 2025에서 테스트된 모든 규모의 Qwen3 모델 성능을 향상시켰습니다. 특히, Qwen3-8B 모델에서 TREK은 AIME 2025 점수를 36.9에서 40.3으로, AIME 2024 점수를 47.9에서 51.1로 (평균@16 기준) 향상시켰습니다. 자체 컨텍스트 변형의 경우 외부 교사 없이도 각각 38.5와 49.6의 성능을 달성했습니다. 에이전트 기반 작업에서는 TREK이 ALFWorld 성공률을 75.8에서 82.8로, ScienceWorld 성공률을 12.5에서 26.7로 향상시켰습니다. 특히 가장 어려운 작업 유형에서 TREK은 학습 초기에 높은 성공률을 달성하는 반면, 외부 도움 없이 GRPO는 비교 가능한 수준에 도달하기 위해 훨씬 더 많은 최적화 단계를 필요로 합니다.

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.