Watermarking › Implementation
Language AI HandbookAll chapters Watermarking › ImplementationOne section · Sections in this chapter7 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement the KGW watermarking scheme from scratch to see exactly how the generation and detection pipelines work. We will use a simplified vocabulary and a toy language model to illustrate the mechanics, then apply it to a realistic scenario.
Setup and Vocabulary
We begin by setting up the vocabulary and hash function for partitioning:
import numpy as np
# Seed for reproducibility
rng = np.random.default_rng(42)
# Simplified vocabulary (actual LLMs have 50k+ tokens)
vocab = [
"the",
"cat",
"sat",
"on",
"mat",
"a",
"dog",
"ran",
"fast",
"slow",
"jumped",
"over",
"fence",
"small",
"large",
"red",
"blue",
"green",
"happy",
"sad",
"sun",
"moon",
"star",
"tree",
"house",
"water",
"fire",
"earth",
"wind",
"rain",
"snow",
"cold",
"warm",
"hot",
"good",
"bad",
"big",
"little",
"old",
"new",
"black",
"white",
"left",
"right",
"up",
"down",
"in",
"out",
"yes",
"no",
]
vocab_size = len(vocab)
word_to_id = {w: i for i, w in enumerate(vocab)}
id_to_word = {i: w for w, i in word_to_id.items()}
# Secret watermark key
SECRET_KEY = "language-ai-handbook-watermark-key-2024"
GAMMA = 0.5 # fraction of vocab in green list
DELTA = 2.0 # logit boost for green tokensVocabulary size: 50 Green list fraction (gamma): 0.5 Logit boost (delta): 2.0 Expected green list size: 25 tokens
With a vocabulary of 50 words and a green fraction of 0.5, roughly 25 tokens will be on the green list at any given step. The partition changes at each step based on the previous token, so the same word can be green at one point in a sequence and red at another.
Green List Generation
The key function maps a context token to a green list using a cryptographic hash:
def get_green_list(
prev_token_id: int, secret_key: str, vocab_size: int, gamma: float
) -> set:
"""
Derive the green list for the current position using a hash of
the previous token and the secret key. Returns the set of token IDs
that are on the green list.
"""
# Create a deterministic seed from the context and key
seed_str = f"{secret_key}:{prev_token_id}"
hash_bytes = hashlib.sha256(seed_str.encode()).digest()
seed = int.from_bytes(hash_bytes[:4], "big")
# Use the seed to deterministically shuffle the vocabulary
local_rng = np.random.default_rng(seed)
shuffled = local_rng.permutation(vocab_size)
# The first gamma fraction of the shuffled vocabulary is the green list
green_size = int(gamma * vocab_size)
return set(shuffled[:green_size].tolist())Green list after 'cat' (first 10): ['a', 'bad', 'cat', 'cold', 'dog', 'earth', 'green', 'house', 'large', 'left'] Green list after 'dog' (first 10): ['a', 'black', 'cat', 'down', 'earth', 'fast', 'fire', 'good', 'green', 'little'] Overlap between the two lists: 15 tokens (60%) Token 'mat' in green list after 'cat': False Token 'mat' in green list after 'dog': False
The overlap between green lists for different context tokens is approximately 50%, consistent with random partitioning. The same token can be green after one context word and red after another, which is why the watermark cannot be removed by simply avoiding specific words. An attacker who observes that "mat" is green after "cat" has learned nothing about whether "mat" will be green after "dog" or "tree."
Watermarked Text Generation
The generation function modifies the logits before sampling. The modification is minimal: we add to the logit of each green-list token and leave red-list logits unchanged. Sampling then proceeds normally from the modified distribution.
def generate_watermarked(
prompt_ids: list[int],
base_logits_fn,
vocab_size: int,
secret_key: str,
gamma: float,
delta: float,
n_tokens: int,
seed: int = 0,
) -> list[int]:
"""
Generate tokens using the KGW watermarking scheme.
At each step, boost the logits of green-list tokens by delta.
"""
gen_rng = np.random.default_rng(seed)
tokens = list(prompt_ids)
for _ in range(n_tokens):
prev_token_id = tokens[-1]
# Get base logits from the language model
logits = base_logits_fn(tokens, gen_rng)
# Get green list for the current context
green_list = get_green_list(
prev_token_id, secret_key, vocab_size, gamma
)
# Boost green-list logits by delta (soft watermarking)
for token_id in green_list:
logits[token_id] += delta
# Convert logits to probabilities via softmax
logits_shifted = logits - logits.max()
probs = np.exp(logits_shifted) / np.exp(logits_shifted).sum()
# Sample next token
next_token = gen_rng.choice(vocab_size, p=probs)
tokens.append(int(next_token))
return tokens[len(prompt_ids) :] # Return only generated tokens
def generate_unwatermarked(
prompt_ids: list[int],
base_logits_fn,
vocab_size: int,
n_tokens: int,
seed: int = 0,
) -> list[int]:
"""Generate tokens without watermarking, for comparison."""
gen_rng = np.random.default_rng(seed)
tokens = list(prompt_ids)
for _ in range(n_tokens):
logits = base_logits_fn(tokens, gen_rng)
logits_shifted = logits - logits.max()
probs = np.exp(logits_shifted) / np.exp(logits_shifted).sum()
next_token = gen_rng.choice(vocab_size, p=probs)
tokens.append(int(next_token))
return tokens[len(prompt_ids) :]
def dummy_lm_logits(
token_ids: list[int], rng: np.random.Generator
) -> np.ndarray:
"""
A simple uniform language model. In practice, this would be the actual
transformer's output logits. Here we use uniform base logits plus a
small random perturbation to simulate natural variation.
"""
base = np.zeros(vocab_size)
# Small random perturbation to avoid perfectly uniform outputs
noise = rng.normal(0, 0.5, vocab_size)
return base + noiseWatermarked : new up good in right little mat white star warm in little over old over left hap... Unwatermarked : right up bad in black old mat left fire rain down old green bad over new sun bad... Tokens generated: 100
Both sequences look like random strings of words because our dummy language model is nearly uniform. In practice, a real language model would produce coherent sentences whose token choices are heavily constrained by grammar and semantics. The key point is that the watermarked sequence's token choices are systematically biased toward the green list at each position, an invisible pattern that accumulates into a detectable signal.
The Detection Procedure
The detector needs only the candidate text tokens, the secret key, and the parameters and threshold. It recomputes the green list at each position (using the same hash function and key used during generation), counts how many tokens fall on the green list, and runs the z-test.
def detect_watermark(
token_ids: list[int],
secret_key: str,
vocab_size: int,
gamma: float,
z_threshold: float = 2.33, # Corresponds to FPR ~0.01
) -> dict:
"""
Run the KGW watermark detector on a sequence of token IDs.
Returns the z-score, p-value, and detection decision.
"""
from scipy import stats
T = len(token_ids)
if T < 2:
return {
"z_score": 0.0,
"p_value": 1.0,
"is_watermarked": False,
"n_tokens": T,
}
# Count green tokens
green_count = 0
for pos in range(1, T):
prev_token_id = token_ids[pos - 1]
token_id = token_ids[pos]
green_list = get_green_list(
prev_token_id, secret_key, vocab_size, gamma
)
if token_id in green_list:
green_count += 1
n_tested = T - 1 # First token has no previous context
expected_green = gamma * n_tested
std_green = np.sqrt(n_tested * gamma * (1 - gamma))
z_score = (green_count - expected_green) / std_green
p_value = 1 - stats.norm.cdf(z_score)
is_watermarked = z_score > z_threshold
return {
"z_score": z_score,
"green_count": green_count,
"expected_green": expected_green,
"p_value": p_value,
"is_watermarked": is_watermarked,
"n_tokens": n_tested,
}Detection results (z-threshold = 2.33, FPR ~0.01): Watermarked text: Green tokens : 85 / 101 (expected 50.5) Z-score : 6.866 p-value : 0.000000 Decision : WATERMARKED Unwatermarked text: Green tokens : 50 / 101 (expected 50.5) Z-score : -0.100 p-value : 0.539631 Decision : NOT WATERMARKED
The watermarked text shows a z-score substantially above the threshold, while the unwatermarked text's z-score is close to zero, consistent with the null hypothesis. The p-value for the watermarked text is very small, indicating strong statistical evidence for the presence of the watermark. Notice that the detection function has no access to the generation function or the original model: it operates only on the token sequence and the key. This is what makes the scheme practical for real-world deployment, where the detector may run on a different server from the generator.
Visualizing Detection Power vs. Text Length
The detection power increases with the number of tokens. Let's visualize how the z-score and detection decision evolve as more text is processed:
def simulate_detection_power(
n_tokens_range: list[int],
n_trials: int,
vocab_size: int,
secret_key: str,
gamma: float,
delta: float,
z_threshold: float,
) -> dict:
"""
For each text length, generate n_trials watermarked and unwatermarked
texts and compute TPR and FPR.
"""
tpr_list = []
fpr_list = []
z_wm_means = []
z_uw_means = []
prompt = [0] # Single starting token (token ID 0)
for n_tokens in n_tokens_range:
wm_detections = 0
uw_detections = 0
z_wm_vals = []
z_uw_vals = []
for trial in range(n_trials):
wm_gen = generate_watermarked(
prompt,
dummy_lm_logits,
vocab_size,
secret_key,
gamma,
delta,
n_tokens,
seed=trial * 100,
)
uw_gen = generate_unwatermarked(
prompt, dummy_lm_logits, vocab_size, n_tokens, seed=trial * 100
)
wm_result = detect_watermark(
prompt + wm_gen, secret_key, vocab_size, gamma, z_threshold
)
uw_result = detect_watermark(
prompt + uw_gen, secret_key, vocab_size, gamma, z_threshold
)
wm_detections += int(wm_result["is_watermarked"])
uw_detections += int(uw_result["is_watermarked"])
z_wm_vals.append(wm_result["z_score"])
z_uw_vals.append(uw_result["z_score"])
tpr_list.append(wm_detections / n_trials)
fpr_list.append(uw_detections / n_trials)
z_wm_means.append(np.mean(z_wm_vals))
z_uw_means.append(np.mean(z_uw_vals))
return {
"n_tokens_range": n_tokens_range,
"tpr": tpr_list,
"fpr": fpr_list,
"z_wm_means": z_wm_means,
"z_uw_means": z_uw_means,
}Detection power vs. text length:
Tokens | TPR | FPR | E[z|WM] | E[z|no WM]
----------------------------------------------------------
10 | 0.610 | 0.010 | 2.33 | 0.07
20 | 0.900 | 0.005 | 3.33 | 0.06
30 | 1.000 | 0.010 | 4.15 | 0.10
50 | 1.000 | 0.005 | 5.40 | 0.10
75 | 1.000 | 0.000 | 6.63 | 0.10
100 | 1.000 | 0.010 | 7.62 | 0.06
150 | 1.000 | 0.010 | 9.34 | 0.03
200 | 1.000 | 0.020 | 10.79 | 0.03The table shows how TPR grows from near-chance at 10 tokens toward reliable detection at 100 or more tokens, while FPR stays close to the theoretical value of 0.01. This empirical confirmation that the test is well-calibrated under the null (FPR tracks 0.01 across all lengths) is important: it means the hypothesis testing framework's assumptions are satisfied in this implementation.
Visualization: Detection Power and Z-Score Distributions
The TPR curve rises steeply between 20 and 100 tokens, reaching reliable detection in the 100 to 150 token range for this toy model. The FPR line stays near the 0.01 theoretical value throughout, confirming that the hypothesis test is well-calibrated. The z-score plot tells the same story from a different angle: the watermarked mean z-score grows proportionally to , while the unwatermarked mean stays near zero as expected.
Simulating a Paraphrase Attack
To understand robustness, let's simulate an attacker who replaces a random fraction of tokens. In practice, a paraphrase model would select semantically appropriate replacements, but random substitution gives us a controlled experiment that quantifies the signal decay as a function of substitution rate.
def paraphrase_attack(
token_ids: list[int],
substitution_rate: float,
vocab_size: int,
seed: int = 0,
) -> list[int]:
"""
Simulate a paraphrase attack by randomly substituting a fraction of tokens.
In a real attack, a paraphrase model would be used to substitute tokens
with semantically similar alternatives. Here we use random substitution
as a worst-case upper bound on attack effectiveness.
"""
attack_rng = np.random.default_rng(seed)
attacked = list(token_ids)
for i in range(len(attacked)):
if attack_rng.random() < substitution_rate:
attacked[i] = int(attack_rng.integers(0, vocab_size))
return attackedParaphrase attack robustness (100 tokens, 300 trials):
Substitution Rate | TPR
----------------------------------
0.0% | 1.000
10.0% | 1.000
20.0% | 0.993
30.0% | 0.940
40.0% | 0.710
50.0% | 0.393
Even with 30% of tokens randomly substituted, detection often remains reliable. At 50% substitution, TPR falls sharply to around 0.4: still well above the false positive baseline, but no longer reliable enough for confident attribution. A real paraphrase attack using a language model would be less effective at low substitution rates because synonym substitutions often preserve green-list membership by chance. Two synonyms for the same concept tend to appear in similar contexts, and the hash function assigns them the same green list for any given preceding token. However, at high substitution rates, even semantically aware paraphrasers become effective.
Comparing Watermarked and Unwatermarked Token Distributions
The separation between the two distributions is the signal that the z-test exploits. The unwatermarked distribution is centered near , exactly as the null hypothesis predicts. The watermarked distribution is shifted to the right because the logit boost increases the probability of choosing green tokens. A higher would shift the watermarked distribution further right, increasing detection power at the cost of more noticeable distortion in text quality.
Key Parameters
The main parameters of the KGW watermarking scheme are:
- (gamma): The fraction of the vocabulary assigned to the green list, typically 0.5. Smaller values make detection harder because the expected green count under both hypotheses is lower, reducing the signal-to-noise ratio. They may also cause more quality degradation at positions where the model's best token is always on the smaller green list.
- (delta): The additive logit boost for green tokens, typically 2.0. Higher values increase detection power because each token provides more evidence, but they distort the output distribution more noticeably. The relationship between and perplexity increase is roughly log-linear for moderate values.
- : The detection threshold corresponding to a chosen FPR. Use 2.33 for FPR 0.01, 3.09 for FPR 0.001, and 5.5 for Bonferroni-corrected FPR at one million tests. Higher thresholds reduce false positives but require more tokens for reliable detection.
- Context window: The number of previous tokens used to hash the green list, default 1. Larger windows increase robustness against targeted substitution but complicate detection after text edits and increase computational cost.