기술 요약: TREK – 탐색을 위한 증류, 정교화를 위한 강화 (Distill to Explore, Reinforce to Refine)

1. 문제 정의

본 논문은 복잡한 추론 작업에 적용된 **그룹 상대 정책 최적화(Group Relative Policy Optimization, GRPO)**에서 발생하는 특정 실패 모드를 다룹니다. GRPO는 현재의 정책(학생 모델)이 이미 유용한 솔루션 궤적을 샘플링할 수 있을 때는 효과적이지만, 정답 솔루션 모드가 학생의 현재 온-폴리시(on-policy) 서포트 범위를 벗어나 있는 **어려운 프롬프트(hard prompts)**에서는 종종 정체되는 현상이 발생합니다.

이러한 어려운 체제에서 학생 모델은 그럴듯하지만 구조적으로 유사한 오답 궤적만을 반복해서 샘플링하게 됩니다. 여기서 부족한 것은 보상 신호가 아닙니다(검증기(verifier)는 정답을 찾기만 하면 점수를 매길 수 있습니다). 진짜 문제는 **솔루션 공간의 탐색 범위(exploration coverage)**입니다. 온-폴리시 방식인 표준 GRPO는 스스로 샘플링하지 못하는 행동을 강화할 수 없습니다. 결과적으로, 더 큰 롤아웃 그룹이나 더 날카로운 상대적 이득(relative advantages)을 사용하더라도 학생의 좁은 현재 서포트 영역을 벗어나지 못하는 경우가 많습니다.

기존의 증류(distillation) 및 교사 기반 방법들은 주로 학생이 이미 도달한 궤적에 대한 감독이나 크레딧 할당(credit assignment)을 개선하는 데 집중합니다. 이들은 학생의 자가 탐색이 실패했을 때 유용한 행동을 어떻게 하면 샘플링 가능하게 만들 것인가라는 핵심 과제를 직접적으로 해결하지 못합니다. 즉, **학생의 도움 없는 탐색이 실패할 때, 어떻게 유용한 행동을 샘플링 가능한 상태로 만들 것인가?**가 핵심 과제입니다.

2. 방법론: TREK

저자들은 **TREK (Teacher-Routed Exploration via Forward KL)**를 제안하며, 이는 증류를 단순한 모방이 아닌 **탐색 서포트 확장(exploration support expansion)**으로 재정의하는 단계적 절차입니다.

핵심 철학

TREK는 교사(외부 모델 또는 자기 증강 모델)를 보편적인 행동 목표가 아니라 하나의 **제안 메커니즘(proposal mechanism)**으로 취급합니다. 이 방법은 출력 전용 인터페이스(output-only interface)를 사용하므로, 교사의 로짓(logits), 은닉 상태(hidden states), 또는 확률값이 아닌 검증된 출력 궤적만을 필요로 합니다. 이를 통해 블랙박스 교사, 화이트박스 교사, 또는 추가적인 추론 시간 컨텍스트(예: 성찰, 더 긴 추론 예산, 또는 실패 사례 학습)를 사용하는 동일 모델을 모두 활용할 수 있습니다.

TREK 알고리즘

이 절차는 **프롬프트 수준의 제안 가용성(prompt-level proposal availability)**과 **궤적 수준의 도달 가능성(trajectory-level reachability)**을 분리합니다.

1. 프롬프트 라우팅 (어려운 프롬프트 마이닝):

   • 각 프롬프트 $x$에 대해, 시스템은 $K$번의 롤아웃을 샘플링하여 도움 없는 학생의 통과율 $p_S(x)$를 추정합니다.
   • 만약 $p_S(x) \leq \tau_{low}$라면 (학생이 실패하는 어려운 프롬프트임을 나타냄), 해당 프롬프트는 제안 단계로 라우팅됩니다. 쉬운 프롬프트는 표준 GRPO 스트림에 그대로 남습니다.

2. 제안 생성 및 선택:

   • 제안 소스 $\pi_T$ (예: 더 강력한 모델 또는 추가 컨텍스트를 가진 동일 모델)가 라우팅된 어려운 프롬프트에 대한 후보 솔루션을 생성합니다.
   • 궤적들은 검증기 $V$에 의해 필터링됩니다. 검증된 성공 사례만 유지됩니다.
   • 도달 가능성 필터링 (Reachability Filtering): 검증된 궤적들 중, 시스템은 현재 학생의 가능성(절단된 길이 정규화된 음의 로그 가능도, $d_S$로 측정)에 따라 순위를 매깁니다.
   • 상위 $r$개의 학생 근접 검증 제안들이 유지됩니다. 이는 학생이 현재 능력에서 너무 멀리 떨어진 솔루션을 모방하도록 요구받지 않도록 하여, 학생이 개연성 있게 도달할 수 있는 모드에 집중하게 합니다.

3. Forward-KL 통합 (Warm-Start):

   • 유지된 궤적들은 데이터셋 $D_{scaf}$를 형성합니다.
   • 짧은 Forward-KL 업데이트가 적용됩니다: $\min KL(q_{prop} \parallel \pi_\theta)$.
   • 이 목적 함수는 학생이 검증된 제안 모드에 낮은 확률을 할당하는 것에 대해 페널티를 부여하며, 결과적으로 이러한 솔루션 모드를 학생의 서포트 안으로 "끌어당기는" 역할을 합니다. 이는 표준 증류(종종 $\pi_\theta \parallel q_{prop}$를 최소화하거나 온-폴리시 증류를 사용하는 방식)와는 구별됩니다.

4. 온-폴리시 정교화 (On-Policy Refinement):

   • 짧은 Forward-KL 웜스타트 이후, 프롬프트들은 표준 GRPO 롤아웃 스트림으로 돌아옵갑니다.
   • 이제 검증된 모드를 포함하도록 서포트가 확장된 학생 모델은, 이를 샘플링하고 검증기로부터 상대적 이득을 받아 정책을 더욱 정교하게 다듬을 수 있습니다.

