PaperlayerView exact revision on arXiv

Source-linked AI summary

Simple and Effective Masked Diffusion Language Models

Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, Volodymyr Kuleshov

arXiv:2406.07524v2cs.CLcs.AIcs.LG

TL;DR

Language modeling diffusion methods have lagged behind autoregressive models in log-likelihood. This paper improves simple masked diffusion with effective engineering and a Rao-Blackwellized objective, yielding state-of-the-art diffusion results that approach autoregressive perplexity while enabling efficient generation from encoder-only models.

  • Problem

    Discrete diffusion methods exhibit a performance gap relative to autoregressive models in language modeling, especially in log-likelihood.

  • Method

    The paper combines a well-engineered masked diffusion implementation with substitution-based reverse modeling, a Rao-Blackwellized objective, and efficient samplers.

  • Results

    Masked diffusion models achieve a new state-of-the-art among diffusion models on language-modeling benchmarks and approach autoregressive perplexity within 15-25%.

  • Takeaways & Limitations

    The framework gives encoder-only BERT-style models principled generative capabilities and supports semi-autoregressive sampling.

  • Takeaways & Limitations

    The language-modeling formulation assumes that the forward noising process is applied independently across a sequence, conditioned on sequence latents.

Abstract

from arXiv · show

While diffusion models excel at generating high-quality images, prior work reports a significant performance gap between diffusion and autoregressive (AR) methods in language modeling. In this work, we show that simple masked discrete diffusion is more performant than previously thought. We apply an effective training recipe that improves the performance of masked diffusion models and derive a simplified, Rao-Blackwellized objective that results in additional improvements. Our objective has a simple form -- it is a mixture of classical masked language modeling losses -- and can be used to train encoder-only language models that admit efficient samplers, including ones that can generate arbitrary lengths of text semi-autoregressively like a traditional language model. On language modeling benchmarks, a range of masked diffusion models trained with modern engineering practices achieves a new state-of-the-art among diffusion models, and approaches AR perplexity. We provide the code, along with a blog post and video tutorial on the project page: https://s-sahoo.com/mdlm

1 Introduction

The paper argues that masked diffusion language models can narrow the language-modeling gap with autoregressive methods through effective engineering, a simplified objective, and efficient samplers. The resulting models achieve strong diffusion-model benchmark performance while supporting generative capabilities in encoder-only architectures.

  • Motivation: Prior discrete diffusion methods show a sizable log-likelihood gap relative to autoregressive models in language modeling.The gap is reported when parameter count is controlled.
  • Method: MDLM combines a well-engineered implementation with a substitution-based reverse-process parameterization and a Rao-Blackwellized variational objective.The objective is a weighted average of masked language-modeling losses with improved ELBO tightness and variance.
  • Method: The objective enables BERT-style encoder-only models to support principled generation through efficient samplers, including semi-autoregressive generation.The figure caption also describes ancestral sampling for MDLM.
  • Results: MDLM achieves a new state-of-the-art among diffusion models across language-modeling benchmarks and approaches autoregressive perplexity within 15-25%.The contribution statement names LM1B, OWT, and DNA among the evaluated benchmarks.
  • Results: Simple engineering choices also substantially improve previously weak baselines, while the framework extends to biological sequence modeling.For DNA models, the paper reports similar or higher downstream performance than classical BERT-style training alongside generative capabilities.

2 Background

Diffusion models learn to reverse a forward corruption process, with the reverse model trained through a variational lower bound on data log-likelihood. For discrete data, this framework can operate directly on categorical Markov processes and their continuous-time extensions.

  • Diffusion framework: Diffusion training reverses a forward process that transforms clean data into progressively noisy latent variables.The forward process is denoted q, while z_t represents the noisy version of x at time t.
  • Diffusion framework: The parameterized reverse model p_θ is trained to maximize an evidence lower bound on log-likelihood.With discretization steps, the negative ELBO can be expressed using KL divergences.
  • Discrete diffusion: Discrete diffusion can either apply Gaussian diffusion to continuous embeddings or define a diffusion process directly on discrete structures.The direct discrete approach is represented by categorical transition probabilities.
  • Discrete diffusion: D3PM defines a discrete Markov forward process whose marginals are obtained by multiplying categorical transition matrices across time steps.The process uses q(z_t|z_{t−1})=Cat(z_t;Q_t z_{t−1}) and corresponding cumulative transitions.
  • Discrete diffusion: Continuous-time extensions of discrete diffusion use continuous-time Markov chain theory and connect score-based modeling with ELBO maximization.SEDD is identified as a likelihood-based score-modeling approach enabled by this connection.

