💬 NLP

SSG: Logit-Balanced Vocabulary Partitioning for LLM Watermarking

The paper proposes SSG (Sort-then-Split by Groups), a new vocabulary partitioning method that improves LLM watermarking detectability in low-entropy scenarios, such as code and math generation, by creating logit-balanced subsets to raise the lower bound of watermark strength.

Original authors: Chenxi Gu, Xiaoning Du, John Grundy

Published 2026-04-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Chenxi Gu, Xiaoning Du, John Grundy

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a secret agent trying to leave a "hidden signature" on a letter to prove you wrote it. You want this signature to be subtle enough that a casual reader won't notice, but obvious enough that your agency can spot it later using a special magnifying glass.

This paper, SSG, is about making that hidden signature much easier to find, especially in situations where it’s normally very hard to hide anything.

The Problem: The "One-Choice" Dilemma

Most AI watermarking (like the famous "KGW" method) works by playing a game of "Red Light, Green Light" with words. The AI is told: "When you pick the next word, try to pick a 'Green' word more often than a 'Red' word."

In normal conversation (high entropy), there are hundreds of words to choose from. If the AI picks a "Green" word instead of a "Red" one, it’s easy to notice the pattern.

But imagine you are writing a math equation or a line of computer code (low entropy). In these cases, there is often only one correct answer. If the next word must be x = 5, the AI doesn't have a choice. It can't pick a "Green" word if the only correct word is "Red." The watermark gets "drowned out" because the AI is forced to be accurate, not "green." It’s like trying to hide a secret message in a book where every single word is already printed in bold, permanent ink.

The Solution: SSG (The "Fair Split" Strategy)

The researchers noticed that the problem isn't the "Green/Red" rule itself; it's how we decide which words are Green and which are Red.

In the old way, words were assigned colors randomly. This was like throwing a handful of red and green marbles into a dark room. By pure bad luck, all the "important" marbles (the ones the AI actually wants to pick) might end up being Red. If that happens, the watermark disappears.

SSG (Sort-then-Split by Groups) changes the game. Instead of random assignment, it follows these steps:

The Lineup: It looks at all the possible next words and lines them up from "most likely" to "least likely" (like a leaderboard).
The Buddy System: It takes the top candidates and puts them into pairs.
The Fair Coin Toss: For every pair, it flips a coin: one word becomes Green, and the other becomes Red.

The Analogy: The Buffet Line
Imagine a buffet where there are 100 dishes, but only 2 are actually delicious (the high-probability tokens).

The Old Way (KGW): You assign colors randomly. You might accidentally make both delicious dishes "Red." Now, the "Green" signal is useless because no one is eating the Green food.
The SSG Way: You take the two delicious dishes, pair them up, and say, "One of you is Green, and one of you is Red." Now, no matter what the AI chooses, it is guaranteed to pick a Green word half the time. You have "balanced" the flavors so the signature always shows up.

Why does this matter?

It’s Smarter, Not Louder: Usually, to make a watermark stronger, you have to force the AI to pick "Green" words so aggressively that the AI starts making mistakes (like writing bad code). SSG makes the watermark stronger without making the AI "stupid."
It Works Everywhere: It fixes the problem in "boring" or "strict" tasks like math and coding, while still working perfectly fine for creative writing.
Better Proof: It provides a mathematical guarantee that the "strength" of the watermark won't drop to zero, even when the AI is being very precise.

In short: SSG ensures that even when the AI has its hands tied by logic and math, it can still leave a clear, detectable fingerprint.

Technical Summary: SSG: Logit-Balanced Vocabulary Partitioning for LLM Watermarking

1. Problem Statement

The paper addresses a critical vulnerability in existing Large Language Model (LLM) watermarking techniques, specifically the widely used KGW (Kirchenbauer et al.) scheme. While KGW is effective for high-entropy natural language generation, its performance degrades significantly in low-entropy settings, such as code generation and mathematical reasoning.

In low-entropy tasks, the next-token probability distribution is "sharply peaked," meaning a few tokens hold almost all the probability mass. The KGW scheme relies on random vocabulary partitioning (splitting tokens into "green" and "red" sets based on a hash). In low-entropy scenarios, there is a high statistical probability that all high-logit (high-probability) tokens will fall into the same subset. If the top tokens are all in the red set, the watermark bias cannot be effectively applied without severely damaging the model's output quality, leading to a near-zero "watermark strength" and making the watermark undetectable.

2. Methodology

The authors introduce a new metric and a novel partitioning algorithm to solve this issue.

New Metric: Watermark Strength ( $f_{ws}$ ): The authors formalize "watermark strength" as a token-level metric that quantifies the normalized probability shift toward the green set during prediction. They prove that under random partitioning, the lower bound of this strength is dictated by the next-token distribution, which can drop to near zero in low-entropy tasks.
Proposed Algorithm: SSG (Sort-then-Split by Groups): Instead of random partitioning, SSG ensures the vocabulary is logit-balanced. The process is as follows:
1. Sort: Tokens are sorted in descending order of their logits.
2. Group: Tokens are organized into consecutive groups of size two.
3. Split: For each group, a pseudorandom generator (seeded by a secret key and context) decides which token goes to the green set ( $V_g$ ) and which to the red set ( $V_r$ ). This ensures that for every pair of high-probability tokens, one is green and one is red.
4. Efficiency Optimization (Top- $k$ ): To avoid the computational overhead of sorting the entire vocabulary at every step, the authors propose applying SSG only to the top- $k$ highest-probability tokens and using random partitioning for the remaining vocabulary.
Detection: SSG is compatible with existing detection methods, such as the z-score test or entropy-weighted scoring.

3. Key Contributions

Theoretical Analysis: A formalization of "Watermark Strength" and a proof showing that SSG raises the lower bound of this strength, guaranteeing detectability even when the distribution is highly concentrated.
Algorithmic Innovation: The SSG algorithm, which moves from random partitioning to a logit-balanced approach.
Efficiency Strategy: The introduction of the Top- $k$ SSG variant, which maintains effectiveness while minimizing the latency added by sorting logits.

4. Experimental Results

The authors evaluated SSG across multiple models (Qwen2.5-Coder, DeepSeekMath, LLaMA-3) and tasks (HumanEval, MBPP, GSM8K).

Improved Detectability: SSG consistently improved the True Positive Rate (TPR) and F1-score for KGW, SWEET, and EWD baselines. In many cases, detection performance saw massive gains (e.g., in mathematical reasoning tasks, TPR at 1% FPR increased significantly).
Preserved Text Quality: Unlike increasing the bias ( $\delta$ ) to improve detection, which hurts accuracy, SSG improved detectability without degrading the Pass@1 accuracy (the ability of the model to generate correct code or math answers).
Top- $k$ Trade-off: Experiments showed that even a small $k$ (e.g., $k=2$ or $k=4$ ) provides substantial benefits, while larger $k$ values offer diminishing returns and increase decoding latency.
Robustness: While all watermarks are susceptible to paraphrasing attacks (using GPT-4o-mini), SSG showed strong robustness in specific models like LLaMA-3.

5. Significance

This paper is significant because it shifts the focus of LLM watermarking research from detection algorithms (how we find the watermark) to injection algorithms (how we embed it). By redesigning how the vocabulary is partitioned, the authors have provided a way to make watermarking robust for the most "structured" and "logical" AI outputs—code and math—which are currently the hardest to protect. This is vital for protecting intellectual property and ensuring the provenance of highly technical AI-generated content.