Hey everyone, Alex here. Welcome back to another post on Coding with Alex at sysseder.com.
If you've been following the AI space lately—and let's be honest, who hasn't?—you've probably noticed a massive, looming problem: attribution. As LLMs generate more of the web's content, we are rapidly approaching an existential crisis for search engines, academic integrity, copyright enforcement, and even the training of future models (preventing "model collapse" from training AI on AI-generated data).
How do we distinguish human-written text from machine-generated prose? Traditional pattern-matching and ML-based classifiers are notoriously unreliable, often yielding high false-positive rates on non-native English writers. This is where cryptographic and statistical watermarking comes in. Rather than trying to detect AI writing after the fact, state-of-the-art systems are embedding invisible, mathematically provable signatures directly into the text as it is being generated.
Today, we're going to dive deep under the hood of how AI text watermarking actually works, look at the mathematics behind token distribution manipulation, and write a Python implementation from scratch so you can see the engineering behind the magic.
The Core Concept: Biasing the Next-Token Distribution
To understand watermarking, we first need to take a quick step back to how Large Language Models generate text. At every step of generation, an LLM outputs a list of vocabulary tokens alongside their raw scores, known as logits. These logits are passed through a softmax function to produce a probability distribution over the entire vocabulary.
Normally, the model selects the next token by sampling from this distribution (using parameters like temperature, top-k, or top-p to control randomness).
Watermarking doesn't change the model's weights. Instead, it intercepts this sampling step. The most widely adopted approach—pioneered by researchers like John Kirchenbauer and his team at the University of Maryland—is based on a pseudo-random red-green split of the vocabulary.
The Red-Green Rule
For every token $t$ that the model is about to generate, the algorithm does the following:
- It hashes the previously generated token $t_{-1}$ (or a window of previous tokens) to seed a pseudo-random number generator (PRNG).
- Using this seed, it splits the entire vocabulary into two partitions: a "green list" and a "red list" (typically a 50/50 split).
- It adds a small bias value ($\delta$) to the logits of all tokens in the green list.
- The softmax is then calculated using these modified logits, and the next token is sampled.
Because the green list tokens have their probabilities artificially inflated, the model is highly likely to choose a green token. However, because the green/red split is determined pseudo-randomly based on the previous token, a human writer (or an un-watermarked model) writing the same text would select red and green tokens roughly equally (50% green, 50% red).
By counting how many tokens fall into their respective green lists over a passage of text, we can mathematically calculate the probability that the text was generated by our watermarked model.
Architecture of a Watermarked Generation System
In a production system, watermarking is implemented as a wrapper or a custom logits processor during the decoding loop. Here is a high-level conceptual look at the pipeline:
[Prompt] ---> [ LLM Forward Pass ]
|
[ Raw Logits ]
|
v
[ Pseudo-Random Generator ] <--- [ Previous Token(s) Hash Key ]
|
v
[ Create Green/Red Split ]
|
v
[ Add Bias (δ) to Green List ]
|
v
[ Softmax & Sample Token ] ---> [ Output Token ] ---> (Repeat Loop)
This architecture is incredibly elegant because it requires zero modifications to the underlying neural network. It's computationally inexpensive, adding only a tiny vector addition step to the generation loop.
Implementing a Text Watermarker in Python
Let's build a simplified, working implementation of this algorithm using PyTorch and Hugging Face's transformers library. We'll write a custom logits processor that applies this bias, generates watermarked text, and then writes a detection function to calculate the z-score of the generated text.
Step 1: The Watermark Logits Processor
First, let's write the processor that hooks into the generation phase. This processor dynamically splits the vocabulary based on a hash of the previous token.
import torch
from transformers import LogitsProcessor
class WatermarkLogitsProcessor(LogitsProcessor):
def __init__(self, vocab_size, green_fraction=0.5, bias=2.0, hash_key=42):
self.vocab_size = vocab_size
self.green_fraction = green_fraction
self.bias = bias
self.hash_key = hash_key
self.green_size = int(vocab_size * green_fraction)
def _get_green_list(self, prev_token_id):
# Seed generator deterministically using the previous token ID and our secret hash key
generator = torch.Generator()
generator.manual_seed(int(prev_token_id) ^ self.hash_key)
# Shuffle the vocabulary indices pseudo-randomly
perm = torch.randperm(self.vocab_size, generator=generator)
green_list = perm[:self.green_size]
return green_list
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
# For batch generation, process each sequence in the batch
for i in range(input_ids.shape[0]):
if input_ids.shape[1] == 0:
continue # Skip if no prefix exists yet
prev_token_id = input_ids[i, -1].item()
green_list = self._get_green_list(prev_token_id)
# Apply bias to green list tokens
scores[i, green_list] += self.bias
return scores
Step 2: Detecting the Watermark (The Z-Score Test)
To detect the watermark, we analyze the text, reconstruct the green/red split for every single step of the generation, and count how many times the actual token chosen was in the green list.
Under the null hypothesis (that the text is human-written or un-watermarked), the number of green tokens $S$ in a text of length $N$ follows a binomial distribution with success probability $p = 0.5$ (assuming our green fraction is 50%). For a reasonably large $N$, this can be approximated by a normal distribution. We calculate the z-score as:
$$z = \frac{S - N \cdot p}{\sqrt{N \cdot p \cdot (1 - p)}}$$If the z-score is high (e.g., $> 4.0$), the probability of this happening by chance is virtually zero, confirming the text is watermarked.
import math
class WatermarkDetector:
def __init__(self, vocab_size, green_fraction=0.5, hash_key=42):
self.vocab_size = vocab_size
self.green_fraction = green_fraction
self.hash_key = hash_key
self.green_size = int(vocab_size * green_fraction)
def _get_green_list(self, prev_token_id):
generator = torch.Generator()
generator.manual_seed(int(prev_token_id) ^ self.hash_key)
perm = torch.randperm(self.vocab_size, generator=generator)
return set(perm[:self.green_size].tolist())
def detect(self, token_ids):
if len(token_ids) < 2:
return 0.0, 0 # Not enough tokens to detect
green_count = 0
total_tokens = len(token_ids) - 1 # We can only test tokens that have a precursor
for i in range(1, len(token_ids)):
prev_token = token_ids[i-1]
current_token = token_ids[i]
green_list = self._get_green_list(prev_token)
if current_token in green_list:
green_count += 1
# Calculate Z-score
p = self.green_fraction
expected_mean = total_tokens * p
expected_var = total_tokens * p * (1 - p)
z_score = (green_count - expected_mean) / math.sqrt(expected_var)
return z_score, green_count
Testing the Watermark in Practice
If you plug this implementation into a text generation loop using GPT-2 or Llama, you will quickly see its power.
A completely standard model output will yield a z-score of around $-1.5$ to $1.5$ (purely random distribution around the mean). However, the watermarked output—even with a modest bias ($\delta = 1.5$) that doesn't visibly degrade the text quality—will easily output z-scores of $6.0$ or higher on short paragraphs of 100 words. A z-score of $6.0$ corresponds to a false positive rate of roughly 1 in a billion!
The Developer's Dilemma: Security vs. Usability
As cool as this math is, watermarking isn't a silver bullet. As developers designing and implementing these systems, we have to navigate several trade-offs:
1. Text Quality vs. Detectability
The parameter $\delta$ (bias) is a double-edged sword. If you set $\delta$ too high, the model is heavily restricted in its vocabulary choices. This results in repetitive phrasing, awkward word choices, or outright hallucinations. If you set $\delta$ too low, the watermark becomes incredibly difficult to detect in shorter snippets of text.
2. The "Paraphrasing" Vulnerability
The watermark relies entirely on the local context (e.g., the previous token $t_{-1}$). If an end-user takes the watermarked text and runs it through a different, un-watermarked model (like a local Mistral instance) with a prompt like "Paraphrase this text," the token transitions change entirely. This destroys the green-list alignment and completely breaks the watermark detection.
3. Cryptographic Secret Management
In our Python example, we used a simple XOR with a hash_key. In a real-world enterprise deployment, if the key used to generate the pseudo-random splits is compromised, bad actors can easily reverse-engineer the green lists and strip out the watermarks by intentionally swapping green tokens for red synonyms. Protecting this hashing key is just as critical as protecting your private PKI keys.
Wrapping Up
AI text watermarking is a fascinating bridge between natural language processing, statistics, and cryptography. For developers building LLM wrappers, educational software, or enterprise compliance tools, understanding how these signatures are injected—and how easily they can be bypassed or preserved—is becoming an essential skill set.
What are your thoughts on watermarking? Do you think it's a viable long-term solution for tracking AI-generated content, or will the cat-and-mouse game of paraphrasing tools always stay one step ahead? Let me know in the comments below!
Until next time, keep coding.
— Alex