Conditional Inference Mismatch in Structured Diffusion Language Models
Abstract
Structured diffusion language models are trained by denoising partially observed sequences, often with access to both past and future context. Many downstream uses, however, require prefix-conditioned inference: predict the next answer from a prompt while future tokens are latent. We study this train/test conditional mismatch as a probabilistic inference problem. The core observation is an entropy gap: future-conditioned denoising conditionals are lower-uncertainty than the causal prefix-conditioned marginal required at test time, and recovering the latter requires marginalizing over exponentially many latent futures. Iterative denoising is an approximate inference procedure for this marginalization, but can amplify uncertainty when it conditions on model generated futures. We support this with one-step conditional likelihood diagnostics, full denoising traces, and controlled synthetic tasks. Finally, we introduce a simple structured corruption prior for masked diffusion LMs: mask later positions more often, so training more frequently denoises suffix tokens from cleaner prefixes while retaining the same tractable tokenwise diffusion-ELBO form for the modified process. The intervention partially reduces the uncertainty gap and improves context-conditioned prediction, suggesting that objective structure is an important inductive bias for prefix-conditioned inference.