Watermarking Diffusion Language Models with Learned Foresight and Balanced Bands

1 Introduction

• DLMs commit tokens in non-left-to-right order, so a watermark must choose before a token’s green list is known while also shaping neighbors’ lists.In LLaDA-8B-Instruct’s default order, 39% of tokens are committed while part of their context remains masked.Source: §1 · p. 1
• The unit-aware objective credits distinct (context, token) pairs once, removing rewards for repeated units that the DLM-KGW tilt contains.It tilts candidates for their own greenness and the greenness they induce at positions that use them as context.Source: §1 · p. 2
• Existing methods use neighbors’ pre-placement beliefs, although placing a candidate changes the most likely token of a masked neighbor in about three quarters of uncertain commits.The two components learned here are the vocabulary partition and the beliefs used for foresight.Source: §1 · p. 2
• Learned bands separate competing candidates while a keyed hash still selects the green band, preserving the exact null distribution.A foresight predictor, initialized from the generator’s top layers and trained by distillation, estimates neighbors’ beliefs after each candidate is placed.Source: §1 · p. 2
• At δ = 2, learned bands and the predictor raise TPR@1% by 15 and 6 points, and by 21 together.Their z-score gains also add across four further keys; at high strengths, the method detects almost every response while matching the unwatermarked text’s judge score.Source: §1 · p. 2

2 Preliminaries

• A masked DLM predicts a vocabulary distribution at each masked position and commits selected tokens, which then stay fixed.The decoder can commit a position before some left neighbors and after some right neighbors.Source: §2 · p. 2
• Green-red watermarking adds strength δ to green-token logits, and a key-based detector counts green tokens to test for excess over chance.The detector’s z-score is approximately standard normal without the watermark.Source: §2 · p. 2
• Detectors count each distinct (context, token) pair, or unit, once because a repeated unit adds no new evidence.DLM-KGW instead tilts candidates according to expected greenness at their own and neighboring positions.Source: §2 · p. 2
• The overview separates offline module learning, generation-time tilting, and detection, which needs only the key and bands.The tilt combines the candidate’s own and foresight terms.Source: Figure 2 · p. 3

3.1 Commits as evidence maximization

• A candidate earns evidence from its own greenness and the greenness it induces at the next two dependent positions; repeated units receive no reward.The unit-aware correction sets the green probability of an already-seen unit to γ, including at dependents.Source: §3.1 · p. 3

3.2 Learned bands

• Learned bands assign each token to a band while a key selects the green band, preserving an exact null distribution for text independent of the key.The partition is trained to spread likely candidates across bands, which improves the tilt's evidence gain.Source: §3.2 · pp. 3–4

3.3 Foresight predictor

• A foresight predictor estimates neighbors’ beliefs after a candidate is placed, using a distilled network rather than the same current beliefs for every candidate.Exact post-placement beliefs require a generator pass per candidate; the predictor processes candidates in one batched pass.Source: §3.3 · p. 4

3.4 Decoding and detection

• At each decoding step, the method tilts likely candidates, recomputes the low-confidence winner until it carries its predicted foresight, then detects distinct units in the final text.Special tokens are treated as red, and the detector computes the green-count z-score and exact p-value.Source: §3.4 · p. 4

Data and models

• Watermarking and evaluation use LLaDA-8B-Instruct on 600 prompts spanning long-form QA, finance QA, and instruction following, with up to 256 generated tokens.Bands and the predictor are learned on disjoint prompts from the same sources; generation uses 32-token blocks at temperature 0.5 and LLaDA’s low-confidence order.Source: §4 · pp. 4–5

Baselines

• The comparison uses the official implementations of DLM-KGW and dgMARK, with DLM-KGW’s main configuration and greedy or top-3 sampling for dgMARK.DLM-KGW uses γ = 0.25 and strengths δ ∈ {2, 3, 4, 5}; dgMARK watermarks the decoding order.Source: §4 · p. 5

