Reasoning with Latent 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 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.

• The paper demonstrates that integrating latent token reasoning in diffusion language models improves sequence-level reasoning and planning.
• The proposed Masked Diffusion Model uses bidirectional and permutation-based approaches to jointly predict observed and latent tokens, outperforming standard baselines.
• Enhanced sample quality and interpretability in tasks like compositional generalization suggest significant practical applications for controlled and structured language generation.

Reasoning with Latent Tokens in Diffusion LLMs

Introduction

The paper "Reasoning with Latent Tokens in Diffusion LLMs" (2602.03769) investigates the integration of latent variable modeling within the framework of diffusion-based LLMs (DLMs). The primary motivation is to improve LLMs' reasoning and planning abilities by leveraging a latent sequence of tokens, promoting more globally coherent generation and enhanced context integration. The authors explore masked and bidirectional diffusion mechanisms and propose a framework for reasoning with partially observed and latent variables in the context of discrete sequence modeling.

Model Architecture and Methodology

The paper introduces Masked Diffusion Models (MDMs) for language, drawing parallels to score-based diffusion in continuous domains but adapted for the discrete token space of language. The model operates by gradually corrupting input text via masking, and then learning to reconstruct masked tokens through an iterative denoising process. The key innovation is the modeling and inference over latent tokens, which are not directly supervised or observed during the training process.

Two primary variants are discussed:

• MDM (Bidirectional): In this variant, the model predicts a set of latent tokens jointly alongside the observed sequence. The bidirectional architecture facilitates the use of both left and right context during token prediction, allowing for more flexible and semantically rich representations. Latent positions are predicted jointly, facilitating global planning and reasoning.
• SIDM/SCDM Baselines: These variants restrict forward passes and prediction to observed tokens, omitting latents from the forward computation. This contrasts with MDM's joint modeling of both observed and latent tokens, effectively disentangling the roles of latent variables in planning and reasoning.

The framework leverages permutation-based factorization, with $\pi$ denoting an ordering over token positions. The model can flexibly condition on arbitrary sets of context, facilitating both autoregressive and non-autoregressive generation modes.

Empirical Results

The paper presents extensive ablation studies and quantitative comparisons between MDMs and standard baselines such as masked LLMs and prior diffusion-based transformers. Strong numerical results are highlighted in tabled summaries throughout the paper, demonstrating:

• Improvement in sequence-level reasoning metrics, with MDMs outperforming autoregressive and partially masked models on tasks requiring long-range coherence and multi-step reasoning.
• Enhanced sample quality as measured by perplexity and downstream task accuracy, validating the hypothesis that joint latent reasoning confers benefits for both generation and comprehension.
• Models incorporating latent token reasoning show superior performance on compositional generalization benchmarks, reflecting improved modeling of underlying structure in language data.

Theoretical Implications

The paper's approach provides a principled pathway towards integrating latent variables and explicit planning in transformer-based LMs. By allowing the model to reason with unobserved latent tokens, the architecture more closely aligns with probabilistic graphical modeling (PGM) traditions while retaining the scalability and generality of modern transformers. This enables a form of amortized, structured reasoning that is more resistant to exposure bias and local minima common in greedy or left-to-right decoding schemes.

The work challenges the traditional autoregressive and masked paradigms by positing joint global reasoning as a necessary condition for robust, interpretable, and sample-efficient language modeling. This has implications for applications demanding controlled generation, text planning, or discrete structural prediction.

Practical Implications and Future Directions

The practical upshot of latent reasoning in DLMs is multifold:

• Improved Planning and Control: The ability to assign and reason over latent tokens enables explicit control over generation length, structure, and sub-task allocation.
• Enhanced Interpretability: Latent tokens provide a structural decomposition of the generation process, potentially making LLM decisions more interpretable and debuggable.
• Extensibility to Multimodal and Compositional Tasks: The presented approach is well-suited to domains such as program synthesis, chain-of-thought reasoning, and multimodal integration where explicit latent structure is central.

Future research directions could include hierarchical extension of the latent structure, integration with external world models or memory modules, and examination of latent diffusion in cross-lingual or domain adaptation scenarios. There is also potential for joint training with reinforcement learning objectives to further improve planning and reasoning fidelity.

Conclusion

This paper presents a coherent and technically rigorous approach to infusing latent variable reasoning within diffusion-based LLMs, achieving strong empirical gains on complex reasoning tasks. By bridging advances in diffusion modeling with structured latent variable methods, it sets a foundation for further advances in generative language modeling, controllability, and interpretability. The framework opens up new lines of research in hybrid symbolic-neural reasoning architectures and structured generative modeling for natural language.