Unigram Watermarking

Context-independent green/red list watermarking where the green list is fixed at initialization and does not depend on preceding tokens.

Example code: 📓 View On GitHub

Author: Zhao et al.

Theory

Unlike KGW/Maryland, which recomputes a new green/red partition at every token position based on the preceding \(h\) tokens, Unigram uses a single fixed partition determined at initialization by hashing a secret key.

The vocabulary \(V\) is split into a green set \(G\) (fraction \(\gamma\)) and red set \(R\) (fraction \(1-\gamma\)) using a SHA-256 hash of the secret hash_key. During generation, green tokens receive a logit bias \(\delta\):

\[\text{logit}_{wm}(v) = \text{logit}(v) + \delta \cdot \mathbf{1}_{v \in G}\]

Detection counts how many generated tokens fall in \(G\) and computes a z-score:

\[z = \frac{|s|_G - \gamma T}{\sqrt{\gamma(1-\gamma)T}}\]

where \(|s|_G\) is the count of green tokens and \(T\) is the total token count.

Note

Because the green list is context-independent, Unigram watermarks are more robust to text edits (insertions, deletions, reordering) than context-dependent methods like KGW. The trade-off is that a fixed partition can introduce more detectable artifacts in the token distribution.

Paper reference

Zhao, X., Ananth, P., Li, L., & Wang, Y.-X. (2024). Provable Robust Watermarking for AI-Generated Text. ICLR 2024. https://arxiv.org/abs/2306.17439

Example code

import os
os.environ["VLLM_USE_V1"] = "1"
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"

from vllm import LLM, SamplingParams
from vllm_watermark.core import (
    DetectionAlgorithm,
    WatermarkedLLMs,
    WatermarkingAlgorithm,
)
from vllm_watermark.watermark_detectors import WatermarkDetectors

llm = LLM(model="meta-llama/Llama-3.2-1B", enforce_eager=True, max_model_len=1024)

wm_llm = WatermarkedLLMs.create(
    model=llm,
    algo=WatermarkingAlgorithm.UNIGRAM,
    seed=42,
    ngram=1,
    gamma=0.5,
    delta=2.0,
    hash_key=15485863,
)

detector = WatermarkDetectors.create(
    algo=DetectionAlgorithm.UNIGRAM_Z,
    model=llm,
    ngram=1,
    seed=42,
    gamma=0.5,
    delta=2.0,
    hash_key=15485863,
    threshold=0.05,
)

prompts = ["Write about machine learning applications"]
sampling_params = SamplingParams(temperature=1.0, top_p=0.95, max_tokens=128)
outputs = wm_llm.generate(prompts, sampling_params)

for output in outputs:
    generated_text = output.outputs[0].text
    result = detector.detect(generated_text)
    print(f"Watermarked: {result['is_watermarked']}, P-value: {result['pvalue']:.6f}")