Stopping Computation for Converged Tokens in Masked Diffusion-LM Decoding
Authors: Daisuke Oba
Danushka Bollegala
Masahiro Kaneko
Naoaki Okazaki
Venue: ICLR 2026 Poster (2026)
Keywords: diffusion language models
compute efficient sampling
skipping compute
adaptive attention
Primary area: generative models

Abstract: Masked Diffusion Language Models generate sequences via iterative sampling that progressively unmasks tokens. However, they still recompute the attention and feed-forward blocks for every token position at every step---even when many unmasked tokens are essentially fixed, resulting in substantial waste in compute.
We propose **SureLock**: when the posterior at an unmasked position has stabilized across steps (our *sure* condition), we *lock* that position---thereafter skipping its query projection and feed-forward sublayers---while caching its attention keys and values so other positions can continue to attend to it.
This reduces the dominant per-iteration computational cost from $O(N^2d)$ to $O(MNd)$ where $N$ is the sequence length, $M$ is the number of unlocked token positions, and $d$ is the model dimension.
In practice, $M$ decreases as the iteration progresses, yielding substantial savings.
On LLaDA-8B, SureLock reduces algorithmic FLOPs by 30--50\% relative to the same sampler without locking, 
while maintaining comparable generation quality.
We also provide a theoretical analysis to justify the design rationale of SureLock:  monitoring only the local KL at the lock step suffices to bound the deviation in final token probabilities. Our project page is available at https://daioba.github.io/surelock.

Published as a conference paper at ICLR 2026
STOPPING COMPUTATION FOR CONVERGED TOKENS IN MASKED DIFFUSION-LM DECODING
Daisuke Oba" Danushka Bollegala 2.3
Masahiro Kaneko*
Naoaki Okazaki' "MBZUAI
Institute of Science Tokyo UUniversity of Liverpool
*Amazon
ABSTRACT
Masked Diffusion Language Models generate sequences via iterative sampling that progressively unmasks tokens. However, they still recompute the attention and feed-forward blocks for every token position at every step-even when many unmasked tokens are essentially fixed, resulting in substantial waste in compute. We propose SURELOCK: when the posterior an unmasked position has sta- bilized across steps (our sure condition), we lock that position- -thereafter skip- ping its query projection and feed-forward sublayers-while caching its attention keys and values so other positions can continue to attend to it. This reduces the dominant per-iteration computational cost from (N)dd to ((MNd) where N is the sequence length, M is the number of unlocked token positions, and d is the model dimension. In practice, M decreases as the iteration progresses, yielding substantial savings. On LLaDA-8B, SureLock reduces algorithmic FLOPS by 30- 50% relative to the same sampler without locking, while maintaining comparable generation quality. We also provide a theoretical analysis justify the design rationale of SureLock: monitoring only the local KL at the lock step suffices to bound the deviation in final token probabilities. Our project page is available https: //aioba qithub. .....n.... to/surelck
INTRODUCTION
Discrete diffusion language models (DLMs) generate text by iteratively denoising discrete se- quence over T steps (Li et al. 2022; Schiff et al. 2024); While its formulation varies (e.g.. token swap, insertion, or masking). widely used family operates through masking and unmasking Masked Diffusion Language Models (MDLMs) (Sahoo et al. 2024 Shi et al.) 2024, Nie et al.) 2025 Arriola et al., 2025). Unlike autoregressive (AR) decoding (aaswani et | (017]017)) 2017)-whose per-step compute naturally grows by one query token via a KV-cache Vy-cache-sttndarrd diffusion-style sampling repeatedly rempliingrpeeaedllll positions at every even for tokens that have already been unmasked and considered to be sta- ble. Hence, the per-block cost is dominated by computing the attention scores QK" and applying them to V, i.e., O(N2d) where d is the model dimension. It leads to substantial waste in compute.
recomputes self-attention and per-token feed-forward sublayers for all N token
To address the computational challenges of DLM sampling, prior work has mainly progressed along two axes: Temporal approaches shrink the step count-e.g., parallel/dilated unmasking and staged or learned
tthereby requiring fewer refinement (Luxembourg et al. 2025) Israel et al.] 2025 Wei et al. 2025); Reuse approaches reduce per-step compute by reusing or approximating
K/V vectors and partially updating hidden features across steps (Ma et al.) 2025 Liu et al. 2025b
Wu et al. 2025a). These choices chiefly either reduce the step count 7 or amortize work across steps by reusing intermediate states; they do not alter the within-step spatial granularity: each step cost remains O d) even in late iterations.
still issues N query rows, so the
Beyond reducing the step count T or reusing K/V across steps, we pursue a largely orthogonal axis: permanently and monotonically deactivating token positions as the sampling unfolds. We propose SURELOCK: once token's 's posterior has stabilized, we lock that position-we we cache K/V vectors
*Corresoonding author. Email: daí suke. oba @n 1p comp. act In this paper, we position to refer the fixed token slot in the length-N sequence; position exists even while its token remains masked.

---