Metrics

• Detection is measured by TPR at calibrated 1% and 0.1% FPR, alongside log-perplexity, judge ratings, mean z-score, and repetition-loop share.Quality uses Qwen3-8B log-perplexity and Gemma-4-12B-it ratings; comparisons also use equal log-perplexity and equal KL.Source: §4 · p. 5

5.1 Comparison with prior DLM watermarks

• At δ = 2 and 3, ours detects 12.7 and 4.5 points more responses at 1% FPR, and 11.9 and 10.3 points more at 0.1% FPR, than DLM-KGW.These significant differences are averaged over five keys on 300 prompts; at equal log-perplexity, all eight differences favor ours, four significantly.Source: Figure 3 · p. 5, §5.1 · p. 5, Table 1 · p. 5
• At δ = 4 and 5, both methods detect 96–98% of responses, while ours receives judge scores of 9.32 and 9.26 versus 9.12 for unwatermarked text and 8.62 and 8.15 for DLM-KGW.The paper attributes most of the quality gap to DLM-KGW rewarding repeated units that the detector counts only once.Source: §5.1 · p. 6

5.2 The two learned components

• Learned bands raise mean z-score by 0.41–0.47 at every strength without raising log-perplexity, while the predictor raises it by 0.10–0.28 with either partition.At δ = 2, bands detect 14.7 points more responses at 1% FPR; the predictor detects 5.8 points more, significantly, while its TPR changes by less than a point at δ = 3 and 4.Source: §5.2 · p. 6
• Together, the learned components yield 20.5 and 4.8 points more detections at 1% FPR at δ = 2 and 3, with no significant change in log-perplexity or judge score.Their mean z-score gains are close to additive, and they require 30% and 22% less KL for the same mean z-score.Source: §5.2 · p. 7

5.3 Why post-placement beliefs matter

• On 2,397 held-out decision states, placement changes a masked neighbor’s most likely token in 74% of cases; the predictor recovers that token in 61% and ranks the exact best candidate first in 54%.Current beliefs rank the exact best candidate first in 28% of states and explain none of exact foresight’s variance.Source: §5.3 · p. 7

5.4 Other tasks, another model, robustness

• On code and QA, ours detects 7.5 points more programs and 5.0 points more answers on tasks used for learning, and as many or more on unseen tasks, without significant losses in pass@1 or accuracy.The learned components were trained on code and QA prompts; the evaluation includes held-out MBPP and MMLU-Pro and unseen HumanEval and GPQA-Diamond.Source: §5.4 · p. 7, Table 3 · p. 8
• On Dream, transferring LLaDA’s partition detects 57.3% and 75.5% of responses at δ = 3 and 4, versus 42.0% and 56.5% for DLM-KGW; Dream-trained components reach 60.0% and 78.3%.The predictor raises mean z-score over current-belief foresight by 0.22 and 0.15, both significant.Source: §5.4 · p. 7
• Under 20% random deletions or insertions, ours detects about as many responses as DLM-KGW in the same order, and its AUC stays above 96% under every token edit.Under 20% substitutions, it detects 6.7 points fewer; after paraphrasing, both token-level watermarks retain AUC of 67–75%.Source: §5.4 · p. 8

• Prior token-level DLM watermarks use a keyed-hash partition and, when looking ahead, neighbors’ beliefs before candidate placement.Other approaches use masked neighbors’ likely tokens, revealed neighbors, or adjacent unmasked positions; another line watermarks decoding order, sampling noise, key patterns, or signal vectors.Source: §6 · p. 8
• Green-red watermark alternatives include distortion-free schemes, list-sharing methods, learned token-specific biasing, and fixed or random vocabulary lists.Examples include WatME, Balance-Marking, Token-specific watermarking, SIR, Unigram-Watermark, and distribution-preserving random segments.Source: §6 · p. 8
• Ours learns a partition that separates competing decoding candidates while a keyed hash still selects each context’s green band.Detection needs only the key and partition.Source: §6 · p. 8

