Rethinking Token Prediction: Tree-Structured Diffusion Language Model
This addresses a bottleneck in training discrete diffusion language models efficiently under constrained resources, offering a domain-specific improvement for NLP researchers and practitioners.
The paper tackles the inefficiency of full-vocabulary token prediction in discrete diffusion language models, which consumes significant parameters and memory, by proposing a tree-structured diffusion model that reduces classification dimensionality and reallocates parameters to attention blocks. The result is a method that halves peak GPU memory usage while matching state-of-the-art perplexity performance under the same parameter budget.
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.