Published as a conference paper at ICLR 2026
musked
decoded
ressed wy eached KV
locking action
step-wisc compute
attended other cached
prompt
prompt
Haseline
(b) Surel oek
Figure Iterative sampling by a normal sampler and SureLock. (a) Baseline consistently
recomputes attention scores and FFN sublayers for every token position at every step even after the marginal tokens have become unmasked. (b) SureLock permanently stops recomputing for locked positions once these positions locked. Via cached K/V, other tokens still attend to locked tokens.
and therefore skip Q-projection and poketoken feed-forward sublayers. Active token positions can still attend to the locked ones via the cached K/ (Sec. (2) The per-block cost becomes O(M,Nd) attention and (M,dd for FFN, instead of O(Nad) and O(Nd)), yielding a monotonically decreasing per-step compute profile stability; a criterion whether the per-step KL divergence of the generative probability distribution falls below threshold In order to justify using KI. as the locking signal, derive closed-form bound, which links a per-step KI. at locking time to the terminal log-prob deviation This axis is complementary the Temporal and Reuse approaches; existing methods continue to operate the remaining active token positions.
M, shrinks. We lock token positions based their local
Our work is most closely related to selection-based methods (e.g., DLLM-CACHE (Liu al. 20255)) lower per-step compute by updating only a subset of positions. SURELOCK, by contrast, answers different question-wwhe to remove from compute permanently, instead of what to com- pute now-so the active position set contracts over time; their selection target of step-wise compute can be tapered down monotonically.
We evaluate SURELOCK 2025] in Sec. 3 Across diverse decoding settings (e.g., sequence length, locking threshold), per- step algorithmic FLOPs are monotonically decreasing. On continuation generation with WikiText- 103 (Merity et al. 2016) and instruction following with MT-Bench (Zheng et al.) 2023), SURELoCK reduces algorithmic FLOPs up to 50% in runs without expense of generation quality.
representative MDLMs-LLASE BaaaDA--B-Base/Innsss Base/-Instruct" (Nie et al.
2
SURELOCK
We consider MDLMs that iteratively denoise length-N sequence, producing intermediate quences over steps = ..... step 6, an L-block Transformer yields token-wise logits =({) and posteriors pt (1) softmax(z(")) at token position together with hidden states (5) e Rd. Let M, C [N] denote the masked positions (prediction targets) and M, = [N] \ M, the unmasked ones. In the standard dLLM sampler, each step recomputes self-attentin and FFN for all tokens.
In the following, we explain the proposed method SURELOCK (Algorithm [.. Once a token's posterior has stabilized at step, it locks that token position to stop step-wise computations at all subsequent steps by bypassing sublayers and caching K/V values (Sec. [.I]. Our locking criterion is based the step-wise KI. divergence of the posterior (Sec. 2.2 We explain the rationale of using the KL. as locking criterion from a theoretical perspective (Sec.)
Pt is deliherately conservative and is intended a design rationale, not calibrated predictor for some specific models and datasets
2

---

Published as a conference paper at ICLR 2026
Algorithm SureLock
Require: sequence length N, total steps T, confidence threshold percentile m. KL threshold e
Require:
vocabulary set V. number of layers L, an unmasking policy UpdateMASK())
1: State: boolean lock {0, 1 (init: all 0), caches C. for (k, v) (init: None)
2: State: previous posteriors Pt-1 (init: None), masked indices M₁ (init: [N]
3: State: frozen block-input
Rxxd (init: prompt embeddings)
4: Notation: X₁ E RR Nd denotes the model input at t (from embeddings or previous step output)
5: for = to T do
one diffusion step
6: A < {i lock;=00 =
Le < {i lock; = =1}
D active and locked position at t
7: Initialize block input x! [A,] < X,[A,]. x [C,] t- #[C,]
8: Compute on active positions for f= to L
Ar:
9:
Q[A:], K [r.,, v[A[] + Proiq, (LN(x!))A) (Ar)))
b main projections > assemble
11:
K [ + [[A,,, [C] t- C.k(((]) same for Vall
12:
( [A]] f FFN(Attn((IA)) Kall, Vall))
> attention with FFN D locked rows pass through
13:
20 [A.] + 10 [Ar], x [C.) t [C0
14:
end for Obtain posteriors P. [a[AA < softmax (Proj (x(++(((A))))
15:
16:
Unmasking: M, UpdateMASK(Pp.. Me).
M, [N MM b unmask with posterior
17:
Score on active positions E Ar:
18:
Ug f- max.ev Pt (v) for E A.
(i)
(6)
D confidence value
19:
D + (L (1pp-)
(8) if t>> else 00
step-wise KL
20:
Locking candidates: 2, + As n Me
> must be active & unmasked b threshold over candidate positions
21:
fm t Percentile (((()) IjEZ, m)
(1)
22:
Lock:
Fr - E Z, u1 < 0m A D < < [ock(FFl + 1; Ce [}F + Ke, VeNFFl AA > update locked indices and cache F[Fe] + [F)
23:
P locking evaluation
24:
25:
> freeze block-input for future steps
26: end for
2.1 PERMANENTLY STOPPING WISE COMPUTE AND CACHING K/V
Once a token i is deemed converged step t*, we permanently eliminate its position from per- step compute: we cache its K/ K/V vectors, and bypass its query projection and per-token FFN at all subsequent steps. Let be the set of indices locked by step t, and define the active set A; == M, W C, (unmasked and not yet locked). At step f we assemble the queries [[A, and run a variable- length attention kernel against the full key'value tables K'lll and / where K K'all [G,), / [c] are read from cache. This yields attention outputs only for active positions E Ai,, and we apply the FFN sublayers only to them. Moreover, for locked indices i E L, we keep their predictions fixed, i.c., Pi + Per and skip their FFN computation. Note that locked positions continue to be attended by other tokens via eached K/V.
2.2 CRITERION FOR LOCKING: STEP-WISE KL DIVERGENCE
Our locking rule is driven primarily by local KL; we will justify the validity of this design in Theo- rem inn Sec. 2.3] We also apply a confidence gating as a secondary safeguard to prefer more confi- dent tokens with peaked posteriors. Disabling the confidence gate leaves the theorem unchanged.
Primary: Local KL divergence.
For position i at step f, we define the one-step divergence as D($))
A KT. (8) (12(2))
( -1 max Pt (v) and gm((u) the empirical
Optional: Confidence gate. Let uncertainty u m-th percentile of { ( j E Ae). When enabled, the gate accepts token position as a locking candidate, if ut S gm (u),, i.e., top-m% confidence among active tokens.
(6)
3

---

Published as a conference paper at ICLR 2026
Locking rule.
We lock token position i at step t", if
S and, if the confidence gate is enabled, 4 S 9m (u+-).
The primary criterion alone suffices Theorem Upon locking we cache (Ki , ( and remove position from As thereafter. Unless otherwise noted, the criterion is based on the raw posterior before any temperature scaling. Therefore, SureLock is independent of sampling temperature in the categorical sampling (See appendix for discussion on this choice).
2.3 DESIGN JUSTIFICATION
The purpose of this section is design-theoretic: we justify the use of local KL threshold locking (freezing) criterion by proving that it upper-bounds the terminal error of the log- log-probability in closed form. The controller based SURELOCK locks position i at step t", bypassing Q-projection and FNN sublayers, and reusing its cached K/ K/V thereafter (Alg. []. We establish that the resulting error of the terminal log-probabilities, relative to alternative identical sampler without locking, is bounded by 5
CaallVE under some assumptions, where Crall depends on pperator-norm constants of the model and is the threshold for determining lock (i.e., D( <e). Therefore, the local KL criterion immediately converts to an explicit design for the terminal error (e(6) = 8² /²/C.. In
the following, Pj and 2,{+) denote the posterior and logits at position after step in the no-lock run, while Pt denotes the corresponding posterior when i is locked at tt.
Note. We emphasize that we here justify the use of local KL threshold as a locking criterion by proving that it upper- upper-bounds bounds the terminal error in closed form, rather than to predict an empirical error value for given specific setting.
Standing assumptions for Theorem 1 Theorem relies four main assumptions. We empha- size that these are simplifying regularity conditions introduced for analytical tractability, and that Theorem should be interpreted as idealized model of the behavior we observe empirically. For Assumption A2, which may appear rather strong, Appendix C shows that its implications do not deviate substantially from the practical situations. Let the logit and f() log softmax(2). Fix constants 0, Lam > 0, and E (0, 1). For any position and step 8, the following hold:
(A1) Locking semantics. Once is locked t*, excluded from subsequent re-masking.
(A2) Geometric tail contraction. DL - p D2 for s >>> t²
(A3) One-step logit smoothness. ||2 - (8) ll2 S L !D (derivation in Appendix D
(A4) Log-softmax Lipschitzness. [1/(z) - f(x))l> S Lam 2/112.
Theorem (Locking error bound and closed-form threshold). Fix position that is unlocked up to t and then locked (Alg. T.. Under (A1)-(44) for any terminal time T > tt,
|log PT log loggl"" I SCaan D
Crail
Lam L/(1- Vp).
In particular, if the locking test enforces D S e, then the terminal cog-probability error is at most
8 = Crail VE, so the closed-form threshold is
=(6) = 82 /
Proof. By (Al), locking token position i t* permanently stops the i's step-wise compute, so for
all /. 2 1a we have pt Pi- and log PT = log p.a Therefore
logpT (4) log p Io llog PT (5)
log (3) Il = -11f(2(9)- f(2(0)
By the triangle inequality and telescoping over = t" +1 +,......
T
1 (1)
1l -
Σ
[|
2 112-
850
4

---

Published as a conference paper at ICLR 2026
Applying (A3) term-wise yields
T 2/ 2 S L Σ
12
D
s=++1
Under (A2), D S pp-11-1 D therefore
D S --1-6*
V D
and
Σ 826
D..0
1 vp
Finally, by (A4):
log Pf (8)
log (9)
11 S Lam
1|
(8)
112 S
Ctail
D
which proves the claim and the stated c(6).
2.4 COMPUTATIONAL COMPLEXITY: ALGORITHMIC FLOPS
This paper reports algorithmic FLOPs, counting only GEMMs (Q/K/V/Out projections, QK and AV. and FFN sublayers) Let Ngen the number of positions reserved for generation at initial step, N Nurongi the number of tokens for the prompt, N₂, == max (mmompi) + Ngen be the sequence length of batch b; S the number of sampling steps; B the batch size; d the model dimension; I. the number of layers; H the number of attention heads; dA = d/Hf H the head size; Hkv the number of K/V heads; du the FFN dimension; T₁ BN,, the number of token positions; Ae,o the active index at t; and Me. = Ar,bl We count algorithmic FLOPs as follows:
Baseline. For a step f, algorithmic FLOOPs per batch are constant across
Ftt basc.b (4BHN, 2d, df
2BNdd
4BNg Ha (Bdddd
TAV
Out
K
FFN
SureLock. Algorithmic FLOPs exhibit temporal dynamics since step-wise computation only for active positions E At,s, which changes across l:
|Ar Fbsse BN,
Mr.b
Ftt peop.b
base.b
Th
3 EXPERIMENTS
This section reports our empirical evaluation. We consider two kinds of basic tangs-la mod- eling and instruction following. In each run with different settings, profiling the per-batch sequence length N₂, and the number of active positions Me. |Arr at each step, we report computational complexity (algorithmic FLOPs). We also assess any quality degradation by locking behavior of SURELOCK. Unless otherwise noted, we set the number of sampling steps 5 = Ngens and fix the percentile for the optional confidence gate (Sec. 2.2 p to m=20%. The block length for the semi-aggressive generation option is set to N₂ Nge yielding fully parallel configuration. Temperature and classifier-free guidance scale are set to 0 by default.
these tasks induced
3.1 MASKED DIFFUSION LANGUAGE MODELS AND DATASETS.
We evaluate leading open-weight and representative MDLMs with 8 billion parameters: LLaDA- 8B-Base and LLLDAA-B--Inst (Nie et al.) 2025).
We utilize the following datasets: WikiText- 103 (Merity et al. 2016) for language modeling and MT-Bench (Zheng et al. 2023) for instruction-followiin benchmarking. These benchmarks evaluate
'Element-win operations (activations. gating. LN, softmax. RoPE) and cache 1/O/packing are lower-order relative to the dominant O(N3d) d) attention and O(N,d FFN matmuls and occur similarly across methods, so omitting them does affect relative comparisons.
5

---

Published as a conference paper at ICLR 2026
N gen
4 F
Fuase
Fyp
F Fruse
64
Se-4
.548
9.017e+11
4.936++11
0.547x
64
5c-3
.506
9.017c+11
4.565c+11
0.506x
64
5c-2
.482
9.017c+11
4.342c+11
0,482x
256
Se-4
.496
3,629e+12
1.801e+12
0.496x
256
5e-3
.482
3,629c+12
1.750c+12
0,482x
256
5e-2
.455
3.629e+12
1.725e+12
0,475x
1024
5e-4
.494
1,492e+13
7.366e+12
0.494x
1024
5e-3
.492
7.333e+12 1.492++13
1,492e+13
0,492x
1024
5e-2
.490
7.315e+12
0.490x
Table Algorithmic FLOPs, is the micro-averaged ratio of the number of active token positions.
complementary aspects: WIKITEXT-103 stresses token-level modeling on broad, category- agnostic text, whereas MT-BENCH evaluates cestruction-folloiinn ability and response helpfulness across eight categories (e.g.. "writing". "coding"); Together, they cover both core modeling capacity and practical pnstruction-folloiing behavior. Data usage settings vary by experimental set, so we explain them in the following as needed. Detailed settings are in Appendices G
3.2 RESULTS
Experiment-I:
Computational Complexity. We first evaluate the impact of SURELOCK the reduction of algorithmic FLOPs, using L.aDD--8--nnstuu on a fixed set of 32 single-turn prompts from MT-Bench, which we sampled uniformly from each category (Appendix, G). We use subset of ANT-BENCII to efficiently sweep diverse inference settings. We set batch size B = 4, Ngen {64, 256, 1024]. {5e 5e 3, 5e 2), and run record the actual number of active token positions Mr = A sequence length No per-batch per-step.
that the computational complexity steadily decreased compared to the same sampler without applying SURELOCK. The smaller the locking criterion threshold s, the lower the reduction rate; it aligns with our design intent where & is set for tightening the locking test (Sec. 2.2). Moreover, on the scale from 5e-2 to 5e-4, is less influential to the reduction of computational complexity. We also report the micro- averaged ratio of the number of active positions; f = E, E, Mr, C. B N, just as a reference to intuitively grasp how many active rows have been reduced. We see that is roughly aligned with the reduction ratio of algorithm FLOPs, suggesting that reducing the number of computable token positions leads to reducing computational load.
Experiment-2: Temporal Dynamics of Computational Complexity. Figure 2 shows the dy- namics computational complexity associated with the progression of sampling steps; logging data points was obtained from the same runs as the experiment-1. We can see that, as sampling pro- gresses, the ratio algorithmic FLOPs decreases at an accelerating rate. This is probably because surrounding most tokens are unmasked in the later steps (M, << N) leading to the stability of local- KL. We put the results on the runs with different settings for (10, 40) % in the Appendix
Trade-off between Efficiency and Generation Quality. We evaluate LLaDA- 8B-Base on Wikitext- 103, where we prompt a fixed set of text fragments, requiring the model to generate the continuation, and report Gen.-PPL measured by an external AR --de--LL 8B (Grattafiori et al., 2024). More details are in Appendix ! We also evaluate LLaDA-8BB-Instruct on MT-Bench with single-turn settings; we report the averaged LLM-ass--juu a----- score from gpt-4o for the generated responses. Other details are described in Appendix []] We set batch size B = 4, Ng.n {64, 128, 256, 512]. E = (5e 4, -3 thoughout datasets, and N, = 64 in WikiText- 103. Table 2 and Table 3 show that while computational complexity is indeed decreasing, performance is being maintained at competitive levels. For the instruct model, the MT-Bench score is essentially unchanged across all settings (at most pt. despite the reduced compute. The base model fol- lows the same trend in most configurations, but exhibits relatively larger degradation (up to 2 1.21 Gen.-PPL) for short generation lengths (Ngen E (64, 128)). It suggests that, in language modeling where prompts are likely to be open-ended, choosing very short Ngen can lead to less desirable be- havior under SURELOCK. While this can affect creative tasks like novel generation, such tasks
Experiment-3:

---

Published as a conference paper at ICLR 2026
1.0
20op baee
0.4
Max: 64, 0,0005
Mosi 256, 0.05
Mys 64, 0.0015
Masi 1024, a: 0.1005
Now
Mosi 1024,
1024,
256, 0.005
0.0
256
512
1024
Sampling Steps
Figure 2: Step-wise FLOPSs ratio. Ratio of step-wise algorithmic FI.OPs, i.e., prop/ /Fomee tently decreases as steps proceed, explaining later-step savings of computational cost.
base" consis-
Nea steps
I Gen.-PPLiase
$ Foase
Gen.-PPLprop
4 Ppop
64
32
5e-4
5.7537
4.488e+11
6.0782 (1.06x)
3.138e+11 11 (0.70x)
32
5e-3
5.7537
4.488c+11
6,2177 (1.08x)
2.788c+11 (0.62x)
64
5e-4
3,4813
8.976c+11
4.5722 (1.31x)
5,203e+11
(0,58x)
64
64
5e-3
3,4813
8.976c+11
4.7596 (1.37
4.664c+11 (0.52x)
128
128
5e-4
2,6006
1.799e+12
2.9202 (1.12x)
9.585e+11
(0.53x)
128
128
5e-3
2.6006
1.799e+12
3.1042 (1.19x)
8.964e+11
(0.50x)
128
64
5e44
3.1604
8.997e+11
3.8371
(1.21x)
5.620e+11
(0.63x)
128
64
5e-3
3.1604
8.997c+11
3.7353 (1,18x)
5.083e+11 (0.57x)
256
128
5e-4
2.3428
1.808e+12
2.6092 (1.11x)
1012e+12 (0.56x)
256
5e-3
2.3428
1.808e+12
2.6687
(1.14x)
9,497e+11
(0.53x)
256
256
5c-4
2.0127
3.616c+12
2.0486
(1.02x)
1.845c+12 (0.51x)
256
256
5c-3
2.0127
3.616c+12
2.2103 (1.10x)
1.779c+12 (0.49x)
512
256
5c-4
1,6293
3.650c+12
1.7134 (1.05x)
1.919c+12 (0.53x)
512
256
5e-3
1,6293
3.650c+12
1.7526 (1.08x)
1.851c+12 (0.51x)
512
512
5c-4
1,4658
7.301c+12
1.5248 (1.04x)
3.643c+12 (0.50x)
512
512
5e-3
1.4658
7.301e+12
1.5444 (1.05x)
3.588e+12 (0.49x)
Table 2: Generation quality of continuation sequences with and without SURELOCK on LLaDA- 8B-Base using WikiText- Gen.-PPI. is the micro averaged PPL evaluated with LLaMA-3-8B.
rarely run with very short outputs. As an additional check for more demanding downstream task, we also conduct small- scale evaluation LLLDDASB-Instruuc on the HumanEval code ande-generation benchmark (Chen et al. 2021). Here, correctness is assessed by executing the generated code against unit tests, where even small local errors can cause failure. Under a representative decoding config- uration, SURELOCK shows no deterioration in Pass@I while providing substantial compute reduc- tion, i.e., ~ 0.5xr (Appendix indicating that the observed changes in Gen.-PPL largely correspond to benign surface-level variations that do not break sentence structure semanties. We also provide different analyses in Experiment-6 on how to interpret Gen.-PPL.
Experiment-4: Runtime Performance. We evaluate End-to-end Token/Sec (E2E-TPS), which is the sustained decoding throughput of the model across multiple batches, and temporal dy- namics for step-wise TPS. See Appendix K for more details on the metrics. We follow the settings in Experiment-1 for the dataset and the model. We set batch size B [1, 2,,,,4 8), Ngen = (64, 256, 1024). ε = 5e - Figure shows that SURELOCK can achieve greater runtime gains in compute- bound areas e.g., N> 256. B>1. However, runtime gain was observed under relatively light computational load settings (N,-64, B=1); this differs from the trend in computa- tional complexity. By design, SureLock replaces dense, uniform computation with more irregular token sparse computation, so some implementation-Iewe overheads (c.g., irregular cache accesses) offset part of the FLOP savings, especially for smaller moderately sized models. As result,
7

---

Published as a conference paper at ICLR 2026
N gen steps
Scorebase
Fasse
pcorpprop
prop
64
32
Se-4
3.9
4.515e+11
3.8 (0.9%))
3.153e+11 (0.698x)
64
5c-3
3.8
4.515c+11
2.947c+11 3.7 (0.97x
(0.653x)
64
5c-4
4.2
9.030c+11
(1.00%
5.658c+11
(0.627x)
64
64
5e-3
4.2
9.030c+11
(1.02x)
5,304c+11 (0,587x)
128
5e-4
4.4
9.047c+11
4.2 (0.96x)
5,700e+11
(0.630x)
128
64
5e-3
4.3
9.047e+11
4,2 (0.98))
5,406c+11 (0.598x)
128
128
5e-4
4.8
18809+12
4,9 (1.02x)
1.044e+12 (0.577x)
128
128
5e-3
4.8
1,809c+12
5.1 (1.06x)
1.000e+12 (0.553x)
256
128
Se-4
4.2
1817e+12
(1.05x)
1,042e+12 (0.573x)
256
128
5e-3
4.3
18817++2
(0.988)
1.007e+12 (0.554x)
256
256
5c-4
4.4
3.634c+12
(1.02))
1,969c+12 (0.542x)
256
256
5c-3
3.634c+12 4.4
4.7 x1.07))
1.920c+12 (0.528x)
512
256
Se-4
4.0
3.667e+12
41.00088
1.971e+12 (0.537x)
512
256
5e-3
4.0
3.667c+12
(1.05%
1,934c+12 (0,528x)
512
512
5e-4
4.5
7.334c+12
4.7 (1.04)
3.819c+12 (0.521
512
512
5c-3
4.4
7.334c+12
4.8 (1.09x)
3.777c+12 (0.515x)
Table 3: Quality changes in generated responses by L.LDD-Instruct Scoresase and Score indi- cate the overall score for the MT--Bench with single-turn settings evaluated with gpt-4o.
3.0
Newi 64,
Myx: 256, B:
Nymi 64,
Nax: 256, B:
1024
1.36
1.39
0OM
1.3
2.5
Mini 64,
Mas: 1024,
Generation length No00
Sprop Shase
Nyni 64, B:
Mys: 1024, B:
Neni 256,
Mas:
pSSpol TPP
2.0
Newi 256, B:
1.1
256
1,23
1.30
1.35
1.30
1.5
-1.0
0.9
1.0
0.73
1.17
1.22
1.19
0.5
2
8
32
64
128
256
512
1024
Batch Size B
Sampling Steps
(a) Ratio of End-to-end TPS
(b) Ratio of step-wise E2E-TPS
Figure 3: Throughput behavior with SURELOCK: (a) end-to-end TPS ratio across different Ngen and
batch size B.; (b) per-step TPS ratio increasing as sampling progresses.
the TPS gains we report from our vanilla Py implementation are conservative; Appendix L analyzes the cause overheads in detail and discusses the possibility of hardware-specific opti-
mizations (c.g., fused kernels) that could further align FLOP reductions with wall-clock speedups. See Figure for temporal dynamics of step-wise TPS. We can see that the gain in local runtime increases as the step progresses, which is consistent with the computational trends (Figure 2)
Experiment-S: Orthogonality to Existing Approaches for FLOPs/Runime Improvement. As already discussed in Section Temporal-, Reuse-, and Selection- -ased approaches have optimiza- tion axes orthogonal to SureLock. While the number of queried positions remains proportional to N at every step, SureLock reduces the set of active positions that these operate on. Therefore, while we do not treat these acceleration methods as primary baselines, demonstrating that SureLock com- plements them in terms of FLOPs/Runtime improvements is still beneficial. For this purpose, implement a surrogate version of selection- based method (Liu al.) 2025b) that computes only on fixed fraction l of active tokens each step, and report algorithmic #L.Ps/Runtime under three set-
8

---

Published as a conference paper at ICLR 2026
SureLock Selection (k- 0.8)
t pPSpr TPSbase /TPSroppTPSSaae
FFrop/Fbase
1.30x
0.54x
1.18x
0,64
1.73x
0.28
Table 4: Algorithmic FLOPs/Runtime for different acceleration methods, evaluated using LLaDA- 8B-Instruct on MT-Bench with Ngen 256.
Nge
steps
Gen.-PPI. base
4 Fbase
4 Gen.-PPl. -prop ppop
Frop
64
64
5e-8
3,4813
8,976e+11
4.0033 (1.14x)
5,930e+11 (0.66x)
64
64
5e-7
3,4813
8,976++11
4.3489 (1.25x)
5,430e+11 (0.60x)
64
5e-4
3,4813
8.976c+11
4.5722 (1.31x)
5.203c+11 (0.58x)
64
5e-3
3.4813
8.976c+11
4.7596 (1.37x)
4.664c+11 (0.52x)
Table 5: Generation quality of continuation sequences with and without SURELOCK on LLaDA- 8B-Base on WikiText-103 for different values sf
tings: selection-based (k: 0.8), ii) SureLock, and iii) their combination (i.e., selecting the fixed fraction k = 0.8 of positions among positions not yet locked). Note that, for focused comparison, we did not search over for optimal acceleration, so the results of the selection baseline should be interpreted as conservative; our goal here is to illustrate the relative behavior. Thblee]ssows that the combination yields additional acceleration over either component alone, supporting our claim that SureLock and selection-based sparsification cooperate rather than compete.
Experiment-6: Room for Optimization of Locking Threshold. In our main experiments (Ta- ble[[and and Table 3). we deliberately used a fixed set of global locking thresholds across all configu- rations, rather than tuning it per setting, aiming to reveal the intrinsic behavior of SureLock under wide range of generation length. This choice makes it clear that the method is more effective in compute-heavy settings (longer outputs and more reverse steps), while the configuration is relatively aggressive for short generations. However, as Theorem 1, by tightening E, we can reduce the final distributional error. Here, we demonstrate the potential for e-optimization by sweeping E. We adopt the settings of short sequence language modeling using LLaDA-SB-Bal B-Basse here, which exhibited rel- atively lower quality in the main experiments. Table 5 shows that, even for short generations, an optimal & exists that minimizes quality degradation while still providing nontrivial acceleration.
Example Analysis. Feeding the same question sampled from MT-Bench the LLaDA-8B- Instruct, Figure shows the responses from the same sampler with and without SURELOCK. With settings achieving approximately 0.6x of computational complexity on average (See Table 3 we confirm that not only the quantitative score but also the qualitative quality is comparable to the baseline, with minor, localized variations. While locking might seem to introduce disrupted generation, the impact is limited. Furthermore, we confirmed in code generation that these minor differences do not disrupt the sentence syntax (Appendix D). More examples are in Appendix N
4 RELATED WORK
Computational savings for DLMs. Prior efforts to decrease the sampling cost of discrete diffu- sion LMs (DLMs) largely follow two axes- Temporal and Reuse. Temporal methods shrink the number of sampling steps T via parallel/dilated unmasking and edit- updates (Ghazvininejad
al. 2019 Gu et al. 2019, Luxembourg et al. 20255 Israel al. 2025) Huang & Tang 2025| Wei et al. 2025). ) or by adopting higher-order distilled samplers from diffusion literature, e.g., DPM- Solver (Lu et al. 2022). Reuse methods amortize work across steps by reusing or approximating intermediate states e.g., K/V caching or approximation and selective feature refresh (Ma et al. 2025 Liu et al. 2025b; b Wu et al. 2025a), Selection based partial compute reduces per-step com- plexity by updating only fixed number of subset positions, sometimes with periodic full recompu- tation for stability (Liu et al. 20256] Wu et al.] 2025a).. These two axes primarily reduce how many reuse between steps; they typically leave
steps take how much previous computation

---

Published as a conference paper at ICLR 2026
Question: Benjamin went to bookstore and purchased a variety of books. He bought 5 copies of sci-fi novel, each priced at $20, 3 copies of history book priced at $30 each, and 2 copies of a philosophy book for $45 each. What the total cost of his purchases?.
Baseline (score: 10)
SureL.ock (score: 10)
Benjamin bought 5 copies of sci-fi novel, each priced $20, so he spent * $20 $100. bought 3 copies of a history book, each priced at $30, so he spent 3 & $30 = $90. He bought copies of philosophy book, cach priced S45, so he spent $45 $90. The total cost of his purchases is $100 + $90 + $90 = $280. Summing up.
Benjamin bought copies of a sci-fi novel for $20 each, so total cost
* $20 $100. He also bought 3 copies of a history book for S30 each, so the total cost is 3 * $30 = $90. Lastly, he bought 2 copies of philosophy book for $45 cach, so the total cost is * $45 $90. The total cost his purchases is $100 + $90 + $90 $280. Con
Figure 4: Comparison of responses between Baseline vs. SURELOCK on -LaDA-8B-Insttrr with 5e 4, Nge 128, S = 128. The question is sampled from MTT-hecch with question id= 119.
the within-step active-query count constant to the end of the sampling step, so per-step compute remains bounded by N even late in sampling. Crucially, these are complementary to SURELOCK: by focusing only on the gradually contracting active positions, they could yield even greater savings.
DLMs on longer sequences. Scaling DLMs to longer contexts has been actively explored such span/block-wise masking (Arriola al. 2025). adaptive masked token insertion (Ariola al.) 2025), DLMs suite Rotary Position Embedding (RoPE) Liu et al. (2025a). This line of research di- rections yields diverse applications such as reasoning on massive code bases. By contrast, decoding longer sequences by DLMs generally requires more sampling steps. making the stationary per-step computational cost more severe. Advancing techniques such as SURELOCK, which monotonically reduces this stationary compute as sampling proceeds, could facilitate DLMs' research on longer contexts, which lags behind compared AR models (Wu et al. 20256). i.e., thousands vs. millions.
CONCLUSION
We introduced SURELOCK, a method for iterative decoding in masked-diffusion LMs (MDLMs) that locks converged token positions and thereafter bypasses their query projection and per-token FFN, yielding a monotonic reduction in per-step compute. Locked positions remain fully attendable via cached K/V vectors. The core design is local KL locking test; we justified this approach by deriving a closed-form link between the lock-time local KL and an error bound on the terminal log- pog-prbbability. On language modeling and cestruction-followwee tasks, SURELOCK achieves large reductions in algorithmic FLOPs while maintaining task quality.
Future work. Our approach is orthogonal to step-count reduction and inter-step reuse techniques, as well as selection approaches (Sec. [] such methods can operate on the progressively shrinking active positions induced by SURELOCK. Experiment shows that combining with an orthogonal ap- proach improves acceleration; full co-optimization is left for future work. Our "converge-then thhn-kokk"" idea is modality-annostic but our KL-bbased formulation assumes discrete posteriors; continuous version would require explicit local uncertainty models and new theoretical treatment.
Limitations. We report efficiency in terms of algorithmic FLOPs, a kernel- and hardware-agnostee proxy for theoretical savings that decouples the sampling algorithm from wall-clock implementa- tions. Theorem is used only to justify the KI-bbased locking rule rather than to produce numeric thresholds; such model- and data-specific tuning is left for future work. Our locking rule is based on step-wite wise distributional stability in the posterior (Sec. (.2). which does not formally guarantee that semantics are fully determined, while we observe no serious degradation in our benchmarks probably because SURELOCK only locks confidently unmasked positions (Appendix Eb
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
This paper aims to reduce computational complexity in inference time for discrete diffusion models, which may have the potential significantly impact real-world applications. Furthermore, this research does not require diffusion models to generate harmful or sexual content. We believe our work does not directly contribute to malicious use such purposes.
REPRODUCIBILITY STATEMENT
More detailed descriptions are provided the Appendices. We provide more details the project page: attps://daicban surelock
Experimental settings are included in Sec. 3
ACKNOWLEDGEMENT
This work was partially supported by JSPS KAKENHI Grant Number 25H01137 and JST K Pro- gram Japan Grant Number JPMJKP24C3.
Danushka Bollegala holds concurrent appointments Professor at University of Liverpool and as an Amazon Scholar. This paper describes work performed at the University of Liverpool and is not associated with Amazon.
REFERENCES
Marianne Arriola, Aaron Gokaslan, Justin Chiu, Zhihan Yang, Zhixuan Qi. Jiaqi Han, Sub- ham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autore- gressive and diffusion language models. The Thirteenth International Conference on Learning Representations, 2025.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374 2021.
Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask predict Parallel decoding of conditional masked language models. arXiv preprint arXiv:1904.09324, 2019.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprini arXiv:2407.2 2024.
Jiatao Gu, Changhan Wang, and Junbo Zhao. Levenshtein transformer. Advances neural infor- mation processing systems, 32, 2019.
Chihan Huang and Hao Tang. Ctrldiff: Boosting large diffusion language models with dynamic block prediction and controllable generation. arXiv preprint a4545 2025.
Daniel Israel, Guy Van den Broeck, and Aditya Grover. Accelerating diffusion Ilms via adaptive parallel decoding. arXiv preprint arXiv:2506. arXIv:2506.004133 2025.
Xiang Li. John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion- Im improves controllable text generation. Advances in neural information processing systems, 35: 4328-4343, 2022.
Xiaoran Liu, Zhigeng Liu, Zengfeng Huang, Qipeng Guo. Ziwei He, and Xipeng Qiu. Longllada: Unlocking long context capabilities diffusion llms. arXiv preprint arXiv:2506. /4429, 2025a.
Zhiyuan Liu, Yieun Yang, Yaojie Zhang, Junjie Chen, Chang Zou, Qingyuan Wei, Shaobo Wang, and Linfeng Zhang. dilm-cache: Accelerating diffusion large language models with adaptive caching. arXiv preprint arXiv:2500.06295 2025b.

---

Published as a conference paper at ICLR 2026
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in neural information processing systems, 35:5775-517 5787, 2022.
Omer Luxembourg, Haim Permuter, and Eliya Nachmani. Plan for speed-dilated scheduling for masked diffusion language models. arXiv preprint arXiv:2506. 19037, 2025.
Xinyin Ma, Runpeng Yu, Gongfan Fang, and Xinchao Wang. dkv-cache: The cache for diffusion language models. arXiv preprint arXiv:2505. /5781, 2025.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint D0XXvv:6099000000 2016.
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. arXiv preprint arXiv:2502.09990 099.00 2025.
Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Advances Neural Information Processing Systems, 37:1301 17:130136-33018 2024.
Yair Schiff, Subham Sekhar Sahoo, Hao Phung, Guanghan Wang. Sam Boshar, Hugo Dalla-torre, Bernardo P de Almeida, Alexander Rush, Thomas Pierrot, and Volodymyr Kuleshov. Simple guidance mechanisms for discrete diffusion models. arXiv preprint arXiv:2412. 2024.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. Simplified and general- ized masked diffusion for discrete data. Advances neural information processing systems, 37: 103131-103161 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
Qingyan Wei, Yaojie Zhang, Zhiyuan Liu, Dongrui Liu, and Linfeng Zhang. Accelerating dif- fusion large language models with slowfast: The three golden principles. arXiv preprint arXiv:2506. 10848, 2025.
Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie. Fast-dllm: Training-free acceleration of diffusion Ilm by enabling kv cache and parallel decoding. arXiv preprint arXiv:2505.22616. 2025a.
Yuhao Wu, Yushi Bai, Zhiqing Hu, Shangqing Tu, Ming Shan Hee, Juanzi Li, and Roy Ka-Wei Lee. Shifting long-context lims research from input output. arXiv preprint arXiv:2503.04723, 20256.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. Judging llm-as-a-judge with mt-bench and chatbot arena. Advances in neural information processing systems, 04446595-46662 2023.
12

---

Published as a conference paper at ICLR 2026
A APPENDIX
B USAGE OF LARGE LANGUAGE MODELS
We used LLMs to assist with writing paragraphs and searching literature. However, we retain full responsibility for the content in this paper.
C EMPIRICAL INVESTIGATION OF ASSUMPTION A2
The assumption of monotonic decrease of the step-wise KL divergence in the reverse step Theo- rem iss, in principle, theoretical. However, it is informative to gain better understanding of how behaves practice; we visualize the dynamics of the step-wise KL divergence Figure [ We used LLaDA-8B IB-Instt with = 5e 4, N gen 128, and S 128 on MT--eech's 's 16 samples. We can see that the Step-wise KL indeed monotonically decreases as the sampling proceeds; we can infer that A2 is not merely a theoretical assumption, and it does not significantly deviate from the practical situation.
1.0
Step-wise KL (avg.)
0.6
0.4
0.2
0.0
60
120
Sampling Steps
Figure 5: Dynamies of Step-wise KL divegence. Averaged step-wise KI. divergence measured using LLaDA-8B-Instre on MT-Bench prompts.
D DETAILED EXPLANATION OF ASSUMPTION A3.
We show that the one-step deviation of the logit vector (0th position i controlled by the previous step's drift on local posterior measured by KL divergence. Throughout, let p! E - be the token posterior position and step D(()) KL(,," p(")) and let E' VV denote the embedding matrix (rows are token embeddings). We use the block norm ||2,1 defined by || Σ, |ıح for X = [T1,...,
Reclaiming Assumption A3. There exists constant 0, depending only on operator norms of the network weights and softmax/L Lipschitz constants, such that for all steps and positions i,
| ( 222 < L DDD
Below we prove a slightly stronger global bound depending on E, V D(U) and then specialize it the local form equation I
Jusing a locality factor.
D. FROM POSTERIOR DRIFT TO EXPECTED-EMBEDDING DRIFT
For any gr E A V-1 define the expected input (embedding) z(q) = ET E Rdmodwl Then
(a(g)) x(r)]l2 || r)
E l1g
r|12
|||| // r/11-
(A.3.1)
13

---

Published as a conference paper at ICLR 2026
By Pinsker, 11q wv-rl11 r111 $ (2KL(g) with Lemb == | 2/222222222 we have
| - z(r)lle & Lemb v /KL(g))r)
(A.3.2)
Applying equation A.3.2 successive steps position
=(221)) xpp22),
| S Lemb v D
(A.3.3)
D.2 SINGLE-HEAD ATTENTION: LIPSCHITZ BOUND
Consider single-head attention with parameters (Wq, WK. Wy) and key/query dimension d₃. For input X [i,,..., xn] T the head output position is
o =
(X) = softmax
1 Vix
(Wox))Wxx)))
(A.3.4)
1=1
For two inputs X and X' = X - AX, a standard three-term decomposition gives
Do = o,(x) a'y Wr, E(ass alyy) Wv o'
E(ars -b1) Wv Azy.
(IT)
(m
(A.3.5)
We upper bound (T) and (IT) and subsume (TTT) into them (yielding conservative bound).
Term (I) By the triangle inequality.
[1]) [Wvll2 E aay | & IWV Σ A ||2 || 1| ||
(A.3.6)
Softmax Jacobian bound. Let a softmax(s) (s) and Jam (s) diag(a) aa It is classical that
sup (s)ll2 2.
(A.3.7)
achieved at two-point distributions. Hence, with mean value theorem, for score vectors s; and s/c,
||a: af. 2a S ||s; s/1l2,
|a: - -dalle S v llaco a'ol12-
(A.3.8)
Score difference.
Write the (i,j) score as
Sij Vole
(Wqz,, WKzz).
(A.3.9)
Then
Asy 1 (Wo Wx2e)) (WQTa Wx (WoAz, W Wx
[[a 2 Wllll (|An; 22 | ll2 + || | || ll |
(A.3.10)
Vdz
Assume uniform radius bound | ll. || 122 < R. for all j (true if inputs are bounded and LN is used). Summing equation A.3.10 over j and using yields
and ||AX||2,1 E, Ar, |
|Asi 12
A2
[Wolla|WWlll (Vn Rs |Axzll2 + R. ||X||2, + |Ax /22 | ddh
(A.3.11)
14

---

Published as a conference paper at ICLR 2026
Term (II. Using the triangle inequality and operator norm of Wv, have
| Σ(a osy)) Wra'sl, s [| Σ oay - dgyl ||
(A.3.12a)
Assuming | S Rg for all j. we obtain
(11)ll2 s [|Wvll2 Rs ||av a',1l1.
(A.3.12b)
Using the norm inequality ||- S vūll 111,,
[( & [[Wv l2 R2 vn lla: a'all22-
(A.3.12c)
Finally, applying the softmax Jacobian bound ||av ale S 2 ||s; - st 12 (cf. equation A.3.8 A.3.88 arrive at
we
|(IT)]|2 [IWylla y R. ||s: $4112-
(A.3.12)
Substituting equation A.3.
into equation A.3.12 we obtain
(1(1))l22
[ l R.
WollaW Wx
(Vn R. |Ax 12 - R2 11 + |Az, |X 12,1
(A.3.13a)
Using |Ar ||X||2 2,,, this simplifies
[I(12)ll S
||Wvl12- [WallalWakl 2/4
nRt
|A
(A.3.13b)
Therefore, combining equation A.3.6 and equation A...13bb
obtain the single-head bound
Ao,ll2 S Ast | Aat
I|Wvll2
[[oll [W lla
R Vd
(A.3.13)
2
D.3 BLOCK COMPOSITION: MHA + FFN RESIDUAL LN
Let one Transformer block be
B(X) X +FFN(MHA(LN((X))
(A.3.14)
Assume Lipschitz bounds for each component (w.rt. || ll2,1)
[ILN(X) - LN(X'))la. < LullX - [MHA(U)- MHA(U"))221 Amhall
U'112.19
(A.3.15)
and
FFF FFN(V"))2. S A A+IV V'll2,
(A.3.16)
For multi-head attention with & heads concatenated and an output projection Wo, one may take
Amha V ||oll
max A (head)
(A.3.17)
By the residual form equation A.3.14
[B(X) B(X') < X - X' 12,1 + Ag Amha Lin ||x - -X'1la1 = X - X'112,1, (A.3.18)
with
Lonk 1 Ar Amba Lin.
(A.3.19)
For stack of blocks.
[Bx(X)= - BH (X')ll2,1 S L x|x-|xx X'112,1
(A.3.20)
Finally, with readout z; = W.hi, we get
[|x((X) 2 (X')l2 ||Wol2 [l (X) (X)- ( (X'Ila,
|Woll2 Lite ||X X'112,1 (A.3.21)
Lnot
15

---

Published as a conference paper at ICLR 2026
D.4 PUTTING THINGS TOGETHER
Applying equation A to successive sampler states (Xs, Xs-1) and using equation A.3.3
|| 20-1112 1l2 S Luet ||x Xx- Luet Σ |
S Luet Lemb Σ
V
D a-11
j
(A.3.22)
Define Lall Luet Lemb- Equation A.3.22 is the global for equation !
|2 - - lla I Laul Σ
D
(A.3.23)
Late in iterative sampling the re-masking rate decreases and many positions become stable. Empir- ically one can upper bound the tail contribution at step s by
EV D jyfi
(i)
D
(0,
(A.3.24)
8-11
where K - an observable ratio. Substituting equation A.3.24
(1)
into equation A.3.23
| - (8)
ll2 S
Lanl (1 + x
D
Lan -0
D
(A.3.25)
=
This is exactly Assumption A3 in the main text.
Constants glance.
L = |Ellav2 x Wolll (1 + Ay Amha Lin) x(1 =
Amha S
|| max (head) head
Trema
Lust
with (bbaad) given by equation A.3.13 and Ar || Will2.
E SEMANTIC STABILITY ON KL-BASED LOCKING
Our locking rule is driven by distributional stability; lock token position when the KL diver- gence between successive posteriors at that position becomes small and remains small across steps. Conceptually, however, low KL between successive posteriors does not strictly guarantee that the semantics at that position are fully determined. In principle, one can imagine a position whose distribution remains almost unchanged across steps while probability mass oscillates between two near-synonymous tokens (e.g.. "happy" and "joyful"), which would still yield a small step-wise KL.
We regard such cases essentially benign. Choosing one of several near-synonymous options has little impact on the overall semantic content of the sentence; in this scenario, locking effectively commits one paraphrase among small set of semantically equivalent alternatives.
More concerning would be different cases when the model is still undecided between semantically distant options, such as "he" vs. "she", "yes" vs. "no", "positive" "negative". In that case, prematurely locking could have much larger impact on the downstream continuation.
In practice, however, the combination of the masked diffusion sampler and our locking policy makes this latter situation rare. First, standard masked diffusion schedulers unmask [MASK] positions in order of confidence:: at each step, they select fixed (or adaptive) number of positions with the highest posterior probability and convert them from [MASK] to concrete tokens. Positions where the model is genuinely uncertain between semantically distant options typically have relatively flat posteriors and are therefore unmasked in later steps. Second, SureLock only applies the KL-based locking rule to positions that have already been unmasked. Thus, by the time a position both unmasked and classified as "stable" by criterion, the model usually assigns high probability to narrow set of semantically consistent tokens, and any residual uncertainty is mostly between close paraphrases rather than contradictory alternatives.
16

---

Published as a conference paper at ICLR 2026
Empirically, this intuition is supported by benchmarks. Across wide range of configurations in Tables 2 and 6 (varying tasks, generation length, number of diffusion steps, and other hyperpa- rameters), enabling locking does not lead to systematic severe degradation in either language modeling metrics or instruction-following performance relative to the unlocked baseline. That means the differences or phrasing around a locked position rather than global breakdowns of topic, discourse structure, or adherence the input instruction. Appendix N provides side-by-side examples of baseline vs. SureLock outputs under our default settings.
do observe are typically localized, i.e., slightly different word choices
F INTERACTION BETWEEN SAMPLING TEMPERATURE AND SURELOCK
In our experiments, we fixed the sampling temperature T to 0 to isolate the effect of SureLock. Im- portantly, however, our locking mechanism itself does not rely on a particular sampling temperature.
The locking criterion (Sec. 2.2 is computed on the raw posterior before temperature scaling, i.c., using the distribution obtained from the raw logits (equivalently, a temperature т, when non-zero, is applied only in the categorical sampling step that materializes discrete tokens, not the distributions used for the KL-based stability test. Consequently, for fixed threshold the set of positions judged "stable" by SureLock is independent of т: increasing changes the diversity sampled outputs, but does not by itself delay or accelerate locking timing,
softmax). The sampling
One could alternatively apply the same temperature transform to the posteriors before computing KL, i... base the stability test on Pf and po-1 rather than Pr and Pt-1- Mathematically, this defines lemperaturre-deeeddenl divergence that changes which parts of the distribution differences emphasized, e.g., flattening posteriors makes KL. smaller and relaxes the notion of "stability".
In practice, however, this mainly reparameterizes the sensitivity of the locking test and plays quite similar role to adjusting the threshold æ itself. Therefore, for clarity and simplicity, we keep the locking criteria lemperature-agnosti and let control how tolerant SureLock is to distributional changes, while temperature controls decoding diversity.
G DATASET SETTINGS FOR EXPERIMENT-L
To efficiently sweep a large configuration space, in experiment-1 to profile algorithm FLOOPs, we used focused subset of MT-BBench rather than the full set. For each of the eight categories e.g., "writing" and "coding," we use the first four single-turn prompts-32 prompts in total-so that diverse phenomena may appear. Prompts are used unedited and are preprocessed with the LLADA- 8B- In-Inutttct tokenizer using its default chat template; no system message is injected.
H ADDITIONAL RESULTS FOR EXPERIMENT-2
We show additional results for different settings of m-percentile used for the optional confidence gate 2.2 we set m7 (10%, 40%), which differ from the one the main text i.e., 20%. The trend in temporal dynamics of FLOPs ratio does not change significantly with m, suggesting m is not a particularly sensitive parameter.
I DETAILED SETTINGS FOR EXPERIMENT-3
I.1 ON WIKITEXT-103
We extracted small subset from Wikitext-103 (Merity al. 2016). which contains diverse text, for the continuation generation task. First, we removed blank lines from the dataset loaded from https: htt/hugggnngfaan the first 120 records at odd indices. Next, we extracted the first 64 tokens from each record to form the prompt for the continuation. We here used LLaDA-8B-BAL
co /ataaattttaaaaeeeeee www.eooc///eeoooeoooooooooooeeee http://ales.orree/.kkittttt wikitext]] then extracted
tokenizer (Nie al. 2025).
For each prompt x, we generate a sequence y of fixed length Ng- then compute Gen.-PPL us- ing external autoregressive language model 0. This work uses Llama-3-88 (Grattafiori et al.
17

---

Published as a conference paper at ICLR 2026
1.0
Mavi 64, 0.0055
0.6
Navi 68, 0.005
Ny-i 64,
0.4
Nori 256, et 0.0005
Mgui 256, at 0.005
Maxi 256, e: 0.05
0.2
My-+ 0,0005
Muu: 0,00s
My.:
Sampling Steps z
Figure 6: Step-wise FLOPs ratio. Algorithmic FLOPs ratio step t i.c., F / base (and the averaged per-step active row ratio (micro)
E, M N. consistently decreases steps proceed, explaining later-step savings of computational cost. m-percentile is 10%.
Masi 0.0005
0.6
Masi: 0.005
Ny-: 64, 0.05
0.1
Mai: 256, xt 0.0085
My.: 256, 0.005
Masi 256, a: 0.05
Npvi 1024, e: 0.0005
Apui: 1024, v: 0.005
Nasi 1021, c: 0.05
Sampling Steps
Figure 7: Step-wise FLOPs ratio. Algorithmic FLOPs ratio step i.e., Fpmpp/ Fooss (and the averaged per-step active row ratio (micro) F; E, M N₁, consistently decreases as steps proceed, explaining later-step savings of computational cost. m-percentile is 40%.
representative autoregressive language model. We define Gen.-PPL perplexity micro- averaged over the dataset; we first compute the micro-averaged negative bog-likel NL.mmmero Σ, NLL, / lual, then PP..micro exp( (N...iicro). In the report, we used natural logarithms.
2024
I.2
ON MT-BENCH
We evaluated ALaDA-8B-Instren (Ne et al.) 2025) using all 80 records MTTbench Zheng et al. (023) with single- Annle-turn configuration used the first prompt for each record. Essentially, using the LLaDA-SB-Instre tokenizer, we applied its default chat template to each prompt to form input. We then fed these prompts the model generate responses.
Response evaluation was conducted using the LLM-as-a-judde format. Specifically, we used the of- ficial MT-Bench evaluation repository hhtpp:: / httpttuttttttt/ggg/// com/ 1m sys / FastChat/treee main/ wast / 1allm___u_pp as-is. As the judge model, we employed OpenAl's gpt-4o.
CHALLENGING SET FOR EXPERIMENT-3
We conduct additional experiments for Experiment-3 as a challenging set. We use code genera- tion, where it is examined by running generated code in sandbox environment report pass rate
18

---

Published as a conference paper at ICLR 2026
(Pass 1). Therefore, it makes it stringent test of whether premature locking induces harmful failures. We evaluate HumanEval's first 70 examples (Chen et al. 2021) under a representative configuration, i... LLaDA-8B- -antt-et- with N gen S = 256, as used in Experiment-3, Under the setting, the baseline achieves pass@) @ of 18.6, while Surel achieves 18.6-20.3 with about 0.5x efficiency in algorithmic FLOPs. We benchmarks where finite surface fluctuations directly impact performance, reducing computational cost while maintaining original performance.
that SureLock operates even in code generation
Ngon steps
1 Pass @ lusse J Fbase
t Pass@ ppop
Frop
256
256
5e-4
18.6
3.648e+12
20.3 (1.09x)
1.716e+ (0.470x) 1.676e+12 (0.459x)
256
256 5e-3
18.6
3.648e+12
18.6 (1.00x)
Table 6: Quality changes in generated responses by LLaDA-Instrun Pass base and Pass @ prop indicate the pass rate (%) of the test cases for HumanEval.
K RUNTIME METRICS
E2E-TPS measures the sustained decoding throughput of the model across multiple batches, includ- ing all inter-batch gaps and host/device launch overheads incurred during decoding. Concretely, after preparing batches we synchronize the device and start single wall-clock timer immediately before the first decoding call, and stop it after the last decoding finishes. The numerator is the total number of unmasked tokens produced by the iterative sampling process.
E2E-TPS #(total unmasked tokens]
wall-clock time
Step TPS at sampling step is computed as the number of newly unmasked tokens at that step divided by the wall-clock time, aggregating across batches:
StepTPS(() (newly unmasked tokens at step f}
At
where At is the wall-clock time for that step t. In both baseline and SURELOCK, timing uses CUDA events with synchronization.
L ALGORITHMIC CHARACTERISTICS ON RUNTIME
While our method yields substantial reductions in algorithmic FLOPs, the corresponding wall-clock speedups are more modest in some settings. This gap arises from several implementation-level overheads.
Reusing locked tokens requires irregular reads and writes to the computed kernels (e.g., Q/K/V) and caches using non-contiguous indices (Algorithm I; Line10-13,, causing scattered, poorly coalesced memory accesses and making the system memory-bound, Even when many tokens are skipped, large amount of K/V values must still be transferred between the cache and the compute kernels (Algorithm 1; Line 11), so cache traffic can dominate latency in non-compute-boure regimes. In addition, each diffusion step incurs a fixed cost for managing active and frozen tokens, such as computing posterior KL values, applying thresholds, constructing index sets, and updating the cache (Algorithm 1, Lines 17-25), which becomes non neglneglie once the per-step arithmetic workload is small. Finally, when only a small subset of tokens remains active, generic dense attention and FFN kernels operate on thin slices the sequence dimension (Algorithm 1, Lines 10 and 12), resulting in low thread/warp occupancy and underutilization of the hardware.
Our current implementations are based on standard PyTorch operators and off-the-shelf dense ker- nels commodity GPUs, without the of custom kernels, kernel fusion, specialized cache layouts. Consequently, the reported wall-clock improvements should be viewed as conservative.
In the future, believe it will be possible to further bridge the gap between computational com- plexity reduction and real-time speed improvements by introducing senddware-spccfii optimizations that are orthogonal to our algorithmic contribution.
19

---

Published as a conference paper at ICLR 2026
In particular, one could exploit (i) cache layout redesign and periodic compaction that active tokens are packed into jroock-)contguuuus regions to restore coalesced loads and stores from the K/V caches; (īi) dedicated fused kernels that, given the current active position set, jointly perform index gathering, attention and FFN computation, and writeback into the compact cache within a single launch, thereby reducing kernel-launch overhead and intermediate global-memory traffic; and (iii) overlap of cache operations with compute by issuing asynchronous copies and prefetches of cached K/V segments for the next step while computing attentions and FFN for the current step, effectively hiding much of the cache-access latency.
M OPTIONAL UNLOCKING
In SURELOCK, once locked, a token position is excluded from re-masking by construction (Sec. However, rare context shifts can render a cached posterior stale. Therefore, allowwng lightweight unlocking to detect context drift within a small finite budget and returning tokens to the active set could also be considered as a useful option. the following, we discuss an implementation for that option.
After every steps, we evaluate a budgeted proxy on locked rows only (e.g.. variable-length atten- tion with a reduced subgraph). This yields proxy uncertainty ū, and distributional drift from the locked time 5.¹)) (L)))) IpE (. Let 6; qm (ue) be the top-m% uncertainty among active tokens at step We unlock locked token i if any holds:
D.0
Sunlock
t
1-
1
D Dintervll-
(2)
Dinterval is set for suppressing failure to converge due to excessive oscillation in the locking and unlocking behavior. Furthermore, after an unlock, position i cannot re-locked for D steps and must satisfy stricter KL threshold pe (p S 1) when re-locking.
On unlock at step f, we invalidate cached K/V for and reintroduce its Query projection and FFN sublayers from step 1+1; the current step is unchanged. Amortized probe overhead attention- dominant FLOPs trade-off with performance. However, the subgraph used for probing can be reduced to size negligible compared to the dominant cost, depending its size.
Let ( log Pt
log P; (6), ||2 be deviation from no-lock decode. Unlocking cannot increase A for L' 2 because the row is computed by the same subgraph as baseline thereafter. Therefore, the error bound (Theorem Sec. 2.3) holds up to the unlock time; after that, deviation non-
increasing.
N DECODED EXAMPLES
We provide additional response examples for MT-Bench. We randomly selected one question per category.
Compose engaging travel blog post about recent trip to Hawaii, highlighting cultural experiences and must-see attractions.
Question:
SURELOCK (score: 3)
Baseline (score: 1)
Title: "A Journey Through Paradise: Unveiling
Title: "Unveiling the Paradise: A Cultural Jour- ney Through Hawaii"
Hidden Gems of Hawaii"
Introduction: As the sun set over the lush land- scapes of Hawaii, felt
Introduction:
Figure 8: Quality comparison of responses between Baseline vs. SURELOCK response. The exam-
ple question uses the record from MT-bench with question id= 81. Responses are from LLaDA-8B- Instrut with KL threshold 5e 4. Ngen 128, S 128.
20

---

Published as a conference paper at ICLR 2026
Question: Embrace the role of Sheldon from "The Big Bang Theory" as we delve into our conversation. Don't start with phrases like "As Sheldon". Let's kick things off with the following question: "What your opinion on hand dryers?"
Baseline (score: 2)
SURELOCK (score: Hand dryers.
Hand dryers are useless.
Figure 9: Quality comparison of responses between Baseline vs. SURELOCK response. The exam- ple question uses the record from MT-bench with question id= 92. Responses are from LLaDA-8B- Instrut with KL threshold g 5e - 4, Ngen 128, S = 128.
Question: Thomas very healthy, to go to hospital every day.
What could be the reasons?
Baseline (score: 3)
here could several reasons why Thomas has to go to the hospital every day, despite being very healthy. Some possible reasons include:
SURELOCK (score: 4)
There could several why Thomas has to go hospital every day despite being healthy. Some possible reasons include:
1. Chronic illness: Thomas may have chronic illness that requires daily medical attention, such as diabetes, cancer, or a heart condition.
1. Chronic illness: Thomas might have chronic illness that requires daily medical at- tention. 2. Medical treatment: He could be undergoing a medical treatment that requires daily visits to the hospital. 3. Recovery period: Thomas might be recovering from injury surgery that requires daily monitoring and care. Medical condition: He could be managing medical condition that requires daily moni- toring and treatment. 5. Medical procedure: Thomas might be undergoing medical proce- dure that requires daily visits the hospital.
2. Injury recovery: Thomas may be recover- ing from a serious injury surgery that requires daily monitoring and care.
3. Medical treatment: Thomas may undergo- ing a medical treatment that requires daily visits to the hospital, such as chemotherapy or radia- tion therapy.
4. Medical procedure: Thomas may under- going medical procedure that requires daily visits to
Figure 10: Quality comparison of responses between Baseline vs. SURELOCK response. The exam- plc question uses the record from MT-bench with question id= 103. Responses are from LLaDA- 8B-Insttut with KI. threshold 5e 4, Ngen 128, 128.
Question: When rolling dice, what is the probability that you roll total number that is least
Baseline (score: 2)
possible outcomes when rolling two dice. We need to count the number of outcomes where the sum of the dice at least
There are
SURELOCK (score: 3)
There are x 6 = 36 possible outcomes when rolling two dice. We will count the number of outcomes where the sum is at least 3. -umm (1,2). (2.1) Sum 4: (1.3), (2.2). (3.1) Sum = 5: (1,4), (2.3), (3,2), (4.1) Sum = (1,5). (2.4). (3.3).
3.
The possible sums 4, and 6.
Sum 3: outcomes (1,2). (2.1) Sum of 4: 3 outcomes (1,3), (2,2). (3.1) Sum of 5: 4 outcomes (1,4), (2.3), (3,
Figure 11: Quality comparison of responses between Baseline vs. SURELOCK response. The exam- ple question uses the record from MT-bench with question id= 114. Responses are from LLaDA- 8B-Instrut with KL threshold 5e - 4, N 128, S = 128.
21

---

Published as a conference paper at ICLR 2026
Question: Write function to find the highest common ancestor (not LCA) of two nodes a binary tree.
Baseline (score: 3)
SURELOCK (score:
Here is Python function to the highest common ancestor (HCA) of two nodes in a bi- nary tree:
Here's Python function to find the highest common ancestor (HCA) of two nodes in a bi- nary tree:
class TreeNode: def (self, valό, selfleft
class TreeNode: def _init_ val=0, selfleft
left=None, right=None): self.val left self.right right
left=None. right=None): self.val
left self.right = right
def find hca(root, q): if root is None: re- turn None if root root left find hea p. q) right find rig(root.right,
def find find hca(root, p. q): if root is None: re- turn None if root turn root left find hea((rot..ll,,,, p. q) right find hig(root.right,
root q:
root qi
Figure 12: Quality comparison of responses between Baseline vs. SURELOCK response. The exam- ple question uses the record from MT-bench with question id= 125. Responses from LLaDA- 8B-Instrut with KI. threshold = 5e - 4, Ngen 128, S = 128.
Question: Please read the paragraph below and count how many times the words "Amazon", "river", and "you" appear. Please present the results in the format of "word, number of appearances" with each word separate line. Sort the lines in order of the number of appearances. The Amazon, mesmerizing expanse of nature's wonders, is home the legendary Amazon River. Flowing through awe-inspiring landscapes like Amazon rainforest, the river weaves its way through Brazil, Colombia, and Peru, giving life to countless creatures. From the mighty jaguars prowling the Amazon jungle to the vibrant macaws soaring above the canopy, this remarkable region teems with biodiversity. Deep within the river's s currents, magnificent pink river dolphins gracefully glide alongside piranhas and electric cels. Along the riverbanks, you'll find bustling cities like Manaus, where the urban meets the wild, and lquitos, gateway to the heart of the Amazon rainforest. As you venture further, the Amazon River reveals hidden gems like the captivating Anavilhanas Archipelago, a mosaic of islands brimming with rare species. Embark on an adventure, explore the enchanting Amazon River, and immerse yourself in world teeming with life and untamed beauty.
Baseline (score: 3) Amazon, 2 river, you,
SURELOCK (score: 3) Amazon, river, you,
Quality comparison of responses between Baseline vs. SURELOCK response. The exam- ple question uses the record from MTT-bench with question id= 136. Responses are from LLaDA- 8B-Insttru with KL. threshold 5e 4, Ngen 128, S 128.
Figure 13:
22

---

Published as a conference paper at ICLR 2026
Question: The city of Vega intends to huild a bridge that will span the Vegona River, covering a distance of 1.8 kilometers. The proposed location falls within scismically active area that has experienced several high-magnitude earthquakes. Given these cireumstances, what would be the best approach constructing the bridge?
SURELOCK (score:
Baseline (score: 7)
Given seismically active nature of area, the hest approach to constructing the bridge would be to use flexible and resilient design that withstand the effects of earthquakes. This could include using a combination of con- crete and steel, as well as incorporating seismic isolation systems such as base isolation or rub- ber bearings. Additionally, the bridge should be designed be flexible and able to absorb and dissipate energy from seismic waves, rather than trying to resist them. It is also impor- tant to conduct thorough research and testing to evaluate the bridge's performance under various earthquake scenarios ensure that bridge remains safe and stable in the face of seismic activity.
Given the seismically active nature of the area, the best approach to constructing the bridge would be to that can withstand This could involve using a combination mate- rials such as steel, concrete, and timber, as well incorporating features such flexible joints and base isolators allow the bridge move in response to seismic activity. Additionally, the bridge should be designed to be lightweight and compact to minimize the risk of damage dur- ing earthquake. Finally, the bridge should be designed and constructed according to seismic codes and regulations to ensure safety and stability in the event of carthquakes.
flexible and resilient design
forces of earthquakes.
Figure 14: Quality comparison of responses between Baseline vs. SURELOCK response. The exam- ple question uses the record from MTT-bench with question id= 147. Responses are from LLaDA- 8B- with KL threshold 5e 4, Ngen 128, S 128.
Question: Which methods did Socrates employ to challenge the prevailing thoughts of his time?
Baseline (score: 6)
SURELOCK (score: 6)
Socrates employed several methods to chal- lenge the prevailing thoughts of his time. One of the most famous methods the Socratic method, which involved asking series of ques- tions to uncover underlying assumptions and heliefs of his listeners. This method encouraged his listeners to think critically and question their beliefs.
Socrates employed several methods to chal- lenge the prevailing thoughts of his time. One of the most famous methods method, which involved asking series of ques- tions encourage someone to think critically and examine their own beliefs. This method was used to uncover the underlying assumptions and belief's of those around him.
the Socratic
Another method Socrates used was concept of the "good life." which believed was the only way to achieve true happiness. He argued that the good life involved living in accordance with reason and virtue, rather than pursuing terial wealth or power.
Another method Socrates used was use irony. He would often use irony to highlight the absurdity of certain heliefs or idcas, and to en- courage others question their own assump- tions.
power of dia- logue to challenge the prevailing beliefs of his time. He would often engage in debates
Finally, Socrates also
Finally, Socrates used the concept of the "just city" to challenge the
Figure 15: Quality comparison of responses between Baseline SURELOCK response. The exam- ple question uses the record from MT-bench with question id= 158. Responses are from LLaDA- SB-Instrut with KL threshold 5e 4, .ggnn 128, S = 128.
23


Reviews:
[{"review_id": "cUHvgaAffC", "reviewer": "Reviewer_XG2i", "summary": "This paper introduces SURELOCK, a novel method to accelerate the iterative decoding process of Masked Diffusion Language Models (MDLMs). The core problem addressed is the computational redundancy in standard MDLMs, where self-attention and feed-forward networks are recomputed for all token positions at every decoding step, even for tokens that have already converged. SURELOCK proposes to dynamically identify and \"lock\" these converged tokens. The locking criterion is based on the step-wise KL divergence of a token's posterior probability distribution falling below a certain threshold. Once a token position is locked, its Query and FFN computations are permanently skipped in subsequent steps, while its Key and Value vectors are cached to allow other active tokens to continue attending to it. This reduces the dominant per-iteration computational cost. The authors provide a rigorous theoretical analysis to justify their use of a local KL-divergence criterion, proving that it bounds the terminal log-probability error. Empirical evaluations on LLaDA-8B demonstrate that SURELOCK achieves a 30–50% reduction in algorithmic FLOPs while maintaining comparable generation quality on language modeling and instruction-following tasks.", "questions": "1. Your locking criterion is based on the stability of the posterior distribution. Does this \"distributional stability\" always coincide with \"semantic stability\"? For instance, could a token position's distribution remain stable (low KL divergence) while oscillating between two semantically similar but different tokens (e.g., \"happy\" and \"joyful\")? If so, would locking it prematurely harm the final semantic nuance of the generated text?\n\n2. The experiments were conducted with a default temperature of 0. How does sampling temperature interact with the SURELOCK mechanism? Intuitively, a higher temperature would lead to flatter, less certain posterior distributions, potentially delaying the locking of tokens and reducing the efficiency gains. Have you investigated this trade-off, and is there an optimal interplay between temperature and the locking threshold ε?\n\n3. Theorem 1 provides a bound on the final error. From a qualitative perspective, how do errors from a potentially premature lock manifest in the generated text? Do they tend to be localized to the area around the locked token, or can they cause cascading failures that affect the global coherence of the sequence? Understanding the typical failure modes would be very insightful.\n\n4. The SURELOCK principle of locking converged elements seems generalizable to other modalities where diffusion models are used, such as image or audio generation. Have you considered the applicability of this method to continuous data domains? What would be the main challenges in adapting the KL-divergence-based criterion from discrete token posteriors to continuous data representations?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1.  **Well-defined Problem and an Intuitive Solution:** The paper effectively identifies a clear source of inefficiency in MDLMs—the redundant computation for already stable tokens. The proposed solution, SURELOCK, is intuitive and directly targets this issue by progressively reducing the set of actively computed tokens.\n\n2.  **Principled Algorithm Design with Theoretical Backing:** A key strength of this work is that it goes beyond a purely heuristic approach. The authors provide a theoretical justification (Theorem 1) for their locking criterion, linking the local, step-wise KL divergence to a bound on the final log-probability error. This analysis provides a solid rationale for the algorithm's design and increases confidence in its stability.\n\n3.  **Solid Empirical Validation:** The experimental evaluation is well-conducted. The authors report both theoretical efficiency gains (algorithmic FLOPs) and practical runtime performance (TPS), which is a good practice. The method is tested on both a foundational language modeling task and a more practical instruction-following task. The use of strong external LLMs for quality evaluation (LLaMA-3, GPT-4o) adds credibility to the claims that generation quality is largely preserved.", "weaknesses": "While the paper is strong overall, there are a few areas where the analysis could be deepened or clarified.\n\n1.  **Empirical Grounding of Theoretical Assumptions:** The theoretical analysis relies on several assumptions, particularly (A2) \"Geometric tail contraction,\" which states that the KL divergence decays at a geometric rate. This is a fairly strong assumption. The paper would benefit from an empirical investigation into how well this assumption holds for the models and tasks tested. Without this, the practical applicability of the derived error bound remains somewhat abstract.\n\n2.  **Gap between Algorithmic FLOPs and Wall-Clock Time:** The paper acknowledges that the significant reduction in algorithmic FLOPs does not always translate to a proportional reduction in wall-clock time, especially in non-compute-bound scenarios. While the analysis of this gap is appreciated, a more in-depth discussion on the specific implementation overheads (e.g., managing the cache, irregular memory access) and the potential for hardware-specific optimizations would make the work more impactful from a practical systems perspective.", "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": "pP8B1vTexZ", "reviewer": "Reviewer_ZEiy", "summary": "This paper proposes SURELOCK, an efficient sampling method for Masked Diffusion Language Models that reduces algorithmic FLOPs. The method is based on a simple yet effective idea: once tokens stabilize during the diffusion process, they can be \"locked,\" allowing their query projection and feed-forward sublayers to be skipped. The locking criterion is determined using step-wise KL divergence. Experimental results demonstrate that SURELOCK reduces algorithmic FLOPs by up to 50% without compromising generation quality.", "questions": "Masked Diffusion Language Models are already known for faster inference compared to autoregressive models. Is trading off generation quality for better inference efficiency a worthwhile choice for Masked Diffusion Language Models?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is simple yet effective, with experimental results demonstrating a substantial reduction in FLOPs without compromising generation quality.\n\n2. Theoretical and experimental analyses further confirm the soundness and robustness of the proposed method.", "weaknesses": "1. While the experiments demonstrate the proposed method's effectiveness, incorporating relevant baselines for comparison would further validate the soundness of this work.\n\n2. The experiments are limited to basic tasks such as language modeling and instruction following. Results on more complex tasks would further validate the method's effectiveness.", "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": "9gYW5njkAX", "reviewer": "Reviewer_BReR", "summary": "The paper proposes SureLock, a decoding-time method for masked diffusion language models (MDLMs) that permanently locks token positions once their posterior is locally \"stable\", thereby skipping Q-projection and FFN for those positions while caching K/V. This changes the dominant per-step cost from $O(N^2 T)$ to $O(MNT)$, where $M$ is the number of still-active positions. The paper establishes a local step-wise KL threshold (with optional confidence gating) for locking, and a theoretical bound for the terminal log-probability error. Experiments on LLaDA-8B (Base/Instruct) show 30–50% algorithmic FLOPs reduction with small changes in generation quality.", "questions": "Can you estimate the constants in Assumptions (A2)-(A4) so that Theorem 1 becomes more explicit and interpretable?\n\nThe proposed locking is monotone, is it possible or beneficial to consider later free to \"locked\" tokens? This might matter for short length generations.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The study is very relevant and practical for discrete diffusion language model deployment.\n\nThe algorithm has simple, orthogonal develop towards temporal/reuse accelerations and integrates with K/V reuse and selective refresh. These are novel development in my opinion.\n\nThe algorithm is clearly presented and the paper is well written.", "weaknesses": "The bound in Theorem 1 relies on many assumptions, such as geometric tail contraction and Lipschitzness. These conditions are not carefully verified.\n\nThe base model shows non-trivial PPL degradation for short outputs ($N_{\\rm gen}$ small). It may be interesting to propose adaptive locking schedules.\n\nEmpirical results are on LLaDA-8B Base/Instruct and two benchmarks. Evaluating long-context discrete diffusion models or other tasks would make the results more convincing and probe the robustness of the 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": "eBxrbrwsmc", "reviewer": "Reviewer_kJZ5", "summary": "SureLock locks converged tokens during masked diffusion LM decoding, skips their query projection + FFN for the rest of the sampling steps, and just reuses cached K/V so other tokens can still attend to them. This is supposed to drop per-step cost from O(N2d) to O(MNd), where M is the number of still-active tokens. Authors report ~$30–50$% FLOP reduction on LLaDA-8B with comparable quality. Authors also demonstrate TPS gains relative to the baseline for large batches and longer generation.\n\nOverall, I like the idea -- it is simple and intuitive,  but the contribution feels modest without wall-clock comparisons to prior methods that tackle the same problem -- inference latency. I rank this paper at 4 for now and open to raising my score if the authors add solid baselines (and ideally more shorter generation results).", "questions": "- Do authors have quantitative results for the optional unlocking variant? (Appendix H)?\n- Since a locked token cannot be revised, can you share a few cherry-picked bad cases where premature locking leads to noticeable quality degradation?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Simple idea (to its credit), training free\n- Proposed method results in TPS  speed ups for large-batch/long-generation settings.\n- The design-theoretic bound helps motivate KL as a convergence signal (even if not yet operationalized)", "weaknesses": "- Table 2, perplexity for WikiText-103: it seems that this method results in considerable perplexity drop even for 256 tokens length generation (~4-8%) and even more for shorter context.\n- For low-batch, short-generation $(N <= 256)$ proposed method shows little to no wall-clock benefit.Note that $N=256$\nis not very short (it’s a few sentences), so the absence of speedup here is concerning for interactive use.\n- Despite method is positioned as orthogonal to existing diffusion-LM accelerators (e.g., KV-cache reuse/adaptive recompute, fewer denoising steps, slow/fast schedulers), I do believe that the experiments lack direct wall-clock comparisons to at least one prior work. This makes it hard to judge the contribution.", "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)