• Learning the partition and post-placement foresight beliefs raises detection, with additive gains, while unit-aware counting keeps the quality of unwatermarked text.The keyed green-band selection preserves an exact null distribution, and detection needs only the key and bands.Source: §7 · p. 8

• The watermark loses most of its signal under paraphrasing.Under random deletions and insertions, it detects about as many responses as DLM-KGW.Source: Limitations · p. 9
• Context width trades detection without edits against robustness to token substitutions.Source: Limitations · p. 9
• The predictor takes about 29 GPU-hours to train once per generator and makes generation about 1.5 times slower than current-belief foresight.Exact foresight is at least 2.15 times slower.Source: Limitations · p. 9
• Evaluation uses open-model judges rather than human raters, and attacks targeting the secret partition are left for future work.Two judge families’ scores correlate at 0.84 over responses and 0.99 over methods.Source: Limitations · p. 9

Ethical considerations

• Watermarking can support attribution and deter misuse, but it can also enable tracking, and false positives can harm people whose writing is flagged.The authors report false-positive rates on human-written text and advise against acting on a single detection without other evidence.Source: Ethical considerations · p. 9

A Experimental protocol

• Evaluation uses 600 prompts, paired prompt-level comparisons, and bootstrap confidence intervals; detection thresholds are empirical quantiles of unwatermarked responses.For paired comparisons at equal log-perplexity or KL, the reference configuration is linearly interpolated within each bootstrap resample.Source: A Experimental protocol · pp. 12–13

B Proofs

• The proofs establish that the one-commit objective is maximized by the Gibbs distribution and that the sampler optimizes commit by commit, not over the full decoding.For green counts, distinct contexts yield independent uniform bands under a random key, supporting the stated null distribution.Source: B Proofs · p. 13

C.1 Learning the partition

• On held-out records, the learned partition’s average gain was 0.432, versus 0.413 for a band network and 0.402 for a keyed hash.Learned bands held 24–26% of model probability mass each, and alternative random initializations reached nearly the same gain but produced partitions barely more similar than random.Source: §C.1 · pp. 13–14

C.2 Decoding

• The detector counts each distinct context-token unit once and computes an exact p-value from the null distribution across distinct contexts.Special tokens are excluded from detection and treated as red during tilting; crediting them as ordinary tokens shortened responses by 17–35 tokens on average in a screen.Source: §C.2 · p. 14, §C.3 · p. 15

C.3 Detection and baselines

• For learned bands, the green-count null distribution is exact over a random key, and its upper tail at the observed count gives the p-value.The detector re-tokenizes responses and counts each distinct context-token unit once; counting only the first occurrence of each context would discard over a tenth of units.Source: §C.3 · p. 15

D The foresight predictor

• On 1,329 evaluation decision states, the predictor ranked the exact best candidate first in 54% and explained 47% of exact-foresight variance, versus 27% and −5% for current beliefs.On validation records, neighbor top-1 agreement rose from 55.9% after one epoch to 61.2% after five, compared with 25.5% for current beliefs; the deployed checkpoint reached 60.8%.Source: D The foresight predictor · p. 16

E Null distribution and calibration over keys

• At δ=2/3, ours exceeded DLM-KGW by 12.7/4.5 points in TPR@1% and 11.9/10.3 points in TPR@0.1% under thresholds pooled over 100 keys.With each key’s own thresholds, the respective advantages were 9.9/5.1 and 13.4/13.2 points. Learned bands kept human-text green rates within half a point of 25% on every corpus, while the unkeyed learned list produced 28–31%; thresholds were pooled over keys because fixed-key calibration varies with shared text units.Source: E Null distribution and calibration… · pp. 17–18, Table 10 · p. 18
• Key-dependent score shifts make single-key calibration unreliable, so the authors use pooled scores; the official key’s random-order threshold is 2.31 versus 2.37 pooled over keys.With detectors’ own nominal p-value thresholds, ours exceeds DLM-KGW by 21.2 and 5.7 points in nominal TPR@1% at δ = 2 and 3; dgMARK’s threshold uses 600 unwatermarked scores, and at 0.1% FPR it is essentially their maximum.Source: E Null distribution and calibration… · p. 18

