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

Thank you for your thoughtful and effortful review! We clarify our theoretical and experimental results below. However, we note your appraisal about clarity differs from the other reviewers and your questions concern the foundations of continuous-time Markov chains. We recognize that readers would benefit from some familiarity with continuous-time Markov chains. We'll add additional mathematical background to the appendix and note Chapter 2 of Liggett 2010 and the τLDR paper provide excellent primers in this established area.

Conceptual and Mathematical Concerns Regarding the Transition Matrix $K$ The definition and usage of matrix $K$ (as described in lines 129-136 and Proposition 4.1) raise several conceptual issues:

Once $\Delta t$ is sampled, the first transition occurs at time $t+\Delta t$, so it is unclear why $x_t$ and $x_{t+\Delta t}$ could be the same.

Our definition and usage of the matrix $K$, including self-transitions, is standard (ex Chapter 2 of Liggett 2010). In our first discussion of the Gillespie algorithm we defined $\Delta t$ as the transition times. However, Prop 4.1 shows that we can simulate from a Markov process defined by $\mathcal L$ even if there isn’t a transition at each $\Delta t$. This is important for applying SCUD to processes which have $\mathcal L$ with non-constant diagonals. The beginning of Sec 5 describes how decreasing the probability of $x_t=x_{t+\Delta t}$ leads to more schedule conditioning.

The proof of Proposition 4.1 appears to implicitly assume that $\Delta t$ is infinitesimally small, but this assumption is neither stated nor justified.

Prop 4.1 simply uses the formal definition of the infinitesimal generator $\mathcal L$ (Equation 2.4 of Liggett 2010) – the limit $t\to 0$ in not an assumption about the time between transitions. Note we are working in continuous time, so there is no minimum $\Delta t$ and the markov process can be uniquely described by a matrix $\mathcal L$ which represents the infinitesimal flow at a particular time. The formal connection between $\mathcal L$ and the Markov process is described in depth in chapter 2 of Liggett 2010.

When sampling $x_{t+\Delta t}$ from a categorical distribution parameterized by $K_{x_t’}$, it is unclear why this process would be independent of the sampled $\Delta t$, especially when $\Delta t$ governs the time horizon of transitions.

This is a property common to all Markov processes. The “event” times are drawn from exponential distributions which are “memory-less”, and then $K$ describes the behaviour at each event. The sampling algorithm we describe is known as the Gillespie algorithm.

Incorrect Claim About Equivalence with Uniform Transition: The statement in lines 205–209 is misleading. Even under uniform transition schedules, transition and masking processes are not equivalent, even if the transition schedule is known. This is because under uniform transitions, the same token position can undergo multiple transitions, which breaks the claimed equivalence with masking.

Lines 205-209 give an informal argument and the formal argument is contained in Proposition 5.1.

Intuitively, the reason the “multiple transitions” doesn’t represent a meaningful difference between the two classes of models is that a position that has transitioned uniformly just once has lost all of its information – therefore when trying to predict $x_0$, knowing that a token in $x_t$ has been transitioned once or many times presents the same information.

Notation and Mathematical Clarity Several mathematical definitions and notations are unclear or inconsistent:

In Figure 3, the role of is ambiguous. If this represents a model prediction, is it a denoised sample or a score? In discrete diffusion, the model typically estimates log-ratio scores [1].

In Proposition 4.4, it appears the authors directly model $x_0$. However, prior work [1] emphasizes estimating the log-ratio between adjacent states, which is not discussed here.

Figure 3 represents predicting the “denoised” letter $x_0$. We will make this more clear in Section 4.2. The parameterization we use is the most common (used by continuous diffusion models as well as D3PM, EvoDiff, and masking diffusion models). However we are familiar with the SEDD (your citation [1]) parameterization and, in fact, one of our central results, Proposition 5.2, actually describes the connection to our parameterization; it shows that the two parameterizations are equivalent.

In Equation (5), it is not clear whether the schedule variable $s$ is discrete or continuous. Since it's drawn from an exponential distribution, one would expect it to be continuous.

Line 155 defines $s_t$ as the number of events that have occurred so far, not the event times $S$. It is distributed as a Poisson as described in Algorithm 2 for example. We will make this more clear.

The notation $\tilde x$ is used without definition in Proposition 5.1.

$\tilde x_{0, \theta}$ is the prediction of the denoised sequence. We will restate the definition in Sec 5 for clarity.

Proposition 5.2 lacks context and should clearly describe the connection to SEDD loss. The statement about equivalence needs to be explained more rigorously.

