PaperlayerView exact revision on arXiv

Source-linked AI summary

Hierarchical Masked Diffusion Language Models

Parul Gupta, Satwik, Avni Aggarwal, Bhavesh Gurnani, Arnab Kumar Mondal, Dinesh Khandelwal, Dinesh Garg, Parag Singla

arXiv:2610.29255v1cs.CLcs.AI

TL;DR

Standard masked diffusion uses a shared schedule and does not encode which positions should be resolved first, despite natural task-specific priorities. HMDLM adds hierarchy to positions and uses level-specific schedules with a two-stage prediction process. The paper reports gains across language generation and Sudoku compared with diffusion baselines.

  • Problem

    Standard masked diffusion uses a shared noise schedule without explicitly encoding which positions should be resolved first, although many tasks have natural positional priorities.

    Source: §1 · pp. 1–2

  • Method

    HMDLM augments positions with hierarchy levels, predicts hierarchy before token identities, and uses a separate noise schedule for each level.

    Source: §1 · p. 2, §4 · p. 4, §4.1 · p. 5

  • Results

    HMDLM reports gains over diffusion baselines across language and structured tasks, including up to 8% higher accuracy than MDLM for the most difficult Sudoku cells.

    Source: §1 · p. 2, §6 · p. 9, §5.2 · p. 9

  • Takeaways & Limitations

    The experiments support hierarchy-aware generation as a promising direction for diffusion-based language models, with added computation from the two-stage design.

    Source: §6 · p. 10

    The two-stage design adds computation compared with the baseline approach.

    Source: §6 · p. 10

Abstract

from arXiv · show

Masked diffusion language models have emerged as an alternative to autoregressive models, enabling parallel token generation. However, standard masked diffusion language models use a shared noise schedule across positions without explicitly encoding which positions should be resolved first. Many tasks exhibit a natural priority among positions: document titles and section headings can guide the generation of body text, while Sudoku cells determined directly from the initial clues can help resolve more difficult cells. To incorporate this notion of priority (or hierarchy) among token positions, we introduce the Hierarchical Masked Diffusion Language Model (HMDLM), which augments each token position with a hierarchy level available as part of the training data, and requires zero or minimal additional annotations. It factors generation into two passes: the first predicts a hierarchy distribution over masked positions, the second predicts token distributions conditioned on it. Giving each hierarchy level its own noise schedule induces a soft, controllable preference for revealing higher-priority positions earlier while keeping decoding fully parallel. Our formulation is mathematically grounded, and we re-derive the forward and reverse processes and the NELBO for this factorization. Our evaluation shows that for the task of language generation on LM1B, HMDLM reduces perplexity by more than 1.5 points over Anchored Diffusion Language Model (ADLM), a recently proposed strong diffusion baseline; on WikiHow, we see a drop in perplexity of up to 4.4 pts on ADLM and up to 7 pts against MDLM. An LLM judge on WikiHow prefers HMDLM generations (the two variants combined) more than $2\times$ as often as the baseline. On the task of solving Sudoku puzzles requiring multi-step reasoning, it improves solution accuracy by up to $8.6\%$ absolute points over existing diffusion baselines.

1 INTRODUCTION

Masked diffusion generates tokens in parallel, but standard models use a shared masking schedule without encoding which positions should be denoised first. HMDLM adds hierarchy to positions and evaluates the approach across language and structured benchmarks.

1 INTRODUCTION

  • Standard masked diffusion treats positions symmetrically, while tasks such as document and Sudoku generation have natural priorities among positions.For Sudoku, priority belongs to the cell position rather than the digit it will contain.Source: §1 · pp. 1–2
  • HMDLM associates each token position with a hierarchy variable, then predicts hierarchy before token identity during reverse diffusion.It uses a separate noise schedule per hierarchy level and derives a NELBO for training.Source: §1 · p. 2
  • Across LM1B, WikiHow, and Sudoku, experiments report gains over diffusion baselines in perplexity, generation quality, and solution accuracy.LM1B comparisons include MDLM, SEDD, ADLM, and DiffusionBERT; WikiHow compares with ADLM and MDLM, while Sudoku compares with MGDM and MDLM.Source: §1 · p. 2
  • The paper contributes a hierarchy-aware diffusion formulation, a two-stage training and inference process, a re-derived ELBO, and benchmark experiments.The first stage predicts a distribution over hierarchy levels; the second predicts tokens given the hierarchy prediction.Source: §1 · p. 2

