PaperlayerView exact revision on arXiv

Source-linked AI summary

Why Masking Diffusion Works: Condition on the Jump Schedule for Improved Discrete Diffusion

Alan N. Amin, Nate Gruver, Andrew Gordon Wilson

arXiv:2506.08316v2cs.LGstat.ML

TL;DR

Masking diffusion often outperforms discrete diffusion models with structured forward processes, despite those processes encoding data-specific inductive biases. This paper explains masking’s advantage through its built-in transition schedule and introduces SCUD, which conditions on that schedule. SCUD versions of structured processes outperform masking on images and proteins.

  • Problem

    Discrete diffusion models with structured forward processes have been outperformed by masking diffusion, leaving unclear whether those processes’ inductive biases help model fit.

    Source: §1 · p. 1, §6.2 · p. 7

  • Method

    SCUD conditions the backward process on the event schedule and learns which previous state to predict at each event.

    Source: §1 · p. 1, §4.1 · p. 5

  • Results

    SCUD with structured forward processes outperforms masking and achieves state-of-the-art diffusion-model fits on image and protein data.

    Source: §1 · p. 2, §6.2 · p. 7, §7 · p. 9

  • Takeaways & Limitations

    The paper concludes that incorporating schedule information lets diffusion models use structured forward processes while outperforming masking.

    Source: §7 · p. 9

    The choice of schedule information S remains a design tradeoff, including how to represent multiple mutation types and what information the de-noiser should receive.

    Source: §7 · p. 9

Abstract

from arXiv · show

Discrete diffusion models, like continuous diffusion models, generate high-quality samples by gradually undoing noise applied to datapoints with a Markov process. Gradual generation in theory comes with many conceptual benefits; for example, inductive biases can be incorporated into the noising Markov process, and access to improved sampling algorithms. In practice, however, the consistently best performing discrete diffusion model is, surprisingly, masking diffusion, which does not denoise gradually. Here we explain the superior performance of masking diffusion by noting that it makes use of a fundamental difference between continuous and discrete Markov processes: discrete Markov processes evolve by discontinuous jumps at a fixed rate and, unlike other discrete diffusion models, masking diffusion builds in the known distribution of jump times and only learns where to jump to. We show that we can similarly bake in the known distribution of jump times into any discrete diffusion model. The resulting models - schedule-conditioned discrete diffusion (SCUD) - generalize classical discrete diffusion and masking diffusion. By applying SCUD to models with noising processes that incorporate inductive biases on images, text, and protein data, we build models that outperform masking.

1 Introduction

Masking diffusion outperforms more structured forward processes despite its simple uniform corruption. The paper attributes its advantage to matching the forward transition schedule in the backward process, then applies that principle to structured processes that outperform masking on images and proteins.

1 Introduction

  • Masking diffusion’s superior performance is attributed to matching the known forward transition schedule, rather than establishing that masking is the optimal forward process.Masking uses a uniform-rate process that transforms tokens into a masking token; it had outperformed proposed gradual and state-dependent processes [2] [23].Source: §1 · p. 1
  • SCUD builds the forward transition-time distribution into the backward process, so the model learns where to transition rather than both when and where.The figure also reports that conditioning on more schedule information improves likelihood on CIFAR-10 with B = 128 states.Source: §1 · p. 1, Figure 1 · p. 2
  • Applying SCUD to structured, state-dependent processes yields models that outperform masking and achieve state-of-the-art results on images and proteins.With a uniform forward process, applying SCUD recovers masking diffusion.Source: §1 · p. 2

2 Background

Discrete diffusion learns a backward Markov process that maps noise toward data while fitting its trajectories to a chosen forward process. Forward-process choices include uniform transitions, masking, and processes designed to incorporate data structure.

2 Background

  • Discrete diffusion starts from an easy-to-sample noise distribution and learns a backward Markov process whose trajectories match those of a forward process from data toward noise.The objective is an evidence lower bound (ELBO), maximized when forward and backward trajectory distributions match.Source: §2 · p. 2
  • A discrete Markov process is specified by a B×B infinitesimal generator, whose off-diagonal entries give transition rates and whose diagonal entries give total rates out of each state.The transition time is sampled from an exponential distribution determined by the outgoing rate; the next state is selected according to the transition rates [8].Source: §2 · pp. 2–3
  • State-of-the-art discrete diffusion models have detectable differences between forward and backward transition rates, whereas masking builds the known forward rate into its backward process.Figure 2 plots the estimated backward transition rate minus the forward rate.Source: Figure 2 · p. 3
  • Uniform processes transition to any state at a constant rate, masking processes transition to a masking state at a constant rate, and other processes can encode data-specific inductive biases.The cited examples of structured processes target text, images, and proteins [2].Source: §2 · p. 3

