WeDLM: Efficient Diffusion Language Models

• WeDLM is a dual innovation that applies a principled watermarking technique for non-autoregressive generation, ensuring reliable provenance detection.
• It employs a fast parallel decoding framework using causal attention and key-value caching, significantly accelerating inference while maintaining quality.
• The combined approach optimizes token recovery and practical scalability, facilitating efficient deployment on diverse benchmarks and real-world applications.

WeDLM (Watermarked & Efficient Diffusion LLM) is the collective name for two closely related innovations in Diffusion LLMs (DLMs): a principled watermarking technique for non-autoregressive generation, and a fast parallel decoding framework reconciling DLMs with standard causal attention and Transformer prefix KV caching. The watermarking algorithm (WeDLM watermark, (Gloaguen et al., 29 Sep 2025)) enables the reliable identification of DLM-generated sequences, while the decoding system (WeDLM causal DLMs, (Liu et al., 28 Dec 2025)) allows practical deployment of diffusion-based models at scale with competitive or superior inference speeds and quality. Both capitalize on the flexibility of diffusion-style token recovery—unmasking any subset of positions at a time, in arbitrary order—while addressing the challenges of context uncertainty (for watermarking) and efficient key-value caching (for decoding).

1. Background: Diffusion LLMs and Their Operational Challenges

Diffusion LLMs generate discrete text sequences through iterative denoising starting from a fully masked string. At each step, a DLM applies a forward pass to compute independent logits $p_t^{(i)}$ for each token position over the extended vocabulary $\overline{E}=E\cup\{\text{mask}\}$, samples unmaskings, and then randomly remasks tokens for the next step. The generation algorithm admits arbitrary unmasking orders and can unmask positions even when some of their context tokens have not yet been determined (Gloaguen et al., 29 Sep 2025). This property offers hardware parallelism advantages over standard Autoregressive LLMs (ARLMs), which generate strictly left-to-right and preserve deterministic context for each step.

However, bidirectional attention—the standard in many prior MDLMs—couples masked prediction slots, breaking prefix-key/value cache (KV cache) compatibility, and impeding parallel throughput in practice (Liu et al., 28 Dec 2025). Existing AR watermarking approaches also rely on full left-to-right context, which is unavailable in early diffusion recovery steps.

2. WeDLM Watermarking: Formulation and Algorithmic Components

The WeDLM watermarking algorithm addresses the context ambiguity intrinsic to DLM unmasking. It frames watermarking as a constrained optimization:

Given

• initial factorized DLM output $p=(p_t)_{t=1}^L$,
• context-hash functions $H_t:\overline{E}^L\rightarrow \mathcal{D}$,
• a green-list bitmap $G\in\{0,1\}^{|\mathcal{D}|\times |E|}$ (drawn i.i.d. Bernoulli($y$)), and
• a KL divergence constraint $\text{KL}(q_t\|p_t)\leq \epsilon$,

the goal is to find a new distribution $q$ maximizing the expected green-token ratio:

$$q^* = \arg\max_{q\in \Delta(\overline{E})^L} \mathbb{E}_{w\sim q}\left[f(w)\right],\quad\text{subject to } \forall t:\mathrm{KL}(q_t\|p_t)\leq\epsilon$$

where $f(w)=\frac{1}{L}\sum_{t=1}^L G_{H_t(w),w_t}$ (Gloaguen et al., 29 Sep 2025).

The solution has an exponential tilt (logit shift):

$$q_t^*(u)\propto p_t(u)\cdot \exp(\lambda\cdot \alpha_t(u))$$

where $\alpha_t$ is the gradient of the expected green-token ratio $J(q)$ with respect to $q_t$, and $\lambda$ is a Lagrange multiplier determined via KKT/Lagrangian conditions (Theorem 3.1 (Gloaguen et al., 29 Sep 2025)). In practice, a single fixed-point iteration suffices for a strong watermark. The process involves:

1. Initialization $q^{(0)}\gets p$
2. For $n=1$: iterate for each $t=1,\dots,L$
   • Compute $h_t\gets \text{HashProb}(p^{(i)},t)$
   • Calculate $\alpha_t=G^\top h_t$
   • Update $$q_t^{(i+1)}(u)\propto p_t^{(i)}(u)\cdot\exp(\delta \alpha_t(u))$$

Context hash schemes include SumHash and MinHash, with efficient implementations. Expectation Boost and Predictive Bias components in $\alpha_t(u)$ respectively prefer tokens likely to be green given partial context and those which will increase future green-token likelihood (Eq. 10 (Gloaguen et al., 29 Sep 2025)).

3. Efficient Parallel Decoding: Causal Attention and Topological Reordering

The second WeDLM innovation is a parallel DLM decoding paradigm compatible with standard Transformer causal attention and KV caching (Liu et al., 28 Dec 2025). It leverages a topological reordering such that observed (unmasked) tokens are physically left-prefixed, while logical token positions are retained through decoupled positional encodings (e.g., RoPE). Each masked slot thus attends to all observed tokens using a strict causal mask:

