Technical Summary: Guidance Watermarking for Diffusion Models

Problem Statement

The rapid advancement of diffusion models has enabled the generation of high-quality, photorealistic images, raising significant concerns regarding the traceability and identification of AI-generated content. Regulatory bodies and industry stakeholders require robust mechanisms to distinguish synthetic media from real images to prevent model collapse (caused by training on AI-generated data) and to warn users of non-real content.

Existing watermarking strategies for diffusion models generally fall into two categories, both with limitations:

1. Post-hoc Watermarking: Embedding a signal into the image after generation. While mature, these methods treat the generated image as a "host" and often fail to integrate the watermark into the semantic content of the generation process, potentially limiting robustness.
2. In-Generation Watermarking: Modifying the generation process itself. Current approaches include:
   • Seed-based methods (e.g., Tree-Rings, Gaussian Shading): Embedding patterns in the initial noise seed. These often struggle with geometric attacks (e.g., cropping, rotation) and can alter the semantic content or composition of the image.
   • VAE-based methods (e.g., Stable Signature): Fine-tuning the Variational Autoencoder (VAE) at the end of the diffusion process. While effective, this focuses watermark energy on high-frequency details, making it vulnerable to low-pass filtering (e.g., JPEG compression). Furthermore, it requires retraining or fine-tuning the model.

A critical challenge remains: converting robust post-hoc watermarking schemes into in-generation methods without retraining the diffusion model or the watermark decoder, while maintaining image quality and diversity.

Methodology

The paper proposes Guidance Watermarking, a method that embeds watermarks during the diffusion process by guiding the sampling trajectory using gradients from an off-the-shelf watermark decoder.

Core Mechanism

The method treats the watermark decoder as a differentiable classifier. Instead of generating an image and then checking for a watermark, the diffusion process is steered to generate images that the decoder intrinsically identifies as watermarked.

1. Gradient-Based Guidance:
   The standard denoising step $\epsilon_\theta(z_t, t)$ is modified by subtracting a term proportional to the gradient of a loss function $L$ with respect to the latent $z_t$:
   $$\hat{\epsilon}(z_t, t) := \epsilon_\theta(z_t, t) - \omega \sqrt{1 - \bar{\alpha}_t} \nabla_{z_t} \log L(z_t)$$
   Here, $\omega$ controls the guidance strength.

2. Loss Function:
   The loss $L$ is defined to minimize the angle between the feature vector extracted from the generated image $\phi(x_0(z_t))$ and the secret watermark vector $u_m$. For zero-bit detection, this is equivalent to maximizing the cosine similarity:
   $$L(z_t, u_m) := 1 - \cos(\phi(x_0(z_t)), u_m)$$
   Minimizing this loss reduces the p-value for a watermarked image, increasing detection probability.

3. Robustness via Augmentation:
   To ensure robustness against attacks the decoder was not originally trained on, the gradient computation incorporates a set of image transformations $\mathcal{T}$ (e.g., JPEG compression, brightness/contrast adjustments, cropping). The gradients from these augmented views are aggregated (using the PCGrad algorithm to handle conflicting gradient directions) to update the noise estimate:
   $$\hat{\epsilon}_{\mathcal{T}}(z_t) := \epsilon_\theta(z_t) - \sqrt{1 - \bar{\alpha}_t} \text{Agg}(\{\nabla_{z_t} \log L(z_t, u_m; T) \mid T \in \mathcal{T}\})$$
   This allows the method to inherit the robustness of the pre-trained decoder and extend it to new transformations without retraining the decoder.

4. Optimizations:

   • Partial Guidance: Guidance is applied only during the final $T_w$ steps of the diffusion process to reduce computational cost.
   • Gradient Clipping: To prevent visual artifacts, the gradient norm is clipped, and the guidance strength is controlled via hyperparameters ($\omega, \eta, \tau$).

Key Contributions

1. First Guidance-Based In-Generation Watermarking: The paper presents the first method to embed watermarks directly into the diffusion process using gradient guidance, converting any post-hoc scheme into an in-generation one.
2. No Retraining Required: The method does not require fine-tuning the diffusion model or the watermark decoder. It utilizes existing, off-the-shelf decoders.
3. Enhanced Robustness: By incorporating augmentations into the gradient computation, the method improves robustness against attacks (including those the decoder was not trained on) without retraining the detector.
4. Balanced Trade-off: The approach strikes a balance between seed-based schemes (which alter semantic content) and VAE-based/post-hoc schemes (which add invisible signals). It preserves the semantic diversity and quality of the generation while embedding the watermark.

Experimental Results

The method was evaluated on three diffusion models (Stable Diffusion 2, Flux-1.0, Sana) using two state-of-the-art detectors (Stable Signature and VideoSeal).

• Image Quality: Quantitative metrics (FID, CLIP score) show that the watermarked images are qualitatively indistinguishable from non-watermarked images. Visual inspection confirms that semantic content, tone, and composition remain largely preserved, unlike the drastic changes seen in some seed-based methods.
• Robustness (Multi-bit & Zero-bit):
  • Geometric Attacks: The method significantly outperforms seed-based schemes (Tree-Rings, Gaussian Shading) against geometric attacks like cropping and rotation, inheriting the robustness of the pre-trained decoders.
  • Valuometric Attacks: It maintains high robustness against JPEG compression and brightness/contrast changes.
  • Adversarial Attacks: The method demonstrates superior robustness against VAE purification, regeneration, and average attacks compared to Tree-Rings and Gaussian Shading.
• False Alarm Rate: The method achieves extremely low false alarm rates (high $-\log_{10}(PFA)$) at high detection probabilities ($PD=0.9$), outperforming both post-hoc and other in-generation baselines.
• Efficiency: While generation requires 2–13 additional steps, decoding is instantaneous (single forward pass), making it up to 50 times faster at detection time compared to seed-based methods that require reverse diffusion steps.

Significance and Claims

The paper claims that Guidance Watermarking offers a principled, modular solution for securing generative AI. Its primary significance lies in:

• Complementarity: It works alongside existing techniques (like VAE fine-tuning) and can be applied to any diffusion model without architectural changes.
• Adaptability: It allows the immediate deployment of robust watermarking by leveraging existing, well-trained decoders, avoiding the resource-intensive process of retraining models or decoders for new attack vectors.
• Practicality: By preserving image quality and diversity while ensuring high detectability and low false positives, it addresses the critical need for traceable AI content in regulatory and security contexts.

The authors acknowledge limitations, noting that robustness depends on the visual content of the generated image and that generation time increases slightly. However, they emphasize that the trade-off is favorable given the elimination of retraining requirements and the significant gains in detection speed and robustness.