Source-linked AI summary
Representation-based Masked Diffusion Model
Yangrong Hu, Ding Huang, Xueyu Zhou, Jian Huang
TL;DR
Parallel masked-token sampling can neglect dependencies among positions, a concern for coherent text generation. RMDM conditions masked diffusion on a global semantic representation aligned to a Gaussian prior. It improves generation quality, especially in few-step regimes, and reports higher throughput at matched quality in an A800 evaluation.
Problem
Existing parallel samplers approximate the joint distribution over masked tokens with independent token-wise predictions, despite residual dependencies among them.
Source: §1 · p. 1, §3.1 · p. 3Method
RMDM uses pretrained encoder representations aligned to a Gaussian prior as a global condition for a masked diffusion model.
Source: §1 · p. 2, §3.2 · pp. 4–5Results
Takeaways & Limitations
RMDM’s shared global semantic information couples masked-token updates and supports simultaneous generation, particularly when sampling uses few steps.
Source: §5 · pp. 8–9Experiments focus on unconditional generation at GPT-2-comparable model scales and computational budgets; transfer to substantially larger models and conditional generation tasks remains unexamined.
Source: Limitations · p. 9
Abstract
from arXiv · showhide
Masked Diffusion Models (MDMs) have emerged as a compelling paradigm for language modeling, offering the capability for efficient parallel text generation. However, existing parallel sampling methods typically update multiple masked tokens independently and ignore the complex mutual dependencies among the masked tokens. This independent updating mechanism lacks global coordination and might lead to incoherent outputs. To address this limitation, we propose Representation-based Masked Diffusion Model (RMDM), a framework that leverages the text representation to explicitly encode global semantics and help to parallel update tokens more precisely. Specifically, we first encode text into a continuous semantic space using a pretrained encoder and learn an invertible transformation that normalizes the representation distribution to a Gaussian prior, facilitating efficient sampling during generation. Conditioned on this latent semantic representation, we train a masked diffusion model to learn the conditional text distribution, where the representation serves as global semantic guidance to coordinate parallel token updates and faithfully approximate the target distribution. Empirical results demonstrate that RMDM significantly improves generation quality, particularly in aggressive few-step sampling regimes.
1 Introduction
Independent parallel sampling can miss dependencies among masked tokens, while RMDM uses a sequence-level semantic representation to coordinate their updates. It aligns pretrained representations with a Gaussian prior so a latent can be sampled at generation time.
1 Introduction
- Independent sampling treats masked-token positions as conditionally independent, although natural language can retain dependencies from agreement, long-range constraints, multi-word expressions, and entity consistency.The resulting sequences can be locally plausible yet globally inconsistent.Source: §1 · p. 1
- Figure 1 depicts Stage I aligning encoder-derived latents with a Gaussian prior and Stage II training RMDM to reconstruct corrupted text conditioned on the latent.The latent is sampled from z ∼ N(0, I) at generation time.Source: Figure 1 · p. 2, §1 · p. 2
- RMDM uses a global latent representation of sequence semantics to reduce residual dependence and make parallel prediction closer to the true joint conditional distribution.Pretrained text representations provide this global information; MeanFlow maps their empirical distribution to a Gaussian prior for sampling at generation time.Source: §1 · p. 2
- RMDM improves sample quality over standard MDMs at a fixed number of sampling steps and achieves approximately 3.6× higher measured throughput at matched generation quality in the A800 evaluation.The pretrained encoder and MeanFlow network construct training latents but are not invoked during generation.Source: §1 · p. 2
2 Preliminaries
The preliminaries describe continuous diffusion and MeanFlow as background for the paper’s two-stage latent construction, then introduce masked diffusion as a discrete denoising process. They also illustrate how a global latent can account for masked-token correlations.
2.1 Continuous Diffusion and MeanFlow
- Continuous diffusion corrupts data toward Gaussian noise, while flow matching learns a velocity field and MeanFlow models average velocity to enable one-step recovery.For MeanFlow, z0 can be recovered from noise z1 in one step; its training identity uses the stop-gradient operator sg.Source: §2.1 · pp. 2–3
2.2 Masked Diffusion Models
- Masked diffusion progressively replaces clean tokens with an absorbing [MASK] state according to a decreasing noise schedule, and the denoiser learns to invert this corruption.The denoising model is trained from a data-log-likelihood-derived objective to predict categorical distributions at masked positions.Source: §2.2 · p. 3
- Figure 2 frames the global latent z as a common cause for masked tokens, with conditioning on z explaining away their correlation.Source: Figure 2 · p. 3
3 Methodology
RMDM addresses the mismatch between token-wise training and parallel sampling by conditioning on a global semantic latent. Its framework aligns pretrained representations to a Gaussian prior, trains a conditional masked diffusion model, and decomposes generation error into three components.
3.1 The Conditional Dependency Gap
- Independent token-wise reconstruction and factorized parallel sampling leave dependencies among masked tokens unmodeled, creating a Conditional Dependency Gap.The gap vanishes exactly when masked tokens are conditionally independent given the context and latent, making independent conditional draws equivalent to sampling the joint block.Source: §3.1 · pp. 3–4
3.2 Representation-based Masked Diffusion Model
- A global latent z captures shared sequence information, reducing residual cross-token dependence and bringing parallel prediction closer to the true joint conditional distribution.RMDM obtains z from a pretrained encoder and uses MeanFlow to align representations with a Gaussian prior that can be sampled at inference.Source: §3.1 · p. 3, §3.2 · pp. 4–5
3.3 Theoretical Analysis: Decomposing the Generation Gap
- The generation-error decomposition identifies prior mismatch, remaining token dependencies, and token-wise prediction error as the three sources addressed by RMDM.The framework targets these terms through MeanFlow alignment, global semantic representations, and masked diffusion training, respectively.Source: §3.3 · pp. 5–6
3.4 Model Architecture
- Both MeanFlow and the masked diffusion model use a Diffusion Transformer, with z injected into every block through Adaptive Layer Normalization.Latent bottleneck projection and injected Gaussian noise are used to reduce overfitting to training latents and improve generalizability.Source: §3.4 · pp. 6–7
4 Experiments
RMDM improves generation quality over diffusion baselines, with larger GenPPL gains when sampling steps are scarce, and reaches matched quality at higher throughput. Evaluation covers unconditional generation on LM1B and OpenWebText, zero-shot perplexity, and an ablation of the latent’s source.
4.1 Experimental Setup
- Evaluation follows unconditional-generation protocols on LM1B and OpenWebText, with zero-shot perplexity and generation-quality assessments across additional test sets.Generation quality uses GPT-2 GenPPL and two independent LLM judges; GenPPL may not fully reflect human-perceived quality.Source: §4.1 · p. 7
4.2 Main Results
- At matched-quality operating points on one A800 GPU (batch size 8, L = 1024), RMDM achieves approximately 3.6× higher throughput than MDLM, despite being about 10% slower at equal step counts.Source: Table 4 · p. 8, §4.2 · p. 8
- RMDM consistently improves GenPPL over MDLM and SEDD, with the gap widening as the sampling budget decreases; both LLM judges rank it strongest across budgets.The independent second judge shows the same overall ranking trend as the primary judge.Source: §4.2 · p. 8
- After training on OpenWebText, RMDM improves zero-shot perplexity over diffusion baselines across diverse test sets.The reported zero-shot perplexities for diffusion models are upper bounds.Source: Table 3 · p. 7, §4.2 · p. 8
4.3 Ablation Study
- Replacing the sequence-dependent latent with an independent Gaussian latent performs consistently worse, especially with few sampling steps.The ablation keeps the RMDM interface, conditioning mechanism, and sampling schedule unchanged.Source: §4.3 · p. 8, Table 5 · p. 8
5 Conclusion
RMDM adds a global continuous latent to masked diffusion models to represent sequence-level semantics and restore intertoken dependencies. Experiments show improved generation quality, especially in few-step regimes, and approximately 3.6× higher measured throughput at matched quality.
5 Conclusion
- RMDM uses a MeanFlow-based representation-learning mechanism to align pretrained encoder embeddings with a Gaussian prior, so generation samples the latent once without running the encoder or MeanFlow network.Conditioning on the latent couples masked-token updates through shared global semantic information and reduces the Conditional Dependency Gap.Source: §5 · p. 9
- Experiments on LM1B and OpenWebText show improved generation quality, especially in aggressive few-step regimes, and approximately 3.6× higher measured throughput at matched quality.The throughput comparison is from the paper’s A800 evaluation.Source: §5 · p. 9
Limitations
The experiments focus on unconditional generation at GPT-2-comparable model scales, leaving transfer to substantially larger models and conditional generation for future work.
Limitations
- The evaluation covers unconditional generation at GPT-2-comparable scales; transfer to substantially larger models and conditional generation remains unexamined.The authors leave both extensions for future work.Source: Limitations · p. 9
Ethics Statement
The paper reports no new human-subject research or data collection and evaluates on public benchmarks, while situating RMDM among discrete, continuous, and latent-augmented diffusion work.
Ethics Statement
- The work studies masked-diffusion statistical modeling using public benchmarks, without human subjects, new data collection, or additional training beyond its experimental setup.The authors report no ethical risks beyond those inherent to the underlying language models and benchmark data.Source: Ethics Statement · p. 9
A Related Works
- RMDM samples one aligned latent and keeps it fixed during discrete denoising, unlike CCDD, which jointly evolves continuous representations and discrete tokens [Zhou et al., 2025].Related work also covers transition-based and score-based discrete diffusion, continuous-space language models, and other latent-augmented approaches.Source: A Related Works · p. 13
B.1 Proof of Proposition 3.1
The appendix proves a condition for eliminating residual conditional dependence and decomposes joint divergence into interpretable terms.
B.1 Proof of Proposition 3.1
- The residual conditional dependency gap is zero exactly when the conditional distributions in its definition are equal almost surely, yielding the stated factorization.This follows from KL non-negativity.Source: §B.1 · p. 14
B.2 Proof of Proposition 3.2
- The joint-KL decomposition separates residual conditional dependence from token-wise prediction error under the stated factorization assumptions.The residual-dependence terms include T(xU | z) and T(xM | xU, z).Source: §B.2 · p. 14
C Experimental Details
The experiments use dataset-specific tokenization and pretrained embeddings, with a shared DiT-based training recipe; both tested encoder variants outperform matched baselines across sampling budgets.
C Experimental Details
- LM1B uses BERT-Base embeddings compressed to 128 dimensions, while OWT uses 32-dimensional Qwen3-Embedding-0.6B outputs.LM1B sequences are length 128; OWT sequences are length 1024.Source: C Experimental Details · p. 16
- Both backbones use a 12-layer, 768-hidden-dimension DiT, with 1M training iterations on LM1B and 50K on OWT.Both use batch size 512, AdamW, a 3 × 10^-4 learning rate after 2.5K warm-up iterations, and EMA decay 0.9999.Source: C Experimental Details · p. 16
Encoder sensitivity
- Both Qwen3 and ModernBERT encoder variants outperform their matched baselines across sampling budgets on GenPPL and LLM-judge scores, with smaller gains for ModernBERT.The ModernBERT comparison uses an approximately 60K-step matched MDLM baseline, while the Qwen3 results use 50K-step models.Source: C Experimental Details · p. 16
D LLM-judge Evaluation
The evaluation supplements GenPPL with LLM judges because evaluator perplexity may miss aspects of perceived text quality. Two independent judges assess standalone generated passages, and their relative rankings agree.
D LLM-judge Evaluation
- LLM judging complements GenPPL, which may miss discourse coherence, degenerate repetition, and natural readability.Each generated passage is judged as a standalone sample.Source: D LLM-judge Evaluation · p. 16
- For each method and sampling budget, both judges evaluated all 256 generated passages; Gemini-3.1-Flash-Lite served as an independent judge, scoring five criteria on a 1–10 scale.The criteria were fluency, coherence, repetition, readability, and overall quality.Source: D LLM-judge Evaluation · p. 17
- Tables 8 and 9 report mean overall-quality scores with 95% confidence intervals, and the judges agree in relative rankings.Source: D LLM-judge Evaluation · p. 17, Table 8 · p. 17, Table 9 · p. 18
E Inference Efficiency
RMDM improves throughput at matched quality despite retaining about 90% of MDLM’s throughput at a fixed step count. Efficiency measurements use one A800 GPU and the DDPM sampler.
E Inference Efficiency
- Efficiency was measured after warm-up on one NVIDIA A800 80GB GPU with DDPM, batch size 8, and sequence length 1024.RMDM and MDLM share a DiT backbone; the pretrained encoder and MeanFlow network are not invoked during inference.Source: E Inference Efficiency · p. 18
- 3.6× throughput improvement at GenPPL ≤50: RMDM reaches it in 256 steps (287 tok/s), versus MDLM’s 1024 steps (80 tok/s).At matched quality, RMDM requires roughly one quarter as many steps; reported improvements across the quality range are 3.5–3.6×.Source: Table 10 · p. 18, E Inference Efficiency · p. 18
- At the same step count, RMDM retains approximately 90% of MDLM’s throughput.Source: E Inference Efficiency · p. 18
F Training Cost
The reported RMDM training cost exceeds MDLM’s in the uncached implementation. Online representation extraction contributes to the measured time, and caching was not used because of storage constraints.
F Training Cost
- For 50K training steps on four A800 GPUs, MDLM took approximately 3 days and 5 hours, while RMDM decoder training took approximately 5 days.Source: F Training Cost · p. 18
- The fixed-corpus encoder representations could in principle be cached to remove most repeated computation, but storage constraints prevented caching in the reported runs.Source: F Training Cost · p. 18
- The 10K-step MeanFlow alignment stage took approximately 1 day and 15 hours, including online encoder-representation extraction.The authors characterize these times as implementation costs, not inherent lower bounds.Source: F Training Cost · p. 18
G.1 Detailed Numerical Example with Gaussian Distribution
The section spans a Gaussian example of global-latent conditioning and OpenWebText generations, as well as passages on Bitcoin development, community coordination, an extension, and a workshop about development standards.
G.1 Detailed Numerical Example with Gaussian Distribution
- The Gaussian example follows a causal graph in which global latent z explains correlations between masked tokens that partial context xU cannot capture.With z fixed, only independent noise terms remain, making the conditional independence assumption valid for parallel generation.Source: §G.1 · pp. 18–19
- Without z, partial context xU does not fully recover the latent, leaving masked tokens dependent and creating information loss under independent parallel sampling.The passage identifies the resulting positive dependency gap as mutual information.Source: §G.1 · p. 19
- The residual correlation falls from σ2/(1 + σ2) to 0 when the global latent z is introduced.Conditioning on z makes the masked tokens conditionally independent, so their conditional dependency gap vanishes.Source: §G.1 · p. 19
G.2 Generated Samples on OpenWebText
- The appendix provides OpenWebText generated samples at four sampling budgets, including a 1024-step sample set.The supplied passages provide sample text but do not support a quality comparison across budgets.Source: §G.2 · p. 20, §G.2.1 · p. 20
G.2.1 1024 Sampling Steps
- The Bitcoin community is described as working closely with a development extension intended to make development easier and help members understand Bitcoin’s future.Source: §G.2.4 · p. 22
- The Light Architecture’s development is given as the reason for organizing a Bitcoin Community workshop to discuss Bitcoin Vision development standards.Source: §G.2.4 · p. 22