Window-Based Token Pruning in Diffusion LMs

• Window-based token pruning is a strategy that confines inference to a dynamic local window, substantially reducing the computation needed compared to full-sequence attention.
• It employs a dual-window approach with token partitioning, active and buffer windows, and phased KV caching to optimize inference without retraining.
• Empirical evaluations on billion-parameter models demonstrate up to 6.6× speedup with minimal accuracy loss, showcasing effective trade-offs between speed and performance.

Window-based token pruning is an inference optimization strategy for diffusion LLMs (DLMs), which constrains computation to a local dynamic window over the token sequence. By exploiting the locality in the denoising process and leveraging the temporal stability of token representations, this method achieves substantial reductions in both computation and memory usage relative to standard full-sequence attention. This technique is exemplified by the Window-Diffusion algorithm, which introduces a dual-window approach with token partitioning, key-value (KV) state caching, and a phased update policy to accelerate inference without degrading model output quality (Zuo et al., 28 Jan 2026).

1. Mathematical Foundation and Token Partitioning

Let $S$ denote the maximum sequence length, and $T_0$ the number of DLM diffusion steps. At step $t$, the latent sequence is $x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$, with previously decoded tokens fixed and the remainder as masked tokens. Token positions are partitioned at each diffusion step according to their functional status and locality:

• $\mathcal{D}^{(<p)}$: tokens decoded in all prior phases (context tokens, KV-cached).
• $\mathcal{D}_p^{(t)}$: tokens decoded within the current phase (may require further stabilization).
• $\mathcal{A}^{(t)}$: "active tokens" currently refined, lying within the internal window (window length $a$).
• $\mathcal{B}^{(t)}$: buffer tokens—those in the external window of length $w$ that are neither currently active nor decoded this phase (buffer size $T_0$0).
• $T_0$1: far-field, undecoded tokens outside the external window; these are pruned and excluded from attention and computation.

Mathematically, at each inference step: $T_0$2 The internal (active) window $T_0$3 of size $T_0$4 advances as tokens within it are finalized, sliding rightward within the bounded external window $T_0$5 until the window’s contents have been decoded or a phase termination criterion is met.

2. Algorithmic Implementation

The window-based token pruning and caching algorithm, as implemented in Window-Diffusion, operates by repeatedly:

1. Initializing with a prompt and applying masks to undecoded tokens.
2. Entering phase $T_0$6 by defining the current external window $T_0$7, where $T_0$8 is the cumulative decoded prefix length.
3. Performing a refresh step (at the start of each phase or every $T_0$9 steps), recomputing forward passes over all previously decoded tokens and the current window to update the KV cache.
4. For each normal inference step within a phase:
   • Identify the active set $t$0 and buffer set $t$1.
   • Compute new Query projections for active tokens.
   • Attend over cached KV states for the concatenated context and buffer sets; far-field tokens are ignored.
   • Update predicted values (logits), select next token values, and refresh the KV cache for changed tokens.
   • Terminate the phase when the window is fully decoded or a step limit is reached, increment $t$2, and begin the next window.

Pseudocode:

$\mathcal{D}_p^{(t)}$2

3. Computational Complexity and Speedup

Full-sequence inference at each diffusion step operates at $t$3 cost due to self-attention over the entire sequence. Window-based pruning modifies this as follows:

• Refresh step: $t$4, where $t$5 is the number of decoded tokens.
• Normal steps: Each incurs $t$6 cost ($t$7, active and buffer tokens included).
• Averaged over a refresh cycle: $t$8 For $t$9, $x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$0, $x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$1, the dominant complexity per step reduces to $x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$2, yielding an idealized speedup: $x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$3 With $x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$4, halving the window results in a four-fold speed improvement. Memory requirements are similarly reduced to $x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$5.

4. Empirical Evaluation

