Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking

TL;DR

MDM-Prime augments masked diffusion models by representing each token as an L-element sub-token sequence under a base- $\sqrt[L]{C}$ encoding, then performing masking and reverse-diffusion at the sub-token level. The resulting partial mask states (e.g. 0m halfway between 00 and 01) eliminate idle reverse-diffusion steps — sampling steps where the sequence is unchanged — taking the idle-step ratio from 36.77% (MDLM, L=1) down to 0.25% (L=5). At those higher utilizations MDLM-Prime hits 15.36 PPL on OpenWebText, the first MDM-only result to beat an AR transformer (17.54) and the EDLM-coAR / BD3-LM hybrids (~17.58) without any autoregressive formulation. On images it also pulls discrete generative modeling to parity with continuous DDPM / StyleGAN+ADA (CIFAR-10 FID 3.26).

Key claims

Standard MDM wastes ~37% of its 1024 reverse-diffusion steps on idle updates where the sequence is unchanged between consecutive timesteps, because tokens can only flip masked → unmasked once [§1, Fig. 1].
Representing each of $C$ token classes as an $L$ -element sub-token sequence under a base- $\sqrt[L]{C}$ encoding, then masking sub-tokens independently, introduces $L^{(L-1)}\cdot C^{(L-1)/L} + L$ intermediate states per token that interpolate between fully-masked and fully-unmasked [§3.1, Prop. A.3].
The sub-token diffusion objective remains a valid variational upper bound on negative log-likelihood when the decoder models the joint sub-token distribution rather than the factorized one, because injective encodings make some sub-token sequences invalid; the parameterization sets those logits to zero [§3.2, Eq. 7, Prop. A.4].
A factorized sub-token decoder, while simpler, degrades the sampling distribution monotonically as $L$ grows because it imposes an extra independence assumption across sub-tokens [§3.2, Fig. 5(a)].
Adapting a standard MDM to MDM-Prime requires only an embedding-layer change: $L$ separate sub-token embeddings of width $d/L$ are concatenated to form a $d$ -dim token embedding; the downstream DiT backbone is unchanged [§3.2, Fig. 6].
MDLM-Prime (L=5) achieves OpenWebText perplexity 15.36, beating AR (17.54), EDLM-coAR / BD3-LM hybrids (17.58), and the baseline MDLM (22.98) — the first MDM-only approach to surpass ARM without any autoregressive formulation [Table 2].
The idle-step ratio (ISR) drops from 36.77% at L=1 to 0.25% at L=5; the paper proposes ISR as an a-priori indicator of likelihood-modeling capacity and offers a heuristic for picking $L$ in Appendix A.2.4 [§4.1, Table 2, Prop. A.1].
Zero-shot generalization from OWT training: MDLM-Prime beats AR on LAMBADA (24.44 vs 51.28), PTB (53.77 vs 82.05), WikiText (23.86 vs 25.75), and ArXiv (24.71 vs 41.73); underperforms AR on AG News [Table 1].
On images, MDM-Prime hits FID 3.26 on CIFAR-10 (matches StyleGAN+ADA, beats DDPM 3.17 marginally and prior discrete generators like Discrete-FM 3.63) and FID 6.98 on ImageNet-32, the first discrete generator to beat continuous DDPM/ScoreFlow at that resolution [Tables 3-4].
Sub-token positions in the image domain carry semantic meaning — the first sub-token of a base-encoded pixel determines the high bits of intensity — enabling controllable imputation by keeping only the first sub-token of a conditioning image and letting MDM-Prime fill the rest [§4.2, Fig. 8].
Lower ISR correlates with reduced sensitivity to NFE during sampling: high-L MDM-Prime models keep similar FID between NFE=128 and NFE=1024 while the L=1 baseline degrades sharply [§4.2, Fig. 7].

Method