Proposition 5.2 is formally stated and proved in the appendix; there are links in the pdf between the statement in the main text and the proof. However the SEDD loss is cumbersome to write, so we simply moved its statement to the appendix.

The empirical validation could be strengthened by including comparisons against more competitive or recent baselines such as SEDD or LLADA, which are more aligned with the proposed approach in terms of model class and task setup.

Our experiments aimed to confirm our theoretical results in practice and leverage them to interrogate some commonly held beliefs around discrete diffusion models. Nevertheless, while not necessary for proving our main claims, our implementation of our SCUD uniform uses the SEDD architecture and can therefore be interpreted as a comparison to SEDD.

In brief, to compliment our theoretical contributions, our goal was to demonstrate our theoretical results in practice, and leverage the theory to reinvestigate structured processes in the literature. In the case of D3PM’s Gaussian and EvoDiff’s BLOSUM, SCUD leads to state-of-the art results, but for D3PM’s sparse graph process SCUD leads to less of an improvement. We will make our goal clearer in future drafts.

In particular, our goal in the experiments section was to

1. Show that schedule conditioning improves performance. This was validated by comparison to previous sota models in theory, images, and proteins. Technical considerations meant we were limited to uniform $\mathcal L$, where the superiority of masking was already established.

2. Redo previous experiments that showed that structured processes performed worse than masking. Our theory suggested these were unfair – masking had the advantage of schedule conditioning while structured processes did not. We redid these fairly by building SCUD versions of structured models. But we are not necessarily advocating for using these specific structured processes. And we of course do not expect every structured forward process will lead to an improvement in performance.

For images and proteins, we saw that Gaussian and BLOSUM processes improve the fit to the data, achieving state-of-the-art fit. For language models, we saw a very small improvement from the graph structure suggested in the D3PM paper, potentially suggesting there is room to see the sorts of improvements from BLOSUM and Gaussian processes but in language by building an improved model.

Furthermore, our models trained on LM1B were trained using the SEDD architecture, so we expect SCUD uniform to be nearly identical to SEDD masking. Therefore one could consider our results a comparison of SEDD / SCUD masking and SCUD structured in both runtime and performance. Finally, models in MDLM and SEDD were trained longer on LM1B than our models, so their performances could not be directly compared. We reasoned that training longer to compare with those models would not impact our claims (1) or (2) since the structured process performs so similarly to SEDD, so we did not do this comparison. Instead we took the language example as an opportunity to demonstrate the scaling of SCUD to huge $K$. This opens the door to more creative forward processes for language making use of other structured $K$.

Finally, we compare to more recent discrete diffusion models in the protein case in our response to Reviewer hLSS (search for "New state-of-the-art protein results on DPLM").

Thank you for your thoughtful review! We are glad you appreciated the theoretical content and address your concerns about the experiments below with a more in depth experiment on using pre-trained weights and DPLM. With this result, we achieve state-of-the-art on both images and proteins.

First let us recap the goals of our experiments to describe how SCUD compares to MDLM and SEDD and why we chose to do further experiments on DPLM.

In brief, to compliment our theoretical contributions, our goal was to demonstrate our theoretical results in practice, and leverage the theory to reinvestigate structured processes in the literature. In the case of D3PM’s Gaussian and EvoDiff’s BLOSUM, SCUD leads to state-of-the art results, but for D3PM’s sparse graph process SCUD leads to less of an improvement. We will make our goal clearer in future drafts.

In particular, our goal in the experiments section was to

1. Show that schedule conditioning improves performance. This was validated by comparison to previous sota models in theory, images, and proteins. Technical considerations meant we were limited to uniform $\mathcal L$, where the superiority of masking was already established.

2. Redo previous experiments that showed that structured processes performed worse than masking. Our theory suggested these were unfair – masking had the advantage of schedule conditioning while structured processes did not. We redid these fairly by building SCUD versions of structured models. But we are not necessarily advocating for using these specific structured processes. And we of course do not expect every structured forward process will lead to an improvement in performance.

For images and proteins, we saw that Gaussian and BLOSUM processes improve the fit to the data, achieving state-of-the-art fit. For language models, we saw a very small improvement from the graph structure suggested in the D3PM paper, potentially suggesting there is room to see the sorts of improvements from BLOSUM and Gaussian processes but in language by building an improved model.

Furthermore, our models trained on LM1B were trained using the SEDD architecture, so we expect SCUD uniform to be nearly identical to SEDD masking. Therefore one could consider our results a comparison of SEDD / SCUD masking and SCUD structured in both runtime and performance. Finally, as you mentioned, models in MDLM and SEDD were trained longer on LM1B than our models, so their performances could not be directly compared. We reasoned that training longer to compare with those models would not impact our claims (1) or (2) since the structured process performs so similarly to SEDD, so we did not do this comparison.

