Robust watermarking for diffusion models based on STDM and latent space fine-tuning

Abstract

Diffusion models (DMs) have demonstrated remarkable capabilities in generating high-quality images, but their potential for disseminating harmful misinformation raises significant concerns. Although reversible watermarking techniques can trace AI-generated images to their source models by embedding watermarks in the latent space, existing methods suffer from two critical drawbacks: (i) limited embedding capacity hinders unique model identification, and (ii) information loss during latent-space re-encoding compromises robustness, exacerbating the inherent trade-off between capacity and robustness. To address these limitations, we propose a novel watermarking framework based on Spread Transform Dither Modulation (STDM) that embeds watermarks into intermediate latent vectors during the diffusion process. Our approach operates in three key steps: (i) executing the standard diffusion process to obtain an intermediate latent vector, (ii) embedding watermarks into the mid-frequency DCT coefficients of this vector using ring-shaped STDM modulation, and (iii) completing the diffusion process to generate the final watermarked image. For watermark extraction, we employ a finely tuned VAE encoder to map the image back to latent space, followed by DDIM inversion and STDM-based extraction. Furthermore, we introduce a joint fine-tuning strategy that optimizes both the encoder and decoder of the diffusion model using watermarked latent vectors, significantly enhancing robustness. Experimental results demonstrate that our method achieves a maximum watermark embedding capacity of 256 bits while maintaining a high extraction accuracy of 98%. The proposed approach exhibits remarkable robustness against various attacks, with significant improvements over baseline methods.