多輪在策略蒸餾與前綴重放
Multi-Turn On-Policy Distillation with Prefix Replay
July 16, 2026
作者: Baohao Liao, Hanze Dong, Christof Monz, Xinxing Xu, Li Dong, Furu Wei
cs.AI
摘要
我們研究面向智能體任務的在線策略蒸餾(OPD),其中大型語言模型智能體與環境進行多輪互動,學生模仿教師在這些多輪互動歷史中的行為。完全在線的 OPD 成本高昂,因為每次更新都需要學生透過環境進行新的 rollout,並在訪問過的歷史狀態上查詢教師。我們提出重放前綴在線策略蒸餾(ReOPD),這是一種脫離環境的替代方法,重複利用預先收集的教師軌跡作為重放前綴:學生在選定步驟行動,而教師提供密集的逐步驟監督,無需執行新的環境互動。我們表明多輪 OPD 引入了一個前綴陷阱:使歷史更符合學生在線策略能提高與學生的相關性,但可能在其目標不可靠的歷史狀態上查詢教師。這造成了學生佔用分佈與教師可靠性分佈之間的雙向偏移。ReOPD 通過將多輪 OPD 視為一個可靠性感知的前綴分佈設計來解決這一問題,並採用簡單的逐步衰減取樣排程來實現,該排程強調早期、低偏移的前綴。在數學推理、Python 和搜尋環境,以及多種教師和學生模型規模下,ReOPD 保持或提高了 OPD 級別的準確率,在學生訓練期間使用零工具呼叫,並且每次 rollout 至少比 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.