WaterMod: Modular Token-Rank Partitioning for Probability-Balanced LLM Watermarking

AI summary

Overview

Research area: AI-generated text provenance and watermarking for large language models (LLM watermarking).

Technical level: Intermediate. The paper assumes familiarity with logit-based decoding, softmax probabilities, Shannon entropy, and z-score hypothesis testing, though each is defined in a preliminaries section.

Scope: The paper introduces WaterMod, a watermarking framework that partitions a probability-sorted vocabulary using modular arithmetic (rank mod k) to support both zero-bit attribution and multi-bit payload embedding, and evaluates it on natural language continuation, mathematical reasoning, and code generation.

What This Paper Is About

Logit-based watermarking typically splits the vocabulary into a pseudorandom green list and red list at each decoding step and boosts the green list's logits, but this random split can put the highest-probability token in the red list and harm fluency. WaterMod instead sorts tokens by the model's own probability and partitions the resulting ranks by rank mod k, which spreads adjacent (semantically similar) tokens across different classes so a high-probability token always remains available. The same modular rule supports both binary attribution (k=2) and multi-bit payload embedding (k>2), with log2 k bits embeddable per decoding step.

Key Contributions

1. Probability-ranked modular partitioning. WaterMod replaces heuristic or hash-based green/red vocabularies with a collision-free, deterministic partition of ranks by rank mod k. The paper states this needs no synonym dictionaries, hashing tricks, or prompt engineering, and yields class sizes that differ by at most one.
2. Entropy-adaptive zero-bit scheme (k=2). An entropy-driven gate selects either the even or odd parity as the green list, using p_odd = (H_t / H_max)^H_scale, guaranteeing at least one high-probability token stays in the green class.
3. Unified multi-bit extension (k>2). A payload digit d selects the residue class satisfying rank mod k = d, embedding one base-k digit (equivalently log2 k bits) per decoding step, with majority voting over color-position tallies recovering digits as a form of error correction.
4. Combined detection and recovery in one pass. The same color-position tally table supports both z-score-based attribution and payload retrieval, with the z-score computed as z = (G - T·p0) / sqrt(T·p0(1-p0)) and p0 = 1/k.

Main Findings

• Zero-bit fluency on C4: WaterMod achieves the lowest perplexity among compared methods at 12.58, with an AUROC of 87.09. SynthID-Text reaches 12.77 perplexity and the highest C4 AUROC at 94.36; LSH has AUROC 88.03 but perplexity 26.19, described as more than twice WaterMod's.
• Zero-bit math reasoning on GSM8K: WaterMod records accuracy 53.83 and AUROC 100. Compared to SynthID-Text, described as the strongest competing method by AUROC, WaterMod improves accuracy by 13.06%.
• Zero-bit code generation on MBPP+: WaterMod reaches AUROC 82.66, which the paper says outperforms the next-best AUROC baseline KGW by 14.12%, while improving pass@1 by 23.07%. Relative to LSH (AUROC 30.72, highest pass@1 41.30), WaterMod improves AUROC by 169.07%.
• Multi-bit versus MPAC: On C4, WaterMod records perplexity 10.87 and AUROC 98.02 versus MPAC's 10.88 and 97.78. On GSM8K, WaterMod reaches accuracy 40.33 and AUROC 96.94 versus MPAC's 31.77 and 95.05, a 26.94% accuracy improvement. On MBPP+, WaterMod reaches AUROC 98.29 versus MPAC's 48.40 (a 103.07% relative improvement) and pass@1 26.20 versus 20.60, a 27.18% improvement.
• Robustness to paraphrasing: Under a ChatGPT (gpt-4o-2025-04-14) paraphrase attack on GSM8K, WaterMod's mean z-score falls from 14.89 to 9.95, while human-written text scores 0.09. AUROC on watermarked text drops to 99.95 (from 100.00), which the paper reports as a 0.05 absolute drop.
• Entropy differences across tasks: The paper reports that GSM8K and MBPP+ exhibit lower token-level entropy than C4 under both Shannon and spike entropy, making watermark insertion harder in structured domains.
• Entropy variant trade-off: Using spike entropy instead of Shannon entropy "substantially improves watermark detection performance," while Shannon entropy delivers superior task-specific utility (reported in Appendix D).

Methodology in Plain English

At each decoding step the model produces a probability for every vocabulary token. WaterMod sorts those tokens from highest to lowest probability. In the zero-bit case, it labels ranks alternately even and odd. A gate based on the entropy of the distribution decides which parity becomes the "green" class: when the distribution is sharp (low entropy), the even group is favored so the top token stays usable; when it is flat (high entropy), the odd group becomes more likely because more candidates are interchangeable. A small constant bias is then added to the logits of the chosen class, and the next token is chosen deterministically (argmax) in the reported experiments.

Detection reverses the process: using the same secret key and the previous token, the detector reconstructs which parity should have been green at each position, counts how many generated tokens landed in that class, and converts the count into a z-score against a null hypothesis in which each class is equally likely.