Low strengths

• Table 13 reports ours-minus-DLM-KGW paired differences on WaterBench-600 with LLaDA-8B-Instruct in low-confidence and random orders, at the same δ and equal log-PPL, with paired-bootstrap 95% confidence intervals.Bold estimates have confidence intervals excluding zero. Dashes indicate ours’ mean log-PPL is below DLM-KGW’s range for equal-log-PPL comparisons or that the judge was not run in random order.Source: Table 13 · p. 20

Repetition loops

• For DLM-KGW at δ = 2–5, looping responses have mean z-scores 0.4–1.1 lower than non-looping responses, and loops lower mean log-perplexity for every method.Under the appendix loop rule, DLM-KGW’s looping share rises from 4.8 to 11.7% with δ, versus 3.7% without a watermark; ours stays at 1.5–2.0%.Source: F Additional WaterBench results · p. 18, Table 14 · p. 20

Second judge and similarity

• The two judges’ quality scores correlate at 0.84 across 6,600 responses and their method means at 0.99 across 11 methods.Qwen3-32B rates ours above DLM-KGW at every strength, by 0.12, 0.27, 0.49, and 0.83 at δ = 2 to 5, and at least as high as unwatermarked text.Source: F Additional WaterBench results · p. 19

Repetition loops

• The two judges agree closely: quality scores correlate at 0.84 over 6,600 responses and their means at 0.99 across 11 methods; similarity correlates at 0.85 over 7,200 pairs.Source: F Additional WaterBench results · p. 19

Per task

• Ours detects more responses than DLM-KGW on ELI5 and FiQA at δ = 2 and 3, and as many at δ = 4 and 5, when both detect 99–100%.On AlpacaFarm, neither method detects consistently more; DLM-KGW lengthens responses at higher strengths, while ours keeps their length.Source: F Additional WaterBench results · p. 20

Random order

• In random order at δ = 3, ours matches DLM-KGW in TPR@1% and exceeds it by 7.0 points in TPR@0.1%; at δ = 4, it detects slightly fewer responses, not significantly.At equal log-perplexity at δ = 4, ours detects 6.6 and 20.3 points more in TPR@1% and TPR@0.1%, respectively; δ = 3 cannot be compared at equal log-perplexity.Source: F Additional WaterBench results · p. 20, Table 13 · p. 20

Unit-aware score and decoding choices

• The unit-aware score raises mean log-perplexity by about 0.03, while its effect without loops is at most 0.016 and not significant.It reduces repetition loops, whose log-perplexity is low.Source: G Paired contrasts of the crossed… · p. 20, Table 18 · p. 23

Additivity

• The components’ interaction in mean z-score is −0.07, −0.10, and +0.10 at δ = 2, 3, and 4, none significant.The interaction is defined as the predictor’s gain with learned bands minus its gain with the keyed hash.Source: G Paired contrasts of the crossed… · p. 21

Second judge

• The two components together change judge scores by −0.07 to +0.11 under either judge, never significantly; the only significant decrease is −0.15 for learned bands with current-belief foresight at δ = 4 under Gemma.Qwen3-32B assigns −0.08 to that same setting, not significantly.Source: G Paired contrasts of the crossed… · p. 21, Table 20 · p. 24

Without the unit-aware score

• With the plain score and current-tilt order, learned bands and the predictor together raise mean z-score by 0.80 and 0.68 at δ = 3 and 4, respectively.The learned partition also lowers judge score by 0.17 and 0.22; ours is rated 0.18 and 0.34 higher than the plain-score, current-tilt variant.Source: G Paired contrasts of the crossed… · p. 21, Table 19 · p. 24

