# SynthID Text

This repository provides a reference implementation of the SynthID Text
watermarking and detection capabilities for the [research paper][nature-paper]
published in _Nature_. It is not intended for production use. The core library
is [distributed on PyPI][synthid-pypi] for easy installation in the
[Python Notebook example][synthid-colab], which demonstrates how to apply these
tools with the [Gemma][gemma] and [GPT-2][gpt2] models.

## Installation and usage

The [Colab Notebook][synthid-colab] is self-contained reference implementation
that:

1.  Extends the [`Gemma3ForCausalLM`][transformers-gemma3] and
    [`GPT2LMHeadModel`][transformers-gpt2] classes from
    [Hugging Face Transformers][transformers] with a [mix-in][synthid-mixin] to
    enable watermarking text content generated by models running in
    [PyTorch][pytorch]; and
1.  Detects the watermark. This can be done either with the [Frequentist
    detector][synthid-detector-frequentist] which requires no training, or with
    the more powerful [Bayesian detector][synthid-detector-bayesian] that
    requires [training][synthid-detector-trainer]. The Frequentist detector
    returns an exact binomial p-value for the null hypothesis that the text was
    not watermarked with the given key, so no length-dependent threshold needs
    to be tuned.

The notebook is designed to be run end-to-end with either Gemma 3 270M IT or
GPT-2, and runs best on the following runtime hardware:

*   Gemma 3 270M IT: Any GPU will work, such as a T4.
*   GPT-2: Any runtime will work, though a High-RAM CPU or any GPU will be
    faster.

Generation and g-value computation run in PyTorch and use the GPU when one is
available. The Bayesian detector is trained and scored with JAX, which is
installed as a CPU build by default; the detector is small, so CPU is
sufficient. To train it on an NVIDIA GPU instead, install the `cuda` extra, e.g.
`pip install synthid-text[notebook,cuda]`.

NOTE: This implementation is for reference and research reproducibility purposes
only. Due to minor variations in Gemma models across
implementations, we expect minor fluctuations in the detectability and
perplexity results obtained from this repository versus those reported in the
paper. The subclasses introduced herein are not designed to be used in
production systems.

NOTE: This implementation utilizes the 20-round Threefry 2x32 counter-based
PRNG for generation-time context hashing and direct Bernoulli(0.5) g-value bit
extraction. Threefry is a well-studied PRNG designed for statistical quality,
but it is not a vetted cryptographic primitive and keys are limited to 63 bits.
No guarantees of cryptographic security are provided.

### Local notebook use

The notebook can also be used locally if installed from source. Using a virtual
environment is highly recommended for any local use.

```shell
# Create and activate the virtual environment
python3 -m venv ~/.venvs/synthid
source ~/.venvs/synthid/bin/activate

# Download and install SynthID Text and Jupyter
git clone https://github.com/google-deepmind/synthid-text.git
cd synthid-text
pip install '.[notebook-local]'

# Start the Jupyter server
python -m notebook
```

Once your kernel is running, open the `.ipynb` notebook file to execute it.

### Running the tests

The source installation also includes a small test suite to verify that the
library is working as expected.

```shell
# Create and activate the virtual environment
python3 -m venv ~/.venvs/synthid
source ~/.venvs/synthid/bin/activate

# Download and install SynthID Text with test dependencies from source
git clone https://github.com/google-deepmind/synthid-text.git
cd synthid-text
pip install '.[test]'

# Run the test suite
pytest .
```

## How it works

For a description of tournament watermarking, speculative sampling with
dual-key watermarking, detection, and recommended tests, with pseudocode, see
[docs/tournament.md](./docs/tournament.md).

### Defining a watermark configuration

SynthID Text produces unique watermarks given a configuration, with the key
components being the secret `key: int` (in `[0, 2**63)`) and tournament
`depth: int` (number of tournament layers, default 30).

The structure of a configuration is described in the following `TypedDict`
subclass, though in practice, the [mixin][synthid-mixin] class in this library
uses a static configuration.

```python
from typing import TypedDict

import torch


class WatermarkingConfig(TypedDict):
    ngram_len: int
    key: int
    depth: int
    context_history_size: int
    device: torch.device
    skip_first_ngram_calls: bool
```

### Applying a watermark

Watermarks are applied by a [mix-in][synthid-mixin] class that wraps the
[`GemmaForCausalLM`][transformers-gemma],
[`Gemma3ForCausalLM`][transformers-gemma3] (text-only Gemma 3 models such as
Gemma 3 270M IT) and [`GPT2LMHeadModel`][transformers-gpt2] classes from
Transformers, which results in three subclasses with the same API that you are
used to from Transformers.
Remember that the mix-in provided by this library uses a static watermarking
configuration, making it unsuitable for production use.

The mix-in requires Transformers v5 (`transformers>=5.16`). The watermark is
applied after all of the logits processors and sampling warpers that
`generate()` builds (temperature, top-k, top-p, repetition penalty, ...), so it
watermarks exactly the distribution that would otherwise be sampled from.
Watermarked generation requires `do_sample=True` and `top_k >= 2`; top-k bounds
the number of candidate tokens that are watermarked at each step. The default
configuration sets `skip_first_ngram_calls=True`, so the first `ngram_len - 1`
generated tokens are not watermarked; detection never scores these tokens.