For multi-bit payloads, the same idea generalizes to k classes. A b-bit message is written as base-k digits, a keyed pseudorandom function picks which digit position to embed at each step, and only the logits of tokens whose rank satisfies rank mod k = d receive the bias. The detector reconstructs the target class at each step, tallies observed colors per digit position, recovers each digit by majority vote, and computes a z-score with null probability 1/k.

Experimental setup: all methods use the same model, Qwen-2.5-1.5B, with deterministic decoding, a maximum of 400 tokens for C4 and 600 for GSM8K and MBPP+, on a single NVIDIA RTX 3090 GPU with 24GB of memory. Zero-bit runs use H_scale = 1.2, bias δ = 1.0, and a green list ratio of 0.5; multi-bit runs use δ = 2.5, base k = 4, and 16-bit payloads, chosen because that supports over 65,000 unique identifiers. C4 uses 500 sampled instances, GSM8K uses the 1,319-instance test split of a dataset with 8,000 problems, and MBPP+ uses 378 problems each with around 100 test cases.

Why This Matters

The paper frames watermarking as a practical route to complying with regulatory disclosure requirements, noting that the EU AI Act requires general-purpose AI outputs to be identified as such, with disclosure obligations expected to take effect by 2026, and that the mark should resist common post-processing and be algorithmically verifiable. It situates the work alongside industry efforts such as OpenAI's acknowledged internal text-watermark detector under evaluation for ChatGPT and Google DeepMind's SynthID-Text.

Real-world applications:

• Provenance and attribution: flagging text as machine-generated for regulatory disclosure.
• Leak tracing: embedding 16-bit payloads to encode over 65,000 unique identifiers, which the paper links to tracking leaked fine-tuning data or identifying a specific model instance behind disinformation.
• Content platforms and publishing: detecting synthetic news-style continuations, the C4 scenario the paper simulates.
• High-stakes technical domains: maintaining detection in mathematical and code-generation outputs, where low entropy makes embedding harder.

Industry relevance: the paper argues that because WaterMod relies only on probability scores the model already outputs, it avoids external synonym resources, embedding hashing, and prompt engineering, and that the approach is model-agnostic and generalizable to other open-source LLMs, with code and data released at a public GitHub repository.

Future Directions

• Automated entropy-scaling selection. The paper sets H_scale = 1.2 in zero-bit experiments and explicitly leaves automated discovery of the optimal entropy scaling factor for future work, noting it could be tuned per domain entropy level.
• Choosing between Shannon and spike entropy. Appendix D shows spike entropy improves detection while Shannon entropy improves task utility, leaving the question of how to select or combine them.
• Scaling and generalizing the multi-bit regime. The evaluation uses base k=4 with 16-bit payloads; behavior at other bases, payload sizes, and capacities is not reported.
• Broader attack and model coverage. The robustness test uses a single ChatGPT paraphrase attack on GSM8K; other post-processing operations, adversarial attacks, and other open-source LLMs are not evaluated in the reported content.

Target Audience

Researchers and practitioners working on LLM watermarking, AI content provenance, and regulatory compliance for generative models; NLP and ML engineers who need to embed and detect provenance signals without degrading output quality; and policy- or industry-facing readers tracking how watermarking methods might satisfy disclosure mandates such as the EU AI Act.

Authors’ abstract

Large language models now draft news, legal analyses, and software code with human-level fluency. At the same time, regulations such as the EU AI Act mandate that each synthetic passage carry an imperceptible, machine-verifiable mark for provenance. Conventional logit-based watermarks satisfy this requirement by selecting a pseudorandom green vocabulary at every decoding step and boosting its logits, yet the random split can exclude the highest-probability token and thus erode fluency. WaterMod mitigates this limitation through a probability-aware modular rule. The vocabulary is first sorted in descending model probability; the resulting ranks are then partitioned by the residue rank mod k, which distributes adjacent-and therefore semantically similar-tokens across different classes. A fixed bias of small magnitude is applied to one selected class. In the zero-bit setting (k=2), an entropy-adaptive gate selects either the even or the odd parity as the green list. Because the top two ranks fall into different parities, this choice embeds a detectable signal while guaranteeing that at least one high-probability token remains available for sampling. In the multi-bit regime (k>2), the current payload digit d selects the color class whose ranks satisfy rank mod k = d. Biasing the logits of that class embeds exactly one base-k digit per decoding step, thereby enabling fine-grained provenance tracing. The same modular arithmetic therefore supports both binary attribution and rich payloads. Experimental results demonstrate that WaterMod consistently attains strong watermark detection performance while maintaining generation quality in both zero-bit and multi-bit settings. This robustness holds across a range of tasks, including natural language generation, mathematical reasoning, and code synthesis. Our code and data are available at https://github.com/Shinwoo-Park/WaterMod.