Simple and Effective Masked Diffusion Language Models

In plain words

Most language models write one token at a time, left to right. Diffusion models, which are famous for images, work differently: they start from noise and clean it up in steps. Applied to text, the "noise" is a sequence where some words have been hidden behind a blank mask token. Earlier work found that this kind of text diffusion trailed ordinary left-to-right models.

This paper argues that the gap was partly an implementation problem. The model is a plain encoder that looks at the whole sequence at once and guesses the hidden words, the same job as BERT-style masked language modelling. Training hides a random share of the words, asks the model to fill them in, and weights the losses in a principled way so they add up to a valid bound on how likely the text is. Generation starts from a row of blanks and fills them in over a series of steps, and the model can be run in chunks to write text of any length, one block after another.

The point for an engineer is that a masked-language-model recipe, trained and sampled this way, can generate text and comes closer to autoregressive models on likelihood benchmarks, while filling many positions at each step.

TL;DR

• Masked diffusion language model (MDLM): corrupt text by masking tokens, train an encoder-only transformer to recover them, and sample by unmasking from an all-mask start.
• The continuous-time objective is a weighted average of masked-token cross-entropy losses, and the paper's parameterization (SUBS) keeps unmasked tokens fixed and forbids the mask as an output.
• Author-reported: diffusion perplexity bounds beat earlier diffusion LMs on LM1B and OpenWebText, but still trail autoregressive baselines; the semi-autoregressive sampler is about 25–30 times faster than SSD-LM in the paper's setup §5.1.
• Relevant if you want a non-autoregressive or bidirectional text generator, or a BERT-style encoder that can also generate. The cost is many sampling steps and likelihood numbers that are upper bounds.

Problem

Discrete diffusion models for text, such as D3PM, lagged autoregressive models on log-likelihood under the then-standard setup §1. The paper argues that masking-only diffusion is simpler and that training recipes, not the framework, caused much of the gap §1; §3.5.1; §6.

Without this method, a diffusion LM needs a general discrete transition matrix and continuous-time Markov chain theory to get a tight bound §6, "Comparison to CTMC". Standard autoregressive sampling is also sequential and cannot fill positions in parallel §1.

Method

Forward process. Each token is independently replaced by the mask token with a probability set by a schedule α_t that decreases from about 1 at t=0 to about 0 at t=1 §3.1, eq. 4. Once masked, a token stays masked §3.2.1. The posterior of the clean token given a partly masked state is given in closed form §3.2.1, eq. 6.

q(\mathbf{z}_t \mid \mathbf{x}) = \mathrm{Cat}\big(\mathbf{z}_t;\ \alpha_t \mathbf{x} + (1-\alpha_t)\mathbf{m}\big) \qquad [\text{eq. 4}]

Reverse model (SUBS). The denoiser returns log-probabilities over the vocabulary. The paper enforces two properties §3.2.3, eq. 7:

1. the mask token is never an output (its logit is pushed to −∞);
2. an unmasked token is carried over unchanged (its distribution is a point mass on itself).

Both are implemented as overwriting the network's logits, not as a learned behaviour diffusion.py.

Objective. Starting from the discrete-time diffusion ELBO, the paper removes the reconstruction term (it vanishes with carry-over) and the prior term (it vanishes under full masking at t=1) §3.3; app. B. The continuous-time limit is §3.4, eq. 10:

