Reducing Pretraining-Generation Mismatch in Diffusion Language Models

The paper identifies a pretraining-generation mismatch in diffusion language models (dLLMs). Autoregressive (AR) 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. The paper states: "During pretraining, a diffusion language model may see masked tokens distributed across an entire sequence rather than a clean prompt followed by an unknown continuation. During generation, especially for prompt continuation and benchmark evaluation, the model often receives a clean prefix and must denoise the unknown continuation without corrupting the prompt."

The authors 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.

The key design: "The left part of the sequence is trained with an AR loss, while the future suffix is corrupted and reconstructed with a diffusion objective. This gives the model the same kind of reliable left context that AR models exploit, while retaining denoising for the unknown continuation."

Preliminaries: Standard masked diffusion training samples a corruption pattern and trains a model to recover masked tokens from partially observed context. The native diffusion loss is: L diff = −E x,m Σ i:m i=1 log p θ(x i x̃, m).

PCD objective: Each sequence is partitioned into an autoregressive prefix and a denoising suffix. For a sampled prefix boundary k, the first k tokens remain clean (x̃ 1:k = x 1:k). The remaining suffix is corrupted according to a diffusion schedule. The model optimizes a weighted sum of prefix AR and suffix MDM losses:

L pcd = −E x,k,m [(λ ar/N ar) Σ i=2 k log p θ(x i xk:m i=1 log p θ(x i x 1:k, x̃ k+1:T, m)]

Here N ar = max(k−1, 1) and N m = max(Σ i>k m i, 1) normalize the two components, and (λ ar, λ mdm) control the AR-to-MDM loss ratio. The paper uses λ ar = λ mdm = 1 by default.

Intra-sample mixing: AR prefix prediction and MDM suffix denoising coexist inside the same sequence. For each training example, a prefix ratio ρ ∼ q(ρ) is sampled, and the prefix boundary is set to k = min(T−1, max(1, ⌊ρT⌋)). The effective mask is m i pcd = s i m i, so all prefix tokens stay visible for AR prediction and only suffix tokens can be selected as MDM targets.

Inter-sample mixing: Mixes objective types across examples in the same minibatch. For each sample, an objective indicator z ∼ Bernoulli(p pcd) is drawn. If z=1, the example uses intra-sample PCD; otherwise, it uses the native dLLM stable objective. The batch objective mixes the two losses: L = E z[zL intra + (1−z)L diff].

Variational view: The paper formalizes the context-distribution gap. Let C be the clean prefix, Y the suffix, C̃ and Ỹ corrupted versions. The information gap is:

V nat* − V pcd* = H(Y C̃, Ỹ, M) − H(Y C, Ỹ, M) = I(Y; C C̃, Ỹ, M) ≥ 0

This shows PCD removes a corrupted-prefix context-shift term from the suffix-denoising objective.

The paper evaluates checkpoints trained from LLaDA2-Mini and Qwen-1.7B backbones across reasoning, coding, math, and knowledge benchmarks. The most important comparison is against a same-family native dLLM stable baseline trained and evaluated with the same infrastructure. The LLaDA2-Mini evaluation uses a 12-benchmark suite including NQ, TriviaQA, MMLU-Pro, OmniMath, MATH, GSM8K, MBPP, HumanEval-FIM, LCBench, BBH, KorBench, and AutoLogi.

LLaDA2-Mini controlled comparison (400B stable budget):

• Native dLLM stable: Avg.6 = 60.30

• PCD stable: Avg.6 = 62.86 (+2.56 points, 4.2% relative gain)

Matched +50B decay continuation:

• Native dLLM decay: Avg.6 = 62.87

• PCD decay: Avg.6 = 64.54 (+1.67 Avg.6 above native decay)

The paper states: At the matched 400B stable budget, PCD improves the native dLLM stable baseline by +2.56 Avg.6 points. Figure 4 shows PCD improves over the same-family native baseline by +0.99 Avg.6 at 100B, +1.33 at 200B, +2.60 at 300B, and +2.56 at 400B.

Qwen-1.7B mechanism comparison (50B continued pretraining):

• Intra-sample-only row gives the strongest controlled Avg.6 gain: +4.86 points (14.2% relative gain)

• Inter-sample-only mixing at p pcd = 0.50: +4.70

• Mixed PCD: +4.66

The paper concludes: "This pattern supports the central mechanism: same-sample clean-prefix suffix denoising is the primary alignment signal, while inter-sample mixing should be viewed as a conservative regularizer or distribution-control knob rather than a replacement for the intra-sample structure."

Chat-SFT transfer: The PCD-derived SFT checkpoint reaches 67.39 Avg. at epoch 5, compared with 66.53 for the best native-SFT checkpoint. The gain is already visible at epoch 3 (67.13 vs. 65.12).

Benchmark prompt-prefix statistics: For the five public main-suite benchmarks, the aggregate per-example prefix ratio is 96.8%; the median rendered prompt contains 754 tokens, while the median gold completion contains only 12 tokens. Under a native random-corruption objective with an illustrative corruption rate r=0.3, the median prompt would contain 226 expected corrupted prefix tokens, and the probability that the entire prefix survives clean is effectively zero at these lengths.

Domain averages: At the matched 400B budget, PCD improves the native baseline on all 12 benchmarks, yielding a +2.48-point gain in their overall average. The effect is broad across QA, general reasoning, math, and code domains.

Prefix ratio sweep: The Qwen sweep peaks at the intermediate [0, 0.75] range rather than at the full range. The paper notes: Very small prefix ratios may fail to provide enough clean-prefix context. Very large prefix ratios may reduce the amount of diffusion denoising signal.

WSD interaction: PCD is designed as a stable-phase intervention. The paper states: "PCD is designed to fit inside a WSD-style AR-to-diffusion conversion pipeline rather than replace it... The stable phase is the most natural place to intervene because it consumes the largest continued-pretraining budget and establishes the dominant diffusion objective."

The paper concludes: "These results suggest that diffusion language models can retain their denoising formulation while recovering part of the clean-prefix behavior that makes AR models effective for continuation. More broadly, the pretraining context distribution should be treated as a first-class design choice for block-diffusion language models, especially when the target use case is prompt-conditioned generation."

The paper acknowledges: "Most large-scale rows are single training runs because 100B–400B continued-pretraining budgets are expensive; future work should test seed sensitivity at smaller scale and confirm the strongest Qwen ablation settings at larger scale. Execution-based coding benchmarks and logic benchmarks can also be sensitive to timeout settings, prompt format, and answer extraction."

Improvements to AI systems:

1. Add a clean-prefix conditioning mechanism to diffusion language models. During pretraining, partition each sequence into an autoregressive prefix (trained with next-token prediction) and a denoising suffix (trained with masked diffusion). This ensures the model learns to preserve clean prompt context while reconstructing unknown continuations, directly addressing the pretraining-generation mismatch.

2. Implement intra-sample mixing of AR and diffusion objectives. For each training example, sample a prefix boundary ratio (e.g., from [0, 0.75]) and apply AR loss to the first k tokens and diffusion loss to the corrupted suffix. This gives the model reliable left-context supervision without sacrificing parallel denoising capability.

3. Add inter-sample mixing as a regularizer. Mix PCD-style examples with native diffusion examples in the same minibatch (e.g., 50% each) to maintain distributional robustness and prevent overfitting to the clean-prefix structure, while keeping the intra-sample alignment as the primary signal.

4. Modify the attention mask and corruption mask jointly. Ensure prefix tokens are always visible (never masked) during suffix denoising, and only suffix tokens are eligible for corruption targets. This prevents the model from seeing corrupted prompt tokens during training, which would otherwise weaken prompt-conditioned generation.

5. Use a variational information-gap objective. Optimize the model to minimize the conditional entropy gap between native and PCD objectives, effectively removing the corrupted-prefix context-shift term from the suffix-denoising loss.

What the improved AI system can do:

• Generate text from a clean prompt with higher fidelity. The model no longer corrupts the prompt during training, so at inference it can reliably condition on the full user-provided prefix without degradation.

• Achieve +2.5 to +4.9 average benchmark gains (across reasoning, math, coding, and knowledge tasks) compared to native diffusion baselines at the same training budget, with consistent improvements across all 12 evaluated benchmarks.

• Maintain parallel decoding speed (unlike autoregressive models) while recovering AR-like prompt-conditioning behavior, making it suitable for latency-sensitive applications like code completion and interactive assistants.

• Transfer better to supervised fine-tuning. The PCD-pretrained model shows faster SFT convergence and higher final performance (e.g., 67.39 vs. 66.53 average score) compared to native diffusion pretraining.

• Handle long prompts with short completions effectively. The system is robust even when the prompt is 96.8% of the total sequence length, a scenario where native diffusion models fail because random corruption would almost certainly damage the prompt.

• Serve as a drop-in replacement for existing diffusion language models in WSD-style training pipelines, requiring only changes to the training objective (attention mask, corruption mask, and label construction) without architectural modifications or new inference modes.