# Rethinking Token Prediction: Tree-Structured Diffusion Language Model

arXiv: https://arxiv.org/abs/2604.03537 · HTML: https://mlsanity.com/paper/2604.03537/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion reasoning and test-time scaling, one of 17 directions in LLM reasoning.
- Builds on 23 of its 29 references that are in the map.
- First posted on arXiv 4 Apr 2026.

## Placed in

- LLM reasoning / [Diffusion reasoning and test-time scaling](https://mlsanity.com/lr/diffusion-reasoning-and-test-time-scaling/)
- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 23 of 23)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [A Continuous Time Framework for Discrete Denoising Models](https://mlsanity.com/paper/2205.14987/) (2022)
- [A Survey on Model Compression for Large Language Models](https://mlsanity.com/paper/2308.07633/) (2023)
- [Between words and characters: A Brief History of Open-Vocabulary Modeling and Tokenization in NLP](https://mlsanity.com/paper/2112.10508/) (2021)
- [Analyzing Transformers in Embedding Space](https://mlsanity.com/paper/2209.02535/) (2022)
- [Generalized Interpolating Discrete Diffusion](https://mlsanity.com/paper/2503.04482/) (2025)
- [SynCode: LLM Generation with Grammar Augmentation](https://mlsanity.com/paper/2403.01632/) (2024)
- [Efficient softmax approximation for GPUs](https://mlsanity.com/paper/1609.04309/) (2016)
- [ZipLM: Inference-Aware Structured Pruning of Language Models](https://mlsanity.com/paper/2302.04089/) (2023)
- [Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking](https://mlsanity.com/paper/2505.18495/) (2025)
- [CompAct: Compressed Activations for Memory-Efficient LLM Training](https://mlsanity.com/paper/2410.15352/) (2024)
- [SentenceKV: Efficient LLM Inference via Sentence-Level Semantic KV Caching](https://mlsanity.com/paper/2504.00970/) (2025)
- [Evaluating and Explaining Large Language Models for Code Using Syntactic Structures](https://mlsanity.com/paper/2308.03873/) (2023)
- [Next Semantic Scale Prediction via Hierarchical Diffusion Language Models](https://mlsanity.com/paper/2510.08632/) (2025)

## Built on by (most cited first, 1 of 1)

- [Denoising Hierarchical Representations: Joint Continuous Diffusion for Language Modeling](https://mlsanity.com/paper/2610.08738/) (2026)

## Abstract (arXiv)

Discrete diffusion language models have emerged as a competitive alternative to auto-regressive language models, but training them efficiently under limited parameter and memory budgets remains challenging. Modern architectures are predominantly based on a full-vocabulary token prediction layer, which accounts for a substantial fraction of model parameters (e.g., more than 20% in small scale DiT-style designs) and often dominates peak GPU memory usage. This leads to inefficient use of both parameters and memory under constrained training resources. To address this issue, we revisit the necessity of explicit full-vocabulary prediction, and instead exploit the inherent structure among tokens to build a tree-structured diffusion language model. Specifically, we model the diffusion process with intermediate latent states corresponding to a token's ancestor nodes in a pre-constructed vocabulary tree. This tree-structured factorization exponentially reduces the classification dimensionality, makes the prediction head negligible in size, and enables reallocation of parameters to deepen the attention blocks. Empirically, under the same parameter budget, our method reduces peak GPU memory usage by half while matching the perplexity performance of state-of-the-art discrete diffusion language models.