AsyncOPD: 온폴리시 증류(On-Policy Distillation)는 얼마나 오래된 데이터를 허용할 수 있는가?

온폴리시 증류 (On-Policy Distillation, OPD)는 교사(teacher)의 피드백에 따라 학생(student) 모델이 자신의 롤아웃(rollouts)을 통해 학습하는 방식으로, 대규모 언어 모델 (Large Language Model, LLM)의 사후 학습 (post-training)에서 점점 더 중요해지고 있습니다. 그러나 강화학습 (Reinforcement Learning, RL)과 마찬가지로, OPD는 온폴리시 시스템의 병목 현상에 직면합니다. 추론 워크로드의 경우 롤아웃이 학습 시간의 대부분을 차지할 수 있기 때문입니다. 비동기식 학습 파이프라인 (Asynchronous training pipelines)은 롤아웃 생성과 학습자 (learner) 업데이트를 분리함으로써 이 병목 현상을 완화할 수 있지만, 이 과정에서 오래된 정책 데이터 (stale-policy data)가 도입됩니다. 기존 연구들이 비동기 RL에서의 오래된 데이터(stale data)를 연구해 왔으나, OPD에서의 영향은 아직 충분히 탐구되지 않았습니다.

본 논문에서는 비동기 OPD에서의 데이터 노후화 (staleness)에 대한 첫 번째 체계적인 연구를 제시합니다. 특히 교사 피드백이 로컬 KL 손실 (local KL losses)을 통해 구현되고, 전체 어휘의 교사 로짓 (full-vocabulary teacher logits)을 저장하거나 전송하는 비용이 너무 커서 유한한 교사 점수 캐시 (finite teacher-score caches)가 필요한 실질적인 환경에 초점을 맞춥니다. 첫째, 우리는 KL 방향이 오래된 데이터 문제를 변화시킨다는 것을 보여줍니다. 교사 가중치 기반의 순방향 KL (teacher-weighted forward KL)은 오래된 롤아웃에 더 강건한 반면, 학생 가중치 기반의 역방향 KL (student-weighted reverse KL)은 취약합니다. 둘째, 이러한 취약한 역방향 KL 사례에 대해, 비동기 RL을 안정화하기 위해 설계된 방법들이 OPD의 노후화를 완화할 수 있는지 연구합니다. 실험 결과, 이러한 방법들은 학습 시점에 현재 학생 모델을 사용하여 역방향 KL 신호를 재계산하는 더 간단한 OPD 전용 대리 목적 함수 (surrogate)보다 성능이 개선되지 않았습니다. 셋째, 유한한 교사 점수 캐시가 희소 및 샘플링된 역방향 KL OPD 추정기 (estimators)에 대해 편향-분산 트레이드오프 (bias-variance tradeoff)를 어떻게 생성하는지 분석합니다. 이는 단일 샘플 분산을 줄이면서도 MC 수정 가능성 (MC correctability)을 유지하는 다중 샘플 몬테카를로 (multi-sample Monte Carlo, MC) 방식을 유도합니다. 마지막으로, 이러한 추정기 선택을 바탕으로 구축된 완전 비동기식 OPD 학습 파이프라인인 AsyncOPD를 제시하고 오픈 소스로 공개합니다. 실험 결과, AsyncOPD는 엄격한 동기식 학습 대비 유사한 정확도를 유지하면서도 학습 처리량 (throughput)을 $1.6 imes$에서 $3.8 imes$까지 향상시킵니다.