2 RELATED WORK

Masked diffusion work has explored denoising objectives and ordering strategies, but typical models use uniform masking schedules. HMDLM places task-dependent hierarchy in the diffusion state and uses it to shape denoising.

2 RELATED WORK

  • Masked diffusion models iteratively reconstruct partially masked sequences, but typically use uniform masking schedules that ignore task hierarchy.The literature includes approaches such as SEDD, DiffusionBERT, MDLM, and large-scale systems including LLaDA.Source: §2 · p. 2
  • Prior work improves denoising order through learned ordering, planner-aware training, planned generation, or sampling-trajectory changes.Examples include inference-time ordering (Kim et al., 2025), planner-aware training (Peng et al., 2025), and planned generation (Israel et al., 2025).Source: §2 · pp. 2–3
  • Unlike anchor-based and planning approaches, HMDLM assigns multiple task-defined hierarchy levels to positions and uses level-specific schedules in forward corruption and reverse denoising.Anchored Diffusion Language Models (Rout et al., 2025) predict anchor-token distributions before missing tokens; Planned Diffusion (Israel et al., 2025) constructs a plan before generating spans.Source: §2 · p. 3

3 BACKGROUND

The background defines discrete tokens and masked diffusion notation, then describes forward corruption, reverse denoising, and the training objective. These elements provide the baseline process that HMDLM extends.

3.1 NOTATIONS

  • The notation represents each discrete token as a one-hot vector and treats [MASK] as the Dth category in the token vocabulary.The diffusion state x_t denotes the process state at time step t.Source: §3.1 · p. 3

3.2 DISCRETE DIFFUSION MODELS

  • The forward process gradually corrupts clean data over T steps by replacing tokens with [MASK], from zero noise at t = 0 to maximum noise at t = T.The masking schedule α_t decreases over time, with α_0 = 1 and α_T = 0.Source: §3.2 · p. 3
  • The reverse process iteratively denoises toward an unmasked sequence, using a neural network to predict clean-token distributions from noisy inputs.The denoising model is trained using the Negative ELBO.Source: §3.2 · pp. 3–4

4 HIERARCHICAL MASKED DIFFUSION LANGUAGE MODEL

HMDLM augments token positions with hierarchy states and uses hierarchy-dependent masking schedules in forward and reverse diffusion. Its two-stage model predicts hierarchies before token identities, with separate hierarchy and token terms in the ELBO.

4 HIERARCHICAL MASKED DIFFUSION LANGUAGE MODEL

  • HMDLM adds a hierarchy state to each token position and uses hierarchy-dependent masking schedules in both forward and reverse diffusion.A token and its hierarchy are masked or revealed jointly.Source: §4 · p. 4

4.1 FORWARD AND REVERSE PROCESS

  • Higher-hierarchy positions are masked more slowly forward and receive greater reveal probability during reverse diffusion, inducing a soft generation order.The reverse reveal probability averages level-specific reveal rates under the predicted hierarchy distribution.Source: §4.1 · pp. 4–5

4.2 ELBO

  • The discrete ELBO expands by the KL chain rule into hierarchy and token terms, modeled by the architecture’s two stages.The paper also derives a continuous-time ELBO limit.Source: §4.2 · p. 5

4.3.1 TWO STAGE MODEL

  • At each reverse step, Stage 1 predicts hierarchy distributions for masked positions, then Stage 2 predicts token distributions conditioned on the revealed hierarchy states.Training uses gold hierarchy and token pairs for unmasked positions; inference keeps each revealed pair fixed thereafter.Source: §4.3.1 · pp. 5–6

