PCD: Fixing the Pretraining-Generation Mismatch in Diffusion Language Models

Forum topic · ✨步子哥 · 2026-08-11

Summary

This article explains a paper titled "Reducing Pretraining-Generation Mismatch in Diffusion Language Models" by Xiaocheng Lu, Huabin Liu, Song Guo, and Jianguo Li (arXiv: 2608.09424). It highlights that standard diffusion language models (dLLMs) such as LLaDA and MDLM randomly mask tokens across the entire sequence during pretraining, but at inference time only the continuation is masked while the prompt stays clean. This Pretraining-Generation Mismatch means dLLMs never learn the exact task they will face at inference. The proposed fix, PCD (Prefix-Conditioned Diffusion), aligns the training interface with the inference interface by keeping the prompt clean via causal attention, masking only the continuation, and splitting the loss into an autoregressive term on the prompt and a denoising term on the continuation. PCD is a training-only modification and requires no changes to inference. Experiments on LLaDA2-Mini and a Qwen-1.7B backbone report average gains of about 4.2% and 14.2% respectively across six benchmarks, supporting interface alignment as a systematic improvement direction for non-autoregressive language models.

Paper

Title: Reducing Pretraining-Generation Mismatch in Diffusion Language Models
Authors: Xiaocheng Lu, Huabin Liu, Song Guo, Jianguo Li
arXiv: https://arxiv.org/abs/2608.09424

Key points

1. An overlooked training-inference mismatch

Autoregressive language models benefit from an implicit alignment: training and inference both use a clean left context to predict the next token. Diffusion language models (dLLMs) such as LLaDA, MDLM, and Plaid break this assumption. During standard dLLM pretraining, masking is globally random, so prompt tokens and continuation tokens are masked symmetrically. At inference, however, the prompt is fully clean and only the continuation is masked and denoised.

| Aspect | Pretraining | Inference | |---|---|---| | Prompt state | Randomly masked | Fully clean | | Continuation state | Randomly masked | Fully masked, then iteratively denoised | | Prompt-continuation relation | Symmetric | Asymmetric (prompt is a condition) |

This mismatch is termed Pretraining-Generation Mismatch: dLLMs are trained on a broader, noisier task ("recover randomly masked tokens in a randomly masked context") rather than the exact inference task ("recover a continuation given a clean prompt").

2. PCD: Prefix-Conditioned Diffusion

PCD aligns the training interface with the inference interface through three coordinated changes:

1. Attention mask. Prompt tokens use causal attention (each prompt token attends only to earlier prompt tokens), while continuation tokens use bidirectional attention over the prompt and other continuation tokens. Prompt representations become clean, ordered, and uncontaminated. 2. Corruption mask. Only the continuation is masked; the prompt stays clean throughout training. 3. Label construction. The loss splits into two parts:

Autoregressive loss on the prompt.
Diffusion (denoising) loss on the continuation.

3. No new components at inference

PCD is a pure training-objective modification. Inference uses the same dLLM decoding procedure as before—no draft model, no auxiliary heads, no decoding-strategy changes. This contrasts with speculative decoding or multi-head decoding methods that change inference.

The authors separate intra-sample prefix conditioning (the primary contribution) from inter-sample objective mixing (an optional knob). The signal comes from sample-level interface alignment, not batch-level mixing.

4. Experimental results

LLaDA2-Mini (continued pretraining): average improvement of +4.2% relative (+2.56 absolute) across six benchmarks.
Qwen-1.5B (training from scratch): average improvement of +14.2% relative (+4.86 absolute).

The consistent direction across two backbones suggests PCD is a general dLLM training-objective improvement rather than a model-specific trick.

5. Why this direction matters

Diffusion language models are a leading non-autoregressive paradigm because they offer parallel decoding (potentially order-of-magnitude faster inference), global sequence planning, and natural infilling/edit support. Yet dLLMs have long lagged autoregressive baselines on standard benchmarks. PCD attributes part of that gap to training-inference interface misalignment and recovers 4-14% of performance through a training-only fix.

6. Honest evaluation

Modest absolute gains. ~0.4-0.8 absolute points per benchmark.
Small-scale validation. Only LLaDA2-Mini and a Qwen-1.5B variant; behavior at 7B+ scale is unverified.
No public code release at the time of writing, which limits reproducibility and follow-up work.
Mixed training regimes. LLaDA2-Mini uses continued pretraining; Qwen-1.5B is trained from scratch. Gains appear in both settings.

7. Core takeaway

> Diffusion language models have never been trained on the exact task they perform at inference. PCD aligns the training interface to the inference interface, recovering 4-14% of performance without changing inference.

The broader lesson for non-autoregressive modeling: explicitly verify that training and inference tasks coincide—interface alignment is not free once you leave the autoregressive paradigm.

Source

Paper link: https://arxiv.org/abs/2608.09424