Reasoning with Latent Tokens in Diffusion Language Models

TL;DR

The paper investigates why diffusion language models excel at reasoning and global coherence by identifying latent tokens—predicted but not decoded auxiliary tokens that participate in predictions. It introduces latent token modulation to trade inference speed against sample quality and demonstrates that increasing latent tokens yields stronger performance on reasoning tasks and lower perplexity in language modeling. Extending the approach to autoregressive models via multi-token prediction (AR-MTP) yields substantial gains, narrowing or surpassing diffusion-model advantages on several constraint-satisfaction benchmarks. The work proposes latent tokens as a general mechanism for enhanced lookahead and global coherence across diffusion and autoregressive paradigms, with practical implications for adjustable inference efficiency and reasoning capability.

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 they require more computation at inference time. We trace this trade-off to a key mechanism: diffusion models are trained to jointly predict a distribution over all unknown tokens, 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 distribution of undecoded tokens. We interpret these as latent tokens and introduce a method for modulating their number, demonstrating empirically that this enables 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, while arising naturally in diffusion, represent a general mechanism for improving performance on tasks requiring global coherence or lookahead.

Figures (8)

• Figure 1: Latent tokens--jointly predicted at inference time but not decoded--improve reasoning performance in both diffusion and autoregressive models. On Sudoku, accuracy increases with the number of latent tokens across both paradigms. (Blue) A semi-causal diffusion model (SCDM) with controllable latent computation interpolates between a fast, independent-prediction model (SIDM) and an accurate, joint-prediction model (MDM). (Orange) Standard autoregressive models (AR) struggle on constraint satisfaction tasks; equipping them with latent tokens via multi-token prediction (AR-MTP) yields dramatic improvements while remaining in the autoregressive paradigm.
• Figure 2:
• Figure 3:
• Figure 4: Performance of a SCDM on four reasoning tasks as a function of latent token count. Under most settings, performance improves as more latent tokens are used for prediction. Varying the number of latent tokens has the effect of interpolating between SIDM ($n=0$), a fast model using no latent tokens, and MDM ($n=L$), a stronger model using all available latent tokens.
• Figure 5: Generative perplexity of an SCDM-like model from sahoo2025esotericlanguagemodels as a function of latent token count. Increasing the number of latent tokens provides a tunable trade-off between sampling speed and sample quality.