Reducing Pretraining-Generation Mismatch in Diffusion Language Models

2026-08-10

AI Summary

This work addresses the training-inference mismatch in diffusion language models (dLLMs), which arises from randomly corrupting both prompts and continuations during pretraining, leading to inconsistent conditioning. To resolve this, the authors propose Prefix-Conditioned Diffusion (PCD), a novel pretraining objective that applies autoregressive supervision over the known prefix while performing unbiased denoising exclusively on the unknown suffix. By employing tailored attention masks, corruption masks, and label construction strategies, PCD achieves local alignment between training and inference objectives—closely mirroring real-world prompt-conditioned generation scenarios—without modifying the decoder architecture or inference pipeline. Experiments demonstrate that PCD significantly outperforms native dLLM baselines on LLaDA2-Mini and Qwen-1.7B, yielding an average relative improvement of 4.2% (+2.56 points) across six benchmarks and a substantial 14.2% (+4.86 points) gain in key ablation comparisons.

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.

Problem

• diffusion language models
• pretraining-generation mismatch
• prompt-conditioned generation
• prefix corruption
• continuation gap

Innovation

• diffusion language models
• prefix-conditioned generation
• pretraining-generation alignment
• PCD
• denoising objective