3. 주요 기여

1. 실패 모드 식별: 본 논문은 어려운 추론 프롬프트에 대한 GRPO의 별도 병목 현상으로서 "누락된 탐색 서포트(missing exploration support)"를 식별했습니다. 이는 보상 희소성(reward sparsity)과는 다른 문제입니다.
2. 탐색으로서의 증류: 증류를 이미 샘플링된 궤적에 대한 감독을 밀도 있게 만드는 것이 아니라, 학생의 도달 가능한 서포트를 확장하기 위한 관점으로 공식화했습니다.
3. TREK 절차: 프롬프트 난이도와 궤적 도달 가능성을 분리하고, 검증된 제안을 공고히 하기 위해 짧은 Forward-KL 단계를 사용하는 단순한 라우팅 알고리즘을 제시합니다.
4. 블랙박스 호환성: 검증된 출력만을 요구함으로써, 내부 상태에 접근할 수 없는 환경에서도 외부 블랙박스 교사, 화이트박스 교사, 또는 자기 성찰 루프를 사용할 수 있도록 하여 범용성을 높였습니다.
5. 학습 효율성: 가장 어려운 작업 유형에 대해 RL이 시작되기 전에 검증된 솔루션 모드를 주입하는 것이 학습 성능에 도달하는 데 필요한 최적화 단계를 크게 단축함을 입증했습니다.

4. 실험 결과

저자들은 수학적 추론(AIME 2024/2025) 및 에이전트 계획 작업(ALFWorld, ScienceWorld)에서 TREK을 평가했습니다.

수학적 추론 (Qwen3 모델)

• 설정: 1.7B, 8B, 14B Qwen3 모델을 대상으로 테스트했습니다. 제안은 DeepSeek-V4(외부) 또는 추가 컨텍스트를 가진 동일한 Qwen3 모델(자기 컨텍스트)로부터 제공되었습니다.
• 성능: TREK은 모든 규모에서 직접 GRPO보다 일관되게 우수한 성능을 보였습니다.
  • Qwen3-8B: DeepSeek-V4 제안을 사용하여 AIME 2025를 36.9%에서 40.3%로, AIME 2024를 47.9%에서 51.1%로 향상시켰습니다.
  • 자기 컨텍스트 변형: 외부 교사 없이도, 자기 컨텍스트 변형은 AIME 2025를 38.5%, AIME 2024를 49.6%로 개선했습니다.
• Ablation (절제 실험): Forward-KL 통합을 온-폴리시 증류(OPD) 스타일의 감독으로 대체했을 때 더 낮은 성능을 보였으며, 이는 Forward-KL 목적 함수가 어려운 프롬프트의 모드 커버리지를 확장하는 데 더 우수함을 확인시켜 줍니다.

에이전트 작업 (Qwen2.5-7B-Instruct)

• ALFWorld: 성공률이 75.8%(GRPO)에서 82.8%(DeepSeek-V4를 사용한 TREK) 및 80.4%(자기 컨텍스트)로 증가했습니다.
• ScienceWorld: 성공률이 12.5%(GRPO)에서 26.7%(DeepSeek-V4를 사용한 TREK) 및 23.4%(자기 컨텍스트)로 두 배 이상 증가했습니다.
• 작업 난이도 분석: 이득은 GRPO의 기본 성능이 가장 낮았던 가장 어려운 작업 유형(예: "Heat & Place", "Examine in Light")에 집중되었습니다. 이는 이 방법이 탐색 병목 현상을 타겟팅함을 확인해 줍니다.
• 학습 역학 (Training Dynamics): TREK은 GRPO에 비해 훨씬 이른 시점에 높은 성공률을 달려냈습니다. 예를 들어, ALFWorld에서 TREK은 스텝 20까지 60% 성공률을 돌파한 반면, GRPO는 50% 미만에 머물렀으며 동일한 수준에 도달하기 위해 약 5배 더 많은 스텝이 필요했습니다.

5. 의의 및 주장

본 논문은 TREK이 보상 획득 궤적을 샘플링할 수 없는 능력이 주요 병목인 온-폴리시 추론 시스템을 개선하기 위한 일반화 가능한 프레임워크를 제공한다고 주장합니다.

• 상호 보완성: TREK은 학생의 현재 서포트 내에서 작동하는 기존 방법들(토큰 가중치나 크레딧 할당 개선 등)과 상호 보완적인 관계로 제시됩니다. TREK은 그보다 앞선 단계에서 작동하여 서포트 자체가 존재하도록 보장합니다.
• 효율성: 주요 장점은 반드시 점진적인 천장(asymptotic ceiling)을 높이는 것이 아니라(GRPO도 엄청난 컴퓨팅 자원을 투입하면 유사한 수준에 도달할 수 있음), 어려운 솔루션 모드를 발견하는 데 필요한 학습 예산을 획기적으로 줄이는 것입니다.
• 실용성: 교사의 내부 상태가 아닌 검증된 출력에 의존함으로써, 블랙박스 교사나 자기 성찰 루프가 흔히 사용되지만 내부 모델 상태에는 접근할 수 없는 실제 배포 환경에서도 견고하게 작동합니다.

저자들은 이 "탐색을 위한 증류, 정교화를 위한 강화" 패러다임이 다른 사후 학습 프레임워크(예: DAPO, GSPO) 및 지속적 기술 습득 시나리오로 확장될 수 있으며, 검증된 제안을 통해 새로운 능력을 도입한 후 온-폴리시 강화를 통해 이를 안정화하는 구조적인 방법을 제공한다고 결론짓습니다.