Эффективные по данным авторегрессионно-диффузионные языковые модели посредством дистилляции на политике

Аннотация

Мы исследуем преобразование авторегрессионных языковых моделей (ARLM) в диффузионные языковые модели (DLM). В отличие от обучения с нуля, в предыдущих работах каузальное внимание в ARLM заменяется двунаправленным, после чего полученная модель обучается с использованием DLM-целевой функции. Однако такие подходы вызывают два распределённых сдвига. Во-первых, переход от цели предсказания следующего токена к DLM-цели может привести к потере знаний, приобретённых ARLM в процессе обучения. Во-вторых, стандартные DLM страдают от несоответствия между обучением и инференсом, поскольку функция потерь на этапе обучения определяется на случайно замаскированных последовательностях, а не на траекториях, порождаемых на инференсе при декодировании на основе уверенности. Для решения обеих проблем мы предлагаем диффузионную языковую модель на политике (OPDLM), в которой для преобразования ARLM в DLM применяется дистилляция на политике (OPD). В частности, OPDLM обучается с помощью самодистилляции на политике, где ученик (ARLM с двунаправленным вниманием) генерирует собственные траектории, а учитель (исходная замороженная ARLM) передаёт свои знания, предоставляя целевые логиты для этих траекторий. Благодаря прямому обучению в стиле on-policy, OPDLM устраняет несоответствие между обучением и инференсом в DLM, а дистилляция от исходной модели повышает сохранение знаний, полученных от ARLM. Эмпирические результаты показывают, что OPDLM требует в 15–7000 раз меньше обучающих токенов, демонстрируя высокую производительность на широком спектре задач. OPDLM позволяет избежать непомерно высоких затрат на предварительное обучение DLM и позиционирует преобразование DLM как форму постобучения 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.