3 Learning when and where to transition

Discrete diffusion must match both when transitions occur and where they go, but published models can mismatch the known transition schedule. The paper proposes matching that schedule explicitly, which explains masking’s advantage and lets structured forward processes improve model fit.

3 Learning when and where to transition

  • Published state-of-the-art diffusion models show detectable forward–backward transition-rate differences, although the forward transition-time distribution is often known.The backward process should match the forward process in both when transitions occur and where they go.Source: §3 · p. 3
  • The ELBO can be decomposed into a term for matching transition destinations given the schedule and a KL term for matching schedule distributions; setting q(S) = p(S) removes the latter.The schedule S is the set of transition times, and the remaining learned task is where to transition.Source: §3 · pp. 3–4
  • Schedule conditioning lets structured forward processes outperform masking by incorporating the transition-schedule information that masking already builds in.The paper reports improved model fits on image and protein data when schedule conditioning is combined with structured forward processes.Source: §3 · p. 4

4 Scheduled conditioned diffusion (SCUD)

SCUD conditions the backward process on event schedules, allowing it to learn predecessor states while using the forward process’s schedule distribution. For high-dimensional data, the method uses per-dimension schedules and a factorized predictor.

4 Scheduled conditioned diffusion (SCUD)

  • SCUD incorporates information about when transitions occur into discrete diffusion, conditioning the backward model on an event schedule.The method introduces constant-rate latent events that often produce transitions, allowing schedules to be used even when transition times themselves have a complex distribution.Source: §4 · p. 4

4.1 Conditioning on event schedules

  • Given the endpoint and sampled event schedule, SCUD predicts the previous state at each event using the current state and the number of events so far.The objective is minimized when the model’s conditional predecessor distribution matches the forward process’s conditional distribution.Source: §4.1 · p. 5

4.2 High dimensional data

  • For high-dimensional data, SCUD simulates a separate schedule for each dimension and factorizes its predictor over dimensions.The per-dimension event counts form a vector, and the resulting objective sums over dimensions.Source: §4.2 · pp. 5–6

5 Connection SCUD to classical and masking discrete diffusion

SCUD uses γ to control how much transition-schedule information the model receives, connecting classical discrete diffusion and masking diffusion as limiting cases.

5 Connection SCUD to classical and masking discrete diffusion

  • γ sets the event rate relative to its minimum, controlling how much schedule information is incorporated into SCUD.At γ = 1 the event rate is slowest; as γ approaches 0 it grows without bound.Source: §5 · p. 6

5.1 Connection to masking diffusion

  • With uniform L and γ = 1 − 1/D, SCUD is equivalent to masking diffusion because the event schedule identifies which token positions are masked.Integrating out the schedule conditional on the mask gives the masking diffusion objective.Source: §5 · p. 6, §5.1 · p. 6

5.2 Connection to classical discrete diffusion

  • As γ approaches 0, the SCUD objective converges to the SEDD loss, recovering classical discrete diffusion.In this limit, the model predicts the infinitesimal change at time t; the loss is also equivalent to that of τLDR.Source: §5.2 · p. 7

6 Results

SCUD combines transition-schedule information with forward-process structure, improving model fit on images and proteins; a language case study finds a smaller gain. The results also show that uniform SCUD interpolates between classical and masking diffusion, while a sparse construction makes a structured language process practical to evaluate.

6.1 Connection to other models

  • On CIFAR-10, uniform SCUD interpolates between uniform discrete diffusion and masking diffusion as γ sweeps from 0.1 to 1.The models use pixels with B = 128 values.Source: §6.1 · p. 7

6.2 SCUD allows models to leverage structure in the forward process

  • SCUD with structured forward processes achieves state-of-the-art model fits on images and proteins, outperforming masking and classical structured diffusion.Figure 4 compares model fit in bits per dimension on CIFAR-10 and perplexity on UniRef50; it labels uniform and structured processes.Source: Figure 4 · p. 8, §6.2 · p. 8
  • For CIFAR-10, SCUD with a Gaussian forward process beats both classical Gaussian diffusion and schedule-conditioned uniform diffusion.The paper attributes classical masking diffusion’s advantage over Gaussian diffusion to schedule conditioning outweighing the inductive-bias benefit.Source: §6.1 · p. 7, §6.2 · p. 8
  • On UniRef50, SCUD with BLOSUM outperforms masking and classical BLOSUM diffusion, achieving state-of-the-art model fit.The BLOSUM process uses rates of amino-acid mutations observed in nature.Source: §6.2 · p. 8

