Watermarking Diffusion Language Models

In plain words

Diffusion language models write text by filling in blanks, and they can fill the blanks in any order, not just left to right. A watermark is a hidden statistical bias in the word choices, so that a detector can later tell the text came from the model. The known watermarks for left-to-right models need the words before each choice, and a diffusion model often does not have them yet when it fills a later blank.

The paper's answer is to bias each blank according to the expected effect over the words that are still unknown. It also favours word choices that make the blanks next to them easier to bias. Detection stays the same simple statistical test, so anyone can check a text without running the model. On two open diffusion models, the authors report that the watermark is found reliably with a small cost in text quality.

TL;DR

• Red-green watermarks, which bias each word toward a "green" list chosen from the preceding context, do not transfer directly to diffusion language models, because the context is often not yet decoded §3.1 §4.1.
• At each diffusion step the method adds a logit tilt equal to a strength times the gradient of an expected green-count energy, computed over the distribution of context hashes §3.2 eq. 2 eq. 4 alg. 2.
• Detection is the red-green binomial test on the green-token count, after removing repeated (context, token) pairs §3.2.
• Author-reported: at the recommended settings, detection reaches 99% TPR@1 on LLaDA and Dream. Log PPL and judge scores are similar to a naive baseline at matched settings, but accuracy effects vary by setting tab. 1.
• The strength values in the paper's table cannot be reproduced from the released configs. See Paper vs code before running experiments.

Problem

A red-green watermark needs the context of each token to pick its green list. In a diffusion model, a token is often sampled before some of its context is fixed. Applying the watermark only where the context is already known leaves few green tokens, so detection is weak §3.1 §4.1.

The watermark must therefore act on the distribution of the unknown context, and it must work regardless of the order in which tokens are unmasked §3.1. The method also assumes that a token is not part of its own context (no self-hashing) §3.1 §3.2. A simple self-hashing variant is evaluated in app. H.

Method

Objective. The paper states the watermark as a constrained optimisation. It maximises the expected green ratio of sampled sequences, where the green ratio is the mean colour of the (context hash, token) pairs. The constraint is that the KL divergence from the model distribution stays within a budget epsilon at every position §3.1 eq. 1.

Energy. The expectation factorises, under the no-self-hashing assumption, into an energy J(q). J is a sum over positions of the hash distribution dotted with the green matrix G and the token distribution §3.1 eq. 2. G is a binary matrix with one row per hash value and one column per token.

Optimal tilt. The optimal distribution at each position is the model distribution multiplied by exp(delta times alpha), where alpha is the gradient of J with respect to that position's distribution. Delta is set by the KL budget. In logit space this is an additive tilt of delta times alpha §3.1 eq. 4. The theorem that gives this solution is in §3.1.

Practical parameterisation. The implicit equation is solved with a single fixed-point iteration §3.2 §4.1. The paper compares two ways to set the strength. The epsilon-parameterisation uses bisection on the KL budget app. G.2. The delta-parameterisation uses a constant delta and is recommended for the main experiments §4.3 app. A.1.

Hashes. Two local hash functions are used. SumHash takes the sum of the token ids at the context offsets C. Its hash distribution is the convolution of the per-offset distributions §3.2 eq. 5 eq. 6. MinHash takes the minimum of permuted token ids. Its hash distribution is a difference of survival functions §3.2 eq. 7 eq. 8. The main experiments use SumHash §4.1.

Green matrix. Each entry of G is drawn independently from a Bernoulli distribution with success probability gamma = 0.25. This differs from the correlated Bernoulli lists of the original red-green method, which keep the green-list size exact app. A.1. The paper describes G as a pre-generated matrix indexed by hash and token alg. 1.

Per diffusion step. At each denoising step, the algorithm: (i) computes the hash distribution from the current token distributions; (ii) evaluates the energy on the top-k tokens and top-k hashes; (iii) takes the gradient of the energy with respect to each token distribution; (iv) tilts the distribution by delta times that gradient, and renormalises alg. 2 §3.2. Top-k is 50 in the main experiments §4.1 app. A.1.

# One watermarked diffusion step  [alg. 2] [§3.2]
# Input:  logits (B, L, V); delta (strength); k (top-k size)
# Output: booster (B, L, V); the sampler adds it to the logits

def watermark_booster(logits, delta, k):
    # Tempered distribution per position  [code: src/dlm_watermark/
    # watermarks/diffusion_watermark.py#L315-L319]
    probs = softmax(logits / temperature)                    # (B, L, V)
    # Distribution of each position's SumHash context  [eq. 6]
    hash_probs = convolve_context_probs(probs, context_offsets)
    # Energy J over the top-k tokens and top-k hashes  [eq. 2]
    energy = expected_green_count(probs, hash_probs, k)      # (B,)
    # Gradient of the energy with respect to each token distribution  [§3.1]
    alpha = gradient(energy, probs)                          # (B, L, V)
    # Scalar strength from delta and the context size  [§3.2] [app. A.1]
    # The code divides by |C| + 1; the appendix text says |C|
    strength = delta / (len(context_offsets) + 1)
    # Booster added to the logits  [eq. 4] [alg. 2]
    return strength * alpha

