Self-Speculative Masked Diffusions
Authors: Andrew Campbell
Valentin De Bortoli
Jiaxin Shi
Arnaud Doucet
Venue: ICLR 2026 Poster (2026)
Keywords: mask diffusion
generative models
speculative decoding
speculative sampling
LLM
Primary area: generative models

Abstract: We present self-speculative masked diffusions, a new class of masked diffusion generative models for discrete data that require significantly fewer function evaluations to generate samples. Standard masked diffusion models predict factorized logits over currently masked positions. A number of masked positions are then sampled, however, the factorization approximation means that sampling too many positions in one go leads to poor sample quality. As a result, many simulation steps and therefore neural network function evaluations are required to generate high-quality data. We reduce the computational burden by generating \emph{non-factorized} predictions over masked positions. This is achieved by modifying the final transformer attention mask from non-causal to causal, enabling draft token generation and parallel validation via a novel, model-integrated speculative sampling mechanism. This results in a non-factorized predictive distribution over masked positions in a single forward pass. We apply our method to GPT2 scale text modelling and protein sequence generation, finding that we can achieve a ~2x reduction in the required number of network forward passes relative to standard masked diffusion models.

Published as a conference paper at ICLR 2026
SELF-SPECULATIVE MASKED DIFFUSIONS
Andrew Campbell; Valentin De Bortoli, Jiaxin Shi! Arnaud Doucet Google DeepMind
andrew. Mndrew.campbelll Bout look. com
(vdebortoli, arnauddoucet @google, com com
ABSTRACT
We present self--peeulatiie masked diffusions, a new class masked diffusion gen- erative models for discrete data that require significantly fewer function evaluations to generate samples. Standard masked diffusion models predict factorized logits over currently masked positions. A number of masked positions are then sampled, however, the factorization approximation means that sampling too many positions in one go leads poor sample quality. As result, many simulation steps and therefore neural network function evaluations are required to generate high-quality data. We reduce the computational burden by generating non-factorized predictions over masked positions. This is achieved by modifying the final transformer atten- tion mask from non-causal causal, enabling draft token generation and parallel validation via novel, model-integraut results a non-factorized predictive distribution over masked positions in a single forward pass. We apply our method to GPT2 scale text modelling and protein sequence generation, finding that we can achieve 2 reduction the required number of network forward passes relative to standard masked diffusion models.
speculative sampling mechanism. This
INTRODUCTION
Generative models for discrete data are at the core of a wide range of modern deep learning systems from chatbots (Gemini Team 2023 OpenAl Team 2023 Anthropic 2024) to biological foundation
models (Hayes et al.) 2025 Wang et al.) 2025 Madani et al. 2023 These models generate data through an iterative process. each step revealing only small portion of the final tokens. Autore- gressive (AR) models (Sutskever et al. 2014 Brown et al. 2020) reveal single token in each generation step. doing so in a left-to-right ordering. Alternatively, masked diffusion models (MDMs) and any-order AR models (Hoogeboom et al. 2022, Shih et al. [022] Sahoo et al. 2024 Shi et al. 2024; Ou et al., (202)) reveal multiple tokens at each step and can operate using any given generation ordering. The ordering flexibility is beneficial for applications without inherent left-to-right structure such as protein sequences (Alamdari et al. 2023] Gruver al. 2024 Wang et al. 2024).
During update, standard MDMS use neural network to output factorized predictive distribution over currently masked positions. Revealing more than one token from this distribution incurs approximation error, data distributions typically do not factorize that way; see e.g. (Zheng et al.) 2025). This imposes an upper limit on how many tokens can be revealed per step without loss of quality, meaning large number of network forward passes required to generate full datapoint.
Our objective is to reveal multiple masked tokens concurrently using a nontfaccod predictive distribution. Inspired by self-speculative sampling (Zhang al. 2024 Elhoushi al. 2024 Liu
et al. 2024), we propose generating a draft sequence using a network subset, then validating it in parallel with the full transformer. By employing speculative sampling (Leviathan et al. 2023 Chen
ct al. (202),, we ensure the accepted token sequence is distributed according to non-factorized target distribution defined by the full capacity network.
Using self-speculative sampling for MDMs provides a non-factorized predictive distribution that can be sampled efficiently, but requires solving two key challenges. Firstly, speculative verification
Work done during a Google DeepMind internship
Now at Meta Superintelligence Labs

---

Published as a conference paper at ICLR 2026
requires causal transformer but standard MDMs use non-causal transformer. We solve this with a novel hybrid non-causal and permutation informed causal architecture (Pannatier et al. [2024). building in residual connections that the causal target distribution learns to improve over the non-causal draft distribution. Second, in contrast to standard speculative sampling (Leviathan et al. 2023), the non-causal layers create a target distribution which depends on the speculative sampling accept/reject sequence. We theoretically characterize this dependence by deriving log-likelihood lower bound for this model class.
We demonstrate our method on GPT2-scale text modelling and protein sequence generation. On the Open Web and UniRef50 datasets, we achieve a 2x reduction in the number of function evaluations (NFE) required for given sample quality relative to standard MDMs.
2
BACKGROUND
2.1 MASKED DIFFUSIONS AS ANY-ORDER AUTOREGRESSIVE MODELS
We assume our data æ = (x) D) has 1 dimensions and S possible categories, i.c. E {I...., .S. Given the values of x on some subset of these dimensions, masking style generative models train neural network to predict the remaining unobserved dimensions. This can be expressed mathematically in two ways. The first, MDMs, define continuous time discrete-valued Markov
chain (x, 1:12 [ee] on the state-space {1..... S, S + 1).2. where an additional mask token M, M + introduced; see e.g. (Ou et al. 2025 Sahoo et al.) 2024 Shi et al. 2024) for details.
This Markov chain, initialized at anj = x." is such that
p(a"l=") = (1 at (0 =x +0002 = M.,
(1)
for a non-decreasing noising schedule ot [o, 1] such that on 0. an 1. The time-reversal of this noising process defines a generative model consisting of random events spaced between with each event revealing a single masked token. Individually simulating these D events is computationally prohibitive when D is large. Hence, many inference techniques instead reveal multiple masked tokens simultaneously using factorized predictive distribution over the masked dimensions given the current partially masked context a.!D. This is parameterized by neural network with parameters / as,
1 t= and
po (a 1:D
(.:D) xt
darf= po(x xt (
II a 8(x =x
(2)
As shown by Austin et al. related to any-order AR models which directly consider the unmasking events, doing away with the time variable L. Let a be a permutation of the numbers (1....,,D) D). with aii j) = (a(1), (i 1).....((1) (5)) representing the i-th to the j-th element of this permutation. This permutation is usually picked uniformly at random. Given some subset of the dimensions of x, any-order AR models predict the remaining dimensions by parameterizing the predictive distribution pe (i++1)) (()x)) AR diffusion models (Hoogeboom et al. 2022) use a factorized parameterization of this distribution similar to MDMs
2021): Hoogeboom et al. (2022): Ou et al. (2025), MDMs are closely
+ Fo (1:0) II
d=i+1 Po (w) (1:1))
where the notation signifies Po (i) (I:))) is parameterized with transformer using non-causal (++) any-to-any attention mask. Conditioning only the (1x)) dimensions is achieved by masking out the x (i++1: dimensions at the input the transformer. The distribution is factorized t>
over the components of (i++:DD We refer p as a non-causal distribution.
An alternative approach is to parameterize the predictive distribution autoregressively, using a network architecture that can operate over a given input ordering (Yang et al. 2019 Pannatier al. 2024). i.c.
(2)) (1a))) = II
d+i+ P (d) ((dd
where the P notation signifies that (it is parameterized with a transformer using causal (-> attention mask ensure it only attends to (1:d-1), We provide overview on
2

---

Published as a conference paper at ICLR 2026
Algorithm 1 Mask Generative Model Sampling With Non-Causal Pa
i i 0 Number of tokens currently revealed
a ~ plo) Draw the generation ordering
while i D do
1o ~ p(k])) x. (++1itt)) ~ Po (1 (1:)))
Number of tokens to reveal, k 2 1
i i end while
attention masks in Appendix A The distribution is autoregressive over the components of (i+1:D)) We refer to as a causal distribution.
Mask generative models with non-causal Po are sampled using the multi-step process in Algorithm Firstly, a generation ordering a is drawn from p(r). This step is explicit for any-order AR models whereas for standard MDMs is implicitly picked uniformly at random. Then, within each update step, a number of tokens are revealed in the ordering specified by a. Any-order AR models typically set & to be some constant number of reveals per step, whereas standard MDMs reveal random number of tokens per step. The number of reveals, k. is drawn from some distribution p(k:ji) that can be derived from the noise schedule Or and depends the current number of revealed tokens, i.
2.2 SPECULATIVE SAMPLING
Speculative sampling is method developed to accelerate inference for large language models (LLMs) (Chen et al.] [0233] Leviathan et al. 2023). The total cost of LLM inference is memory-bound, i.e., dominated by the cost of moving the model weights from high- baabhdhndwitthh memory to the accelcrator cache. Since cost is dominated by memory operations, applying the transformer to sequence has similar latency to a single token, despite increased arithmetic computations in the former case. Speculative sampling exploits this by creating draft sequence using a smaller model, which is then efficiently verified in parallel with a single forward pass the larger target LLM. Verification results in the acceptance of a subset of the draft tokens, such that the accepted sequence is still distributed according to the probability distribution defined by the target model.
Formally, let the draft model be (1:dd) and the target model be (ar We aim draw samples from by first sampling sequence of tokens from and accepting some number of them. We let the context of already generated tokens be x 1: and let W > be some window size. To sample new tokens, we first draw the draft sequence, gi+l: ~ ~ p( x @.+1id-)). We then, in parallel, obtain the probabilities that the target model assigns to the draft tokens 1i+1s {g(ā") [2 @i+1ad-1) () Starting from the (i + 1)-th token, we accept each token with probability min (1)
(ii-ldd1)/p(((((i ) Let the first position we reject be j. We then resample & from the adjusted distribution r(00|1 (( x max (0, q(0' | (0)11111)-- - -p@'' ( (((,,11)11))) One can prove that the resulting sequence aitl:j is distributed according to the target distribution g( (=):)) (Leviathan et al. 2023
SELF-SPECULATIVE MASKED DIFFUSIONS
To sample from MDMs efficiently, we want to reveal many tokens k per update step (Algorithm
D.. The standard approach computes the factorized distribution P (zt(((+)))))) (Equation using one single forward pass and samples k tokens from it. If A is large, sample quality degrades due to the conditional independence approximation. Ideally, we would sample from the fon-factorized distribution P Ps(x= (+++(++))))) (1=))) (Equation 4). However, naively sampling l tokens autoregressively from P& requires forward passes, negating any computational savings.
To circumvent this, we use speculative sampling to efficiently sample from P (x+((+))))) (1:)). using the non-causal model as the draft and the causal model the target. To avoid needing two models, Section 3.T describes a novel hybrid transformer with both can-causal and causal layers, enabling drafting and verification in one forward pass. We then describe our training objective
3

---

Published as a conference paper at ICLR 2026
Mon-sauesd Couscl Iiteeen
Me (a
Posittion Information
Speoutation reaak>
Non Causal Tronsfor mer
Causal Trans rormer
Permutation
cendio
"hararding"
wurllll (.00)
fozording
<nosko
88
"like"
bozor diing
"Spero
grese
2ue5a
e6
e(t)ee
=[2)+s
0[3)-2
a(4)m4
=13)=3
D[B)=t
guees
hozording
tion
Fosition Encodiing
Toput Feowdttiig
Counni
litdddn Stotte
Concotonation
Oncering
ratwanting
Kon-Cousat Hidden Stotes
Figure I: Hybrid non can-causl/ccaual transformer during training on the sentence "Speculation is like hazarding a guess" which is partially corrupted in the ordering given by a 6) [6, 5, 2, 2, 4, 1]. We always mask the final elements in the sequence, here the final 3 elements in the sequence, [4,3, 1]. At inference time, the causal transformer uses draft tokens rather than real tokens. Empty circles represent concatenation whilst circles with a plus represent addition.
(Section 3.2) and sampling algorithm (Section 3.3 This combined architecture induces shifting speculative sampling target, which we characterize theoretically in Section 3.4 and discuss practical sampling improvements Section 3.5
3.1 ARCHITECTURE
We need to parameterize both non-causal factorized distribution P (x (an (1:i)) and causal AR distribution P p0 (1:4)) for any permutation a. We propose a hybrid architecture that starts with set of non-causal blocks parameterizing Pa similar standard MDMs. We then use set of causal blocks operating on the hidden states from the previous layers to parameterize p These causal blocks follow the o-GPT architecture (Pannatier et al. 2024). which models any ordering by operating on the permuted sequence with additional positional encodings. Our full architecture is given in Figure we now detail each component.
Non-Causal Blocks. These layers follow the standard MDM architecture: masked/unmasked tokens and positional encodings passed into any-to-any attention layers. The factorized distribution
sa Po( (i+1:() (0(= is computed from the output hidden states; the track corresponding to o
predicting the token in position ajj) for j > i. For example, in Figure 1 t p (x) (5))0)(()((()) ((=(())))) is
ordering (h3 in this case since (5) 3). The notation (1:))) denotes that tokens (1:a) were revealed as conditioning signal that was subsequently processed by the non-causal layers with parameters 0.
computed from the hidden state on the track corresponding to the fifth position
Causal Blocks. We parameterize the causal distribution via the o-GPT architecture: left-to-right attention masked causal blocks applied to the permuted sequence. Unlike non-causal blocks, the track corresponding position ojj) predicts the token for the next position in the sequence, a + This is standard for causally masked transformers; since the full unmasked sequence is input, predicting the token a(j) from track (j)) would be trivial. Because we operate on any ordering a, each track needs information positional encodings (Pannatier et al. 2024). For example, in Figure the (5)) track uses positional encodings for a (5) =3 and o (6)
both its own position, o(j), and the next position o 1), provided via
to predict the token "Speculation" in position (66) =
In our architecture, we also input the non-causal hidden states corresponding to the current and next position into the causal blocks, following the same reasoning as for the positional encodings. For the first position, or (1). we set the causal distribution equal to the con-causal distribution as no extra information is available.
4