$$M^{\text{causal}}_{ij} = \begin{cases} 1 & j\leq i \ 0 & \text{otherwise} \end{cases}$$

Training utilizes this reordering to expose masked predictions to all available context within a causally masked computation graph. The weighted cross-entropy loss is applied per mask:

$$\mathcal{L}(\theta) = -\mathbb{E}_{\gamma,x_0,M}\left[\frac{1}{\gamma}\sum_{j=1}^{N_m}\log P_\theta(x_0^{(m_j)}|\tilde{x}_{<N_o+j},\tilde{p}_{<N_o+j})\right]$$

Algorithmically, TopologicalReorder sorts observed and masked indices, constructs new token and position sequences $(\tilde{x},\tilde{p})$, and applies standard causal Transformer forward passes.

4. Streaming Parallel Decoding and Practical Deployment

WeDLM decoding at inference maintains a fixed-size window $W$ composed of both mask and filled slots, reorders the window before each forward pass, and commits any contiguous left prefix of filled tokens to the global KV cache. Mask selection for sampling uses entropy-based thresholds with distance penalty:

• For mask-slot $i$, compute $H_i=-\sum_n p_i(n)\log p_i(n)$,
• $\tilde H_i=H_i+\lambda d_i$ (where $d_i$ is slot distance to leftmost mask),
• select those with $\tilde H_i\leq\tau$.

Pseudocode sketch:

• Prefill prompt to obtain initial KV cache.
• Initialize $W$ mask tokens.
• While window not empty:
  • Reorder so filled tokens precede masks.
  • Forward pass, extending KV as needed.
  • Sample and fill low-entropy masks.
  • Refill to maintain window size.

Immediate commitment of filled tokens preserves streaming cache utilization and avoids block stalls typical in prior block-diffusion frameworks (e.g., SDAR, NBDiff).

5. Detection Methodology and Theoretical Guarantees

WeDLM watermark detection matches ARLM binomial or z-score tests (Gloaguen et al., 29 Sep 2025):

1. For each token, recompute hash $h_t=H_t(w)$.
2. Count $S$ = number of positions with $G_{h_t,w_t}=1$.
3. Under the null hypothesis, $S\sim\text{Binomial}(L,y)$ with $y$ the green-list probability.
4. Compute a one-sided $p$-value $p = P_{\text{Bin}(L,y)}(S'\geq S)$.

Critical properties:

• The one-sided binomial test controls FPR at the chosen significance level.
• Detection is exact for i.i.d. Bernoulli green lists (see App. F (Gloaguen et al., 29 Sep 2025)).
• Theoretical guarantees (Theorem 3.1) establish the optimality and uniqueness of the exponential-tilted solution.

6. Empirical Evaluation and Benchmark Performance

Benchmark results span watermark strength, detection robustness, model quality, and inference speed:

Model	Baseline TPR@1%	WeDLM TPR@1%	Log PPL Impact	GPT4 Impact
LLADA-8B, C={−1}	0.63	0.99	+0.34 → 1.90	8.43 vs 8.48
DREAM-7B, C={−1,1}	0.74	0.99	+0.24 → 2.18	7.85 vs 7.94

• Watermark TPR@1%FPR >99% with Δlog PPL ≤0.4 (Gloaguen et al., 29 Sep 2025).
• Robustness: TPR@1%FPR remains >90% after random edits up to 30% of tokens; superior performance under adversarial paraphrasing or infilling tasks compared to AR baselines.
• Sequence-length scaling: >95% TPR@1%FPR with ~50 tokens at δ=4, while baseline requires ~350.
• In decoding, WeDLM-8B achieves ≈3× speedup on GSM8K and up to 10× in low-entropy regimes against vLLM-served AR baselines (Liu et al., 28 Dec 2025).

Benchmark	Qwen2.5-7B	Qwen3-8B	LLaDA-8B	Dream-7B	WeDLM-7B	WeDLM-8B
ARC-C	89.93	92.66	81.14	88.40	90.70	92.92
GSM8K	79.23	85.97	71.80	75.97	84.76	90.20
HumanEval	59.14	68.90	31.71	20.12	68.90	75.00

WeDLM models perform comparably to AR baselines on complex reasoning and open-ended benchmarks while delivering substantial speedups.

7. Practical Integration and Significance

WeDLM architectures require no custom kernels and integrate directly with vLLM-like inference engines, leveraging causal attention and prefix-extendable KV caches:

• Standard position ID management maintains logical sequence ordering.
• Mask tokens do not leak future context under causal masking.
• No additional KV refreshes needed—window commits extend cache naturally.
• Applies to any engine supporting dynamic masks, per-step KV extend, and position IDs for RoPE.

Significance:

• Diffusion-style parallel generation can, under strictly causal reordering, outperform optimized AR engines on real hardware (Liu et al., 28 Dec 2025).
• The WeDLM watermark provides a robust, practical mechanism for provenance detection in DLM outputs with strong guarantees and empirical reliability (Gloaguen et al., 29 Sep 2025).

A plausible implication is that WeDLM principles may generalize to other diffusion-based sequence models in domains where hardware efficiency and provenance watermarking are critical.