mlgenerate.comVisual Generative ML
All Models/text/Masked Diffusion LM

Masked Diffusion Language Model (MDLM / LLaDA)

Diffusion2024 · Cornell / Stanford / Renmin Univ.

Continuous-time discrete token masking + iterative bidirectional unmasking

Execution Paradigm:Iterative Parallel Confidence Unmasking
Parallel Multi-Token Unmasking per Step
Interactive Stage (INFERENCE)All target generation slots start as [MASK] tokens.
MDLM Inference: Parallel Confidence UnmaskingStep 1 of 3

The canvas begins fully masked. At each step, all remaining [MASK] positions are predicted in parallel. The highest-confidence tokens are locked into the sequence, progressively anchoring the generation.

Active Generation Canvas (Length T = 8)Unmasked: 0 / 8
Slot 0
[MASK]
Pending
Slot 1
[MASK]
Pending
Slot 2
[MASK]
Pending
Slot 3
[MASK]
Pending
Slot 4
[MASK]
Pending
Slot 5
[MASK]
Pending
Slot 6
[MASK]
Pending
Slot 7
[MASK]
Pending
Step 1 of 3

1. Initialize Full [MASK] Canvas

Concatenate the user prompt with T_gen absorbing [MASK] tokens representing the entire unknown response canvas at t = 1.0.

💡 Takeaway: Response length T_gen is allocated upfront as a parallel canvas.

Active Tensor Dimensions & Memory Shapes

x_t[B, P + T_gen]
Prompt + all-masked canvas
sample_mdlm_iterative.py
1
def generate_mdlm(prompt, model, gen_len=64, steps=8):
2
xt = torch.cat([prompt, full((B, gen_len), MASK_ID)], dim=1)
3
for k in reversed(range(1, steps + 1)):
4
probs = model.bidirectional_forward(xt, k / steps).softmax(-1)
5
conf, pred_tokens = probs.max(dim=-1) # Confidence per slot
6
unmask_idx = select_top_k_confident(conf, is_masked=(xt == MASK_ID), k=gen_len // steps)
7
xt = torch.where(unmask_idx, pred_tokens, xt)
8
return xt
Click code line to jump to corresponding stepPyTorch Reference
Why This Architecture Matters

Enables arbitrary-order generation, bidirectional infilling, and controllable step-vs-quality trade-offs in language modeling without causal left-to-right constraints.

Runtime Complexity & Bottlenecks

Training Objective:Time-Weighted Masked Cross-Entropy ELBO: ∫ (α′_t / (1 - α_t)) CE(x_0, p_θ(x_t)) dt

Inference Bottleneck:O(K) full-sequence bidirectional passes