6.3 Scaling case study: Language

  • Classical diffusion and SCUD have similar run times, within 10% of each other, while SCUD can support model decisions that are computationally intractable for classical diffusion.For LM1B’s 30,522-token vocabulary, the classical loss requires computing exp(tL), which is prohibitively expensive; SCUD uses a sparse transition construction.Source: §6.3 · p. 9
  • On language, adding nearest-neighbor structure to SCUD improves fit by a small amount, unlike the substantial improvements reported for images and proteins.The authors suggest the language process could be improved.Source: §6.2 · p. 7, §6.3 · p. 9

7 Discussion

The discussion attributes masking diffusion’s performance to its use of transition-schedule information and argues that SCUD lets structured forward processes outperform masking. It also flags tradeoffs from conditioning on more information and notes that sampling was not evaluated.

7 Discussion

  • SCUD combines schedule information with structured forward processes, enabling models with inductive biases to outperform masking.The paper reports state-of-the-art diffusion model fits on image and protein data.Source: §7 · p. 9
  • Conditioning on more information in S makes inference for x̃0,θ harder and can make the third ELBO term fail to converge.The discussion gives specifying every forward-process mutation as an example.Source: §7 · p. 9
  • The paper improves ELBOs without evaluating how its changes affect sampling.The authors describe sampling research as potentially complementary future work.Source: §7 · p. 10

A Proofs of results

The appendix proves results used in the paper’s discrete-diffusion and SCUD derivations. Its arguments cover Markov-process representation, event schedules, the objective, and a masking parameterization.

A Proofs of results

  • The appendix rewrites the trajectory bound by conditioning on S and reduces the conditional forward transition to the first event.The proof notes that p(x1|x0, S) = p(x1|x0, s1).Source: A Proofs of results · pp. 12–13
  • Any Markov process with generator L can be simulated using L = r(K − I), for r ≥ max_b −L_b,b and K = L/r + I.The appendix identifies I as the identity matrix.Source: A Proofs of results · p. 12
  • Given the event schedule S, the distribution of the preceding state depends on S through the event count at that time.The proof uses the Markov property to reduce dependence on later states and the schedule.Source: A Proofs of results · p. 12
  • The event times on [0,t] form a Poisson process with density μ(s) = rβs, and the proof derives the conditional density of events near t′.The expected event count in a short interval is expressed using the schedule-conditioned event count and βt′.Source: A Proofs of results · p. 14
  • When the denoiser predicts x0, the schedule is useful only to identify corrupted tokens, so the prediction can condition on their indicator mask.The appendix denotes that mask by mt = st > 0.Source: A Proofs of results · p. 14
  • As γ approaches zero, the objective in Eqn. 6 converges to the limiting expression given in the appendix.The appendix states this convergence while suppressing the dependence of s̃θ on xt and t.Source: A Proofs of results · p. 15

B Details of method

The appendix details SCUD’s ELBO estimation, parameterization, sampling, and schedule selection, then describes structured forward processes for images, language, and proteins.

B.1 Algorithm for estimating ELBO

  • SCUD’s ELBO estimate uses an unbiased estimator, with the stationary distribution estimated by spectral decomposition or power iteration when the transition matrix is large.Algorithm 1 summarizes the estimator using Proposition 4.4.Source: §B.1 · pp. 15–16

B.2 Parameterization

  • The parameterization conditions a neural network on the corrupted sample and schedule, predicting quantities used to construct event-level reverse transitions.The posterior factorizes across dimensions, supporting per-dimension predictions.Source: §B.2 · p. 16

B.3 Sampling

  • To reduce evaluations during generation, SCUD reverses multiple events at once, using a k-Gillespie-style procedure and a fixed evaluation budget.The sampler starts from the stationary distribution and samples the event schedule; events are denoised in batches.Source: §B.3 · p. 17

B.4 Choosing the rate

  • SCUD chooses β_t so schedule-conditioned mutual information between clean data and corrupted data decreases linearly over time.The rate is computed with a Newton root finder; as γ approaches zero, the integrated rate converges to the classical discrete-diffusion choice.Source: §B.4 · pp. 17–18

C Structured processes

  • Structured processes can use stationary distributions beyond the uniform distribution, unlike prior constructions that enforced doubly stochastic transition matrices.The processes are designed for continuous-time Markov models.Source: C Structured processes · p. 18

C.1 Gaussian process for images

  • For images, the Gaussian process assigns larger transition rates to pixel values that are closer together.The rate form uses exp(−200(i−j)^2/B), with 200 selected in small-scale experiments.Source: §C.1 · p. 19

C.2 Nearest neighbour process for language

  • For language, a sparse nearest-neighbor process uses embedding similarities among the 1,000 most frequent words, which cover 95% of data tokens.Occasional transitions across the neighbor graph were added because particles often took a long time to reach stationarity.Source: §C.2 · p. 19

