cs.CLApr 4, 2026

Rethinking Token Prediction: Tree-Structured Diffusion Language Model

Authors: Zihao Wu, Haoming Yang, Juncheng Dong, Vahid Tarokh

Organizations: Department of Electrical & Computer Engineering, Duke University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.

Figures & tables

Figure 1: Generation process (left): Our algorithm features in-level child prediction rather than standard token prediction, delivering substantial efficiency gains while achieving improved performance. Token tree (right): Child prediction is enabled by a principled token tree, in which each state evolves strictly between adjacent levels.
ModelTrain. toks.Val. PPL ( ↓\downarrow )Gen. PPL ( ↓\downarrow )
GPT2 †\daggerunk.23.40–
Llama110M (retrain.) †\dagger262B16.11–
SEDD ( Lou et al., 2023 ) †\dagger262B≤\leq 24.10–
MDLM-small ( Sahoo et al., 2024 ) ∗*131B≤\leq 27.39163.7
GIDD+-small ( von Rütte et al., 2025 ) ∗*131B≤\leq 25.82170.2
HDLM-small ( Zhou et al., 2025 ) ∗*131B≤23.25\leq\mathbf{23.25}148.0
Table 1: Validation and generative perplexity on OWT. Bold and underline denote best and second best. Adopted from Zhou et al. (2025) ∗* ; adopted from corresponding previous work †\dagger .
Figure 2: Memory effiency of TDLM. (Left) Throughput and peak memory comparison (denoted as size of the circle) across diffusion language models at different scales. (Right) Validation negative ELBO and average Flops/s of small scale models during training using full capacity of four 24 GB GPUs.
Figure 3: Left: Validation nagative ELBO of various branching factor KK and cluster size ratio, plotted over training iterations. Solid lines indicate varying KK and fixed cluster size ratio, while dashed lines indicate fixed KK but varying cluster size. Middle: each bar indicates the cumulative validation negative ELBO obtained from a different tree construction, while each stack indicates the contribution of a particular tree height. (top stack indicates root level) Right: raw validation negative ELBO of each height level under different tree constructions.
Figure 4: Smoothed validation negative ELBO of different weight schedules. The weight schedules are plotted in the subfigures (Left: exponential; Right: linear).
ModelInference Steps (Top/Bottom Level)Gen. PPL ( ↓\downarrow )
TDLM-base ( H=2H=2 )64 / 448142.1
128 / 384140.3
192 / 320140.7
256 / 256138.0
320 / 192141.7
384 / 128149.2
Table 2: Generation perplexity with different allocations of the total 512 inference steps using TDLM-base model.
Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Smoothed validation ELBO for joint modeling with different branching factor KK and neighborhood length LL .
Figure 6: Level-wise weighting schedules used in our ablations.
Qualitative Examples: Generated Text With 512 Sampling Steps
Main difference is that much like the economic ladder movement sparked in Ferguson, that small town, citizens concerned about police misconduct, from labor activists to hardcore progressive institutions, all gathered. The message of income linkage was all at the center, the widespread Streeck discontent, communities that felt much better off than individuals once were. Their voice was more relevant than activists giving rise to the message, but the main driver was inherent inequality, and that is a large part of what was the gap between poverty and income. You likely didn’t see a bold income linkage message from the many pro-dissinity Democrats who urge economic polarization and there just weren’t many answers calling for policy that’s harmful for the needy, either. A message: Income linkage led by affirmative Democratic leaders might not resonate with low-income individuals who are housing survivors in their lives and have deeper ties to the Left where their previous-class status still allows. Further to these aspects, I argued that if the income scale were off-course for theoretical days it would have been left to the working and middle class– (these elements that even the with the most political willingness and ability feel marginalized by left-progressive leadership long deserve at the bottom of the income scale). The message might have expanded and supersized by taking some finer steps that made sense in fighting intolerance and poverty. The message would have seen a left step up if centrists opted to blame the plight of some low-income people on others, but even then the progressive movement likely wouldn’t have topped paychecks among all the lowest-income earners. We would have made policies anti-poverty if we appealed to only the least, those groups with more at stake in engaging in poverty reduction discourse. We needed to reject these political science to recognize that certain anecdotes were more uncomfortable in the middle of (status.) Burrow specifically foresaw that the Obama administration’s opportunities to act were controlled by the anti-poverty speculation we are now consistently dealing with. (Especially Washington policymakers who this moment have been willing to overlook a cohesive progressive agenda.) We dispelled the decision to make legal requirements for the poor more economically viable, welfare ambassador time was spent personally with meaning more cash for community organizations and the race to a grass of good is surely over. Burrow is right: we were policies that hold the stronger voice for the affirmative in the larger segment that fails to involve message reaching the poor or status. Income disparities discourse were not extraneous to ending
Gender advocacy on workers relationship prime in life should be: The daunting struggle of life. Women with single mothers, women who are single professionals, and, in particular, busy men most easily ought to prefer leading a wonderful life to, as we have written previously, women imposed by more extraordinary circumstances and accomplures. Yet, both men run more well than women, while gender attitudes on differences of adolescence break up entirely between the two and into one getting more firmly briefed, the other is often emotionally closer to men. Both seek a lot farther from the ideal of life for all, while both demographic groups still survives. As for Scientific opinion about adolescence—in short is “Live Long, Die Long.” In commonly the wisdom and insight of teen Vikings readers that Dr. Dana Reeer had compared to the scientific evidence, like childhood drug toxicity, or thyroid problems. or, even more expound by experts, research that finds that children and teens, even people in their late 30s — and not adults — may have been most prone to physical, emotional, violence, and hostile home environment. The signs of adolescent health and well-being set off a mutual struggle between those fears and attitudes of teens that belongs to both sexes, Reeer writes: For those who have research find to adults, after 30 or young adulthood, think of body mass index and life expectancy as such fifth percent degree, while teens themselves are in the 70s. The sciences, attitudes and mentations of adults, critical focus into childhood undo the accumulated patterns of girls and women now and over their lives and and place an increased value on women’s emotional identity and well-being. Regardless of expenditures and people’s incomes, women are more hands-on on their lives than men, and more women are able to resolve the baseline blurred in emotions nor behaviors, Dr. Reeer writes for the Disbalance, as well as to address the differences between men and women. But to add to that indignant attitude uttered by both genders toward people who are later jealous, more girls when single are able to avoid focus as prepared for these first impulses of adulthood, and more able to weave an emotion removed longer from ignoring the differences of youth and partner than dispatch on keeping sides of the root cause of neglect that momentarily affects so much of your baseline survivability. But you cannot only focus on emotions and critical emotion and filling comfort gaps as we have written in Getting started on Enthusiasm Approaches, Mistakes, Trans
Appendix
Table 3: Qualitative examples generated at 512 sampling steps.
Qualitative Examples: Generated Text With 1024 Sampling Steps
Greenfield, for his part, cited internet-based super PAC reports of second-hand contributions, judged by pundits, up to $22 million in the last election from super PACs. But he cited the Independent Spending Council considered an election spending association, in recognition that super PACs stand foremost as top-spending” groups on political activities, that registered firms should accept their spending and conducting organizations. Super PACs now have the Federal Election Commission and redistricting Democrats with the recent Supreme Court decision as their reins, agree. Republican leaders like the LP cheered the Citizens United ruling and observed that Obama’s federal campaign finance laws are comfortably toward being struck down, again. “I think of how the Republican presidential candidates see a positive step here in Citizens United,” said Republican National Committee Chairman David Walker Jr. R-Ohio, of clogging Obama’s individual mandate “wizards” toward the “enormancy and ignorance” of the Capitol Hill office of Senate Majority Leader Harry Reid Jr., D-California. But even he’s not even going to seek to gut the opposing precedent of Citizens United. In the protest of Democrat leaders stepping to the party after Long’s announcement included Capitol Hill’s Jack Reed (R., Virginia). A chief, conservative font included: Rep. Nancy Pelosi, Chairwoman of the U.S. Securities & Exchange Commission. Stern said there has been financial consensus in the House of Representatives to get him to recuse himself from benefiting from secret money used by super PAC groups for political activity and ban groups and members from voting against them. The coalition created to protest against Tuesday’s decision was organized between Long and Pai, stars and various top actors, including the Allison Wives magazine run, who via the LP. They endorsed the Mega Finance statement as “an inclusive, politically empowered multicultural coalition mixed with an intelligent record of interventions on the racial and sexual diversity of the political financing landscape.” The firm includes Mega Finance Chairman Mark Michael Ficloma; Richard Mache Perez, the chief economist of Prudential Advisors; Taylor Jin, VP Senior Chairman at Federal Bank of Chicago and Hallie Douglas (MO), Chairman of National Rentali Corp.; and Jacqueline Rosen, president and financial advisor
The health care system is also very problematic. Private insurance has received an ample chunk too little for it to open drug facilities in some parts of the market, and Private rated providers cannot at first take return, on subscriber money to open free rehab facilities. Fortunately, there is now part of the Affordable Care Act’s tax breaks for certificates in Pharmacology Management Education (POAC) and other health care necessities. Obamacare’s subsidies will directly support researchers at all federal levels for reduction in traffic deaths due to opioids. The subsidies will help testing of ”planning techniques” and drug-education programs for inmates seeking outpatient treatment. One possibility for the largest benefit of ministerial prisons specifically when considering drug punishments is the ability to assist inmates, who want to be helped, into treatment via co-program of forbearance to take some psychapeptic drugs. Programming based on a lesser degree of forbearance is also typical for most medical insurance companies to purchase some inexpensive prescription benefit. Better still, however, is changes in medical insurance laws. A much newer short-term insurance plan can render treatment effective in varies and drug abusers by removing expensive treatment for those patients. Often, additional long-term treatment is preceded by a medical spending of outreach premiums 36% or much more reduced, which lead to changes in daily legal killers that propelled opioid users and changed drug symptoms. Sometimes, it is appropriate to limit ministerial-lock policies by other means, such as if TTE+ or tertiary palliative, while purpose that do not necessarily require its expected derive. Based on federal offender registries stat is about 1.2 million Americans die from overdose each year. The interpretation system of drug abuse from those that have really no delusion is in addition to multiple consistent requirements: ward off the black-and-white idea of sufficient drug funding. However, the recent BBC documentary study showed that 80% of prison opioids have far less major issues than even a first year’s “white” virus in the commonly-ignored norm. It found that the rate of past reports of adverse effects seen for every 5,000 IHA members is 12%, and the 2015 estimated rate of urine inhaled hazardous is 3.9%. Those crucial findings which reveal Bad-ass Prison Revenge will likely lead to a very restrictive medical norm. The public suffers from a future in which 4 drug prison gang and 5 giro-club spread for most drug dealing and where new convictions, convicts not released, are quickly giving drug
Appendix
Table 4: Qualitative examples generated at 1024 sampling steps.
More in this topic

