Stopping Computation for Converged Tokens in Masked Diffusion-LM Decoding

The paper introduces SureLock, a method for Masked Diffusion-LMs that dynamically locks converged token positions to skip redundant computations while caching their keys and values, thereby reducing per-iteration complexity from $O(N^2d)$ to $O(MNd)$ and achieving 30–50% FLOP savings without compromising generation quality.

The Gist

Imagine you are directing a massive, chaotic play where the script is being written in real-time. You have a cast of N actors (the tokens in a sentence), and every single second (every step of the generation process), you ask every single actor to:

1. Listen to what everyone else is saying.
2. Think about their next line.
3. Shout their new line out loud.

Even if an actor has already decided on their line and is 100% sure they won't change it, the director (the AI) still forces them to listen, think, and shout again just to be safe. This is incredibly exhausting and wasteful. This is how current "Diffusion Language Models" work. They keep re-calculating everything, over and over, even for parts of the story that are already finished.

Enter "SURELOCK": The "Sit Down and Be Quiet" Rule.

The paper introduces a clever new method called SURELOCK. Think of it as a smart director who realizes that once an actor is confident in their line, they don't need to keep rehearsing.

Here is how it works, broken down into simple analogies:

1. The "Stable" Actor (The Locking Mechanism)

In the middle of the play, the director checks the actors. If an actor has been saying the same thing for a few seconds and their confidence is high (the "posterior has stabilized"), the director says:

"Okay, you've got your line. Sit down. Stop thinking. Stop shouting. Just stay there."

This is the Lock. The actor is "locked" in place. They stop doing the hard mental work (the heavy math called "Feed-Forward" and "Query Projection").

2. The "Ghost" Presence (Caching K/V)

But here's the magic trick: Even though the actor is sitting down and not doing any work, they are still part of the scene.

The director keeps a "ghost recording" of what that actor said (their Keys and Values). The other actors who are still standing and working can still look at this recording and say, "Oh, I see what that guy said, so I'll adjust my line accordingly."

So, the locked actors don't do any work, but they still influence the story. The active actors can still "attend" to them.

3. The Result: A Shrinking Workforce

As the play goes on, more and more actors get locked.

• Start of the play: Everyone is standing and working. (High cost).
• Middle of the play: Half the cast is sitting down. (Medium cost).
• End of the play: Only the few actors writing the final punchlines are standing. The rest are sitting quietly, watching. (Very low cost).

Because the number of people actually doing work shrinks as the sentence gets finished, the computer saves a massive amount of energy. The paper shows this can cut the computing work by 30% to 50% without making the story any worse.

The "Safety Net" (The Math Behind the Magic)

You might worry: "What if the actor sits down too early and then realizes they made a mistake?"

The authors built a Safety Net. They don't just lock an actor because they look confident; they lock them only when their line stops changing at all. They use a mathematical test called KL Divergence (think of it as a "Change Detector").

• The Rule: If an actor's line changes by less than a tiny, tiny amount (a threshold called $\epsilon$), we know they are truly done.
• The Guarantee: The paper proves mathematically that if you wait until the change is this small, the final story will be almost identical to the one where everyone kept working. The "error" is so small it's practically invisible.

Why is this a big deal?

Imagine you are baking a cake.

• Old Way: You keep stirring the batter, tasting it, and adjusting the oven temperature for every single minute, even after the cake is already baked and sitting on the counter.
• SURELOCK Way: Once the cake is baked and the timer goes off, you stop stirring. You just let it sit there while you finish frosting the unbaked parts of the cake.

Summary

SURELOCK is a technique for AI that says: "Once a word is settled, stop wasting energy on it."

It saves the computer from doing unnecessary math, speeds up the generation process, and keeps the quality of the writing just as good as before. It's like turning off the lights in the rooms of a house that no one is using, while keeping the lights on in the kitchen where the cooking is still happening.

Technical Summary

Here is a detailed technical summary of the paper "Stopping Computation for Converged Tokens in Masked Diffusion-LM Decoding" (SURELOCK), published as a conference paper at ICLR 2026.

1. Problem Statement

Masked Diffusion Language Models (MDLMs) generate text by iteratively denoising a sequence of tokens over $T$ steps. Unlike autoregressive (AR) models, which naturally cache Key/Value (K/V) states and only compute for the next token, standard MDLM sampling recomputes self-attention and feed-forward (FFN) sublayers for all $N$ token positions at every step.

• The Bottleneck: Even when many tokens have been "unmasked" (decoded) and their posterior probabilities have stabilized, the model continues to perform full $O(N^2d)$ attention computations and $O(Nd^2)$ FFN computations for these fixed tokens.
• The Waste: This results in substantial computational redundancy, as the model repeatedly processes tokens that are effectively "converged" and unlikely to change in subsequent steps.
• Limitations of Prior Work: Existing acceleration methods focus on reducing the total number of steps ($T$) or reusing K/V vectors across steps (Reuse approaches). However, these methods do not reduce the per-step spatial complexity; they still compute attention for all $N$ positions, maintaining the $O(N^2d)$ cost until the very end.

2. Methodology: SURELOCK

