Exploring text generation watermarking

Anthropic recently announced that in compliance with a new EU AI generated content law, they will be watermarking content generated using their models. While there are differing opinions on whether this is a good or bad thing as well as excellent interactive explainers on how it works, I was interested in the underlying mechanics behind it. Anthropic cites SynthID-Text as their methodology which builds upon work done here.

Text watermarking

To summarize, at each token position, the vocabulary is partitioned into a green/red list where the goal is to bias the generated token towards the green list. A key phrase/string is used as an RNG seed so the generated green/red lists can be deterministically derived.

I made a minimal sample to explore the process

import torch
import hashlib
from transformers import AutoTokenizer, Qwen3_5ForCausalLM, LogitsProcessor


class Watermarking(LogitsProcessor):
    def __init__(self, secret_key: str, gamma: float = 0.5, delta: float = 4.0):
        self.secret_key = secret_key
        self.gamma = gamma
        self.delta = delta

    def __call__(
        self,
        input_ids: torch.LongTensor,
        scores: torch.FloatTensor,
    ) -> torch.FloatTensor:
        batch_size, vocab_size = scores.shape

        for i in range(batch_size):
            prev_token_id = input_ids[i, -1].item()
            # Compute a unique seed using the previous token and secret
            seed_material = f"{self.secret_key}-{prev_token_id}".encode()
            digest = hashlib.sha256(seed_material).digest()
            # torch seeds need to fit in 64 bits — take the first 8 bytes
            seed = int.from_bytes(digest[:8], "big")

            # Compute a mask for green token ids
            generator = torch.Generator(device="cpu")
            generator.manual_seed(seed)

            perm = torch.randperm(vocab_size, generator=generator)
            cutoff = int(vocab_size * self.gamma)
            green_ids = perm[:cutoff]

            mask = torch.zeros(vocab_size, dtype=torch.bool)
            mask[green_ids] = True

            # Boost the probability of green token ids
            scores[i, mask] += self.delta

        return scores


def main():
    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    model_name = "Qwen/Qwen3.5-4B"

    model = Qwen3_5ForCausalLM.from_pretrained(model_name, device_map={"": device})
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    messages = [
        {
            "role": "user",
            "content": "What is the distance between earth and the sun?",
        }
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        enable_thinking=False,  # Disables the <think> block
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    )
    inputs = inputs.to(device)

    outputs_watermarked = model.generate(
        **inputs,
        max_new_tokens=70,
        do_sample=True,
        temperature=0.7,
        logits_processor=[
            Watermarking(
                secret_key="updownleftrightba",
            )
        ],
    )
    outputs_plain = model.generate(
        **inputs, max_new_tokens=70, do_sample=True, temperature=0.7
    )

    print(
        "Watermarked:",
        tokenizer.batch_decode(outputs_watermarked, skip_special_tokens=True)[0],
    )
    print(
        "Plain:      ",
        tokenizer.batch_decode(outputs_plain, skip_special_tokens=True)[0],
    )


if __name__ == "__main__":
    main()

Running some test generations colored red/green for visualization purposes. It's pretty easy to see the biasing in action.

Prompt: Summarize the plot of the hound of the baskervilles. Gamma: 0.5, Delta: 2.0

A gamma of 0.5 partitions the vocabulary evenly, while the delta represents how strongly to bias the green list. In this example, there's clearly a bias towards the green list which we can validate for by reversing the process. Generation quality is consistent but what's interesting is how nearly imperceptible it us to the reader. The model maintains coherence despite the change in probability distribution.

Since this only skews the probability distribution and does not remove tokens from consideration, other weighting techniques can be layered on top of it without much effort. For example, enforcing a grammar/schema to remove tokens from consideration.


JSON_SUMMARY_PATTERN = regex.compile(r'\{"summary": "[^"\n]+"\}')


class JsonLogitsProcessor(LogitsProcessor):
    """Rudimentary JSON-schema-constrained decoding: at each step, keep only the
    tokens whose text is a valid partial match for the target pattern."""

    def __init__(self, tokenizer, pattern: regex.Pattern, prompt_len: int):
        self.tokenizer = tokenizer
        self.pattern = pattern
        self.prompt_len = prompt_len
        self._token_texts: list[str] | None = None

    def __call__(
        self,
        input_ids: torch.LongTensor,
        scores: torch.FloatTensor,
    ) -> torch.FloatTensor:
        batch_size, vocab_size = scores.shape

        if self._token_texts is None:
            self._token_texts = [
                self.tokenizer.decode([token_id]) for token_id in range(vocab_size)
            ]

        for i in range(batch_size):
            prefix = self.tokenizer.decode(input_ids[i, self.prompt_len :])
            mask = torch.ones(vocab_size, dtype=torch.bool)

            if self.pattern.fullmatch(prefix):
                # Schema already satisfied — only allow the sequence to end
                mask[self.tokenizer.eos_token_id] = False
            else:
                for token_id, token_text in enumerate(self._token_texts):
                    if self.pattern.match(prefix + token_text, partial=True):
                        mask[token_id] = False

            scores[i, mask] = float("-inf")

        return scores

In this case I apply a JSON schema to force the output to begin with {"summary":. This can easily be applied after applying the watermarking to enforce our schema.

When tabulating the results, you may want to exclude fixed tokens with not much choice (JSON punctuation etc) to avoid skewing the results.

The key here is maintaining control over the probability distribution generated and more intelligently sampling/partitioning your vocabulary. This opens up a lot of doors to do things like sample or watermark differently depending on the topic of the generation, or potentially encode hidden information in the generated content.

Encoding messages in watermarks

Since we control the sampling process encoding a message is a matter of encoding a recognizable pattern that can be read by the decoding party. My first pass at this involved further splitting 2 partitions representing 2 bits. The decoder reads green list items then reads the message based on which sub partition the target token falls into.

class WatermarkWithEncodedMessage(LogitsProcessor):
    def __init__(
        self,
        secret_key: str,
        gamma: float = 0.5,
        delta: float = 1.0,
        message: str | None = None,
        message_delta: float = 1.0,
    ):
        self.secret_key = secret_key
        self.gamma = gamma
        self.delta = delta
        # Repeats across the whole generation (bit_index = step % len(message_bits)),
        # so the decoder can recover each bit by majority vote over many observations.
        self.message_bits = text_to_bits(message) if message else None
        self.message_delta = message_delta
        self.step = 0

    def __call__(
        self,
        input_ids: torch.LongTensor,
        scores: torch.FloatTensor,
    ) -> torch.FloatTensor:
        batch_size, vocab_size = scores.shape

        target_bit = None
        if self.message_bits:
            target_bit = self.message_bits[self.step % len(self.message_bits)]

        for i in range(batch_size):
            prev_token_id = input_ids[i, -1].item()
            ids = green_ids(self.secret_key, prev_token_id, vocab_size, self.gamma)
            # Boost the probability of green token ids
            scores[i, ids] += self.delta

            if target_bit is not None:
                # Additionally boost whichever green half encodes the current
                # message bit, without touching the overall green/red split.
                half = len(ids) // 2
                signal_ids = ids[half:] if target_bit else ids[:half]
                scores[i, signal_ids] += self.message_delta

        self.step += 1
        return scores


This is obviously error prone, considering that the text can be altered. So I encode the message within the text in a loop and decode the bits based off of a majority voting system as a crude error correction mechanism.

Encoding 'hello world' as the message resulted in a 97% recovery accuracy encoded over 600 tokens.

Decoding over text that was not watermarked predictably returns garbage output.

The full source of my experiment is available here. As stated earlier, this was just a first pass exploring the content. There's a lot of room for improvement and refinement.