Diffusion Language Models

Diffusion language models generate and refine text through iterative denoising, enabling bidirectional context and parallel token updates.

See all 471 papers

Explore similar work

CardsList
Feb 11, 2026

Just on Time: Token-Level Early Stopping for Diffusion Language Models

Diffusion language models generate text through iterative refinement, a process that is often computationally inefficient because many tokens reach stability long before the final denoising step. We introduce a training-free, token-level early stopping approach that identifies convergence independently at each position. Our method leverages lightweight signals derived from the model's predictions and local context to dynamically determine when individual tokens can be finalized. This yields adaptive per-token freezing without task-specific fine-tuning, substantially reducing the total number of diffusion steps required. Across diverse benchmarks, spanning mathematical reasoning, general question answering, and scientific understanding, our approach achieves substantial efficiency gains while preserving generation quality.

Zakhar Kohut, Severyn Shykula, Mykola Vysotskyi +7

SoftServe Inc.

Diffusion Language Model InferenceDiffusion Language Model DecodingDiffusion Model AccelerationDiffusion Language Models
AskPDF
May 14, 2026

Factorization-Error-Free Discrete Diffusion Language Model via Speculative Decoding

Discrete diffusion language models improve generation efficiency through parallel token prediction, but standard X0X_0 prediction methods introduce factorization errors by approximating the clean token posterior with independent token-wise distributions. This paper proposes Factorization-Error-Free Discrete Diffusion Language Modeling (FeF-DLLM), which replaces independent clean-token prediction with an exact prefix-conditioned factorization of the clean posterior to better preserve token dependencies. To reduce the sequential cost introduced by prefix conditioning, FeF-DLLM further incorporates speculative decoding within diffusion denoising, accelerating inference while maintaining the parallel prediction and re-masking properties of DLLMs. Theoretically, we prove that FeF-DLLM generates from the true joint distribution and derive its expected acceleration ratio. Experiments on GSM8K, MATH, HumanEval, and MBPP demonstrate that our method improves accuracy by an average of 5.04 percentage points while achieving an average inference speedup of 3.86×3.86\times.

