OSI: One-step Inversion Excels in Extracting Diffusion Watermarks

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.

💡 Research Summary

The paper addresses a critical bottleneck in training‑free diffusion watermarking, namely the costly multi‑step diffusion inversion required to recover the initial noise that carries a Gaussian Shading (GS) watermark. While GS embeds a binary watermark into the sign pattern of the initial latent noise, extracting it traditionally demands 50 forward‑backward passes of the diffusion model to reconstruct the exact noise vector before taking its sign. This “sledgehammer” approach is computationally heavy and unnecessary because the watermark extraction task only needs the discrete sign information, not the precise continuous values.

To solve this, the authors propose One‑step Inversion (OSI), a learning‑based framework that reformulates watermark extraction as a binary sign classification problem. OSI consists of two trainable components: (1) an encoder Encψ that maps a watermarked image I back into the latent space, producing an estimated initial latent ẑ₀, and (2) a classifier OSIθ that takes ẑ₀ and outputs a probability p for the positive sign. The predicted mask is obtained by thresholding p at 0.5. Training optimizes a combined loss: binary cross‑entropy (BCE) between the ground‑truth sign label y (derived from the original sign mask m) and p, plus a mean‑squared error (MSE) term that forces ẑ₀ to stay close to the true initial noise z₀. The BCE term directly maximizes watermark accuracy, while the MSE term preserves latent consistency and prevents the classifier from over‑fitting to spurious patterns.

Training data are synthetically generated using the diffusion pipeline itself. For each sample, a random noise z_T ∼ N(0, I) is drawn; its signs become the ground‑truth mask m. The diffusion process then produces the 0‑step latent z₀ and the final image I after decoding. This yields a triplet (I, z₀, z_T) that provides both the input image and the exact sign labels without any manual annotation. The authors augment these samples with resizing, blurring, and cropping to improve robustness to common image distortions. Crucially, OSI is initialized from the pretrained diffusion backbone (for the classifier) and the VAE decoder (for the encoder), leveraging the model’s inherent ability to invert diffusion, which accelerates convergence.

Extensive experiments are conducted on Stable Diffusion 2.1 (the primary benchmark), as well as SDXL, SD 3.5, and DiT backbones, covering a variety of schedulers (DPM‑Solver, DDIM, Euler‑discrete). Across all settings, OSI achieves more than a 20× reduction in FLOPs and a 25× speedup in wall‑clock time (0.06 s per image versus ~1.5 s for the 50‑step baseline). Accuracy metrics show TPR@FPR=1e‑6 consistently at 0.998–1.000, bit‑accuracy improvements of 5–30 % over GS, and a doubling of payload rate (bits per image). Even under adversarial conditions—added noise, JPEG compression, color perturbations—OSI maintains high recovery rates, indicating that the learned inversion effectively increases the effective channel capacity of the diffusion process.

The authors also reinterpret the GS system through a communication‑theoretic lens, separating the “channel” (the diffusion generation‑inversion pipeline) from the “coding” (the encryption scheme such as repetition codes or pseudo‑random codes). Prior work focused on improving coding while leaving the channel fixed, thus being bounded by the channel’s intrinsic capacity. OSI, by training a more accurate inversion model, effectively expands the channel capacity, allowing higher payloads and lower error rates without altering the cryptographic coding. This perspective suggests future directions: combining OSI with advanced error‑correcting codes (e.g., LDPC, polar codes) could push the system close to Shannon limits.

In summary, OSI delivers a practical, scalable solution for extracting training‑free diffusion watermarks: (1) it reframes extraction as a discrete classification task, eliminating unnecessary regression; (2) it leverages pretrained diffusion components for rapid initialization; (3) it uses large synthetic datasets with precise sign labels; and (4) it demonstrates robustness and generality across models, schedulers, and attack scenarios. The method makes real‑time, low‑overhead watermark verification feasible for large‑scale image generation services, positioning OSI as a foundational building block for future copyright‑preserving diffusion ecosystems.