UNIFUSION: Adapting Autoregressive Language Models into Discrete Diffusion under a Unified Reverse-Rate Objective
2026-07-27 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors propose a new method called ourmethod that directly adapts GPT-2 models to a type of noise process where all tokens can change during generation, unlike previous methods that used masked noise. They connect several existing diffusion language models by framing their training losses under a unified mathematical objective, allowing them to switch between different noise types. Their approach improves the balance between how well the model predicts words and the diversity of its output across different sampling steps. Experiments show ourmethod outperforms comparable models in text generation quality and accuracy on several benchmarks.
autoregressive language modeldiffusion language modeluniform-noise diffusionmasked diffusionKullback--Leibler divergenceGPT-2generative perplexityunigram entropyWinoGrandecontinual pre-training
Authors
Xiaoyi Jiang, Jingyuan Li, Yixuan Jiang, Wei Liu, Yi Zhu, Zuoqiang Shi, Pipi Hu
Abstract
Existing methods mainly adapt pretrained autoregressive (AR) language models to masked diffusion, whereas we directly adapt them to uniform-noise diffusion, where every token remains editable during sampling. However, adapting AR checkpoints across corruption kernels remains challenging because existing DLMs use different objectives and prediction parameterizations. We establish connections among SEDD, MDLM/GIDD, M2S, and Neural CTMC by expressing their conditional losses as a single generalized Kullback--Leibler objective over model reverse rates. We further derive conversions from clean-token predictions to concrete-score, posterior-mean, and exit-rate/jump parameterizations, yielding a shared \(x_0\) interface that supports switching between mask and uniform kernels. Building on these connections, we propose \ours{}, a simple continual pre-training approach for directly adapting pretrained GPT2 checkpoints to uniform-noise diffusion. Through systematic evaluation of 124M- and 355M-parameter models, we show that \ours{} steadily improves the trade-off between generative perplexity (GenPPL) and unigram entropy as the sampling budget increases from 16 to 256 steps. At 256 steps, \ours{}-S and \ours{}-M achieve GenPPL/entropy pairs of \(97.783/5.2626\) and \(71.516/5.6669\), respectively; no evaluated model at the same scale simultaneously outperforms \ours{} on both metrics. At both scales, \ours{} also achieves the highest WinoGrande, SIQA, and BBH accuracy among the compared diffusion models.