4.3.2 PER-LEVEL NOISE SCHEDULE AND HIERARCHY EMBEDDINGS

  • Each hierarchy level receives its own masking schedule, so more informative, higher-level tokens are masked more slowly and unmasked earlier.Learned hierarchy embeddings are added to token embeddings before they enter the backbone.Source: §4.3.2 · p. 6

4.4 TRAINING AND INFERENCE

  • Training applies each position’s hierarchy-specific masking schedule and passes the resulting state through hierarchy prediction followed by token prediction.Inference starts from a fully masked sequence and iteratively samples hierarchy states to determine which positions are revealed.Source: §4.4 · p. 6

5 EXPERIMENTS

Experiments assess HMDLM on language modeling with LM1B and WikiHow and on Sudoku solving, using perplexity, generation-quality measures, judge preferences, and puzzle accuracy. HMDLM performs strongly against diffusion baselines, with results depending on the hierarchy assignment and task.

5 EXPERIMENTS

  • The evaluation covers LM1B and WikiHow language generation and Sudoku solving, using perplexity and generation-quality measures for language and accuracy for puzzles.LM1B models are trained from scratch; WikiHow models are fine-tuned from pretrained checkpoints, and Sudoku models are trained from scratch.Source: §5 · p. 7, §5.1.1 · p. 7, §5.1.2 · p. 8

5.1.1 LM1B DATASET

  • At the 65B-token budget on LM1B, HMDLM has the best model and generative PPL among diffusion baselines, with generative PPL nearly 95 points below MDLM.Entropy is 5.876 versus 6.022 for MDLM, with no clear evidence of severe diversity collapse; MAUVE is also higher than MDLM, and the judge prefers HMDLM 56% to 44%.Source: Table 1 · p. 7, Table 2 · p. 7, §5.1.1 · p. 8

5.1.2 WIKIHOW DATASET

  • On WikiHow, HMDLM-P has the best generative PPL across all three lengths, while HMDLM-F has the best model PPL and is closest to the reference distribution in generative PPL.The LLM judge preferred HMDLM-P 40.75% and HMDLM-F 31.88% of the time, versus ADLM 15.25% and MDLM 12.12%.Source: §5.1.2 · pp. 8–9, Table 4 · p. 8
  • On WikiHow, HMDLM-F has better model PPL, MAUVE, and entropy closer to the reference than HMDLM-P at all lengths; the authors hypothesize key tokens matter more than positions for this dataset.They note this is domain-dependent and suggest position-based hierarchy may be essential for tasks such as Sudoku.Source: §5.1.2 · p. 9

5.2 SUDOKU DATASET

  • On Sudoku, HMDLM beats the baseline in solve rate and cell accuracy for both 25M and 6M models, with gains up to 8% on the hardest level-0 cells.For the 25M model, 45.6% of 287 puzzles solved by HMDLM but not MDLM involved MDLM committing early to a search cell that HMDLM filled later, when fewer candidates remained.Source: §5.2 · p. 9, Table 5 · p. 9

6 CONCLUSION AND LIMITATIONS

HMDLM jointly diffuses tokens and hierarchy labels, then factors reverse prediction into hierarchy and token stages. The conclusion presents hierarchy-aware generation as promising while noting the added computation of its two-stage design.

6 CONCLUSION AND LIMITATIONS

  • 18k Sudoku-hard puzzles form Table 6’s test set, with 6M-parameter models trained for 100k steps and evaluated over three sampling seeds.The table reports mean ± standard deviation and uses T = 81 with confidence-based sampling.Source: Table 6 · p. 9
  • HMDLM assigns priority through level-specific noise schedules and predicts hierarchy before token identity, enabling structured, controllable generation.The conclusion describes the design as lightweight and modular, while noting the two-stage computation cost.Source: §6 · p. 9

