접두사 재생을 활용한 다중 턴 온-정책 증류

초록

우리는 에이전트 작업을 위한 온-정책 증류(OPD)를 연구하며, 여기서 LLM 에이전트는 다중 턴에 걸쳐 환경과 상호작용하고, 학생은 이러한 다중 턴 상호작용 이력을 통해 교사를 모방한다. 완전 온라인 OPD는 각 업데이트마다 새로운 학생 롤아웃을 환경에서 수행하고 방문한 이력에서 교사 쿼리를 필요로 하므로 비용이 많이 든다. 우리는 Replayed-Prefix On-Policy Distillation(ReOPD)을 제안한다. 이는 사전 수집된 교사 궤적을 재생된 접두사로 재사용하는 환경 외 대안으로, 학생은 선택된 단계에서 행동하고, 교사는 새로운 환경 상호작용을 실행하지 않고 밀집된 단계별 지도를 제공한다. 우리는 다중 턴 OPD가 접두사 함정을 도입함을 보인다. 즉, 이력을 더 학생-온-정책으로 만드는 것은 학생과의 관련성을 높이지만, 교사의 대상이 신뢰할 수 없는 이력에서 교사를 쿼리할 수 있다. 이는 학생 점유와 교사 신뢰성 사이에 양면 분포 이동을 생성한다. ReOPD는 다중 턴 OPD를 신뢰성 인식 접두사 분포 설계로 취급하고, 초기의 낮은 이동 접두사를 강조하는 간단한 단계 감쇠 샘플링 일정으로 이를 구현함으로써 이 문제를 해결한다. Python을 사용한 수학적 추론 및 검색 환경에서 여러 교사 및 학생 모델 규모에 걸쳐, ReOPD는 OPD 수준의 정확도를 유지하거나 개선하며, 학생 훈련 중 도구 호출을 전혀 사용하지 않고, 롤아웃당 OPD보다 최소 4배 빠르다. 따라서 ReOPD는 비용이 많이 드는 에이전트-환경 상호작용을 재사용 가능한 오프라인 자원으로 전환하여, 도구, 작업 및 환경 전반에 걸친 확장 가능한 증류를 가능하게 한다.

English

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 4times faster per rollout than OPD. ReOPD therefore turns expensive agent-environment interaction into a reusable offline resource, enabling scalable distillation across tools, tasks, and environments.