프리픽스 리플레이를 사용한 다중 회전 온폴리시 증류

Multi-Turn On-Policy Distillation with Prefix Replay

본 연구는 LLM 에이전트가 환경과 여러 단계에 걸쳐 상호작용하고, 학생 모델이 이러한 다중 단계 상호 작용 이력을 기반으로 교사 모델을 모방하는 에이전트 기반 작업에서의 온폴리시 증류(OPD)를 다룹니다. 완전한 온라인 OPD는 각 업데이트마다 새로운 학생 모델의 환경 시뮬레이션과 방문된 이력에 대한 교사 모델 쿼리가 필요하기 때문에 비용이 많이 듭니다. 우리는 Replayed-Prefix On-Policy Distillation (ReOPD)을 제안합니다. 이는 사전 수집된 교사 모델의 경로를 재사용하는 오프라인 방식입니다. 학생 모델은 선택된 단계에서 행동하고, 교사 모델은 새로운 환경 상호 작용 없이 각 단계별로 상세한 지침을 제공합니다. 다중 회전 OPD는 '프리픽스 트랩'이라는 문제를 야기하며, 이는 학생 모델의 정책이 이력과 더 관련성을 높이지만, 학생 모델의 목표가 신뢰할 수 없는 이력에 대한 교사 모델 쿼리를 발생시킬 수 있습니다. 이는 학생 모델의 상호 작용 분포와 교사 모델의 신뢰성 사이의 양방향적인 분포 변화를 초래합니다. ReOPD는 다중 회전 OPD를 신뢰성을 고려한 프리픽스 분포 설계로 간주하고, 초기 단계에 더 높은 가중치를 부여하는 간단한 스텝 감쇠 샘플링 스케줄을 사용하여 이를 구현합니다. 수학적 추론(Python) 및 다양한 교사/학생 모델 크기를 갖는 검색 환경에서 실험 결과, ReOPD는 OPD 수준의 정확도를 유지하거나 향상시키고, 학생 모델 훈련 중 도구 사용을 0으로 만들며, OPD보다 최소 4배 빠른 속도로 학습할 수 있습니다. 따라서 ReOPD는 비싼 에이전트-환경 상호 작용을 재사용 가능한 오프라인 자원으로 전환하여, 다양한 도구, 작업 및 환경에 대한 확장 가능한 증류를 가능하게 합니다.

We study on-policy distillation (OPD) for agentic tasks, where an LLM agent interacts with an environment over multiple turns and a student imitates a teacher over these multi-turn interaction histories. Fully online OPD is costly because each update requires fresh student rollouts through the environment and teacher queries at visited histories. We propose Replayed-Prefix On-Policy Distillation (ReOPD), an off-environment alternative that reuses pre-collected teacher trajectories as replayed prefixes: the student acts at selected steps, while the teacher provides dense per-step supervision without executing new environment interactions. We show that multi-turn OPD introduces a prefix trap: making histories more student-on-policy improves relevance to the student, but can query the teacher on histories where its target is unreliable. This creates a two-sided distribution shift between student occupancy and teacher reliability. ReOPD addresses this by treating multi-turn OPD as a reliability-aware prefix distribution design and implements it with a simple step-decaying sampling schedule that emphasizes early, lower-shift prefixes. Across mathematical reasoning with Python and search environments over multiple teacher and student model scales, ReOPD preserves or improves OPD-level accuracy, uses zero tool calls during student training, and is at least 4$\times$ faster per training step than OPD. ReOPD therefore turns expensive agent-environment interaction into a reusable offline resource, enabling scalable distillation across tools, tasks, and environments.