3 Simple Masked Diffusion Models

MDLM focuses on absorbing-state masking diffusion and combines a SUBS reverse-process parameterization with Rao-Blackwellized objectives and engineering choices. The resulting objective is a weighted average of MLM losses, while the framework supports encoder-only models and principled generation.

  • Absorbing-state masking diffusion is selected because it has consistently achieved the best performance among prior discrete diffusion processes.
  • The Rao-Blackwellized objective analytically exploits these substitutions, reducing training variance and improving ELBO tightness.
  • The forward process gradually transfers probability mass from clean tokens to a mask state, which remains absorbing once reached.
  • SUBS enforces zero probability for masking outputs and carries already-unmasked tokens unchanged through reverse diffusion.
  • The continuous-time extension yields a diffusion loss invariant to the functional form of the noise schedule.
  • For sequences, the objective becomes a weighted average of MLM losses, connecting diffusion training with BERT-style encoder models and principled generation.
  • A well-engineered implementation, including careful tokenization, substantially improves masked diffusion performance and previously weak baselines.

4 Inference and Sampling in Masked Diffusion Language Models

MDLM generates sequences by discretizing reverse diffusion and can skip redundant denoising steps when no tokens become unmasked. Its semi-autoregressive sampler extends sequences by reusing generated tokens as prefixes.

  • The reverse sampler starts from a masked sequence and iteratively constructs token sequences by discretizing the reverse diffusion process.
  • Unmasked tokens remain unchanged during reverse diffusion, allowing the process to preserve completed positions.
  • When a denoising step reveals no new tokens and the network is time-independent, MDLM can skip that step and save a denoising-network call.
  • The semi-autoregressive sampler generates arbitrary-length text by using previously generated tokens as a prefix for another reverse-diffusion round.

5 Experiments

Experiments show that MDLM performs strongly as both a generative and representation model across language and biological-sequence benchmarks. It improves diffusion-model likelihoods, supports faster semi-autoregressive generation, preserves downstream performance, and benefits from specific implementation choices.

  • 5.1 Masked Diffusion Language Models: MDLM outperforms previous diffusion methods on LM1B and OWT while narrowing the perplexity gap with autoregressive models.On LM1B, it improves the SEDD perplexity bound by 17% and comes within 14% of an AR baseline; the OWT trend is similar.
  • 5.1 Masked Diffusion Language Models: MDLM consistently outperforms the SEDD parameterization in zero-shot evaluation and sometimes achieves better perplexity than AR models.The evaluated unseen datasets include PTB, Wikitext, LM1B, Lambada, AG News, and Scientific Papers.
  • 5.1 Masked Diffusion Language Models: BERT fine-tuned with MDLM gains generative capability without degrading GLUE performance.On C4, the MDLM fine-tuned model reaches a PPL upper bound of 35 versus 78 for pretrained BERT, while downstream GLUE performance has no degradation.
  • 5.1 Masked Diffusion Language Models: MDLM generates 2048-token sequences with better generative perplexity and approximately 25-30x faster decoding than SSD-LM.The comparison uses 200 sequences on a single 3090 GPU, with MDLM generated using L′ =512.
  • 5.2 Masked Diffusion DNA Models: On DNA modeling, MDLM outperforms other diffusion schemes and preserves or improves downstream genomic benchmark performance after generative fine-tuning.The genomic evaluation covers eight regulatory-element classification tasks; perplexity and downstream-performance correlation is weaker for biased genome subsamples.
  • 5.3 Ablation Analysis: Ablations show that streamlined implementation choices matter: removing carry-over increases perplexity by 1.5 points, while discrete T =1000 worsens it by 0.1.Further removing zero masking produces no meaningful perplexity change.

6 Related Work

