OSI: One-step Inversion Excels in Extracting Diffusion Watermarks

TL;DR

This work tackles the inefficiency of extracting Gaussian Shading watermarks in diffusion-generated images by reframing watermark extraction as a one-shot sign classification problem. It introduces One-step Inversion (OSI), a learnable extractor initialized from the diffusion backbone and trained on synthesized noise–image pairs with a sign-classification objective, achieving ~20x faster extraction and higher accuracy with doubled payload capacity. The approach unifies a robust taxonomy via a communication-system perspective, showing OSI generalizes across diffusion backbones, schedulers, and cryptographic schemes. Empirically, OSI outperforms multi-step inversion across SD2.1/XL/3.5 models, maintains robustness under distortions and advanced attacks, and remains adaptable to broader diffusion watermarking settings with favorable amortized compute at scale.

Abstract

Watermarking is an important mechanism for provenance and copyright protection of diffusion-generated images. Training-free methods, exemplified by Gaussian Shading, embed watermarks into the initial noise of diffusion models with negligible impact on the quality of generated images. However, extracting this type of watermark typically requires multi-step diffusion inversion to obtain precise initial noise, which is computationally expensive and time-consuming. To address this issue, we propose One-step Inversion (OSI), a significantly faster and more accurate method for extracting Gaussian Shading style watermarks. OSI reformulates watermark extraction as a learnable sign classification problem, which eliminates the need for precise regression of the initial noise. Then, we initialize the OSI model from the diffusion backbone and finetune it on synthesized noise-image pairs with a sign classification objective. In this manner, the OSI model is able to accomplish the watermark extraction efficiently in only one step. Our OSI substantially outperforms the multi-step diffusion inversion method: it is 20x faster, achieves higher extraction accuracy, and doubles the watermark payload capacity. Extensive experiments across diverse schedulers, diffusion backbones, and cryptographic schemes consistently show improvements, demonstrating the generality of our OSI framework.

Figures (9)

• Figure 1: Overview of the proposed OSI method. For generation, (a) we embed Gaussian Shading style watermark, concealing watermark in the signs of initial latent. For extraction, (b) we adopt learnable one-step inversion (OSI) to directly predict the sign mask, while (c) previous methods first perform multi-step diffusion inversion to reconstruct continuous latent $\hat{z}_T$ and then extract the sign of $\hat{z}_T$.
• Figure 2: A communication system perspective on Gaussian Shading watermarking systems.
• Figure 3: Comparison of user number across repetition count settings. Numbers above the bars display the corresponding bit accuracies. The higher accuracy of OSI yields an exponentially larger addressable user number than GS gaussianshading_2024.
• Figure 4: Training loss comparison between different initializations of UNet parameters. Random initialization encounters training collapse and fails to converge.
• Figure 5: Visualization of different adversarial image distortions. (a) Clean watermarked image. (b). 80% random drop. (c) 60% random crop. (d) Resize to 25% and restore. (e) JPEG with $QF=25$. (f) Brightness with $f=6$. (g) Gaussian blur with $r=4$. (h) Gaussian Noise with $\mu=0,\sigma=0.05$. (i) Median blur with $k=7$. (j) Salt & Pepper noise with $p=0.05$.