AsyncOPD: Насколько устаревшей может быть дистилляция на основе текущей политики?

Аннотация

Дистилляция на политике (On-Policy Distillation, OPD) обучает студента на его собственных розыгрышах, направляемых обратной связью от учителя, и приобретает всё большее значение для пост-тренировки больших языковых моделей (Large Language Models, LLM). Как и в случае обучения с подкреплением (Reinforcement Learning, RL), OPD сталкивается с системным узким местом, присущим обучению на политике: для задач рассуждений розыгрыши могут доминировать во времени тренировки. Асинхронные конвейеры обучения могут ослабить это узкое место, отделяя генерацию розыгрышей от обновлений обучающегося, однако такое разделение вносит устаревшие данные политики. Хотя в предыдущих работах изучались устаревшие данные при асинхронном RL, их влияние на OPD остаётся малоизученным. Мы представляем первое систематическое исследование устаревания в асинхронной OPD, сосредоточившись на практическом сценарии, где обратная связь от учителя реализуется через локальные потери KL, а логиты учителя по всему словарю слишком дороги для хранения или передачи, что требует использования конечных кэшей оценок учителя. Сначала мы показываем, что направление KL изменяет проблему устаревших данных: прямое KL с весами учителя (teacher-weighted forward KL) более устойчиво к устаревшим розыгрышам, в то время как обратное KL с весами студента (student-weighted reverse KL) уязвимо. Во-вторых, для этого уязвимого случая обратного KL мы изучаем, могут ли методы, предназначенные для стабилизации асинхронного RL, смягчить устаревание в OPD. В наших экспериментах они не превосходят более простой суррогат, специфичный для OPD: пересчёт сигнала обратного KL при текущем студенте в момент обучения. В-третьих, мы анализируем, как конечные кэши оценок учителя создают компромисс между смещением и дисперсией для разреженных и семплированных оценщиков обратного KL в OPD. Это мотивирует использование многосэмплового метода Монте-Карло (Multi-Sample Monte Carlo, MC), который сохраняет корректируемость MC при снижении дисперсии одного сэмпла. Наконец, мы представляем и публикуем в открытом доступе AsyncOPD — полностью асинхронный конвейер обучения OPD, построенный на основе выбранных оценщиков. Эксперименты показывают, что 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.