Modèles de langage auto-régressifs à diffusion économes en données par distillation sur politique

Résumé

Nous étudions la transformation des modèles autorégressifs (ARLM) en modèles de langage à diffusion (DLM). Plutôt que de pré-entraîner un modèle de zéro, les travaux précédents remplacent l'attention causale dans les ARLM par une attention bidirectionnelle, puis entraînent le modèle résultant à l'aide d'un objectif DLM. Cependant, ces approches introduisent deux décalages distributionnels. Premièrement, le passage d'un objectif de prédiction du token suivant à un objectif DLM peut faire perdre les connaissances acquises par l'ARLM lors de l'entraînement. Deuxièmement, les DLM standard souffrent d'un décalage entre l'entraînement et l'inférence, car la fonction de perte d'entraînement est définie sur des séquences masquées aléatoirement plutôt que sur les trajectoires rencontrées lors de l'inférence, produites par un décodage basé sur la confiance. Pour relever ces deux défis, nous introduisons un modèle de langage à diffusion sur politique (OPDLM), dans lequel la distillation sur politique (OPD) est utilisée pour la transformation ARLM-DLM. Concrètement, OPDLM est entraîné via auto-OPD, où l'étudiant, un ARLM à attention bidirectionnelle, génère ses propres trajectoires, et l'enseignant, l'ARLM original figé, distille ses connaissances en fournissant des logits cibles sur ces trajectoires. En s'entraînant directement de manière sur politique, OPDLM élimine le décalage entre entraînement et inférence dans les DLM, tandis que la distillation à partir du modèle original améliore la rétention des connaissances de l'ARLM. Les résultats empiriques montrent qu'OPDLM nécessite 15 à 7 000 fois moins de tokens d'entraînement tout en obtenant des performances solides sur une grande variété de tâches. OPDLM évite le coût prohibitif du pré-entraînement DLM et positionne la transformation DLM comme une forme de post-entraînement d'ARLM.

English

We study the transformation of autoregressive models (ARLMs) into diffusion language models (DLMs). Rather than pretraining from scratch, prior work replaces the causal attention in ARLMs with bidirectional attention and then trains the resulting model using a DLM objective. However, these approaches incur two distribution shifts. First, transitioning from a next-token prediction objective to a DLM objective can discard knowledge acquired by the ARLM during training. Second, standard DLMs suffer from a train-inference mismatch, as the training loss is defined on randomly masked sequences rather than the trajectories encountered at inference produced by confidence-based decoding. To address both challenges, we introduce an On-Policy Diffusion Language Model (OPDLM) in which On-Policy Distillation (OPD) is employed for ARLM-to-DLM transformation. Specifically, OPDLM is trained via self-OPD, where the student, an ARLM with bidirectional attention, generates its own trajectories, and the teacher, the original frozen ARLM, distills its knowledge by providing target logits on these trajectories. By training directly in an on-policy manner, OPDLM eliminates the train-inference mismatch in DLMs, while distillation from the original model enhances knowledge retention from the ARLM. Empirical results demonstrate that OPDLM requires 15x to 7,000x fewer training tokens with strong performance across a wide variety of tasks. OPDLM avoids the prohibitive cost of DLM pretraining and positions DLM transformation as a form of ARLM post-training.