Tree-Structured Diffusion Language Model

• TDLM is a family of diffusion models that infuse hierarchical structure into generative and denoising processes for both natural language and code.
• It leverages AST-guided span masking for code and a balanced K-ary vocabulary tree for language, reducing computational cost while maintaining model strength.
• Empirical evaluations show TDLM achieving improved syntactic validity, parameter efficiency, and enhanced scalability compared to traditional full-vocabulary diffusion models.

The Tree-Structured Diffusion LLM (TDLM) is a family of discrete diffusion-based LLMs that introduces hierarchical structure into the generative and denoising processes for natural language and code. TDLM leverages tree-based representations either of code syntax (via abstract syntax trees) or of vocabularies (via hierarchical token clustering) to achieve improved syntactic fidelity, parameter efficiency, and computational scalability, while preserving or advancing upon the modeling power of traditional full-vocabulary diffusion LLMs.

1. Discrete Diffusion Language Modelling Foundations

Diffusion LLMs operate by framing text generation as an iterative denoising process, where a syntactically or semantically meaningful input $x_0$ is corrupted through a forward stochastic process to produce a noisy intermediate state $x_T$, and a learned parameterized model recovers the original $x_0$ from $x_T$ via a sequence of reverse steps. Standard approaches, such as Masked Diffusion LLM (MDLM) and Generalized Iterative Denoising Diffusion (GIDD), treat each token as a state in a flat vocabulary of size $|V|$, and at each step must compute a full softmax prediction over all possible tokens. For vocabularies of practical size (e.g., $|V| \approx 50,000$), this softmax output layer dominates both parameter count and GPU memory bandwidth during training and inference (Wu et al., 4 Apr 2026).

Tree-structured diffusion models offer two distinct augmentations:

• Syntax-aware code generation: Leveraging the tree structure of code (ASTs) to guide the forward corruption and reverse denoising processes, preserving grammatical boundaries (Zeng et al., 2 Aug 2025).
• Tree-structured token prediction: Employing a hierarchical K-ary tree over the vocabulary, such that prediction at each diffusion step factorizes into a product of small branching decisions rather than a flat softmax, drastically reducing resource requirements (Wu et al., 4 Apr 2026).

2. AST-Guided TDLM for Code Generation

In TDLMs for code, the code portion of each sequence is parsed into an abstract syntax tree $G = (V, E)$. Each node $v \in V$ corresponds to a token interval $(s_v, e_v)$ covering a syntactically meaningful subtree. Forward diffusion proceeds by masking not individual tokens, but complete subtree spans, sampled via span-level probabilities that match a desired token-masking rate $\varepsilon_t$ in expectation. The corruption model for a given AST span $x_T$0 of length $x_T$1 is

$x_T$2

If $x_T$3, all tokens in $x_T$4 are masked. The total number of masked tokens is thus controlled in expectation: $x_T$5 where $x_T$6.

Reverse denoising uses a Transformer decoder $x_T$7 conditioned on timestep $x_T$8 to predict the original content of masked tokens. Critically, no changes are made to the backbone architecture, but the noise patterns—being AST-aware—encourage the model to internalize hierarchical and grammatical code structure. The training loss is the cross-entropy (equivalently, ELBO) between ground-truth and predicted tokens over masked positions, averaged over sampled $x_T$9 and maskings (Zeng et al., 2 Aug 2025).

3. Tree-Structured Token Prediction for Efficient Diffusion

Standard diffusion LMs require full-vocabulary prediction at every step, which is computationally prohibitive at large scale. TDLM for general language instead exploits the hierarchical structure of modern subword vocabularies by constructing a balanced $x_0$0-ary tree $x_0$1 whose $x_0$2 leaves correspond bijectively to subword tokens. Tree construction is typically done through recursive $x_0$3-means clustering of pretrained token embeddings, with path padding as needed for uniform depth $x_0$4 (Wu et al., 4 Apr 2026).

For token $x_0$5, the root-to-leaf ancestry is $x_0$6, with $x_0$7 (nodes at level $x_0$8). In the generative process, forward diffusion refines each token in $x_0$9 coarse-to-fine stages, with individual noise schedules per tree level:

$x_T$0

At any moment $x_T$1 the latent $x_T$2 has support only on $x_T$3. Reverse denoising proceeds in $x_T$4 stages, each predicting which child of the current node corresponds to the ground-truth token's path, thus factorizing full softmax prediction into a sequence of $x_T$5-way (often $x_T$6) decisions.

The training objective decomposes into a sum of cross-entropies (ELBO terms) for each level and is proportional to

$x_T$7

