[논문리뷰] TREK: Distill to Explore, Reinforce to Refine

링크: 논문 PDF로 바로 열기

메타데이터

저자: 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

---

1. Key Terms & Definitions (핵심 용어 및 정의)

• GRPO (Group Relative Policy Optimization): 언어 모델의 추론 성능을 향상시키기 위해, 여러 응답 그룹 내에서 상대적인 보상을 사용하여 정책을 최적화하는 On-policy 강화학습 기법입니다.
• Support Expansion: 학습 과정에서 모델이 보상(Reward)을 받을 수 있는 유용한 행동(Trajectory)을 샘플링할 수 있도록, 정책의 확률 분포 범위를 확장하는 과정을 의미합니다.
• Forward-KL: 제안된 방법론에서 학생 모델이 검증된 정답 모드(Solution modes)를 잘 포함하도록 하기 위해, 제안 분포와 학생 모델 간의 KL Divergence를 최소화하는 정규화 및 학습 목적함수입니다.
• Student-proximal Verified Proposals: 외부 제안 모델이나 추가적인 Context를 통해 생성된 정답 경로 중, 현재 학생 모델이 수행 가능한 범위와 가장 유사하거나 가까운 경로를 선별한 데이터입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 GRPO 학습 중 발생하는 핵심적 한계인 '탐색 공간의 부족(Inadequate exploration)' 문제를 해결하고자 합니다. 기존 GRPO는 학생 모델이 이미 보상을 얻을 수 있는 유용한 궤적을 충분히 샘플링하고 있을 때 최적화가 원활하지만, 난이도가 높은 프롬프트에서는 모델이 실패한 궤적에 갇혀 더 나은 해결책을 스스로 탐색하지 못하는 문제가 있습니다. [Figure 2] 이러한 경우 문제의 본질은 보상 희소성(Reward sparsity)이 아니라, 보상을 받을 수 있는 행동 자체가 학생 모델의 현재 확률 분포(On-policy support) 밖에 존재한다는 점입니다. 저자들은 단순한 모방 학습(Imitation)을 넘어, 정답 궤적을 모델이 스스로 생성할 수 있도록 유도하는 새로운 Distillation 기반 탐색 전략이 필요하다고 주장합니다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

본 논문은 TREK (Teacher-Routed Exploration via Forward KL)이라는 단계적 학습 프레임워크를 제안합니다. 이 방법론은 먼저 학생 모델의 낮은 정답률(Pass rate)을 기반으로 어려운 프롬프트를 식별하고, 해당 프롬프트에 대해서만 외부 제안원(Proposal source)으로부터 검증된 해결책을 획득합니다. [Figure 2] 이후 획득된 경로 중 학생 모델과 가장 근접한 궤적을 선별하여, 짧은 Forward-KL Warm-start 학습을 통해 학생 모델의 분포를 해당 정답 모드로 이동시킵니다. 마지막으로 다시 일반적인 GRPO 정제(Refinement) 과정을 거쳐 최종 정책을 완성합니다. 실험 결과, Qwen3-8B 모델 기준 AIME 2025 벤치마크에서 기존 GRPO 대비 정답률을 36.9%에서 40.3%로, AIME 2024에서는 47.9%에서 51.1%로 유의미하게 향상시켰습니다. [Figure 1] 또한 에이전트 작업인 ALFWorld에서도 성공률을 75.8%에서 82.8%로 높였으며, 특히 난이도가 높은 작업 유형에서 학습 초기 단계에 성공률을 빠르게 끌어올리는 탐색 가속화 효과를 확인하였습니다. [Figure 3]

4. Conclusion & Impact (결론 및 시사점)

본 연구는 Distillation을 단순히 행동 모방이 아닌, '탐색 지원을 위한 지원(Exploration support expansion)'의 관점에서 재해석하여 LLM의 추론 및 에이전트 성능을 개선하는 새로운 접근법을 제시합니다. TREK은 특정 모델 구조에 의존하지 않고 외부 Black-box 모델이나 추가적인 Context를 제안원으로 유연하게 활용할 수 있어 범용성이 높습니다. 결론적으로 이 연구는 강화학습이 정체되는 학습 초기 및 고난도 프롬프트 구간에서 효율적인 탐색 방향을 제공함으로써, 전체 학습 과정의 데이터 효율성을 높이고 모델의 잠재적 성능 상한선을 효과적으로 확보할 수 있음을 입증했습니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.