New state-of-the-art protein results on DPLM

Since 1) the EvoDiff’s BLOSUM process lead to a greater boost in performance than DPLM’s language graph process; and 2) DPLM starts with pretrained weights from ESM3 making it an interesting test case to show that SCUD can also leverage pre-trained weights; we elected to do further experiments in the protein setting. Note however that DPLM was trained with a reweighted ELBO as an objective, and developed a specialized sampling procedure that does not sample from the learned distribution; therefore, its performance as an unconditional generative model, which is what we are interested in this work, may not reflect the performance on design tasks the authors demonstrated in their paper. Below we report DPLM perplexities and sample quality according to how DPLM was trained – as a discrete-time masking diffusion model with $\Delta t=1/500$.

In new experiments, we started with DPLM medium architectures, added FiLM layers to the pretrained ESM models and trained with BLOSUM classical, uniform SCUD / masking and BLOSUM SCUD. During the rebuttal period we managed to train on 4.3 million tokens; unfortunately this is less than 20% of the training budget from DPLM. Nevertheless, we show that SCUD BLOSUM indeed achieves the best likelihoods. For sample sample quality we measured by mean predicted foldability (pLDDT) of 100 sequences of length 200) and used 500 model evaluations. We see that SCUD is the best among our model implementations and is competitive with DPLM. Note DPLM was trained for much longer, which is known to increase the confidence of model predictions – indeed the sample quality of our models are still improving. We will update these numbers for models trained for longer in the final draft.

Thank you for your thoughtful review! We clarify our positive scaling and experimental results, provide numbers for the computational complexity of SCUD, and clarify how SCUD is distinct from previously observed time-invariance. We hope that you will consider raising your score if your concerns are addressed!

C1 / Q3: Computational efficiency in practice.

C2: GPU vectorization issue.

SCUD is nearly as computationally efficient as classical discrete diffusion, even at large scale, and our implementation performs very similar operations on the GPU, incurring no more than 10% overhead. The exception is the text example with a structured forward process where classical discrete diffusion was impossibly expensive to train! SCUD actually enables training, with only a minor overhead from the simulation of the forward process: >1000% overhead for classical diffusion to about 20%. And as discussed in the appendix, we trained our language models for 2 days each on $1.1 \times 10^{10}$ tokens.

The computational complexity of SCUD and classical discrete diffusion both depend on the neural network architecture and the forward process chosen. In most cases, even in our large scale experiments, both models had a less than 10% difference in their operations – $s_t$ is treated nearly the same way that $t$ is treated in classical diffusion, and $K^{s_t}$ is calculated via an eigendecomposition nearly the same way $e^{-\log\alpha(t)\mathcal L}$ is in classical diffusion. The GPU operations are nearly the same, except a vector broadcast when dealing with the extra dimension of $s_t$ rather than $t$ – there is no poor memory access.

The exception is the text example where the alphabet size is huge, but computational complexity is an issue for both SCUD and classical diffusion. In text models $K^{s_t}$ is too large to compute via an eigendecomposition; and simulating the forward process by calculating $x_0^TK^{s_t}$ for each position is not naively parallel on a GPU. But the same issue seemingly also applies to the calculation of $x_0^Te^{-\log\alpha(t)\mathcal L}$ via eigendecomposition in classical diffusion. As a result, previous methods restrict to $\mathcal L$ for which one can perform an analytical eigendecomposition (uniform and masking in SEDD) or save precomputed matrices in discrete time (D3PM). But when one can analytically calculate $x_0^Te^{-\log\alpha(t)\mathcal L}$, they can just as easily calculate $x_0^TK^{s_t}$! Therefore SCUD does not suffer from inherent GPU-incompatibility.

However, in Section 6.3 we study a structured process designed in previous work which was found to perform worse than masking in discrete time. To study if this poor performance was a property of the process or of schedule conditioning, we attempted to implement it in SCUD; note we are investigating, not necessarily advocating the use of this process. Now in this case, $x_0^Te^{-\log\alpha(t)\mathcal L}$ was impossible to calculate analytically, so it cannot be implemented in SEDD or MDLM. Similarly we could not compute $x_0^TK^{s_t}$ using eigendecomposition, so we had to perform many iterative matrix multiplications. To make this as GPU-compatible as possible, we flattened all $x_0$, sorted them by decreasing $s_t$, iteratively applied $K$ via in-place batched sparse matrix multiplication to slices like $x_0[:M]$, and unsorted them. This indeed involved $\max(s_t)$ sequential operations and was not as ideal on the GPU as parallel eigendecomposition. It added a roughly 20% overhead but could likely greatly be improved with caching and parallel-scan algorithms.

