This is the official code for the paper:
BitMark: Watermarking Bitwise Autoregressive Image Generative Models
NeurIPS 2025
TL; DR: We introduce a bitwise, radioactive watermark for bitwise image generative models and show that it remains robust against various attacks while preserving high image quality and generation speed.
State-of-the-art text-to-image models like Infinity generate photorealistic images at an unprecedented speed. These models operate in a bitwise autoregressive manner over a discrete set of tokens that is practically infinite in size. However, their impressive generative power comes with a growing risk: as their outputs increasingly populate the Internet, they are likely to be scraped and reused as training data—potentially by the very same models. This phenomenon has been shown to lead to model collapse, where repeated training on generated content, especially from the models’ own previous versions, causes a gradual degradation in performance. A promising mitigation strategy is watermarking, which embeds human-imperceptible yet detectable signals into generated images—enabling the identification of generated content. In this work, we introduce BitMark, a robust bitwise watermarking framework. Our method embeds a watermark directly at the bit level of the token stream during the image generation process. Our bitwise watermark subtly influences the bits to preserve visual fidelity and generation speed while remaining robust against a spectrum of removal techniques. Furthermore, it exhibits high radioactivity, i.e., when watermarked generated images are used to train another image generative model, this second model’s outputs will also carry the watermark. The radioactive traces remain detectable even when only fine-tuning diffusion or image autoregressive models on images watermarked with our BitMark. Overall, our approach provides a principled step toward preventing model collapse in image generative models by enabling reliable detection of generated outputs.
Please first download the pretrained weights for the respective models by following the READMEs in the corresponding folder.
BitMark is defined in extended_watermark_processor.py.
The robustness evaluation is in tools/robustness_test.py and to run it, it requires a path of clean images and a path of watermarked images and then computes the TPR@1%FPR for all images in the watermarked path against all attacks.
Our novel BitFlipper attack is applied by running flipper.py
Infinity generates in a per scale fashion, where each scale can be understood as corresponding to a different resolution.
Install uv:
curl -LsSf https://astral.sh/uv/install.sh | shInstall dependencies in the root folder of this project:
uv syncIf you want to run the Jupyter notebook starter (recommended), specify bitmark as environment in ./scripts/infer_infinity.ipynb.
To download the required model weights, run:
uv run download_models.py --output-dir './weights'Otherwise, the file tools/comprehensive_infer.py offers a way to run the Infinity model and watermark images using our BitMark. The file can be executed by calling:
uv run tools/comprehensive_infer.py --model_path "./weights/Infinity/infinity_2b_reg.pth" --vae_type 32 --vae_path "./weights/Infinity/infinity_2b_reg.pth" --add_lvl_embeding_only_first_block 1 --model_type "infinity_2b" --seed 0 --watermark_scales 2 --watermark_delta 2 --watermark_context_width 2 --out_dir "./" --jsonl_filepath = "captions.json"
where --jsonl_filepath leads to a json, which contains the prompts used for image generation. The delta and sequence length (watermark_context_width) can be adapted if wished.
For our radioactivity experiments, we generated 1,000 watermarked images from MS-COCO prompts, watermarking all scales with delta = 2.
We fine-tuned Infinity for 1 epoch at a learning rate of 1e-4 on the 1,000 watermarked images. The fine-tuning script is located at ./Infinity/finetune_for_radioactivity.sh. An example of the JSONL format required by Infinity is provided in ./Infinity/finetune_jsonl_example.
Note: When loading a fine-tuned model, the flag
args.enable_model_cache=1must be set during inference.
We trained SD2.1 on the same dataset using the Diffusers text-to-image training script with the following hyperparameters:
--pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1-base \
--train_data_dir=$dataset_path \
--use_ema \
--resolution=512 \
--center_crop \
--train_batch_size=4 \
--gradient_checkpointing \
--num_train_epochs=5 \
--mixed_precision="fp16" \
--learning_rate=1e-04 \
--enable_xformers_memory_efficient_attention \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--checkpointing_steps 1000000 \
--seed 1 \
Note: SD2.1 is apparently not anymore supported by huggingface. At the current time, the following model should still be supported: Manojb/stable-diffusion-2-1-base
If you find our work useful, please cite our paper:
@inproceedings{
kerner2025bitmark,
title={BitMark: Watermarking Bitwise Autoregressive Image Generative Models},
author={Louis Kerner and Michel Meintz and Bihe Zhao and Franziska Boenisch and Adam Dziedzic},
booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
year={2025},
url={https://openreview.net/forum?id=VSir0FzFnP}
}The code is licensed under an MIT license. It relies on code and models from other repositories. See the next Acknowledgements section for the licenses of those dependencies.
This project builds on the following repositories:
We thank these great works!