Green‑red watermarks generate a green list from a secret key derived from preceding tokens and bias token selection toward that list. Diffusion language models often decide on a token before its preceding tokens appear, and the chosen token later reshapes the green lists of its neighbors. Consequently, prior work can only tilt each selection probability to increase the expected number of green tokens at the current and neighboring positions.
We extend this objective in two ways: first, we learn vocabulary bands that serve as green lists, separating competing candidates while keeping the null distribution exact; second, we train a foresight predictor that estimates how each candidate shifts the predictions of its neighbors.
Moreover, we count green tokens exactly as the detector does: each distinct context‑token pair contributes at most one green token, so repetitions earn no extra credit.
Experiments on the WaterBench benchmark with LLaDA‑8B‑Instruct show that each learned component raises the z‑score at equal distortion, and their gains are nearly additive. With a strong watermark, detection reaches at least 96% of responses at a 1% false‑positive rate while matching the LLM‑judge score of unwatermarked text. Detection improves further on Dream and is comparable or better on code and QA tasks without a noticeable drop in task accuracy.
Review