WeaveMark

通过编码载荷扩展实现鲁棒且可扩展的多比特 LLM 水印的参考实现。

WeaveMark 在保持 token 分布无偏的同时，将 k 比特消息嵌入到 LLM 生成的文本中。它结合了编码载荷扩展（每 token 在 ℓ 层上承载 κ 个码字比特，并采用上下文打乱的层→比特分配）、基于每比特投票余量的软判决 ECC 解码，以及无偏多层重加权（每个方向与一次上下文抛硬币结果进行 XOR）。专用的零比特层从上下文中划分出来而非从消息中划分，从而提供与消息无关的存在性检测。

目录结构

weavemark/            库
  watermark.py        WeaveMark logits 处理器（嵌入）
  extraction.py       extract_bits, detect_zerobit
  prf.py              带密钥的 PRF（嵌入/提取共用）
  data.py             数据集加载、提示词准备、JSONL 记录
  device.py           CUDA 检查
  ecc/                Golay、Reed-Muller、软判决解码器
generate.py           生成带水印 / 普通文本
detect.py             提取 + 零比特 z-score + PPL
evaluate_downstream.py  摘要 / 翻译质量评估
perplexity.py         PPL 评分器
attacks/              替换攻击、DIPPER
run_*.py              流水线（生成+检测、攻击、下游任务）
data/OpenGen.jsonl    OpenGen 提示词（--dataset opengen）；见下方说明

data/OpenGen.jsonl（标准的公开 OpenGen 基准）因体积原因未包含在本归档中；请将其放在 data/ 目录下以使用 --dataset opengen。c4 和 openwebtext 数据集从 Hub 流式加载，无需本地文件。

请从仓库根目录运行所有脚本，以便 import weavemark 能够正常解析。

安装

需要 Python 3.10–3.12 和 CUDA GPU（模型通过 bitsandbytes 以 4-bit 加载；不支持 CPU 路径）。4-bit LLaMA-3-8B 约需 16 GB 显存。

cd WeaveMark
python -m venv .venv && source .venv/bin/activate     # 或：conda create -n weavemark python=3.11

在安装 requirements.txt 之前，请先安装 CUDA 版本的 PyTorch —— 直接执行 pip install torch 在 Windows 上会得到仅支持 CPU 的 wheel，且该 wheel 必须包含适用于你 GPU 的内核（RTX 50 系列 / sm_120 需要 CUDA ≥ 12.8 上的 torch ≥ 2.7；较旧的 cu121 wheel 最高支持 sm_90，运行时将失败）。请从 https://pytorch.org 选择匹配的命令。

pip install torch --index-url https://download.pytorch.org/whl/cu128   # 示例：Blackwell
pip install -r requirements.txt

检查 GPU 是否确实能运行内核（仅 is_available() 返回 True 是不够的 —— 即使 wheel 缺少适用于该架构的内核，它也会返回 True）：

python -c "import torch; t=torch.zeros(8).cuda().normal_(); print(torch.__version__, torch.cuda.get_device_capability(), (t*t).sum().item())"

LLaMA-3 和 C4 受访问限制：请执行 huggingface-cli login 并接受其条款。OpenGen 随仓库附带。使用 CUDA_VISIBLE_DEVICES=<i> 选择 GPU。

快速开始

python generate.py --method weavemark --ecc_method rm --random_message \
  --dataset c4 --num_test 100 --max_new_tokens 200 \
  --bpt 10 --num_layers 10 --window_size 2 --prob_delta 1.0 \
  --num_zerobit_layers 0 --do_sample

python detect.py --data_dir output_dump/<run_folder> --detect --stride 25

detect.py 会写入 detailed_eval_*.csv 和 summary_eval_*.csv；关键列是 success_rate，以及使用 --zerobit_threshold 时的 z_success_rate。run_main.py 串联了生成 + 检测（请编辑其顶部的常量）。

数据集：c4 / openwebtext 从 Hub 流式加载并截断为 --prompt_len 个词；opengen 读取 data/OpenGen.jsonl 并将每个前缀截断为 --prompt_len 个 token。

参数

ECC 消息→码字：golay 12→24，rm 16→32，dual_golay 24→48，dual_rm 32→64。none 以未编码方式嵌入 --message。

论文设置：多比特追踪使用 --bpt 10 --num_layers 10 --window_size 2 --prob_delta 1.0 --num_zerobit_layers 0，ECC 按消息长度选择；组合设置添加 --num_zerobit_layers 2 --enable_zerobit；纯零比特为 --bpt 0 --num_layers 0 --num_zerobit_layers 10 --enable_zerobit。

攻击与下游任务

python attacks/substitution.py --data_dir output_dump/<run> --ratios 0.1 0.2 0.3
python attacks/dipper_attack.py --data_dir output_dump/<run> --lex_diversity 20 --order_diversity 20
python detect.py --data_dir output_dump/<run> --text_file attacked_text_10.jsonl --detect --stride 999

run_substitution.py / run_dipper.py 串联了攻击 + 整篇文本检测（--stride 999 对每个片段进行整体评估）。run_downstream.py 运行摘要（CNN/DailyMail）和翻译（WMT16 en→ro）任务，并使用 BERTScore + ROUGE/BLEU 进行评估。

说明

• window_size = 2：无偏设计禁止复用上下文，因此更大的窗口会减少跳过损失（见附录）。
• 4-bit 量化下的绝对 PPL 会因重复而偏低；请将其视为相对指标。
• no_watermark 运行会复用相同的密钥/参数来测量误报率。

引用

@inproceedings{weavemark,
  title  = {Robust and Scalable Multi-bit LLM Watermarking via Coded Payload Spreading},
  author = {Anonymous},
  year   = {2026},
  note   = {Under review}
}

致谢与许可证

attacks/dipper.py 改编自 martiansideofthemoon/ai-detection-paraphrases （该文件的许可证适用于该文件）。以 MIT 许可证发布（见 LICENSE）；发布前请设置版权持有人。