---

Published as a conference paper at ICLR 2026
Output. At the causal output, for the prediction of token in position o(j + 1), connection adding the non-causal hidden state for position a(++1) c.g. h7 for the token "Speculation" in Figure [ This ensures the causal target learns improve over the non-causal distribution. It also aligns the draft and target distributions, increasing the speculative acceptance rate. The causal output from our hybrid architecture depends both the parameters of the non-causal blocks (0) and the causal blocks (Φ). write this as Pe ( (2x==))0(x( (1+11-t ))) specifying that tokens (1:s)) were processed by B in the non-causal blocks, whilst future tokens of (x (++1d-1)) were input only into the causal blocks.
residual
Training and Inference. During training, we take a true token sequence, apply masks and input this into the non-causal blocks with the model predicting the true token values in the masked positions. We input the permuted full true token sequence into the causal blocks without masks, with the causal attention operation ensuring each track can only attend to tokens prior to it in the permuted sequence.
During inference, let the currently revealed set of tokens be (1:1), We pass (1:a) into the non- causal blocks, with any positions not yet revealed being represented with mask token. For the causal blocks, which take a full real token sequence as input (no mask tokens), we pass x (1:i) as well as draft tokens generated by the non-causal blocks for future positions, as we will see in Section 3.3
3.2 TRAINING OBJECTIVE
When the revealed context to the non-causal blocks is (1:)), our architecture parameterizes both the non-causal distribution
II T Polo (0(x
(5)
and the causal distribution
((zz))))))
-II== [[+=++11 (x) (1=(4)0((x0 (zz (zx ((ti-1)))
(6)
We would like both the non-causal and causal distributions approximate the corresponding true conditional distributions of the data distribution
(o(a (1ar0))))) ~ Paata (1)) (1x
(7)
Po.e (1)) (0 (()))))) of (ittdd 1!)) Pdata (d)) (1:d-11)
(8)
We note that in Equation g there are multiple causal distributions that approximate the same true conditional Pdata (a. n (1:d-1)) for different values of which specifies the computation path taken to parameterize Po.0- We discuss this further in Section 3.3 We train 0 and & by jointly D that normalizes by the number of
maximizing the standard cross entropy loss with a weighting D-c
masked positions D - (Hoogeboom et al. 2022)
D D- (log Pol (x ((x0))))) (1 + log Po. ( (tt1:))))) (1xx)))) .0)]
-=
where the expectation is taken over Pdata (x) pl p(1) with o sampled uniformly and p(i) following a noising schedule similar to MDMs with p D) 0. The first part of the loss is mathematically equivalent to the MDM loss (Shi et al. 2024 Sahoo et al., 2024), as shown by Ou et al. (2025). The second part of the loss can be seen
standard AR cross-entropy loss computed over the masked - tokens under random order UUria et al. 2014 We highlight that evaluating both p and Pe.s for the objective takes only one single forward pass of the hybrid network. In Appendix [ we carry out FLOP analysis for the extra overhead our architecture versus a standard transformer, finding only a 0.98% increase in FLOPs for standard settings.
Training curves (Section 5.1 show that the causal distribution Pa.s achieves significantly lower loss than the non-causal draft distribution Pe. as it can model missing tokens with a non-factorized distribution. We note that architecture is compatible with both fine-tuning and full full-training. In fine-tuning we can freeze the con-causal backbone to a pretrained masked diffusion model and train only the causal blocks top, as we demonstrate in Section 5.3
+
5

---

Published as a conference paper at ICLR 2026
Algorithm 2 Self-Speculative Masked Diffusion Sampling
it 0 Number of tokens currently revealed
a a~ p()) Draw the generation ordering
while i <<D do
o (b+1:: ~ Po ( ((x001)))) Sample draft tokens
Compute (P (a (10)) ((
for d d=i==ii-
do
U ~ 10),1)
if U < min (1, Po ( ((0 (( ((000000000 then
Accept draft token (a())) ar(())
else
Reject (d) and resample, (d( p())(()) where
(())(()) x max (0, P9 o (d) ((x (( (( - - Exit for loop and set (d) (d)) end if
end for < d end while
Update the current number of tokens revealed
3.3 SAMPLING PROCEDURE
We sample our hybrid architecture using Algorithm 2 The general procedure follows Algorithm I revealing tokens in update steps until the sequence is complete. In a given update step (with i tokens revealed), the non-causal blocks first generate draft tokens for all unknown positions, es (I++1:) ~ Po (())((+))))) (((x(()))) The causal blocks then give target probabilities for the Finally, speculative sampling inner loop accepts some of the drafted tokens. Appendix [ details more general windowed procedure.
drafted tokens {P& a (ir(d) ((x))))) ((
Our procedure bears some similarity to self-speculative (Zhang al. 2024 Elhoushi et al. 2024
Liu et al., 2024) and Medusa (Cai et al. 2024) style approaches where an initial subset of the network Is used to parameterize draft distribution which is then verified by the full transformer. However, using self-speculation for MDMs causes additional complexity because the target dis- tribution given by Equation 6 changes depending on the number of revealed tokens. Consider the targets Pa (x (1z)(()) and Pa.a (x+(++1:)) (Ixa()))) o) for > i. The fac- tors for fixed d differ; i.e., importantly, we have Pa.s( (1x)()))) (1 * Po Pa.a (1))))))((
((x((i+d))))) This is because the non-causal blocks operate on different inputs (one sequence with revealed tokens, the other with j). So the hidden states passed to the causal blocks differ, causing the target distribution for position (dd to change throughout the generation trajectory. This is not an issue for existing self-speculative models restricted to a single left-to-right ordering. All transformer blocks causal in this case and never observe future tokens, meaning the target remains unchanged during generation.
3.4 THEORETICAL RESULTS
Let pe. (I= be the probability of generating x (1:D) in the ordering a under Algorithm 2 It is non-trivial to reason about (a, (xx(-)1) because the inner loop target distribution changes ++ whenever a rejection occurs. Furthermore, the alignment between Pg and Pe.s that determines + the rejection probability and both Pa and Pe.s influence po, (1:D)) (o). Therefore, to caleulate + Pa.a (1:))) it would seem that one would have to sum the contributions of Pe and Pa.e over all combinations of possible numbers tokens revealed during each inner loop which would be intractable for D >> Instead, we show in Proposition 3.1 that po. (I:) can be tractably calculated with D forward passes and O(D²) operations using recursive decomposition. The proof is provided in Appendix C.T
+

---

Published as a conference paper at ICLR 2026
Proposition 3.1. Consider the sampling scheme defined by Algorithm 2 For given ordering o, let (r denote the event that the token in position a (d) was accepted and let Ro(d) denote the event
it was rejected and resampled. The distribution of samples output by Algorithm 2 pe, p (1: (r),
then admits the following decomposition (with (D(D+): == 0, pa. (x (1:0) R 1) 1)
Po. (x. (1:)))) (x (
po. (2) Re po p (dtti.)) 1olat1:D)/ Ro(d))
(10)
where (o. Re((1)) = - a
pa (1 1 (2)(
x pe Aolkd-1) 1) 12 (0 (1:k-1), 11( (11)
Po. (x Ao(d+1:D) (1.d), Re(d) and (kr Ao(ked-1), Re oo (1: Ro((-11) correspond to the likelihood associated with a single inner loop speculative sampling procedure. Their 4s directly calculable forms in terms of Pa and Pa. are given in Appendix C.I Further, pa. (1)()))) can be computed with D neural network forward passes and O(D2) operations.
Proposition 3.1 is a dynamic programming recursion where the total likelihood (e. (o) is split up using the simpler object Pe, pe.ş (1xdd) Re) which can be written as a recursion on its previ- to d. With tractable expression for pa (1.D()))))
ous values pe.d ( (I:-1: R=) for we can now obtain an Evidence Lower BOund (ELBO) on log &o (z),
Po. (x)) 2 Ep(e) loggo. ( (1))))))
(12)
We could use this ELBO a training objective but favored instead the objective in (9) which is simpler and computationally much cheaper. Another quantity of theoretical interest is the number of inner loops of Algorithm 2 (each corresponding to one network forward pass) required to generate sample, which defines the expected computational cost to generate a given datapoint. This requires knowing the total rejection count, ND, which can be found again through recursive decomposition of the model likelihood. In Appendix C.2 derive expression for po (N (1.D), o)) for given generation and ordering.
3.5 SAMPLING IMPROVEMENTS
Algorithm 2 assumes each sampling inner loop corresponds to one forward pass of the non- causal blocks and a single verification pass through the remaining causal blocks. We can vary the sample quality-
trade-off by rumning the speculative inner loop multiple times per single forward pass of the non-causal blocks. The non-causal forward pass first creates draft Po Po(x ((( probabilities on these draft tokens using the causal layers. One speculative sampling loop is run to accept a number of these draft tokens finishing with rejection and a resampling of a token. In the next speculative sampling inner loop we re-use the non-causal hidden states but recompute the target probabilities on the remaining draft tokens, noting the resampling operation will change the target probabilities for subsequent tokens compared to the first speculative sampling inner loop. This procedure is repeated for a user specified number of inner loops, the full algorithm de- scribed in Appendix [ In experiments, the vast majority of the network is non-causal, this procedure greatly increases efficiency. To know when to stop generating tokens with the same
distribution. Draft tokens are then sampled and we compute the target
draft P Po (a: (1:))) we use windows (Leviathan et al. 2023 Chen et al. 2023) to set maximum number of tokens we can reveal. We discuss window choices in Appendix D
RELATED WORK
Several speculative sampling works have explored using single network for both the draft model and target model (Zhang et al. 2024 Elhoushi et al.) [2024] Liu et al., 2024). termed self-speculative decoding. These differ from our work by using causal draft and target layers in strict left-to-right
7

---

Published as a conference paper at ICLR 2026
setting. In the same setting, non-causal factorized predictions were explored by Cai et al. (2024) using a mostly causal stack with single-layer factorized prediction heads to draft tokens. We flip this: our network is mostly non-causal with a small final causal layer. Finally, non-causal draft model and entirely separate causal target model was explored for audio generation in Ziv et al. (2024). We
combine these two into a single architecture reducing deployment complexity.
Causal models for any-order AR training were explored by Yang et al. (2019), who introduced XLNet, a specialized two-stream architecture enforcing correct causal dependencies for different orderings. Further architectures used a standard single stream causal architecture and input ordering information through double positional encodings (Pannatier al., 2024) or extra position tokens Pang et al. ).25). Pannatier et al. (2024) describe speculative sampling procedure but its validity is not established and the draft model is not described. An any-order causal transformer baseline reported by HHoogeboom et al.] (2022) performed poorly in the MDMs context. We improve on the pure any-order causal model by letting re-use computations from a powerful non-causal transformer.
Guo and Ermon (2025) perform self-speculative decoding within any-subset AR model, which, given an arbitrarily located prompt, in-fills missing tokens a left-to-right fashion. In contrast, our model can condition
arbitrarily located prompt and in-fill using any ordering too, more akin to standard MDMs. They rely on the XLNet architecture which faces optimization difficulties for generation tasks with short prompts and large numbers of missing tokens requiring optimization tuning such as a masking schedule warm-up (Yang et al., 2019). Our architecture does not require any mask schedule warm-up and operate in the fully masked regime with prompt.
In the MDMs context, previous works have also aimed at speeding up sampling. Similar to our work, some aim to introduce dependencies between the newly sampled tokens via an energy-based model (Xu et al. 2025) or a latent variable mixture model (Hayakawa et al., (025). We instead opt for speculative sampling to sample from a con-factorized target in one shot. Liu et al. (2025) use an AR model to derive a target, however, their predictive distribution must be sampled autoregressively for each denoising step. [Deschenaux and Gulcehre (2025) use distillation speed up sampling, however,
their student model remains factorized, limiting potential speedups.
5 EXPERIMENTS
5.1 TEXT8
We first validate our method on the small-scale text8 dataset (100M chars, 27-token vocabulary). We trained a 150M parameter semi-causal transformer (11 non-causal blocks, 1 causal block) using the objective given Equation 9 In Figure 2 we show the training losses split by the non-causal, log Pa (x=((())))) (l:)), and causal, log Pa. (x (i++()) I:D) ((x (1:))) o) components. The total loss is the sum of these two components. We see that in the initial stage of training the non-causal and causal loss exactly track each other, attributable to our use of a residual connection to the causal output in our architecture. After around 104 steps, the causal block then learns make use of the additional context available to it and dramatically reduces its training loss. This exemplifies the extra capacity a causal distribution has to fit the training data distribution due moving beyond the factorization assumption. With our speculative sampling procedure, we can efficiently produce samples from this more powerful causal distribution.
We now investigate our model's 's seliciincy-saamppe quality trade-off. We compute the NFE during sampling by defining 1 NFE as a full 12-block forward pass. A standard speculative step (11 non- causal causal pass) counts as 1 NFE. Running the causal block multiple times increases this fractionally (e.g., non-causal + 7 causal passes = 18/12=1.5 .. NFF). For MDMs, an update is NFE unless tokens change value in which case the update step could have been skipped and this is counted as 0 NFE. We used this best case analysis for MDM to provide a strong bascline.
To measure sample quality we calculate spelling accuracy by generating D 256 characters and then computing the proportion of words within the sample that also appear in the training dataset; a word being defined any number of lowercase characters between two whitespace characters. To create the accuracy-NFE tradeoff, we vary the window size and number of draft-verify steps (each requiring one forward pass of the causal block) undertaken per pass of the non-causal blocks. We
8

---

Published as a conference paper at ICLR 2026
4.0
0.90
causal noncausal
3.5
Spelling Accuracy
3.0
Speculative Mask Diffusion
2.5
0.80
2.0
0.75
1.5
0.70
10a
10*
10a
10c
10
20
30
40
50
70
100
Training Steps
NFE
Figure 2: Causal (Po..) and non-causal (p,)
Figure 3: Spelling accuracy versus number of function evaluations (NFE) on the Text8 dataset for speculative approach and mask diffusion.
training losses on the text8 dataset.
utilize a cosine shaped window (similar to the cosine discretization grid proposed in Shi et al. 2024) with varying time step parameter, the derivation of which is given in Appendix D
We compare against MDM using the implementation from Shi et al. (0024) with a varying number
of timesteps used simulate the generative process. We plot our results in Figure 3 We that
our approach is able to achieve higher spelling accuracy at lower NFE versus standard MDMs. For example in the low NFE range, our approach is able achieve greater than 2x reduction in NFE.
5.2 OPENWEBTEXT
We demonstrate our method OpenWebTexx using GPT2-scale, 150M parameter, 12-layer transformer with RoPE positional encodings (Shi et al.) 2024), setting the first 11 layers as non-causal and the final layer causal. We use generative perplexity (measured by GPT2) supplementing it with unigram token entropy to ensure sample diversity (as generative perplexity can be cheated by low-temperature sampling). Compared to an MDM baseline (Shi et al. 2024) with an identical architecture, our method (Table D achieves the same generative perplexity with half the NFE, while maintaining similar sample entropy for all NFE levels.
metric,
We also compare with Self-Distillation Through Time (SDTT) (Deschenaux and Gulcehre 2025), where student model is trained on coarse time grid match teacher MDM sampled fine time grid. SDTT achieves very low GPT2 NLLs, far below the teacher MDM model using much larger NFE. Investigating the entropy of the produced samples reveals that SDTT samples lower entropy than other methods. This mode-seeking is likely caused by truncation errors in the SDTT teacher model sampling (Zheng al., 2025). In contrast, our approach reduces NFE while maintaining sample diversity equivalent to the baseline.
Table 1: GPT2 NLLs (in nats per dim) and unigram token entropy (nats). For each method a metric- NFE tradeoff curve created by varying sampling parameters. Values at each NFE are read off by linearly interpolating between the two nearest points. For comparison we include results for mask diffusion and Self Distillation Through Time (SDTT) (Deschenaux and Gulcehre) 20255.. We also include ablations, removing the residual output connection and a 10 non-causal 2 causal block model.
Method
GPT2 NLL (L)
Entropy
NFE
64 NFE
128 NFE
256 NTE
32 NFB
64 NIE
128 NFE
256 NFE
Masked Diffusion
5.50
5.27
5.13
5.05
5.72
5.70
5.68
5.67
Speculative (ours)
5.28
5.12
5.05
5.02
5.72
5.70
5.68
5.66
SDTT
3.70
3.46
3.30
3.18
5.31
5.25
5.18
5.11
No output residual
5.36
5.16
5.10
5.05
5.74
5.70
5.68
5.67
10ne-2c layers
5.34
5.16
5.06
5.03
5.71
5.68
5.67
5.66

---

Published as a conference paper at ICLR 2026
Finally, we perform two architectural ablations. First, removing the output residual connection (see Figure worsens the GPT2 NLL NFE trade-off. We hypothesize that the residual connection both makes the target easier for the single causal layer to learn and aligns the draft/target distributions, boosting acceptance rates. Second, moving from an 11 non-causal blocks and causal block to a 10 non-causal blocks and 2 causal blocks architecture also worsens the GPT2 NILL NFR trade-off demonstrating that the best balance between the power of the draft distribution and power of the target distribution is achieved with only a single causal block applied the non-causal stack.
5.3 PROTEIN SEQUENCE MODELLING
We now demonstrate our method on non-text data by modeling amino acid sequences from the UniRef50 dataset, following settings from Wang et al. 2024) (40M sequences, max length 1022). To demonstrate applicability to pretrained models, we take the 150M parameter, 30-layer ESM2-based model from Wang al. (2024) (Lin al. 2023). freeze its parameters and add single causal block (carrying hidden states from the final non-causal block), training only this extra head.
We measure sample quality using the average pLDDT from ESMFold (Lin et al. 2023); higher PLDDT indicates higher folding confi- dence, suggesting the sequence better follows the natural distribution. We compare this to the original Wang et al.] (2024) non-causal model, sampled using the standard MDM algorithm (omitting resampling, temperature scaling, or confidence scoring for a consistent baseline). Our results (Figure 4 show better sample quality-NFF: trade-off than the standard algo- rithm, in particular achieving for high PLDDT. This showcases the benefit of adding just a single causal block frozen pre- trained network. We note that algorithms that also choose can achieve higher pL.DDT val- ues, for example Peng et al.] (2025), however we here focus on the standard MDM formulation with selected uniformly at random.
44
42
Speculative Mask Diffusion
36
34-
20
30
50
100
200
NFE
2x speed-up
Figure 4: pLDDTSs versus NFE for mask diffusion and our speculative method. The mean PLDDT is computed over 512 samples with the standard error of the mean represented by the shading.
6 DISCUSSION
We presented self-speculative masked diffusions, a novel architecture and sampling scheme for MDMs that reduces the number of network forward passes required for given sample quality compared to standard MDM approaches. On variety of datasets, from text to protein sequences, our method consistently improves computational efficiency, reducing the number of network evaluations by up to 2x compared to MDMs while adding only negligible computational overhead (a 0.98% increase in FLOPs in standard settings). Further work could explore the natural combination of our method with compute-intensive inference-scaling techniques, such as re-masking corrector steps, in order to reach a new level of model reasoning capability for a fixed compute budget.
REFERENCES
Alamdari, S., Thakkar, N., den Berg, R., Lu, A. X., Fusi, N., Amini, A. P., and Yang. K. K. (2023). Protein generation with evolutionary diffusion: sequence is all you need. In NeurlPS 2023 Generative Al and Biology Workshop.
Anthropic (2024). The Claude model family: Opus, Sonnet, Haiku. Technical Report.
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and Van Den Berg, R. (2021). Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems.
10

---

Published as a conference paper at ICLR 2026
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners. In Advances in Neural Information Processing Systems.
Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T. (2024). Medusa: Simple LLM inference acceleration framework with multiple decoding heads. International Conference on Machine Learning.
Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L.., and Jumper, J. (2023). Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.013/6
Deschenaux, J. and Gulcehre, C. (2025). Beyond autoregression: Fast LLMs via self-distillation through time. In International Conference Learning Representations.
Elhoushi, M., Shrivastava, A., Liskovich, D., Hosmer, B.. Wasti, B., Lai, L., Mahmoud, A., Acun, B., Agarwal, S., Roman, A., et al. (2024). LayerSkip: Enabling early exit inference and self-speculative decoding. In Annual Meeting of the Association for Compulational Linguistics.
Gemini Team (2023). Gemini: family of highly capable multimodal models. arXiv preprint arXiv:23/2.11805
Gruver, N., Stanton, S., Frey, N., Rudner, T. G., Hotzel, I., Lafrance- Vanasse, J., Rajpal, A., Cho, K.,, and Wilson, A. G. (2024). Protein design with guided discrete diffusion. In Advances in Neural Information Processing Systems.
Guo, G. and Ermon, S. (2025). Reviving any-subset autoregressive models with principled parallel sampling and speculative decoding. arXiv preprint 25X4:25044.0456.
Hayakawa, S., Takida, Y., Imaizumi, M., Wakaki, H., and Mitsufuji, Y. (2025). Distillation of discrete diffusion through dimensional correlations. In Advances in Neural Information Processing Systems.
Hayes, T., Rao, R., Akin, H., Sofroniew, N. J., Oktay, D., Lin, Z., Verkuil, R., Tran, V. Q., Deaton, J., Wiggert, M., et al. (2025). Simulating 500 million years of evolution with a language model. Science, page eads0018.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L.,, Hendricks, L. A., WelbI, J., Clark, A., et al. (2022). Training compute-optimal large language
models. arXiv preprint arXiv:2203.1 15556.
Hoogeboom, E., Gritsenko, A. A., Bastings, J., Poole, B., Berg, R. d., and Salimans, T. (2022).
Autoregressive diffusion models. In International Conference on Learning Representations.
Leviathan, Y., Kalman, M., and Matias, (2023). Fast inference from transformers via speculative decoding. In International Conference on Machine Learning.
Lin, Z., Akin, H., Rao, R., Hie, B., Zhu, Z., Lu, W., Smetanin, N., Verkuil, R., Kabeli, O., Shmueli, Y., et al. (2023). Evolutionary-scale prediction atomic-level protein structure with a language model. Science, 379(6637):1 123-11110
Liu, A., Broadrick, O., Niepert, M., and Broeck, G. V. d. (2025). Discrete copula diffusion. In International Conference on Learning Representations.
Liu, F., Tang, Y., Liu, Z., Ni, Y., Han, K., and Wang, Y. (2024). Kangaroo: Lossless self-speculative decoding via double early exiting. In Advances Neural Information Processing Systems.
Madani, A., Krause, B., Greene, E. R., Subramanian, S., Mohr, B. P., Holton, J. M., Olmos, J. L.,
Xiong, C., Sun, Z. Z., Socher, R., et al. (2023). Large language models generate functional protein sequences across diverse families. Nature Biotechnology.
OpenAl Team Team
(2023). GPT-4 technical report. arXiv preprint arXiv:2303.08774.
Ou, J., Nie, S., Xue, K., Zhu, F., Sun. J., Li, Z., and Li, C. (2025). Your absorbing discrete diffusion secretly models the conditional distributions of clean data. International Conference on Learning Representations.
11

---

Published as a conference paper at ICLR 2026
Pang, Z., Zhang, T., Luan, F., Man, Y., Tan, H., Zhang, K., Freeman, W. T., and Wang, Y.-X. (2025). Randar: Decoder-only autoregressive visual generation in random orders. In Proceedings of the Computer Vision and Pattern Recognition Conference.
Pannatier, A., Courdier, E., and Fleuret, F. (2024). o-GPTs: A new approach autoregressive models. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases.
Peng, F. Z., Bezemek, Z., Patel, S., RectoR- Bector-Brooks, J., Yao, S., Bose, A. J., Tong, A., and Chatterjee, P. (2025). Path planning for masked diffusion model sampling. arXiv preprint arXiv:2502.033 03500.0440..
Sahoo, S. S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, J. T., Rush, A., and Kuleshov, V. (2024). Simple and effective masked diffusion language models. In Advances in Neural Information Processing Systems.
Shi, J. Han, K., Wang, Z., Doucet, A., and Titsias, M. K. (2024). Simplified and generalized masked diffusion for discrete data. In Advances Neural Information Processing Systems.
Shih, A., Sadigh, D., and Ermon, S. (2022). Training and inference on any-order autoregressive models the right way. In Advances in Neural Information Processing Systems.
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. (2024). Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:12706: 127063.
Sutskever, I., Vinyals, O., and Le, Q. V. (2014). Sequence to sequence learning with neural networks. In Advances Neural Information Processing Systems.
Uria, B., Murray, I., and Larochelle, (2014). deep and tractable density estimator. In International Conference on Machine Learning.
Wang, X., Zheng. Z., Ye, F., Xue, D.,, Huang, S., and Gu, Q. (2024). Diffusion language models versatile protein learners. In International Conference on Machine Learning.
Wang, X., Zheng, Z., Ye, F., Xue, D., Huang, S., and Gu, Q. (2025). DPLM-2: A multimodal diffusion protein language model. In International Conference on Learning Representations.
Xu, M., Geffner, T., Kreis, K., Nie, W., Xu, Y., Leskovec, J., Ermon, S., and Vahdat, A. (2025). Energy-based diffusion language models for text generation. In International Conference on Learning Representations.
Yang, Z., Dai, Z., Yang. Y., Carbonell, J., Salakhutdinov, R., and Le, Q. V. (2019). XLNet: general- ized autoregressive pretraining for language understanding. In Advances in Neural Information Processing Systems.
Zhang, J., Wang, J., Li, H., Shou, L., Chen, K., Chen, G., and Mehrotra, S. (2024). Draft & verify: Lossless large language model acceleration via self-speculative decoding. In Annual Meeting of the Association for Computational Linguistics.
Zheng, K., Chen, Y., Mao, H., Liu, M.-Y., Zhu, J., and Zhang, Q. (2025). Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. In International Conference Learning Representations.
Ziv, A., Gat, I., Lan, G. I.,, Remez, T., Kreuk, F., Défossez, A., Copet, J., Synnaeve, G., and Adi, Y. (2024). Masked audio generation using a single non-autoregressive transformer. In International Conference Learning Representations.
12

---

Published as a conference paper at ICLR 2026
Appendix
This appendix is organized as follows. In Section A
present illustration of the different attention mechanisms considered in this paper. Section B details the complete speculative sampling algorithm, including the windowing technique used in our experiments. Section C then provides proofs for all of our propositions, deriving the likelihood of our model when sampled using Algorithm
2 and deriving the distribution over the number of rejection steps within Algorithm [2] In Section D we discuss choices for the window function used during sampling including the derivation of cosine style window. Section we provide a FLOP analysis of the additional overhead of our
self-speculative architecture and in Section [F] we discuss the influence of hyperparameters on our method. Finally, details for all our experiments are provided Section G
A ILLUSTRATION OF ATTENTION MECHANISMS
In order to give an intuitive illustration of the different mechanisms at play in masked diffusion, classical AR models and AR models with random ordering we present the attention matrix (with its mask) in Figure 5 We consider the text sequence from Figure] I "Speculation is like hazarding a guess".
Figure 5: From left to right: The attention mechanism based on the example described in Figure From left right: any-to-any attention masks, standard left-to-right attention masks for AR models and a causal attention mask applied to the ordering from Figure [ Rows correspond to the token making the query (attending). Columns correspond the token providing the key (being attended to).
The first case on the left of Figure 5lis the non-causal any-to-any attention mask for standard MDMs. Each token position is able attend the values in all other positions represented by blue squares. The correct dependencies in the P p0(x"" ((x distribution are induced through masking the tokens that are not present in the conditioning information (1:a).
The second case in the middle of Figure is a causal left-to-right attention mask for classical AR models that only consider linear left-to-right ordering, a (1,2.... D}. Each row corresponds to a token making a query against the values along the top. For example, the token "Speculation" can attend only to itself whereas the token "like" can attend to all of "Speculation", "is". ","", and "like". Since token can attend to its own value, this track of the transformer predicts the value in the next position in the ordering, e.g. the track corresponding to the token "Speculation" " predicts the token "is" and the track corresponding to the token "like" predicts the token "hazarding" We note this is different to the non non-causal attention mask case on the left of Figure 5 where each track predicts the token in its own position.
13

---

Published as a conference paper at ICLR 2026
Finally, represent the case of causal attention mask applied to an arbitrary ordering the right of Figure [
We consider the ordering of Figure [
(1) guess
(2) a
a a33 2 is
(4) hazarding
(55) 3 like
a (6)
Speculation
Each token can only attend to itself and those tokens that precede it in the ordering. For example token "a" which 2-nd in the a ordering can attend to tokens "a" and "guess" whereas token "like" which is 5-th in the ordering can attend to all of "guess", "a", "is". "huzarding" and "like". As for the classical AR case, since token can attend to itself, we have that track of the transformer predict the token that comes next in the ordering. For example, the track corresponding token "a" predicts "is" " and the track corresponding to token "like" predicts "Speculation"" ". Since the attention mask imparts the required conditional dependencies, no explicit mask tokens are required in the input to this style of transformer block. We note that in practice, this style of attention is achieved with a standard left-to-right attention matrix but operating on the permuted sequence which is the scheme depicted in Figure
B FULL SAMPLING ALGORITHM
In this section we present the full sampling algorithm with additional hyperparameters as detailed in Section 3.5 We require window function W (i) that takes in the current number of revealed tokens and outputs the maximum number of tokens allowed be revealed for this current pass of the non-causal blocks. We discuss choices for W(i) in Appendix D We also require a constant N that is the number of draft-verify steps to perform per forward pass of the non-causal blocks. The number of causal forward passes per non-causal forward pass is then equal to N.
Algorithm 3 Self-Speculatiie Sampling Masked Diffusion Sampling Full Procedure
i Number of tokens currently revealed
a ~ p(o) Draw the generation ordering
while i< D do x (+++imi((+(++
~ po ( (# (++1 ((), ((x)
Non-causal forward pass
- Revealed tokens counter for internal loop for = to N do Compute ( ((x= ( (( m s=jj++
min(++W(i.D) Causal forward pass
for d = j+ to min(i ((i,, D) do U ~ 1))
if U min (1, Po, (d) ((x)(1:)) (( ((2 (a ((x=(1))) then
Accept draft token (d) + (d))
else
Reject (d) and resample (dd)) (()) where
(())(()) max (o, Po (#old) ((z0())) (( ) Set (d) (dd)
Exit "d = + to min(i + ((1), D)" for loop
end if
end for
j +-
Update internal counter of current number of tokens revealed
end for
j Update outer counter of current number of tokens revealed end while
Algorithm 3 consists of an outer loop each corresponding to one forward pass of the non-causal layers and an inner speculative sampling loop with N iterations. When the current number of
14

---

Published as a conference paper at ICLR 2026
revealed tokens is i, the outer loop starts with forward pass of the non-causal blocks to obtain the draft distribution for the currently masked tokens up to a maximum length specified by the window function, P ( W ()) (0))(()))))))) The algorithm then begins the N inner speculative sampling loops.
The inner speculative sampling loop begins with forward pass of the causal layers re-using the non-causal hidden states that were computed on (1:a), In addition, the causal layers get access to the values of the drafted tokens. The causal probabilities are then of the form
(Po.s (d) 0(x (l:))) <( 1)) (i)
(13)
With these target probabilities, speculative sampling procedure can then be run which will accept some number of the draft tokens. Once a rejection occurs, that draft token is resampled from the resample probability distribution
p(arld)) x max (o, (0.Po,s (o (l.) o( ((
(14)
The token (0(()) is set to this new value that has been sampled from (())(()) The next iteration of speculative sampling can then begin. We perform another forward pass of the causal layers using the same non-causal hidden states as before. However, due to the resampling. one of the draft tokens has changed value and hence the causal probabilities are different than the previous speculative sampling loop. We repeat the speculative sampling procedure with these new causal target probabilities, noting the non-causal draft probabilities remain unchanged however.
These speculative sampling loops repeat until either N loops have been completed or min (i+ (++((). D) tokens have been revealed. In the next outer iteration, the non-causal layers are applied to the currently revealed tokens including those that were just revealed in the prior N speculative sampling inner loops and the algorithm repeats.
C PROOFS
C.I PROOF OF PROPOSITION 3.1
We aim to find the likelihood assigned a given datapoint 1: under the generative model defined 6a by models pa. pg: sampling Algorithm 2 and ordering a. Let us first start by assuming that we have fixed linear ordering a = {l..... D} for simplicity. We let A denote the event that for the d-th dimension, we accepted the draft token in the speculative sampling procedure. Let R correspondingly denote the event that in the d-th dimension we rejected the draft token.
We note that if we have event Ad, then for dimension d + 1 we will be continuing in the same speculative sampling loop as for dimension d. If we have event Rg, on the other hand, we will exit the speculative sampling loop at dimension d and then enter into the next iteration of the outer loop in Algorithm
2 The significance of this outer loop iteration change is that the draft distribution changes
to Pe(x (xd+1:P1((zz !d))), with d > - meaning additional context is now revealed to the non non-causal blocks. In other words, once we have rejection, then the draft distribution will update with all the new tokens that were produced in the previous speculative sampling run. Therefore, in order to know the correct distribution that is being used as the draft, we need to know when our last rejection was to know what conditioning information is being input into the non-causal blocks.
from p (x +1: +1:110((
In order make headway into this problem, we utilize recursive decomposition approach exploiting the fact that the likelihood a and A is independent of all accept/reject decisions prior to the most <+ recent reject decision. This is because P and P depend only on the currently revealed tokens
2: and not the exact path taken to generate those tokens.
In our proof, we will let yd be the result of the accept/reject step for dimension d i.e. yd E (Ad, Rd).. As shorthand, we will sometimes simply write A to mean yd = A and Rd to mean yd Rdd
15

---

Published as a conference paper at ICLR 2026
We start by introducing some marginalization over the accept/reject decisions,
po,a(xx:.7) I:D)
(15)
= =E pa 1:D yl:D)
(16)
YED
Σ po. 1:D yl:D-1 , + E
(2 yl:D-1, ,RD) RRD)
(17)
VER-1
yI-D-1
=
Σ pe 1:0 ( yl:b-1 ,AP) + -po,s(x"1) 1:D RR))
(18)
= Σ po,a(x 1:D yI.D- A-1 ,AD) + Σ
po.s yl.D-2
( y1:D-2,RD-1 ap-1,AD)
(19)
YLD
YLD
+ po. ( 1:D .RD)
(20)
= Σ
pe. 1:D y1:D=2, D-1 A + pa
L:D RRD )++++ + pa. (RD)
(21)
yID-2
(22)
D
=20,0 1:D ,A1:D)= +
Σ
(0,0(x11
1:D Rd Adt1:D)
(23)
d=1
D
+ -Po,olx1:1) A1:D)
Σ
pa
lid pa Ad+1:D/a lol Rd)
(24)
d=1
The distribution Pa. (m Ad+1:D andd R2)) is easy to calculate because it is the probability of generating the given d++1DD I:D and all of them being accepted when the last rejection was Rd. Since the last rejection is the same for all the + 1: D dimensions, the conditioning information for the draft model is not changing since all the dimensions being generated in the same inner speculative sampling procedure.
Specifically, we have
D II k kmd+1
pe (0. d+1:D Ad+1:D x l:d Rd)
(A.
1:k-1 Rd. A A2t1a-1)
(25)
We now introduce Lemma C.I that allows us to compute the joint probability having an acceptance in position & and the output token having value x
Lemma C.I. For single speculative sampling accept/reject step with target distribution p and draft distribution p. we have that the joint probability over the output token and that acceptance occurs is
p(æ, A) min(F(x), (x)),
(26)
while the joint probability over the oulput token and that a rejection occurs is
p(x, R) = max(0, ((x) - (x))
(27)
Proof. A single speculative sampling step consists of 1) Sampling draft token & (#t, 2)
Accepting & with probability min (1, (())
3) If rejected, resampling token à from the
((a)
residual distribution Pres (0) max(0, p (()). For an accepted token æ, the joint probability p(x, A) is simply the probability of drafting x and then accepting it:
((x)
p(x, A) f(x)) min
1,
min( (xx))
(28)
p(x)
16

---

Published as a conference paper at ICLR 2026
For rejected and resampled token it must be drawn from (z)) after some draft token is rejected. The total probability of rejection, P(R), is the expected rejection rate over all possible draft tokens:
P(R) = Σ (()) ( min (1. f(()
(p)))))))))))))(())) p()))
(29)
Σ (Ø()) -min(((2), p(t))
(30)
Σ max(0, (()) - ((t)
(31)
Σ max(0, ( (()) as E - = 0.
(32)
Finally, given p(æ, R) = P(R) Pres (a:), we obtain
a R) = P(R) max(O,, p(x) (xx) max(0, p(x) - ((x)).
(33)
P(R)
This completes the proof.
can now continue with our proof Proposition 3.1
namely by using Lemma C.1
we have
D
II k=d+1
(A
ax 1o Rd 1d+I-=1)
(34)
II min (Fola"10(olid)) Po (2n) (0(xlid), ((xd+1ik-1))) k=d+1
(35)
Therefore, calculate Pa.s (x Adt1: Rd) in O(D) operations.
We now deal with the term pa.s idd Rd... We perform recursive decomposition,
Pa. (a .12e)
(36)
Σ (a lid ylid-1 124)
(37)
pidd11
Σ (o lid ylid Ant 10-1,,10 1R + Σ pa l:d yldd= Addd B²a)
(38)
yiid2
yiid=
Σ (0 ylidd2 ad- R²)) + (o lid R Rd
(39)
ylid22
Σ (x 2
lid ylid2 Ad-1 ).d) + po, ( 2id-11 Bel po. (x. Hrll 1i pel-1)
(40)
Σ (o lid ylid-3 Au-2 Ad-1 Rdd) +
Σ
(a ylid-3, Rdd fd- Rd))
(41)
Vis--
viid3
t lid-1
ped (eddl o00" Red ad 1id-d 1e--1)
(42)
Σ 1o, lid ylid33 Ad-2 Ad ) Fia-3
(43)
+ +P8 (x 1:d-2 ( -1 l:d, اd-1 R2 (x 1id-22 120-2)
(44)
+ Pa (a 1id-1 1 1 (ac" Radl 1id-d1 12
(45)
(46)
Pa sa (x l:k-11 1 (o Alid-1 1e 12h-1),
(47)
17

---

Published as a conference paper at ICLR 2026
where ( R 1 and pa.øl (xd Ad Rd x lid=1 -RR-1) pa (2" R 1 BB--))
by definition.
We can efficiently compute all the required values of po. (aktd Alid-1 Heth lo 1ik- Rbi-1) for all
d 2 k. We first split each term
oo (x kid Ahid-1 R lx lik 1)
(48)
- (xx kid-1 Ahad-1 loc 1:k-1 R (p- Re 1e lidid11 Akid-1 R(-11).
(49)
By Lemma C.1
we have that
po kid-1 Akid-1 1 1 1
(50)
d-1
po A 211-1 Ahil-1 1)
(51)
I=
d-1
min (Po(z')0(z)) ((xa)))(('"))) ()))) (x 0(xz lill ((x
(2')0(x
(52)
Therefore, we can just keep extending this product and obtain all of
(po, (ar A Ahad-1 1:k-1 R² Rz-1)} D
d=
(53)
in just O(D) operations.
To obtain the required p (ar Ak:d-1, Rd 1ik--1 R²-1) distributions in Equation 49 to compute the additional rejection factor, for which we again make use of Lemma C.I
need
Pe, (xe"
12 11 1: 1:d-1 A (, 11²-11
(54)
(o. ( (00)0(())) (:k-)))), oiz ()) P (0(2"1002)
max
(55)
We can therefore obtain
(oo knd Ahid-1 Ahid-1 1d lik-1 Rd- 1)} })2
(56)
d=k
in O(D) operations. Allowing to vary as well, we can obtain
(pos((o knd A 1t 11 p lik
112-11 kad
(57)
in O(D2) operations.
Once, we have of these values pre-computed we can use our formula for pa.s (x Red in Equation 47 to compute all of
( 1:d n²t) d=1
D
(58)
in O(D²) operations. The total cost so far is then O(D2 + D2) O(D2).
The final point of order is obtaining po (x 1:D 1: A1:D). By Lemma C.1 we have
D
1x: A
II
Po.o (ad Ad Al:-11)
(59)
II min (po ((0)): Do (x (z=)0((), (zx1td-1))
(60)
18

---

Published as a conference paper at ICLR 2026
We now have all the required quantities to compute the full likelihood
po. po,0( 1:D Al:D)
Σ
pa Id R p (wd+1:D Adt1:D loo Rd).
(61)
Computing all the (ar ntd) quantities cost O(D2) operations. Computing Po.e (Id+1:: Ad+1:D lid Rd)) cost O(D) each so computing all these required factors is
another (DD)) cost. Computing po.o(x (x! A Al:D) is just O(D). Therefore, the total cost of getting the full likelihood O(D)).
An important note to count how many forward passes the neural networks are required. The required distributions are
P (0(2)((())))) (0 o(x b(x++ldd)))
Vd > K
(62)
Ad
(63)
which in total only O(D) forward passes. So although we need to perform O(D²) total operations, there only ODD) expensive neural network operations required.
Now remove the assumption that the ordering is fixed when calculating the likelihood. Let us assume some distribution over orderings, o ~ p(o) that we use to calculate our likelihood e.g. uniform. A bound the log-likelihood is therefore
log ( :D) 2 Ep(r) lo. ((l:)()))
(64)
For a given a, we can calculate log pa. (x: (o)) as above, just assuming we generate in the ordering given by Therefore, and calculating log Pe. (x ( for numerical stability.
obtain estimate of lower bound the log-likelihood by sampling
(o). To compute the log of sums of probabilities, we use LogSumExp
19

---

Published as a conference paper at ICLR 2026
C.2 NUMBER OF REJECTIONS
Here we derive an expression for the distribution over the number of rejections during the speculative sampling procedure given in Algorithm 2 The number of forward passes required to generate x I: then one more than the number of rejections. We now formally state and prove our result.
Let N be the number of rejections occurring in the algorithm up until dimension d in the ordering. The posterior of N given a datapoint and specific ordering is
Proposition C.2. Consider the sampling scheme defined by Algorithm 2
(11D) ,0) po,0(2 (1::).)11) (.P.,
(e 2o(1:
(65)
po.o(x1:D10
where
la I:D
(N (o) (20,0(x01:0 1:D ,,,, (N))+
(66)
Σ1
10,0(x 11-1 _Nd-1lo)x
(67)
d=1 Na-1
D0,0( ,AD .ND p _id-1, No-1 .0)
(68)
with
po l:d R², No = (x.1.d 1:d Al:d-1 , Nd
(69)
po,olanlik Rk- No 1o)
(70)
k= Nk-1
pa. Abid-1 Rd, 1o R=-1 1 _
.0) (71)
where the quantities Pa.o (w (wd::D, Ad:D, No x 2dd11 1,dd 1 Add1,a)) Pa.s (xd Al.d-1 Re Ndlo)
and po (a A kid- Rd N lik- Rk-1, Nk-1, .) are directly calculable in terms of to and
P 0.
Proof. For notational simplicity, in the following we will drop the conditioning with understood that I:D (1:D) and similarly for other random variables. All quantities conditioned on the same ordering, o, and so this does not affect our results.
By Bayes' rule, quantity of interest is given by,
Po.0l (ND 2:: Pa, (x .D) Po.0 (x: I)
(72)
The denominator can be obtained from Proposition 3.1 but we need to establish expression for the numerator. We will use the same recursive decomposition technique to derive the numerator. The central object of the recursion will be
po (2 Rd (N) (20 Rd Rd (wdd Rd),
(73)
where we can represent pe (N (Nd Rdd) just a D-dimensional probability vector (as the proba- bility of more than D rejects is 0). We can obtain (artd (dd Re)) from the same calculations in
20

---

Published as a conference paper at ICLR 2026
Proposition 3.1
We first write po.o(x!:., (N)) in terms of po.o(x (x RR, ,N,,,
po,0(a". .ND
(74)
=Pe,pl= 1:D , .ND)= + po, RR. .ND)
(75)
=poppl0 1:1 A- AA" ,ND) p RD-1 AA, ,ND) + p(xl ,R", ND)
(76)
=ooib(00 I:D D-1 AD, ,ND) + Σ p( R ND-1, .AA,ND, + p (x1:D) RRD .DD)
ND-1
(77)
-o, D-1 AD .
(78)
+
Σ (n 1:D-1 RD-1, 10-1,ND--0 ND-11 po p0, Ao, ND 1 RD-1, ND-1)
(79)
NP
+ I:D RD NDD
(80)
-o, 1:D A AD2D: 2:D ND) + p0 1:D. , AD-1:D ,DD)
(81)
(a 1:D-1 RD-1 ND ppo AD, ND 1D , RD-1 - (.a1 (t-NN-
+
E
(82)
+ I:D RP,, ,ND)
(83)
po, I.D. AD-2:D .ND) + Σ (o, 1:D. RD-2 ND-2 ai-aad (aN
(84)
ND
+
pn 1:D-1 nD-1 ND-1 (0 do da
no-1 Ia 1.D11
(D-1
(85) )
ND-1
(86)
+ (a I: RD, (ND)
P I: A DD22:D (DD
(0)
(87)
milpoaaa (2 AD-1D
Σ
Pa I:D- RD-2 ND-2 20022000,000
B NDIza:--2 RD-3 2D-2,ND-3
ND
(88)
ND 2
Σ
(n
1:D-1 RD-1
AP 2:D-1 RD-1 ND-1)
(69)
ND-1
+ Po 1:D
(90)
(91)
po (2. 1:d-1 Rdl-1 Nd-1 )o ADD ND 1o Rd-1 Nrd-1)
(92)
Nd
+ Po I:D RD. .ND),
(93)
where pe (x. RO, No) 1.0 and p8, (x AL:D, No [2:1:0 R.. N2) Pe, (a: A1:D ND) by definition. This is tractable to compute because we can compute the inner term as
(x AdDD 10 Rd-1 NN--1
(94)
(xt -ik-- A Rd-1, Ne-1) Po.o (ND) A Rd-1, ( Nd-1).. (95)
The distributions over " A are just the same as those in the proof of Proposition 3.1 as the
extra conditioning over the number of rejections up to dimension 1, Ad does not provide anymore relevant information than is already provided in the conditioning Rd-1 The distribution in Equation 93 is new however and is given by
pe.0 p (N (N I.D A Rd-1 Nd-1) S(ND Nd-13
(96)
because we condition on all accepts between and D so D is the same as Nd-1
21

---

Published as a conference paper at ICLR 2026
Now we find recursion for pe.0(x ld, R .4)
(e.a( R Nº
(97)
= 2:d Ad-1 Rd, N4) + Pa Rd-1, RR-,,,, RRRNN No)
(98)
= 1id Ad-1 Rd, .4A- Σ po.olaohid, lid Rd-1, Nd-1 R² No0
(99)
= 2:d ,d- , , .No)
(100)
Σ P0,0 (2 2:1-1 nd-11, No- po.o n², 1o no--1 Nd-1)
(101)
Nd-1
-Po,0lx:1dd 1id
1d-2id-1, R², No + pa (a: ne- Ad-1, , Re,, ,T)
(102)
+ Σ (0, 1:d-1 ,,00,, _d-1,, N 1 a Re, No 1:d-1 Re- (__=,,,d1 (Nd-1)
(103)
Nd-1
=Pe P2 l:d
Ad Ad-2nd-1,
2:d-1 Re, No) + Σ (al hi Rd-2, Na- Ad-1 ned No)
(104)
Na-2
pe, 1:d-1 1,0-,,,--111 _1--1, 1 _No-)))poodee (x0 Re, No lid-1 ldô11 N
(105)
No
Po (x 1:d
ad-2:d-1_ Rd, R²,NN) Noo)
(106)
(0 1:d-2 Rd_2, ,, No po (2 1:d Ad-1 Rd I Rd- (d
(107)
Nd-2
Σ
po Ra-1, R )0,,0 Rd, Iid R Nd-1)
(108)
Nd-1
=
(109)
- 1:d 1:d-1 Re .Nd)
(110)
h 1k-11, (k-1))o, "onnlla k:d Alod-1 1 ne,, No 2ike- pk-11 1
(111)
All these terms can be calculated with simple adjustments from quantities we have used before to get the full likelihood po..(2x1: Firstly, we have
Po. (a ltd A 1:d-1 Rdd, Ie",Ná) = pa.o (x 1:d
Iid-1,
_Ie pp."))o,o 1:d A 11 102-1,100 R²a)
(112)
Po (xc ltd A 1:dd1 _e [==)6((==11)
(113)
Secondly, we also have
(a kid A 1 Rdd Nd Nod 1ik-1 Rk-11, (NN-1
(114)
=pa kid Ahud-1 1 Rdd x lik-
R
poo (N
x lid Rd And-1 _R _ |N--
(115)
(x kid Ak:d-1 Rd a 1:k-
R (
k + +1}.
(116)
Substituting both Equation| 113 and Equation 116 into Equation
111 gives
(z. ttd Rdd, Nd) = (2 1ldd11 Rd ) 1}
(117)
(2 Rk-1 NNN,,
(118)
Nk-1
po, po kid Alid-1 Ret 2 1sk- H1th ())')8( = -K +
+1}
(119)
Now, by re-arranging Equation 73 we obtain
Po. po (Nd12 l:d Red) = po (20 1:d, Re, No)
(120)
po (1x (a.!.d) lid Rd)
22

---

Published as a conference paper at ICLR 2026
So the overall update to get pa.s (NV I aiiid RR) will be
po. po. (x l:d, A Alid-1, R
Pe, Pe. 1:d R²²
6( = 13+
(121)
p8.0 l (olidd Rd)
po. (2e 1
Σ
Po. (N
10211-11 1bb-11x
(122)
po.ol Rdd)
k= N
pa kid Akid-1 1etl 1o ,P 186-2)6(N2 -N-11 + 1}
(123)
Once all of {po.o (N 1:d ) D-1 have been computed we are then able to use the recursion in
Equation 93 to compute po.s I.D ND). The final step then substitutes this value as well as the value of p I:D) from Proposition 3.1 into Equation 72 to obtain po.o P|z (D:)) for this given value of a.
D CHOICES OF WINDOW FUNCTION
Algorithm requires a window function W(1) that sets the maximum number of tokens that are allowed to be revealed for the current forward pass of the non-causal blocks when i tokens have so far been revealed. In our experiments, we found that monotonically increasing functions work best for (i)). We hypothesize that this is because at the beginning of sampling, when there is much current context for the model to reason over, each new token drastically reduces the possible space of what the final sample can look like, meaning the model needs to make the initial token choices carefully. However, at later steps when the majority of positions have been revealed, the final values of the tokens are strongly determined by the previous context. The model then has less uncertainty over what their values are and can reveal many in one go.
simple example of a window function is linear,
W (i) ii+1 Linear window.
(124)
In our experiments, however, found that cosine shaped window works the best, similar the finding for the schedule in MDMs (Shi et al.) 2024).
In order to derive cosine shaped window, aim to emulate masked diffusion process that is being sampled according to a cosine schedule. For this masked diffusion process, we let ag be the proportion of positions that are masked. Then Or cos (3 (1-7)) where т is a uniform time variable that has a uniform discretization between 0 and 1. We assume that the r step size is Ar. If the current uniform time is T and we consider moving the uniform time down to
AT (as we integrate from r = 1 pure noise to T = 0 clean data). The initial expected proportion of masks for the masked diffusion process at = cos (ī (1 T)). After the update step, the expected proportion of masks is ar-Ar = cos (3 (1 - + AT)).
When using this style of window for speculative sampling, we will use the expected number of revealed tokens during one update step for the mask diffusion process the maximum number of allowable tokens to be revealed during this speculative update step. When we enter the speculative sampling update step,
do not have a time variable and we first need estimate what the equivalent time would be for masked diffusion process with this proportion of mask tokens. We do this by calculating the current proportion of mask tokens in the sample and use this as an estimate of ar, We then convert the a, estimate into an estimate for 7 using the inverse equation,
2
T
cos -1 (a7).
(125)
π
We then calculate the expected proportion of masks after the update step by doing
Cr-Ar = cos
(~ (1 - (1-++Ar))
(126)
23

---

Published as a conference paper at ICLR 2026
The value of Ar is hyperparameter of this window function. The final form for window function then
D
a, D
(127)
2
5 1
cos
cos-((a))
(a+):
(128)
π
( = D
cos
2
π
(1 (1-T)) - cos
a ( (1- + Ar)))
(129)
where L.] is the floor function.
E FLOP ANALYSIS
Our analysis follows that of Hoffmann et al. (2022) Appendix F. The vanilla result holds for both autoregressive models and MDMs since these two differ only in the attention mask which does not affect the total FLOPs.
We make the following definitions: let be the base hidden dimension, F be the forward hidden dimension, H be the number of heads, K be the number of dimensions in the key, V be the vocabulary size and S be the sequence length. Hoffmann et al. (2022) then provide the following relations for the number of FLOPs in each part of transformer terms of C, F, H, K, V.S.
Embedding 2SVC
Single layer attention:
QKV projection = 6SCKH
K@Q = 25"KH
Softmax = 3H
Softmax @ query reduction = 2S*KH
Linear 2SKHC
Dense block = 4SCF
Final logits 2SCV
The total FLOPs for forward pass of vanilla transformer network is then Total vanilla FLOPs embedding num layers (attention + dense) logits.
In our seff-speculative architecture, the FLOPs calculation within the blocks unchanged however we additionally include a projection step before the causal block to include the current and next state information. We re-use the input token embeddings from the standard input the transformer and in our experiments we used rotary positional encodings to encode the position information, splitting the channel dimension in half between the current position and next position. Therefore the only extra FLOPs incurred at this stage are those projecting the concatenated current position hidden state, next position hidden state and token embedding.
A projection operation is of the form h = Wæ where x R W E Rexe and h R To compute the matrix multiplication there are Ga multiplications and Can additions for each element in h meaning total FLOP count of 2c,. In our case, we concatenate two hidden states together and token embedding and project back down to C giving a FLOP count of 2 x 3C x = 6C per token. In addition, we add on the non-causal hidden states to the output causal hidden states before projecting them to the causal distribution logits which is C add operations, resulting in an extra C FLOPs per token. For the architecture settings used in our experiments, FLOPs for the vanilla architecture versus the extra overhead FLOPs incurred by our method. Our transformer Open WebText used the following values: 768, V 50, 257, K 64, H 12,
compare the
24

---

Published as a conference paper at ICLR 2026
F = 3072, S 1024, num layers 12. The vanilla transformer FLOPs are then
Embedding 7.9c10
Single layer attention:
QKV projection 3.6c9
k@Q = 1.6e9
Softmax 3.7e7
Softmax @ query reduction = 1.6c9
Linear 1.2c9
Total = 8c9
Dense block 9.7e9
Final logits 7.9c10
Total vanilla FLOPs 7.9e10 12 x (8e9 9.7c9) + 7.9+10 3.7e11
The additional overhead FLOPs in our method are S x (6C2 + C) 3.6c9. Therefore the overhead FLOPs are only 0.98% of the value of the total FLOPs for a forward pass of the entire network and so are insignificant when compared the 2x reduction in total number of forward passes required to generate a datapoint.
F HYPERPARAMETER INFLUENCE ON PERFORMANCE
In this section we discuss the influence of the hyperparameters of our overall self-speculative approach on sample quality and inference time.
In terms of the number non-causal vs causal layers, Table finds that if the number of causal blocks is increased to 2 and non-causal blocks decreased to 10 then we find worse trade-off between sample quality and NFE likely due to stronger target not making up for the weaker draft model. We expect this trend to continue to higher numbers of causal blocks.
For the number of verification steps per draft and the ÅT parameter of the cosine window, we found in Figure 3 that increasing At results in more tokens being accepted per draft step. However, this increased limit on the number of revealed tokens can result occasions where too many tokens are revealed from weak target early in generation causing worse generation quality. We can see this by focusing on the spelling accuracy NFE tradeoff in the regime where AT is varied and the number of verification steps per draft is held constant at 1.
AT
Spelling Accuracy
NFE
0.01
0.91
80
0.02
0.90
44
0.04
0.88
28
0.083
0.87
21
Table 2: Table highlighting the specific influence of At on spelling accuracy on the text8 experiment in Section 5.I
We observe that initially as At is increased away from 0.01, there is large reduction in NFE with only a marginal decrease in spelling accuracy however, the spelling accuracy becomes worse as AT is increased further and more tokens are allowed to be revealed per step.
We also find in Figure 3 that increasing the number of verification steps per draft has a similar effect to AT in that it allows more tokens to be revealed for each pass of the non-causal blocks which can decrease sample quality due to more tokens being revealed at the start of sampling when the model has less information the final datapoint resulting in weaker predictive distribution.
The ultimate hyperparameter selection that a practitioner should use will depend on their use case and their requirements sample quality and latency. Should latency be major factor, then AT and the number verification steps per draft should be increased up until the point that sample quality becomes unsatisfactorily low.
25

---

The learning rate is then decayed using cosine Published as a conference paper at ICLR 2026
G EXPERIMENTAL DETAILS
In this section we provide a detailed description of our experimental setups.
G.1
TEXT8
For both the mask diffusion and speculative model, use a 12 layer transformer with 12 heads and 768 hidden dimension following the architecture used by Shi et al. (2024). The model was trained for IM steps with a batch size of 256. We use learning rate schedule with 2000 linear warm-up steps with maximum learning rate value of 3 x 10 The schedule over the IM training steps. We use dropout rate of 0.05 within the transformer and apply weight decay with parameter 0.03. For the mask diffusion baseline, we use cosine noise schedule as used by Shi et al. (2024).
To calculate the spelling accuracy we take 1024 samples each of 256 characters and aggregate all samples together to compute the overall accuracy. To compute the number function evaluations, general use best case analysis. If during update step any character changes then this counted towards the total number function evaluations. If no character updated then this function evaluation could have been skipped and this does not count the total number of function evaluations. Note that this is calculated independently for each element the batch and so if one element in the batch changes but another does not, then only the NFE counter for the element in the batch that changed is incremented. This method is applied to count NFE for both mask diffusion and speculative mask diffusion. The overall NFE is found by averaging the NFE count for all individual items in the batch.
When sampling the mask diffusson baseline, in order avoid the truncation issue described in Zheng et al. (2025), we first sample Zp from the model's 's denoising distribution and then randomly select some number of currently masked positions to switch value to the corresponding Zo value. The probability of selecting masked position to be revealed is given by the noise schedule. This avoids computing the probability as combination of both the reveal probability and the token value probability which results in the truncation problem.
We use the cosine window for speculative sampling and the settings for this experiment are given in Table
We provide example outputs from our speculative model and the masked diffusion baseline for both low and high NFE in the following text blocks.
Num draft/verify steps per non-causal
Cosine window parameter AT
0.01
0.02
1
0.04
1
0.083
2
0,083
3
0.125
4
0.167
Table 3: Speculative sampling settings used for the Text8 experiment.
26

---

Published as a conference paper at ICLR 2026
Speculative Sampling NFE 80 Samples:
o zero zero one commercial access on me schillwoltern specifies that things in the language under such definition it can be slow in advantage and athiff they are able to integrate with it as well pe is indicated by typical application cabl
experiment there may often be even scarce eia personas only go and forefall seattle cowards from a more developed family or variety of limitations to create a wash covering nocturnes lady treatment us have an elevated pharmacy or the same thing see
gally protective and has also been prohibiting of rappression tag fades with this imitative of this in focus on being instanced by classes and user outcomes flag lists effect derives from careful crime in legalisation paw systems in proctisation mys
s chimney fertile oil and ocour riots half of the mountain and the largest are factory is three zero zero silica stuff paper which is lowered by a battery five six zero houses a four zero six three ft standard
pound one eight zero bw
Speculative Sampling NFF 11 Samples:
e depending on the sex of the mentioner in the teacher of the berth bettalion in israel preliminary betterment a poem defined the suppositity of euce it ichta profandous this is consistent with the use of sa an a sp and as most serious design of
tabe darna ab gilfara ibn orgaq permus ibn bur dubna starbe ktak elevna and central rebuilt posic dy dibnese jewel freeware the stone can never exist into portugu one and their national museum or ftr bilet vete diordivine wheelway from kapurblake ev
artnnock niap dead sigh one nine nine one development kingbook of self line fm one nine nine one time zony one nine nine two fives aces one nine nine three snum land up two zero zero zero as carl hund one nine nine one time hand go in the sounder convent
e could be seen baokino orpit and kuapan at kagabi school or rkrts if the new generals evolve investigation bp assaulta front and niphi orph ieris moyer the oscris the american laday atm of the iconmen ion goud bo
contain sound specialization the
27

---

Published as a conference paper at ICLR 2026
Mask Diffusion Sampling NFE 81
Samples:
even nine seven oman kratter american short choir one nine six three jhann henrich plehrbach german writer nine six eight robert hoffa english guitarist one nine six four puesere people from tanzanya one hour flies for the fhi to spy and turbay sighter
n day world goal levels shafe a trade distributed by the number of dual trade buyers the ullh carrier or shuit trills even their broad range of equipment with the ammo buy prior to very widely success and preyet trademarks and are required to use softwar
deeco can walk for gyms ginger grass sand rassing into the ground berlin beerstorming pavelos grauss eildedetwith cmnx inday scalting sauce z quix es sensitive sugar bream g bex morzoni foguamore gobimarun ych di mwa mitzazu tersitz wahrer wiestas go
menus to soon promote the music of communicative psychology to exhibit person responses habitual medicare practices controls medical secrets of vox from an extract qu nt fundamental episode an outbreak of given distance between him mojo and prago
Mask Diffusion Sampling NFE 11 Samples:
like oods and also the consulted as in wightrone people at Ihs weight ancient ageiplrf am fart liadone monk pisa cri taa pri bhe water and tet tendar ones tenemsarians theomanteon pemadeonias semen litlachlism the moimosa gyllopeneros koa sikog noun
four zero three zazueng cardel dating the adjective moa line on a fly polaszin moritz the toe pared barrels of berthe wiesel two zero zero two the tendency of oor fiff ben lertwitt you go thigs paddunta blue in initudige case and similar ouareeti n
ine con oupgore a has a e tea the us true commercia stlinen lare in britain oreitnope anciths type of brysosma produces prass the consist of d flats and instead oion the on ess tresoio d flute as circular e when the prvss arose h
f nmig sammeirm unatrowan pass hby hour aegilon abb chodiva fighteng tucsetohouelon fanua first mmanda howered rose formseon aldupo agrenm satimehascof melarimtelan cathodic hypertoroiform by aolelliggragnae pe miligrfin a ane entesteilllaclon ba
28

---

Published as a conference paper at ICLR 2026
G.2 OPENWERTEXT
For our mask diffusion baseline, self-speculatiie model, we switch the last block to causal resulting in an non-causal block I causal block architecture. Our training setup follows that Sh et al. (2024), training for 1M iterations with a batch size of 256 on 64 TPUv3 chips. The training losses are given in Figure 6 Each element in the batch is sequence of tokens of length 1024. We the dropout rate O and
the 12 block transformer from Shi et al.
(002)). As our
weight decay value of 0.03. For our noising schedule during training, we use the cosine noise schedule following Shi al. 2024). The learning rate has 2000 step linear warmup to a maximum value of 3 10 which is then decayed using the cosine schedule.
As for the text8 experiment, we trace out the sample quality NFE trade off curve by varying the number of speculative inner loops per non-causal forward pass and by varying the ÅT parameter of the cosine window. We provide our best settings Table
To compute the GPT2 NLL., we generate samples of length 1024. When computing the unigram entropy, we calculate for each sentence individually and then average the values. Specifically, for each sample we create a histogram for the frequency of each token and then compute the entropy as E0ppo Pi log p; where p; is the observed probability of token with index i.
For the SDTT baseline (Deschenaux and Gulcehre 2025), we run the author provided code available at https: /ttp://gittbbb.. ctm.. com/ jdeschena/ tdtt sdtt The authors find their KLD objective to perform rounds in their work). We therefore compute the baseline with the provided KLD round 7 model. We run the generated samples through our GPT2 NLL computation pipeline to ensure consistency in our comparison.
the best with the maximum speed-up observed on the final round of self-distillation
As for the text8 experiment, we again provide example outputs from our speculative model and masked diffusion baseline for both low and high NFE in S.2.ionG.2.II
12
11
causal
noncausal
10
9
Loss
8
7
5
4
103
10
10
10
Training Steps
Figure 6: Training losses the OpenWebText dataset.
Num draft/verify steps per non-causal
Cosine window parameter AT
1
0.002
2
0.005
3
0.01
4
0.02
8
0.04
10
0.083
Table 4: Speculative sampling settings obtain the NFE tradeoff curve for WebTee.
29

---

Published as a conference paper at ICLR 2026
G.2.1 SAMPLES
Speculative Sampling NFE 376
Sample:
At the end of the day, attend Children's Marnhalle lunch and forget great Italian dad get it because of rubber cup for cup for trophies, electronics and my accent my son, worn a sporting back shirt, is still around 1-2-years in his playing book.
Need say, a boy in the world who hasn't 't yet got the football on his back? At least the city of Holger and Magna Frankos can give him better sense football are all those and they could giving him at 14 years old the idea.
He is said to fail at Neilskad, the peak of the Waratahs in Germany but he is seen an Arsenal, trophy for the club, hoping to have made it once), shout continental enthusiasm.
MORE: Milk's cash back: Nick Scholes gets lifetime deal knowing he's won the Cup
last season showed Stellar, serious tissue character and even Sir Alex Ferguson believes it shows promise. During ISG-season you pass around seeing the first triage in deal to give Obama an iPad that he had re- portedly wished one herself, with the once overlooked-for two. June certainly seemed to the flash kicker of the season.
I'd written about Michael Zae in his Playboy section in 2012. With the cleanest voice, he told the New York Journal the shower about decade ago connected to the 'parent, one guardianship Eve' tween images. turned round and promised ball to the father.
(... continued)
Speculative Sampling NFE 28 Sample:
too tiny and completely makes shitty idea separate from format and fashion! Time, time. Harmonically takes it wrong. Yes, it probably will be ready for next console. You can shout the praises of Xbox Foundation to the Xbox Wire Editor; blew Microsoft repeatedly on Twitter. Less beast you' 've just you into: just heard buzz about the last Q& A Time. It's brilliant developer grants and... followed shoutout. The question will be back. Anyway.
Xbox one: Games are on sale now in the near future?
Xbox: Well, chances are it stays on pace - there another two and these are tremendous results. With console at home no problem nowadays... not least but for the older generations, the kids, they just have a safe home to play games and play things. For them, that all other applications are there, like the apps you use for them.
The soon phone have a convenient service, connected windows, the ability to can choose when to hide etc - lot things there still; the roof of the gate that will replace those from dialup-more examples. The screen that flights will soon have could be unique in different airports... These will all do it again. I don't know some drops! am sure. Achievements in the industry store possibility of pulling away from its competitors with the intention to be overpriced. Rift-free. There are also other kinds of content there. Teachers could get them demos.
Xbox Sure some kind of multiplayer software here confirmed Ultra HD.
(... continued)
30

---

Published as a conference paper at ICLR 2026
Mask Diffusion NFE 372
Sample:
The French government's Energy Commission the only objective publication on the entire portfolio of different technologies, now based on those Paris principles," she added. Then we had all the rain-drops, even the good of' blue-heam wide lights, and she went couch-merschrs like corrint swims was her game, "the science is right only one year" She knows the planet's most wit victim has no gift (apply wind and solar together for second?), It suddenly can both be bought unjust US v. EU, UK whose latest efficiency. Their point is response them that they never asked for, or was just fahanately naive. They re acting like the President, not about Trump but Caracas, Italy for Italy.
harbordmlunk "The administration has appointed committee n for ex- ploration, and governments will have approve for safety. Either way, it will be based on the successful strategies against over-fueled energy. Friend welcomed. But there is every reason cause for worry under the current French approach to climate change."
Edwin DBrowse/flickr "We think coal in its peak was worth less clean en- ergy," writes V. "Australia Nuclear Research Institute(AER) promises that A$20 billion renewables 2020 Project, handsome bet on Australian renewables, an ironically a direct rival to Japan."
What About China?
One person who wondered why is Obama to lend his support to the US side like President Nicolas Hollande, to tell him if President of France agrees to tougher stance against Chinese on climate change plans.
(... continued)
Mask Diffusion NFE 28 Sample:
regime of suing the other for money made.
Finally, there are several questions on policy and many times the question how calculation fees will be justified beyond income cannot be satisfied for something other than income. wwwndooieettttWeeeeeee. Whether you would have thank for this paper that unfortunately should never have been presented, digital looks like good choice dump your test results for those means to loose one ast cage and to free any item that may be interesting. The experimental method (zap DXA) was used from 10 to 40 (chboys age 14 to 19), and the the purpose in this test was not the aim of nudging monkeys to challenge themselves, they were experiment. As the rats were given a task for every 100kg i ar em 3 cube and an array of weights being: 119.80 cm 18 meters) once 8.70 kg ( + kg for pill bottle, ( + for 0) carried art, they were randomly assigned to identical set bareweight draws on silver cans, with the averageweight art being ran with Colored paint ( (A.g., half weight). All monkeys were given fair-weight version of the design of an artwork, over small pickup pot, and thereafter, all color and depth was measured on the pickup material and the design was assessed between 0 and 13 grams cm2. All monkeys were given half weight instructions
100% of the test artists with no strength in water (e.g. were average adult monkeys and were limited when bench compared with the rose at day (ca. 73 9) and sun (45 deg m 3 ). .For more check out includ- ing meemarrs' bathroom in Carroll Gardens [ of PDF The undefended experiments intended for the researchers (... continued)
31

---

Published as a conference paper at ICLR 2026
G.3 UNIREF50
We build our experiments on top of the codebase provided by Wang et al. 2025), https: github com/ httt:/yytdaanc///eeeee dplm We take the airkiiiii/iaaaaaaaaa hnpp///dddl.mmm.mmmmmmmmmmm dp 15 Dm as our base masked diffusion model which is based the 150M parameter ESM2 network (Lin et al. 2023). The standard network contains 30 non-causal blocks and we add an additional block with the same architecture as the previous 30 however switching to causal attention mask. We interface the non-causal blocks and causal block using the wiring diagram given Figure [ We fine-tune the additional causal block for 500 10a iterations, the training losses shown in Figure 7 As the non-causal blocks are frozen, their average losses remain constant during fine-tuning, the noise being caused by different datapoints being sampled in each batch along with different masking ratios. We see that during fine-tune the-ttningg the causal block learns to use the information provided in the extra tokens revealed to it in order to achieve a lower loss. Thus this provides target distribution that better models the data than the non-causal draft distribution.
The base ESM2 network utilizes rotary position embeddings or RoPE (Su et al., 2024 to encode the sequence positions rather than position encodings that are concatenated to the input as shown Figure | RoPE rotates query and key values self-attention by angles given by that track's position in the sequence. We would like to use the double embedding system from o-GPT with RoPE, where both the current position and the next position in the ordering is encoded. To achieve this we split the RoPE channels in half, with the first half encoding the current position and the second half encoding the next position in the sequence.
0.830
0.825
Loss
0.820
0.815
causal
noncausal
0.810
2
4
5
Training Steps (x10²)
Figure 7: Training losses on the UniRe150 dataset.
32


Reviews:
[{"review_id": "78oje3IZ2f", "reviewer": "Reviewer_Fyeq", "summary": "This paper proposes self-speculative masked diffusion,  a new class of masked diffusion generative models for discrete data that require significantly fewer function evaluations to generate samples. In self-speculative masked diffusion, the neural network is decomposed into a large bidirectional diffusion model and a small any-order causal autoregressive decoding head. The large bidirectional model produces candidates, while the small decoding head serves as a verifier. Through self-speculative approaches, the authors are able to get a 2x sampling speed boost with barely any quality decrease.", "questions": "* In Figure 1, do the dots connecting input tensors mean addition? I am assuming this because the token embeddings and positional encodings are connected by dots, but I think this should be more explicitly indicated.\n* For the input to the causal parts, are the hidden states from the non-causal model necessary? I think from a theoretical perspective, only token embeddings and positional embeddings are enough. Can the authors explain, theoretically or empirically, why these hidden states are needed, especially given the residual connections to the output of the causal part?\n* My biggest question is, during sampling, how do the authors handle the causal attention? If it still remains a causal structure, if I understand correctly, then in the case of $k\\geq 2$, i.e., parallel decoding, the mask tokens in the last are attending to mask tokens before them, leading to a different situation from the model's training scenario. Can the authors provide any justifications on why this does not ruin the generation quality?\n* The biggest contribution of this work is to accelerate inference without loss of quality. Can the authors provide any comparison with other acceleration algorithms, such as Fast-dLLM [1]?\n* The authors do not point out if any part of the model is frozen during training, so I assume no parts are actually frozen. However, in the protein experiments, the authors freeze the backbone and only trained a causal head. What will the results be if the backbone is also finetuned?\n\n---\n**References**\n\n[1] Wu, C., Zhang, H., Xue, S., Liu, Z., Diao, S., Zhu, L., ... & Xie, E. (2025). Fast-dllm: Training-free acceleration of diffusion llm by enabling kv cache and parallel decoding. arXiv preprint arXiv:2505.22618.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* This paper is practically significant. I think overall this technique is very promising for real-world deployment. \n* The empirical results are impressive. With barely any loss of generation quality, self-speculative masked diffusions achieve a 2x acceleration, making this almost a free lunch.\n* I think overall the paper is well-written, with a clear motivation, a solid method, real innovations, and abundant theoretical analysis.", "weaknesses": "* The sampling algorithm is not explained thoroughly enough, especially regarding the attention patterns. Please see Question 3 for more details.\n* I think Section 3.4 is less accessible to readers. If the meaning of this proposition can be explained more intuitively, readers will better understand it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mYmBvWsuZJ", "reviewer": "Reviewer_Joif", "summary": "The paper proposes Self-Speculative Masked Diffusions (SSMD), a masked-diffusion generator for discrete data that aims to reduce sampling cost. The key idea is a hybrid transformer: most layers are non-causal to produce a factorized draft, followed by a small causal head that verifies and (via a specialized speculative step) yields a non-factorized distribution over the masked positions in a single forward pass. A theoretical treatment derives a likelihood decomposition for the induced sampler and an ELBO. Empirically, the method reports roughly 2× fewer network forward evaluations (NFE) than a standard MDM at matched sample quality on text (Text8, OpenWebText) and protein sequence generation (UniRef50 + ESMFold pLDDT).", "questions": "see weakness.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Principled modeling & architecture.** The hybrid non-causal/causal stack with σ-aware causal masking is clearly specified and lets the causal head “strictly improve” the draft through a residual connection, aligning draft/target and plausibly raising acceptance rates. The training objective jointly optimizes both heads in a single pass.  \n2. **Sound theoretical analysis.** The paper characterizes the sampler’s distribution under changing targets and derives a tractable decomposition for likelihood and an ELBO; this is rare for speculative schemes beyond left-to-right AR. \n3. **Clear empirical win in NFE.** On OpenWebText, the method matches or improves GPT-2 NLL at **~½ the NFE** relative to a standard MDM baseline, with similar token entropy; ablations support the importance of the residual and “11 non-causal + 1 causal” design. \n4. **Protein application with minimal intrusion.** The protein experiment keeps a frozen ESM2-based backbone and trains just the causal head, showcasing plug-in potential.", "weaknesses": "1. **Reported “2×” is an NFE metric, not end-to-end latency.** The paper motivates gains by memory-bound inference and single-pass verification, but practical wall-clock speed depends on kernel fusion, batch size, cache behavior, acceptance rates, and framework overheads. The text emphasizes that longer sequences do not necessarily mean worse latency because the cost is memory-bound, yet the experiments still operationalize efficiency largely as NFE. It’s unclear how much of the 2× NFE translates into **real time** speedups on GPUs for moderate batch sizes typical in protein generation.  \n2. **Protein benchmark strength and comparability.** The protein study uses **average pLDDT from ESMFold** on sequences sampled from UniRef50. While pLDDT is a useful proxy, the field has converged on stronger metric foldability. Compared to \"Path Planning for Masked Diffusion Model Sampling\", which reports much stronger protein benchmarks with higher-quality samples and broader evaluation, the current results look considerably weaker. This makes it difficult to assess the real progress or competitiveness of this approach in protein generation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jeNJwaTHZZ", "reviewer": "Reviewer_fgUG", "summary": "The paper introduces \"self-speculative masked diffusions,\" a novel method to accelerate the sampling process of masked diffusion models (MDMs) for discrete data like text and protein sequences. The core problem addressed is the slow generation speed of standard MDMs, which stems from their factorized output distribution. To overcome this, the authors propose a hybrid transformer architecture that integrates speculative sampling directly into the model. During sampling, a single forward pass through the non-causal part generates a full draft. This draft is then efficiently verified in parallel against the target distribution from the causal part. This allows the model to accept and reveal a large number of tokens per step, significantly reducing the total NFE. The authors demonstrate the effectiveness of their method on text (OpenWebText) and protein generation (UniRef50) tasks, achieving a ~2x reduction in NFE compared to standard MDMs while maintaining or improving sample quality.", "questions": "- Could you provide insight into the wall-clock time for a single forward pass of your hybrid model compared to the baseline MDM and autoregressive model?\n- What is the practical advantage compared to autoregressive+Medusa, in aspects of implementation difficulty, generation speed and performance?\n- How does the acceptance rate of speculative tokens evolve during a typical generation process?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper presents a clever and significant contribution by combining the any-order generation flexibility of MDMs with the speed benefits of speculative sampling. Addressing the sampling efficiency of MDMs is a critical research problem, and the proposed solution is both novel and effective.\n- The hybrid non-causal/causal architecture is elegant. Using the non-causal component as a drafter and the causal component as a verifier within a single, end-to-end trainable model is a clean design. The inclusion of a residual connection from the non-causal to the causal output ensures the target distribution is a strict improvement over the draft, which likely boosts the acceptance rate.\n- The authors validate their method on multiple diverse and relevant tasks (small-scale text, large-scale text, and protein modeling). The results are consistently strong, showing a clear improvement in the sample quality vs. efficiency trade-off. The authors also ensure a comprehensive evaluation by measuring both GPT2 NLL and Entropy.", "weaknesses": "- Unclear Advantages to Fully Autoregressive Approaches. While native autoregressive models are memory-bound, they achieve better performance than MDMs and can be accelerated by speculative decoding. Meanwhile, the paper is strongly inspired by methods like Medusa, and the paper's 2x speedup is also similar to its speedup to LLMs. It is unclear whether true advantages exist by adapting speculative decoding to MDMs.\n- Limited Discussion on Wall-Clock Time and Computational Overhead. The primary metric for efficiency is the Number of Function Evaluations (NFE). While NFE is a good hardware-agnostic proxy for speed, it doesn't capture the full picture. The proposed hybrid model is slightly larger and computationally more complex per forward pass than a baseline MDM of the same depth. A discussion on the actual wall-clock speed-up and the overhead per forward pass would strengthen the paper's practical claims.\n- Sensitivity to New Hyperparameters. The method introduces several new hyperparameters, including the split of non-causal vs. causal layers (e.g., 11-1), the number of verification steps per draft, and the parameters of the cosine windowing schedule. While the authors present a working configuration, a more detailed analysis of the model's sensitivity to these choices would be beneficial for future practitioners wanting to apply this method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WIGU7UUXWQ", "reviewer": "Reviewer_V8GM", "summary": "The work presents an architecture combining a masked diffusion model that drafts multiple tokens in parallel, and an any-order autoregressive model which evaluates their likelihood to perform speculative sampling. This is done by adding a light causal Transformer layer on top of the non-causal layers of a masked diffusion model (along with positional encodings as in $\\sigma$-GPT). This allows for faster generation than standard any-order autoregressive models, while mitigating the bias that masked diffusion models incur from sampling multiple tokens in parallel (since the independently factorized distribution they sample from is replaced by an autoregressive one). The method is shown to achieve a better NFE/performance trade-off when training from scratch on the Text8, and OpenWebText datasets, as well as when the causal layer is added (and trained) to a frozen pretrained masked diffusion model for protein sequences.", "questions": "1. How does this compare to drafting and evaluating likelihood in two separate models? \n2. How does adding more causal layers impact performance?\n3. In equation 7, why is the weighting justified for the autoregressive loss term?\n4. Does training with the ELBO in equation (10) offer any benefits in the performance of trained models (ignoring computational expense) over the chosen loss?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The work is well-written and flows clearly\n2. The paper is clearly motivated and tackles an important problem in the bias that masked diffusion models have when sampling multiple tokens in parallel.\n3. Experiments clearly demonstrate the impact of the causal layer in better matching the target distribution with fewer NFEs than a masked diffusion model (Figure 2, Table 1, Figure 4)\n4. The broad applicability of the method to pretrained masked diffusion models is demonstrated empirically and convincingly (Figure 4). In my opinion this is a critical use-case of the work.", "weaknesses": "1. Figure 3 evaluates adherence to the target distribution by examining the spelling accuracy (presence of words that also occur in the training set). I think also reporting the perplexity of a larger/better fit model (as in Table 1) would help make this more convincing \n2. The baselines used appear to not use common heuristics for masked diffusion sampling, such as confidence based unmasking. Is the method compatible with such techniques. And if so, what is the impact on the NFE/performance trade-off.\n3. Minor grammatical/spelling mistakes:\n    - Line 233 “oowhere”\n\nI recommend an accept due to the comprehensive and clearly investigated advantages of the approach. The most important issue I have with the work is point 2 above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)