Xun Fang, Yunchen Li, Hang Yuan +1

East China Normal University · Beijing Zhongguancun Academy

Speculative Decoding for Diffusion Language ModelsDiffusion Language Model InferenceDiffusion Language Model DecodingDiffusion Language Models
AskPDF
Oct 1, 2026

Hierarchical Continuous Diffusion Language Models

Discrete diffusion language models offer a compelling alternative to autoregressive generation for tasks demanding bidirectional reasoning and global constraint satisfaction. Yet they share a structural bottleneck: when decoding in parallel, each token is sampled independently from its marginal, severing the statistical dependencies among the tokens decoded together. Continuous diffusion language models avoid this by denoising a shared continuous state, but their denoiser sees only that state, so nothing ties it to a valid token configuration until it is finally decoded. To address this, we propose Hierarchical Continuous Diffusion Language Models (HC-DLM), which couple discrete token generation with a continuous latent trajectory in a single, principled denoising process, whose training objective is derived from a variational bound on the token likelihood. In contrast to recent methods that attach continuous context to a self-contained discrete chain, HC-DLM makes the latent the only persistent generative state: tokens are read out from it at every step and feed back as a scaffold for the next latent update. On structured reasoning (Sudoku), mathematical planning (Countdown) and language modeling (LM1B), HC-DLM improves over discrete and continuous diffusion baselines at matched model size, in puzzle accuracy on Sudoku and Countdown and in generative perplexity on LM1B. Project page: https://hc-dlm.github.io/.

Hui Ren, Zihan Li, Chang Liu +2

University of Illinois Urbana-Champaign · Amazon.com, Inc.

Structured Reasoning in Language ModelsDiffusion Language Model DecodingDiffusion Language ModelsContinuous Diffusion Language Models
AskPDF