Moins est plus : rollout avec arrêt précoce pour la distillation sur politique

Résumé

La distillation en ligne (on-policy) est récemment apparue comme une alternative prometteuse à l'imitation standard au niveau des séquences, entraînant un modèle étudiant en évaluant ses propres générations avec un modèle enseignant. Cependant, nous observons un problème de « dégradation de l'enseignant hors ligne » (Off-policy Teacher Decay) dans ce paradigme : pour les tokens ultérieurs, avec la trajectoire antérieure de l'étudiant comme contexte (hors ligne par rapport à l'enseignant), la capacité de l'enseignant à produire un score correctif se dégrade, et peut retomber dans un comportement de complétion de tokens appris lors de la phase de pré-entraînement. Nous vérifions empiriquement ce problème et proposons l'Arrêt Précoce de Génération (Early Stopping Rollout, ESR) pour y remédier : une stratégie de distillation simple mais efficace qui se contente de limiter la génération aux premiers tokens de réponse. Nous montrons que l'ESR surpasse les performances de la distillation en ligne complète (OPD) quelle que soit la taille du modèle, la famille, les tâches et le régime d'entraînement, et présente une efficacité GPU et une stabilité d'entraînement bien supérieures, en particulier dans des scénarios de familles de modèles croisées. Nous étudions plus en détail le mécanisme derrière cette performance surprenante et découvrons des effets d'« alignement en cascade » et d'« engagement en sous-mode » de l'ESR qui pourraient expliquer son efficacité et même parfois sa capacité à surpasser les performances du modèle enseignant. De plus, nous montrons que cette stratégie de sélection de tokens basée sur la position ne peut être entièrement expliquée par la divergence KL et les signaux d'entropie.

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.