The paper positions masked diffusion as a simpler discrete alternative to broader diffusion and score-based frameworks, while adding principled generation to encoder-only models. It distinguishes its approach through efficient arbitrary-length sampling and biological-sequence modeling.

  • Comparison to D3PM: Masked diffusion is a strict subset of D3PM, specializing the forward process to masking rather than general discrete noise.The paper attributes improvements over D3PM to SUBS, objective simplification, and training recipes, with engineering contributing the most.
  • Comparison to Score Estimation: Unlike score-based approaches that typically rely on advanced CTMC theory, MDLM uses a simpler formulation while remaining compatible with score-based techniques.The paper also extracts an MDLM score for use with score-oriented samplers, parameterizations, denoising designs, and guidance.
  • Comparison to BERT: MDLM provides a principled way to make BERT-style encoders generative through randomized masking rates, rather than Gibbs sampling or ad-hoc methods.Its objective connects masked diffusion with masked language modeling while retaining encoder-style representation learning.
  • Comparison to Latent Diffusion LMs: MDLM derives a tight, principled lower bound on log-likelihood, unlike the cross-entropy objective used by Discrete Flow Matching.The comparison concerns discrete structures rather than Gaussian diffusion over word embeddings.
  • Concurrent Works: MDLM differs from concurrent masked-diffusion work by supporting semi-autoregressive decoding and arbitrary-length text generation.The cited comparison states that concurrent methods sample fixed-length sequences and evaluate only on NLP datasets, whereas this work also models biological sequences.

7 Conclusion

The paper concludes that a well-engineered masked diffusion implementation with a simple variational objective reaches state-of-the-art diffusion perplexities and converts BERT-style encoders into generative models. It also rigorously derives a related continuous-time ELBO and notes that language modeling retains the field’s inherent risks and opportunities.

  • 7 Conclusion: A well-engineered masked diffusion implementation with a simple variational objective attains state-of-the-art diffusion perplexities on language benchmarks.The conclusion frames this result as part of an exploration of masked diffusion.
  • 7 Conclusion: The framework provides an efficient way to convert BERT-style encoders into generative models.This capability is presented alongside the masked diffusion likelihood results.
  • 7 Conclusion: The paper presents a rigorous CTMC-based derivation for an ELBO previously described without an explicit derivation.It further shows that the formula reduces to the proposed continuous-time ELBO under SUBS parameterization.

Appendices

The appendices derive masked-diffusion posteriors, objectives, and continuous-time formulations, then simplify the NELBO through SUBS and Rao-Blackwellization. They also describe the training algorithm and establish that carry-over unmasking improves the likelihood estimate.

  • Posterior and forward-process derivation: The appendix derives true and approximate posteriors for absorbing-state diffusion before simplifying them for the masked case.The forward process transfers part of the probability mass to the masking prior during diffusion.
  • Reverse-process derivation: The reverse-process derivation separates unmasked and masked states, yielding categorical distributions supported only on the original token and the mask.The appendix explicitly analyzes the cases z_t=x and z_t=m.
  • Rao-Blackwellized objective: The diffusion loss is zero when z_t is unmasked, and the simplified NELBO is obtained by applying zero-masking and carry-over-unmasking Rao-Blackwellization.Carry-over unmasking removes the masking indicator because the model reconstructs x exactly for unmasked inputs.
  • Rao-Blackwellized objective: Carry-over unmasking makes the reconstruction term zero for unmasked inputs, producing a better likelihood estimate than the preceding simplified expression.The paper says this improvement is empirically verified in Table 8.
  • Training algorithm: Algorithm 1 trains MDLM by sampling sentences and time steps, corrupting tokens, and taking gradient steps on the resulting objective.The supplied algorithm excerpt identifies sentence sampling, uniform time sampling, and gradient descent as training steps.
  • Continuous-time formulation: The continuous-time formulation is derived by taking the limit of the discrete-time process as T →∞ and relating transition and rate matrices.Under SUBS parameterization, the resulting formula reduces to MDLM’s simplified continuous-time NELBO.

D.7 Diffusion DNA Models