Exact foresight

• Exact foresight raises mean z-score over the predictor by 0.15, 0.31, and 0.14 at δ = 2, 3, and 4, respectively, at the same δ and about the same KL.Source: G Paired contrasts of the crossed… · p. 21, Table 21 · p. 24

Foresight at equal KL

• At equal KL at δ = 3, predicted foresight adds 0.94 to mean z-score versus no foresight, while current-belief foresight adds 0.80.At δ = 4, both use more KL than the no-foresight cell at its strongest setting, so those comparisons extrapolate the latter linearly.Source: G Paired contrasts of the crossed… · p. 21

Detection–quality frontier

• At equal mean z-score, learned bands, the predictor, and both together require 22%, 18%, and 30% less KL at δ = 2, and 16%, 12%, and 22% less at δ = 3.At equal log-perplexity, both together detect 4.7 and 5.1 points more at δ = 3 and 4; at equal KL, they raise mean z-score by 0.48 at δ = 2 and 3.Source: G Paired contrasts of the crossed… · p. 22, Figure 5 · p. 22, Table 17 · p. 22

Further keys

• Across four further keys at δ = 3, learned bands and the predictor together raise mean z-score by 0.50 with no significant log-perplexity change; their 0.02 interaction is not significant.The learned-partition gain varies more across keys than the predictor’s gain.Source: G Paired contrasts of the crossed… · p. 24

Context width

• At δ = 4, two-token contexts give the highest mean z-score and lowest log-perplexity; wider contexts are less robust, while one-token contexts are most robust but detect less without edits.The analysis uses no foresight, and its nominal TPR@1% robustness rates are not comparable with the other appendix robustness results.Source: H Design ablations · pp. 24–25, Table 23 · p. 25

Number of bands

• With eight rather than four bands, KL is 46.3 rather than 64.3 nats per response at δ = 4, and mean z-score is 0.17 higher at equal KL, significantly.At equal mean log-perplexity, eight bands are 0.19 lower, not significantly; the paper retains four bands.Source: H Design ablations · p. 25, Table 24 · p. 26

Foresight weight and predictor gate

• At δ = 4, increasing foresight weight from 0.25 to 0.5 raises mean z-score from 5.01 to 5.93 at unchanged mean log-perplexity, while increasing it to 1 raises log-perplexity by 0.18 and KL by 40 nats per response.At δ = 3, the predictor’s gain over current beliefs grows from 0.05 at λ = 1/4 to 0.26 at 1/2 and 0.37 at 1, at about the same KL; a gate of 0.02 avoids most predictor calls.Source: H Design ablations · pp. 25–26, Table 25 · p. 26

Foresight beyond the dependents

• Extending foresight to offsets 1–6 raises expected evidence by 12% at δ = 3 and 11% at δ = 4, but raises KL by 7% and 6%, so evidence per KL grows only about 4%.Adding the two dependents to the own term raises expected evidence by about 30%; the longer horizon would require predicting beliefs farther away.Source: H Design ablations · p. 26

Evaluation scope

• HumanEval and GPQA-Diamond are out of distribution for the code/QA-trained components because no training prompts were available for those tasks.The evaluation uses zero-shot LLaDA-8B-Instruct; code is scored by execution against tests and QA by parsed answer accuracy.Source: I Code, question answering and Dream · pp. 26–27

Code and question answering

• On MBPP, ours detects more responses than DLM-KGW by 7.5 and 10.0 points with the code/QA-trained and WaterBench-trained components, respectively, both significantly.With code/QA training, detection rises on HumanEval, MMLU-Pro, and GPQA-Diamond by 4.9, 2.1, and 3.0 points, but falls on MBPP by 2.5 points; most per-task differences are within sampling noise.Source: I Code, question answering and Dream · p. 27, Table 26 · p. 27
• Code/QA training halves the MMLU-Pro accuracy drop from 7.9 to 4.3 points below unwatermarked responses, with neither difference significant.On GPQA-Diamond, unwatermarked accuracy is 20.7%, below the 25% random-guessing rate, so differences there mostly reflect noise.Source: I Code, question answering and Dream · p. 27, Table 26 · p. 27

