Minder is meer: Vroegtijdig stoppen van uitrol voor on-policy distillatie

Samenvatting

Op-beleidsdistillatie is recentelijk naar voren gekomen als een veelbelovend alternatief voor standaard sequentieniveau-imitatieleren, waarbij een student wordt getraind door zijn eigen gegenereerde rollouts te scoren met een lerarenmodel. We observeren echter het 'Off-policy Teacher Decay'-probleem in dit paradigma: voor de latere tokens, met de eerdere trajectorie van de student als context die off-policy is ten opzichte van de leraar, neemt het vermogen van de leraar om een corrigerende score te produceren af, en kan het terugvallen op token-voltooiingsgedrag dat in de pre-trainingsfase is geleerd. We verifiëren dit probleem empirisch en stellen Early Stopping Rollout (ESR) voor om het te verhelpen: een eenvoudige maar effectieve distillatiestrategie die simpelweg de rolloutgeneratie beperkt tot de eerste <N> antwoordtokens. We tonen aan dat ESR zowel de prestaties van volledige rollout OPD overtreft voor verschillende modelgroottes, -families, taken en trainingsregimes, en een veel hogere GPU-efficiëntie en trainingsstabiliteit vertoont, vooral in scenario's met verschillende modelfamilies. We onderzoeken verder het mechanisme achter deze verrassende prestaties en ontdekten het 'Cascading Alignment'- en 'Sub-mode Commitment'-effect van ESR dat mogelijk verklaart waarom het effectief werkt en soms zelfs de prestaties van het lerarenmodel overtreft. Daarnaast tonen we aan dat deze positiegebaseerde tokenselectiestrategie niet volledig verklaarbaar is door KL-divergentie en entropiesignalen.

English

On-policy distillation has recently emerged as a promising alternative to standard sequence-level imitation, training a student by scoring its own rollouts with a teacher model. However, we observe ``Off-policy Teacher Decay'' problem in this paradigm: for the later tokens, with student's earlier trajectory as context that is off-policy to the teacher, the teacher's ability to produce a corrective score would decay, and may fall back to token-completion behavior learned in the pre-training stage. We empirically verify this problem, and we propose Early Stopping Rollout (ESR) to fix it: a simple yet effective distillation strategy that simply restricts the rollout generation to the first response tokens. We show that ESR both surpasses the full rollout OPD performance across model size, family, tasks and training regime, and exhibit much higher GPU efficiency and training stability, especially under cross model family scenarios. We further investigate the mechanism behind this surprising performance and discovered "Cascading Alignment" and "Sub-mode Commitment" effect of ESR that may explain why it works effectively and even sometimes exceeding the teacher model performance. Besides, we show that this position-based token selection strategy cannot be fully explainable by KL divergence and entropy signals.