Reducing Pretraining-Generation Mismatch in Diffusion Language Models
2026-08-10 • Computation and Language
Computation and Language
AI summaryⓘ
The authors explain that traditional autoregressive language models train by predicting the next word from a clean beginning, matching how they are used during text generation. Diffusion language models, which generate text by gradually fixing corrupted tokens, can mix up the clean start and continuation during training, making it harder to generate text based on a prompt. They propose a new training method called Prefix-Conditioned Diffusion (PCD) that keeps the prompt clean while only applying diffusion to the rest, better matching how generation works. Their experiments show PCD improves model performance on several benchmarks without changing how the models generate text.
autoregressive language modelsdiffusion language modelspretraining objectiveprompt-conditioned generationprefix conditioningtoken denoisingattention maskLLaDA2-MiniQwen-1.7Bblock-diffusion models
Authors
Xiaocheng Lu, Huabin Liu, Song Guo, Jianguo Li
Abstract
Autoregressive language models align training and use: generation conditions on a clean prompt, and training predicts future tokens from clean left context. Diffusion language models offer parallel denoising, but native dLLM pretraining can randomly corrupt prompt and continuation tokens together, weakening the clean-prefix interface needed for prompt-conditioned generation. We identify this mismatch for prompt continuation and propose PCD (Prefix-Conditioned Diffusion), a pretraining objective that combines AR prefix supervision with no-shift suffix denoising. At the training-objective level, PCD changes the attention mask, corruption mask, and label construction in continued pretraining; it does not require an autoregressive decoder, verifier, or new inference mode. By supervising the clean-prefix side autoregressively and applying diffusion only to the unknown continuation, PCD makes the local training interface resemble how block-diffusion models are queried at evaluation time. We further separate intra-sample prefix conditioning from inter-sample objective mixing, allowing us to identify the local alignment signal separately from the optional batch-level mixing knob. Across LLaDA2-Mini and Qwen-1.7B backbones, PCD consistently improves over same-family native dLLM stable baselines, reaching a 4.2% relative gain on the main LLaDA2-Mini six-benchmark average (+2.56 points) and a 14.2% relative gain in the primary Qwen mechanism comparison (+4.86 points). These results suggest that aligning the pretraining context distribution with prompt-conditioned generation can recover a measurable part of the dLLM continuation gap without changing inference.