A.1 JOINT GENERATIVE MODEL pθ(x0, h0)

  • At each position, masking the token also masks its hierarchy label, and unmasking reveals both together.The joint diffusion state is (x_t, h_t).Source: §A.1 · p. 13

A.3 STEP 2 — APPLY JENSEN’S INEQUALITY

  • The derivation applies Jensen’s inequality to the joint model and expands its log probability into terminal-prior and reverse-step terms.It then telescopes the forward-process terms to express the objective through reverse-model and forward-posterior comparisons.Source: §A.3 · p. 13, §A.4 · p. 14, §A.5 · p. 14

A.5 STEP 4 — TELESCOPE THE FORWARD TERMS

  • The terminal prior factorizes into token and hierarchy priors, matching the fully masked terminal state used in the derivation.The passage says the reverse model is also factorized, following the hierarchy-first, token-second order.Source: §A.2 · p. 13, §A.5 · p. 15

A.2 STEP 1 — INTRODUCE THE JOINT DIFFUSION PROCESS

  • The reverse process first predicts the preceding hierarchy state, then predicts the preceding token conditioned on that hierarchy state.The token prediction uses h_(t−1), which carries the hierarchy information needed for that step.Source: §A.2 · p. 13, §A.5 · p. 15

A.6 STEP 5 — DECOMPOSE THE JOINT KL TERM Lt (FULL DERIVATION)

The appendix decomposes the joint reverse-process KL according to HMDLM’s hierarchy-first factorization. It then expands the hierarchy and token terms into reveal-weighted losses over positions.

A.6 STEP 5 — DECOMPOSE THE JOINT KL TERM Lt (FULL DERIVATION)

  • The joint KL splits into a hierarchy KL and a token KL, matching Pass 1’s hierarchy prediction and Pass 2’s token prediction conditioned on hierarchy.The decomposition applies the KL chain rule in the same variable order as the reverse model.Source: §A.6 · p. 15

A.7 STEP 6 — SIMPLIFY THE TERMINAL TERM LT

  • At the fully masked terminal state, the prior matches the forward distribution, so the terminal KL term is dropped.The passage states that every token and hierarchy label is fully masked at the terminal time.Source: §A.7 · p. 15

A.9 STEP 8 — EXPAND THE HIERARCHY KL (PASS 1)

  • Because positions factorize independently in the forward process and hierarchy head, the hierarchy KL becomes a sum of per-position KLs.The Pass-1 distribution assigns probabilities over the K hierarchy classes at each position.Source: §A.9 · p. 16
  • The hierarchy loss is a reveal-weighted cross-entropy on Pass 1’s distribution over hierarchy classes.For the gold class, the reveal rate cancels between numerator and denominator, leaving cross-entropy on the predicted class distribution.Source: §A.9 · pp. 16–17

A.10 STEP 9 — EXPAND THE TOKEN KL (PASS 2)

  • The token loss is a reveal-weighted cross-entropy on Pass 2’s softmax, weighted by the true hierarchy class’s reveal probability.Only masked positions contribute; unmasked positions have zero KL.Source: §A.10 · p. 18, DISCRETE →CONTINUOUS · p. 19

A.11 STEP 10 — FINAL TRAINING LOSS

The training procedure combines token, hierarchy, and consistency losses. It samples a subset of positions with gold hierarchy during training so the token model also encounters masked hierarchy context, as it does at inference.

A.11 STEP 10 — FINAL TRAINING LOSS

  • The final training objective combines token, hierarchy, and consistency loss components, with tunable coefficients replacing tied weights in practice.The implementation describes optimizing a reweighted sum of the same three terms.Source: §A.11 · p. 20
  • Training samples a subset of masked positions to receive gold hierarchy, while the remaining positions retain masked hierarchy for the token prediction stage.The authors added this procedure because training with gold hierarchy everywhere did not yield the expected inference performance, where the second stage sees masked hierarchy.Source: §A.11 · p. 20
  • The algorithm first predicts hierarchy and then samples tokens at revealed positions conditioned on the sampled hierarchy.Committed positions are frozen and not revisited in later reverse steps.Source: §A.11 · p. 21

