OmniOPD : Distillation sur politique sans logits via vérification spéculative

Résumé

La Distillation sur Politique (OPD) entraîne un modèle étudiant sur ses propres trajectoires génératives sous un retour dense au niveau des tokens fourni par un enseignant plus fort, atténuant à la fois le décalage de distribution hors politique du Supervised Fine-Tuning (SFT) et l'attribution de crédit éparse de l'Apprentissage par Renforcement (RL). Cependant, l'OPD standard présente deux limitations couplées. Premièrement, elle exige un accès direct aux logits de l'enseignant au niveau des tokens, excluant ainsi une large classe de modèles propriétaires performants du rôle d'enseignant. Deuxièmement, le signal de logit au niveau du token est lui-même fragile, dépendant d'un chevauchement étroit des tokens suivants plausibles entre l'enseignant et l'étudiant, et enclin à amplifier des motifs dégénérés tels que les boucles de répétition. Dans cet article, nous introduisons OmniOPD, un nouveau cadre qui répond à ces deux limitations grâce à un signal de supervision sans logits et au niveau des chunks. OmniOPD remplace l'appariement déterministe des logits par des simulations de Monte Carlo qui approximent les préférences locales de l'enseignant via une métrique continue de similarité sémantique sur des chunks multi-tokens, et concentre cette supervision via un planificateur d'entropie maximale qui n'audite l'étudiant qu'à ses bifurcations de raisonnement à haute incertitude. Un a priori bayésien de Dirichlet-Multinomiale et une ancre KL du modèle de base bornent davantage la variance de l'échantillonnage discret et empêchent l'effondrement de la politique sur les tokens non audités. Sur des benchmarks compétitifs, OmniOPD surpasse l'approche OPD standard de jusqu'à +28,64% en mathématiques, confirmant que la vérification sémantique au niveau des chunks extrait un signal d'apprentissage plus fiable que l'appariement des logits au niveau des tokens, dont la haute densité d'information est compensée par un bruit et une fragilité significatifs. De plus, lorsqu'il est associé à des enseignants boîte noire plus forts tels que Claude-4.5-Haiku et Gemini-2.5-Flash, OmniOPD atteint un gain relatif supplémentaire de +9,54% en mathématiques par rapport à son homologue enseignant à poids ouverts, faisant progresser l'étudiant au-delà des performances du RL auto-exploratoire.

English

On-Policy Distillation (OPD) trains a student model on its own generative trajectories under dense token-level feedback from a stronger teacher, mitigating both the off-policy distribution shift of Supervised Fine-Tuning (SFT) and the sparse credit assignment of Reinforcement Learning (RL). However, standard OPD faces two coupled limitations. First, it requires direct access to the teacher's token-level logits, excluding a broad class of capable proprietary models from serving as teachers. Second, the token-level logit signal itself is brittle, depending on a narrow overlap of plausible next tokens between teacher and student, and prone to amplifying degenerate patterns such as repetition loops. In this paper, we introduce OmniOPD, a novel framework that addresses both limitations through a logit-free, chunk-level supervision signal. OmniOPD replaces deterministic logit matching with Monte Carlo rollouts that approximate the teacher's local preferences through a continuous semantic similarity metric over multi-token chunks, and concentrates this supervision via a peak-entropy scheduler that audits the student only at its high-uncertainty reasoning forks. A Dirichlet-Multinomial Bayesian prior and a base-model KL anchor further bound the variance of discrete sampling and prevent policy collapse across unaudited tokens. Across competitive benchmarks, OmniOPD surpasses the standard OPD approach by up to +28.64% on math, confirming that chunk-level semantic verification extracts a more reliable learning signal than token-level logit matching, whose high information density is offset by significant noise and brittleness. Furthermore, when paired with stronger black-box teachers such as Claude-4.5-Haiku and Gemini-2.5-Flash, OmniOPD achieves an additional +9.54% relative on math over its open-weight teacher counterpart, advancing the student past the performance of self-exploratory RL.