Shallow Diffuse: Robust and Invisible Watermarking through Low-Dimensional Subspaces in Diffusion Models

1. Problem Statement

The rapid proliferation of AI-generated content (AIGC) via diffusion models (e.g., Stable Diffusion, DALL-E 3) has created significant challenges regarding misinformation, copyright infringement, and "model collapse" (degradation of models trained on synthetic data). While watermarking is a primary solution for attribution, existing methods suffer from critical limitations:

• Inconsistency: Methods that embed watermarks into the initial random noise (e.g., Tree-Ring, RingID) often distort the Gaussian noise distribution, leading to significant visual degradation and inconsistent image generation compared to the original model.
• Scenario Limitations: Many robust methods require access to the initial random seed (Server Scenario), making them unusable for post-processing existing images (User Scenario). Conversely, traditional post-processing methods are often fragile against simple attacks like blurring or compression.
• Trade-off: There is a lack of methods that simultaneously achieve high robustness, high visual consistency (imperceptibility), and flexibility across both server and user scenarios without requiring model retraining.

2. Methodology: Shallow Diffuse

The core innovation of Shallow Diffuse is the decoupling of the watermarking process from the image generation process by leveraging the low-dimensional subspace structure inherent in diffusion models.

Key Theoretical Insights

• Low-Rank Jacobian: The authors exploit the finding that the Posterior Mean Predictor (PMP) in diffusion models, denoted as $f_{\theta,t}(x_t)$, exhibits local linearity and its Jacobian matrix $J_{\theta,t}$ is low-rank during the middle timesteps of the diffusion process (specifically $t \in [0.2T, 0.7T]$).
• Null Space Separation: Because the Jacobian is low-rank ($r_t \ll d$, where $d$ is the ambient dimension), a significant portion of any perturbation vector lies in the null space of the Jacobian.
• Decoupling Strategy: By injecting the watermark into a specific timestep $t^*$ where the Jacobian is low-rank, the watermark perturbation $\lambda \Delta x$ is largely ignored by the generation process (since $J_{\theta,t} \Delta x \approx 0$). This ensures the predicted clean image $\hat{x}_0$ remains consistent with the unwatermarked version, while the watermark remains embedded in the residual noise term.

Algorithm Workflow

The method operates in two modes:

1. Server Scenario (Generation):
   • Start with a random seed $x_T$ and text prompt $c$.
   • Run the diffusion process forward to a specific timestep $t^*$ (typically $0.3T$) using Classifier-Free Guidance (CFG).
   • Inject the watermark $\lambda \Delta x$ into the latent representation $x_{t^*}$.
   • Continue the reverse sampling process to generate the final watermarked image $x_0^W$.
2. User Scenario (Post-Processing):
   • Take an existing image $x_0$.
   • Perform DDIM Inversion to map the image back to the latent space at timestep $t^*$ ($x_{t^*} = \text{DDIM-Inv}(x_0, t^*)$).
   • Inject the watermark $\lambda \Delta x$ into $x_{t^*}$.
   • Run the reverse sampling process from $t^*$ to $0$ to generate the watermarked image.

Watermark Design

• Frequency Domain: The watermark is designed in the frequency domain using Discrete Fourier Transform (DFT).
• High-Frequency Focus: Unlike previous methods that modify low-frequency components (causing visible artifacts), Shallow Diffuse targets high-frequency components.
• Pattern: The watermark consists of a circular mask $M$ and a key $W$ composed of concentric rings with values sampled from a Gaussian distribution. This design ensures the watermark is imperceptible (invisible) yet robust.

3. Key Contributions

1. Novel Decoupling Mechanism: The first training-free method to decouple watermarking from the sampling process by exploiting the low-dimensional null space of the diffusion Jacobian.
2. Dual-Scenario Flexibility: Uniquely supports both Server-side (generation-time) and User-side (post-processing) scenarios without requiring access to the initial seed or model retraining.
3. Theoretical Guarantees: Provides rigorous proofs (Theorem 1 and 2) showing that:
   • Consistency: The $\ell_2$ distance between watermarked and original images is bounded by a term dependent on the low rank $r_t$, not the high ambient dimension $d$.
   • Detectability: The watermark remains recoverable with high probability after DDIM inversion.
4. Training-Free: Unlike adversarial optimization methods (e.g., ROBIN), Shallow Diffuse requires no fine-tuning or retraining of the diffusion model.

4. Experimental Results

The authors evaluated Shallow Diffuse against state-of-the-art baselines (Tree-Ring, RingID, Gaussian Shading, Stable Signature, etc.) on Stable Diffusion 2.1.

• Consistency (Visual Quality):
  • Server Scenario: Shallow Diffuse achieved the highest PSNR (35.49), SSIM (0.96), and lowest LPIPS (0.02) among diffusion-based methods, significantly outperforming Tree-Ring (PSNR 16.61) and RingID (PSNR 14.27).
  • User Scenario: It achieved consistency scores nearly identical to the theoretical upper bound (SD w/o WM), far surpassing other diffusion-based methods which suffered from DDIM inversion errors or poor consistency.
• Robustness:
  • Demonstrated state-of-the-art robustness against distortion attacks (JPEG, blur, noise), regeneration attacks (diffusion purification, VAE compression), and adversarial attacks (averaging).
  • Achieved a TPR@1%FPR of 0.93 on average in the server scenario and 0.93 in the user scenario, comparable to the most robust methods (Gaussian Shading, RingID) but without their consistency trade-offs.
• Trade-off Analysis: By adjusting the mask radius, the method allows a controllable trade-off between visual fidelity and geometric robustness, consistently outperforming the Pareto frontier of existing methods.
• Generalization: The method was successfully generalized to Transformer-based models (FLUX), confirming the universality of the low-dimensional subspace hypothesis.

5. Significance

• Trustworthy AI: Shallow Diffuse provides a practical, deployable solution for tracing AI-generated content, crucial for combating misinformation and protecting intellectual property.
• Paradigm Shift: It moves away from the "noise injection" paradigm (which distorts generation) to a "subspace injection" paradigm, proving that watermarks can be hidden in the "blind spots" of the generation process.
• Practical Deployment: Being training-free and compatible with both server and user workflows makes it immediately applicable to existing commercial diffusion pipelines without computational overhead or model retraining.
• Theoretical Foundation: The work establishes a strong theoretical link between the intrinsic low-dimensionality of diffusion manifolds and the feasibility of robust, invisible watermarking.