B EXPECTED MASK FRACTION UNDER HIERARCHICAL LOG-LINEAR SCHEDULE

The appendix derives expected masking and supervised-token budgets under a hierarchical log-linear schedule. It also describes assigning rare tokens to a hierarchy level using their within-sequence frequency.

B EXPECTED MASK FRACTION UNDER HIERARCHICAL LOG-LINEAR SCHEDULE

  • For LM1B, a token is assigned to level 1 when it occurs at most τ times within its own sequence; at τ = 1, that means exactly once.The passage reports 737,634,266 level-1 tokens at τ = 1 and 914,638,961 at τ = 5.Source: B EXPECTED MASK FRACTION UNDER… · pp. 22–23
  • At τ = 1, the reported expected mask fraction is approximately 0.399633; at τ = 5, it is approximately 0.375669.The passage gives the τ = 5 result directly and reports the τ = 1 value in the budget comparison.Source: B EXPECTED MASK FRACTION UNDER… · p. 23
  • The expected supervised-token count is the number of steps times batch size, mean real sequence length, and expected mask fraction.Padding is excluded from the supervised-position budget.Source: B EXPECTED MASK FRACTION UNDER… · p. 23
  • With batch size 512 and mean real length 30.44, HMDLM supervises about 6,228.3 tokens per step at τ = 1 and 5,854.8 at τ = 5, versus MDLM’s 7,784.7.These counts use expected mask fractions for real positions.Source: B EXPECTED MASK FRACTION UNDER… · p. 23

C EXPECTED MASKING FRACTION FOR WIKIHOW

The section describes expected masking rates and step-matched training budgets for WikiHow. It distinguishes level-1 token definitions under the summary and count rules.

C EXPECTED MASKING FRACTION FOR WIKIHOW

  • A response token at level ℓ is masked at time t with probability (1 − ε)t^wℓ, and its expected masking rate is derived from this schedule.The section also defines gℓ as the fraction of tokens assigned to level ℓ, pooled over diffused positions and excluding prompts and padding.Source: C EXPECTED MASKING FRACTION FOR WIKIHOW · p. 23
  • Table 7 reports step-matched WikiHow training budgets at context lengths 512 and 1024, with g1 denoting the fraction of tokens at level 1.Level 1 is the article summary under the summary rule; under the count rule with τ=1, it is any token occurring at most once in its sample.Source: Table 7 · p. 24
  • To match MDLM’s expected number of masked, loss-bearing tokens at the same batch size, HMDLM’s training steps are adjusted.The WikiHow MDLM runs use approximately 29.8 epochs at each context length; the passage gives step counts for both lengths.Source: C EXPECTED MASKING FRACTION FOR WIKIHOW · p. 24

D PER-LEVEL SCHEDULES AND THE INDUCED REVEAL ORDER

The appendix specifies per-level schedules that induce a soft reveal order, derives continuous-time ELBO weights for perplexity evaluation, and gives the model architectures used in experiments.

D PER-LEVEL SCHEDULES AND THE INDUCED REVEAL ORDER

  • While a higher level is active, its predicted positions receive greater reveal mass and are revealed earlier, inducing HMDLM’s soft generation order.The schedules used in experiments satisfy the required ordering of per-step reveal probabilities; the ordering is vacuous when a level is inactive.Source: D PER-LEVEL SCHEDULES AND THE INDUCED… · pp. 24–25

E PERPLEXITY: CONTINUOUS-TIME ELBO WEIGHT

  • For a level with offset δk, the continuous-time ELBO weight is 1/(t − δk) on its active interval and is zero before that interval.At δ0 = 0, the weight recovers the MDLM weight 1/t; evaluation uses the model’s posterior over hierarchy levels because the gold hierarchy is unavailable.Source: E PERPLEXITY: CONTINUOUS-TIME ELBO… · p. 25

