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

TL;DR

WaterMod addresses the need for provenance marks on LLM outputs without compromising fluency. It introduces probability‑balanced, modular token‑ranking via rank modulo $k$, enabling zero‑bit and multi‑bit watermarking that preserves high‑probability tokens while embedding detectable signals. The method achieves strong watermark detection (AUROC) across NLP, math reasoning, and code generation, and demonstrates robustness to paraphrase attacks, with a unified framework that scales from binary attribution to rich payloads. This approach provides a practical, regulator‑friendly pathway for accountable AI with minimal impact on generation quality.

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.

Figures (4)

• Figure 1: Watermark embedding procedure of WaterMod in the zero-bit setting.
• Figure 2: Overview of the message encoding and recovery process in WaterMod under the multi-bit watermarking regime.
• Figure 3: We report the mean token-level entropy computed over LLM-generated outputs for each dataset. Entropy is measured at every decoding step and averaged across all tokens. Both Shannon entropy (left) and spike entropy (right) are presented. Tasks such as mathematical reasoning and code generation show consistently lower entropy than natural language generation, indicating more deterministic token distributions in structured domains.
• Figure 4: Performance comparison of WaterMod using Shannon entropy and spike entropy under the zero-bit watermarking setting. The bar charts on the left present watermark detection performance, while those on the right show task-specific performance. Blue bars represent the performance of WaterMod configured with spike entropy, and green bars correspond to WaterMod using Shannon entropy.