Window-based token pruning via Window-Diffusion was assessed on LLaDA and Dream billion-parameter DLMs using benchmarks including GSM8K, MATH, HumanEval, and MBPP. Notable empirical findings include:

• With $x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$6 or $x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$7, accuracy drops by less than one point compared to the baseline, while block diffusion suffers 5–30 point degradations.
• On Dream-Instruct (w=16, refresh=32), Window-Diffusion achieves 2.3–6.6× speedup with performance matching or exceeding full-sequence inference:

Method	GSM8K (%) @t/s	MATH (%) @t/s	HumanEval (%) @t/s	MBPP (%) @t/s
Dream (full)	81.0 @14.4	39.2 @8.7	55.5 @6.1	58.8 @4.7
DKV-Cache	82.7 @17.4	39.0 @11.2	33.5 @8.1	53.2 @6.4
Fast-dLLM	81.0 @24.9	39.2 @17.4	54.9 @12.9	48.6 @10.0
Window-Diff	82.9 @32.5	38.5 @33.9	58.5 @32.3	55.4 @31.1

• Adaptive decoding (early stopping at $x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$8eos$x^{(t)} = (x_1^{(t)}, x_2^{(t)}, ..., x_S^{(t)})$9) yields up to $\mathcal{D}^{(<p)}$0 latency reduction (e.g., MBPP from 217.8s to 2.2s) with preserved accuracy.

5. Parameterization, Limitations, and Trade-offs

Window-based token pruning introduces a set of tunable parameters:

• Window size $\mathcal{D}^{(<p)}$1: Larger $\mathcal{D}^{(<p)}$2 improves accuracy but reduces throughput.
• Internal window $\mathcal{D}^{(<p)}$3: Smaller $\mathcal{D}^{(<p)}$4 reduces computation per step but can slow convergence if too small.
• Refresh interval $\mathcal{D}^{(<p)}$5: Longer intervals boost throughput but risk cache staleness, which can degrade performance.

Careful hyperparameter tuning is required. Recommended regimes are $\mathcal{D}^{(<p)}$6–$\mathcal{D}^{(<p)}$7, $\mathcal{D}^{(<p)}$8, $\mathcal{D}^{(<p)}$9. Excessively small windows or long refresh cycles cause “drift,” decreasing output quality.

A key property is that full integration with pretrained mask-based DLMs requires no retraining or architecture modifications: the dual-window and pruning logic entirely subsumes standard inference.

6. Extensions and Future Directions

Possible avenues for further development of window-based token pruning include:

• Adaptive resizing of $\mathcal{D}_p^{(t)}$0 and $\mathcal{D}_p^{(t)}$1 conditioned on model confidence or local uncertainty.
• Dynamic refresh scheduling—e.g., triggering a cache refresh when divergence metrics (such as KL divergence between token distributions) exceed a threshold.
• Combining with quantization, distillation, or additional pruning strategies for further computational gains.

A plausible implication is that window-based token pruning forms a useful abstraction for accelerating diffusion-based generative models more broadly, wherever strong locality and temporal stability of intermediate representations apply.

7. Context, Impact, and Practical Integration

Window-based token pruning exploits two empirical properties of DLM inference: (1) the rapid attenuation of dependence on distant masked tokens; (2) the stage-wise temporal stability of representations, which enables aggressive KV reuse. These insights are substantiated by detailed token-level locality and stability analyses (Zuo et al., 28 Jan 2026).

The method has demonstrated robust scaling on billion-parameter models, realizing order-of-magnitude throughput improvements with negligible or even positive accuracy differences across multiple benchmarks and datasets. Since no retraining or architectural changes are required, it is directly deployable for inference acceleration in existing frameworks for large pretrained diffusion LMs.

Potential limitations relate to staying within the operational envelope defined by window size and refresh cadence: under extreme parameter reductions, accuracy can drop. Careful task-specific selection of window and refresh hyperparameters is therefore essential for maximizing utility without compromising generative quality.