F.1 ARCHITECTURE DETAILS

  • The 6M HMDLM adds a level embedding and prediction head, about 5K parameters, while the 25M HMDLM adds about 104K parameters over its MDLM counterpart.The 6M models use a compact three-layer GPT-2 Transformer; the 25M models use a 12-layer diffusion transformer with six hierarchy levels.Source: §F.1 · p. 26

F.2 QUALITATIVE ANALYSIS

The Sudoku qualitative analysis finds that HMDLM tends to defer ambiguous cells until neighboring cells are filled. It also reports improved accuracy on the hardest hierarchy level.

F.2 QUALITATIVE ANALYSIS

  • In Example 1, MDLM predicts R9C6 incorrectly at round 2, whereas HMDLM delays revealing it until the only possible value is 5.HMDLM fills other cells before resolving R9C6.Source: §F.2 · p. 27
  • In Example 2, MDLM incorrectly predicts R7C1 at round 1, while HMDLM predicts neighboring cells first and later gets R7C1 right.The later prediction occurs at round 36.Source: §F.2 · pp. 28–29
  • In 131 of 287 puzzles HMDLM solved but MDLM did not, MDLM’s first error was an early commitment to a search cell that HMDLM filled only after one candidate remained.Across the examples, MDLM tends to unmask ambiguous cells early and err, while HMDLM tends to defer them until neighboring cells are filled.Source: §F.2 · p. 29

F.3 CELL ACCURACY PER HIERARCHY LEVEL

  • HMDLM improves accuracy for level-0 cells, the most difficult cells, by up to 8% over MDLM.Level 0 includes cells forced only after six or more rounds or never forced by the stated rules, and thus requiring search.Source: §F.3 · p. 29, Table 8 · p. 29

G QUALITATIVE ANALYSIS ON WIKIHOW

The WikiHow examples compare generated responses for two how-to prompts, with HMDLM judged more prompt-focused than MDLM and ADLM. The appendix also provides evaluation-prompt and computational-resource details.

G.1 PROMPT: HOW TO REPLACE THE REAR SHOCK ON AN 1150RT

  • For the rear-shock prompt, HMDLM stays on task and refers to relevant components, though its step order is repetitive and somewhat confused.ADLM is largely incoherent and gives physically implausible instructions, while MDLM drifts to unrelated topics and never addresses the shock replacement.Source: §G.1 · p. 30

G.2 PROMPT: HOW TO BUY A CUSTOM FIT CAR COVER

  • For the custom-fit car-cover prompt, HMDLM gives the most relevant response and follows a sensible sequence, despite a few garbled sentences.MDLM gives unrelated and contradictory instructions, while ADLM shifts to garment fitting and is largely incoherent.Source: §G.2 · p. 31

H.1 LM1B PROMPT

  • The LM1B LLM-as-judge prompt asks evaluators to choose the better standalone generated passage or declare a tie, without considering the generating model.The instructions also say not to favor a sample for being longer or penalize it for being brief.Source: §H.1 · p. 31, §4 · p. 31

H.2 WIKIHOW PROMPT

  • The WikiHow LLM-as-judge prompt asks evaluators to select the best of four responses to the same how-to prompt.The provided prompt passage specifies the choice among A, B, C, and D.Source: §H.2 · pp. 31–32

I SOCIETAL IMPACT OF THE MODEL

  • The paper identifies potential benefits in language generation and reasoning applications, while noting that HMDLM could be misused to generate misleading or harmful text.Source: I SOCIETAL IMPACT OF THE MODEL · p. 32

J COMPUTATIONAL RESOURCES

  • LM1B HMDLM training used eight NVIDIA A100 80GB GPUs for approximately four days and completed 2.5M steps at a 65B-token budget.Source: J COMPUTATIONAL RESOURCES · p. 32
  • WikiHow HMDLM training used four NVIDIA A100 80GB GPUs for approximately eight hours; Sudoku training used four GPUs for four hours.Source: J COMPUTATIONAL RESOURCES · p. 32

Related papers

Loading 2610.29255v1…