Reasoning with Undecoded Tokens in Diffusion Language Models
Abstract
Discrete diffusion models have recently become competitive with autoregressive models for language modeling, even outperforming them on reasoning tasks requiring planning and global coherence, but diffusion requires more computation at inference time. We trace this trade-off to a key mechanism: diffusion models are trained to jointly predict all unknown tokens simultaneously, including those that will not actually be decoded in the current step. Ablating this joint prediction yields faster inference but degrades performance, revealing that accurate prediction at the decoded position relies on joint reasoning about the undecoded tokens. We interpret these undecoded positions as latent tokens and introduce a method for modulating their number, achieving a smooth tradeoff between inference speed and sample quality. Furthermore, we demonstrate that latent tokens can be introduced into autoregressive models through an auxiliary multi-token prediction objective, yielding substantial improvements on the same reasoning tasks where they have traditionally struggled. Our results suggest that latent tokens--arising from jointly predicting multiple unknown positions--represent a general mechanism for improving performance on tasks requiring global coherence or lookahead.