With `return_dict_in_generate=True, output_scores=True`, the returned `scores`
are the processed scores after watermarking, so they reflect the watermarked
distribution rather than the model's, and the mix-in warns about this. Use
`output_logits=True` for the model's raw logits, e.g. to compute perplexity.

```python
from synthid_text import synthid_mixin
import transformers
import torch


DEVICE = (
    torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
)
INPUTS = [
    "I enjoy walking with my cute dog",
    "I am from New York",
    "The test was not so very hard after all",
    "I don't think they can score twice in so short a time",
]
MODEL_NAME = 'google/gemma-3-270m-it'
TEMPERATURE = 0.5
TOP_K = 40
TOP_P = 0.99

# Initialize a standard tokenizer from Transformers. Decoder-only models need
# left padding for batched generation.
tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.padding_side = 'left'
# Initialize a SynthID Text-enabled model.
model = synthid_mixin.SynthIDGemma3ForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map='auto',
    dtype=torch.bfloat16,
)
# Prepare your inputs in the usual way.
inputs = tokenizer(
    INPUTS,
    return_tensors='pt',
    padding=True,
).to(DEVICE)
# Generate watermarked text.
outputs = model.generate(
    **inputs,
    do_sample=True,
    max_length=1024,
    temperature=TEMPERATURE,
    top_k=TOP_K,
    top_p=TOP_P,
)
```

### Detecting a watermark

Watermark detection can be done using a variety of scoring functions (see
paper). This repository contains code for the Frequentist and Bayesian
detectors described in the paper. The colab contains examples for how
to use these detectors.

Both detectors below take the generated token ids, with the prompt removed, and
compute the g-values and masks internally. The logits processor used for
detection must have the same watermarking config as the one used for
generation.

```python
from synthid_text import detector_frequentist
from synthid_text import logits_processing

logits_processor = logits_processing.SynthIDLogitsProcessor(
    **synthid_mixin.DEFAULT_WATERMARKING_CONFIG
)

# Get only the generated text from the model's predictions.
outputs = outputs[:, inputs['input_ids'].shape[1]:]
```

The Frequentist detector requires no training. Its `p_value()` function returns,
per example, the p-value for the null hypothesis that the text was not
watermarked with the given key. Small values are evidence of the watermark.

```python
frequentist_detector = detector_frequentist.FrequentistDetector(
    logits_processor, eos_token_id=model.generation_config.eos_token_id
)
p_values = frequentist_detector.p_value(outputs)
```

The Bayesian detector must be trained on watermarked and unwatermarked data
before it can be used. The Bayesian detector must be trained for each unique
watermarking key, and the training data used for this detector model should be
independent from, but representative of the expected character and quality of
the text content the system will generate in production.

```python
from synthid_text import detector_bayesian

# Lists of [batch_size, length] token id arrays, with prompts removed:
# responses generated with the watermark, and unwatermarked text tokenized with
# the same tokenizer.
tokenized_wm_outputs = ...
tokenized_uwm_outputs = ...

detector, selection_loss = detector_bayesian.BayesianDetector.train_best_detector(
    tokenized_wm_outputs=tokenized_wm_outputs,
    tokenized_uwm_outputs=tokenized_uwm_outputs,
    logits_processor=logits_processor,
    eos_token_id=model.generation_config.eos_token_id,
    torch_device=DEVICE,
)
```

Once the Bayesian detector is trained, use the `detector.score()` function to
generate a per-example score indicating if the text was generated with the given
watermarking configuration. Score values will be between 0 and 1, with scores
closer to 1 indicating higher likelihood that the text was generated with the
given watermark. You can adjust the acceptance threshold to your needs.

```python
bayesian_scores = detector.score(outputs)
```

The g-values and the mask of scored positions are also available directly:

```python
g_values, mask = logits_processor.compute_g_values_and_mask(
    outputs, eos_token_id=model.generation_config.eos_token_id
)
```

## Human Data

We release the human evaluation data, where we compare watermarked text against
unwatermarked text generated from the Gemma 7B model. The data is located in
`data/human_eval.jsonl`. To get the prompts used for generating the responses,
please use the following code.

```python
import json
import tensorflow_datasets as tfds

ds = tfds.load('huggingface:eli5/LFQA_reddit', split='test_eli5')
id_to_prompt = {}
for x in ds.as_numpy_iterator():
  id_to_prompt[x['q_id'].decode()] = x['title'].decode()

full_data = []
with open('./data/human_eval.jsonl') as f:
  for json_str in f:
    x = json.loads(json_str)
    x['question'] = id_to_prompt[x['q_id']]
    full_data.append(x)
```

## Dual-Key Speculative Decoding