Dream results

• On Dream, ours reaches nominal TPR@1% of 60.0% and 78.3% at δ = 3 and 4, versus 42.0% and 56.5% for DLM-KGW, with both differences significant.Most of the detection difference comes from response validity; short responses are 13% of ours versus 44% of DLM-KGW at δ = 4.Source: I Code, question answering and Dream · pp. 27–28, Table 27 · p. 28

Dream components

• On Dream, the predictor raises mean z-score over current-belief foresight by 0.22 and 0.15 at δ = 3 and 4, and nominal TPR@1% by 4.5 and 2.2 points, with the latter not significant.The Dream-trained partition detects about as many responses as the transferred LLaDA partition; neither difference is significant.Source: I Code, question answering and Dream · p. 28, Table 27 · p. 28

J Out-of-distribution writing tasks

• The out-of-distribution evaluation uses 300 prompts each from C4 RealNewsLike news continuations, Dolly creative-writing instructions and WritingPrompts, with components learned on WaterBench-style prompts.Thresholds are calibrated on unwatermarked responses to the same prompts; dgMARK uses its own unwatermarked top-3 responses.Source: J Out-of-distribution writing tasks · p. 28
• The authors attribute the smaller high-strength gain to distribution shift: news and story prompts yield candidates and post-placement beliefs unlike those used to train the components.They state that extending genre coverage requires adding unlabeled prompts to training, without labels or human annotation.Source: J Out-of-distribution writing tasks · p. 28
• Across 900 out-of-distribution prompts, ours exceeds DLM-KGW TPR@1% by 6.0 points at δ = 2, but trails by 3.7 at δ = 4.The log-perplexity difference is −0.010, 0.012 and 0.061 at δ = 2, 3 and 4; repetition loops rise from 14% to 26% for DLM-KGW but stay between 10% and 12% for ours.Source: J Out-of-distribution writing tasks · p. 28, Table 28 · p. 29

K Robustness

• With 20% token deletions or insertions, ours detects about as many responses as DLM-KGW; with substitutions, it detects 6.7 points fewer, while AUC stays above 96% under each edit.The substitution gap is attributed to an edited token changing units at three positions rather than two with two-token contexts.Source: Table 29 · p. 29, K Robustness · p. 29
• After paraphrasing, only 2.2–8.5% of responses are detected at 1% FPR; ours and DLM-KGW retain AUC of 67–75%, versus 52–58% for dgMARK.Paraphrasing uses Qwen2.5-32B-Instruct with the dgMARK prompt.Source: Table 29 · p. 29, K Robustness · p. 29

L Resources

• Experiments use LLaDA-8B-Instruct and Dream-v0-Instruct-7B; partition and predictor training prompts come from WaterBench sources, ELI5, FiQA and AlpacaFarm.Predictor data collection takes about 23 GPU-hours and training about 6; learning the partition takes a few minutes on one CPU core.Source: L Resources · pp. 29–30
• The authors use artifacts for non-commercial research where terms allow it, release code under the MIT license, and do not redistribute human-written answers or generated texts.The code of Gloaguen et al. (2026) states no license; the authors use it only to reproduce the baseline and do not redistribute it.Source: L Resources · p. 30

M Judge prompts

• Figure 6 gives quality-judge instructions, while Figure 7 gives similarity-judge instructions; quality scores average style and consistency, and similarity compares with the same prompt’s unwatermarked response.Replies missing the three quality scores are regenerated up to twice, and all 19,032 judged responses received those scores.Source: M Judge prompts · p. 30, Figure 6 · p. 31, Figure 7 · p. 31