The forward process treats each data sequence $x \in V^N$ as a sub-token sequence $\tilde x \in \tilde V^{NL}$ via an injective encoding $h: V \to \tilde V^L$ (e.g. base-2 binary encoding when $L = \log_2 C$ ). Sub-tokens are masked independently under the standard MDM Bernoulli schedule $\alpha_t$ , so each original token transits a state tree with $L+1$ levels of partial-mask states between fully-masked and fully-unmasked. The reverse process unmasks one sub-token at a time, which is why a 4-class token can be resolved as two binary decisions (mm → m1 → 11) rather than one 4-way classification.

The decoder models $p_\theta(\tilde x_0 | \tilde x_t)$ as a joint distribution over the $L$ sub-tokens at each position, with logits explicitly zeroed for invalid combinations (sub-token tuples that don’t correspond to any $v \in V$ under $h$ ). To preserve the carry-over parameterization (set unmasked output logits equal to the input), the joint must put probability mass only on sub-token tuples whose unmasked entries match the input — implemented as a precomputed boolean filter applied to the final logits. The encoder uses $L$ separate sub-token embeddings of width $d/L$ concatenated into a single $d$ -dim token embedding; the rest of the diffusion transformer (RoPE DiT, following MDLM) is unchanged. Training is the standard MDM ELBO objective applied at the sub-token level; a single hyperparameter $L$ controls the tradeoff between idle-step reduction (favored by large $L$ ) and the difficulty of joint decoding (favored by small $L$ ).

Results

OpenWebText perplexity (1M training steps, GPT-2 tokenizer, modified DiT backbone):

MDLM-Prime L=5: 15.36 PPL, ISR 0.25%.
MDLM-Prime L=8: 15.48 PPL, ISR 0.03%.
MDLM-Prime L=4: 15.62, L=3: 16.36, L=2: 17.90.
ARM (sahoo2024simplifieddiff): 17.54.
EDLM-coAR (xu2025ebmdiff): 17.58.
BD3-LM (arriola2025block): 20.73.
SEDD (lou2024sedd): 24.10.
MDLM baseline: 22.98 PPL, ISR 36.77% [Table 2].

Zero-shot PPL on OWT-trained checkpoints across 7 evaluation corpora — MDLM-Prime L=4 generally wins; the headline gap is LAMBADA where it scores 24.44 vs ARM 51.28 [Table 1].

CIFAR-10 FID: MDM-Prime 3.26 (NFE=512), vs Discrete-FM 3.63 (NFE=1024), Tau-LDR 3.74, baseline MDM 4.66, DDPM 3.17, StyleGAN+ADA 3.26 [Table 3]. ImageNet-32 FID: MDM-Prime 6.98 (NFE=1024), vs baseline MDM 7.91, Continuous-FM 5.02, ScoreFlow-VP 8.34, DDPM 16.18 [Table 4].

Why it’s interesting

This sits squarely in the wiki’s ongoing “can DLMs catch AR?” thread but attacks the problem from a representation angle rather than a serving angle. Introspective Diffusion Language Models (I-DLM) and WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference convert an existing AR model into a DLM that serves under causal-attention infrastructure; MDM-Prime instead rewrites the diffusion process itself to eliminate idle steps and trains from scratch — and it’s the first paper on the wiki to beat the AR teacher on perplexity at the same architecture/data, where the serving-recipe papers settle for parity. The L>1 sub-token decomposition is also a different lever than The Diffusion Duality‘s uniform-state-via-Gaussian-projection trick: both make the discrete reverse process more expressive, but Duo gets there by lifting to continuous space and consistency-distilling, while MDM-Prime stays purely discrete and gets the win from intermediate states with no distillation. Worth tracking against The Design Space of Tri-Modal Masked Diffusion Models (Apple’s tri-modal absorbing MDM, which trades sampling speed for the Chinchilla-better data-efficiency law) — MDM-Prime targets the same absorbing-state weakness (idle steps) and addresses it without giving up the data-efficiency story. Open follow-up: how does MDM-Prime compose with consistency distillation (Duo) or with the I-DLM serving trick — does the partial-mask state tree survive aggressive few-step distillation?