Modelos de Lenguaje de Autoregresivo a Difusión Eficientes en Datos mediante Destilación On-Policy

Resumen

Estudiamos la transformación de modelos autorregresivos (ARLM) en modelos de lenguaje de difusión (DLM). En lugar de preentrenar desde cero, trabajos previos reemplazan la atención causal en los ARLM con atención bidireccional y luego entrenan el modelo resultante utilizando un objetivo de DLM. Sin embargo, estos enfoques incurren en dos cambios de distribución. Primero, la transición de un objetivo de predicción del siguiente token a un objetivo de DLM puede descartar conocimiento adquirido por el ARLM durante el entrenamiento. Segundo, los DLM estándar sufren un desajuste entre entrenamiento e inferencia, ya que la pérdida de entrenamiento se define sobre secuencias enmascaradas aleatoriamente, en lugar de las trayectorias encontradas en la inferencia producidas por decodificación basada en confianza. Para abordar ambos desafíos, introducimos un Modelo de Lenguaje de Difusión On-Policy (OPDLM), en el cual se emplea Destilación On-Policy (OPD) para la transformación de ARLM a DLM. Específicamente, OPDLM se entrena mediante auto-OPD, donde el estudiante, un ARLM con atención bidireccional, genera sus propias trayectorias, y el profesor, el ARLM original congelado, destila su conocimiento proporcionando logits objetivo sobre estas trayectorias. Al entrenar directamente de manera on-policy, OPDLM elimina el desajuste entre entrenamiento e inferencia en los DLM, mientras que la destilación del modelo original mejora la retención del conocimiento del ARLM. Los resultados empíricos demuestran que OPDLM requiere de 15 a 7,000 veces menos tokens de entrenamiento con un rendimiento sólido en una amplia variedad de tareas. OPDLM evita el costo prohibitivo del preentrenamiento de DLM y posiciona la transformación de DLM como una forma de post-entrenamiento de 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.