This leads to an exponential reduction in logit dimensionality and parameter cost for the prediction heads, enabling deeper transformer backbones within the same parameter budget.

4. Algorithmic Workflow

AST-Guided TDLM Diffusion Step

1. Take input sequence $x_T$8 and extract AST spans from code $x_T$9.
2. For timestep $|V|$0, compute mask budget $|V|$1.
3. Mask complete AST spans according to sampled $|V|$2; if unmasked tokens remain, mask additional tokens randomly to meet $|V|$3.
4. Create corrupted sequence $|V|$4 with $|V|$5 in masked positions.
5. Model is trained to predict $|V|$6 (or $|V|$7) from $|V|$8, learning to reconstruct entire subtrees (Zeng et al., 2 Aug 2025).

Vocabulary Tree-Structured TDLM

1. Construct balanced $|V|$9-ary tree over vocabulary.
2. For each training batch, maintain for each token its ancestry through the tree.
3. At inference or learning, reverse denoise by predicting the appropriate child at each level, using minimal-parameter tree-heads attached to the transformer outputs (Wu et al., 4 Apr 2026).

5. Empirical Evaluation and Results

Code-Oriented AST-Guided TDLM

Experiments on 150,000 samples from OpenCodeReasoning (with 1,000 validation) evaluate pass@1 (exact match plus unit-test) on HumanEval and MBPP, as well as syntactic validity by Python parsing.

Model	HumanEval pass@1 (512)	HumanEval pass@1 (1024)	MBPP pass@1 (512)
Random Masking	31.71%	33.54%	31.13%
AST-Token Masking	31.71%	28.66%	24.51%
TDLM (ours, AST span-masking)	32.93% (+1.22%)	36.59% (+3.05%)	33.07%

Qualitative cases indicate TDLM's improved semantic correctness and generalization, particularly in tasks requiring reasoning about set equality rather than superficial matches. As input length increases, token-level masking degrades (due to fragmentation of sub-expressions), whereas span-level AST masking preserves grammatical structure and robustness (Zeng et al., 2 Aug 2025).

Vocabulary-Tree TDLM

Experiments on OpenWebText, comparing DiT-style models with sequence length 512 and ~131 billion token updates:

Model	Val PPL	Gen PPL	Layers/Params	Head Params
MDLM-small	≤27.39	163.7	12 / 92M	(full softmax)
TDLM-small	≤25.50	159.3	17 / 130M	0.4M (tree)
HDLM-base	≤19.22	139.9	24 / 321M	(full softmax)
TDLM-base	≤18.95	138.0	27 / 361M	0.4M (tree)

Peak GPU memory usage is halved, output projection parameter cost drops by ≳100×, and available capacity can be reallocated to deepening the transformer backbone, all while matching or exceeding state-of-the-art perplexities (Wu et al., 4 Apr 2026).

Ablations show shallow trees ($|V| \approx 50,000$0) with high branching factor ($|V| \approx 50,000$1) yield optimal trade-offs; unbalanced or uneven schedule allocations do not improve performance.

6. Design Considerations and Limitations

Implementation of TDLM with vocabulary trees requires only modifying the prediction head and tree construction, leaving the transformer backbone unchanged—a one-time cost. However, performance depends critically on tree quality; excessively deep trees (small $|V| \approx 50,000$2) accumulate loss over more levels, reducing efficacy. Tree joint modeling of multi-token paths offers modest gains only at small scale.

Inference speed is impacted by the $|V| \approx 50,000$3 iterative diffusion, but the reduction in output dimensionality enables throughput ∼1.25× that of conventional DLMs for equal diffusion steps. For AST-guided TDLM, improvements are tied to structured domains (code), with no evidence provided for natural-language syntax trees (Wu et al., 4 Apr 2026, Zeng et al., 2 Aug 2025).

7. Research Impact and Outlook

TDLM demonstrates that infusing hierarchical or syntactic structure into diffusion-based sequence models affords substantial gains in both modeling fidelity and computational efficiency. In code generation, AST-aware span masking enables the model to respect grammatical boundaries, improving syntactic correctness and generalization, especially for long or deeply nested constructs. For general language modeling, tree-factored prediction reduces both parameter and activation-memory footprint, facilitating more efficient scaling under resource constraints.

A plausible implication is that future language modeling research will further explore structured denoising pathways, including syntax and ontology-aware diffusion, and dynamically adaptive tree construction, to blend the interpretability and inductive bias of structure with the flexibility of parametric models. Limitations regarding tree depth, construction quality, and diffusion step complexity remain open areas for improvement and investigation (Zeng et al., 2 Aug 2025, Wu et al., 4 Apr 2026).