The authors propose SURELOCK, a method that dynamically and permanently deactivates ("locks") token positions once their posterior distributions stabilize.

Core Mechanism

1. Locking Condition: A token position $i$ is locked at step $t^*$ if its local posterior distribution has stabilized. The primary criterion is the step-wise Kullback-Leibler (KL) divergence between the current posterior $p_t^{(i)}$ and the previous posterior $p_{t-1}^{(i)}$.
   • Condition: $D_t^{(i)} = \text{KL}(p_t^{(i)} \| p_{t-1}^{(i)}) \leq \epsilon$.
   • A secondary "confidence gate" can be applied to ensure the token is also highly confident (low entropy).
2. Permanent Deactivation: Once locked:
   • The position is removed from the set of active tokens ($A_t$) for all future steps.
   • The model skips the Query projection and Feed-Forward Network (FFN) computations for this token.
   • The Key ($K$) and Value ($V$) vectors for the locked token are cached.
3. Continued Attention: Crucially, active (unlocked) tokens can still attend to locked tokens. The attention mechanism uses the cached $K/V$ vectors for locked positions while computing fresh $Q/K/V$ only for the active subset.

Algorithmic Complexity

• Baseline: $O(N^2d)$ for attention and $O(Nd^2)$ for FFN per step.
• SURELOCK: If $M_t$ is the number of unlocked (active) tokens at step $t$ (where $M_t \leq N$):
  • Attention cost: $O(M_t N d)$ (Active queries attend to all keys/values).
  • FFN cost: $O(M_t d^2)$ (Only active tokens pass through FFN).
  • As $t$ increases, $M_t$ monotonically decreases, leading to a significant reduction in computational load.

Theoretical Justification

The paper provides a theoretical bound (Theorem 1) proving that the error in the final log-probability of a token, caused by locking it early, is bounded by the square root of the KL threshold used at the locking step:
$$\|\log p_T^{(i)} - \log \hat{p}_T^{(i)}\|_\infty \leq C_{tail} \sqrt{\epsilon}$$
This establishes a closed-form relationship between the locking threshold $\epsilon$ and the maximum deviation in the final output distribution, justifying the use of local KL as a stopping criterion.

3. Key Contributions

1. Orthogonal Optimization Axis: SURELOCK introduces a new dimension of optimization (monotonic reduction of active positions) that complements existing step-count reduction and K/V reuse methods.
2. Theoretical Guarantee: Derivation of a closed-form error bound linking the local KL divergence at the lock step to the terminal log-probability error.
3. Efficient Implementation: A mechanism that caches K/V vectors to allow active tokens to attend to locked ones, ensuring the model's global context is preserved without recomputing the locked tokens' internal states.
4. Empirical Validation: Extensive experiments on LLaDA-8B models demonstrating significant FLOP reductions without compromising generation quality.

4. Experimental Results

The authors evaluated SURELOCK on LLaDA-8B-Base and LLaDA-8B-Instruct using WikiText-103 (language modeling) and MT-Bench (instruction following).

• Computational Efficiency:

  • FLOP Reduction: SURELOCK reduced algorithmic FLOPs by 30% to 50% compared to the baseline sampler.
  • Temporal Dynamics: The ratio of active tokens ($M_t/N$) decreases monotonically as sampling progresses, leading to accelerating savings in later steps (Figure 2).
  • Throughput: End-to-end tokens-per-second (TPS) improved by up to 1.39x in compute-bound scenarios (large batch sizes and long sequences), though gains were modest in memory-bound or small-batch scenarios due to implementation overheads (irregular memory access).

• Generation Quality:

  • Instruction Following: On MT-Bench, the LLM-as-a-judge scores remained virtually unchanged (difference $\leq 0.1$ points) across various settings.
  • Language Modeling: On WikiText-103, the generated perplexity (Gen.-PPL) showed minimal degradation. Shorter generations ($N_{gen}=64$) showed slightly higher PPL increases, suggesting the method is most effective for longer contexts.
  • Code Generation: On the HumanEval benchmark, SURELOCK maintained the Pass@1 score (18.6% vs 18.6-20.3%) while halving the computational cost, proving that locking does not break syntactic or logical structures.

5. Significance

• Scalability for Long Contexts: As Diffusion LMs struggle with long-context generation due to high per-step costs, SURELOCK offers a pathway to scale these models to longer sequences by making the per-step cost adaptive and decreasing.
• Hardware Agnosticism: The method reduces the theoretical algorithmic complexity (FLOPs), which is a hardware-agnostic metric. While current PyTorch implementations show conservative wall-clock speedups due to memory access patterns, the paper outlines clear paths for hardware-specific optimizations (e.g., fused kernels, contiguous cache layouts) to fully realize the speed gains.
• Complementary Nature: SURELOCK can be combined with other acceleration techniques (e.g., parallel unmasking, K/V reuse, or selection-based sparsification) to achieve multiplicative efficiency gains.

In summary, SURELOCK effectively addresses the "wasted compute" problem in masked diffusion decoding by intelligently identifying converged tokens and freezing their computation while preserving their contribution to the global attention mechanism, achieving substantial efficiency gains with negligible impact on output quality.