ShapeMark: Robust and Diversity-Preserving Watermarking for Diffusion Models

ShapeMark is a robust and diversity-preserving watermarking method for diffusion models that encodes watermark bits into structured noise patterns rather than individual values and employs a randomization design to prevent fixed patterns, thereby achieving state-of-the-art robustness against lossy scenarios while maintaining high generation quality.

The Gist

Imagine you are an artist who can instantly create beautiful paintings using a magical machine (a Diffusion Model). These paintings look so real that people can't tell they were made by a computer. But there's a problem: if someone steals your painting, sells it, or claims they made it, how do you prove it's yours?

You need a watermark. But here's the catch:

1. It must be invisible: You don't want a giant "PROPERTY OF YUQI" stamp ruining the art.
2. It must be unbreakable: If someone crops the image, compresses it for the web, or adds a filter, the watermark shouldn't disappear.
3. It must not ruin the art: The machine should still be able to paint different beautiful pictures every time, not just the same one over and over.

Existing methods tried to solve this by hiding tiny secrets in the "static" (noise) the machine uses to start painting. But they had two big flaws:

• The "Fragile Pixel" Problem: They hid the secret by changing the exact brightness of a single pixel. If you squint or the image gets blurry, that tiny change disappears, and the secret is lost.
• The "Stuck Record" Problem: To make the secret stronger, they repeated the same pattern over and over. This made the machine start painting the same background noise every time, killing the creativity and variety of the images.

Enter ShapeMark. Think of it as a new, super-smart way to hide a secret in the chaos.

The Magic Trick: "The Shuffle" instead of "The Change"

Instead of changing the value of a single pixel (like changing a red dot to a blue dot), ShapeMark changes the order of the dots.

1. The "Sorting Hat" Analogy (Structural Encoding)

Imagine you have a giant bag of marbles of different sizes (the noise).

• Old Method: You try to hide a secret by painting a tiny dot on a specific marble. If someone rubs the marble, the dot fades.
• ShapeMark: You sort the marbles into piles based on size (smallest to largest). Then, you hide your secret by shuffling the order of the piles.
  • The Secret: "The small pile is on the left, the big pile is on the right."
  • Why it works: Even if someone smudges the marbles or the image gets blurry, the relative order (small vs. big) usually stays the same. You don't need to see the exact color of a marble; you just need to know which pile is which. This makes the watermark super robust against damage.

2. The "Random Seat Swap" Analogy (Payload-Debiasing)

Here is the second problem: If you always shuffle the piles in the exact same way for the same secret, the machine might start looking for that specific shuffle pattern, making every image look slightly similar (boring).

• ShapeMark's Solution: Before you shuffle the piles, you do a random seat swap that changes every single time you make a picture.
  • Imagine you have a secret code for "Red." Usually, you put the Red pile in Seat A. But today, you randomly swap Seat A with Seat Z. Tomorrow, you swap Seat A with Seat B.
  • The secret is still "Red," but the location of the Red pile is different every time.
  • The Result: The machine creates a totally unique, diverse image every time, but because you know the "seat swap" rule (the key), you can still find the secret later. This keeps the art diverse and fresh.

How to Find the Secret (Decoding)

When you want to check if an image is yours:

1. You take the image and run it backward through the machine to get the original "bag of marbles" (the noise).
2. You look at the order of the piles. Did they follow your secret shuffle pattern?
3. Even if the image was cropped or compressed, the pattern of the shuffle is usually still there, like recognizing a song even if the radio is static-filled.

Why is this a Big Deal?

• It's Tough: It survives JPEG compression, cropping, blurring, and even noise. It's like a tattoo that stays visible even if you get a sunburn.
• It's Diverse: It doesn't force the AI to be boring. It lets the AI create millions of unique variations while still carrying your invisible ID card.
• It's Fair: It doesn't require retraining the AI model. It just changes how the machine starts its process.

In short: ShapeMark is like hiding a secret message in the dance steps of a crowd, rather than writing it on one person's forehead. Even if the crowd gets jostled (distorted) or the person moves (randomized), you can still see the dance pattern and know who they are.

Technical Summary

Here is a detailed technical summary of the paper "ShapeMark: Robust and Diversity-Preserving Watermarking for Diffusion Models."

1. Problem Statement

Diffusion models have revolutionized high-fidelity image synthesis, but their widespread use has created urgent challenges regarding intellectual property (IP) protection and content provenance. While Noise-as-Watermark (NaW) methods have emerged as a promising solution by embedding watermarks into the initial noise latent before generation, existing approaches suffer from a critical trade-off:

• Fragility: Most NaW methods rely on value encoding, where watermark bits are mapped to specific sampled values (e.g., sign or magnitude). This is highly fragile because lossy post-processing (compression, cropping, noise) and imperfect diffusion inversion make precise value recovery unreliable.
• Diversity Loss: To improve robustness, some methods repeatedly embed watermarks or create fixed noise patterns. This introduces payload-dependent spatial regularities, reducing the visual diversity of generated images and making watermark artifacts perceptible.

The core challenge is to design a watermarking scheme that is robust against distortions and inversion errors while preserving the high generation diversity and visual fidelity inherent to diffusion models.

2. Methodology: ShapeMark

ShapeMark addresses these challenges by shifting from value-based encoding to structural encoding and introducing a payload-debiasing randomization mechanism. The pipeline consists of two main stages:

A. Structural Encoding (SE)

Instead of encoding bits into individual noise values, SE encodes them into the relative structural relationships between groups of noise elements.

1. Separability-Guided Template Construction:
   • The noise latent is ranked by magnitude.
   • Indices are split into $Q$ quantile bins (groups of elements with similar magnitudes).
   • Within each bin, indices are shuffled using a secret key $\kappa$ and bundled into blocks.
   • Groups are formed by aligning one block from each quantile bin. This ensures that each group contains elements from statistically distinct magnitude regions.
2. Payload-Conditioned Permutation:
   • Watermark bits are mapped to a permutation codebook.
   • For each group, the blocks are rearranged (permuted) according to the payload bits.
   • Key Insight: The watermark is defined by the order of blocks within a group, not their absolute values. Since the permutation pattern is stable even if individual values are perturbed by noise or inversion errors, this approach is inherently robust.

B. Payload-Debiasing Structural Randomization (PDSR)

SE introduces a deterministic structural pattern based on the payload, which can lead to fixed spatial artifacts and reduced diversity. PDSR mitigates this:

• Mechanism: After SE, a global permutation is applied to the block positions. This permutation is derived from a public nonce (unique per image) and the secret key, but is independent of the payload.
• Effect: It decouples the payload identity from a fixed spatial location. The same watermark payload can correspond to diverse noise realizations, preventing the model from learning a specific "watermarked" spatial pattern.
• Invertibility: Since PDSR is a pure permutation, it is exactly reversible during decoding if the nonce is known.

C. Decoding Process

1. Inversion: The queried image is inverted back to the noise latent space using standard DDIM inversion.
2. Undoing PDSR: Using the public nonce and secret key, the verifier reverses the PDSR permutation to restore the SE-aligned structure.
3. Matching: The verifier reconstructs the reference groups (based on the key) and compares the observed block ordering in the recovered latent against the shared permutation codebook to decode the payload or verify identity.

3. Key Contributions

• Structural Encoding (SE): A novel encoding scheme that embeds watermarks into the permutation of noise blocks across quantile bins. This shifts the watermark signal from fragile absolute values to robust relative structures.
• Payload-Debiasing Structural Randomization (PDSR): A lightweight strategy that randomizes block positions per image without altering values. This eliminates payload-induced spatial biases, ensuring high generation diversity and imperceptibility.
• Robustness without Diversity Sacrifice: The method achieves state-of-the-art robustness against lossy attacks (JPEG, cropping, noise, etc.) while maintaining generation diversity comparable to unwatermarked models.

4. Experimental Results

The authors evaluated ShapeMark on Stable Diffusion v2.1 using the MS-COCO and SDP datasets, comparing it against post-processing, model-fine-tuning, and other NaW baselines.

• Robustness: ShapeMark achieved a True Positive Rate (TPR) of 0.999 at a False Positive Rate (FPR) of $10^{-6}$ under various distortions. It maintained a per-bit recovery accuracy of 98.70% even after severe attacks, outperforming existing NaW methods.
• Diversity: ShapeMark achieved the highest generation diversity among all methods (LPIPS 0.7338), significantly outperforming other NaW approaches (e.g., Tree-Ring: 0.6941, Gaussian Shading: 0.6322). This confirms that PDSR successfully prevents diversity collapse.
• Visual Quality: The method maintained competitive image quality metrics (CLIP score and FID) with no statistically significant degradation compared to the base model.
• Scalability: The system demonstrated robustness even at high payload capacities (up to 2048 bits), though accuracy naturally decreased slightly as capacity increased.

5. Significance

ShapeMark represents a significant advancement in AIGC provenance:

• Practical Deployment: By operating purely on the noise latent without modifying model weights or the final image, it is easily integrable into existing diffusion pipelines and compatible with various models.
• Solving the Trade-off: It successfully breaks the traditional trade-off where robustness comes at the cost of diversity. This makes it suitable for multi-tenant generative services where both IP protection and user experience (diversity) are critical.
• Resilience: The structural approach provides a new paradigm for watermarking that is resilient to the inevitable noise introduced by diffusion inversion and real-world image processing, offering a reliable mechanism for tracing AI-generated content.