The DNA experiments adapt masked diffusion to biological sequences and assess both training choices and sampling behavior. The appendix reports schedule-invariant likelihoods with variance differences, while noting that genomic benchmark distributions may not match pretraining data.

  • Dataset: The DNA experiments pre-train Caduceus MLM on the HG38 human reference genome using character- or base-pair-level tokenization.The training split contains 35 billion tokens, with segments extended to a maximum length of 1,048,576.
  • Architecture: Caduceus uses a bidirectional Mamba-based state-space backbone designed to preserve reverse-complement equivariance.This architectural bias reflects the symmetry of double-stranded DNA molecules.
  • Training details: Diffusion fine-tuning lowers learning rates for Mamba-based models from 8e-3 to 1e-3 and omits timestep embeddings because they have minimal generative impact.The omission is supported by the paper’s reference to an ablation in Table 12.
  • Benchmark caveat: Downstream genomic benchmark distributions may differ from pretraining and perplexity data, so perplexity and downstream performance may not correlate.The benchmark subset emphasizes human-data tasks such as enhancers, promoters, and open chromatin regions.
  • ELBO invariance: Different noise schedules yield the same likelihood, but the log-linear schedule has the lowest per-data-point variance among those tested.The appendix attributes equal likelihoods to theoretical invariance while distinguishing schedules by variance.
  • Sampling efficiency: Sampling-time comparisons vary diffusion steps T and batch size while measuring wall-clock generation time for 64 samples on one A5000 GPU.The figure compares AR, SEDD, and MDLM with or without caching.

E.3 LM1B ablations

The ablation compares discrete and continuous diffusion-time evaluation for MDLM and examines training-loss behavior on OpenWebText. Continuous-time evaluation is associated with decreasing NLL as T approaches infinity, while MDLM shows lower training-loss variance than SEDD.

  • NLL consistently decreases as T →∞ in MDLM evaluations.The comparison evaluates discrete and continuous diffusion steps, with MDLM trained using T =∞.
  • MDLM achieves lower variance loss during training than SEDD on OpenWebText.Training spans 1M steps, corresponding to 524B tokens.
  • Figure 3 tracks training NLL across 1M gradient steps on OpenWebText without value smoothing.NLL is logged every 1K steps over 524B tokens.

E.5 Time-conditioning ablation on OWT

The OWT ablation evaluates whether time conditioning affects MDLM performance and examines samples at two diffusion-step settings. Time conditioning has minimal impact on perplexity in this evaluation.

  • Time conditioning has minimal impact on perplexity for MDLM trained on OWT.The ablation is reported in Table 12 after 1M training steps on 524B tokens.
  • The OWT time-conditioning assessment uses training over 1M steps, corresponding to 524B tokens.
  • The study presents unconditional MDLM samples generated with context length L=1024.Samples are shown for T ={1000,10000}.

E.6.1 T = 1000

The T = 1000 examples contain unconditional samples with sports-centered narratives and news-like passages. The outputs include football teams, players, games, playoffs, and a separate hospital-related report.

  • The samples include passages about wins, losses, rankings, injuries, turnovers, and playoff qualification.
  • Several samples generate football narratives involving teams, players, games, and playoff outcomes.Examples mention the Saints, Panthers, Cardinals, South Carolina, and Brees.
  • A separate sample presents a hospital-related news passage involving the Coast Guard and public questions.
  • Other generated text discusses criminal charges, police investigations, university officials, attorneys, and alleged evidence.

E.6.2 T =10000

The T =10000 examples include passages about government fiscal measures and economic management alongside conversational and autobiographical text. The economic passages discuss crisis response, debt, borrowing, demand, and growth.

  • Several samples describe temporary government measures intended to address an economic crisis and fiscal situation.
  • Other passages discuss borrowing costs, debt repayment, revenue budgets, and fiscal consolidation.One passage reports $440 billion in debt issues, while another mentions a $5.2 billion annual budget cost.
  • The economic text connects temporary measures with growth, output, external demand, internal demand, and government credibility.

NeurIPS Paper Checklist

The checklist covers reproducibility, experimental reporting, ethics, societal impacts, and whether the paper’s claims reflect its contributions and scope.

  • The checklist must remain in the submission and is placed after the references and optional supplemental material without counting toward the page limit.
  • Authors should provide concise justifications for each checklist answer, including when answering “No” or “NA.”
  • Authors are asked whether the abstract and introduction accurately reflect the paper’s contributions and scope.
  • Theory Assumptions and Proofs: The checklist asks whether theoretical results include complete assumptions and proofs, or whether the paper has no theoretical results.
  • Experimental Result Reproducibility: Experimental reproducibility requires disclosure of hyperparameters, settings, compute resources, and appropriate error bars or other statistical-significance information.
  • Open access and societal impact: The paper must address open access to data and code, responsible release safeguards, and potential positive and negative societal impacts.
Loading 2406.07524v2…