C.3 BLOSUM process for protein

  • For proteins, the BLOSUM-based process models transitions between amino acids using related-protein frequencies and restricts transitions from special tokens to canonical amino acids.The process generator is set to L = K − I.Source: §C.3 · p. 19

D Experimental details

The appendix notes the hardware used for training and labels Figure 5 as the BLOSUM process K.

D Experimental details

D.1 Samples from image SCUD

The appendix labels an image-sample figure and describes an ELBO evaluation setup for pretrained protein models.

D.1 Samples from image SCUD

  • Figure 6 presents samples from a SCUD Gaussian model trained on CIFAR-10.Source: Figure 6 · p. 20

D.2 Pretrained protein weights

  • The protein experiment trains 150M-parameter ESM2 models with pretrained UniRef2 weights and compares them with DPLM.The DPLM comparison uses a similar strategy.Source: §D.2 · p. 20
  • DPLM is treated as a 500-timestep discrete-time masking diffusion model to evaluate its ELBO using the formula from Austin et al. [2].Source: §D.2 · p. 20

D.3 Transition rates in Fig. 2

The appendix reports that schedule conditioning improves uniform corruptions and accommodates structured protein processes with pretrained weights. It also details model configurations, computational complexity, and an extension of schedule conditioning to flow matching.

D.3 Transition rates in Fig. 2

  • The transition-rate analysis simulates CIFAR-10 forward trajectories and backward samples for τLDR, while D3PM backward transition probabilities are converted to rates.For τLDR, 2,000,000 forward trajectories and 1,000 backward samples are used; D3PM uses 1,000 backward samples.Source: §D.3 · pp. 20–21
  • Schedule conditioning improves uniform corruptions and accommodates structured protein forward processes with a pretrained ESM2 backbone.Table 2 describes the protein experiment as replacing the main-text architecture with pretrained ESM2.Source: Table 2 · p. 21

D.4 Images

  • Image models replace time conditioning with FiLM conditioning on the number of events, and the large models train on 2.6 × 10^9 images.The large models use a 32-layer UNet; the Figure 1 models use an 8-layer UNet and 2,048 function evaluations for image generation.Source: §D.4 · p. 21

D.5 Language

  • The language models use a diffusion transformer with FiLM layers conditioned on the schedule and train on 1.1 × 10^10 tokens for two days.Source: §D.5 · p. 21

D.6 Protein

  • Protein models use the small CARP architecture with schedule-conditioned FiLM layers and train for two days on 1.1 × 10^11 tokens.The cited D3PM baseline used 3.3 × 10^11 tokens.Source: §D.6 · p. 21

D.7 Computational complexity

  • SCUD showed no substantial computation difference in the experiments, and sparse or low-rank K can make matrix products practical for large vocabularies.The stated big-O costs of the neural predictions are identical; replacing time with the schedule adds negligible overhead for the chosen architectures.Source: §D.7 · p. 22

E Extending flow matching to SCUD

  • Schedule conditioning extends to flow matching, where training a SCUD model can provide samples from a large class of schedule-conditioned flow-matching models.Source: E Extending flow matching to SCUD · p. 22

E.1 Schedule conditioned flow matching (SCUM)

  • SCUM trains on samples corrupted by s events and can select the event-transition distribution at test time without retraining the denoising model.Its reverse step averages the transition matrix over the estimated clean-data posterior.Source: §E.1 · pp. 22–23

E.2 SCUD is SCUM

SCUD is a special case of SCUM for a particular transition-kernel choice. The section also describes alternative sampling kernels and stochasticity choices available after training.

E.2 SCUD is SCUM

  • With a particular choice of K, SCUM’s simulated trajectories match SCUD, whose event-denoising distribution is the one approximated in Algorithm 2.The authors identify SCUD as SCUM with that choice of K.Source: §E.2 · pp. 23–24
  • After training SCUD’s predictor of x0 from (xs, s), the transition kernel Ks|x0 can be changed at test time.This allows the trained predictor to be used with different SCUM sampling choices.Source: §E.2 · p. 24

E.3 Examples of SCUM

  • Alternative SCUM kernels must preserve SCUD’s transition marginals; one option moves excess mass toward deficient states, following Campbell et al. [4].The kernel must have positive entries and rows summing to 1; the mass-transfer choice moves mass in proportion to the deficits.Source: §E.3 · p. 24
  • A stationary detailed-balance generator can add stochasticity, and varying η lets sampling stochasticity be optimized against likelihood.After training, η can be optimized to obtain desirable samples.Source: §E.3 · p. 24
  • In the high-dimensional case, SCUD predicts each dimension’s x0 conditional, and predictions conditioned on the other noisy dimensions can be transformed into it.The paper describes this difference as a matter of parameterization, not a dependence on the specific choice of Ks|x0.Source: §E.3 · p. 24
Loading 2506.08316v2…