\mathcal{L}_{\text{NELBO}} = \mathbb{E}_{q,\,t\sim U[0,1]}\left[ \frac{\alpha_t'}{1-\alpha_t}\,\log\big\langle \mathbf{x}_\theta(\mathbf{z}_t,t),\,\mathbf{x}\big\rangle \right] \qquad [\text{eq. 10}]

For a sequence of L tokens, the same objective is a sum over positions, each with the same weight §3.5, eq. 11. Unmasked positions contribute zero because of carry-over §3.5, "Note". The objective is therefore a weighted average of masked-token losses, with the weight set by the schedule §3.5.

Noise schedule. The main text keeps α_t general. The experiments use a log-linear schedule app. D.4 in which the masking probability at time t is p_m = t app. D.2. Under that schedule, the weight in eq. 10 equals 1/t. This is our derivation from eq. 10 and p_m = t; the paper does not write it out. The code uses the same weight diffusion.py.

Training (Algorithm 1). alg. 1 (app. B) samples a sequence and a time t, masks each token independently with probability t, and takes a gradient step on the weighted loss over all positions. The sampler that picks t uses a low-discrepancy scheme across the batch app. D.3.

# One MDLM training step  [alg. 1] [code: diffusion.py#L847-L916]
# Input:  tokens (B, L) of token ids; pad_mask (B, L), 1 for real tokens
# Output: scalar loss, averaged over real tokens

def training_step(tokens, pad_mask):
    # Draw one time per sequence, stratified across the batch  [app. D.3]
    times = stratified_uniform(B, low=eps, high=1.0)   # (B,)
    # Mask each token independently with probability t  [app. D.2]
    is_masked = uniform(B, L) < times[:, None]         # (B, L)
    noisy = where(is_masked, MASK_ID, tokens)          # (B, L)
    # Predict per-position log-probabilities  [§3.2.3]
    logits = backbone(noisy, times)                    # (B, L, V)
    log_probs = subs(logits, noisy)                    # (B, L, V)
    # Log-probability of the clean token at each position  [§3.5]
    true_logp = gather(log_probs, tokens)              # (B, L)
    # Weight from eq. 10 under the log-linear schedule, 1 / t  [eq. 10]
    weight = 1.0 / times[:, None]                      # (B, 1)
    # Unmasked positions are 0 here because SUBS carries them over  [§3.2.3]
    per_token = -weight * true_logp                    # (B, L)
    per_token = per_token * pad_mask                   # drop padding
    return sum(per_token) / sum(pad_mask)              # token average

SUBS as code. The mask logit is set to a very negative finite value, the logits are normalised, and each clean position is overwritten to a point mass on its observed token §3.2.3, eq. 7 diffusion.py.

# SUBS parameterization of the denoiser  [§3.2.3] [eq. 7]
# [code: diffusion.py#L261-L277]
# Input:  logits (B, L, V) from the backbone; noisy (B, L)
# Output: log_probs (B, L, V), a normalised distribution per position

def subs(logits, noisy):
    # The mask token is never a clean output  [§3.2.3]
    logits[:, :, MASK_ID] = NEG_LARGE         # code uses -1e6, paper -inf
    # Normalise over the vocabulary  [§3.2.3]
    log_probs = log_softmax(logits, axis=-1)
    # Clean positions copy their observed token  [§3.2.3]
    is_clean = noisy != MASK_ID
    log_probs[is_clean] = NEG_LARGE
    log_probs[is_clean, noisy[is_clean]] = 0.0
    return log_probs

Ancestral sampling. Generation starts from L mask tokens and runs N steps from t=1 down to a small ε §4.1. In the code, the mask probability at time t is 1 − exp(−σ(t)), which is t under the log-linear schedule diffusion.py. A masked token stays masked with probability m_s/m_t, where m = 1 − α, and otherwise becomes a token drawn from the model §4.1; §3.2.2. A final pass takes the argmax at every position diffusion.py.

# Ancestral sampling with a log-linear schedule  [§4.1] [app. D.2]
# Input:  batch, length L, number of steps N
# Output: tokens (batch, L)

def sample(batch, length, steps):
    # Start from a row of mask tokens  [§4.1]
    tokens = full((batch, length), MASK_ID)
    # Mask rates run from 1 down to a small eps  [§4.1]
    times = linspace(1.0, eps, steps + 1)
    for i in range(steps):
        # Log-linear mask rate is m = t  [app. D.2]
        log_p = subs(backbone(tokens, times[i]), tokens)
        # A masked token stays masked with probability m_s / m_t  [§3.2.2]
        stay = times[i + 1] / times[i]
        # Unmask with probability 1 - stay, drawing from the model  [§4.1]
        is_mask = tokens == MASK_ID
        unmask = is_mask & (uniform(batch, length) < 1.0 - stay)
        drawn = sample_categorical(log_p)
        tokens = where(unmask, drawn, tokens)
    # Final argmax at every position  [code: diffusion.py#L690-L697]
    return argmax(subs(backbone(tokens, 0.0), tokens), axis=-1)

Semi-autoregressive generation. To write text longer than the training window, the model generates a block of L′ tokens while keeping the last L − L′ tokens of the previous output as a clean prefix. The prefix is carried over at every step §4.2. The repository implements this sliding-window scheme with its own indexing, which commits the first L′ positions of each window and carries the rest diffusion.py. Check the indexing before reusing it.

# Semi-autoregressive generation of arbitrary length  [§4.2]
# Input:  window length L, stride L' < L, number of rounds R
# Output: tokens of length L + R * L'

def semi_ar(window, stride, rounds):
    # First block: a full window generated from masks  [§4.2]
    output = sample(batch=1, length=window, steps=N)
    for round in range(rounds):
        # Last L - L' tokens become a clean prefix  [§4.2]
        prefix = output[-(window - stride):]
        # Prefix is copied through every step; the rest are masks  [§4.2]
        block = concat(prefix, full((stride,), MASK_ID))
        block = sample_with_prefix(block, prefix)
        # Append the newly generated L' tokens  [§4.2]
        output = concat(output, block[-stride:])
    return output

Implementation details from the paper. The denoiser is a diffusion transformer (DiT) with timestep conditioning and rotary position embeddings, not the T5 architecture used by D3PM §3.5.1. The paper notes that tokenization matters, and that computing the objective only at masked indices is numerically more stable than materialising transition matrices §3.5.1. Each masked position is drawn independently in a step §4.1. When the model ignores time, a step that unmasks nothing reuses the cached denoiser output instead of a new network call §4.1. The code disables caching when time conditioning is on diffusion.py.

Hyperparameters

Values are stated in the paper unless marked "code". Settings from the code are the defaults in configs/, not necessarily what every paper run used.

Results

Author-reported. All diffusion numbers are upper bounds on perplexity (or log-likelihood) and are not directly comparable with exact autoregressive likelihoods unless the caption says so. Setup: LM1B and OWT with a 12-layer, 768-dimensional DiT-style model §5.1; app. D.4. Test-set perplexity on LM1B and held-out OWT documents §5.1. Baselines are the retrained AR model, SEDD, D3PM and others tab. 1.

LM1B test perplexity tab. 1. Parameters for diffusion models are 110M unless noted.

OpenWebText test perplexity tab. 2. Models trained on 262B tokens tab. 2 caption.

Zero-shot perplexity on OWT-trained models tab. 3. Caption says models trained for 524B tokens; see Pitfalls.

Diffusion values are upper bounds tab. 3 caption.

GLUE tab. 4. MosaicBERT pretrained on C4, then fine-tuned with MDLM app. D.6. Metrics: F1 for QQP and MRPC, Spearman for STS-B, accuracy otherwise, m/mm for MNLI.

On C4 validation perplexity, the text reports AR 22, pretrained BERT ≤78 and BERT + MDLM-FT ≤35 §5.1, "Downstream Task Evaluation".

Semi-autoregressive generation tab. 5. 200 sequences of 2,048 tokens on one 3090 GPU. Generative perplexity under pretrained GPT-2 §5.1.

DNA perplexity tab. 6. Caduceus MLM fine-tuned on HG38 with each objective §5.2; app. D.7. Diffusion values are upper bounds.

Genomic Benchmarks tab. 7. Top-1 accuracy across 5-fold cross-validation. The ± value is the max−min across 5 random seeds tab. 7 caption. Bold marks the best value per row, as in the source table, which also italicises the second best. Fine-tuning and sweep details are in app. D.7.

The caption describes the Mamba column as AR and the other columns as a pretrained Caduceus model fine-tuned with each diffusion parameterization. The MDLM column is the paper's method. Caduceus MLM is the masked-LM objective, not a diffusion one.

Ablations on LM1B tab. 8. Test PPL with the discrete-T variant at T = 1000 and std over 5 evaluation seeds.

Wall-clock sampling tab. 10. Minutes to generate 64 samples on one A5000 GPU app. E.2.

Discrete-T evaluation of a continuous-time model tab. 11. OWT, MDLM trained with T = ∞ and evaluated with discrete T app. E.3.

Time conditioning on OWT tab. 12 app. E.5.

Noise schedules tab. 9 app. E.1. The paper reports the same likelihood (bits per dimension, 3.30) for four schedules, with different per-datapoint variance: log-linear 1.81, cosine 3.30, cosine-squared 3.30, linear 7.57.

Paper vs code

Code: the official repository at commit c112c526d193436838c98d81455ee51f90309470, Apache-2.0. The repository has training configs for the OWT and LM1B settings, and a sampler. It has no scripts for the DNA or GLUE experiments, so those results cannot be checked against code.

1. Log-linear schedule. The appendix writes σ(t) = −log(1 − t) app. E.1. The code uses σ(t) = −log(1 − (1 − ε)t) with ε = 1e-3 noise_schedule.py. The training time is also drawn from [ε, 1], not [0, 1] config.yaml diffusion.py.
2. Mask logit. The paper sets the mask logit to −∞ §3.2.3. The code uses the finite value −1e6 diffusion.py diffusion.py. The effect after softmax is negligible, but a port to exact −∞ changes the arithmetic path.
3. Time conditioning. The main table's OWT entry, 23.21 tab. 2, matches the appendix row labelled "with time conditioning" tab. 12. The code default turns time conditioning off config.yaml, and in that mode the time input is set to zero diffusion.py. The OWT MDLM script does not override it train_owt_mdlm.sh, so that run reproduces the 23.05 row, not 23.21.
4. Low-discrepancy sampler. The paper describes stratification over the minibatch of N samples app. D.3. The code stratifies over the batch that reaches the sampler diffusion.py. The training script uses a per-device batch of 16 and accumulates gradients to reach 512 train_owt_mdlm.sh config.yaml. The stratification therefore looks like it runs over each micro-batch, not the global batch of 512. This is our reading of the code, not something the code states.
5. Training-only settings not in the paper. The code uses an EMA of weights with decay 0.9999 config.yaml, gradient clipping at 1.0 config.yaml, weight decay 0 config.yaml, and Adam betas 0.9/0.999 with ε = 1e-8 config.yaml. The paper does not state these for the language-model runs app. D.4. Reported numbers may depend on the EMA weights, which the evaluation uses by default config.yaml.
6. Config fields that do nothing. configs/noise/loglinear.yaml sets σ_min = 1e-4 and σ_max = 20 loglinear.yaml. The log-linear class takes only ε and ignores these fields noise_schedule.py noise_schedule.py.
7. Padding in the loss. On LM1B, padded positions are noised and passed to the backbone without an attention mask, but excluded from the loss dataloader.py diffusion.py. The paper says sequences are padded and truncated app. D.4 and does not say how padding enters the loss.
8. Normalisation of the loss. The code divides the summed per-token loss by the count of real tokens in the batch diffusion.py. The paper writes a per-sequence sum §3.5, eq. 11 and does not state the normalisation.
9. Semi-autoregressive steps. The code fixes the step size at dt = 0.001, so 1,000 steps per block diffusion.py. The paper gives no step count for this sampler §4.2. The code also commits the first L′ positions of each window diffusion.py, which is a sliding-window formulation of the paper's prefix scheme.
10. Final denoising. The code takes an argmax after the last step by default config.yaml diffusion.py. The paper does not mention a final argmax.
11. Generative-perplexity evaluator. The paper uses "a pre-trained GPT-2" §5.1 without a size. The code's default evaluator is GPT-2 large config.yaml.
12. Caching. The code disables the cache whenever time conditioning is on diffusion.py. The paper says the speed-up applies when the model ignores time §4.1.
13. LM1B MDLM script. The repository has LM1B scripts for AR and D3PM train_lm1b_ar.sh train_lm1b_d3pm.sh but none for LM1B MDLM. The LM1B MDLM settings in the paper app. D.4 are therefore not verified by code.
14. Backbone naming. The code's dit backbone is the one the paper describes config.yaml. The dimamba backbone also exists in the repository diffusion.py, but no DNA script or config is included.

No discrepancies found in the reconstruction-term removal, the carry-over and zero-mask substitutions, or the continuous-time weight diffusion.py.

Pitfalls

• Token counts disagree within the paper. OWT at 1M steps is 262B tokens in the experiments section and appendix D.2 §5.1 app. D.2, but 524B in the zero-shot table caption tab. 3 and in appendix E.4 app. E.4. The 262B figure follows from the stated count: 1M steps × 512 × 1,024 × 0.5. The 524B figure omits the 0.5 factor for masking. For the representation fine-tuning, §5.1 says 327M tokens, which matches 5k steps × 512 × 128 ≈ 328M, while appendix D.6 says 32M. Treat 327M as the likely value, but check the paper.
• Headline gap claim. The introduction says the model approaches AR perplexity within 15–25% §1. The LM1B section says "within 14%" §5.1. Against the table, the 327B-token MDLM bound is 23.00 vs 20.86 for AR, about 10% higher, and the 33B-token bound is 27.04 vs 22.32, about 21% higher. The percentages in the text do not match a single clear comparison in the table, so recompute before quoting them.
• Bounds vs exact likelihood. Diffusion numbers are upper bounds tab. 1 tab. 2 tab. 3 tab. 6. The AR numbers are not labelled as bounds. Compare the two with care.
• Validation is one random draw. The likelihood evaluation uses a single Monte-Carlo estimate of t per example app. D.1. The validation loop uses the same stratified sampler as training diffusion.py diffusion.py. Reported numbers therefore depend on the random draw and on batch composition. Spread across seeds appears only in the LM1B ablation tab. 8 and the DNA tables tab. 6 tab. 7; the main LM and OWT tables give no spread.
• Table and equation numbers. Table and figure numbers here follow the order of the LaTeX source, and equation numbers follow the order of labelled equations in the main text. Float placement may change this in the rendered PDF, so check against the PDF before citing numbers.
• Sampling cost. Sampling needs one network call per step, so the step count multiplies cost. The OWT training script sets 1,000 sampling steps train_owt_mdlm.sh, and the caching figure varies T from 100 to 10,000 app. E.2. The cache is reused only when no token changes in the whole batch, and it is off when time conditioning is on diffusion.py. The paper's 2× speed-up is reported at batch size 1 app. E.2.
• Semi-autoregressive quality depends on the stride. Only one stride (L′ = 512 for L = 2048) is reported §5.1. Quality and speed at other strides are not reported.
• Carry-over assumption. The paper's zero-reconstruction-loss step depends on the carry-over substitution holding exactly §3.3; app. B. The paper notes that a model that learns this behaviour instead of having it built in would need a different objective §3.3.
• Log-linear weight. The per-token weight is 1/t under the log-linear schedule, which grows without bound as t → 0. The code bounds t from below at ε = 1e-3 diffusion.py, so the largest weight is near 1,000.
• Results with unclear settings. The DNA runs use 5-fold CV with a per-task sweep over batch size and learning rate, and the repository has no DNA code to check them against app. D.7. The GLUE evaluation uses batch size 16 for the Hugging Face script, while MDLM fine-tuning uses 512 app. D.6.

Lineage

• Builds on: D3PM (Austin et al., 2021), which MDLM specialises to masking-only processes §6, "Comparison to D3PM". Not linked: no arXiv id in the bibliography. SEDD (Lou et al., arXiv:2310.16834), whose log-linear schedule and transformer architecture MDLM reuses app. D.4 §5.1.
• Compared against: Diffusion-LM, Plaid, SSD-LM (semi-autoregressive), DiffusionBERT, and the retrained AR baselines tab. 1 tab. 5 tab. 6 §5.1.
• Concurrent work with a similar simplified objective: Shi et al. (2024) and Ou et al. (2024). Neither has an arXiv id in the bibliography; not linked §6, "Concurrent Works".
• Caduceus (Schiff et al., arXiv:2403.03234) is the DNA backbone used for the DNA experiments §5.2.
• Superseded by: not known from content already in this wiki.