Rethinking Token Prediction: Tree-Structured Diffusion Language Model

• URL: http://arxiv.org/abs/2604.03537v1
• Date: Sat, 04 Apr 2026 01:19:12 GMT
• Title: Rethinking Token Prediction: Tree-Structured Diffusion Language Model
• Abstract summary: We build a tree-structured diffusion language model with latent states corresponding to a token's ancestor nodes.<n>Our method reduces peak GPU memory usage by half while matching the perplexity performance of state-of-the-art discrete diffusion language models.
• Score: 17.568195527860155
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Learn from your own latents and not from tokens: A sample-complexity theory [53.5821824211418]
  We study networks trained to predict their emphown latent representations of related views or masked regions.<n>We show that latent prediction achieves this with a number of samples emphconstant in $L$, up to logarithmic factors.<n>This suggests that explicit stacking such as H-JEPA is largely redundant.
  arXiv  Detail & Related papers  (2026-05-26T22:16:42Z)
• The Efficiency Gap in Byte Modeling [101.18202046105718]
  Two alternative paradigms have challenged the use of subword tokenization and autoregressive ordering.<n>We show that removing these structural priors incurs a significant computational cost.<n>Our findings suggest that future modality-agnostic designs must incorporate alternative structural biases to maintain viable scaling trajectories.
  arXiv  Detail & Related papers  (2026-05-13T03:03:30Z)
• Memory-based Language Models: An Efficient, Explainable, and Eco-friendly Approach to Large Language Modeling [0.4411777886421431]
  We present memory-based language modeling as an efficient, eco-friendly alternative to deep neural network-based language modeling.<n>It offers log-linearly scalable next-token prediction performance and strong memorization capabilities.
  arXiv  Detail & Related papers  (2025-10-25T14:34:18Z)
• Context-level Language Modeling by Learning Predictive Context Embeddings [79.00607069677393]
  We introduce textbfContextLM, a framework that augments standard pretraining with an inherent textbfnext-context prediction objective.<n>This mechanism trains the model to learn predictive representations of multi-token contexts, leveraging error signals derived from future token chunks.<n>Experiments on the GPT2 and Pythia model families, scaled up to $1.5$B parameters, show that ContextLM delivers consistent improvements in both perplexity and downstream task performance.
  arXiv  Detail & Related papers  (2025-10-23T07:09:45Z)
• Latent Thought Models with Variational Bayes Inference-Time Computation [52.63299874322121]
  Latent Thought Models (LTMs) incorporate explicit latent thought vectors that follow an explicit prior model in latent space.<n>LTMs demonstrate superior sample and parameter efficiency compared to autoregressive models and discrete diffusion models.
  arXiv  Detail & Related papers  (2025-02-03T17:50:34Z)
• Causal Estimation of Memorisation Profiles [58.20086589761273]
  Understanding memorisation in language models has practical and societal implications. Memorisation is the causal effect of training with an instance on the model's ability to predict that instance. This paper proposes a new, principled, and efficient method to estimate memorisation based on the difference-in-differences design from econometrics.
  arXiv  Detail & Related papers  (2024-06-06T17:59:09Z)
• Towards A Unified View of Sparse Feed-Forward Network in Pretraining Large Language Model [58.9100867327305]
  Large and sparse feed-forward layers (S-FFN) have proven effective in scaling up Transformers model size for textitpretraining large language models. We analyzed two major design choices of S-FFN: the memory block (a.k.a. expert) size and the memory block selection method. We found a simpler selection method -- textbftextttAvg-K that selects blocks through their mean aggregated hidden states, achieving lower perplexity in language model pretraining.
  arXiv  Detail & Related papers  (2023-05-23T12:28:37Z)
• Compressing Sentence Representation with maximum Coding Rate Reduction [0.0]
  In most natural language inference problems, sentence representation is needed for semantic retrieval tasks. Due to space and time hardware limitations, there is a need to attain comparable results when using the smaller model. We demonstrate that the new language model with reduced complexity and sentence embedding size can achieve comparable results on semantic retrieval benchmarks.
  arXiv  Detail & Related papers  (2023-04-25T09:23:43Z)
• A Sparsity-promoting Dictionary Model for Variational Autoencoders [16.61511959679188]
  Structuring the latent space in deep generative models is important to yield more expressive models and interpretable representations. We propose a simple yet effective methodology to structure the latent space via a sparsity-promoting dictionary model.
  arXiv  Detail & Related papers  (2022-03-29T17:13:11Z)
• Low-Rank Constraints for Fast Inference in Structured Models [110.38427965904266]
  This work demonstrates a simple approach to reduce the computational and memory complexity of a large class of structured models. Experiments with neural parameterized structured models for language modeling, polyphonic music modeling, unsupervised grammar induction, and video modeling show that our approach matches the accuracy of standard models at large state spaces.
  arXiv  Detail & Related papers  (2022-01-08T00:47:50Z)
• Adaptive Semiparametric Language Models [17.53604394786977]
  We present a language model that combines a large parametric neural network (i.e., a transformer) with a non-parametric episodic memory component. Experiments on word-based and character-based language modeling datasets demonstrate the efficacy of our proposed method.
  arXiv  Detail & Related papers  (2021-02-04T11:47:03Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.