The `speculative_engine` module provides a standalone implementation of
dual-key tournament watermarked speculative sampling. It accelerates generation
by coupling a target model with a faster draft model that shares its tokenizer,
watermarking accepted draft tokens with `key_draft` and replacement/bonus
tokens with `key_target`. Averaged over watermarking keys, each generated token
follows the target model's top-k distribution, as in autoregressive top-k
sampling. For an interactive walkthrough that visualizes token origins,
acceptance rates, and detector $p$-values, see the
[speculative decoding notebook in Colab][synthid-speculative-colab].

```python
from synthid_text import detector_frequentist
from synthid_text import speculative_engine
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained("gpt2-xl")
target_model = AutoModelForCausalLM.from_pretrained("gpt2-xl").to(device)
draft_model = AutoModelForCausalLM.from_pretrained("distilgpt2").to(device)

# Watermarking runs on the CPU (watermark_device defaults to CPU), while the
# models run on `device`.
config = speculative_engine.DualKeyWatermarkingConfig(
    ngram_len=5,
    key_draft=42,
    key_target=84,
    depth=15,
    temperature=0.7,
    top_k=40,
    draft_len=4,
    target_device=device,
    draft_device=device,
)

engine = speculative_engine.DualKeySpeculativeEngine(
    target_model=target_model,
    draft_model=draft_model,
    config=config,
)

# Generate watermarked text with dual-key speculative decoding. The engine
# generates a single sequence at a time (batch size 1).
prompt_ids = tokenizer(
    "Write a short story about a lighthouse.", return_tensors="pt"
).input_ids
output = engine.generate(
    prompt_ids,
    max_new_tokens=100,
    eos_token_id=target_model.generation_config.eos_token_id,
)
print(tokenizer.decode(output.sequences[0]))

# Statistical detection aggregating draft and target key signals. Only the
# generated tokens are scored; output.sequences also contains the prompt.
generated_ids = output.sequences[:, prompt_ids.shape[1]:]
detector = detector_frequentist.DualKeyFrequentistDetector.from_config(
    config, eos_token_id=target_model.generation_config.eos_token_id
)
p_value = detector.p_value(generated_ids)[0]
print(f"p-value: {p_value:.4e}")
```

## Citing this work

```bibtex
@article{Dathathri2024,
    author={Dathathri, Sumanth and See, Abigail and Ghaisas, Sumedh and Huang, Po-Sen and McAdam, Rob and Welbl, Johannes and Bachani, Vandana and Kaskasoli, Alex and Stanforth, Robert and Matejovicova, Tatiana and Hayes, Jamie and Vyas, Nidhi and Merey, Majd Al and Brown-Cohen, Jonah and Bunel, Rudy and Balle, Borja and Cemgil, Taylan and Ahmed, Zahra and Stacpoole, Kitty and Shumailov, Ilia and Baetu, Ciprian and Gowal, Sven and Hassabis, Demis and Kohli, Pushmeet},
    title={Scalable watermarking for identifying large language model outputs},
    journal={Nature},
    year={2024},
    month={Oct},
    day={01},
    volume={634},
    number={8035},
    pages={818-823},
    issn={1476-4687},
    doi={10.1038/s41586-024-08025-4},
    url={https://doi.org/10.1038/s41586-024-08025-4}
}
```

## License and disclaimer

Copyright 2024 DeepMind Technologies Limited

All software is licensed under the Apache License, Version 2.0 (Apache 2.0);
you may not use this file except in compliance with the Apache 2.0 license.
You may obtain a copy of the Apache 2.0 license at:
https://www.apache.org/licenses/LICENSE-2.0

All other materials are licensed under the Creative Commons Attribution 4.0
International License (CC-BY). You may obtain a copy of the CC-BY license at:
https://creativecommons.org/licenses/by/4.0/legalcode

Unless required by applicable law or agreed to in writing, all software and
materials distributed here under the Apache 2.0 or CC-BY licenses are
distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
either express or implied. See the licenses for the specific language governing
permissions and limitations under those licenses.

This is not an official Google product.

[flax]: https://github.com/google/flax
[gemma]: https://ai.google.dev/gemma/docs/model_card
[gpt2]: https://huggingface.co/openai-community/gpt2
[jax]: https://github.com/jax-ml/jax
[pytorch]: https://pytorch.org/
[nature-paper]: https://doi.org/10.1038/s41586-024-08025-4
[synthid-colab]: https://colab.research.google.com/github/google-deepmind/synthid-text/blob/main/notebooks/synthid_text_huggingface_integration.ipynb
[synthid-speculative-colab]: https://colab.research.google.com/github/google-deepmind/synthid-text/blob/main/notebooks/synthid_text_speculative_decoding.ipynb
[synthid-pypi]: https://pypi.org/project/synthid-text/
[synthid-detector-bayesian]: ./src/synthid_text/detector_bayesian.py
[synthid-detector-frequentist]: ./src/synthid_text/detector_frequentist.py
[synthid-detector-trainer]: ./src/synthid_text/detector_bayesian.py
[synthid-mixin]: ./src/synthid_text/synthid_mixin.py
[transformers]: https://github.com/huggingface/transformers
[transformers-gemma]: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma/modeling_gemma.py
[transformers-gemma3]: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma3/modeling_gemma3.py
[transformers-gpt2]: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py