AsyncOPD: 온-정책 증류는 얼마나 오래되어도 되는가?

초록

온-정책 증류(OPD)는 교사 피드백의 지도를 받아 학생 모델이 자체 롤아웃(rollout)을 통해 학습하는 방식으로, 대규모 언어 모델(LLM) 사후 학습(post-training)에서 점점 더 중요해지고 있다. 그러나 강화 학습(RL)과 마찬가지로 OPD도 온-정책 시스템 병목 현상에 직면하는데, 추론 워크로드의 경우 롤아웃이 훈련 시간을 지배할 수 있기 때문이다. 비동기 학습 파이프라인은 롤아웃 생성과 학습자 업데이트를 분리함으로써 이러한 병목 현상을 완화할 수 있지만, 이 과정에서 지연 정책(stale-policy) 데이터가 도입된다. 선행 연구에서 비동기 RL의 지연 데이터를 다루었지만, OPD에서의 영향은 아직 충분히 탐구되지 않았다. 본 연구는 비동기 OPD에서 지연성(staleness)에 대한 최초의 체계적 연구를 제시하며, 특히 교사 피드백이 지역 KL 손실을 통해 구현되고 전체 어휘에 대한 교사 로짓을 저장하거나 전송하는 비용이 매우 커서 유한한 교사 점수 캐시가 필요한 실용적 환경에 초점을 맞춘다. 첫째, KL 방향이 지연 데이터 문제를 변화시킨다는 점을 보인다. 교사 가중 순방향 KL은 지연된 롤아웃에 더 강건한 반면, 학생 가중 역방향 KL은 취약하다. 둘째, 취약한 역방향 KL 사례에 대해 비동기 RL을 안정화하기 위해 설계된 방법이 OPD의 지연성을 완화할 수 있는지 연구한다. 실험 결과, 이러한 방법들은 학습자 시간에 현재 학생 모델 하에서 역방향 KL 신호를 재계산하는 더 간단한 OPD 특화 대체 방법보다 개선되지 않았다. 셋째, 유한한 교사 점수 캐시가 희소 및 샘플링 기반 역방향 KL OPD 추정기에 대해 편향-분산 트레이드오프를 어떻게 생성하는지 분석한다. 이는 MC 수정 가능성(MC correctability)을 유지하면서 단일 샘플 분산을 줄이는 다중 샘플 몬테카를로(MC)를 동기 부여한다. 마지막으로, 이러한 추정기 선택을 바탕으로 구축된 완전 비동기 OPD 학습 파이프라인인 AsyncOPD를 제시하고 오픈소스로 공개한다. 실험 결과, AsyncOPD는 엄격한 동기식 학습에 비해 학습 처리량을 1.6배에서 3.8배까지 향상시키면서도 유사한 정확도를 달성한다.

English

On-policy distillation (OPD) trains a student on its own rollouts guided by teacher feedback and is becoming increasingly important for large language model (LLM) post-training. Like reinforcement learning (RL), however, OPD faces an on-policy systems bottleneck, as rollouts can dominate training time for reasoning workloads. Asynchronous training pipelines can alleviate this bottleneck by decoupling rollout generation from learner updates, but doing so introduces stale-policy data. While prior work has studied stale data in asynchronous RL, its effects in OPD remain underexplored. We present the first systematic study of staleness in asynchronous OPD, focusing on a practical setting where teacher feedback is implemented through local KL losses and full-vocabulary teacher logits are too expensive to store or transfer, necessitating finite teacher-score caches. We first show that KL direction changes the stale-data problem: teacher-weighted forward KL is more robust to stale rollouts, whereas student-weighted reverse KL is vulnerable. Second, for this vulnerable reverse-KL case, we study whether methods designed to stabilize asynchronous RL can mitigate OPD staleness. In our experiments, they do not improve over a simpler OPD-specific surrogate: recomputing the reverse-KL signal under the current student at learner time. Third, we analyze how finite teacher-score caches create a bias-variance tradeoff for sparse and sampled reverse-KL OPD estimators. This motivates multi-sample Monte Carlo (MC), which preserves MC correctability while reducing one-sample variance. Finally, we present and open-source AsyncOPD, a fully asynchronous OPD training pipeline built from these estimator choices. Experiments show that AsyncOPD improves training throughput by 1.6times to 3.8times over strict synchronous training while reaching comparable accuracy.