Q2: In the language case, only D3PM with a different tokenizer is compare with and recent SOTA text diffusion models are missing. Can you explain why?

We cannot perform the same comparison because SCUD actually enabled training the first structured continuous time language model – it was impossible to implement a model trained using the structured forward process with classical discrete diffusion as described above. Nevertheless, we used exactly the architecture from SEDD to perform our investigation in Sec 6.2.

In brief, to compliment our theoretical contributions, our goal was to demonstrate our theoretical results in practice, and leverage the theory to reinvestigate structured processes in the literature. In the case of D3PM’s Gaussian and EvoDiff’s BLOSUM, SCUD leads to state-of-the art results, but for D3PM’s sparse graph process SCUD leads to less of an improvement. We will make our goal clearer in future drafts.

In particular, our goal in the experiments section was to

1. Show that schedule conditioning improves performance. This was validated by comparison to previous sota models in theory, images, and proteins. Technical considerations meant we were limited to uniform $\mathcal L$, where the superiority of masking was already established.

2. Redo previous experiments that showed that structured processes performed worse than masking. Our theory suggested these were unfair – masking had the advantage of schedule conditioning while structured processes did not. We redid these fairly by building SCUD versions of structured models. But we are not necessarily advocating for using these specific structured processes. And we of course do not expect every structured forward process will lead to an improvement in performance.

For images and proteins, we saw that Gaussian and BLOSUM processes improve the fit to the data, achieving state-of-the-art fit. For language models, we saw a very small improvement from the graph structure suggested in the D3PM paper, potentially suggesting there is room to see the sorts of improvements from BLOSUM and Gaussian processes but in language by building an improved model.

Furthermore, our models trained on LM1B were trained using the SEDD architecture, so we expect SCUD uniform to be nearly identical to SEDD masking. Therefore one could consider our results a comparison of SEDD / SCUD masking and SCUD structured in both runtime and performance. Finally, models in MDLM and SEDD were trained longer on LM1B than our models, so their performances could not be directly compared. We reasoned that training longer to compare with those models would not impact our claims (1) or (2) since the structured process performs so similarly to SEDD, so we did not do this comparison.

Instead we took the language example as an opportunity to demonstrate the scaling of SCUD to huge $K$. This opens the door to more creative forward processes for language making use of other structured $K$.

C3: Foundation model concerns.

SCUD’s reliance on $s_t$ does not make it any more susceptible to this issue than classical discrete diffusion. As well, we believe this issue can be solved simply by picking a good parameterization for your model.

Note classical discrete diffusion suffers from the same issue – a radical change in the schedule would strongly modify its dependence on $t$. The expectation in both SCUD and classical discrete diffusion is masking, because it is time-invariant. But actually any discrete diffusion model may be parameterized to be time and schedule invariant (described below), and can therefore be transferred between different schedules, just like masking; with the parameterization below, any SCUD or discrete diffusion model can work on any schedule. Therefore we do not expect SCUD to be any more vulnerable to domain shift than classical discrete diffusion models.

Q1: Your argument that masking diffusion inherently "knows when to transition" is equivalent to saying masking diffusion denoisers are time-invariant.

The connection to the time-invariance of masking models is an interesting connection, however, we make a separate and novel observation with SCUD – indeed $t$ is replaced by $s_t$ in SCUD, it does not go away.

Although not discussed in other literature, we believe the time-invariance of masking diffusion is simply a matter of parameterization, and in fact any discrete diffusion model can easily be made time invariant. We save the details for a future in-depth discussion, but, briefly, this is because one can show that, for the target distribution $p(x_0^d\mid x_t)$, the set of vectors $((p(x\_t^{d}\mid t, x\_0^d=b))\_b)\_d$ are sufficient statistics. The neural network can therefore be parameterized to take just these vectors as the argument, making it depend on no other variables, in particular making it invariant to time and even the particular forward process. When we consider masking diffusion, the vector $(p(x_t^{d}\mid t, x_0^d=b))_{b}$ is a one-hot encoding when $x_t$ isn’t a masking token and is uniform otherwise, recovering the canonical time-invariant parameterization. The same logic can also be used to parameterize SCUD models to be independent of $s_t$ by replacing $t$ with $s_t$.