Watermarking a Tiny GPT: 91 Lines, No Frameworks
This new voyage starts with a rather intriguing question: can you watermark outputs generated by an LLM?
Not by attaching metadata. Not by inserting some secret phrase into the response. But by somehow leaving a statistical signature in the text itself invisible to the reader, yet detectable by someone who knows what to look for.
It turns out you can. Scott Aaronson proposed an elegant theoretical approach to this, and Google DeepMind has since deployed a related idea for text through SynthID. But reading about a statistical watermark and actually understanding how one gets into generated text are two very different things. So I wanted to build one the hard way, or perhaps the simple way: pure Python, no machine-learning frameworks, and as little code as possible inspired by Andrej Karpathy's microgpt.
Luckily, Andrej Karpathy had already given me the perfect ship, hence no need to build one from scratch.
Enter microgpt
Earlier this year, Andrej Karpathy published microgpt, a complete GPT implementation in roughly 200 lines of pure Python with no dependencies. Dataset, tokenizer, autograd, transformer, attention, Adam, training, inference the whole voyage.
His goal was to strip GPT down until only the essential algorithm remained. I loved that philosophy, so rather than build another model, I kept microgpt intact and asked a different question:
If this is all we need to generate text with a GPT, how can we watermark its outputs? and where exactly does it go?
The answer turns out to be remarkably small.
How Does an LLM Generate Text?
Before watermarking anything, there's one thing worth being precise about. An LLM doesn't simply decide "the next character is a." Instead, the model produces scores for every possible next token, and softmax turns those scores into a probability distribution. microgpt's vocabulary is just 27 tokens the letters a-z plus a beginning/end token so at one point during generation it might produce something like:
token probability
'a' 58.83%
'e' 16.28%
'i' 16.06%
'o' 2.54%
'y' 2.19%
...Then comes a tiny but extremely important line in Karpathy's inference code:
token_id = random.choices(
range(vocab_size),
weights=[p.data for p in probs]
)[0]That line samples the next token from the probabilities, and we repeat:
predict → probabilities → sample → next token → predict → ...A production LLM is enormously larger, but fundamentally the same loop remains. And this little sampling step is exactly where a watermark can enter.
The Idea
Suppose instead of using ordinary randomness to sample the next token, we use randomness derived from a secret key:
KEY = b"microgpt-demo-key"
That's a hardcoded string, chosen for readability, not security. In practice, this should be a cryptographically secure random key, and the detector should not know it. A real deployment would generate KEY with something like secrets.token_bytes(32), not a string you can read in the source file the whole scheme only works for as long as the key stays secret.
The key doesn't tell the GPT what to write it still produces exactly the same probabilities. Instead, the key generates a deterministic but seemingly random number for each candidate token:
GPT
↓
token probabilities
↓
secret-key pseudorandomness
↓
sample
↓
next tokenDo this repeatedly and something interesting happens: the individual choices still look perfectly normal, but across enough tokens they develop a statistical relationship with the secret key. That relationship is the watermark.
Gumbel-Max Enters the Grand Line
This is where Aaronson's idea becomes elegant. For every candidate token i, we have a probability p_i assigned by the GPT and a pseudorandom number r_i generated from the secret key and the recent token context. We pick the next token with:
x_t = argmax_i log(r_i) / p_iIt looks intimidating. The code isn't this is the actual watermark.py from the repository, unchanged:
def uniform(context, token, key=KEY):
"""secret key + context + candidate token -> deterministic U(0, 1)"""
msg = ','.join(map(str, context[-CONTEXT:])) + f'|{token}'
digest = hmac.new(key, msg.encode(), hashlib.sha256).digest()
integer = int.from_bytes(digest[:8], 'big')
return (integer + 0.5) / 2**64 # never exactly 0 or 1
def sample(probabilities, context, key=KEY):
"""Keyed Gumbel-max sampling: argmax log(U_i) / p_i."""
scores = []
for token, p in enumerate(probabilities):
u = uniform(context, token, key)
scores.append(math.log(u) / p if p > 0 else -math.inf)
return max(range(len(scores)), key=scores.__getitem__)
Given the same secret key, context, and candidate token, uniform() always returns the same pseudorandom value that's an HMAC-SHA256 digest, not a coin flip. No PyTorch, no watermarking library, no model modification. The GPT hasn't changed at all. We changed how we use its probabilities.
One Token Under the Microscope
Instead of treating the algorithm as a black box, I froze one generation step and inspected it. The repository actually ships a microscope() helper in watermark.py that prints exactly this table for a single token decision. Here's real output from a training run:
token model p keyed U log(U)/p
------------------------------------------------
'a' 0.5883 0.7111903 -0.5793 <-- selected
'e' 0.1628 0.4810648 -4.4936
'i' 0.1606 0.3999150 -5.7064
'o' 0.0254 0.0220344 -150.4391
'y' 0.0219 0.4344855 -38.0628
'l' 0.0127 0.8099715 -16.5407
'r' 0.0077 0.2737357 -168.6167
'<BOS>' 0.0073 0.5730758 -76.7202
Two completely different things are happening here. The GPT supplies model p. The watermark supplies keyed U. Those two values combine to decide the winner, and here's the counterintuitive part: the GPT's probability distribution hasn't changed at all. It still believes 'a' has probability 58.83%. Only the source of randomness used to sample from that distribution has changed.
Same GPT, Different Sampler
Running the same trained model two ways ordinary random.choices() versus the keyed sampler produced this:
prompt normal watermarked
--------------------------------------
yuh yuha yuhan
dio dion dioeen
xav xavinn xavia
jor jori jorsen
jua juan juale
era eranan eran
phi philel phian
sam samen samir
pho phoran phonn
emm emman emman
Notice the last row: normal and watermarked sampling produced the exact same name, emman. There isn't a special "watermarked vocabulary" or a hidden character being inserted. The model is the same. Only the sampler changed.
But Where Is the Watermark?
You can't look at samir and point to the watermark there's no letter carrying it. It emerges statistically across many token choices, the way a biased coin only reveals its bias after many flips.
Because the detector knows the secret key, it can reconstruct the pseudorandom value tied to every token that was actually generated, and accumulate a score across them:
S = sum_t -log(1 - r_t)Under ordinary, non-watermarked text, each token contributes about 1 to that sum on average:
E[-log(1-r)] = 1So normal text should hover around a mean score of roughly 1. Watermarked generation preferentially selects tokens tied to unusually favourable keyed values, and the score rises.
detector.py implements this with math.log1p(-uniform(...)) (a numerically stable form of -log(1-U)) and turns the accumulated score into a p-value via a log-sum-exp Gamma-tail calculation, so it doesn't quietly underflow to zero on strong evidence. It also de-duplicates repeated token contexts across samples a repeated context reuses the same keyed random values, so counting it twice would manufacture fake confidence.
Does It Actually Work?
After training microgpt for 1,000 steps on 32,033 names (final loss 2.6497), I collected a few hundred generated tokens and ran the detector myself:
tokens mean score p-value
---------------------------------------------------------
normal text 308 0.962 7.459e-01
watermarked text 322 2.105 6.844e-53
wrong secret key 322 0.984 6.094e-01
Normal text and text checked with the wrong key both hover around a mean of 1, exactly as the null hypothesis predicts. Correctly watermarked text jumps to a mean of 2.105, with a p-value of 6.844 × 10⁻⁵³ an extraordinarily small probability of seeing evidence this strong by chance. The detector isn't asking "does this look AI-generated?" It's asking something much more specific: do these token choices have the statistical relationship expected from this particular secret key? Using the wrong key destroys the signal completely, which is exactly what you'd want from something claiming to detect a keyed watermark rather than generic "AI-ness."
91 Lines
The repository is deliberately boring:
gpt-watermark/
├── README.md
├── input.txt
├── microgpt_watermark.py (microgpt + the watermark hookup)
├── watermark.py (46 lines the keyed sampler)
└── detector.py (45 lines the statistical detector)
Unlike a from-scratch reimplementation, Karpathy's microgpt isn't split out into its own file here it lives inside microgpt_watermark.py alongside the integration that swaps in the watermarked sampler, which is why the whole repository is only four small files. The part that actually matters the watermark and the detector is 91 lines of pure Python.
Why Keep It This Small?
Because this isn't meant to be a production watermarking library there are much more sophisticated schemes for that. The purpose here is understanding. Karpathy describes microgpt as containing the algorithmic essence of a GPT, and I wanted to preserve that spirit for watermarking: strip it down until you can see exactly where it enters generation and why detection is possible.
logits
↓
softmax
↓
probabilities
↓
sampling ← watermark goes here
↓
next tokenOnce you see that, a lot of the mystery disappears.
The Catch
We haven't found the One Piece just yet. The watermark lives in the particular sequence of token choices produced during generation. Change enough of those choices and you damage the statistical signal rewrite the text, paraphrase it heavily, translate it into another language and back. The original token sequence changes, and detection can become much weaker. Token-level watermarking gives us something remarkable, but not indestructible.
Can We Watermark Meaning?
Suppose an LLM writes:
The pirate sailed across the sea searching for treasure.
And someone rewrites it as:
A sailor crossed the ocean in pursuit of riches.
Almost every token changed. The meaning didn't. Could an AI leave a signature in the idea itself one that survives paraphrasing, translation, rewriting? Could we watermark semantics rather than tokens? Aaronson has pointed to this as one of the deeper open questions raised by LLM watermarking, and there's active theoretical work exploring what is and isn't possible. I don't know the answer. But that's what makes the next island interesting.
Acknowledgements
This experiment stands very deliberately on other people's work. Andrej Karpathy's microgpt is the foundation his ability to strip complicated systems down to their algorithmic essentials is what made this possible, and the philosophy I tried hardest to be inspired by and to follow. Scott Aaronson's work on LLM watermarking provides the central idea explored here: replacing some of the randomness used during generation with keyed pseudorandomness to create a detectable statistical signal while preserving the model's token probabilities. And Google DeepMind's SynthID demonstrates that text watermarking isn't merely a theoretical curiosity related ideas can be engineered and deployed in real generative AI systems.
My contribution here is much smaller: I wanted to understand it, so I built the smallest version I could.
Where It Lives
The code is open-source, dependency-free, and small enough to read in one sitting:
Getting started is three lines:
git clone https://github.com/berba-q/gpt-watermark
cd gpt-watermark
python3 microgpt_watermark.pyNo pip install, no requirements.txt just the standard library and about a minute of training on a laptop.
"The sea is vast. The data is deep. But sometimes the most interesting treasure is finding out exactly where the randomness lives."
Griffiths, Log 003