Modelos de Linguagem Autoregressivo-para-Difusão Eficientes em Dados via Destilação On-Policy

Resumo

Estudamos a transformação de modelos autorregressivos (ARLMs) em modelos de linguagem de difusão (DLMs). Em vez de pré-treinar do zero, trabalhos anteriores substituem a atenção causal nos ARLMs por atenção bidirecional e, em seguida, treinam o modelo resultante com um objetivo de DLM. No entanto, essas abordagens incorrem em duas mudanças de distribuição. Primeiro, a transição de um objetivo de predição do próximo token para um objetivo de DLM pode descartar conhecimento adquirido pelo ARLM durante o treinamento. Segundo, DLMs padrão sofrem de uma incompatibilidade treino-inferência, pois a perda de treinamento é definida em sequências mascaradas aleatoriamente, em vez das trajetórias encontradas na inferência, produzidas por decodificação baseada em confiança. Para enfrentar ambos os desafios, introduzimos um Modelo de Linguagem de Difusão On-Policy (OPDLM), no qual a Destilação On-Policy (OPD) é empregada para a transformação de ARLM para DLM. Especificamente, o OPDLM é treinado via auto-OPD, onde o aluno, um ARLM com atenção bidirecional, gera suas próprias trajetórias, e o professor, o ARLM original congelado, destila seu conhecimento fornecendo logits alvo nessas trajetórias. Ao treinar diretamente de forma on-policy, o OPDLM elimina a incompatibilidade treino-inferência nos DLMs, enquanto a destilação do modelo original melhora a retenção de conhecimento do ARLM. Resultados empíricos demonstram que o OPDLM requer de 15 a 7.000 vezes menos tokens de treinamento, com desempenho robusto em uma ampla variedade de tarefas. O OPDLM evita o custo proibitivo do pré-treinamento de DLMs e posiciona a transformação de DLM como uma forma de pós-treinamento 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.