Decomposition. With a single previous-token context, the tilt splits into two parts. The booster is the red-green boost averaged over the distribution of the previous token. The greenify term favours tokens that make the next position's hash more likely to be green §3.3 eq. 9 eq. 10. For an autoregressive model, the method reduces to the red-green watermark §3.3 app. E.

Detection. The detector is the red-green detector. For each token it looks up the colour of its (context hash, token) pair in G. It removes repeated (context, token) pairs, then runs a one-sided binomial test of the green count against gamma §3.2. The main text reports detection as TPR at 1% FPR §1.

Strength under a KL budget (only for the epsilon-parameterisation).

# Strength under a KL budget, by bisection  [app. G.2]
# Input:  logits (B, L, V); alpha (B, L, V); probs (B, L, V); epsilon
# Output: strength (B, L, 1)

def strength_for_kl_budget(logits, alpha, probs, epsilon):
    # Variance of alpha under the original distribution, per position
    mean_alpha = sum(probs * alpha, axis=-1, keepdims=True)
    variance = sum(probs * (alpha - mean_alpha) ** 2, axis=-1, keepdims=True)
    # Search range: from 0 up to 2 * sqrt(2 * epsilon / variance), capped
    low = zeros_like(variance)
    high = minimum(2 * sqrt(2 * epsilon / variance), 1e4)
    for _ in range(16):
        # Candidate strength  [app. G.2]
        middle = (low + high) / 2
        # KL between the tilted and the original distribution  [app. G.2]
        tilted = softmax((logits + middle * alpha) / temperature)
        kl = sum(tilted * log(tilted / probs), axis=-1, keepdims=True)
        # Shrink the range toward the side that keeps KL within budget
        if kl > epsilon:
            high = middle
        else:
            low = middle
    # Return the lower bound so that KL never exceeds the budget  [app. G.2]
    return low

Hyperparameters

Results

Author-reported. Setup: 600 WaterBench long-answer prompts app. A.1; LLaDA-8B-Instruct and Dream-v0-Instruct-7B; detection is TPR@1 (TPR at 1% FPR) §1 §4.1. Average accuracy is the mean over MMLU, HellaSwag, ARC-C, GPQA, MBPP and GSM8K, at T = 0.1 §4.1 tab. 1 app. C.5.

Source: tab. 1. "Baseline" is the naive adaptation that applies the red-green watermark only where the context is fixed §4.1. The paper reports that log PPL and the GPT-4o judge are similar for the two methods at matched settings §4.1. Per-benchmark accuracies are in app. C.5.

Detection length. For LLaDA with C = {-1} and delta = 4, detection grows with length. At about 50 tokens the method matches the baseline at about 350 tokens §4.1 fig. 2.

Robustness §4.2 fig. 3. For C = {-1}, delta = 4 and 300-token texts, detection stays strong until up to 30% of the words are deleted or substituted. Context-aware substitution hurts less than for the baseline. Paraphrasing and back-translation reduce detection, and longer texts recover a stronger signal.

Ablations §4.3 fig. 4. The hash choice has no significant effect. Using both the booster and greenify terms beats either alone. More fixed-point iterations give a very small gain at a linear cost. The epsilon-parameterisation gives much weaker detection than the delta-parameterisation. The paper suggests the KL constraint is an imperfect proxy for quality.

Infilling §4.1. The paper reports 99% TPR@1 with the same settings on DreamOn, an infilling model.

Paper vs code

Commit 380a4bef75879c28ebcfaf954f1974deb0d54980 of https://github.com/eth-sri/diffusion-lm-watermark, cloned to raw/code/2509.24368. The repository is from the paper's authors' lab (the paper links it), and it has no LICENSE file.

