SSG: Logit-Balanced Vocabulary Partitioning for LLM Watermarking

Abstract: Watermarking has emerged as a promising technique for tracing the authorship of content generated by LLMs. Among existing approaches, the KGW scheme is particularly attractive due to its versatility, efficiency, and effectiveness in natural language generation. However, KGW's effectiveness degrades significantly under low-entropy settings such as code generation and mathematical reasoning. A crucial step in the KGW method is random vocabulary partitioning, which enables adjustments to token selection based on specific preferences. Our study revealed that the next-token probability distribution plays an critical role in determining how much, or even whether, we can modify token selection and, consequently, the effectiveness of watermarking. We refer to this characteristic, associated with the probability distribution of each token prediction, as \emph{watermark strength.} In cases of random vocabulary partitioning, the lower bound of watermark strength is dictated by the next-token probability distribution. However, we found that, by redesigning the vocabulary partitioning algorithm, we can potentially raise this lower bound. In this paper, we propose SSG (\textbf{S}ort-then-\textbf{S}plit by \textbf{G}roups), a method that partitions the vocabulary into two logit-balanced subsets. This design lifts the lower bound of watermark strength for each token prediction, thereby improving watermark detectability. Experiments on code generation and mathematical reasoning datasets demonstrate the effectiveness of SSG.

• The paper introduces a logit-balanced vocabulary partitioning technique (SSG) that systematically distributes high-logit tokens to boost watermark detection in LLM outputs.
• It employs a sorted, pairwise splitting method on top-k tokens and applies a fixed bias to the green set, ensuring robust statistical watermarking even in low-entropy generation.
• Empirical results indicate significant improvements in detection metrics for code synthesis and text tasks without compromising generation quality.

SSG: Logit-Balanced Vocabulary Partitioning for Robust LLM Watermarking

Introduction and Motivation

Text watermarking for LLM outputs is increasingly critical for content provenance and model authorship attribution, particularly given regulatory and proprietary concerns. The canonical KGW watermarking scheme—partitioning vocabulary pseudorandomly and boosting logits of a "green" set—performs well in high-entropy contexts but suffers severe detectability degradation in low-entropy tasks such as code synthesis and mathematical reasoning. The principal limitation originates from the random partitioning step: when high-probability tokens are clustered within a single set, the capacity to inject and later detect the watermark via statistical methods is effectively nullified, especially when the next-token distribution is sharply peaked.

To address this, the proposed SSG (Sort-then-Split by Groups) algorithm introduces a logit-balanced partitioning method. The green and red vocabulary sets are crafted so that high-logit tokens are distributed evenly, maximizing the watermarking "strength"—a token-level shift in probability mass measurable across all entropy regimes. This approach provably increases the lower bound of watermark strength and thus enhances statistical detectability compared to KGW.

SSG Methodology

Algorithmic Design

The SSG injection scheme fundamentally modifies the partitioning phase of watermarking:

1. Sorting and Grouping: All vocabulary tokens are sorted in descending order of predicted logits at each decoding step.
2. Pairwise Splitting: For the top-k tokens (where $k$ is a small integer to balance cost and efficacy), consecutive pairs are deterministically split between the green and red sets based on groupwise pseudorandom assignment seeded by a fixed secret.
3. Logit Bias Application: A fixed bias $\delta$ is added to the logits of the assigned green tokens. Remaining vocabulary tokens (outside top-k) are partitioned randomly as in KGW, maintaining computational efficiency.
4. Sampling and Decoding: Next-token generation continues as usual with the modified probabilities.

This careful groupwise partition ensures both the highest-logit tokens are never grouped within a single set and that the minimum attainable watermark strength (as measured by normalized probability shift) cannot collapse to zero except in pathological, fully-deterministic cases.

Figure 1: Influence of top-k on SSG performance.

Computational Trade-offs

Applying SSG to only the top-k tokens achieves most of the detection gains at a fraction of the computational cost of a full-vocabulary sort. Empirical results show marginal improvements in detection for $k > 2$, with diminishing returns as $k$ increases. Thus, the recommended configuration for practical deployment is $k=2$ or $k=4$, balancing runtime and watermark strength.

Theoretical Analysis: Watermark Strength

Watermark strength $f_{ws}(\delta, p_g)$ is defined as the normalized difference in the sum probability of the green set before and after biasing. The function is provably bounded away from zero under SSG's partitioning, except when the token distribution is extremely concentrated on a single token (i.e., $p_1 \to 1$). In this regime, no logit-based watermarking is robust—watermark signal is inherently indistinguishable from the model's prediction.

Figure 2: Curve of Watermark Strength over $p_g$ and $\delta$.

In all settings where more than one token shares significant probability mass, SSG enforces a tight lower bound on watermark strength, thus guaranteeing that the injected statistical footprint remains detectable under classical statistical tests (e.g., z-score).

Experimental Evaluation

Low-Entropy Tasks

Experiments on code generation (HumanEval, MBPP) and mathematical reasoning (GSM8K) with models such as Qwen2.5-Coder-7B and DeepSeekMath-7B demonstrate that SSG-augmented variants of KGW, SWEET, and EWD watermarking algorithms substantially increase detection metrics (TPR, F1, z-score) versus their vanilla counterparts without observable reduction in generation quality (i.e., Pass@1).

Token-Level Watermark Strength Distributions

Empirical analysis reveals that SSG greatly reduces the proportion of tokens with near-zero watermark strength across generated sequences, even in severely low-entropy tasks, thereby ensuring more uniform and reliable detection footprints sequence-wide.

Figure 3: Examples of the watermark strength distributions on Qwen2.5-Coder-7B and LLaMA3-8B for code generation tasks.

High-Entropy Compatibility and Robustness

SSG shows consistent benefits on high-entropy text tasks (e.g., C4, CNN/DailyMail) as well, underlining the method's generality. Regarding paraphrase attacks, SSG confers substantial robustness improvements on some model-task pairings but is not universally robust, indicating residual challenges in adversarial settings.

Practical and Theoretical Implications

The introduction of logit-balanced vocabulary partitioning in SSG represents a significant shift: watermark detectability is now primarily limited by the intrinsic entropy of the model's predictive distribution, not by the artifacts of partitioning strategy. In real-world deployments—where regulatory and IP enforcement requires robust attribution even for LLM outputs in structured, low-diversity formats—SSG provides meaningful improvements without notable degradation in output utility or fluency.

However, reliance on prompt-conditional logit information for detection limits the approach when original prompts are unavailable or contaminated. Thus, future directions include prompt-independent watermarking, more sophisticated logit bias allocation exploiting model internals, and greater robustness to attack scenarios such as paraphrasing or token-level mutation.

Conclusion

The SSG watermarking algorithm provides a robust, theoretically justified improvement over prior logit-based watermarking schemes, particularly in challenging low-entropy LLM generation settings. By enforcing logit balance in vocabulary partitioning, SSG raises the minimum achievable watermark strength, leading to significantly better detection across a wide variety of tasks and models while preserving text quality and decoding efficiency. SSG thus sets a new standard for reliable, computationally efficient watermarking in LLM applications where authorship verification and output tracing are operationally essential.

Figure 2: Curve of Watermark Strength over $\delta$0 and $\delta$1.