1. Strength values do not match. The paper gives delta = 4 for C = {-1} and delta = 5 for larger contexts §4.1. The appendix says delta is divided by |C| app. A.1. The code divides by |C| + 1 [code: src/dlm_watermark/watermarks/diffusion_watermark.py#L349-L354]. All main configs set delta: 2.0 ourWatermark_llada8b_instruct.yaml ourWatermark_dream7b_instruct.yaml ourWatermark_DreamOn7b_instruct.yaml. Under the code's rule the effective strength is 1.0 for C = {-1} and 2/3 for C = {-1, 1}. Under the appendix's rule the paper's values are 4 and 2.5. No setting in the repository reproduces the paper's stated delta. Which value produced Table 1 is not stated.
2. Dream remasking. The paper says random remasking for both models app. A.1. The Dream config uses remasking: origin [code: configs/main/Dream/ourWatermark_dream7b_instruct.yaml#L10]. In the code, "origin" transfers each masked position with probability 1 minus s over t, the same for all positions at a step [code: src/dlm_watermark/models/dream_generation_utils.py#L436-L438]. This is not the same as the paper's random strategy. Which one produced the Dream results is not stated.
3. Green list. The paper describes G as a pre-generated matrix alg. 1. The code generates the Bernoulli bits on demand with a stateless hash (splitmix64 constants, seed 42) [code: src/dlm_watermark/watermarks/random_utils.py#L40-L44] [code: src/dlm_watermark/watermarks/random_utils.py#L72-L77] [code: src/dlm_watermark/watermarks/diffusion_watermark.py#L96-L104]. The distribution is the same. The individual bits are not, so detection with another implementation needs this exact hash.
4. Detector output name. For Bernoulli lists, detect returns the green fraction under the key z_score, not a z-score [code: src/dlm_watermark/watermarks/diffusion_watermark.py#L427-L435] [code: src/dlm_watermark/watermarks/diffusion_watermark.py#L459-L461]. The p-value is under p_value.
5. Consistent points. The bisection matches the appendix: 16 iterations, the bound 2 * sqrt(2 * epsilon / variance), the cap at 1e4, and the lower bound returned diffusion_watermark.py app. G.2. Top-k of 50 on both tokens and hashes matches §3.2 [code: configs/main/Llada/ourWatermark_llada8b_instruct.yaml#L21]. The delta-parameterisation is the default in the configs [code: configs/main/Llada/ourWatermark_llada8b_instruct.yaml#L19]. Detection removes repeated (context, token) pairs §3.2 [code: src/dlm_watermark/watermarks/diffusion_watermark.py#L405-L414].

Pitfalls

• Strength scale. Do not copy delta = 4 from the paper into the config. The code's effective strength for C = {-1} is delta / 2 [code: src/dlm_watermark/watermarks/diffusion_watermark.py#L349-L354]. Calibrate the strength on your own texts with the detector you deploy.
• Temperature. The code adds the booster to the raw logits, and the sampler divides by the temperature afterwards [code: src/dlm_watermark/watermarks/diffusion_watermark.py#L191] [code: src/dlm_watermark/models/llada.py#L12-L22]. At T = 0.5 the tilt in the tempered distribution is therefore twice delta times alpha. The paper's equation does not say which distribution p is eq. 4. Whether this matches the paper's intent is not stated.
• The name delta means two things. With enforce_kl the parameter is the KL budget epsilon diffusion_watermark.py. Without it, the same parameter is the strength. The class default is enforce_kl = true, while the configs set false [code: src/dlm_watermark/watermarks/diffusion_watermark.py#L18] [code: configs/main/Llada/ourWatermark_llada8b_instruct.yaml#L19].
• Top-k truncation. The energy keeps only the top-k tokens and the top-k hashes at each position diffusion_watermark.py diffusion_watermark.py. The gradient therefore reaches only those entries through that position's own term. The paper calls top-k a practical choice §3.2, and its effect is ablated in app. C.2.
• Cost on all positions. The code applies the booster to every position from the first masked one, and its docstring says this is not optimised for positions about to be unmasked [code: src/dlm_watermark/watermarks/diffusion_watermark.py#L292-L293]. Cost is O(n L |C| |Sigma| log |Sigma|) with SumHash §3.2. The paper says the overhead is minimal app. I.
• Float32 FFT. SumHash convolutions run in float32 whatever the model dtype, and the result is cast back utils.py utils.py. The code does not clamp the output, so tiny negative values can appear in the hash distribution. The effect on the watermark is not measured in the paper.
• Accuracy cost varies. LLaDA accuracy falls from 59.4 to between 53.7 and 56.0 for the recommended watermark settings, and Dream from 50.4 to between 47.0 and 50.1 tab. 1. The paper calls the accuracy effect model-dependent and notes that some baselines lose less accuracy but detect less §4.1.
• Short texts. detect returns an empty dict when the input is not longer than the context diffusion_watermark.py. Repeated (context, token) pairs are dropped, so effective length is less than the token count diffusion_watermark.py.
• Self-hashing. Offset 0 must not be in C for the main method §3.2. A self-hash variant is in app. H.
• Remasking is part of the setup. Detection figures depend on the remasking strategy. With entropy remasking the appendix chooses delta = 3 for C = {-1} so that TPR@1 is about 1.0 at 300 tokens app. C.5.

Lineage

• Builds on: red-green watermarks of Kirchenbauer et al. (kgw in the paper's bibliography; no arXiv id given there); SumHash and MinHash are the local-hash variants of kgw2, which the paper compares with §3.2; the deduplication of detection follows "Three Bricks" [arXiv:2308.00113] §3.2.
• Compared against: the naive autoregressive adaptation §4.1; Unigram and PatternMark app. B; AAR, KTH and Gaussmark app. C.3.
• Evaluation protocol: WaterBench [arXiv:2311.07138] app. A.1.
• Models: LLaDA [arXiv:2502.09992] app. A.1; Dream and DreamOn (not linked; arXiv ids not looked up).
• Superseded by: not known.